ARTICLE · INTELLIGENCE

战地情报 · 详情页

来自尧图项目组的一线实战观察与深度解析

策略梯度入门:从REINFORCE推导到实现与调试

策略梯度入门:从REINFORCE推导到实现与调试

最近在啃强化学习,翻到第八章第18页 Policy Gradient 的时候,我明显感觉到自己撞在了一堵墙上。前面 Q-learning、DQN 这些内容,好歹还有“表格”“近似函数”这种看得见摸得着的抓手,可策略梯度一上来就给你满页的期望符号、对数概率、梯度展开,第一遍读下来基本是滑过去的。后来我拿着纸笔把每一步推导拆开重写,又动手写了一个最朴素的 REINFORCE 跑在经典控制环境上,才算把这条技术线真正理顺。这篇就把我对 Policy Gradient 的完整理解倒出来,从“它到底在解决什么问题”讲起,到那条核心公式的推导直觉,再到一份能跑通的极简实现,以及我实际调试中踩过的坑。适合已经了解 Q-learning、正在往策略梯度方法过渡的人,也适合想在项目里快速上手决策控制算法的工程师。

1. 为什么价值函数方法不够用:策略梯度的出场逻辑

1.1 价值函数方法的两道天花板

价值函数方法,包括 DQN 这类深度强化学习算法,核心思路是先估计每个状态-动作对的价值,然后从价值函数里“挑”出最优动作。这套路在游戏场景里确实很能打,但你往现实问题里放就会发现它有两个很普通的限制。

第一个限制是连续动作空间。假设你要控制一个机械臂的关节力矩,每个关节的输出就是一个连续实数,比如 -1 到 1 牛米。价值函数方法需要比较所有动作的 Q 值大小,在连续空间里你没法枚举所有动作——总不能把力矩按 0.01 间隔全部遍历一遍,关节一多维度直接爆炸。现实中的机器人控制、自动驾驶、能源调度,动作空间几乎都是连续的,这是价值函数方法一道实打实的天花板。

第二个限制更隐蔽:价值函数方法天然输出确定性策略,也就是“在某个状态,最优动作是固定的那一个”。但很多问题的最优策略本来就是随机策略。最经典的例子是石头剪刀布,如果策略是确定性的,对手摸到规律之后你就一把都赢不了。再比如不完全信息博弈、信息受限的感知决策场景,随机性本身就是策略的一部分。Q-learning 里的 epsilon 贪心只能算一个临时补丁,它让策略偶尔随机一下是为了探索,但本质上还是把最优动作找出来,而不是真正学会一个带概率分布的决策方式。

1.2 把策略本身参数化成一条可调的曲线

Policy Gradient 换了一个完全不同的思路:既然价值函数方法绕了一道“先估价值、再选动作”的弯,那还不如直接把策略本身参数化,让策略输出每个动作的概率。假设我们用 θ 表示策略参数,那么策略就是一个概率分布 π_θ(a|s),这个分布在状态 s 下给出选择动作 a 的概率。

目标也随之改变。价值函数方法追求的是让 Q 值尽可能逼近真实价值,而 Policy Gradient 直接优化一个更原始的指标——策略的期望累积回报:

J(θ) = E_{τ~π_θ}[R(τ)]

这里 τ 表示一条完整轨迹,也就是从初始状态开始,策略与环境交互到结束的全过程;R(τ) 是这条轨迹上所有折扣奖励的总和。整个目标的含义非常直觉:我们希望找到一组参数 θ,让“按这个策略跑出来的轨迹平均回报”最大。

到这里,问题变成了一个标准的“黑箱优化”问题:决策过程内部包含环境转移、状态反馈,我们往往不知道环境的完整动态模型,但可以通过采样来估计这条轨迹的回报。剩下的核心就是:在不知道环境模型的情况下,如何求出 J(θ) 对参数 θ 的梯度?这正是策略梯度方法整条推导线的起点。

