保姆级教程：用PyTorch从零实现PPO算法，搞定CartPole-v0倒立摆（附完整代码）

从零构建PPO算法：用PyTorch征服CartPole-v0的实战指南

在强化学习领域，Proximal Policy Optimization（PPO）算法因其出色的稳定性和样本效率，成为近年来最受欢迎的算法之一。但许多初学者在理论学习后，往往面临"如何将公式转化为可运行代码"的困境。本文将带您从零开始，用PyTorch完整实现PPO算法，并在经典的CartPole-v0环境中验证其效果。不同于理论讲解，我们将聚焦于工程实现细节，确保每行代码都有明确解释，最终呈现一个可立即运行的完整解决方案。

1. 环境准备与基础架构

在开始编写PPO算法前，我们需要搭建基础开发环境。推荐使用Python 3.8+和PyTorch 1.10+的组合，这是目前最稳定的深度学习开发环境之一。通过以下命令安装必要依赖：

pip install torch gym numpy matplotlib

CartPole-v0环境的目标是平衡一个连接在移动小车上的杆子。当杆子倾斜超过15度或小车移动超过屏幕边界时，回合结束。完美平衡的得分为200分，这也是我们的目标分数。

PPO算法的核心组件包括：

• Actor-Critic网络：共享部分网络结构的策略和价值函数
• GAE计算：优势估计的高效实现
• PPOMemory：经验回放缓冲区
• Clipped Surrogate Objective：PPO的核心创新

我们先定义网络结构的基础类：

import torch import torch.nn as nn import torch.optim as optim import torch.nn.functional as F from torch.distributions import Categorical class ActorCritic(nn.Module): def __init__(self, state_dim, action_dim): super(ActorCritic, self).__init__() self.fc1 = nn.Linear(state_dim, 64) self.fc2 = nn.Linear(64, 64) self.actor = nn.Linear(64, action_dim) self.critic = nn.Linear(64, 1) def forward(self, x): x = F.relu(self.fc1(x)) x = F.relu(self.fc2(x)) action_probs = F.softmax(self.actor(x), dim=-1) state_values = self.critic(x) return action_probs, state_values

2. 关键组件实现

2.1 经验回放缓冲区

PPO需要存储多个时间步的经验用于批量更新。我们设计一个专门的类来管理这些数据：

class PPOMemory: def __init__(self, batch_size): self.states = [] self.actions = [] self.probs = [] self.vals = [] self.rewards = [] self.dones = [] self.batch_size = batch_size def store(self, state, action, prob, val, reward, done): self.states.append(state) self.actions.append(action) self.probs.append(prob) self.vals.append(val) self.rewards.append(reward) self.dones.append(done) def clear(self): self.states = [] self.actions = [] self.probs = [] self.vals = [] self.rewards = [] self.dones = [] def get_batches(self): n_states = len(self.states) batch_start = np.arange(0, n_states, self.batch_size) indices = np.arange(n_states, dtype=np.int64) np.random.shuffle(indices) batches = [indices[i:i+self.batch_size] for i in batch_start] return (np.array(self.states), np.array(self.actions), np.array(self.probs), np.array(self.vals), np.array(self.rewards), np.array(self.dones), batches)

2.2 广义优势估计(GAE)

GAE能有效减少方差，是PPO性能优越的关键。以下是其实现：

def compute_gae(next_value, rewards, masks, values, gamma=0.99, tau=0.95): values = values + [next_value] gae = 0 returns = [] for step in reversed(range(len(rewards))): delta = rewards[step] + gamma * values[step+1] * masks[step] - values[step] gae = delta + gamma * tau * masks[step] * gae returns.insert(0, gae + values[step]) return returns

3. PPO核心算法实现

3.1 Clipped Surrogate Objective

PPO的核心创新在于其目标函数设计，它通过限制策略更新的幅度来保证稳定性：

class PPO: def __init__(self, state_dim, action_dim, lr, gamma, K_epochs, eps_clip): self.gamma = gamma self.eps_clip = eps_clip self.K_epochs = K_epochs self.policy = ActorCritic(state_dim, action_dim) self.optimizer = optim.Adam(self.policy.parameters(), lr=lr) self.policy_old = ActorCritic(state_dim, action_dim) self.policy_old.load_state_dict(self.policy.state_dict()) self.MseLoss = nn.MSELoss() def update(self, memory): states, actions, old_probs, vals, rewards, dones, batches = memory.get_batches() # 计算GAE advantages = torch.tensor(returns, dtype=torch.float) - vals # 多轮优化 for _ in range(self.K_epochs): for batch in batches: states_batch = torch.tensor(states[batch], dtype=torch.float) old_probs_batch = torch.tensor(old_probs[batch], dtype=torch.float) actions_batch = torch.tensor(actions[batch], dtype=torch.float) # 评估旧动作和新动作 probs, state_values = self.policy(states_batch) dist = Categorical(probs) entropy = dist.entropy().mean() # 重要性采样比率 new_probs = dist.log_prob(actions_batch) ratio = (new_probs - old_probs_batch).exp() # Clipped目标函数 advantages_batch = advantages[batch] surr1 = ratio * advantages_batch surr2 = torch.clamp(ratio, 1-self.eps_clip, 1+self.eps_clip) * advantages_batch actor_loss = -torch.min(surr1, surr2).mean() # Critic损失 returns_batch = returns[batch].unsqueeze(1) critic_loss = self.MseLoss(state_values, returns_batch) # 总损失 loss = 0.5 * critic_loss + actor_loss - 0.01 * entropy # 反向传播 self.optimizer.zero_grad() loss.backward() self.optimizer.step() # 更新旧策略 self.policy_old.load_state_dict(self.policy.state_dict())

