一文速通 PPO

也就是说：

• 前面的 token 没有直接的任务奖励，但可以有各自的 KL 惩罚；
• 最后一个有效 token 得到完整的 $R_{\text{RM}}$，同时也有自己的 KL 惩罚；
• Prompt token 和 Padding token 通过 mask 排除，不参与这条奖励序列和后续 PPO Loss。

得到逐 token Reward 后，GAE 从最后一个有效 response token 开始向前递推。下面具体看最终奖励是如何影响每个 token 的。

假设回答包含 $T$ 个 token，并且第 $T$ 个 token 后真正终止，那么 $d_T=1$。最后一个位置的 TD Error 为：

由于 $r_T$ 中包含完整的 $R_{\text{RM}}$，所以最终任务奖励首先进入 $\delta_T$。接着从右向左递推：

因此，只看最终任务奖励 $R_{\text{RM}}$ 对前面位置的贡献，它在第 $t$ 个 token 的 Advantage 中对应：

距离结尾越远，影响通常会被 $\gamma\lambda$ 逐步衰减。如果 $\gamma=\lambda=1$，最终奖励会以相同权重影响前面的所有 token；如果 $\lambda<1$，越早的 token 接收到的信用越弱。

一个四 token 的具体例子

为了单独观察最终奖励的传播，暂时忽略 KL 惩罚，并假设 Critic 在所有位置的初始预测都是 $0$：

取 $\gamma=1,\lambda=0.95$。因为只有最后一个 token 有奖励，所以：

然后从右向左计算：

虽然 token 1、2、3 没有直接获得任务奖励，但它们的 Advantage 都受到了最终奖励的影响。这就是 GAE 对终止奖励进行时间信用分配的具体过程。

最后，为每个位置构造 Return target：

在这个简化例子中 $V_{\text{rollout}}=0$，因此：

真实训练中，逐 token KL 惩罚通常不为零，Critic 的 Value 也不为零。因此每个位置都有自己的 $\delta_t$；最终 Advantage 是“当前位置的 TD Error”和“后续所有 TD Error 的衰减和”，不会只是上面这组简单的等比数列。

要点｜最终奖励不是被平均拆分 GAE 没有把 $R_{\text{RM}}$ 平均分成 $T$ 份。它把最终奖励放在最后一个位置，再通过反向递推，让它以不同权重进入前面各 token 的 Advantage。

很多实现还会在 batch 内对 Advantage 做标准化，使其均值接近 $0$、标准差接近 $1$，以提高训练稳定性。

2. Actor Loss：根据 Advantage 调整动作概率

PPO 首先比较当前 Actor 和采集 trajectory 时的 Old Policy 对同一个动作给出的概率。为了避免与环境奖励 $r_t$ 混淆，这里把概率比记作 $\rho_t(\theta)$：

梯度关系可以简单记成：

• 有梯度：当前 Actor 重新计算的 log_prob、由它得到的 $\rho_t$ 以及 PPO Loss。梯度沿这条路径传回 Actor 参数 $\theta$；
• 无梯度：rollout 时保存的 old_log_prob 和预先计算的 $\hat A_t$，二者都作为 detach 后的固定训练信号使用。

如果直接使用 $\rho_t\hat A_t$ 更新，Actor 可能一次改变太多。PPO 使用 clipped surrogate objective 限制这种变化：

上式是一套通用模板。代入不同的 Advantage 估计，就得到对应的两个版本。

基础 Monte Carlo 版：

实践 GAE 版：

这个目标在数学上需要被最大化。深度学习框架通常默认最小化 Loss，因此代码里常写成：

其直觉是：

• Advantage 为正时，提高该动作概率，但不允许一次提高过多；
• Advantage 为负时，降低该动作概率，但不允许一次降低过多。

LLM 中的特殊点

LLM 的 Actor Loss 仍然是同一个 PPO clipped objective，只是它作用在 token-level action 上：

• 当前 LLM 重新计算生成 token 的 log_prob；
• 与 rollout 时保存的 old_log_prob 计算概率比；
• 使用对应 token 的 Advantage 计算 clipped loss；
• 最后只对有效 response token 做 masked mean。

Prompt token 和 Padding token 不是这次采样得到的 Action，通常不参与 Actor Loss。

3. Critic Loss：让 Value 更接近 Return

Critic 的任务是让当前预测 $V_\phi(s_t)$ 接近 rollout 数据计算出的价值目标。不同 Advantage 估计方法对应不同的 Critic target。