1.3 为什么梯度方向比随机搜索靠谱

有人可能会问:不就是让平均回报变大吗,我随机给参数加噪声、跑一批轨迹看哪个参数表现好,保留好的参数不就行了?这确实是进化策略一族干的事情,简单直接,但样本效率低到令人发指。参数维度一高,随机搜索在庞大参数空间里漫无目的地试探,绝大多数扰动都踩在无用的方向上。

梯度方法的优势在于,它通过策略概率对参数的导数告诉你“在当前参数附近,往哪个方向调整能让期望回报增长得最快”。即使每次采样的梯度估计有噪声,平均下来它依然指向真实梯度的方向。这就相当于你站在山坡上有根拐杖,每次往前试探一下,脚底传来的坡度信息能引导你往上爬,而不是闭上眼睛瞎跑。策略梯度的整个理论框架,其实就是为了把这根拐杖造出来。

2. 策略梯度核心公式的推导直觉

2.1 log-trick:把梯度挪进期望里面

推导的核心难点在于,轨迹的分布由当前策略决定,采样过程本身不可微。你无法直接对“样本”求梯度,但可以对“采样概率”求梯度。这里用到一个在统计里非常常见的技巧,叫做对数概率梯度恒等式:

∇_θ π_θ(a|s) = π_θ(a|s) · ∇_θ log π_θ(a|s)

这个恒等式看着简单,但它解决了策略梯度的大半问题。把左边那个“对概率求梯度”改写为右边“概率值乘上对对数概率求梯度”之后,我们就可以把梯度搬进期望号里面。整个过程大致是:

∇_θ J(θ) = ∇_θ ∫ P(τ|θ)R(τ)dτ = ∫ P(τ|θ)∇_θ log P(τ|θ)R(τ)dτ = E_{τ~π_θ}[∇_θ log P(τ|θ)·R(τ)]

然后进一步展开 log P(τ|θ),轨迹的概率由策略和环境动态共同决定。展开后可以发现,环境动态部分不依赖于参数 θ,求梯度时直接归零,真正留下来的只有策略项:

∇_θ J(θ) = E_{τ~π_θ}[ Σ_{t=1}^T ∇_θ log π_θ(a_t|s_t) · R(τ) ]

这个公式漂亮就漂亮在:环境模型完全不需要出现在表达式里。哪怕你不知道状态怎么转移,只要你能采样轨迹,就能用观测到的轨迹去估计这个梯度。

2.2 REINFORCE:整条轨迹的成绩单

基于上面这个公式,最简单的策略梯度算法就是 REINFORCE,它的执行流程非常朴素:

  1. 用当前策略 π_θ 采样一条完整轨迹;
  2. 计算这条轨迹的折扣总回报 R(τ);
  3. 对轨迹上的每一个时间步,计算 ∇_θ log π_θ(a_t|s_t) 并乘以 R(τ);
  4. 把所有时间步的梯度加总取平均,作为这次采样的梯度估计;
  5. 沿梯度方向更新参数:θ ← θ + α·梯度;
  6. 回到第 1 步,循环迭代。

把公式翻译成人话:每条轨迹就像一张成绩单。如果某条轨迹的总回报是正数,我们就奖励这条轨迹上出现过的所有动作,让它们未来的出现概率都升高;如果总回报是负数,我们就惩罚这些动作,让它们的概率降低。所有动作共享同一张成绩单,一荣俱荣、一损俱损。

2.3 高方差的根源:用少数样本来猜全局

理解了流程,你就能理解为什么策略梯度方法有一个臭名昭著的毛病——方差大。因为我们需要用采样轨迹的回报来估计期望梯度,而每条轨迹的回报本身波动就非常大。同一个策略跑两条轨迹,一条可能拿到了高分,另一条可能早早失败,两者的回报差异巨大。REINFORCE 把整条轨迹的总回报平均地分给每一个时间步,这里其实存在两个问题:第一,轨迹之间的随机性没有做任何平滑;第二,某个动作明明只产生了 3 步之后的效果,却要承担整条轨迹所有奖励的功劳或责任,这一点在技术上叫做 credit assignment 问题。