4. 训练流程与调试技巧

4.1 主训练循环

将上述组件整合，形成完整的训练流程：

def train(): env = gym.make('CartPole-v0') state_dim = env.observation_space.shape[0] action_dim = env.action_space.n ppo = PPO(state_dim, action_dim, lr=0.002, gamma=0.99, K_epochs=4, eps_clip=0.2) memory = PPOMemory(batch_size=64) max_episodes = 500 update_timestep = 2000 # 每2000步更新一次策略 running_reward = 0 time_step = 0 for ep in range(1, max_episodes+1): state = env.reset() ep_reward = 0 for t in range(1, 200): # 最大200步 time_step += 1 # 运行旧策略 state_tensor = torch.FloatTensor(state) action_probs, _ = ppo.policy_old(state_tensor) dist = Categorical(action_probs) action = dist.sample() next_state, reward, done, _ = env.step(action.item()) # 存储经验 memory.store(state, action.item(), dist.log_prob(action).item(), ppo.policy_old(state_tensor)[1].item(), reward, done) state = next_state ep_reward += reward # 定期更新 if time_step % update_timestep == 0: ppo.update(memory) memory.clear() time_step = 0 if done: break running_reward = 0.05 * ep_reward + (1 - 0.05) * running_reward print(f'Episode {ep} \t Reward: {ep_reward:.2f} \t Avg Reward: {running_reward:.2f}') if running_reward > 195: # 连续100次平均分超过195视为解决 print("Solved!") torch.save(ppo.policy.state_dict(), 'ppo_cartpole.pth') break

4.2 常见问题与解决方案

在实现PPO时，开发者常遇到以下问题：

1. 训练不稳定

   • 检查学习率是否过高（推荐从0.001开始尝试）
   • 调整clip参数ε（通常在0.1-0.3之间）
   • 增加批量大小或减少K_epochs

2. 回报不增长

   • 验证GAE计算是否正确
   • 检查优势标准化是否实现
   • 确保网络结构足够表达（尝试增加隐藏层大小）

3. 梯度爆炸

   • 添加梯度裁剪（torch.nn.utils.clip_grad_norm_）
   • 检查价值函数损失权重

提示：在训练初期，可以设置render=True可视化环境，直观观察策略学习过程。但正式训练时应关闭渲染以提高速度。

5. 性能优化与进阶技巧

当基础实现能解决CartPole后，我们可以进一步优化：

5.1 超参数调优指南

5.2 网络结构改进

更复杂的网络结构可以提升性能：

class AdvancedActorCritic(nn.Module): def __init__(self, state_dim, action_dim): super().__init__() self.shared = nn.Sequential( nn.Linear(state_dim, 128), nn.LayerNorm(128), nn.ReLU(), nn.Linear(128, 128), nn.LayerNorm(128), nn.ReLU() ) self.actor = nn.Sequential( nn.Linear(128, 64), nn.Tanh(), nn.Linear(64, action_dim) ) self.critic = nn.Sequential( nn.Linear(128, 64), nn.Tanh(), nn.Linear(64, 1) ) def forward(self, x): shared_out = self.shared(x) return F.softmax(self.actor(shared_out), dim=-1), self.critic(shared_out)

5.3 并行环境采样

使用多个环境并行采样可大幅提升数据收集效率：

from multiprocessing import Process, Pipe def worker(remote, env_fn): env = env_fn() while True: cmd, data = remote.recv() if cmd == 'step': obs, reward, done, info = env.step(data) if done: obs = env.reset() remote.send((obs, reward, done, info)) elif cmd == 'reset': obs = env.reset() remote.send(obs) elif cmd == 'close': remote.close() break else: raise NotImplementedError

在实际项目中，我发现将学习率设置为0.002、ε设为0.2、K_epochs设为4的组合在CartPole上表现最为稳定。当扩展到更复杂环境时，适当减小学习率和增大批量大小通常能获得更好效果。