PPO 深度解析：信赖域优化的艺术与 LLM 对齐实践

1. 引言：PPO 为什么会成为 LLM 对齐的主流？#

1.1 强化学习的三代算法#

在 PPO 之前，Policy Gradient 算法经历了三代演进，每一代都在解决上一代的核心问题：

1
第一代: REINFORCE (Williams, 1992)
2
  ∇J = E[∇_θ log π_θ(a|s) · R]
3
  问题：更新步长难以控制，方差极大，训练极不稳定
4

5
  ↓ 解决方差问题
6

7
第二代: Actor-Critic (Sutton & Barto, 1999)
8
  用 Value Network 估计 baseline，减少方差
9
  问题：仍无步长保证，可能一步跨太大 → 崩溃
10

11
  ↓ 加入信赖域约束
12

13
第三代: TRPO (Schulman et al., 2015)
14
  用 KL 约束限制新旧策略差异
15
  问题：需要共轭梯度求解二阶优化，工程复杂
16

17
  ↓ 简化 KL 约束（不用二阶优化）
18

19
第四代: PPO (Schulman et al., 2017)
20
  用 Clip 机制一阶优化，兼顾稳定性与简单性
21
  成为 LLM 对齐的默认选择

1.2 为什么 LLM 对齐选择 PPO？#

PPO 相比 DPO 有不可替代的优势：

维度	DPO	PPO
在线探索	❌ 纯离线	✅ 在线采样
多轮迭代改进	❌ 一次性	✅ 持续改进
复杂奖励	⚠️ 单标量	✅ 任意奖励函数
工程复杂度	低	高
显存需求	~2× 模型	~4× 模型
代表应用	Llama 3, Mistral	GPT-4, Claude 2, ChatGPT

核心洞察：DPO 是离线的直接偏好优化，PPO 是在线的策略优化。当你的目标是：

追赶 SOTA（最大性能）→ PPO
快速迭代（速度和简单性）→ DPO

Tip

2024 年 DeepSeek-R1 的出现打破了这一格局——GRPO 以 PPO 的在线探索优势 + DPO 的简洁性，在数学推理任务上取得了 SOTA。这将在后文详细讨论。

2. 策略梯度基础：从 REINFORCE 到 Actor-Critic#

2.1 策略梯度定理#

对于一个策略 $\pi_\theta(a|s)$ ，目标是最小化负期望回报（对于 LLM 对齐是最大化期望奖励）：

J(\theta) = \mathbb{E}_{\tau \sim \pi_\theta}\Big[\sum_{t=0}^{T-1} r(s_t, a_t)\Big]

其中 $\tau = (s_0, a_0, s_1, a_1, \ldots)$ 是轨迹（对于 LLM， $s$ 是 token 序列的前缀， $a$ 是下一个 token）。

策略梯度定理给出：

\nabla_\theta J(\theta) = \mathbb{E}_{\tau\sim\pi_\theta}\Big[\sum_{t=0}^{T-1}\nabla_\theta\log\pi_\theta(a_t|s_t)\;G_t\Big]

其中 $G_t = \sum_{t'=t}^{T-1} r_{t'}$ 是从时刻 $t$ 开始的累计回报。

2.2 REINFORCE：基本策略梯度#

1
def reinforce_loss(log_probs, rewards, dones, gamma=1.0):
2
    """
3
    REINFORCE 损失函数。
4
    log_probs: (B, T) 每个 token 的 log π(a_t | s_t)
5
    rewards: (B, T) 每个 token 的奖励
6
    """
7
    T = log_probs.size(1)
8

9
    # 计算 G_t（从后往前累计）
10
    G = torch.zeros_like(rewards)
11
    running_return = torch.zeros(rewards.size(0), device=rewards.device)
12
    for t in reversed(range(T)):
13
        running_return = rewards[:, t] + (1 - dones[:, t]) * gamma * running_return
14
        G[:, t] = running_return
15

16
    # 策略梯度: E[∇log π · G]
17
    # 对每条轨迹的 log prob 求和，乘以 G
18
    loss = -(log_probs * G).sum(dim=1).mean()
19
    return loss

2.3 方差问题的根源#

REINFORCE 的方差极大——因为 $G_t$ 是从随机策略采样的路径上计算的，不同采样路径的 $G_t$ 差异巨大。

两种降低方差的方法：

增加采样数量：理论上有效，但计算成本爆炸
引入 baseline：从 $G_t$ 中减去一个基准 $b(s_t)$ ，不改变期望但降低方差

2.4 Actor-Critic：引入 Value Network#

引入一个 Value Network $V_\phi(s_t)$ 作为 baseline：

A_t = G_t - V_\phi(s_t)

$A_t$ 称为优势函数（Advantage），衡量在状态 $s_t$ 下采取动作 $a_t$ 比”平均水平”好多少。

Actor-Critic 架构：

1
class ActorCritic(torch.nn.Module):
2
    """
3
    简化的 Actor-Critic。
4
    对于 LLM：状态 s 是 token 序列，动作 a 是下一个 token。
5
    """
6
    def __init__(self, llm_backbone):
7
        super().__init__()
8
        self.llm = llm_backbone
9
        hidden_size = llm_backbone.config.hidden_size
10
        self.value_head = nn.Linear(hidden_size, 1)  # Critic
11
        self.action_head = nn.Linear(hidden_size, llm_backbone.config.vocab_size)  # Actor
12

13
    def forward(self, input_ids, attention_mask):
14
        outputs = self.llm(input_ids=input_ids, attention_mask=attention_mask)
15
        hidden = outputs.hidden_states[-1]
