策略梯度(Policy Gradient)是强化学习里一个绕不过去的坎。不管你从 DQN 入门,还是从 Sarsa 入门,最后都会撞上这堵墙——只是早撞晚撞的区别。很多资料一上来就扔出策略梯度定理的公式,然后写一句"证明略",留下我在屏幕前一脸茫然。这篇文章我把推导过程一步一步掰开揉碎写清楚,从轨迹概率怎么写开始,一路推到 REINFORCE,再解释为什么需要 baseline,为什么大家都在用 advantage function。适合刚学完基础概念、对"策略梯度定理"知其然不知其所以然的读者。李宏毅老师的课、Sutton 教材、OpenAI Spinning Up,都可以用这篇当推导参考。
1. 先把基础设定理清楚:MDP、轨迹与目标函数
1.1 在推导之前需要明确的符号与假设
推导策略梯度之前,必须把符号刻在脑子里。强化学习环境通常建模成马尔可夫决策过程(MDP),核心是状态 s、动作 a、状态转移概率 P(s'|s,a)、奖励函数 r(s,a,s') 和折扣因子 γ。策略 π_θ(a|s) 是在状态 s 下选择动作 a 的概率分布,θ 是需要优化的参数。
一条完整交互轨迹写作 τ = (s_0, a_0, r_1, s_1, a_1, r_2, ...),在有限时域里到终止状态 T 结束。轨迹的累计回报 R(τ) = ∑_{t=0}^{T} γ^t r_{t+1},γ 在 0 到 1 之间,表示未来奖励的相对权重。这个设定决定了我们求导时是"对整个轨迹长度求期望",而不是对单个时间步求期望。
有一个细节必须一开始就说清楚:虽然策略是每个状态上动作的概率分布,但 Markov 性质保证了当前动作只影响当前和未来的状态,不影响过去的状态。这个"因果性"是后面把 R(τ) 换成 R_t 的重要依据,现在先放在这里,后面会展开。
1.2 目标函数 J(θ) 为什么写成期望形式
算法要优化的目标函数是 J(θ) = E_{τ~π_θ}[R(τ)],也就是所有可能轨迹的累计回报的期望。它衡量的是:如果把当前策略 π_θ 放进环境里走一遍,平均能拿多少回报。
为什么不用某一条具体轨迹的回报作为目标?因为轨迹本身是随机的。环境转移有随机性,策略采样也有随机性,同一条策略跑两次,拿到的回报往往不一样。优化一个随机变量没有意义,只有优化它的期望才有意义。
期望形式另一个好处是:期望可以用采样来估计。我们没法枚举所有轨迹,但可以让智能体实际和真实环境交互,采很多条轨迹,用它们的平均回报来逼近真实期望。策略梯度推导出的结果正好是一个期望形式,可以用蒙特卡洛采样来近似,这就让"更新参数"这个操作有了一条可落地的路。
1.3 其他目标函数形式与轨迹概率公式
除了期望回报,强化学习里还有其他目标函数写法,比如平均回报 η(π) = ∑_{t} E[r_t],或者固定起始分布下的价值函数 J(θ) = V^{π_θ}(s_0)。这些本质上等价,推导过程也类似。Sutton 的书里用 V^{π}(s_0) 作为目标,但推导思路完全一致——先展开期望,再对参数求梯度。
接下来写出轨迹概率公式,这是整个推导最基础的一步:一条轨迹出现的概率等于初始状态概率乘上所有时间步策略概率和状态转移概率的乘积,写作 P_θ(τ) = p(s_0) ∏_{t=0}^{T} π_θ(a_t|s_t) p(s_{t+1}|s_t, a_t)。把正体 P 读作"该条轨迹在策略 π_θ 和环境动力学共同作用下被采样到的概率"。后面求梯度时,这个乘积公式的对数会被展开成求和形式,每一项都能单独处理。
2. 核心推导过程:梯度如何一步一步走到可计算的形式
2.1 从期望梯度出发:对数似然技巧
现在正式开始推导。目标函数的梯度是 ∇_θJ(θ) = ∇_θE_{τ}[R(τ)]。把期望展开成积分形式:
∇_θJ(θ) = ∇_θ∫P_θ(τ)R(τ)dτ
积分符号和期望符号在这里可以互换:期望就是概率密度乘以回报的积分。但这个积分里有两个地方依赖 θ——轨迹概率 P_θ(τ) 和隐含的回报 R(τ)。直接对积分求导会很麻烦,因为积分里面的概率密度本身是参数化的。
这里用到强化学习推导中最核心的一个数学技巧——对数似然技巧(log-derivative trick):对任意函数 f,有 ∇_θf = f·∇_θlog f,前提是 f > 0。概率密度函数满足这个条件。于是把 ∇_θP_θ(τ) 替换成 P_θ(τ)·∇_θlogP_θ(τ):
∇_θJ(θ) = ∫P_θ(τ)·∇_θlogP_θ(τ)·R(τ)dτ = E_{τ~π_θ}[∇_θlogP_θ(τ)·R(τ)]
从积分形式回到期望形式,这一步完成了从"无法解析求导"到"可以采样估计"的关键转换。这里的核心逻辑是:我们不需要知道环境动态的具体解析式,只需要策略本身——因为策略是我们能控制的,可以求出梯度。而环境动态的部分,在对数轨迹概率展开后会被消掉,这正是下一步要做的事。
2.2 展开轨迹概率的对数:环境动态如何被消掉
把上一节得到的梯度公式里的 ∇_θlogP_θ(τ) 展开。对轨迹概率取对数,乘积变求和:
logP_θ(τ) = logp(s_0) + ∑_{t=0}^{T} logπ_θ(a_t|s_t) + ∑_{t=0}^{T} logp(s_{t+1}|s_t,a_t)
第一项 logp(s_0) 是初始状态分布,由环境决定;第三项是状态转移概率的对数,同样由环境动力学决定。这两项都不含参数 θ,对 θ 求导的结果为 0。真正留下的只有中间那项——所有时间步上策略动作概率的对数之和。
于是梯度公式简化为:
∇_θJ(θ) = E_{τ~π_θ}[∑_{t=0}^{T}∇_θlogπ_θ(a_t|s_t)·R(τ)]
这就是策略梯度定理最原始、最核心的表达式。展开轨迹概率后,环境动态神奇地消失了——我们只需要知道策略的梯度 ∇_θlogπ_θ(a_t|s_t),完全不需要知道状态转移概率 p(s'|s,a)。这个特性是无模型(model-free)强化学习能够成立的根本原因:我们不建模环境,只优化策略本身的参数。
这个公式的直觉可以这样理解:它是在说,如果一条轨迹的回报 R(τ) 是正的,就把这条轨迹上所有动作的概率调高;如果回报是负的,就调低这些动作的概率。调整的幅度正比于回报绝对值,也正比于策略对动作概率的梯度方向。整体就是一个"加权最大似然"——用轨迹回报作为权重,去最大化所有被采到的动作的 log 概率。
2.3 梯度表达式的直观理解:为什么"概率上升幅度正比于回报"
把上一节的公式再品一品:∇_θJ(θ) = E[∑_t ∇_θlogπ_θ(a_t|s_t)·R(τ)]。想象你是一个策略的"管理者",你拿到了很多条轨迹和它们的回报。高回报的轨迹告诉你:"这条轨迹上的动作序列是好的,把它们出现的概率调高";低回报的轨迹告诉你:"这条轨迹上的动作序列不行,把它们出现的概率调低"。
但这里有个容易被忽略的地方:如果当前环境的奖励全为正(比如每一步都有小的正奖励,只在结束时给一个很大的正奖励),那么所有轨迹的回报都是正的,高回报轨迹会得到更大概率提升,低回报轨迹也会被提升——只是幅度小一些。这会导致所有动作概率都有上升趋势,策略的熵容易变大。好在归一化本身会平衡这一点:某个动作概率调高,其他动作概率自动就会被压缩。真正到实现阶段,reward 的偏移还是会影响实际效果,处理办法在后面的第 4 节会展开讨论。
3. 从理论推导到算法实现:REINFORCE 与因果性的优化
3.1 REINFORCE 算法的完整流程
有了上面的梯度表达式,最直接的算法就是 REINFORCE,也叫蒙特卡洛策略梯度。它的做法是:用当前策略 π_θ 在环境里采样一整条轨迹,然后对轨迹上的每一个时间步,用"该时间步的策略梯度 × 整条轨迹的回报"来更新参数。
REINFORCE 算法流程如下:
- 初始化策略参数 θ。
- 用 π_θ 在环境中采样一条完整轨迹 τ。
- 对轨迹上的每个时间步 t:
- 计算该时间步的动作对数概率梯度 ∇_θlogπ_θ(a_t|s_t)。
- 计算累计回报 R(τ)。
- 累加梯度:g_t = ∇_θlogπ_θ(a_t|s_t)·R(τ)。
- 更新参数 θ ← θ + α·∑_t g_t,α 是学习率。
- 回到第 2 步重复,直到收敛。
上面是最朴素的版本。实际代码里,为了稳定性,一般会跑一个批量(batch)的轨迹,把多条轨迹中相同参数的梯度加总平均,再做一次梯度上升。批量大小至少 8 条、16 条,效果会稳很多。后面第 5 节我会给出一段可运行的参考代码。
注意第 4 步用的是梯度上升(因为我们是在最大化回报期望),方向和"最小化损失"的梯度下降相反。如果用通常的深度学习框架,有两种做法:一种是对"负回报 × logπ"求梯度,然后调用的 loss.backward() 正常做梯度下降;另一种是对"回报 × logπ"求梯度,但在实现时手动把梯度取反。推荐第一种,和框架的自动求导结合得更自然。
3.2 关键优化:为什么用回报 R_t 而不是整条轨迹回报 R(τ)
既然公式推导出来的是每一项都乘以 R(τ),那 REINFORCE 为什么可以用 R_t = ∑_{k=t}^{T}γ^{k-t}r_{k+1} 替代?这是策略梯度推导中的一个关键细节,也是很多初学者混淆的地方。
答案是因果性:动作 a_t 只会影响 t 时刻之后的环境状态和奖励,不可能影响 t 时刻之前的奖励。梯度 ∇_θlogπ_θ(a_t|s_t) 乘以前面的奖励 r_1、r_2、...、r_t 的期望为 0,因为那些奖励的期望不依赖于当前动作。既然期望为 0,把它们从求和里删掉,梯度的期望不变。
形式化一点:设 R(τ) = ∑_{k=0}^{T}γ^k r_{k+1},把它拆成两段——k<t 段和 k≥t 段。那么 E[∇_θlogπ_θ(a_t|s_t)·∑_{k<t}γ^k r_{k+1}] = 0,因为 ∑_{k<t}γ^k r_{k+1} 是关于状态转移和过去动作的随机量,这些都与 θ 对当前动作的梯度无关(这里做一个简化的条件期望论证即可)。所以梯度表达式可以改写成:
∇_θJ(θ) = E[∑_t ∇_θlogπ_θ(a_t|s_t)·R_t]
其中 R_t 是从 t 时刻起未来折扣回报,也叫 return-to-go。这个改进的价值在于:它减小了梯度估计的方差。乘的东西越少,噪声越小;乘进去无关的过去奖励项,相当于加了一个与真实目标无关的随机扰动。R(τ) 换成 R_t 之后,方差显著下降,算法更容易收敛,这在实验里的效果非常明显。
3.3 折扣因子 γ 在推导中的位置
读者可能已经注意到,推导时 R(τ) 里挂着 γ^t,而实际实现时通常使用 R_t = ∑_{k=t}^{T}γ^{k-t}r_{k+1},γ 指数从 1 开始(t 时刻自身的奖励不折扣)。这是因为折扣因子γ^k 的两个作用:
第一,它是收敛性保证。无限时域下累计回报可能是无穷大的,折扣因子保证期望回报有界,这是数学上良好定义的需要。第二,它是方差调节器。γ 越小,远期奖励权重越低,传入梯度的噪声越小,但可能牺牲掉一些长期信息;γ 越接近 1,越重视长期回报,但也引入更大的方差。实际使用中 γ 常取 0.99 或 0.995,连续控制任务里为了稳定甚至取 0.95 左右,具体值要看任务的时域长度。
这里有一个推导时容易混淆的点:γ^t 是轨迹总回报自身带的折扣系数,而 R_t 的定义是未来回报的折扣和——它不再含全局的时间基准 t,而是以当前时刻为起点。两个表达殊途同归,本质都是"施加因果性之后把指数拆干净"。推导时建议从 R(τ) 拆成两段开始,最后把后一段写成 R_t,逻辑最顺。
4. 高方差问题与 baseline:为什么减去一个常数梯度期望不变
4.1 策略梯度为什么方差很大
虽然现在梯度表达式已经能用采样估计,但直接使用时方差大得离谱。原因不难理解:策略梯度是蒙特卡洛估计,用的是"几条有限轨迹去估计所有可能轨迹空间的期望"。轨迹空间是指数级的,假设每个状态有 10 个可选动作,一条 20 步的轨迹就有 10^20 种可能。你只采了几十条轨迹,估计出的梯度自然有巨大的噪声。而且回报 R_t 本身的方差可能很大——同一个状态下采取同一个动作,后续环境随机性会让最终累计回报相差很大。
可以类比扔骰子:理论上每个面概率 1/6,你只扔 6 次,估计出来的概率分布波动很大;但如果你每次扔之前先随机决定"多算一次"还是"少算一次",噪声更更大了。实验里的典型表现是:训练曲线剧烈震荡,loss 一会儿上天一会儿入地,有时甚至看不出上升趋势。
4.2 baseline 的构造与无偏性证明
既然方差来自"回报的绝对量级",一个自然的想法是:把回报做一个平移,让正负更对称,但不改变梯度的期望方向。这个平移量就是 baseline,记为 b(s_t)。
考虑新的梯度表达式:E[∇_θlogπ_θ(a_t|s_t)·(R_t - b(s_t))]。我们希望添加 b(s_t) 后梯度期望不变。证明如下:
E_{a~π}[∇_θlogπ_θ(a|s)·b(s)] = b(s)·∑_a ∇_θπ_θ(a|s) = b(s)·∇_θ∑_aπ_θ(a|s) = b(s)·∇_θ1 = 0
这个恒等式非常漂亮。它说的是:梯度的期望里乘上一个只依赖当前状态、不依赖动作的常数,在期望意义下就是 0。所以无论 baseline 取什么值,梯度期望都不会改变,这就是无偏性。
它之所以成立,是因为对概率归一化 ∑_aπ_θ(a|s)=1 求导等于 0。这个看似平凡的恒等式是策略梯度降方差的基石:你可以在不改期望的前提下自由选择 baseline,来减小方差。
4.3 最优 baseline 的推导与常见选择
在保证无偏的前置下,什么 baseline 能让方差最小?设 g = ∇logπ·(R - b),其方差的近似(忽略不同维度间的协方差)正比于 E[(R - b)^2·‖∇logπ‖^2] 减掉一个常数。为了保证期望不变,把 b 当作变量,对 E[(R - b)^2·g^2] 求极小值,得到 b* = E[R·‖∇logπ‖^2] / E[‖∇logπ‖^2]。但这里 g 本身依赖 b,想要精确求 b* 并不容易。
实际中最常见的 baseline 选择是状态价值函数 V(s_t),即 b(s_t) = V(s_t)。代入后得到 R_t - V(s_t),这就是 advantage function A(s_t, a_t) = Q(s_t, a_t) - V(s_t) 的蒙特卡洛估计。直觉上,V(s_t) 表示"从状态 s_t 出发平均能拿多少回报",R_t 是"本次实际跑出来的回报",二者相减得到"这次动作相对平均水平好不好"。这样做梯度就有了正负,好动作往上抬,差动作往下压,训练会更稳定。
另一个常见做法是用 running mean 作为 baseline:在整个训练过程中维护一个回报的滑动平均,用 R_t 减去这个均值。这个实现最省事,效果也不错,但理论依据弱一些。基线还可以用可学习的神经网络来做,比如 A2C/A3C 里使用 critic 网络 V(s_t) 作为 baseline,配合 actor 共享特征提取器,是实践中效果最好的方案之一。
4.4 baseline 参数忽略的推导技巧
在做带 baseline 的推导时,需要注意一个细节:V(s_t) 也依赖参数 θ(如果 V 是同一个参数化的网络),那对 θ 求梯度时 V 也会产生梯度。但我们的推导把它当作常数 baseline 看待,不计算它的梯度。这就是"stop-gradient"的来源。
如果坚持把 V 的梯度也算进去,会出现 actor-critic 算法中 critic 与 actor 目标函数相互纠缠的复杂局面,推导会变得不干净。实践中更常见的做法是:critic 网络独立于一节参数,通过监督回归去逼近 V(s_t),即最小化 (V(s_t) - R_t)^2;actor 则使用 stop-gradient 的 V 作为 baseline。我把这个分开计算的方式理解为"各自的梯度走各自的路":actor 的梯度只来自策略对数概率乘 advantage,critic 的梯度只来自值函数的状态价值估计误差。这样两条梯度的分工明确,训练更稳定。
5. 从推导到实战:REINFORCE 的代码框架与调参细节
5.1 一个可直接运行的 REINFORCE 参考实现(PyTorch)
理论推导讲完了,落到代码上。下面这段代码是一个最简的 REINFORCE 实现,只用 policy 网络,不做 baseline(对应第 3 节的算法),便于理解核心逻辑。
import torch import torch.nn as nn import torch.optim as optim class PolicyNet(nn.Module): def __init__(self, state_dim, action_dim, hidden=64): super().__init__() self.fc = nn.Sequential( nn.Linear(state_dim, hidden), nn.ReLU(), nn.Linear(hidden, hidden), nn.ReLU(), nn.Linear(hidden, action_dim) ) def forward(self, s): return torch.softmax(self.fc(s), dim=-1) 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) def train_one_episode(env, policy, optimizer, gamma=0.99): log_probs = [] rewards = [] state, _ = env.reset() done = False while not done: state_t = torch.as_tensor(state, dtype=torch.float32) probs = policy(state_t) dist = torch.distributions.Categorical(probs) action = dist.sample() log_probs.append(dist.log_prob(action)) next_state, reward, done, truncated, _ = env.step(action.item()) done = done or truncated rewards.append(reward) state = next_state returns = compute_returns(rewards, gamma) # 累积损失 = -∑ logπ(a_t|s_t) * R_t policy_loss = [] for log_p, G in zip(log_probs, returns): policy_loss.append(-log_p * G) loss = torch.stack(policy_loss).sum() optimizer.zero_grad() loss.backward() optimizer.step() return sum(rewards), loss.item()代码里最值得关注的是 loss 构造:-log_p * G里有个负号,因为优化器用的是梯度下降,而我们是要最大化期望回报,所以用负的期望目标作为 loss。G用的是每一步算出来的 return-to-go,对应第 3.2 节的推导结论。对数概率 log_p 是从策略分布里采样出来的动作的概率对数,loss.backward()会自动求出 ∇logπ(a|s) 乘上系数 G 的梯度。
5.2 回报归一化:最简单有效的方差降噪手段
REINFORCE 想跑得动,第一件事就是做回报归一化。代码里把compute_returns产生的 returns 做一下标准化即可,比如returns = (returns - returns.mean()) / (returns.std() + 1e-8)。
为什么要归一化?因为原始回报的绝对量级可能很大也可能很小,且往往是正偏态的。如果所有回报都是正的且数值较大,梯度始终朝"增大所有动作概率"方向走,策略的熵会持续增加——这和最大化期望回报的目标冲突。归一化后回报均值为 0、标准差为 1,一部分梯度走向高回报(正权重),另一部分走向低回报(负权重),训练方向和幅度都更均衡。
注意:归一化会改变梯度的绝对幅度,但不改变梯度方向的比例关系。它本质上属于"重新缩放"操作,会把梯度的量纲统一到一个合理范围。批量训练时可以在 batch 内做归一化,效果比单条轨迹做要好,因为 batch 内的统计量更稳定。
5.3 我在实践中踩过的坑
坑一:回报归一化的时机。如果在一整条轨迹上先做了折扣累计再归一化,和先归一化再乘折扣,两者结果完全不同。正确顺序是:先按折扣因子算好 R_t = ∑γ^(k-t)r_k,再做标准化。不要先把每一步的稀疏奖励归一化,再乘折扣——那会破坏折扣因子的时间语义。
坑二:学习率过大,策略会在几个 episode 内"崩溃"到确定性分布。这个现象的机制是:某个动作恰好采到高回报,梯度把它概率推到接近 1,后续所有采样都集中在这个动作上,再也没有探索机会,策略就卡死在局部最优。解决办法是把学习率调小,REINFORCE 通常需要比 DQN 更小的学习率,1e-3 到 3e-4 的量级比较常见;同时确保采样轨迹数足够,单轨迹更新几乎一定会出问题。
坑三:baseline 本身也是网络时,要记得 stop-gradient。在 PyTorch 里用V.detach(),TensorFlow 里用tf.stop_gradient,否则基线网络产生的梯度会回流到 actor 主干,和策略梯度混在一起,训练动态会非常糟糕。这正是第 4.4 节讨论的问题在代码里的直接对应。
下表总结了几种常见做法和它们对收敛性的影响:
| 做法 | 梯度期望 | 对训练稳定性影响 |
|---|---|---|
| 用 R(τ) 作为权重 | 不变 | 方差大,难收敛 |
| 用 R_t (return-to-go) | 不变 | 明显改善,常配合 γ 使用 |
| R_t - V(s_t) | 不变 | 方差更小,最常使用 |
| (R_t - mean(R_t))/std | 有偏(改方向比例)但常用 | 工程上最稳,提升显著 |
5.4 调试时的诊断指标与观察要点
策略梯度跑不跑得起来,不要只盯着回报曲线。我建议至少同时记录三样东西:每步策略的熵、平均 log_prob、以及回报的均值与标准差。熵反映了策略的探索能力,如果训练刚开始熵就掉到接近 0,说明策略正在过早确定化,准入死胡同的概率很高;如果熵一直很大,说明策略几乎没学到东西,分布接近均匀。平均 log_prob 如果整体持续下降,说明策略越来越确定,这是好事;但如果下降太剧烈,要警惕过拟合到某几条轨迹上。
回报的方差也很关键。如果回报标准差比均值还大,说明要么是环境的随机性太大,要么是采样数太少,这时增大 batch size 比调学习率更有效。有一个经验性观察:对 CartPole 这样的简单环境,批量 8 条轨迹、γ=0.99、学习率 1e-3 就能在几百幕之内达到目标;但对稍微复杂的环境,REINFORCE 就显得力不从心,这时候就要考虑加入 baseline 和 critic,或者直接上 PPO。
5.5 从 REINFORCE 到 PPO:策略梯度家族的自然演化
掌握了策略梯度推导,再看 PPO、TRPO、A2C/A3C 等进阶算法时会轻松很多——它们的核心都建立在策略梯度的基本表达式上。TRPO 是在更新步长上做文章,用 KL 散度约束保证每次更新不翻车;PPO 用 clipped surrogate objective 近似 TRPO 的约束,实现更简单;A2C/A3C 的核心是引入了 critic 网络作为 baseline 且多环境并行采集,从根本上降低梯度的方差。
也就是说,第 2 节推出来的那个核心公式 ∇J(θ) = E[∑_t ∇logπ(a_t|s_t)·A_t],A_t 是 advantage,是大部分现代策略梯度算法的骨架。A_t 的具体算法可以是蒙特卡洛(REINFORCE)、可以是 TD 方式(A2C)、可以是 GAE(PPO 中常见),万变不离其宗。理解了这个推导,后面看 PPO 的论文代码都会顺畅得多。
我的个人体会是,学习策略梯度推导不要跳过"因果性"这一步——它决定了你写代码时天然想到用 R_t,而不是照搬公式傻乎乎乘 R(τ)。另一个体会是 baseline 的推导值得亲手算一遍,把那个恒等式 E[∇logπ·b] = 0 推导三遍,比背十遍公式都有用。实际跑算法时,归一化回报和合适的学习率,比换模型结构带来的收益大得多。弄清楚这些细节,你才算真正跨进了深度强化学习的大门。