【SSL-RL】自监督强化学习：事后经验回放 (HER)算法

"""《HER 方法实现简单流程》
    时间：2024.10.
    作者：不去幼儿园
"""
import numpy as np

class HER:
    def __init__(self, env, agent, replay_buffer):
        self.env = env
        self.agent = agent
        self.replay_buffer = replay_buffer

    def generate_trajectory(self):
        state = self.env.reset()
        done = False
        trajectory = []
        while not done:
            action = self.agent.select_action(state)
            next_state, reward, done, _ = self.env.step(action)
            trajectory.append((state, action, reward, next_state))
            state = next_state
        return trajectory

    def hindsight_replay(self, trajectory):
        # 随机选择替代目标
        final_state = trajectory[-1][3]
        for (state, action, reward, next_state) in trajectory:
            # 替代目标回放
            goal = final_state  # 简单起见，使用最终状态作为替代目标
            new_reward = 1.0 if np.array_equal(next_state, goal) else 0.0
            self.replay_buffer.add(state, action, new_reward, next_state, goal)

    def train(self, num_episodes):
        for episode in range(num_episodes):
            trajectory = self.generate_trajectory()
            self.hindsight_replay(trajectory)
            self.agent.train(self.replay_buffer)

"""《 HER+DDPG 方法实现简单流程》
    时间：2024.10.
    作者：不去幼儿园
"""
import gym
import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import random
from collections import deque

# Hyperparameters
GAMMA = 0.99
TAU = 0.005
ACTOR_LR = 1e-4
CRITIC_LR = 1e-3
BUFFER_SIZE = 1000000
BATCH_SIZE = 64
HER_K = 4  # HER 技术的 k 值

# Actor network
class Actor(nn.Module):
	def __init__(self, state_dim, action_dim):
		super(Actor, self).__init__()
		self.fc1 = nn.Linear(state_dim, 256)
		self.fc2 = nn.Linear(256, 256)
		self.fc3 = nn.Linear(256, action_dim)

	def forward(self, state):
		x = torch.relu(self.fc1(state))
		x = torch.relu(self.fc2(x))
		action = torch.tanh(self.fc3(x))  # 使用 tanh 激活函数将输出限制在 [-1, 1]
		return action

class Critic(nn.Module):
    def __init__(self, state_dim, action_dim):
        super(Critic, self).__init__()
        # 调整 fc1 的输入维度为 5
        self.fc1 = nn.Linear(5, 256)  # 输入维度为 5 (4 for state + 1 for action)
        self.fc2 = nn.Linear(256, 256)
        self.fc3 = nn.Linear(256, 1)

    def forward(self, state, action):
        # 将 state 和 action 拼接，确保维度是 (256, 5)
        x = torch.cat([state, action], dim=1)
        x = torch.relu(self.fc1(x))  # fc1 输入维度现在匹配输入
        x = torch.relu(self.fc2(x))
        x = self.fc3(x)
        return x

# Replay Buffer with HER
class ReplayBuffer:
	def __init__(self, buffer_size, state_dim, action_dim, goal_dim):
		self.buffer = deque(maxlen=buffer_size)
		self.state_dim = state_dim
		self.action_dim = action_dim
		self.goal_dim = goal_dim

	def add(self, state, action, reward, next_state, done, goal):
		# goal 是一个包含两个元素的元组，提取第一个元素作为实际的 goal
		goal = goal[0]  # 仅提取数组部分，忽略空字典
		# print(f"Processed Goal: {goal}")  # 调试输出处理后的 goal

		# 确保所有数据都是 NumPy 数组并具有一致的形状
		state = np.array(state).reshape(-1)
		action = np.array(action).reshape(self.action_dim)
		next_state = np.array(next_state).reshape(-1)
		goal = np.array(goal).reshape(-1)  # 保持 goal 为一维数组

		self.buffer.append((state, action, reward, next_state, done, goal))

	def sample(self, batch_size):
		batch = random.sample(self.buffer, batch_size)
		states, actions, rewards, next_states, dones, goals = zip(*batch)
		# 确保每个维度的数组形状一致
		return (
			np.array(states),
			np.array(actions),
			np.array(rewards),
			np.array(next_states),
			np.array(dones),
			np.array(goals)
		)

	def her_sample(self, batch_size, k):
		states, actions, rewards, next_states, dones, goals = self.sample(batch_size)
		her_states, her_actions, her_rewards, her_next_states, her_dones, her_goals = [], [], [], [], [], []

		for i in range(batch_size):
			future_idx = np.random.randint(i, len(self.buffer), size=k)
			for idx in future_idx:
				new_goal = self.buffer[idx][3]  # 使用未来状态作为新的 goal
				her_states.append(states[i])
				her_actions.append(actions[i])
				her_rewards.append(-np.linalg.norm(next_states[i] - new_goal))  # 使用新目标计算奖励
				her_next_states.append(next_states[i])
				her_dones.append(dones[i])  # 确保 dones 被正确加入
				her_goals.append(new_goal)

		# 返回 6 个值，而不是 5 个
		return (
			np.array(her_states),
			np.array(her_actions),
			np.array(her_rewards),
			np.array(her_next_states),
			np.array(her_dones),
			np.array(her_goals)
		)

	def __len__(self):
		return len(self.buffer)