16
        values = self.value_head(hidden).squeeze(-1)  # (B, L)
17
        logits = self.action_head(hidden)              # (B, L, V)
18
        return logits, values

Actor-Critic 的根本问题

Actor-Critic 没有新旧策略比率的约束，梯度可能步子太大，导致策略崩溃。PPO 的 clip 机制正是为了解决这个问题。

3. TRPO：信赖域优化的理论基础#

3.1 信赖域的核心思想#

直觉：在参数空间中，新旧策略不能离得太远，否则策略会崩溃。

TRPO 的约束形式：

\max_\theta\ \mathbb{E}_t\Big[\frac{\pi_\theta(a_t|s_t)}{\pi_{\theta_{\text{old}}}(a_t|s_t)}\;A_t\Big]

\text{s.t.}\quad \mathbb{E}_t\Big[\text{KL}\big(\pi_{\theta_{\text{old}}}(\cdot|s_t)\,\|\,\pi_\theta(\cdot|s_t)\big)\Big] \leq \delta

$\delta$ 是信任域半径（trust region radius）。

3.2 从约束到惩罚：共轭梯度法#

TRPO 需要求解带约束的优化问题，实际实现中：

用共轭梯度法（Conjugate Gradient）近似求解 KL 约束的二阶近似
线性搜索确定步长

1
def trpo_update(policy_model, old_policy, advantages, kl_target=0.01):
2
    """
3
    TRPO 更新步骤（简化的核心逻辑）。
4
    1. 计算策略梯度
5
    2. 共轭梯度法求搜索方向
6
    3. 线搜索确定步长
7
    """
8
    # === 1. 计算 KL 约束的 Hessian-vector product ===
9
    # H = ∇² KL / ∇θ²，复杂度 O(θ²)
10
    # 共轭梯度法只需要 H·v，不需要显式构造 H
11

12
    # === 2. 计算自然梯度方向 ===
13
    # ∇_natural = H⁻¹ · ∇_policy
14
    search_direction = conjugate_gradient(
15
        hvp_func=lambda v: kl_hessian_vector_product(old_policy, policy_model, v),
16
        b=policy_grad,
17
        max_iter=20,
18
    )
19

20
    # === 3. 线搜索（在信赖域内找最优）===
21
    for alpha in [1.0, 0.5, 0.25, 0.125]:
22
        candidate_theta = theta + alpha * search_direction
23
        if kl_divergence(theta, candidate_theta) < kl_target:
24
            theta = candidate_theta
25
            break

3.3 TRPO 的工程痛点#

痛点	描述
共轭梯度开销大	需要计算 Hessian-Vector Product， $O(\theta^2)$
内存占用高	需要存储 Fisher 信息矩阵或其近似
超参敏感	$\delta$ 的选择对结果影响很大
难以与大规模 LLM 结合	上述问题在亿级参数上被放大

PPO 用 clip 机制一阶近似了这个二阶约束，完美解决了这些问题。

4. PPO：Clip 机制的数学推导#

4.1 重要性比率#

定义重要性比率（Importance Sampling Ratio）：

r_t(\theta) = \frac{\pi_\theta(a_t|s_t)}{\pi_{\theta_{\text{old}}}(a_t|s_t)}

对于 LLM，这意味着每个 token 的概率比：

1
def compute_importance_ratio(
2
    log_probs_new: torch.Tensor,  # π_θ(a_t | s_t)
3
    log_probs_old: torch.Tensor,  # π_{θ_old}(a_t | s_t)
4
) -> torch.Tensor:
5
    """r_t(θ) = π_θ / π_{θ_old}"""
6
    return torch.exp(log_probs_new - log_probs_old)  # (B, T)

4.2 原始策略梯度损失#

如果用原始的策略梯度：

\mathcal{L}^{\text{PG}}(\theta) = \mathbb{E}_t\Big[r_t(\theta)\;A_t\Big]

问题：当 $r_t(\theta)$ 变得很大（超过 1）时，更新步长无法控制——策略可能在一次更新中完全偏离。

4.3 Clipped Surrogate 目标#

PPO 的核心创新是用 clip 操作限制 $r_t$ 的变化范围：

\mathcal{L}^{\text{CLIP}}(\theta) = \mathbb{E}_t\Big[\min\big(r_t(\theta)\;A_t,\ \text{clip}(r_t(\theta), 1-\epsilon, 1+\epsilon)\;A_t\big)\Big]

其中 $\epsilon$ 是超参数（通常 $\epsilon = 0.2$ ）。

4.4 Clip 机制的直觉#

1
假设 A > 0（好的动作，正优势）：
2

3
  r_t(θ)
4
    │
5
 2.0│                    ┌────────── ← clip 上界
6
    │                   /
7
 1.8│                 /
8
 1.6│               /  ← 超出上界部分被 clip
9
 1.4│             /
10
 1.2│           /
11
 1.0│─────────●───────────────────
12
    │         ↑
13
    └─────────┼─────────────── θ_old
14
              │
15
              1 - ε   1.0   1 + ε
16

17
情况分析：
18
  θ 在 [θ_old - ε, θ_old + ε] 内 → r_t 正常，无 clip
19
  θ 超出信赖域（|r_t - 1| > ε）→ r_t 被 clip 到 [1-ε, 1+ε]
20
  → 梯度被限制，不会一步更新太大

4.5 完整的数学推导#

分情况讨论：

情况 1：A > 0（正优势，鼓励这个动作）

