2. PPO算法实现：Actor-Critic架构、GAE优势函数估计、PPO的PyTorch代码实现（离散动作空间）
好，咱们直接进入正题。上一章我们把PPO的理论框架梳理了一遍，这一章我带你手撕代码。说实话，理论看再多，不如自己跑通一个demo来得实在。我个人习惯是，先搭一个最小的可行版本，再逐步优化。今天我们就用PyTorch实现一个离散动作空间的PPO，核心就三块：Actor-Critic网络、GAE优势估计、以及那个关键的clip损失。
2.1 Actor-Critic架构：两个网络各司其职
PPO用的是Actor-Critic架构。说白了，就是两个网络：
- Actor（策略网络）：负责输出动作的概率分布。在离散动作空间里，就是输出每个动作的logits，然后过softmax得到概率。
- Critic（价值网络）：负责评估当前状态的价值，输出一个标量V(s)。
我在项目中遇到过一个问题：很多人把Actor和Critic设计成共享底层特征提取层。这样做参数少、训练快，但有个坑——两个任务的梯度可能会互相干扰。我的建议是：如果状态空间简单（比如CartPole），共享没问题；如果复杂（比如Atari游戏），最好分开，或者至少用两个独立的输出头。
核心要点：Actor输出的是概率分布，不是具体的动作值。采样时用Categorical分布，训练时用这个分布计算log_prob。
来看代码。我习惯把Actor和Critic封装成一个类，方便管理：
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.distributions import Categorical
class ActorCritic(nn.Module):
def __init__(self, state_dim, action_dim, hidden_dim=64):
super().__init__()
# 共享的特征提取层
self.fc1 = nn.Linear(state_dim, hidden_dim)
self.fc2 = nn.Linear(hidden_dim, hidden_dim)
# Actor头：输出动作logits
self.actor = nn.Linear(hidden_dim, action_dim)
# Critic头：输出状态价值
self.critic = nn.Linear(hidden_dim, 1)
def forward(self, x):
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
logits = self.actor(x)
value = self.critic(x)
return logits, value
def get_action(self, state):
"""根据状态采样动作，返回动作和对应的log_prob"""
logits, value = self.forward(state)
dist = Categorical(logits=logits)
action = dist.sample()
log_prob = dist.log_prob(action)
return action.item(), log_prob, value
def evaluate(self, state, action):
"""用于训练时：给定状态和动作，计算log_prob和value"""
logits, value = self.forward(state)
dist = Categorical(logits=logits)
log_prob = dist.log_prob(action)
# 计算熵，用于鼓励探索
entropy = dist.entropy()
return log_prob, value, entropy小技巧：get_action用于采样，evaluate用于训练。这样区分开，逻辑更清晰。我早期写代码时混在一起，调试起来很痛苦。
2.2 GAE优势函数估计：让优势计算更平滑
接下来是GAE（Generalized Advantage Estimation）。为什么要用GAE？你想想看，如果用简单的TD误差（r + γV(s') - V(s)），方差小但偏差大；如果用蒙特卡洛回报（累计折扣奖励），偏差小但方差大。GAE在两者之间做了个折中，通过一个参数λ来控制。
GAE的计算公式其实不复杂：
δ_t = r_t + γ * V(s_{t+1}) - V(s_t)
A_t^{GAE(γ,λ)} = Σ_{l=0}^{∞} (γλ)^l * δ_{t+l}
嗯，这里要注意：实际实现时，我们是倒着算的，从最后一个时间步往前推。这样效率高，而且符合递推关系。
我曾经在实现GAE时犯过一个低级错误——忘记对最后一个状态的V(s_{t+1})做处理。如果episode结束了，V(s_{t+1})应该是0，而不是网络输出的值。这个bug让我排查了一整天，后来发现是done标志没处理好。
def compute_gae(rewards, values, dones, gamma=0.99, lam=0.95):
"""
计算GAE优势函数
rewards: list of rewards
values: list of state values (包括最后一个状态的V(s_{t+1}))
dones: list of done flags
"""
advantages = []
gae = 0
# 倒着算，从最后一个时间步开始
for t in reversed(range(len(rewards))):
# 如果当前episode结束，下一个状态的价值为0
if dones[t]:
delta = rewards[t] - values[t]
gae = delta
else:
delta = rewards[t] + gamma * values[t+1] - values[t]
gae = delta + gamma * lam * gae
advantages.insert(0, gae)
# 计算returns = advantages + values（用于Critic的监督信号）
returns = [adv + val for adv, val in zip(advantages, values[:-1])]
return advantages, returns注意：values的长度比rewards多1，因为包含了最后一个状态的V(s_{t+1})。dones的长度和rewards一致。这个细节很容易搞错。
2.3 PPO的PyTorch完整实现
好了，核心组件都有了，现在把它们组装起来。PPO的训练流程其实就几个步骤：
- 用当前策略与环境交互，收集一批数据（轨迹）
- 计算GAE优势和returns
- 用这些数据更新Actor和Critic，重复K个epoch
- 清空数据，开始新一轮交互
我直接给出完整的训练代码，注释写得很清楚：
class PPO:
def __init__(self, state_dim, action_dim, lr=3e-4, gamma=0.99,
lam=0.95, clip_epsilon=0.2, K_epochs=10, batch_size=64):
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.net = ActorCritic(state_dim, action_dim).to(self.device)
self.optimizer = torch.optim.Adam(self.net.parameters(), lr=lr)
self.gamma = gamma
self.lam = lam
self.clip_epsilon = clip_epsilon
self.K_epochs = K_epochs
self.batch_size = batch_size
# 用于存储轨迹数据
self.states = []
self.actions = []
self.rewards = []
self.dones = []
self.log_probs = []
self.values = []
def select_action(self, state):
"""与环境交互时调用，返回动作"""
state = torch.FloatTensor(state).unsqueeze(0).to(self.device)
action, log_prob, value = self.net.get_action(state)
# 存储到缓冲区
self.states.append(state)
self.actions.append(action)
self.log_probs.append(log_prob)
self.values.append(value)
return action
def store_transition(self, reward, done):
"""环境返回reward和done后调用"""
self.rewards.append(reward)
self.dones.append(done)
def update(self):
"""收集完一个batch的数据后，调用此函数更新网络"""
# 1. 计算GAE优势和returns
# 注意：values需要包含最后一个状态的V(s_{t+1})
values = torch.cat(self.values).squeeze().cpu().detach().numpy().tolist()
# 获取最后一个状态的value（用于GAE计算）
last_state = self.states[-1]
with torch.no_grad():
_, last_value = self.net(last_state)
values.append(last_value.item())
advantages, returns = compute_gae(
self.rewards, values, self.dones, self.gamma, self.lam
)
# 2. 转换为tensor
states = torch.cat(self.states).to(self.device)
actions = torch.LongTensor(self.actions).to(self.device)
old_log_probs = torch.cat(self.log_probs).to(self.device)
advantages = torch.FloatTensor(advantages).to(self.device)
returns = torch.FloatTensor(returns).to(self.device)
# 标准化优势（经验上能稳定训练）
advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
# 3. 多epoch更新
dataset_size = len(states)
for _ in range(self.K_epochs):
# 随机打乱数据
indices = torch.randperm(dataset_size)
for start in range(0, dataset_size, self.batch_size):
batch_indices = indices[start:start + self.batch_size]
batch_states = states[batch_indices]
batch_actions = actions[batch_indices]
batch_old_log_probs = old_log_probs[batch_indices]
batch_advantages = advantages[batch_indices]
batch_returns = returns[batch_indices]
# 计算新策略的log_prob和value
log_probs, values, entropy = self.net.evaluate(batch_states, batch_actions)
# 计算ratio: π_new / π_old
ratios = torch.exp(log_probs - batch_old_log_probs)
# PPO的clip损失
surr1 = ratios * batch_advantages
surr2 = torch.clamp(ratios, 1 - self.clip_epsilon,
1 + self.clip_epsilon) * batch_advantages
actor_loss = -torch.min(surr1, surr2).mean()
# Critic损失（MSE）
critic_loss = F.mse_loss(values.squeeze(), batch_returns)
# 熵正则项（鼓励探索）
entropy_loss = -entropy.mean()
total_loss = actor_loss + 0.5 * critic_loss + 0.01 * entropy_loss
# 反向传播
self.optimizer.zero_grad()
total_loss.backward()
# 梯度裁剪，防止梯度爆炸
torch.nn.utils.clip_grad_norm_(self.net.parameters(), max_norm=0.5)
self.optimizer.step()
# 4. 清空缓冲区
self.clear_buffer()
def clear_buffer(self):
self.states = []
self.actions = []
self.rewards = []
self.dones = []
self.log_probs = []
self.values = []关键参数说明：
- clip_epsilon=0.2：PPO的核心超参数，控制策略更新的幅度。太大容易不稳定，太小收敛慢。
- K_epochs=10：每条数据重复使用的次数。我试过5-15之间效果都不错，再大容易过拟合。
- lam=0.95：GAE的λ参数。0.95是个比较安全的默认值，偏向于降低方差。
2.4 核心逻辑流程图
下面我用一张SVG图把整个PPO的训练流程串起来，方便你理解数据流向：
2.5 避坑指南与经验分享
代码写完了，但跑起来可能还有各种问题。我把自己踩过的坑列出来，你遇到了可以直接对照：
| 常见问题 | 现象 | 解决方案 |
|---|---|---|
| 梯度爆炸 | loss突然变成NaN，或者策略瞬间崩溃 | 加梯度裁剪（clip_grad_norm_），我一般设0.5或1.0 |
| 策略坍塌 | 熵值降到接近0，不再探索 | 加熵正则项（0.01左右），或者降低K_epochs |
| 训练不稳定 | reward曲线剧烈震荡 | 检查GAE的λ参数，试试0.9-0.95；或者减小学习率 |
| 收敛太慢 | reward长期不增长 | 增大batch_size，或者提高学习率到5e-4 |
我的经验：如果你在CartPole上测试，一般100-200个episode就能收敛到满分（200分）。如果跑了几百个episode还没动静，先检查reward的scale——PPO对reward的scale比较敏感，建议归一化到[-1, 1]区间。
最后说一句，代码只是工具，理解背后的原理才是关键。你把这个离散版本的PPO跑通了，连续动作空间无非就是把Categorical分布换成Normal分布，其他逻辑几乎一模一样。好了，动手试试吧。