零基础速通 PPO 学习笔记

学习背景

最近套磁到了港科的一个 RA 岗位，组里在做 Agent 强化学习相关的研究，按照学长的建议先把 PPO 算法系统学习一遍，于是有了这篇笔记。

符号定义

在强化学习中，常用符号约定如下：

$a_t$ ： $t$ 时刻 Agent 采取的动作（action）；
$s_t$ ： $t$ 时刻 Agent 所处的状态（state）；
$\pi$ ：Agent 的策略函数（policy），输入状态，输出每个动作的概率分布 $\pi(a_t \mid s_t)$ ；
$r_t$ ： $t$ 时刻 Agent 采取动作后获得的奖励（reward）；
$\tau$ ：一条轨迹（trajectory），即一段状态-动作序列；
Episode：一次完整的交互过程，从环境初始化开始，到达到终止状态为止；
Rollout：按当前策略实际"跑出来"的一段轨迹数据，不一定是完整的一局，也可以只是其中的一段。

$(s_0,a_0,s_1,a_1 \dots s_T)$

环境的状态转移满足 $s_{t+1} = f(s_t,a_t)$ （确定性环境）或 $s_{t+1} \sim P(\cdot \mid s_t,a_t)$ （随机性环境）。

Return（回报）指从当前时间步到 episode 结束所获得奖励的累积和（或带折扣的累积和）。

训练目标

我们要训练一个策略网络 $\pi_\theta$ ，使其在所有可能的状态下做出动作后，期望回报最大化：

$E (R(\tau))_{\tau \sim P_{\theta}(\tau)} = \sum\limits_\tau R(\tau)P_\theta(\tau)$

\begin{aligned} \nabla E(R(\tau))_{\tau \sim P_\theta(\tau)} &= \nabla \sum_\tau R(\tau)P_\theta(\tau)\\ &= \sum_\tau R(\tau)\nabla P_\theta(\tau)\\ &= \sum_\tau R(\tau)\nabla P_\theta(\tau) \cdot \dfrac{P_\theta(\tau)}{P_\theta(\tau)}\\ &= \sum_\tau P_\theta(\tau)R(\tau) \dfrac{\nabla P_\theta(\tau)}{P_\theta(\tau)}\\ &\approx \frac{1}{N}\sum\limits_{n=1}^N R(\tau^n)\dfrac{\nabla p_\theta(\tau^n)}{P_\theta(\tau^n)}\\ &= \frac{1}{N}\sum\limits_{n=1}^N R(\tau^n)\nabla \log P_\theta(\tau^n)\\ &= \frac{1}{N}\sum\limits_{n=1}^N R(\tau^n)\nabla \log \prod\limits_{t=1}^{T_n}P_\theta(a_n^t|s_n^t)\\ &= \frac{1}{N}\sum\limits_{n=1}^N R(\tau^n)\sum\limits_{t=1}^{T_n}\nabla \log P_\theta(a_n^t|s_n^t)\\ &= \frac{1}{N}\sum\limits_{n=1}^N \sum\limits_{t=1}^{T_n} R(\tau^n)\nabla \log P_\theta(a_n^t|s_n^t)\\ \end{aligned}

$Loss = -\frac{1}{N}\sum\limits_{n=1}^N \sum\limits_{t=1}^{T_n} R(\tau^n) \log P_\theta(a_n^t|s_n^t)$