如果 $r_t < 1 + \epsilon$ ：无 clip， $\mathcal{L} = r_t \cdot A$ （正常增加概率）
如果 $r_t \geq 1 + \epsilon$ ：被 clip 到 $(1+\epsilon)\cdot A$ ，但 min 操作使得额外增大不再被奖励

情况 2：A < 0（负优势，抑制这个动作）

如果 $r_t > 1 - \epsilon$ ：无 clip， $\mathcal{L} = r_t \cdot A$ （负值，减小概率）
如果 $r_t \leq 1 - \epsilon$ ：被 clip 到 $(1-\epsilon)\cdot A$ ，但 min 操作使得过度减小不再被惩罚

1
def ppo_clipped_loss(
2
    log_probs_new: torch.Tensor,    # (B, T)
3
    log_probs_old: torch.Tensor,    # (B, T)
4
    advantages: torch.Tensor,       # (B, T)
5
    clip_eps: float = 0.2,
6
) -> torch.Tensor:
7
    """
8
    PPO Clipped Surrogate Loss。
9

10
    L = min(r(θ)·A, clip(r(θ), 1-ε, 1+ε)·A)
11

12
    物理含义：
13
      - min 操作确保 clip 不增加 loss（不会过度优化）
14
      - 当 A > 0 时，防止 r 变得太大（过度增加好动作的概率）
15
      - 当 A < 0 时，防止 r 变得太小（过度惩罚坏动作）
16
    """
17
    ratio = torch.exp(log_probs_new - log_probs_old)  # (B, T)
18

19
    # 原始目标
20
    policy_loss_1 = ratio * advantages
21

22
    # Clipped 目标
23
    ratio_clipped = torch.clamp(ratio, 1 - clip_eps, 1 + clip_eps)
24
    policy_loss_2 = ratio_clipped * advantages
25

26
    # 取较小的那个
27
    policy_loss = -torch.min(policy_loss_1, policy_loss_2)
28

29
    # 通常对序列求平均（也可以 token-wise）
30
    return policy_loss.mean()

4.6 为什么 min 而不是 max？#

这是一个精妙的设计：保证在 clip 区域外不会过度优化。

考虑 $A > 0$ 且 $r_t(\theta) > 1+\epsilon$ （策略已经大幅增加了这个动作的概率）：

$\mathcal{L}_1 = r_t \cdot A$ 很小（因为 $r_t$ 大导致梯度小）
$\mathcal{L}_2 = (1+\epsilon)\cdot A$ 更小（是上界）

取 $\min(\mathcal{L}_1, \mathcal{L}_2) = \mathcal{L}_2$ ，不再增加这个动作的概率——因为已经过度了。

5. GAE：广义优势估计#

5.1 TD(λ) 的优势函数版本#

对于 LLM 的序列决策（只有一个终态 reward），优势函数的估计是一个关键问题。

TD(0) 估计：

\hat{A}_t^{\text{TD}(0)} = r_t + \gamma\,V(s_{t+1}) - V(s_t)

n步回报：

\hat{A}_t^{(n)} = \sum_{l=0}^{n-1}\gamma^l r_{t+l} + \gamma^n V(s_{t+n}) - V(s_t)

**GAE（Generalized Advantage Estimation）**将这些 n步估计做加权平均：

\hat{A}_t^{\text{GAE}(\lambda,\gamma)} = (1-\lambda)\sum_{n=0}^{\infty}(\lambda\gamma)^n\;\hat{A}_t^{(n)}

5.2 GAE 的递推形式#

GAE 可以用递归方式高效计算：

\hat{A}_t = \delta_t + \gamma\,\lambda\;\hat{A}_{t+1}

其中 $\delta_t = r_t + \gamma\,V(s_{t+1}) - V(s_t)$ 是 TD 误差。

1
def compute_gae(
2
    rewards: torch.Tensor,        # (B, T) 序列奖励
3
    values: torch.Tensor,         # (B, T) Value 网络估计
4
    dones: torch.Tensor,          # (B, T) 是否结束
5
    gamma: float = 1.0,
6
    lam: float = 0.95,
7
) -> tuple[torch.Tensor, torch.Tensor]:
8
    """
9
    GAE 计算。
10

11
    递推公式: A_t = δ_t + γλA_{t+1}
12
    其中 δ_t = r_t + γV(s_{t+1}) - V(s_t)
13

14
    参数：
15
      gamma (γ): 折扣因子。对于 LLM 序列（终态任务），通常 γ=1.0
16
      lam (λ): GAE 平滑参数。
17
        λ=0  → TD(0)，方差低但高偏差
18
        λ=1  →蒙特卡洛，无偏差但高方差
19
        λ≈0.95 → 平衡两者
20
    """
21
    B, T = rewards.shape
22
    advantages = torch.zeros_like(rewards)
23
    last_adv = torch.zeros(B, device=rewards.device)
24

25
    for t in reversed(range(T)):
26
        if t == T - 1:
27
            next_value = 0.0  # 序列结束，无 future value
28
        else:
29
            next_value = values[:, t + 1]
30

31
        nonterminal = 1.0 - dones[:, t]
32
        delta = rewards[:, t] + gamma * nonterminal * next_value - values[:, t]
33

34
        # A_t = δ_t + γλ A_{t+1}
35
        last_adv = delta + gamma * lam * nonterminal * last_adv
36
        advantages[:, t] = last_adv
37

38
    returns = advantages + values  # 用于训练 Value 网络
39
    return advantages, returns

5.3 GAE 超参数的影响#

$\lambda$	偏差	方差	适用场景
0（TD(0)）	高	极低	在线学习，快速估计
0.9 ~ 0.95	中	中	大多数场景，推荐值
1（MC）	0（无偏）	高	短序列，高采样数