方差大带来的直接后果是训练不稳定:明明策略在变好,某一次采样运气差了梯度方向就可能完全跑偏;运气好了,又可能让策略被一次偶然的高分轨迹裹挟。所以,策略梯度算法的全部演进史,基本上就是在跟“方差”这两个字做斗争。

3. 从零实现一个能跑的 REINFORCE

3.1 实验设定与网络结构

我建议第一次实操选一个轻量级的控制任务,比如经典的倒立摆环境(CartPole)。这个环境状态是连续的,动作是离散的,很适合用来验证策略梯度基础实现。网络不需要复杂,两层全连接就够:输入层接状态向量,中间层用 ReLU,输出层给两个动作的 logits,再接一个 softmax 得到动作概率分布。以下是我用的简化结构:

import torch import torch.nn as nn from torch.distributions import Categorical class PolicyNet(nn.Module): def __init__(self, obs_dim=4, hidden_dim=128, act_dim=2): super().__init__() self.net = nn.Sequential( nn.Linear(obs_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, act_dim) ) def forward(self, obs): logits = self.net(obs) return Categorical(logits=logits)

这里的输出是一个 Categorical 分布对象,它会根据 logits 自动算出各个动作的概率,同时提供 sample() 方法和 log_prob() 方法,正好对应 REINFORCE 里需要的“采样”和“算对数概率梯度”。

3.2 训练闭环:采样、算回报、更新

训练循环就是采样一条轨迹,保存每一步的动作对数概率和奖励,然后反向传播更新。核心在于奖励的处理,这一步也是新手最容易写错的地方。折扣回报需要从轨迹末尾往前逆推:

def compute_returns(rewards, gamma=0.99): returns = [] G = 0 for r in reversed(rewards): G = r + gamma * G returns.insert(0, G) return torch.tensor(returns, dtype=torch.float32)

损失函数看起来非常反直觉——为什么最大化回报要写成“最小化某个带负号的式子”?原因是我们习惯用深度学习框架的 loss.backward() 做梯度下降,所以要取负号把“上升”变成“下降”:

loss = -(log_probs * returns).mean() loss.backward() optimizer.step()

其中 log_probs 是整条轨迹上每个动作的对数概率张量,returns 是对应的折扣回报。用乘法对齐之后取平均,就是这个 batch 的近似梯度。

3.3 关键的一步:回报归一化

我强烈建议在算 loss 之前把 returns 做一次标准化,也就是减均值除以标准差:

returns = (returns - returns.mean()) / (returns.std() + 1e-8)

这一步不在原始 REINFORCE 的推导里,但属于实践中的标准做法。因为回报的绝对值大小会影响梯度更新的幅度,不同任务、不同环境的回报尺度差异很大,直接乘上去很容易让更新步子时大时小。归一化之后,训练收敛明显更稳。这个操作可以看做一个最简单的 baseline 技巧,它不会改变梯度的无偏性,但能显著压低方差。

3.4 三个我踩过的坑

第一个坑是用同一条轨迹数据跑多次梯度更新。REINFORCE 的梯度估计依赖当前策略的采样分布,轨迹是从旧策略采出来的,数据和策略已经对不上了。你多更新几步,就是在拿过期的数据训练新策略,很容易发散。正确做法是一条轨迹采完,更新一次,扔掉,重新采样。

第二个坑是学习率贪心。REINFORCE 的梯度估计方差本来就大,学习率稍微调高一点就会剧烈震荡。我第一次跑的时候设了 0.01,曲线像心电图一样疯狂跳动。后来降到 1e-3 才稳定。这个算法真的吃不了大补药。

