策略梯度算法详解：从REINFORCE到PPO的强化学习之路

去年我在某实验室做连续控制实验时被一个简单问题卡了很久用DQN训练一个双关节机械臂动作空间稍微一变宽Q网络就学不动了。后来我换了个思路不再让智能体去估算每个动作的“价值”而是直接让策略网络输出“该做哪个动作的概率”。这个看似简单的转变把我拉进了强化学习里最基础也最核心的方法之一——Policy Gradient。它不依赖值函数逼近的动作枚举天然适配连续动作和高维动作空间是理解后续Actor-Critic、PPO、DPO等大量算法的地基。这篇内容没有复杂的项目背景就是纯粹把Policy Gradient这个主题拆开讲清楚它解决什么问题、数学推导怎么来的、最小可跑通实现长什么样、训练时翻车最多的原因有哪些、以及从它到PPO之间到底改了什么。适合刚学完Q-learning、想往策略搜索方向走的同学也适合工作中需要快速上手策略梯度做简单控制的工程师。1. 为什么卡在“连续动作”时我转向了Policy Gradient1.1 基于价值的方法在连续控制上的天然短板DQN这类基于价值的方法核心是学一个Q(s, a)函数然后通过argmax_a Q(s, a)来选动作。这个逻辑在离散动作空间里非常顺畅比如游戏里上下左右四个动作遍历一遍就行。但遇到连续控制——比如机器人关节力矩是一个连续值、倒立摆需要输出从-1到1的任意推力——你怎么去枚举argmax不可能遍历所有实数。你可能会说可以用一个连续优化器在每一步去搜索最优动作。这确实是一种方案但每一步都跑一遍优化计算开销大到难以落地。而策略梯度直接换了一个视角我不再问“哪个动作价值最大”我直接让网络输出一个动作分布——连续动作就是高斯分布的均值和方法离散动作就是softmax概率——然后通过梯度上升拿奖励最大的方向去调整这些参数。动作是“采”出来的不是“选”出来的。1.2 随机策略带来的额外红利探索不再靠epsilon熟悉DQN的人都知道DQN靠ε-greedy做探索以一定概率随机乱选看运气能不能碰到新状态。而策略梯度天然输出的是概率分布训练初期所有动作的概率都比较均匀本身就是一个探索过程训练后期概率逐渐集中到高回报动作上这又是一个天然的“利用”过程。这个特性在实践中很舒服。我不用去衰减ε参数也不用担心陷入“探索不够”或“利用过少”的两难。你只需要看策略的熵变化——熵从高到低往往意味着策略在逐渐稳定。这是基于价值的算法里没有的直观反馈。2. 策略梯度的核心推导一个公式看懂REINFORCE2.1 目标函数与log导数技巧先说优化目标。我们用参数θ定义一个策略π_θ(a|s)希望最大化期望累计回报J(θ) E_τ ~ π_θ [ Σ_t γ^t r_t ]这里的τ是一条完整轨迹trajectory包含状态、动作、奖励序列。求梯度时把期望展开成概率加权求和剩下关键一步既然期望里带着策略参数不能直接对奖励求导就要用到一个化简技巧——∇π / π ∇log π。于是推导结果非常简洁∇θ J(θ) E_τ [ Σ_t ∇θ log π_θ(a_t|s_t) · R(τ) ]注意这里用的是整条轨迹的总回报R(τ)不是单步奖励r_t。这是很多刚入门的人容易晕的点。为什么是整条轨迹因为REINFORCE推导中我们用一次完整的采样轨迹来估计期望回报所以参数更新沿着整条轨迹的“总结果”走。2.2 REINFORCE的直观理解好轨迹抬概率坏轨迹压概率把上面公式翻译成人话如果一条轨迹的总回报很高那我们就按这条轨迹上所有动作的log概率的梯度方向去增加它们的出现概率如果总回报很低就降低它们的概率。这个逻辑非常朴素像极了“老板表扬整个项目的功劳全组加薪项目搞砸了全组扣绩效”。递推关系也很清晰。策略网络输入状态s_t输出动作分布在a_t处取出log π_θ(a_t|s_t)乘上这条轨迹的累计回报R(τ)对所有时间步求和再去平均多条轨迹就是梯度的蒙特卡洛估计。2.3 一步一个脚印的REINFORCE伪代码整体流程拆成五步用当前策略跑若干个完整episode记录每一步的状态、动作、奖励。对每个episode从结束位置往前倒推计算累计回报G_t r_t γ G_{t1}。构建损失函数L - Σ log π_θ(a_t|s_t) · G_t。加负号是因为我们要梯度上升最大化回报而PyTorch的optimizer.step()默认做梯度下降。反向传播更新策略网络参数。重复。这里有一点要额外强调每一步的log π要乘以对应步的累计回报G_t而不是整条轨迹总回报。虽然标准的REINFORCE推导里写的是R(τ)但实践中几乎都是用带折扣的G_t作为从当前步开始未来累计回报的估计效果更稳定。原因后面第3章会展开。3. 方差失控训练翻车的最常见原因与baseline修复3.1 为什么Policy Gradient的方差那么大蒙特卡洛采样的本质是拿少量样本来估计一个高维积分而轨迹本身的随机性来自两方面策略采样动作有随机性环境转移动态也有随机性。两条轨迹即使初始状态一样、动作一模一样后续状态和奖励都可能差异很大。方差自然就大。更麻烦的是log π梯度的量级受回报大小直接影响。某个episode运气好回报高达几千另一个episode运气差回报接近零。梯度更新步长就会剧烈抖动网络参数来回震荡训练曲线很难看。3.2 baseline不改变期望却能砍掉一大半方差REINFORCE的经典改进是引入基线b(s)把更新中使用的回报替换成G_t - b(s_t)。这个改动最妙的地方是由于Σ_a π(a|s) ∇log π(a|s) 0减去baseline后梯度的期望不变。也就是说它不改变我们优化方向的数学期望只是让每次采样的梯度估计更稳定。原理可以这样理解如果一个轨迹本身只是“环境本身波动大导致的回报高”而不是“策略动作真的好”那减去baseline后这部分高回报中与动作无关的贡献就被抵消了。最常用的baseline是状态价值函数V(s)即“在这个状态下正常情况下我能拿多少回报”。把G_t - V(s_t)当作“动作比平均好多少”的度量这就是优势函数的思想雏形。3.3 实操中的另一个低成本降方差手段回报归一化baseline是理论上的标准方案但在最小实现里还有一个非常实用的小技巧对一批episode的G_t做标准化减均值除标准差。这样做不改变每个episode的相对好坏关系但会把梯度量级拉到相对可控的范围配合固定学习率时更容易稳定。我试过的经验是如果奖励设计得不好比如奖励范围差了一百倍纯REINFORCE几乎训不动加上baseline或归一化之后至少能看出策略在学。做实验时我会先跑通“纯REINFORCE 回报归一化”的版本再决定要不要加真正的critic网络。4. 一个能跑通的PyTorch最小实现和完整参数解读4.1 网络定义与采样逻辑这个示例用经典的CartPole环境动作是离散的左右两挡状态是四维向量。我们用一个两层MLP做策略网络输出两个logits再用softmax转成概率采样动作。import torch import torch.nn as nn import torch.optim as optim import gymnasium as gym class PolicyNet(nn.Module): def __init__(self, state_dim4, hidden_dim128, action_dim2): super().__init__() self.fc nn.Sequential( nn.Linear(state_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, action_dim) ) def forward(self, x): logits self.fc(x) return torch.softmax(logits, dim-1), logits选择动作时我会用torch.distributions.Categorical(logits)来采样这样可以直接调用dist.log_prob(action)计算log概率比自己手动算softmax再取对数方便很多也不容易出错。4.2 训练主循环与回报计算这是整个实现里最容易写错的部分。很多初学者会把log_prob拼成一个大张量再统一乘回报结果维度对不上。我的习惯是用一个列表逐步累积每个时间步单独保存log_prob和reward。def compute_returns(rewards, gamma0.99): returns [] G 0 for r in reversed(rewards): G r gamma * G returns.insert(0, G) return torch.tensor(returns, dtypetorch.float32)训练主循环大致是这样env gym.make(CartPole-v1) policy PolicyNet() optimizer optim.Adam(policy.parameters(), lr1e-3) for episode in range(1000): state, _ env.reset() done False log_probs [] rewards [] while not done: state_t torch.FloatTensor(state).unsqueeze(0) probs, logits policy(state_t) dist torch.distributions.Categorical(logitslogits) action dist.sample() log_probs.append(dist.log_prob(action)) state, reward, terminated, truncated, _ env.step(action.item()) done terminated or truncated rewards.append(reward) if log_probs: returns compute_returns(rewards, gamma0.99) # 回报归一化实验稳定性的关键 returns (returns - returns.mean()) / (returns.std() 1e-8) loss [] for lp, G in zip(log_probs, returns): loss.append(-lp * G) loss torch.cat(loss).sum() optimizer.zero_grad() loss.backward() optimizer.step()这段代码跑在CartPole上通常几百个episode内能拿到200分左右的成绩CartPole的满分线。不过我要提醒一句CartPole本身太简单纯Policy Gradient能看到learning trend但看不出方差问题有多严重。真想体会方差控制的重要性换一个奖励波动大的环境或者把CartPole的reward改成带噪声的形式你会立刻发现纯REINFORCE的训练曲线像心电图一样乱跳。4.3 关键参数的经验值学习率CartPole上用1e-3比较稳超过5e-3容易发散连续控制任务建议从1e-4往下试。γ折扣因子0.99是通用默认值。任务越长γ越接近1任务很短比如几十步0.95也够用。每轮采样episode数REINFORCE每更新一次只用当前episode样本效率低。如果内存允许可以攒一波episode再更新梯度估计更稳。网络大小策略网络不需要很大CartPole这种任务两层128就够了。过大的网络反而容易让策略过早熵坍缩直接收敛到固定动作。5. 从REINFORCE到Actor-Critic与PPO改进到底改了什么5.1 Actor-Critic把baseline换成一个学出来的价值函数REINFORCE加baseline时我们用G_t - V(s_t)降低方差但这里的V如果也是采样的平均就不够准。Actor-Critic的思路是同时训练两个网络Actor输出策略Critic输出状态价值V(s)用Critic的输出去替代baseline。这个改进带来的直接好处是不用等整条episode结束就能更新。REINFORCE必须等轨迹走完才能计算G_t而Critic可以在每一步提供一个估计值结合TD误差做更新。换句话说REINFORCE是“考完试一起评讲”Actor-Critic是“每做一题老师就点评一句”。后者反馈频率高样本利用率自然上去了。5.2 PPO限制每一步别改太猛从REINFORCE到Actor-Critic方差问题缓解了但没有完全解决策略更新的稳定性问题。PPO的做法更直接在目标函数里加一项限制如果新旧策略的比值r_t(θ)离1太远就把这一部分的收益clip掉避免更新步长过大把策略推到悬崖边。代码上只在计算loss时多了一行clamp但效果立竿见影。我自己跑连续性控制任务时同样的环境PPO大概一万步能稳定收敛Actor-Critic要两万步还偶尔发散。这也是现在工业界几乎不用裸REINFORCE的原因——不是它不对是方差控制手段不够工程化。5.3 什么时候还值得用纯Policy Gradient既然PPO更好为什么还要学Policy Gradient我的观点很明确Policy Gradient是所有策略优化算法的“语法”PPO只是加了几条“修辞规则”。你如果连REINFORCE的梯度流和方差问题都没亲自体验过直接上手PPO出了问题根本不知道是该调clip范围还是调GAE的λ。另外有些场景PPO反而别扭——比如你只需要做一次模仿学习行为克隆加一个奖励项用REINFORCE思路几行代码就能算梯度或者你在做策略可视化想观察一个随机策略如何逐渐收紧概率分布裸策略梯度的熵变化曲线比PPO更清晰。理解最原始版本永远是最好的调试起点。6. 策略梯度训练的10个实操细节与常见坑6.1 回报计算里的方向从后往前才算得对倒推计算G_t时必须从episode最后一步开始往前累加不能正着从0开始累加。G_t r_t γ G_{t1}正着算意味着你要“知道未来”而这在在线采样里做不到。写代码时reversed(rewards)这一步别省否则回报序列全是错的。6.2 千万不要在no_grad里更新策略PyTorch有个常见坑为了省显存有些人会在采样时加torch.no_grad()包住前向传播。但如果策略网络同时要被更新log_prob必须保留计算图才能反向传播。我的做法是采样阶段直接调dist.sample()log_prob正常在前向图里只有在需要独立计算基线时才考虑no_grad。6.3 熵正则化别让策略太早“自以为懂了”策略梯度训练早期如果某个动作碰巧拿到高回报策略会迅速把概率集中到那个动作上导致探索停止后面发现更好策略的机会就没了。解决方法是损失函数里加一项-α * H(π(·|s))也就是熵正则化让策略在训练早期保持一定的随机性。我在实践里喜欢用α0.01起步训练到中后期逐步衰减到0。观察策略熵也是个好习惯——如果熵掉到接近0而平均回报还在涨问题不大如果熵掉到0但回报也停了那要警惕过早收敛。6.4 归一化的位置比想象中更讲究回报归一化我建议在整批数据上做而不是每条episode单独做。单独归一化会抹掉不同episode之间的整体好坏差异——本来就差的episode被强行拉成均值以上梯度方向会失真。一整批做标准化既能稳住量级又保留相对好坏比较。6.5 梯度裁剪给策略网络加个“保险丝”连续控制任务里奖励尺度浮动范围大梯度爆炸是家常便饭。torch.nn.utils.clip_grad_norm_(policy.parameters(), 0.5)这行代码几乎是我训练所有策略网络的标配。裁剪之后参数更新不会因为单个异常episode而跳飞。6.6 同一条轨迹里的动作概率是连乘损失求和要小心策略梯度公式里Σ_t log π(a_t|s_t)是一条轨迹上所有动作的对数概率之和。很多人在实现时会把log_prob直接相乘再求梯度这在数值上会下溢成0。用log域求和才是对的积的log等于log的和这也避免了浮点数精度问题。6.7 评估要看平均回报的滑动平均别盯着单次曲线策略梯度的单episode回报曲线波动极大。我看训练状态时只看最近50个episode的滑动平均顺便统计“最好10%的平均”和“最差10%的平均”之间的差距。差距越大说明方差问题越严重越需要baseline或调整奖励尺度。6.8 连续动作空间用高斯分布而不是softmax这一条专门写给做连续控制的人。策略网络的输出层应该输出均值μ和一个对数方差log_std动作从高斯分布中采样。一个容易忽略的细节是log_std初始化为-1~-0.5之间的负值比较合适太小会让策略过早固化太大会让动作乱逛。训练过程中也可以逐渐降低log_std让策略从探索过渡到利用。6.9 环境随机性越强越要多采样再更新如果环境触发条件很随机单episode的回报几乎不可信这时候应该攒一批episode再更新参数。不要觉得慢方差大时频繁更新只会让参数白跳。我习惯先跑50个episode估算一下回报的标准差标准差大于均值的一半时就加倍采样量。6.10 先复现基线结果再改自己的任务环境最后一条听起来像废话但每次还是会有人踩坑跑Policy Gradient之前先在标准环境上复现一遍论文或博客里的基线曲线。这一步能帮你确认自己的代码逻辑没问题再去替换成你自己的奖励函数和状态空间。我见过太多人直接在新环境上跑一看曲线不涨就开始调各种花哨trick最后发现是回报计算写错了——白白浪费一下午。我自己动手写第一个Policy Gradient实现时最深的体会是这个算法表面只有一行公式但真正让它能落地的东西全都在方差控制和工程细节里。理解数学推导是为了知道优化方向对不对而调baseline、归一化、熵正则这些操作才是决定能不能收敛的关键。如果你现在也被某个连续控制问题卡住不妨从这篇最小实现开始跑通一遍再把环境换成自己的实际任务慢慢感受策略梯度从一个朴素想法变成稳定工具的过程。