LLM 对齐中的特殊处理

在 LLM 对齐中，reward 通常只在序列最后一个 token 给出（RM 分数），中间 token 的 reward 为 0。这种情况下 GAE 的 $\lambda$ 参数影响更加显著—— $\lambda$ 越接近 1，奖励信号越能传播到所有 token。

5.4 Value 网络损失#

Value 网络（Critic）的任务是预测期望回报，其损失为均方误差：

\mathcal{L}_V = \mathbb{E}\big[(\hat{R}_t - V_\phi(s_t))^2\big]

其中 $\hat{R}_t$ 是 GAE 返回值（returns）。

通常使用 clip Value 预测来防止 Value 网络更新过大（PPO-Penalty 的做法）：

1
def value_loss(
2
    values_new: torch.Tensor,
3
    values_old: torch.Tensor,
4
    returns: torch.Tensor,
5
    clip_eps: float = 0.2,
6
):
7
    """
8
    Clipped Value Loss（PPO-Penalty 的 Value 损失）。
9
    限制 Value 预测的变化幅度，防止 Value 更新过大导致不稳定。
10
    """
11
    clipped_values = torch.clamp(
12
        values_new,
13
        values_old - clip_eps,
14
        values_old + clip_eps,
15
    )
16
    loss1 = (values_new - returns) ** 2
17
    loss2 = (clipped_values - returns) ** 2
18
    return torch.max(loss1, loss2).mean()

6. PPO 用于 LLM 对齐：完整实现#

6.1 四模型架构#

在 LLM 对齐中，PPO 需要同时运行四个模型：

1
┌─────────────────────────────────────────────────────────────┐
2
│                    PPO 四模型架构                             │
3
├─────────────────────────────────────────────────────────────┤
4
│                                                              │
5
│  ① Policy Model (π_θ)                                       │
6
│     → 待优化的 LLM，梯度更新                                  │
7
│                                                              │
8
│  ② Reference Model (π_ref)                                  │
9
│     → SFT 模型，参数冻结，计算 KL 惩罚                        │
10
│                                                              │
11
│  ③ Reward Model (R_φ)                                        │
12
│     → 偏好数据训练的奖励模型，参数冻结                         │
13
│     → 只在 rollout 阶段使用一次打分                           │
14
│                                                              │
15
│  ④ Value Network (V_ψ)                                      │
16
│     → 与 policy 共享 LLM backbone，仅 value head 可学习       │
17
│     → 在 rollout 和 update 阶段都运行                         │
18
│                                                              │
19
│  显存占用 ≈ 4 × LLM 参数 + embeddings                        │
20
│  → 通常需要对 policy/ref/value 做 LoRA 或量化                 │
21
│                                                              │
22
└─────────────────────────────────────────────────────────────┘

6.2 Rollout：在线采样#

1
def ppo_rollout(
2
    policy_model,
3
    ref_model,
4
    reward_model,
5
    critic_model,
6
    prompts: list[str],
7
    tokenizer,
8
    max_new_tokens: int = 256,
9
    temperature: float = 1.0,
10
    top_p: float = 0.9,
11
    num_return_sequences: int = 4,  # GRPO 风格的 G
12
) -> dict:
13
    """
14
    PPO Rollout 阶段：
15
    1. 用当前策略采样多个回答
16
    2. 用 Reward Model 评分
17
    3. 计算 KL 惩罚
18
    4. 收集用于 PPO 更新的数据
19
    """
20
    # === 1. 用策略采样（可以多个 response per prompt）===
21
    all_responses = []
22
    all_log_probs = []
23
    all_ref_log_probs = []
24
    all_values = []
25
    all_rewards = []
26

27
    for prompt in prompts:
28
        prompt_tensors = tokenizer([prompt], return_tensors="pt").to(policy_model.device)
29

30
        # 如果 num_return_sequences > 1，就是 GRPO-style
31
        if num_return_sequences > 1:
32
            # 批量生成多个 response
33
            prompt_batch = [prompt] * num_return_sequences
34
        else:
35
            prompt_batch = [prompt]
36

37
        with torch.no_grad():
38
            outputs = policy_model.generate(
39
                input_ids=prompt_batch,
40
                max_new_tokens=max_new_tokens,
41
                do_sample=True,
42
                temperature=temperature,
43
                top_p=top_p,
44
                pad_token_id=tokenizer.pad_token_id,
45
            )
46

47
        # === 2. 提取 response 部分 ===
48
        response_tensors = extract_response(prompt_tensors, outputs, tokenizer)
49

50
        # === 3. 计算 log probs ===
51
        # Policy log probs
52
        policy_log_probs = compute_response_log_probs(policy_model, response_tensors)
53

54
        # Reference log probs（用于 KL）
55
        with torch.no_grad():
56
            ref_log_probs = compute_response_log_probs(ref_model, response_tensors)
57

58
        # === 4. Value 估计 ===
59
        with torch.no_grad():
60
            _, values_pred = critic_model(
61
                input_ids=response_tensors["input_ids"],
62
                attention_mask=response_tensors["attention_mask"],
63
            )
64
            values_pred = values_pred[:, :-1].mean(dim=-1)  # 标量 per sequence
65

66
        # === 5. Reward Model 评分 ===
67
        with torch.no_grad():
68
            rewards = reward_model(
69
                input_ids=response_tensors["input_ids"],
70
                attention_mask=response_tensors["attention_mask"],
71
            )
72

73
        # === 6. 构建 KL 奖励塑形 ===