第三个坑是不看策略熵。策略训练到后期可能退化成确定性策略,输出概率变成 1 和 0,熵归零。这种情况下训练曲线可能还在涨,但策略已经失去了探索能力,一旦环境有扰动就会脆断。每次训练循环里顺手打印一下策略熵,能帮你及时发现这个苗头。

4. 常见问题排查与调参实战

4.1 训练不收敛,先查这几个方向

如果奖励曲线连续很多个 episode 都没有爬升趋势,我一般先查 gamma 和学习率。gamma 太大会让长期回报的方差不断累积,太小又会让算法只盯着眼前奖励,学不到长程决策。倒立摆这类短 horizon 任务,0.99 是个比较稳的起点,但真正合适的值还是得自己扫几轮。

另一个容易忽略的问题是环境本身的终止条件。倒立摆的奖励到 500 步会截断,如果你不处理这种“截断”和“真正失败”的区别,算法很容易把 500 步截断当成高分回报,产生虚假的梯度信号。处理办法是在计算回报时对截断状态做特殊标记,或者简单一点,把环境的最大步数设短一点,让算法更频繁地看到完整的失败轨迹。

4.2 方差压缩的三板斧

第一板斧是用 reward-to-go 替代整条轨迹回报。整条轨迹的总回报把早期奖励也算进了当前动作的功劳簿,但从因果上讲,当前动作只影响当前及未来的奖励。所以更合理的做法是让每个动作只对“从当前时刻开始到轨迹结束”的累积回报负责,环境回馈的过往奖励不该算进来。

第二板斧就是 baseline。最省事的 baseline 是回报均值,因为数学上可以证明减去一个不依赖动作的常数不会改变梯度的期望,但能显著降低方差。我前面说的归一化操作就是它的一种实现。

第三板斧是往损失里加策略熵正则项,损失函数变成:

loss = -(log_probs * returns).mean() - beta * entropy.mean()

entropy 是当前策略分布的熵,beta 是一个很小的系数,比如 0.01。加这项的目的是惩罚“过于确定的策略”,强迫策略保留一定的随机性。这不会直接压缩方差,但能防止策略过早收敛到局部确定性策略,间接提升稳定性。

4.3 真正有用的几个 Debug 手段

比起整天盯着平均奖励曲线,我更推荐打印三个量:每个 episode 的策略熵、动作对数概率的梯度的范数、以及回报归一化之后的均值标准差。策略熵如果快速掉到接近 0,先看是不是学习率太大;梯度范数如果忽大忽小,看是不是回报尺度没处理好。

还有一个经验:把每次 episode 的步数画成散点图,比画 loss 曲线直观得多。在倒立摆环境里,episode 步数就是最直接的学习进度条。你会发现这个曲线不是平滑上升的,而是充满毛刺,但那不代表算法坏了——策略梯度本来就允许表现有波动,只要整体趋势往上,就说明方向是对的。我第一次看到这种锯齿状曲线时差点放弃调参,后来才意识到这是策略梯度的工作方式。

5. 写在最后的个人体会

我实操完 REINFORCE 之后最大的收获,不是看到奖励曲线终于爬上去那一刻的兴奋,而是终于明白了为什么策略梯度之后还会演化出 PPO、A3C 这一大堆方法。所有这些改进,本质上都是在回答同一个问题:怎么在保持梯度无偏的前提下,把方差压得更低。理解了 REINFORCE 的缺陷,再去看后续那些算法,会顺畅很多。

所以如果你也正卡在这个章节,我的建议很简单:第一,动手把那个梯度公式从目标函数开始推一遍,别只看结论;第二,把最朴素的 REINFORCE 写在纸上、跑在机器上,哪怕它跑得并不漂亮。这两件事做完,你再看策略梯度这家人的其他方法,心态会完全不一样。最后再分享一个小技巧:打印日志时把“当前策略熵”带上,这是我排查训练异常时最先看的指标,比奖励值可靠得多。

RELATED READING

延伸阅读

更多一线实战笔记与深度复盘,助您持续精进