基础 Monte Carlo 版

直接把实际累计回报作为 target：

因此：

实践 GAE 版

先用 rollout 时保存的 Value 计算 GAE Advantage，再构造 GAE Return：

因此：

这里的 $V_{\text{rollout}}$ 是生成训练目标时使用的旧 Value；$V_\phi$ 是正在被优化的 Critic。训练时把 $\hat R_t^{\text{GAE}}$ 当作固定 target，不让梯度通过它反向传播。

因此，Actor 和 Critic 使用的是同一批 rollout 数据，但学习目标不同：

• Actor 根据 Advantage 学习哪些动作应该更常出现；
• Critic 根据 Return target 学习如何更准确地评价状态。

一些 PPO 实现还会对 Value 的变化做 clipping，避免 Critic 在一次更新中变化过大。

LLM 中的特殊点

LLM 的 Value Model 通常为每个 token 位置输出一个标量 Value，因此 Critic Loss 同样只在有效 response token 上计算，并排除 Prompt 与 Padding。

Actor 与 Critic 可以是两个独立模型，也可以共享同一个 Transformer backbone，再分别连接 policy head 和 value head；这属于工程结构差异，不改变 PPO 的基本公式。

四、总 Loss 和一次完整的 PPO 更新流程

1. 总 Loss 的定义

前面已经分别得到了 Actor Loss 和 Critic Loss。常见的 PPO 总 Loss 可以写成：

其中：

• $L_{\text{actor}}$：PPO clipped loss，使用 Advantage 更新策略；
• $L_{\text{critic}}$：Value loss，使 Critic 预测接近 Return target；
• $\mathcal H(\pi_\theta)$：当前策略的熵，鼓励策略保留一定探索能力；
• $c_v$：控制 Critic Loss 的权重；
• $c_e$：控制熵奖励的权重。

这里把 $L_{\text{actor}}$ 定义成需要最小化的负 clipped objective，因此总 Loss 中直接加上它；熵越大通常越有利，所以熵项前面带负号。

在 LLM 中，每个有效 response token 都会产生自己的 Actor Loss、Critic Loss 和 Entropy，Prompt 与 Padding 通过 mask 排除。各项先分别做 masked mean，再组合成最终的标量 Loss。

说明｜工程实现不一定真的只有一个 Loss 如果 Actor 和 Critic 共享 backbone，通常可以组合成一个 $L_{\text{total}}$，一次反向传播共同更新。如果 Actor 和 Critic 是两个独立模型，也可以分别使用 Actor optimizer 和 Critic optimizer；上面的公式主要表达它们在一次 PPO 更新中的总体目标。

经典 RLHF-PPO 中，KL 惩罚经常已经被加入逐 token Reward。如果某个框架把 KL 写成独立 Loss，则还需要在总 Loss 中加入对应项，具体以代码实现为准。

2. 一次完整的 PPO 迭代

与下图一致，PPO 的一次迭代分成三个阶段：

1
2
3

阶段 A：使用当前策略收集一批 rollout 数据
阶段 B：计算并固定 Advantage 与 Return target
阶段 C：固定这批数据，对它连续进行多轮更新

阶段 C 完成后会停止使用这批旧数据，再用更新后的策略开始下一轮阶段 A；这只是迭代之间的循环衔接，不再单独编号。

阶段 A：收集一批 rollout 数据

首先用当前策略与环境交互，收集固定数量的决策步。这里把决策步总数记作 $N_{\text{step}}$，其中每一步对应一条 state-action transition：

$N_{\text{step}}$ 不是完整轨迹的条数。例如并行运行 $E=16$ 个环境，每个环境收集 $H=128$ 步，那么：

这 2048 个决策步可能来自多条完整 trajectory，也可能包含某些 trajectory 的截断片段。

不要求一次必须刚好跑完一条完整 trajectory：

• 如果 episode 在收满 $N_{\text{step}}$ 步之前终止，就重置环境，继续收集下一个 episode；
• 如果收满 $N_{\text{step}}$ 步时 episode 还没有终止，就暂时截断这一段 rollout；
• 对真正终止的位置，下一状态的 Value 取 $0$；
• 对仅因 rollout 长度而截断的位置，可以使用最后状态的 Value 进行 bootstrap。

每个决策步通常保存：

• State 和实际执行的 Action；
• Reward 与终止标记 done；
• rollout 时旧策略对该动作的 old_log_prob；
• rollout 时 Critic 给出的 old_value。

