"""《MAXQ分解算法实现》
    时间：2024.10.03
    环境：CartPole 
    作者：不去幼儿园
"""
import gym
import numpy as np
import random

# 定义超参数
GAMMA = 0.99
LEARNING_RATE = 0.1
EPSILON_DECAY = 0.995
MIN_EPSILON = 0.1
NUM_EPISODES = 500

# 定义 MAXQ 子任务节点
class MAXQNode:
    def __init__(self, num_actions, is_primitive=False):
        self.num_actions = num_actions
        self.is_primitive = is_primitive  # 是否为原子任务
        self.q_values = {}  # Q 值字典，(state, action) -> Q 值

    def get_q(self, state, action):
        state = tuple(state)  # 将 state 转换为元组
        if (state, action) not in self.q_values:
            self.q_values[(state, action)] = 0.0  # 初始化 Q 值
        return self.q_values[(state, action)]

    def set_q(self, state, action, value):
        state = tuple(state)  # 将 state 转换为元组
        self.q_values[(state, action)] = value

# 创建 MAXQ 分解的顶层任务
class MAXQTask:
    def __init__(self, num_actions):
        self.num_actions = num_actions
        self.subtasks = []
        self.root = MAXQNode(num_actions)  # 顶层任务节点

    def add_subtask(self, subtask):
        self.subtasks.append(subtask)

# 定义智能体
class MAXQAgent:
    def __init__(self, env):
        self.env = env
        self.epsilon = 1.0
        self.maxq_root = MAXQTask(env.action_space.n)  # CartPole 的动作数为 2

        # 添加子任务
        self.pickup_subtask = MAXQNode(env.action_space.n, is_primitive=True)  # 装载任务
        self.dropoff_subtask = MAXQNode(env.action_space.n, is_primitive=True)  # 卸载任务
        self.navigate_subtask = MAXQNode(env.action_space.n)  # 导航任务

        # 将子任务添加到根节点
        self.maxq_root.add_subtask(self.pickup_subtask)
        self.maxq_root.add_subtask(self.dropoff_subtask)
        self.maxq_root.add_subtask(self.navigate_subtask)

    def select_action(self, state, subtask, epsilon):
        # ε-贪婪策略选择动作
        if random.random() < epsilon:
            return self.env.action_space.sample()  # 探索
        else:
            q_values = [subtask.get_q(state, action) for action in range(subtask.num_actions)]
            return np.argmax(q_values)  # 利用当前策略

    def update_q(self, state, subtask, action, reward, next_state):
        # 确保传入的是 MAXQNode 而不是 MAXQTask
        if not isinstance(subtask, MAXQNode):
            subtask = subtask.root  # 传入的是 MAXQTask，则使用其根节点
        next_q_values = [subtask.get_q(next_state, a) for a in range(subtask.num_actions)]
        max_next_q = max(next_q_values)

        # Q 值更新
        current_q = subtask.get_q(state, action)
        new_q = current_q + LEARNING_RATE * (reward + GAMMA * max_next_q - current_q)
        subtask.set_q(state, action, new_q)

    def train(self, num_episodes):
        for episode in range(num_episodes):
            state,_ = self.env.reset()
            done = False
            total_reward = 0

            while not done:
                # 确保传递的是 MAXQNode
                action = self.select_action(state, self.maxq_root.root, self.epsilon)
                next_state, reward, done, _, __ = self.env.step(action)

                # 这里传递的是 MAXQNode，而不是 MAXQTask
                self.update_q(state, self.maxq_root.root, action, reward, next_state)
                state = next_state
                total_reward += reward

            # 逐渐减少 epsilon 以减少探索
            self.epsilon = max(MIN_EPSILON, self.epsilon * EPSILON_DECAY)
            print(f"Episode {episode + 1}: Total Reward: {total_reward}")

# 创建 CartPole 环境并训练智能体
env = gym.make('CartPole-v1')
agent = MAXQAgent(env)
agent.train(NUM_EPISODES)
env.close()

import gym
import torch

# 测试 MAXQ 智能体并显示动画
def test_maxq_agent(agent, env, num_episodes=5):
    for episode in range(num_episodes):
        state, _ = env.reset()
        done = False
        total_reward = 0
        env.render()  # 初始化渲染

        while not done:
            env.render()  # 显示动画
            action = agent.select_action(state, agent.maxq_root.root, epsilon=0.0)  # 使用已学策略选择动作
            next_state, reward, done, _, __ = env.step(action)
            state = next_state
            total_reward += reward

        print(f"测试 Episode {episode + 1}: Total Reward: {total_reward}")
    env.close()

# 创建环境并调用测试函数
env = gym.make('CartPole-v1', render_mode='human')
test_maxq_agent(agent, env)