# DDPG agent
class DDPGAgent:
	def __init__(self, state_dim, action_dim, goal_dim):
		self.actor = Actor(state_dim, action_dim).float()
		self.actor_target = Actor(state_dim, action_dim).float()
		self.critic = Critic(state_dim + goal_dim, action_dim).float()
		self.critic_target = Critic(state_dim + goal_dim, action_dim).float()

		self.actor_optimizer = optim.Adam(self.actor.parameters(), lr=ACTOR_LR)
		self.critic_optimizer = optim.Adam(self.critic.parameters(), lr=CRITIC_LR)

		self.replay_buffer = ReplayBuffer(BUFFER_SIZE, state_dim, action_dim, goal_dim)

	def update(self, batch_size):
		if len(self.replay_buffer) < batch_size:
			return

		# 从经验池中采样
		states, actions, rewards, next_states, dones, goals = self.replay_buffer.her_sample(batch_size, HER_K)

		states = torch.FloatTensor(states)
		actions = torch.FloatTensor(actions)
		rewards = torch.FloatTensor(rewards).unsqueeze(1)
		next_states = torch.FloatTensor(next_states)
		dones = torch.FloatTensor(dones).unsqueeze(1)
		goals = torch.FloatTensor(goals)

		# Actor 更新
		pred_actions = self.actor(states)
		# print(f"states shape: {states.shape}, pred_actions shape: {pred_actions.shape}")
		actor_loss = -self.critic(states, pred_actions).mean()

		self.actor_optimizer.zero_grad()
		actor_loss.backward()
		self.actor_optimizer.step()

		# Critic 更新
		target_actions = self.actor_target(next_states)
		target_q_values = rewards + GAMMA * self.critic_target(next_states, target_actions) * (1 - dones)
		critic_loss = (self.critic(states, actions) - target_q_values.detach()).pow(2).mean()
		self.critic_optimizer.zero_grad()
		critic_loss.backward()
		self.critic_optimizer.step()

		# 软更新目标网络
		self.soft_update(self.actor, self.actor_target)
		self.soft_update(self.critic, self.critic_target)

	def soft_update(self, net, net_target):
		for param, param_target in zip(net.parameters(), net_target.parameters()):
			param_target.data.copy_(TAU * param.data + (1 - TAU) * param_target.data)

	def act(self, state):
		state = torch.FloatTensor(state).unsqueeze(0)
		action = self.actor(state).detach().numpy()[0]
		return action

	def add_to_buffer(self, state, action, reward, next_state, done, goal):
		self.replay_buffer.add(state, action, reward, next_state, done, goal)

# CartPole action mapping to continuous space [-1, 1]
def action_mapper(action):
	return np.array([action * 2 - 1])  # 将 CartPole 的动作 0/1 映射到 [-1, 1]


# CartPole environment with continuous action space
class CartPoleContinuousWrapper(gym.ActionWrapper):
	def __init__(self, env):
		super(CartPoleContinuousWrapper, self).__init__(env)

	def action(self, action):
		# 将连续动作转换为离散动作
		return int(action[0] > 0)

	def reverse_action(self, action):
		return np.array([1.0 if action == 1 else -1.0])

# Main training loop
env = CartPoleContinuousWrapper(gym.make('CartPole-v1'))
agent = DDPGAgent(state_dim=env.observation_space.shape[0], action_dim=1, goal_dim=env.observation_space.shape[0])

NUM_EPISODES = 100
MAX_STEPS = 200

for episode in range(NUM_EPISODES):
	# 兼容不同版本的 Gym, env.reset() 只取需要的 state
	result = env.reset()
	state = result[0] if isinstance(result, tuple) else result  # 兼容不同版本

	goal = env.reset()  # 使用初始状态作为目标
	total_reward = 0

	for step in range(MAX_STEPS):
		action = agent.act(state)
		result = env.step(action_mapper(action))

		# 兼容不同版本的 Gym, env.step() 取前四个返回值
		next_state, reward, done = result[:3]

		agent.add_to_buffer(state, action, reward, next_state, done, goal)
		agent.update(BATCH_SIZE)

		state = next_state
		total_reward += reward

		if done:
			break

	print(f"Episode {episode + 1}, Total Reward: {total_reward}")

# 测试程序
def test_policy(env, policy, num_episodes=10):
	for episode in range(num_episodes):
		result = env.reset()
		state = result[0] if isinstance(result, tuple) else result  # 兼容不同版本
		total_reward = 0
		for step in range(MAX_STEPS):
			action = agent.act(state)
			result = env.step(action_mapper(action))
			# 兼容不同版本的 Gym, env.step() 取前四个返回值
			next_state, reward, done = result[:3]
			total_reward += reward
			if done:
				break
		print(f"测试第 {episode + 1} 集，总奖励: {total_reward}")


# 设置测试CartPole环境
env = CartPoleContinuousWrapper(gym.make('CartPole-v1', render_mode="human"))
test_policy(env, agent)