在 LLM 中，一批 rollout 通常包含 $B_{\text{seq}}$ 个 Prompt 及其回答。每个有效 response token 都是一个决策步，因此总决策步数为：

其中 $T_b$ 是第 $b$ 条回答的有效 response token 数。这里需要同时区分“回答条数” $B_{\text{seq}}$ 和“token-level 决策步数” $N_{\text{token}}$。

要点｜Rollout 阶段不更新参数 采集数据时通常处于推理模式，并使用 no_grad。必须先得到一批由同一版本策略生成的数据，之后才能计算固定的 old_log_prob、Advantage 和 Return target。

阶段 B：计算并固定 Advantage 与 Return target

收集完成后，根据 Reward、rollout Value 和 done 计算：

1
2
3
4
5
6
7

Reward + old_value + done
             ↓
            GAE
       ┌─────┴─────┐
       ↓           ↓
  Advantage    Return target
  训练 Actor    训练 Critic

这些量在接下来的多轮 PPO 更新中保持不变：

• old_log_prob：作为概率比的固定分母；
• Advantage：作为 Actor 的固定训练信号；
• Return target：作为 Critic 的固定监督目标。

它们都应当 detach，不参与反向传播。

阶段 C：对同一批 rollout 连续更新多次

一批 on-policy 数据的采集成本很高，所以 PPO 不会只使用它做一次梯度更新。通常会：

1. 在 GAE 已经计算完成后，打乱训练样本的索引；
2. 将训练样本切分成多个 mini-batch；
3. 对每个 mini-batch 重新计算当前 Actor 的 new_log_prob 和当前 Critic 的 new_value；
4. 计算 Actor Loss、Critic Loss 和 Entropy；
5. 执行 backward() 和 optimizer.step()；
6. 将同一批 rollout 重复训练 $K$ 个 update epochs。

Mini-batch 的划分与数据打乱

GAE 计算完成后，rollout buffer 中的 old_log_prob、Advantage 和 Return target 已经固定，接下来才把数据交给 DataLoader。

设 DataLoader 实际接收的训练单元数量为 $N_{\text{unit}}$，每个 mini-batch 包含 $M_{\text{unit}}$ 个训练单元，那么一个 update epoch 包含：

个 mini-batch。这里的训练单元不一定是一条完整轨迹：

• 普通前馈 PPO 通常把所有 transition 展平，此时 $N_{\text{unit}}=N_{\text{step}}$；
• LLM PPO 通常以一条 Prompt/Response sequence 作为训练单元，此时 $N_{\text{unit}}=B_{\text{seq}}$；
• 每条 LLM sequence 内部包含多个有效 response token，最终 Loss 对这些 token 做 masked mean；
• 使用 sequence packing 或动态 token budget 时，mini-batch 的实际组织方式由 DataLoader 决定。

数据打乱同样发生在 GAE 计算完成之后。计算 GAE 时必须保留同一条轨迹中 $t$ 与 $t+1$ 的关系，才能从末尾向前递推 Advantage。训练目标固定以后，普通前馈 PPO 可以把不同环境、不同轨迹和不同时间步的 transition 全部展平，再随机打乱 transition 索引，使每个 mini-batch 混合来自不同位置的数据。

LLM PPO 通常随机打乱 Prompt/Response sequence 之间的顺序，而保留一条回答内部的 token 顺序。保留内部顺序不是 PPO Loss 的要求：当每个位置的 Loss 已经算出后，求和或求平均与排列顺序无关：

保留 token 顺序是为了高效地重新计算：

Transformer 需要看到正确的前缀 $x,y_{\lt t}$。将完整回答按原顺序送入一次带 causal mask 的前向传播，就能同时得到所有 token 的 new_log_prob 和 Value。理论上也可以给每个 token 单独携带完整 prefix，再任意打乱 token 样本，但这会重复计算大量相同前缀，实际训练通常不会这样做。

因此，一批 rollout 对应的梯度更新次数大约是：

例如：

每个 epoch 有 $2048/256=8$ 个 mini-batch，所以这一批数据一共会执行：

次梯度更新。

虽然 old_log_prob、Advantage 和 Return target 在这 $32$ 次更新中保持不变，但 Actor 和 Critic 的参数在不断改变，所以每个 mini-batch 都需要重新计算 new_log_prob 与 new_value。

要点｜多次利用，但不能无限利用 PPO 是 on-policy 算法。同一批 rollout 可以通过 mini-batch 和多个 update epochs 得到充分利用，但更新次数不能无限增加。策略变化过大后，这批数据会变得过旧，概率比也会偏离 $1$。实践中通常使用较小的 $K$，并可通过 clip fraction 或近似 KL 提前停止更新。