74
        # KL penalty per token：-β * (log π_θ - log π_ref)
75
        kl_penalty = -(policy_log_probs - ref_log_probs)  # (B, T)
76

77
        # 只有最后一个 token 获得真正的 RM 分数
78
        T = response_tensors["input_ids"].size(1)
79
        token_rewards = kl_penalty * 0.0  # 初始化为 0
80
        token_rewards[:, -1] = rewards     # 最后一个 token 获得 RM 分数
81
        # 或者更常用的：token_rewards = -beta * kl_penalty, token_rewards[:, -1] += rewards
82

83
        all_rewards.append(token_rewards)
84
        all_log_probs.append(policy_log_probs)
85
        all_ref_log_probs.append(ref_log_probs)
86
        all_values.append(values_pred)
87

88
    return {
89
        "rewards": torch.cat(all_rewards, dim=0),
90
        "log_probs": torch.cat(all_log_probs, dim=0),
91
        "ref_log_probs": torch.cat(all_ref_log_probs, dim=0),
92
        "values": torch.cat(all_values, dim=0),
93
    }

6.3 PPO 更新#

1
def ppo_update_step(
2
    policy_model,
3
    critic_model,
4
    rollout_data: dict,
5
    beta: float = 0.1,       # KL 系数
6
    clip_eps: float = 0.2,   # PPO clip epsilon
7
    gamma: float = 1.0,      # 折扣因子
8
    lam: float = 0.95,       # GAE lambda
9
    entropy_coef: float = 0.01,  # 熵正则
10
    vf_coef: float = 0.5,
11
    max_grad_norm: float = 1.0,
12
    ppo_epochs: int = 4,     # 同一个 rollout 数据更新 PPO epochs 次
13
    minibatch_size: int = 8,
14
):
15
    """
16
    PPO 更新步骤。
17
    """
18
    rewards = rollout_data["rewards"]       # (B, T)
19
    log_probs_old = rollout_data["log_probs"]  # (B, T)
20
    ref_log_probs = rollout_data["ref_log_probs"]
21
    values_old = rollout_data["values"]    # (B,)
22

23
    # === 1. GAE 计算优势 ===
24
    dones = torch.zeros_like(rewards)
25
    dones[:, -1] = 1.0  # 序列结束标记
26
    advantages, returns = compute_gae(rewards, values_old.unsqueeze(1).repeat(1, rewards.size(1)), dones, gamma, lam)
27

28
    # 标准化优势
29
    advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
30

31
    # === 2. 每个 token 的 log prob（用于策略梯度）===
32
    # 对于 LLM，优势是 per-sequence 的，但 apply 时展平到 token
33
    advantages_flat = advantages.reshape(-1)  # (B*T,)
34

35
    # === 3. PPO Epoch 循环 ===
36
    B, T = rewards.shape
37
    indices = torch.randperm(B * T)  # 打乱所有 token
38

39
    for epoch in range(ppo_epochs):
40
        for start in range(0, B * T, minibatch_size):
41
            end = start + minibatch_size
42
            idx = indices[start:end]
43

44
            # === 3.1 Policy forward ===
45
            # 需要重新 forward 以获取当前策略的 log probs
46
            # （简化版：这里假设 rollout_data 存储了 old_log_probs）
47
            log_probs_new = log_probs_old  # placeholder，实际需要重算
48

49
            # === 3.2 PPO Clipped Loss ===
50
            policy_loss = ppo_clipped_loss(
51
                log_probs_new.reshape(-1)[idx].unsqueeze(0),
52
                log_probs_old.reshape(-1)[idx].unsqueeze(0),
53
                advantages_flat[idx].unsqueeze(0),
54
                clip_eps=clip_eps,
55
            )
56

57
            # === 3.3 Value Loss ===
58
            # 从 rollout 数据中取出对应样本的 returns
59
            values_new = values_old  # placeholder，实际需要重算
60
            vf_loss = ((values_new - returns.reshape(-1)[idx]) ** 2).mean()
61

62
            # === 3.4 熵正则（鼓励探索）===
63
            # entropy = compute_entropy(logits)  # placeholder
64
            # ent_loss = -entropy_coef * entropy.mean()
65

66
            # === 3.5 总损失 ===
67
            loss = policy_loss + vf_coef * vf_loss  # + ent_loss
68

69
            optimizer.zero_grad()
70
            loss.backward()
71
            torch.nn.utils.clip_grad_norm_(policy_model.parameters(), max_grad_norm)
72
            optimizer.step()

6.4 KL 约束的奖励塑形形式#

在 LLM 对齐中，更常用的是奖励塑形（Reward Shaping）版本，而非 PPO-Penalty 的 KL 惩罚项：

r_t^{\text{shaped}} = \begin{cases} R(x, y) - \beta\,\text{KL}\big(\pi_\theta(\cdot|x_t)\,\|\,\pi_{\text{ref}}(\cdot|x_t)\big) & t = T \\ -\beta\,\text{KL}\big(\pi_\theta(\cdot|x_t)\,\|\,\pi_{\text{ref}}(\cdot|x_t)\big) & t < T \end{cases}

1
def compute_kl_reward_shaping(
2
    policy_log_probs: torch.Tensor,   # (B, T)
3
    ref_log_probs: torch.Tensor,      # (B, T)
4
    reward_scores: torch.Tensor,      # (B,) RM 分数
5
    beta: float = 0.1,
6
) -> torch.Tensor:
7
    """
8
    KL 奖励塑形。
9
    每个 token 获得 -β·KL 惩罚，只有最后一个 token 获得 RM 分数。
10
    """