$R(\tau^n) = \sum\limits_{t=t'}^{T_n} \gamma^{t-t'} r_{t'}^n = R_t^n$

Actor-Critic 演员-评论家算法

为了衡量某个动作"相对好坏"，我们让奖励 $R_t$ 减去一个基准值 $B(s_t)$ ，从而降低梯度估计的方差：

$\frac{1}{N} \sum\limits_{n=1}^N \sum\limits_{t=1}^{T_n} (R_t^n - B(s_t^n))\nabla \log P_\theta(a_t^n|s_t^n)$

然而 $R_t^n$ 仅来自一次随机采样，方差很大、训练不稳定。一个自然的改进是引入价值函数：用 $Q_\theta(s,a)$ 表示在状态 $s$ 下采取动作 $a$ 的期望回报，即 动作价值函数；用 $V_\theta(s)$ 表示在状态 $s$ 下的期望回报，即 状态价值函数。

定义 优势函数（Advantage Function） $A_\theta(s,a) = Q_\theta(s,a) - V_\theta(s)$ ，它衡量在状态 $s$ 下，动作 $a$ 相对于"该状态下的平均水平"有多大优势。

将上述策略梯度公式改写为：

$\dfrac{1}{N}\sum\limits_{n=1}^N \sum\limits_{t=1}^{T_n} A_\theta(s_t^n,a_t^n)\nabla \log P_\theta(a_t^n|s_t^n)$

利用 Bellman 关系，可以把动作价值函数与状态价值函数联系起来：

Q_\theta(s_t, a) = r_t + \gamma V_\theta(s_{t+1}) \\ A_\theta(s_t, a) = r_t + \gamma V_\theta(s_{t+1}) - V_\theta(s_t)\\ V_\theta(s_{t+1}) \approx r_{t+1} + \gamma V_\theta(s_{t+2})

经过这一步代换， $A_\theta(s,a)$ 的估计只依赖一个价值函数 $V_\theta$ ，公式整体的复杂度也得以降低。

那么应该向后展开（rollout）多少步再用 $V_\theta$ 截断呢？展开越多步，估计就越接近真实回报，偏差越小，但方差也会越大；反之展开越少，偏差大、方差小。这就是经典的 bias-variance trade-off：

A_\theta^1(s_t, a) = r_t + \gamma V_\theta(s_{t+1}) - V_\theta(s_t)\\ A_\theta^2(s_t, a) = r_t + \gamma r_{t+1} + \gamma^2 V_\theta(s_{t+2}) - V_\theta(s_t)\\ A_\theta^3(s_t, a) = r_t + \gamma * r_{t+1} + \gamma^2 * r_{t+2} + \gamma^3 V_\theta(s_{t+3}) - V_\theta(s_t)\\ \vdots \\ A_\theta^T(s_t, a) = r_t + \gamma * r_{t+1} + \gamma^2 * r_{t+2} + \gamma^3 * r_{t+3} + \cdots + \gamma^T * r_T - V_\theta(s_t)

为了让公式更简洁，引入一个中间量 $\delta_t^V$ ，表示 第 $t$ 步采取该动作所带来的 TD 残差（temporal-difference error）：

\delta_t^V = r_t + \gamma * V_\theta(s_{t+1}) - V_\theta(s_t)\\ \delta_{t+1}^V = r_{t+1} + \gamma * V_\theta(s_{t+2}) - V_\theta(s_{t+1})\\ A_\theta^1(s_t, a) = \delta_t^V\\ A_\theta^2(s_t, a) = \delta_t^V + \gamma \delta_{t+1}^V\\ A_\theta^3(s_t, a) = \delta_t^V + \gamma \delta_{t+1}^V + \gamma^2 \delta_{t+2}^V\\ \vdots

Generalized Advantage Estimation (GAE)

那么"展开几步"这件事到底该怎么选？GAE 给出的答案是：小孩子才做选择，全都要！

GAE 把不同展开步数的优势估计 $A^k$ 用一个权重 $\lambda$ 做指数加权平均，再化简成一个简洁的求和形式：

\begin{aligned} A_\theta^{GAE}(s_t, a) &= (1 - \lambda)(A_\theta^1 + \lambda * A_\theta^2 + \lambda^2 A_\theta^3 + \cdots) \quad\quad \lambda = 0.9: \quad A_\theta^{GAE} = 0.1A_\theta^1 + 0.09A_\theta^2 + 0.081A_\theta^3 + \cdots \\ &= (1 - \lambda)(\delta_t^V + \lambda * (\delta_t^V + \gamma \delta_{t+1}^V) + \lambda^2 (\delta_t^V + \gamma \delta_{t+1}^V + \gamma^2 \delta_{t+2}^V) + \cdots) \\ &= (1 - \lambda)(\delta_t^V (1 + \lambda + \lambda^2 + \cdots) + \gamma \delta_{t+1}^V * (\lambda + \lambda^2 + \cdots) + \cdots) \\ &= (1 - \lambda)\left(\delta_t^V \frac{1}{1 - \lambda} + \gamma \delta_{t+1}^V \frac{\lambda}{1 - \lambda} + \cdots\right) \\ &= \sum\limits_{b=0}^\infty (\gamma\lambda)^b \delta_{t+b}^V \end{aligned}

通过等比数列求和化简，最终得到 GAE 的紧凑形式。

GAE 优势函数本质上是在 $\lambda \to 0$ （高偏差、低方差）与 $\lambda \to 1$ （低偏差、高方差）之间做插值，从而平衡 bias 与 variance。

整理一下，到这里我们得到了三个关键表达式：

\delta_t^V = r_t + \gamma V_\theta(s_{t+1}) - V_\theta(s_t) \\ A_\theta^{GAE}(s_t, a) = \sum_{b=0}^{\infty} (\gamma\lambda)^b \delta_{t+b}^V \\ \frac{1}{N}\sum_{n=1}^N\sum_{t=1}^{T_n} A_\theta^{GAE}(s_n^t, a_n^t)\, \nabla \log P_\theta(a_n^t|s_n^t)

这里的状态价值函数 $V_\theta$ 一般用一个神经网络来拟合（即 critic 网络），可以与策略网络共用主干参数，仅在最后一层分叉为两个 head。

Proximal Policy Optimization (PPO) 近端策略优化算法

在经典的强化学习训练范式里，我们通常一边采集数据、一边更新模型，采过的数据用一次就丢掉——这种做法被称为 on-policy。问题在于，强化学习的环境交互成本往往很高，这样"用一次就扔"显然非常浪费。如果我们能让当前策略 $\pi_\theta$ 复用 旧策略 $\pi_{\theta'}$ 采集的数据进行训练（即 off-policy），训练效率就能显著提升。要实现这一点，关键就是 重要性采样（Importance Sampling）。

重要性采样

我们可以把" $f(x)$ 在分布 $p$ 下的期望"改写为" $f(x) \cdot \frac{p(x)}{q(x)}$ 在另一分布 $q$ （proposal distribution）下的期望"，这样就能用从 $q$ 采样的数据来估计原本在 $p$ 下的期望：

\begin{aligned} \mathbb{E}(f(x))_{x \sim p(x)} &= \sum_{x} f(x) \cdot p(x) \\ &= \sum_{x} f(x) \cdot p(x) \frac{q(x)}{q(x)} \\ &= \sum_{x} f(x) \frac{p(x)}{q(x)} \cdot q(x) \\ &= \mathbb{E}\left(f(x) \frac{p(x)}{q(x)}\right)_{x \sim q(x)} \\ &\approx \frac{1}{N} \sum_{n=1}^{N} f(x_n) \frac{p(x_n)}{q(x_n)}, \quad x_n \sim q(x) \end{aligned}

利用重要性采样，我们就可以把 on-policy 的梯度公式改写为可以复用旧数据的 off-policy 形式。

Off-policy

记 $\theta'$ 为 采集数据时使用的旧策略， $\theta$ 为 当前要优化的策略；优势 $A_{\theta'}^{GAE}$ 由旧策略下的价值网络估计而来。结合恒等式 $\nabla \log f(x) = \dfrac{\nabla f(x)}{f(x)}$ ，可以把策略梯度写成包含 重要性采样比 $\dfrac{P_\theta}{P_{\theta'}}$ 的形式：

\nabla \log f(x) = \frac{\nabla f(x)}{f(x)}

\begin{aligned} &\frac{1}{N} \sum_{n=1}^{N} \sum_{t=1}^{T_n} A_\theta^{GAE}(s_n^t, a_n^t) \nabla \log P_\theta(a_n^t | s_n^t) \\ &= \frac{1}{N} \sum_{n=1}^{N} \sum_{t=1}^{T_n} A_{\theta'}^{GAE}(s_n^t, a_n^t) \frac{P_\theta(a_n^t | s_n^t)}{P_{\theta'}(a_n^t | s_n^t)} \nabla \log P_\theta(a_n^t | s_n^t) \\ &= \frac{1}{N} \sum_{n=1}^{N} \sum_{t=1}^{T_n} A_{\theta'}^{GAE}(s_n^t, a_n^t) \frac{P_\theta(a_n^t | s_n^t)}{P_{\theta'}(a_n^t | s_n^t)} \frac{\nabla P_\theta(a_n^t | s_n^t)}{P_\theta(a_n^t | s_n^t)} \\ &= \frac{1}{N} \sum_{n=1}^{N} \sum_{t=1}^{T_n} A_{\theta'}^{GAE}(s_n^t, a_n^t) \frac{\nabla P_\theta(a_n^t | s_n^t)}{P_{\theta'}(a_n^t | s_n^t)} \end{aligned}

