元强化学习

• 简单直观：通过RNN直接建模任务历史，实现策略的元学习，易于理解和实现。
• 端到端训练：将任务推断与策略更新整合在RNN中，简化训练流程。
• 缺点：
  • 采样效率低：依赖在线策略，需大量交互数据，训练成本高。
  • 任务推断能力有限：RNN难以处理长序列或复杂任务依赖关系，推断能力受限。
• 应用场景：
  • 简单任务序列学习：如快速适应不同游戏规则的小规模任务。
  • 动态环境决策：在任务规则变化不剧烈的场景中，实现策略的动态调整。

• 核心思想：PEARL是一种基于actor-critic框架的元强化学习方法，特别设计用于处理部分可观测环境中的问题。将元学习问题建模为概率推断，通过推断任务潜在表示调整策略。
• 算法流程：通过学习一个环境编码器，提取当前环境的关键特征，并将这些特征用于指导演员网络如何调整策略以最适应新环境。
• 实现代码

import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.distributions import Normal
import gym
from collections import deque, namedtuple
import random
import matplotlib.pyplot as plt# 设置设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")# 经验回放缓冲区
class ReplayBuffer:def __init__(self, capacity):self.buffer = deque(maxlen=capacity)self.Transition = namedtuple('Transition',('state', 'action', 'reward', 'next_state', 'done'))def add(self, state, action, reward, next_state, done):self.buffer.append(self.Transition(state, action, reward, next_state, done))def sample(self, batch_size):transitions = random.sample(self.buffer, batch_size)batch = self.Transition(*zip(*transitions))states = torch.FloatTensor(batch.state).to(device)actions = torch.FloatTensor(batch.action).to(device)rewards = torch.FloatTensor(batch.reward).unsqueeze(1).to(device)next_states = torch.FloatTensor(batch.next_state).to(device)dones = torch.FloatTensor(batch.done).unsqueeze(1).to(device)return states, actions, rewards, next_states, donesdef __len__(self):return len(self.buffer)# 编码器网络 - 学习任务的潜在表示
class Encoder(nn.Module):def __init__(self, state_dim, action_dim, latent_dim, hidden_dim):# state_dim: 状态的维度。它表示每个状态向量的长度。# action_dim: 动作的维度。它表示每个动作向量的长度。# latent_dim: 潜在表示的维度。它表示编码器输出的潜在向量的长度。# hidden_dim: 隐藏层的维度。它表示编码器中隐藏层的神经元数量。super(Encoder, self).__init__()self.encoder = nn.Sequential(nn.Linear(state_dim + action_dim + 1, hidden_dim),nn.ReLU(),nn.Linear(hidden_dim, hidden_dim),nn.ReLU(),nn.Linear(hidden_dim, hidden_dim),nn.ReLU())self.mean = nn.Linear(hidden_dim, latent_dim)  #用于从隐藏表示中计算潜在表示的均值。self.log_std = nn.Linear(hidden_dim, latent_dim)  #用于从隐藏表示中计算潜在表示的对数标准差。self.latent_dim = latent_dimdef forward(self, context):# context: [batch_size, seq_len, state_dim + action_dim + 1]# batch_size: 批处理大小。# seq_len: 序列长度，表示每个任务上下文的长度。# state_dim + action_dim + 1: 每个时间步的状态、动作和奖励的拼接。batch_size, seq_len, _ = context.shape# 处理序列数据h = self.encoder(context.reshape(batch_size * seq_len, -1))   #将输入 context 重塑为 [batch_size * seq_len, state_dim + action_dim + 1]h = h.reshape(batch_size, seq_len, -1)  #将 h 重塑为 [batch_size, seq_len, hidden_dim]# 聚合序列信息（使用平均池化）h = torch.mean(h, dim=1)# 计算均值和标准差mean = self.mean(h)log_std = self.log_std(h)log_std = torch.clamp(log_std, -20, 2)  # 限制标准差范围return mean, log_stddef sample(self, mean, log_std):std = log_std.exp()dist = Normal(mean, std)  #创建正态分布return dist.rsample()  # 重参数化技巧#使用 rsample() 方法从正态分布中采样，这是一种重参数化技巧，用于在反向传播中保持梯度。# 策略网络 - 基于潜在表示和状态生成动作
class Policy(nn.Module):def __init__(self, state_dim, action_dim, latent_dim, hidden_dim, max_action):super(Policy, self).__init__()self.latent_dim = latent_dimself.policy = nn.Sequential(nn.Linear(state_dim + latent_dim, hidden_dim),nn.ReLU(),nn.Linear(hidden_dim, hidden_dim),nn.ReLU(),nn.Linear(hidden_dim, hidden_dim),nn.ReLU())self.mean = nn.Linear(hidden_dim, action_dim)self.log_std = nn.Linear(hidden_dim, action_dim)self.max_action = max_actiondef forward(self, state, z):# 拼接状态和潜在表示state_z = torch.cat([state, z], dim=1)h = self.policy(state_z)mean = self.mean(h)log_std = self.log_std(h)log_std = torch.clamp(log_std, -20, 2)  # 限制标准差范围std = log_std.exp()dist = Normal(mean, std)# 使用重参数化技巧采样动作action = dist.rsample()# 应用tanh squashing将动作限制在[-1, 1]范围内squashed_action = torch.tanh(action) * self.max_action# 计算对数概率（考虑tanh squashing）log_prob = dist.log_prob(action).sum(dim=-1, keepdim=True)log_prob -= (2 * (np.log(2) - action - F.softplus(-2 * action))).sum(dim=1, keepdim=True)return squashed_action, log_prob, meandef get_action(self, state, z):if isinstance(state, tuple):state = state[0]state = torch.FloatTensor(state.reshape(1, -1)).to(device)z = torch.FloatTensor(z.reshape(1, -1)).to(device)with torch.no_grad():action, _, _ = self.forward(state, z)return action.cpu().data.numpy().flatten()# Critic网络 - 评估动作价值
class Critic(nn.Module):def __init__(self, state_dim, action_dim, latent_dim, hidden_dim):super(Critic, self).__init__()# Q1架构self.Q1 = nn.Sequential(nn.Linear(state_dim + action_dim + latent_dim, hidden_dim),nn.ReLU(),nn.Linear(hidden_dim, hidden_dim),nn.ReLU(),nn.Linear(hidden_dim, 1))# Q2架构self.Q2 = nn.Sequential(nn.Linear(state_dim + action_dim + latent_dim, hidden_dim),nn.ReLU(),nn.Linear(hidden_dim, hidden_dim),nn.ReLU(),nn.Linear(hidden_dim, 1))def forward(self, state, action, z):# 拼接状态、动作和潜在表示state_action_z = torch.cat([state, action, z], dim=1)q1 = self.Q1(state_action_z)q2 = self.Q2(state_action_z)return q1, q2def Q1(self, state, action, z):state_action_z = torch.cat([state, action, z], dim=1)return self.Q1(state_action_z)# PEARL主体算法
class PEARL:def __init__(self, state_dim, action_dim, max_action, latent_dim=16, hidden_dim=256,lr=3e-4, gamma=0.99, tau=0.005, beta=1.0, context_batch_size=64,context_length=10):self.state_dim = state_dimself.action_dim = action_dimself.max_action = max_actionself.latent_dim = latent_dimself.gamma = gammaself.tau = tau   #目标网络的软更新系数。self.beta = beta  # KL散度系数self.context_batch_size = context_batch_sizeself.context_length = context_length# 初始化网络self.encoder = Encoder(state_dim, action_dim, latent_dim, hidden_dim).to(device)self.policy = Policy(state_dim, action_dim, latent_dim, hidden_dim, max_action).to(device)self.critic = Critic(state_dim, action_dim, latent_dim, hidden_dim).to(device)self.critic_target = Critic(state_dim, action_dim, latent_dim, hidden_dim).to(device)# 复制目标网络参数self.critic_target.load_state_dict(self.critic.state_dict())# 优化器self.encoder_optimizer = optim.Adam(self.encoder.parameters(), lr=lr)self.policy_optimizer = optim.Adam(self.policy.parameters(), lr=lr)self.critic_optimizer = optim.Adam(self.critic.parameters(), lr=lr)# 上下文缓冲区（用于存储任务相关的转换）self.context = []def reset_context(self):self.context = []def add_to_context(self, state, action, reward, next_state, done):self.context.append((state, action, reward, next_state, done))def get_latent(self):   #从当前任务收集的上下文数据中获取潜在表示if len(self.context) == 0:# 如果没有上下文，返回零向量return torch.zeros(1, self.latent_dim).to(device)# 准备上下文数据states = []actions = []rewards = []for s, a, r, ns, d in self.context:# 确保每个状态是一个数值数组if isinstance(s, tuple):s = s[0]states.append(np.array(s, dtype=np.float32).flatten())  # 将状态展平为一维数组actions.append(np.array(a, dtype=np.float32))rewards.append(np.array(r, dtype=np.float32))# 转换为 NumPy 数组states = np.array(states)actions = np.array(actions)rewards = np.array(rewards)# 转换为张量states = torch.FloatTensor(states).to(device)actions = torch.FloatTensor(actions).to(device)rewards = torch.FloatTensor(rewards).unsqueeze(1).to(device)# 拼接状态、动作和奖励context = torch.cat([states, actions, rewards], dim=1).unsqueeze(0)  # [1, seq_len, state_dim+action_dim+1]# 编码上下文获取潜在表示mean, log_std = self.encoder(context)z = self.encoder.sample(mean, log_std)return zdef select_action(self, state):z = self.get_latent()return self.policy.get_action(state, z)def train(self, replay_buffer, context_buffer, iterations):for it in range(iterations):# 从主回放缓冲区采样state_batch, action_batch, reward_batch, next_state_batch, done_batch = replay_buffer.sample(self.context_batch_size)# 从上下文缓冲区采样任务上下文context_batch = random.sample(context_buffer, self.context_batch_size)# 处理上下文数据context_states = []context_actions = []context_rewards = []for context in context_batch:# 从每个任务的上下文中随机选择一些转换transitions = random.sample(context, min(len(context), self.context_length))cs = torch.FloatTensor([s for s, a, r, ns, d in transitions]).to(device)ca = torch.FloatTensor([a for s, a, r, ns, d in transitions]).to(device)cr = torch.FloatTensor([r for s, a, r, ns, d in transitions]).unsqueeze(1).to(device)context_states.append(cs)context_actions.append(ca)context_rewards.append(cr)# 对每个任务并行处理for i in range(self.context_batch_size):# 准备当前任务的上下文cs = context_states[i].unsqueeze(0)  # [1, seq_len, state_dim]ca = context_actions[i].unsqueeze(0)  # [1, seq_len, action_dim]cr = context_rewards[i].unsqueeze(0)  # [1, seq_len, 1]# 拼接上下文context = torch.cat([cs, ca, cr], dim=2)  # [1, seq_len, state_dim+action_dim+1]# 编码上下文获取潜在表示mean, log_std = self.encoder(context)z = self.encoder.sample(mean, log_std)   #表示当前任务的潜在特征# 扩展潜在表示以匹配批量大小z_expanded = z.repeat(self.context_batch_size, 1)# 计算目标Q值with torch.no_grad():next_action, _, _ = self.policy(next_state_batch, z_expanded)target_Q1, target_Q2 = self.critic_target(next_state_batch, next_action, z_expanded)target_Q = torch.min(target_Q1, target_Q2)target_Q = reward_batch + (1 - done_batch) * self.gamma * target_Q# 获取当前Q估计current_Q1, current_Q2 = self.critic(state_batch, action_batch, z_expanded)# 计算critic损失critic_loss = F.mse_loss(current_Q1, target_Q) + F.mse_loss(current_Q2, target_Q)# 优化criticself.critic_optimizer.zero_grad()critic_loss.backward()self.critic_optimizer.step()# 计算policy损失new_action, log_prob, _ = self.policy(state_batch, z_expanded)Q1, Q2 = self.critic(state_batch, new_action, z_expanded)Q = torch.min(Q1, Q2)# 计算策略损失和KL散度policy_loss = -Q.mean() + (self.beta * log_prob).mean()# 优化policyself.policy_optimizer.zero_grad()policy_loss.backward()self.policy_optimizer.step()# 计算encoder损失（变分下界）# 先验是标准正态分布prior_mean = torch.zeros_like(mean)prior_log_std = torch.zeros_like(log_std)# 计算KL散度kl_div = -0.5 * torch.sum(1 + 2 * log_std - mean.pow(2) - (2 * log_std).exp(), dim=1)kl_div = kl_div.mean()# 计算编码器损失encoder_loss = kl_divprint("encode loss is", encoder_loss, "------", i)# 优化encoderself.encoder_optimizer.zero_grad()encoder_loss.backward()self.encoder_optimizer.step()# 更新目标网络for param, target_param in zip(self.critic.parameters(), self.critic_target.parameters()):target_param.data.copy_(self.tau * param.data + (1 - self.tau) * target_param.data)# 训练函数
def train_pearl(env_name, max_episodes=200, max_timesteps=100, latent_dim=16,    #max_timesteps训练过程中的最大回合数， latent_dim潜在表示的维度hidden_dim=256, batch_size=64, context_batch_size=16,    #hidden_dim训练时的批量大小， context_batch_size上下文批次的大小context_length=10, beta=1.0, eval_freq=10):    #context_length 每个任务上下文的长度，即从每个任务中采样的轨迹长度， beta KL 散度的系数# 创建环境env = gym.make(env_name)state_dim = env.observation_space.shape[0]action_dim = env.action_space.shape[0]max_action = float(env.action_space.high[0])# 初始化PEARL代理agent = PEARL(state_dim, action_dim, max_action, latent_dim, hidden_dim, beta=beta,context_batch_size=context_batch_size, context_length=context_length)# 初始化回放缓冲区，用于训练 Critic 网络和策略网络replay_buffer = ReplayBuffer(int(1e6))# 初始化上下文缓冲区（存储多个任务的上下文），用于更新编码器的参数context_buffer = []# 记录结果rewards_history = []# 训练循环for episode in range(max_episodes):print("max_episodes-->", episode)# 重置环境和代理的上下文state = env.reset()if isinstance(state, tuple):state = state[0]agent.reset_context()episode_reward = 0episode_context = []  #存储当前任务的上下文for t in range(max_timesteps):print("max_timesateps-->",t)# 选择动作action = agent.select_action(state)# 执行动作next_state, reward, done, _, _ = env.step(action)# 保存转换到主回放缓冲区replay_buffer.add(state, action, reward, next_state, done)# 保存转换到当前任务的上下文agent.add_to_context(state, action, reward, next_state, done)episode_context.append((state, action, reward, next_state, done))state = next_stateepisode_reward += rewardif done:break# 将当前任务的上下文添加到上下文缓冲区context_buffer.append(episode_context)# 记录奖励rewards_history.append(episode_reward)# 训练代理if len(replay_buffer) > batch_size and len(context_buffer) > context_batch_size:agent.train(replay_buffer, context_buffer, iterations=1)# 定期评估if episode % eval_freq == 0:avg_reward = evaluate_policy(agent, env_name, max_timesteps, eval_episodes=5)print(f"Episode {episode}, Avg Reward: {avg_reward:.3f}")# 绘制学习曲线plt.figure(figsize=(10, 5))plt.plot(rewards_history)plt.title('PEARL Training Rewards')plt.xlabel('Episode')plt.ylabel('Reward')plt.show()return agent# 评估策略
def evaluate_policy(agent, env_name, max_timesteps, eval_episodes=5):env = gym.make(env_name)avg_reward = 0.for _ in range(eval_episodes):state = env.reset()agent.reset_context()done = Falsewhile not done:action = agent.select_action(state)state, reward, done, _, _ = env.step(action)avg_reward += rewardavg_reward /= eval_episodesreturn avg_reward# 训练PEARL代理
if __name__ == "__main__":env_name = "Pendulum-v1"  # 可以替换为其他连续控制环境agent = train_pearl(env_name, max_episodes=200, max_timesteps=100)