【Hierarchical RL】分层深度Q网络（Hierarchical-DQN）算法

"""《Hierarchical-DQN 实现》
    时间：2024.10.07
    环境：CartPole
    作者：不去幼儿园
"""
import gym
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
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  # 高层更新频率
LOW_LEVEL_UPDATE_FREQUENCY = 1  # 低层更新频率


# Q网络
class QNetwork(nn.Module):
	def __init__(self, input_dim, output_dim):
		super(QNetwork, self).__init__()
		self.fc1 = nn.Linear(input_dim, 128)
		self.fc2 = nn.Linear(128, output_dim)

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


# Hierarchical-DQN 智能体
class HierarchicalDQNAgent:
	def __init__(self, state_dim, action_dim, goal_dim):
		self.high_level_net = QNetwork(state_dim, goal_dim)  # 高层 Q 网络
		self.low_level_net = QNetwork(state_dim + 1, action_dim)  # 低层 Q 网络
		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)
			q_values = self.high_level_net(state)
			return torch.argmax(q_values).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)
			q_values = self.low_level_net(state_goal)
			return torch.argmax(q_values).item()

	def update_high_level(self, state, goal, reward, next_state):
		state = torch.FloatTensor(state).unsqueeze(0)
		next_state = torch.FloatTensor(next_state).unsqueeze(0)
		goal = torch.tensor([goal], dtype=torch.float32)

		q_values = self.high_level_net(state)
		next_q_values = self.high_level_net(next_state).detach()

		# 确保 target_q 维度匹配
		target_q = torch.tensor([reward + GAMMA * torch.max(next_q_values)])
		loss = nn.functional.mse_loss(q_values[0, goal.long()], target_q)

		self.high_level_optimizer.zero_grad()
		loss.backward()
		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), torch.FloatTensor([[goal]])), dim=-1)
		next_state_goal = torch.cat((torch.FloatTensor(next_state).unsqueeze(0), torch.FloatTensor([[goal]])), dim=-1)

		q_values = self.low_level_net(state_goal)
		next_q_values = self.low_level_net(next_state_goal).detach()

		target_q = reward + GAMMA * torch.max(next_q_values)
		loss = nn.functional.mse_loss(q_values[0, action], target_q)

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

	def train(self, env, num_episodes):
		goal_dim = self.high_level_net.fc2.out_features

		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}")

# 测试 Hierarchical-DQN 智能体并显示动画
def test_hdqn_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()

# 测试智能体并显示动画
test_hdqn_agent(agent, env)