11
    kl_per_token = policy_log_probs - ref_log_probs  # (B, T)
12
    token_rewards = -beta * kl_per_token
13
    token_rewards[:, -1] += reward_scores  # 最后一个 token 加上 RM 分数
14
    return token_rewards

7. 进阶技巧：让 PPO 训练更稳定#

7.1 Reward 归一化#

RM 的原始分数方差可能很大（不同回答的 RM 分数可能差 10 倍），导致 GAE 估计不稳定。

1
def normalize_rewards(rewards: torch.Tensor, eps: float = 1e-8) -> torch.Tensor:
2
    """对每个 batch 的 reward 做 z-score 归一化"""
3
    return (rewards - rewards.mean()) / (rewards.std() + eps)

7.2 Reward Clipping#

将 reward 限制在一个范围内，防止极端 reward 破坏训练：

1
def clip_rewards(rewards: torch.Tensor, lo: float = -5.0, hi: float = 5.0) -> torch.Tensor:
2
    return torch.clamp(rewards, lo, hi)

7.3 KL Annealing#

训练初期可以允许策略偏离参考模型多一些（快速学习），后期收紧 KL（稳定收敛）：

1
def get_beta(step: int, total_steps: int, beta_init: float = 0.1, beta_final: float = 0.2) -> float:
2
    """
3
    KL 系数预热。
4
    训练初期 β 较小，后期逐渐增大。
5
    """
6
    progress = step / total_steps
7
    return beta_init + (beta_final - beta_init) * progress

7.4 Value 网络预热#

在 PPO 更新之前，先单独训练 Value 网络几个 epoch，减少初始 Value 预测的巨大误差：

1
def pretrain_critic(critic_model, rollout_data, num_warmup_epochs=3):
2
    """
3
    Value 网络预热：用 rollout 数据的 returns 训练几个 epoch。
4
    """
5
    returns = rollout_data["returns"]
6
    values_old = rollout_data["values"]
7
    for _ in range(num_warmup_epochs):
8
        loss = ((values_old - returns) ** 2).mean()
9
        loss.backward()
10
        optimizer.step()

7.5 早停机制#

监控 KL 散度，当策略偏离参考模型太远时早停：

1
def check_early_stop(
2
    policy_log_probs: torch.Tensor,
3
    ref_log_probs: torch.Tensor,
4
    kl_threshold: float = 0.05,
5
):
6
    """如果平均 KL 超过阈值，跳过此次更新"""
7
    mean_kl = (policy_log_probs - ref_log_probs).mean()
8
    if mean_kl > kl_threshold:
9
        return True  # 早停
10
    return False

8. PPO vs GRPO：两条技术路线#

8.1 GRPO 的核心洞察#

GRPO（DeepSeek-R1）提出：对于 LLM 对齐，不需要 critic 网络。

原因：

LLM 的 action space 是离散的 token（ $|\mathcal{A}| = 3-10$ 万）
每个 prompt 可以采样 $G$ 个回答，用组内排名作为优势

8.2 GRPO vs PPO 详细对比#

维度	PPO	GRPO
Value Network	需要（估计 baseline）	不需要
优势估计	GAE（需 Value 网络）	组内归一化
模型数量	4（Policy + Ref + RM + Critic）	3（Policy + Ref + RM）
显存占用	~4× LLM	~3× LLM
在线探索	是（ε-greedy 或 sample）	是（sample）
稳定性	中（需调 clip、Value loss 等）	较高（无需 Value 网络）
性能	SOTA（大量调参后）	SOTA（DeepSeek-R1 证明）

8.3 GRPO 实现#

1
def grpo_update(
2
    policy_model,
3
    ref_model,
4
    reward_model,
5
    prompts: list[str],
6
    G: int = 8,             # 每个 prompt 采样的回答数
7
    beta: float = 0.04,
8
    clip_ratio: float = 0.2,
9
):
10
    """
11
    GRPO 更新（DeepSeek-R1 风格）。
12
    """
13
    # === 1. 采样 G 个回答 ===
14
    responses = []
15
    policy_log_probs_all = []
16
    ref_log_probs_all = []
17

18
    for prompt in prompts:
19
        prompt_batch = [prompt] * G
20
        # 生成
21
        with torch.no_grad():
22
            outputs = policy_model.generate(
23
                prompt_batch, max_new_tokens=256, do_sample=True, temperature=0.8
24
            )
25

26
        # 提取 response
27
        resp = extract_responses(prompt_batch, outputs)
28

29
        # 计算 log probs
30
        policy_lps = compute_response_log_probs(policy_model, resp)
31
        with torch.no_grad():
32
            ref_lps = compute_response_log_probs(ref_model, resp)
33

34
        responses.extend(resp)
35
        policy_log_probs_all.append(policy_lps)
36
        ref_log_probs_all.append(ref_lps)
37

38
    # === 2. RM 评分 ===
39
    rewards = reward_model.batch_score(responses)  # (len(prompts), G)
40

41
    # === 3. 组内归一化得到优势 ===
42
    mean_r = rewards.mean(dim=-1, keepdim=True)
43
    std_r = rewards.std(dim=-1, keepdim=True) + 1e-8
44
    advantages = (rewards - mean_r) / std_r  # (num_prompts, G)
45

46
    # === 4. GRPO 损失 ===
47
    policy_lps_flat = torch.cat(policy_log_probs_all, dim=0)   # (num_prompts*G,)
48
    ref_lps_flat = torch.cat(ref_log_probs_all, dim=0)
49

50
    # 展平优势（每个 response 一个标量优势）
