pip install gym torch numpy

"""《Hierarchical Actor-Critic (HAC) 算法项目》
    时间：2024.10.11
    环境：CartPole
    作者：不去幼儿园
"""
import gym
import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import random

# 超参数
GAMMA = 0.99
LEARNING_RATE = 0.001
EPSILON_DECAY = 0.995
MIN_EPSILON = 0.1
NUM_EPISODES = 500
HIGH_LEVEL_UPDATE_FREQUENCY = 10  # 高层更新频率

# Actor-Critic 网络
class ActorCritic(nn.Module):
	def __init__(self, state_dim, action_dim):
		super(ActorCritic, self).__init__()
		self.actor = nn.Sequential(
			nn.Linear(state_dim, 128),
			nn.ReLU(),
			nn.Linear(128, action_dim),
			nn.Softmax(dim=-1)
		)
		self.critic = nn.Sequential(
			nn.Linear(state_dim, 128),
			nn.ReLU(),
			nn.Linear(128, 1)
		)

	def forward(self, state):
		action_probs = self.actor(state)
		state_value = self.critic(state)
		return action_probs, state_value

# Hierarchical Actor-Critic 智能体
class HACAgent:
	def __init__(self, state_dim, action_dim, goal_dim):
		# 高层 Actor-Critic
		self.high_level_net = ActorCritic(state_dim, goal_dim)
		# 低层 Actor-Critic
		self.low_level_net = ActorCritic(state_dim + 1, action_dim)
		# 优化器
		self.high_level_optimizer = optim.Adam(self.high_level_net.parameters(), lr=LEARNING_RATE)
		self.low_level_optimizer = optim.Adam(self.low_level_net.parameters(), lr=LEARNING_RATE)
		self.epsilon = 1.0

	def select_high_level_goal(self, state, epsilon):
		if random.random() < epsilon:
			return random.choice([0, 1])  # 随机选择目标
		else:
			state = torch.FloatTensor(state).unsqueeze(0)
			action_probs, _ = self.high_level_net(state)
			return torch.argmax(action_probs).item()

	def select_low_level_action(self, state, goal, epsilon):
		if random.random() < epsilon:
			return random.choice([0, 1])  # 随机选择动作
		else:
			state_goal = torch.cat((torch.FloatTensor(state).unsqueeze(0), torch.FloatTensor([[goal]])), dim=-1)
			action_probs, _ = self.low_level_net(state_goal)
			return torch.argmax(action_probs).item()

	# 修改后的 update_high_level 方法
	def update_high_level(self, state, goal, reward, next_state):
		state = torch.FloatTensor(state).unsqueeze(0)
		next_state = torch.FloatTensor(next_state).unsqueeze(0)

		# 分别提取动作概率和状态值
		_, state_value = self.high_level_net(state)
		_, next_state_value = self.high_level_net(next_state)

		# 确保仅对状态值使用 detach
		next_state_value = next_state_value.detach()

		# 计算TD目标值
		target_value = reward + GAMMA * next_state_value
		loss_critic = nn.functional.mse_loss(state_value, target_value)

		# 计算高层策略损失
		action_probs, _ = self.high_level_net(state)
		log_prob = torch.log(action_probs[0, goal])
		advantage = target_value - state_value
		loss_actor = -log_prob * advantage

		loss = loss_critic + loss_actor
		self.high_level_optimizer.zero_grad()
		loss.backward()
		self.high_level_optimizer.step()

	# 修改后的 update_low_level 方法
	def update_low_level(self, state, goal, action, reward, next_state):
		state_goal = torch.cat((torch.FloatTensor(state).unsqueeze(0), torch.FloatTensor([[goal]])), dim=-1)
		next_state_goal = torch.cat((torch.FloatTensor(next_state).unsqueeze(0), torch.FloatTensor([[goal]])), dim=-1)

		# 分别提取动作概率和状态值
		_, state_value = self.low_level_net(state_goal)
		_, next_state_value = self.low_level_net(next_state_goal)

		# 确保仅对状态值使用 detach
		next_state_value = next_state_value.detach()

		# 计算TD目标值
		target_value = reward + GAMMA * next_state_value
		loss_critic = nn.functional.mse_loss(state_value, target_value)

		# 计算低层策略损失
		action_probs, _ = self.low_level_net(state_goal)
		log_prob = torch.log(action_probs[0, action])
		advantage = target_value - state_value
		loss_actor = -log_prob * advantage

		loss = loss_critic + loss_actor
		self.low_level_optimizer.zero_grad()
		loss.backward()
		self.low_level_optimizer.step()

	def train(self, env, num_episodes):
		goal_dim = 2  # 目标维度设定为2，直接设定

		for episode in range(num_episodes):
			state, _ = env.reset()  # 修改后的reset返回值
			goal = self.select_high_level_goal(state, self.epsilon)  # 高层选择目标
			done = False
			episode_reward = 0
			steps = 0

			while not done:
				steps += 1
				action = self.select_low_level_action(state, goal, self.epsilon)  # 低层选择动作
				next_state, reward, done, _, _ = env.step(action)  # 修改后的step返回值

				# 更新低层
				self.update_low_level(state, goal, action, reward, next_state)

				# 每隔 HIGH_LEVEL_UPDATE_FREQUENCY 更新一次高层
				if steps % HIGH_LEVEL_UPDATE_FREQUENCY == 0:
					new_goal = self.select_high_level_goal(next_state, self.epsilon)
					self.update_high_level(state, goal, reward, next_state)
					goal = new_goal

				state = next_state
				episode_reward += reward

			self.epsilon = max(MIN_EPSILON, self.epsilon * EPSILON_DECAY)
			print(f"Episode {episode + 1}: Total Reward: {episode_reward}")


# 创建 CartPole 环境用于训练，不渲染动画
env = gym.make('CartPole-v1')  # 不设置 render_mod
state_dim = env.observation_space.shape[0]
action_dim = env.action_space.n
goal_dim = 2  # 目标维度设为2（简单划分目标）
agent = HACAgent(state_dim, action_dim, goal_dim)
agent.train(env, NUM_EPISODES)

# 测试 Hierarchical Actor-Critic 智能体并显示动画
def test_hac_agent(agent, env, num_episodes=5):
	for episode in range(num_episodes):
		state, _ = env.reset()  # 修改后的reset返回值
		goal = agent.select_high_level_goal(state, epsilon=0.0)  # 高层选择目标
		done = False
		total_reward = 0
		env.render()

		while not done:
			env.render()
			action = agent.select_low_level_action(state, goal, epsilon=0.0)  # 低层选择动作
			next_state, reward, done, _, _ = env.step(action)  # 修改后的step返回值
			state = next_state
			total_reward += reward

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


# 创建 CartPole 环境用于测试，显示动画
env = gym.make('CartPole-v1', render_mode='human')
test_hac_agent(agent, env)