【Hierarchical RL】动态分层强化学习（DHRL）算法

"""《DHRL算法实现》
    时间：2024.10
    环境：CartPole
    作者：不去幼儿园
"""
import gym
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
import random


# 高层控制器网络
class ControllerNet(nn.Module):
	def __init__(self, state_dim, num_subtasks):
		super(ControllerNet, self).__init__()
		self.fc = nn.Sequential(
			nn.Linear(state_dim, 128),
			nn.ReLU(),
			nn.Linear(128, num_subtasks)
		)

	def forward(self, state):
		return self.fc(state)


# 子任务策略网络
class SubtaskPolicyNet(nn.Module):
	def __init__(self, state_dim, action_dim):
		super(SubtaskPolicyNet, self).__init__()
		self.fc = nn.Sequential(
			nn.Linear(state_dim, 128),
			nn.ReLU(),
			nn.Linear(128, action_dim)
		)

	def forward(self, state):
		return self.fc(state)


# 低层子任务类
class Subtask:
	def __init__(self, policy_net):
		self.policy_net = policy_net

	def select_action(self, state):
		state = torch.FloatTensor(state).unsqueeze(0)
		action_probs = self.policy_net(state)
		action = torch.argmax(action_probs).item()
		return action


# 动态分层强化学习智能体
class DHRLAgent:
	def __init__(self, state_dim, action_dim, num_subtasks):
		self.controller = ControllerNet(state_dim, num_subtasks)
		self.subtasks = [Subtask(SubtaskPolicyNet(state_dim, action_dim)) for _ in range(num_subtasks)]
		self.controller_optimizer = optim.Adam(self.controller.parameters(), lr=0.001)
		self.subtask_optimizers = [optim.Adam(subtask.policy_net.parameters(), lr=0.001) for subtask in self.subtasks]
		self.gamma = 0.99

	def select_subtask(self, state):
		state_tensor = torch.FloatTensor(state).unsqueeze(0)
		subtask_probs = self.controller(state_tensor)
		selected_subtask = torch.argmax(subtask_probs).item()
		return selected_subtask

	def train_subtask(self, subtask_idx, state, action, reward, next_state, done):
		optimizer = self.subtask_optimizers[subtask_idx]
		subtask = self.subtasks[subtask_idx]
		state_tensor = torch.FloatTensor(state).unsqueeze(0)
		next_state_tensor = torch.FloatTensor(next_state).unsqueeze(0)

		# 计算当前 Q 值
		q_values = subtask.policy_net(state_tensor)
		next_q_values = subtask.policy_net(next_state_tensor)

		target_q_value = reward + (1 - done) * self.gamma * torch.max(next_q_values).item()
		loss = nn.functional.mse_loss(q_values[0, action], torch.tensor(target_q_value))

		optimizer.zero_grad()
		loss.backward()
		optimizer.step()

	def train_controller(self, state, subtask_idx, reward, next_state, done):
		state_tensor = torch.FloatTensor(state).unsqueeze(0)
		next_state_tensor = torch.FloatTensor(next_state).unsqueeze(0)

		# 计算高层控制器 Q 值
		controller_q_values = self.controller(state_tensor)
		next_controller_q_values = self.controller(next_state_tensor)

		target_q_value = reward + (1 - done) * self.gamma * torch.max(next_controller_q_values).item()
		loss = nn.functional.mse_loss(controller_q_values[0, subtask_idx], torch.tensor(target_q_value))

		self.controller_optimizer.zero_grad()
		loss.backward()
		self.controller_optimizer.step()


# 训练函数
def train_dhrl(env, agent, num_episodes=1000):
	for episode in range(num_episodes):
		state, _ = env.reset()
		total_reward = 0

		for t in range(200):  # 每个 episode 最多 200 步
			subtask_idx = agent.select_subtask(state)  # 选择当前子任务
			action = agent.subtasks[subtask_idx].select_action(state)  # 选择动作

			next_state, reward, done, _, _ = env.step(action)
			total_reward += reward

			# 训练子任务
			agent.train_subtask(subtask_idx, state, action, reward, next_state, done)
			# 训练高层控制器
			agent.train_controller(state, subtask_idx, reward, next_state, done)

			state = next_state

			if done:
				break

		print(f'Episode {episode}: Total Reward: {total_reward}')


# 环境设置
env = gym.make('CartPole-v1')
state_dim = env.observation_space.shape[0]
action_dim = env.action_space.n

# 实例化 DHRL 智能体并训练
num_subtasks = 2  # 两个子任务，左右移动
agent = DHRLAgent(state_dim, action_dim, num_subtasks)

# 训练智能体
train_dhrl(env, agent, num_episodes=1000)

# 测试函数
def test_dhrl(env, agent, num_episodes=5):
	for episode in range(num_episodes):
		state, _ = env.reset()
		total_reward = 0

		for t in range(200):
			env.render()  # 显示动画

			subtask_idx = agent.select_subtask(state)
			action = agent.subtasks[subtask_idx].select_action(state)

			next_state, reward, done, _, _ = env.step(action)
			total_reward += reward

			state = next_state

			if done:
				break

		print(f'Test Episode {episode}: Total Reward: {total_reward}')

	env.close()


# 测试代理并显示动画
env = gym.make('CartPole-v1', render_mode="human")
test_dhrl(env, agent, num_episodes=5)