51
    adv_flat = advantages.flatten()  # (num_prompts*G,)
52

53
    # PPO-style loss
54
    grpo_loss = compute_grpo_loss(
55
        policy_lps_flat, ref_lps_flat, adv_flat,
56
        beta=beta, clip_ratio=clip_ratio,
57
    )
58

59
    return grpo_loss

8.4 何时选 PPO，何时选 GRPO？#

1
选择 PPO：
2
  ✅ 需要追赶最强性能（GPT-4 级别）
3
  ✅ 有足够的 GPU 显存（支持 4 模型）
4
  ✅ 有能力做细致的超参数调优
5
  ✅ 任务需要复杂的多步探索
6

7
选择 GRPO：
8
  ✅ 数学推理、代码生成（DeepSeek-R1 证明）
9
  ✅ 显存受限（3 模型 vs 4 模型）
10
  ✅ 希望简化工程复杂度
11
  ✅ 快速验证对齐效果

9. 超参数调优指南#

9.1 关键超参数表#

超参数	推荐范围	说明
clip $\epsilon$	0.1 ~ 0.3（默认 0.2）	太小→学习慢；太大→不稳定
KL 系数 $\beta$	0.01 ~ 0.2	太小→偏离原始模型；太大→学习慢
GAE $\lambda$	0.9 ~ 0.95	方差-偏差权衡
折扣 $\gamma$	1.0	LLM 终态任务不需要折扣
PPO epochs	4 ~ 10	多 epoch 用同一 rollout 数据
Mini-batch 数	4 ~ 32	分多批更新更稳定
学习率	1e-6 ~ 5e-6	比 SFT 小 10 倍
熵系数	0 ~ 0.01	鼓励探索，太大→策略退化
Value loss 系数	0.5 ~ 1.0	平衡 Policy 和 Value 学习
Gradient clip	1.0	防止梯度爆炸
Rollout 采样数 $G$	4 ~ 16（GRPO）	越大优势估计越稳定

9.2 调参优先级#

1
第一优先（影响最大）：
2
  1. β（KL 系数）—— 直接控制策略偏离程度
3
  2. clip ε —— 控制每次更新步长
4
  3. 学习率 —— 整体更新幅度
5

6
第二优先：
7
  4. GAE λ —— 优势估计质量
8
  5. Value loss 系数 —— Critic 训练速度
9

10
第三优先（微调）：
11
  6. 熵系数
12
  7. PPO epochs
13
  8. Batch size

9.3 训练不稳定信号与对策#

监控指标	危险信号	对策
**KL(π	π_ref)**	> 0.3 / token
RM reward	突然跌到 0 或负	回滚 checkpoint，检查 RM
Value loss	持续增大	预热 Value 网络，增大 vf_coef
Clip fraction	> 20% 被 clip	增大 clip ε
Policy entropy	快速下降→0	增大熵系数
Gradient norm	> 5	减小学习率，增大 gradient clip

1
def log_training_metrics(step, metrics):
2
    """关键监控指标"""
3
    print(f"Step {step}: "
4
          f"KL={metrics['mean_kl']:.4f} | "
5
          f"Reward={metrics['mean_reward']:.3f} | "
6
          f"ClipFrac={metrics['clip_fraction']:.2%} | "
7
          f"VF={metrics['value_loss']:.4f} | "
8
          f"Ent={metrics['entropy']:.3f} | "
9
          f"GradNorm={metrics['grad_norm']:.2f}")

10. 核心公式汇总#

10.1 策略梯度定理#

\nabla_\theta J(\theta) = \mathbb{E}_{\tau\sim\pi_\theta}\Big[\sum_{t=0}^{T-1}\nabla_\theta\log\pi_\theta(a_t|s_t)\;G_t\Big]

10.2 重要性比率#

r_t(\theta) = \frac{\pi_\theta(a_t|s_t)}{\pi_{\theta_{\text{old}}}(a_t|s_t)}

10.3 PPO Clipped Surrogate#

\mathcal{L}^{\text{CLIP}}(\theta) = \mathbb{E}_t\Big[\min\big(r_t(\theta)\,A_t,\ \text{clip}(r_t(\theta), 1-\epsilon, 1+\epsilon)\,A_t\big)\Big]

10.4 GAE 递推#

\hat{A}_t = \delta_t + \gamma\,\lambda\,(1-d_t)\,\hat{A}_{t+1}

其中 $\delta_t = r_t + \gamma\,V(s_{t+1}) - V(s_t)$ ， $d_t$ 是结束信号。

10.5 Value 损失#

\mathcal{L}_V = \mathbb{E}\Big[\big(R_t - V_\phi(s_t)\big)^2\Big]

10.6 KL 奖励塑形#

r_t^{\text{shaped}} = r_t - \beta\,\text{KL}\big(\pi_\theta(\cdot|s_t)\,\|\,\pi_{\text{ref}}(\cdot|s_t)\big)

10.7 GRPO 组内归一化优势#

A_i = \frac{r_i - \text{mean}(\mathbf{r})}{\text{std}(\mathbf{r})}

10.8 TRPO 信赖域约束#

\mathbb{E}_t\Big[\text{KL}\big(\pi_{\theta_{\text{old}}}(\cdot|s_t)\,\|\,\pi_\theta(\cdot|s_t)\big)\Big] \leq \delta

11. 总结#

11.1 PPO 的核心价值#

