【Hierarchical RL】隐空间分层强化学习（HRL-LS ）算法

"""《HRL-LS算法实现》
    时间：2024.10.13
    环境：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
LATENT_DIM = 2  # 潜在空间维度
HIGH_LEVEL_UPDATE_FREQUENCY = 10  # 高层更新频率
LOW_LEVEL_UPDATE_FREQUENCY = 1  # 低层更新频率

# Encoder 网络，将状态映射到潜在空间
class Encoder(nn.Module):
	def __init__(self, state_dim, latent_dim):
		super(Encoder, self).__init__()
		self.fc1 = nn.Linear(state_dim, 128)
		self.fc2 = nn.Linear(128, latent_dim)

	def forward(self, state):
		x = torch.relu(self.fc1(state))
		latent = self.fc2(x)
		return latent

# Decoder 网络，将潜在空间中的目标解码为实际目标
class Decoder(nn.Module):
	def __init__(self, latent_dim, state_dim):
		super(Decoder, self).__init__()
		self.fc1 = nn.Linear(latent_dim, 128)
		self.fc2 = nn.Linear(128, state_dim)

	def forward(self, latent_goal):
		x = torch.relu(self.fc1(latent_goal))
		goal = self.fc2(x)
		return goal

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

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

# Hierarchical Reinforcement Learning with Latent Space 智能体
class HRLLSAgent:
	def __init__(self, state_dim, action_dim, latent_dim):
		self.encoder = Encoder(state_dim, latent_dim)  # 编码器
		self.decoder = Decoder(latent_dim, state_dim)  # 解码器
		self.high_level_net = ActorCritic(latent_dim, latent_dim)  # 高层策略在潜在空间中操作
		self.low_level_net = ActorCritic(state_dim + state_dim, 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 encode_state(self, state):
		state = torch.FloatTensor(state).unsqueeze(0)
		latent_state = self.encoder(state)
		return latent_state

	def decode_goal(self, latent_goal):
		# 检查 latent_goal 的大小是否为 LATENT_DIM
		if latent_goal.numel() != LATENT_DIM:  # `numel()` 返回张量中元素的总数
			raise ValueError(f"Invalid shape for latent_goal: {latent_goal.size()}, expected size: {LATENT_DIM}")

		# 调整 latent_goal 的形状为 (1, LATENT_DIM)
		latent_goal = latent_goal.view(1, LATENT_DIM)
		return self.decoder(latent_goal)

	def select_high_level_goal(self, latent_state, epsilon):
		if random.random() < epsilon:
			return torch.randn(1, LATENT_DIM)  # 随机选择潜在空间目标，确保大小正确
		else:
			action_probs, _ = self.high_level_net(latent_state)
			goal_index = torch.argmax(action_probs).unsqueeze(0)  # 找到目标索引

			# 将目标索引扩展到潜在空间维度，确保形状为 [1, LATENT_DIM]
			latent_goal = torch.zeros(1, LATENT_DIM)
			latent_goal[0, goal_index] = 1.0  # 将目标索引位置设置为 1，其他为 0
			return latent_goal

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

	def update_high_level(self, latent_state, latent_goal, reward, next_latent_state):
		_, state_value = self.high_level_net(latent_state)
		_, next_state_value = self.high_level_net(next_latent_state)
		next_state_value = next_state_value.detach()

		target_value = reward + GAMMA * next_state_value
		loss_critic = nn.functional.mse_loss(state_value, target_value)

		# 获取动作概率
		action_probs, _ = self.high_level_net(latent_state)

		# 确保 latent_goal 在动作空间的范围内
		latent_goal = latent_goal.clamp(0, action_probs.size(1) - 1)

		# 计算 log_prob，避免索引超出范围
		log_prob = torch.log(action_probs[0, latent_goal.long()])

		advantage = target_value - state_value
		loss_actor = -log_prob * advantage

		loss = (loss_critic + loss_actor).mean()

		self.high_level_optimizer.zero_grad()
		loss.backward(retain_graph=True)  # 保持计算图
		self.high_level_optimizer.step()

	def update_low_level(self, state, goal, action, reward, next_state):
		state_goal = torch.cat((torch.FloatTensor(state).unsqueeze(0), goal), dim=-1)
		next_state_goal = torch.cat((torch.FloatTensor(next_state).unsqueeze(0), goal), dim=-1)

		_, state_value = self.low_level_net(state_goal)
		_, next_state_value = self.low_level_net(next_state_goal)
		next_state_value = next_state_value.detach()

		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(retain_graph=True)  # 保留计算图
		self.low_level_optimizer.step()

	def train(self, env, num_episodes):
		for episode in range(num_episodes):
			state, _ = env.reset()  # 修改后的reset返回值
			latent_state = self.encode_state(state)  # 将状态映射到潜在空间
			latent_goal = self.select_high_level_goal(latent_state, self.epsilon)  # 高层选择潜在空间目标
			goal = self.decode_goal(latent_goal)  # 解码为实际空间中的目标
			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)

				# 更新高层目标
				if steps % HIGH_LEVEL_UPDATE_FREQUENCY == 0:
					next_latent_state = self.encode_state(next_state)
					new_latent_goal = self.select_high_level_goal(next_latent_state, self.epsilon)
					self.update_high_level(latent_state, latent_goal, reward, next_latent_state)
					latent_goal = new_latent_goal
					goal = self.decode_goal(latent_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}")

# 测试 Hierarchical Reinforcement Learning with Latent Space 智能体并显示动画
def test_hrlls_agent(agent, env, num_episodes=5):
	for episode in range(num_episodes):
		state, _ = env.reset()  # 修改后的reset返回值
		latent_state = agent.encode_state(state)  # 将状态映射到潜在空间
		latent_goal = agent.select_high_level_goal(latent_state, epsilon=0.0)  # 高层选择潜在空间目标
		goal = agent.decode_goal(latent_goal)  # 解码为实际目标
		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"Episode {episode + 1}: Total Reward: {total_reward}")

# 运行测试
if __name__ == "__main__":
	env = gym.make('CartPole-v1')
	state_dim = env.observation_space.shape[0]
	action_dim = env.action_space.n

	agent = HRLLSAgent(state_dim, action_dim, LATENT_DIM)
	agent.train(env, NUM_EPISODES)
	env = gym.make('CartPole-v1', render_mode="human")
	test_hrlls_agent(agent, env)