对应地，将期望最大化转换为损失最小化（取负号）：

Loss = -\frac{1}{N} \sum_{n=1}^{N} \sum_{t=1}^{T_n} A_{\theta'}^{GAE}(s_n^t, a_n^t) \frac{P_\theta(a_n^t | s_n^t)}{P_{\theta'}(a_n^t | s_n^t)}

这里有一个 隐含的前提：重要性采样要求 $\theta$ 与 $\theta'$ 不能差太多，否则比值 $\frac{P_\theta}{P_{\theta'}}$ 的方差会爆炸，估计就会失真。换句话说，我们需要给"新旧策略的差距"施加一个约束。PPO 给出了两种约束方式：

PPO-Penalty：在 surrogate 上加 KL 惩罚

如何让训练策略与参考策略不至于偏离太远？最直观的做法是给目标函数加一个 KL 散度惩罚项。KL 散度衡量两个分布的差异：差异越小，KL 越接近 0；差异越大，KL 越大。我们用一个权重 $\beta$ 来控制这一惩罚的强度（实际中 $\beta$ 还会做自适应调整）：

Loss_{\text{ppo}} = -\frac{1}{N} \sum_{n=1}^{N} \sum_{t=1}^{T_n} A_{\theta'}^{GAE}(s_n^t, a_n^t) \frac{P_{\theta}(a_n^t|s_n^t)}{P_{\theta'}(a_n^t|s_n^t)} + \beta\,\mathrm{KL}(P_{\theta}, P_{\theta'})

PPO-Clip：截断重要性采样比

PPO 还有一种更常用的实现，用 截断（clip）重要性采样比 来代替 KL 惩罚，同样起到限制新旧策略偏差的作用。它的目标函数由两部分组成（下方红色与蓝色项），最终取两者的 较小值：

红色部分：原始的 surrogate 目标 $\rho \cdot A$ ，其中 $\rho = \frac{P_\theta(a \mid s)}{P_{\theta'}(a \mid s)}$ 是重要性采样比；
蓝色部分：把 $\rho$ 截断到 $[1-\epsilon,\, 1+\epsilon]$ 区间内之后再乘以 $A$ ；当 $\rho$ 落在区间内时返回原值，落在外面则返回最近的边界值。

对二者取 min 的目的是：当一个动作有正向优势 $A > 0$ 时，限制 $\rho$ 不会被推得过高（避免步子迈太大）；当 $A < 0$ 时，限制 $\rho$ 不会被压得过低。这样既保证了"敢于改进策略"，又避免了"一次更新偏离太多"。

Loss_{\text{ppo2}} = -\frac{1}{N} \sum_{n=1}^{N} \sum_{t=1}^{T_n} \min\left(\textcolor{red}{A_{\theta'}^{GAE}(s_n^t, a_n^t) \frac{P_{\theta}(a_n^t|s_n^t)}{P_{\theta'}(a_n^t|s_n^t)}},\,\textcolor{skyblue}{\mathrm{clip}\left(\frac{P_{\theta}(a_n^t|s_n^t)}{P_{\theta'}(a_n^t|s_n^t)},\, 1-\epsilon,\, 1+\epsilon\right) A_{\theta'}^{GAE}(s_n^t, a_n^t)}\right)

PPO-Clip 由于实现简单、效果稳定，是当前应用最广泛的版本，也是 RLHF/GRPO 等大模型对齐算法的基础。

参考资料

零基础学习强化学习算法：ppo