完成本轮更新：进入下一轮迭代

完成 $K$ 个 update epochs 后：

1. 停止使用这一批 rollout；
2. 更新后的当前 Actor 成为下一轮采样策略；
3. 使用它重新收集一批新数据，并保存新的 old_log_prob；
4. 重复 GAE 与多轮更新过程。

很多实现并不需要真的复制一个独立的 Old Policy 模型。只要在 rollout 时保存动作对应的 old_log_prob，就已经固定了计算概率比所需的旧策略信息。

3. PyTorch 风格伪代码

下面的代码省略了分布式训练、混合精度和 Value clipping，但保留了一次 PPO 迭代的完整主线：

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119

def masked_mean(x, mask):
    mask = mask.to(x.dtype)
    return (x * mask).sum() / mask.sum().clamp_min(1.0)


trainable_parameters = [
    p
    for group in optimizer.param_groups
    for p in group["params"]
]


for iteration in range(num_iterations):
    # ============================================================
    # 阶段 A：使用当前策略收集 N_step 个决策步
    # ============================================================
    actor.eval()
    critic.eval()

    with torch.no_grad():
        rollout = collect_rollout(
            actor=actor,
            critic=critic,
            horizon=N_step,
        )

        # rollout 中包含：
        # states, actions, rewards, dones, mask,
        # old_log_probs, old_values, next_value

        # ========================================================
        # 阶段 B：计算并固定 Advantage 与 Return target
        # ========================================================
        advantages, returns = compute_gae(
            rewards=rollout.rewards,
            values=rollout.old_values,
            dones=rollout.dones,
            next_value=rollout.next_value,
            gamma=gamma,
            gae_lambda=gae_lambda,
        )

        advantages = masked_normalize(
            advantages,
            rollout.mask,
        )

        # 整个更新阶段都保持固定
        old_log_probs = rollout.old_log_probs.detach()
        advantages = advantages.detach()
        returns = returns.detach()

    # ============================================================
    # 阶段 C：同一批 rollout 重复训练 K 个 update epochs
    # ============================================================
    actor.train()
    critic.train()

    for epoch in range(update_epochs):
        for idx in iterate_minibatches(
            # 普通 PPO 可取 transition 数；
            # LLM PPO 通常取本轮的 sequence / packed-sample 数。
            num_items=rollout.num_training_units,
            minibatch_size=minibatch_size,
            shuffle=True,
        ):
            states = rollout.states[idx]
            actions = rollout.actions[idx]
            mask = rollout.mask[idx]

            # 当前参数下重新计算，因此保留梯度
            new_log_probs, entropy = actor.evaluate_actions(
                states,
                actions,
            )
            new_values = critic(states)

            # ---------- Actor Loss ----------
            log_ratio = new_log_probs - old_log_probs[idx]
            ratio = log_ratio.exp()

            adv = advantages[idx]
            surrogate_1 = ratio * adv
            surrogate_2 = torch.clamp(
                ratio,
                1.0 - clip_eps,
                1.0 + clip_eps,
            ) * adv

            actor_loss = -masked_mean(
                torch.minimum(surrogate_1, surrogate_2),
                mask,
            )

            # ---------- Critic Loss ----------
            critic_loss = 0.5 * masked_mean(
                (new_values - returns[idx]) ** 2,
                mask,
            )

            # ---------- Entropy Bonus ----------
            entropy_bonus = masked_mean(entropy, mask)

            total_loss = (
                actor_loss
                + value_coef * critic_loss
                - entropy_coef * entropy_bonus
            )

            optimizer.zero_grad()
            total_loss.backward()
            torch.nn.utils.clip_grad_norm_(
                parameters,
                max_grad_norm,
            )
            optimizer.step()

    # 本批 rollout 到此停止使用。
    # 下一轮会用更新后的 actor 重新采样并保存新的 old_log_probs。

如果 Actor 和 Critic 使用两个独立 optimizer，可以分别对 Actor objective 和 Critic Loss 执行反向传播与参数更新；数据收集、GAE 和多轮 mini-batch 复用的主流程不变。

要点｜PPO 的最终主线 先用当前策略收集一批 on-policy 数据；再固定 old log-prob、Advantage 和 Return target；然后把同一批数据切成 mini-batch，连续更新多个 epochs；最后停止复用旧数据，用更新后的策略重新采样。