1
┌─────────────────────────────────────────────────────────────┐
2
│                    PPO 的三个关键创新                          │
3
├─────────────────────────────────────────────────────────────┤
4
│                                                              │
5
│  1. Clip 机制：无需二阶优化即可实现信赖域约束                  │
6
│     比 TRPO 简单 10 倍，比 Actor-Critic 稳定 10 倍            │
7
│                                                              │
8
│  2. 置信域思想：限制新旧策略比率的变化范围                      │
9
│     防止策略在单次更新中崩溃                                   │
10
│                                                              │
11
│  3. GAE：高效的方差-偏差权衡优势估计                           │
12
│     一个参数 λ 控制整个谱的估计质量                            │
13
│                                                              │
14
└─────────────────────────────────────────────────────────────┘

11.2 算法演进图谱#

1
REINFORCE (1992)
2
    ↓ 方差极大
3
Actor-Critic (1999)
4
    ↓ 无步长保证
5
TRPO (2015)
6
    ↓ 二阶优化太复杂
7
PPO (2017) ←──────────────────────────────┐
8
    ↓                                      │
9
    ├── PPO-Penalty（KL 奖励塑形）           │
10
    ├── PPO with Value Clip                 │
11
    └── Adaptive KL Penalty                 │
12
                                             │
13
                          ┌─────────────────┘
14
                          ↓
15
          GRPO (DeepSeek-R1, 2024)
16
            → 去掉 Critic，用组内排名替代
17
            → 在数学推理任务上达到 SOTA

11.3 工程实践建议#

Warning

PPO 是 LLM 对齐中最复杂的训练范式。在决定使用 PPO 之前，请确认：

你有足够的 GPU 显存（至少 2×80G A100 用于 7B 模型，LoRA 优化）
你有足够的工程能力处理 4 模型协调、early stopping、gradient checkpointing
你的目标是追赶最强性能，而非快速验证

如果不确定，先用 DPO 或 GRPO，在验证有效后再升级到 PPO。

11.4 未来方向#

Adaptive Clip：根据训练动态自动调整 clip ε
V-MPO：基于最大后验的策略优化，比 PPO 更适合离策略
Dreamer/Model-based PPO：结合世界模型，减少样本复杂度
On-policy → Off-policy：PPO 的 off-policy 变体，减少采样浪费

参考资料#

Schulman, J., et al. (2017). “Proximal Policy Optimization Algorithms.” arXiv.
Schulman, J., et al. (2015). “Trust Region Policy Optimization.” ICML.
Schulman, J., et al. (2016). “High-Dimensional Continuous Control Using Generalized Advantage Estimation.” ICLR.
Ouyang, L., et al. (2022). “Training language models to follow instructions with human feedback.” NeurIPS.
Bai, Y., et al. (2022). “Training a Helpful and Harmless Assistant with Reinforcement Learning from Human Feedback.” arXiv.
Shao, Z., et al. (2024). “DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models.” arXiv (GRPO).
Guo, D., et al. (2025). “DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning.” arXiv.
Engstrom, L., et al. (2020). “Implementation Matters in Deep RL: A Case Study on PPO and TRPO.” ICLR.
Irving, G., et al. (2018). “AI safety via debate.” arXiv.
Lambert, N., et al. (2022). “Tune: A Research Platform for Model Tuning.” GitHub.

Federico

PPO 深度解析：信赖域优化的艺术与 LLM 对齐实践

1. 引言：PPO 为什么会成为 LLM 对齐的主流？#

1.1 强化学习的三代算法#

1.2 为什么 LLM 对齐选择 PPO？#

2. 策略梯度基础：从 REINFORCE 到 Actor-Critic#

2.1 策略梯度定理#

2.2 REINFORCE：基本策略梯度#

2.3 方差问题的根源#

2.4 Actor-Critic：引入 Value Network#

3. TRPO：信赖域优化的理论基础#

3.1 信赖域的核心思想#

3.2 从约束到惩罚：共轭梯度法#

3.3 TRPO 的工程痛点#

4. PPO：Clip 机制的数学推导#

4.1 重要性比率#

4.2 原始策略梯度损失#

4.3 Clipped Surrogate 目标#

4.4 Clip 机制的直觉#

4.5 完整的数学推导#

4.6 为什么 min 而不是 max？#

5. GAE：广义优势估计#

5.1 TD(λ) 的优势函数版本#

5.2 GAE 的递推形式#

5.3 GAE 超参数的影响#

5.4 Value 网络损失#

6. PPO 用于 LLM 对齐：完整实现#

6.1 四模型架构#

6.2 Rollout：在线采样#

6.3 PPO 更新#

6.4 KL 约束的奖励塑形形式#

7. 进阶技巧：让 PPO 训练更稳定#

7.1 Reward 归一化#

7.2 Reward Clipping#

7.3 KL Annealing#

7.4 Value 网络预热#

7.5 早停机制#

8. PPO vs GRPO：两条技术路线#

8.1 GRPO 的核心洞察#

8.2 GRPO vs PPO 详细对比#

8.3 GRPO 实现#

8.4 何时选 PPO，何时选 GRPO？#

9. 超参数调优指南#

9.1 关键超参数表#

9.2 调参优先级#

9.3 训练不稳定信号与对策#

10. 核心公式汇总#

10.1 策略梯度定理#

10.2 重要性比率#

10.3 PPO Clipped Surrogate#

10.4 GAE 递推#

10.5 Value 损失#

10.6 KL 奖励塑形#

10.7 GRPO 组内归一化优势#

10.8 TRPO 信赖域约束#

11. 总结#

11.1 PPO 的核心价值#

11.2 算法演进图谱#

11.3 工程实践建议#

11.4 未来方向#

参考资料#

文章分享

Federico