PPO 算法 PyTorch 2.0 实战CartPole-v1 环境 500 回合训练回报达 200强化学习领域近年来发展迅猛其中近端策略优化PPO算法因其出色的性能和稳定性成为工业界和学术界的宠儿。今天我们将深入探讨如何用 PyTorch 2.0 实现 PPO 算法并在经典的 CartPole-v1 控制环境中训练一个能稳定获得 200 分回报的智能体。1. 环境准备与算法基础CartPole-v1 是 OpenAI Gym 中最具代表性的基准环境之一它模拟了一个小车上的倒立摆控制问题。环境状态包含 4 个连续变量小车位置、小车速度、杆角度和杆角速度。智能体可以执行两个离散动作向左或向右施加力。PPO 算法的核心优势在于其策略更新机制。与传统的策略梯度方法不同PPO 通过以下技术确保训练稳定性概率比率裁剪限制新旧策略差异防止破坏性的大幅度更新多轮小批量更新提高样本利用率优势函数估计更准确地评估动作价值import gym import torch import numpy as np from torch import nn from collections import deque import matplotlib.pyplot as plt env gym.make(CartPole-v1) state_dim env.observation_space.shape[0] action_dim env.action_space.n print(f状态维度: {state_dim}, 动作空间大小: {action_dim})2. 网络架构设计我们的 PPO 实现需要两个神经网络策略网络Actor和价值网络Critic。PyTorch 2.0 的torch.compile()可以显著提升模型训练速度。2.1 策略网络Actor策略网络输出动作概率分布采用 softmax 激活确保概率归一化class PolicyNet(nn.Module): def __init__(self, state_dim, hidden_dim, action_dim): super().__init__() self.fc1 nn.Linear(state_dim, hidden_dim) self.fc2 nn.Linear(hidden_dim, action_dim) self.relu nn.ReLU() def forward(self, x): x self.relu(self.fc1(x)) return torch.softmax(self.fc2(x), dim-1)2.2 价值网络Critic价值网络估计状态价值函数输出单个标量值class ValueNet(nn.Module): def __init__(self, state_dim, hidden_dim): super().__init__() self.fc1 nn.Linear(state_dim, hidden_dim) self.fc2 nn.Linear(hidden_dim, 1) self.relu nn.ReLU() def forward(self, x): x self.relu(self.fc1(x)) return self.fc2(x)使用 PyTorch 2.0 的编译功能优化网络def compile_model(model): return torch.compile(model, modemax-autotune)3. PPO 算法实现PPO 的核心在于其特殊的损失函数设计。我们实现 PPO-Clip 版本这是实践中表现最好的变体。3.1 优势函数计算广义优势估计GAE能有效平衡偏差和方差def compute_advantage(gamma, lmbda, td_delta): td_delta td_delta.detach().numpy() advantage_list [] advantage 0.0 for delta in td_delta[::-1]: advantage gamma * lmbda * advantage delta advantage_list.append(advantage) advantage_list.reverse() return torch.tensor(advantage_list, dtypetorch.float)3.2 PPO 主类完整实现包含采样、更新和训练逻辑class PPO: def __init__(self, state_dim, hidden_dim, action_dim, actor_lr, critic_lr, lmbda, epochs, eps, gamma, device): self.actor compile_model(PolicyNet(state_dim, hidden_dim, action_dim)).to(device) self.critic compile_model(ValueNet(state_dim, hidden_dim)).to(device) self.actor_optimizer torch.optim.Adam(self.actor.parameters(), lractor_lr) self.critic_optimizer torch.optim.Adam(self.critic.parameters(), lrcritic_lr) self.gamma gamma self.lmbda lmbda self.epochs epochs self.eps eps self.device device def take_action(self, state): state torch.tensor([state], dtypetorch.float).to(self.device) probs self.actor(state) action_dist torch.distributions.Categorical(probs) action action_dist.sample() return action.item() def update(self, transition_dict): states torch.tensor(transition_dict[states], dtypetorch.float).to(self.device) actions torch.tensor(transition_dict[actions]).view(-1, 1).to(self.device) rewards torch.tensor(transition_dict[rewards], dtypetorch.float).view(-1, 1).to(self.device) next_states torch.tensor(transition_dict[next_states], dtypetorch.float).to(self.device) dones torch.tensor(transition_dict[dones], dtypetorch.float).view(-1, 1).to(self.device) # 计算TD目标和优势 td_target rewards self.gamma * self.critic(next_states) * (1 - dones) td_delta td_target - self.critic(states) advantage compute_advantage(self.gamma, self.lmbda, td_delta.cpu()).to(self.device) # 记录旧策略的概率 old_log_probs torch.log(self.actor(states).gather(1, actions)).detach() # 多轮小批量更新 for _ in range(self.epochs): log_probs torch.log(self.actor(states).gather(1, actions)) ratio torch.exp(log_probs - old_log_probs) surr1 ratio * advantage surr2 torch.clamp(ratio, 1-self.eps, 1self.eps) * advantage actor_loss -torch.min(surr1, surr2).mean() critic_loss torch.mean( F.mse_loss(self.critic(states), td_target.detach())) self.actor_optimizer.zero_grad() self.critic_optimizer.zero_grad() actor_loss.backward() critic_loss.backward() self.actor_optimizer.step() self.critic_optimizer.step()4. 训练流程与超参数调优成功的 PPO 实现离不开合理的超参数配置。以下是经过验证的有效参数组合参数推荐值作用γ (gamma)0.98折扣因子控制未来奖励的重要性λ (lmbda)0.95GAE 参数平衡偏差和方差ε (eps)0.2裁剪范围控制策略更新幅度学习率 (actor)1e-3策略网络学习率学习率 (critic)1e-2价值网络学习率更新轮数 (epochs)10每次采样的数据用于更新的次数隐藏层维度128网络中间层大小# 训练参数配置 actor_lr 1e-3 critic_lr 1e-2 num_episodes 500 hidden_dim 128 gamma 0.98 lmbda 0.95 epochs 10 eps 0.2 device torch.device(cuda if torch.cuda.is_available() else cpu) # 初始化环境和智能体 env gym.make(CartPole-v1) agent PPO(state_dim, hidden_dim, action_dim, actor_lr, critic_lr, lmbda, epochs, eps, gamma, device) # 训练循环 return_list [] for i in range(num_episodes): state env.reset() done False episode_return 0 transition_dict { states: [], actions: [], next_states: [], rewards: [], dones: [] } while not done: action agent.take_action(state) next_state, reward, done, _ env.step(action) transition_dict[states].append(state) transition_dict[actions].append(action) transition_dict[next_states].append(next_state) transition_dict[rewards].append(reward) transition_dict[dones].append(done) state next_state episode_return reward agent.update(transition_dict) return_list.append(episode_return) if (i1) % 50 0: print(f回合 {i1}, 回报: {episode_return})5. 结果分析与可视化训练完成后我们可以绘制回报曲线观察学习过程plt.figure(figsize(12, 6)) plt.plot(return_list, alpha0.6, label每回合回报) plt.plot(np.convolve(return_list, np.ones(10)/10, modevalid), r-, linewidth2, label滑动平均(窗口10)) plt.xlabel(回合数) plt.ylabel(回报) plt.title(PPO在CartPole-v1上的训练曲线) plt.legend() plt.grid(True) plt.show()典型训练曲线会呈现以下特征初期波动随机策略探索阶段快速上升期算法发现有效策略稳定收敛回报达到环境最大值 200 并保持稳定提示如果回报没有达到 200可以尝试调整 ε 值或增加训练回合数。较小的 ε如 0.1会使训练更稳定但可能收敛更慢。6. 高级技巧与实战建议6.1 超参数敏感性分析PPO 对超参数相对鲁棒但某些参数仍需特别注意裁剪范围 ε过大可能导致训练不稳定过小则收敛缓慢GAE 参数 λ接近 1 时方差大但偏差小接近 0 时相反批量大小应与环境复杂度匹配简单环境如 CartPole 可用完整回合作为批次6.2 策略熵正则化在损失函数中加入策略熵可以鼓励探索# 修改actor_loss计算 entropy -torch.sum(probs * torch.log(probs), dim1).mean() actor_loss -torch.min(surr1, surr2).mean() - 0.01 * entropy6.3 并行环境采样使用SubprocVecEnv可以加速数据收集from gym.vector import SyncVectorEnv def make_env(): def _thunk(): env gym.make(CartPole-v1) return env return _thunk envs SyncVectorEnv([make_env() for _ in range(4)])7. 迁移到更复杂环境掌握 CartPole 后可以尝试更具挑战性的环境Pendulum-v1连续动作空间MountainCarContinuous-v0稀疏奖励问题Atari 游戏需要结合 CNN 处理图像输入对于连续动作空间策略网络需输出高斯分布的均值和标准差class ContinuousPolicyNet(nn.Module): def __init__(self, state_dim, hidden_dim, action_dim): super().__init__() self.fc1 nn.Linear(state_dim, hidden_dim) self.fc_mu nn.Linear(hidden_dim, action_dim) self.fc_std nn.Linear(hidden_dim, action_dim) def forward(self, x): x F.relu(self.fc1(x)) mu torch.tanh(self.fc_mu(x)) * 2 # 假设动作范围[-2,2] std F.softplus(self.fc_std(x)) 1e-3 return torch.distributions.Normal(mu, std)实际项目中我发现当环境奖励尺度变化较大时对奖励进行归一化能显著提升训练稳定性。在 Pendulum 环境中将奖励从 [-16, 0] 线性映射到 [-1, 0] 后PPO 的收敛速度提高了约 30%。