# 第 25 课：Reward、Logprob 与 KL Penalty

这一课不是再讲一遍 PPO 或 GRPO，而是把第 21-24 课反复出现的三个核心量统一整理清楚：`reward`、`logprob`、`KL penalty`。看懂这一节，后面读 `train_agent.py` 时就不会把 reward model、old logprob、reference model、approx KL 混在一起。

## 目录

- [0. 本节主线](#l25-mainline)
- [1. 本节要懂的 7 个原理](#l25-principles)
- [2. 完整原理：三个量分别管什么](#l25-complete-principle)
- [3. 源码阅读顺序图](#l25-reading-order)
- [4. MiniMind 源码走读](#l25-source-walkthrough)
- [5. 本节必须会写 / 暂时不要求](#l25-must-write)
- [6. 手写模块复盘](#l25-handwrite)
- [7. 实验验证](#l25-experiment)
- [8. 阶段组装](#l25-stage-assembly)
- [9. 本节检查](#l25-check)
- [10. 下一课](#l25-next)

<a id="l25-mainline"></a>
## 0. 本节主线

MiniMind 的 RL 训练里，这条链路反复出现：

```text
rollout policy 生成 response
-> 保存 old_logp，也就是采样当时每个 response token 的 logprob
-> reward model / 规则函数给整条 response 一个 sequence reward
-> PPO 把 sequence reward 放到最后一个有效 token，再用 critic + GAE 得到 token advantage
-> GRPO 把同 prompt 多条 response 的 reward 做 group mean/std，得到 sequence advantage
-> 训练时用 current policy 重新计算 policy_logp
-> 用 frozen reference model 计算 ref_logp
-> ratio = exp(policy_logp - old_logp)
-> reference KL penalty = exp(ref_logp - policy_logp) - (ref_logp - policy_logp) - 1
-> mask 掉 padding / EOS 后无效 token
-> policy loss 反向更新 policy
```

一句话：

```text
reward 告诉模型哪条回答好；
logprob 连接 policy 参数并形成 ratio；
KL penalty 约束 policy 不要离 reference 太远。
```

这三个量的角色不能混：

```text
reward:
  来自外部评价或规则，通常不反传梯度。

policy_logp:
  来自当前正在训练的 policy，是策略梯度的主要路径。

old_logp:
  来自 rollout 时的旧 policy，用来算 ratio。

ref_logp:
  来自冻结 reference model，用来做 KL 约束或 DPO 偏好基准。
```

<a id="l25-principles"></a>
## 1. 本节要懂的 7 个原理

| 原理 | 要理解什么 | 源码位置 |
|---|---|---|
| reward 是 sequence 级信号 | MiniMind 先对整条 response 打分，不是每个 token 单独打分 | `trainer/train_ppo.py:52-76`, `trainer/train_grpo.py:37-68` |
| reward model wrapper 只负责打分 | `LMForRewardModel.get_score` 把对话历史和回答交给 reward model，并裁剪分数范围 | `trainer/trainer_utils.py:160-177` |
| old logprob 来自 rollout | 采样 response 时同时保存每个 response token 的 logprob | `trainer/rollout_engine.py:23-36`, `trainer/rollout_engine.py:71-92` |
| current/ref logprob 在 update 时重算 | 训练时对同一批 `output_ids` 重算 policy/ref logprob | `trainer/train_ppo.py:117-139`, `trainer/train_grpo.py:97-105` |
| reward 到 advantage 的算法不同 | PPO 用 critic + GAE；GRPO 用 group reward 标准化 | `trainer/train_ppo.py:140-155`, `trainer/train_grpo.py:121-124` |
| KL 有不同用途 | `approx_kl` 监控 policy-old 距离；reference KL penalty 约束 policy-ref 距离 | `trainer/train_ppo.py:186-205`, `trainer/train_grpo.py:132-143` |
| DPO 的 reference 不是 token KL | DPO 用 chosen/rejected 的 reference logratio 作为偏好基准 | `trainer/train_dpo.py:34-49` |

学完本节，你应该能独立说清楚这些 tensor：

```text
old_logps:        rollout policy 采样时的 response token logprob
policy_logps:     当前 policy 对同一 response token 的 logprob
ref_logps:        frozen reference 对同一 response token 的 logprob
rewards:          每条 response 或 trajectory 的外部分数
advantages:       reward 经过 PPO/GRPO 规则转换后的训练权重
ratio:            exp(policy_logps - old_logps)
ref_kl_penalty:   exp(ref_logps - policy_logps) - (ref_logps - policy_logps) - 1
mask:             哪些 response token 参与 loss
```

<a id="l25-complete-principle"></a>
## 2. 完整原理：三个量分别管什么

### 2.1 Reward：回答好不好

MiniMind 的 PPO/GRPO 都先生成完整 response，再计算一个标量 reward：

$$
R_i \in \mathbb{R}
$$

这个 reward 可以来自：

```text
规则项：
  长度是否合理、thinking 格式是否合理、重复惩罚等。

reward model：
  用外部模型给 prompt/response 打分。
```

它不是直接对 policy 参数求导的东西。更准确地说：

```text
reward 是训练信号来源；
policy_logp 才是把这个信号传回模型参数的路径。
```

PPO 会把 sequence reward 放到最后一个有效 response token 上：

$$
r_{i,t}
=
\begin{cases}
R_i, & t = T_i \\
0, & \text{otherwise}
\end{cases}
$$

然后用 critic 的 value estimate 做 GAE，得到 token-level advantage：

$$
A_{i,t}
$$

GRPO 不训练 critic，而是同一个 prompt 采样多条 response：

$$
R_{i,1}, R_{i,2}, \dots, R_{i,G}
$$

再做组内标准化：

$$
A_{i,j}
=
\frac{
R_{i,j} - \mu_i
}{
\sigma_i + \epsilon
}
$$

所以：

```text
PPO:
  sequence reward -> last token reward -> GAE -> token advantage

GRPO:
  group sequence rewards -> group normalization -> sequence advantage
```

### 2.2 Logprob：policy 更新的入口

对一个已经生成出来的 token：

$$
a_t
$$

current policy 的 logprob 是：

$$
\ell_t^\theta
=
\log \pi_\theta(a_t \mid s_t)
$$

rollout old policy 的 logprob 是：

$$
\ell_t^{old}
=
\log \pi_{old}(a_t \mid s_t)
$$

reference model 的 logprob 是：

$$
\ell_t^{ref}
=
\log \pi_{ref}(a_t \mid s_t)
$$

MiniMind 里这些 logprob 都是通过同一个基本操作拿到的：

```text
logits -> log_softmax -> gather 目标 token id
```

即：

$$
\log p(a_t)
=
\operatorname{log\_softmax}(\operatorname{logits}_t)_{a_t}
$$

注意 logprob 的对齐关系：

```text
logits[:, k] 预测 input_ids[:, k + 1]
```

所以 response 第 `j` 个 token 的 logprob，要从：

```text
prompt_len - 1 + j
```

这个 logits 位置取。

### 2.3 Ratio：current policy 相对 old policy 变了多少

PPO/GRPO 的 ratio 是：

$$
r_t
=
\frac{
\pi_\theta(a_t \mid s_t)
}{
\pi_{old}(a_t \mid s_t)
}
=
\exp
\left(
\ell_t^\theta - \ell_t^{old}
\right)
$$

它回答的问题是：

```text
当前 policy 对当时采样出来的这个 token，比 old policy 更想生成还是更不想生成？
```

如果：

```text
ratio > 1
```

说明 current policy 比 old policy 更偏向这个 token。

如果：

```text
ratio < 1
```

说明 current policy 比 old policy 更不偏向这个 token。

ratio 进入 PPO/GRPO/CISPO 的 policy loss；它不是 reward，也不是 KL。

### 2.4 Reference KL penalty：别离 SFT reference 太远

MiniMind 的 PPO/GRPO 使用同一个 token-level reference KL 近似：

$$
\Delta_t
=
\ell_t^{ref} - \ell_t^\theta
$$

$$
\operatorname{KLPenalty}_t
=
\exp(\Delta_t) - \Delta_t - 1
$$

源码里对应：

```python
kl_div = ref_logp - policy_logp
per_token_kl = torch.exp(kl_div) - kl_div - 1
```

它的作用是：

```text
policy 可以因为 reward 往某些方向移动；
但不要离 frozen reference 太远。
```

### 2.5 Approx KL：PPO 早停用，不是 reference KL

PPO 里还有一个：

$$
\operatorname{approx\_kl}
=
\frac{1}{2}
\left(
\ell_t^\theta - \ell_t^{old}
\right)^2
$$

MiniMind 用它判断当前 policy 相对 rollout old policy 是否更新过猛：

```text
如果 approx_kl 太大，就停止本轮 PPO 多 epoch 更新。
```

所以要区分：

| 名称 | 比较对象 | 主要用途 |
|---|---|---|
| `ratio` | current policy vs old policy | policy loss 的概率比 |
| `approx_kl` | current policy vs old policy | PPO update early stop |
| `reference KL penalty` | current policy vs reference model | 防止偏离 SFT reference |
| DPO `ref_logratio` | reference chosen vs rejected | 偏好基准，不是 token KL penalty |

<a id="l25-reading-order"></a>
## 3. 源码阅读顺序图

建议按这个顺序读：

```text
1. trainer/trainer_utils.py:160-177
   先看 reward model wrapper 返回什么。

2. trainer/train_ppo.py:52-76
   看 PPO 的规则 reward + reward model score 如何相加。

3. trainer/rollout_engine.py:23-36
   看 rollout 阶段如何保存 old per-token logprob。

4. trainer/train_ppo.py:117-139
   看 PPO 如何定位 response token，并计算 ref logprob。

5. trainer/train_ppo.py:140-155
   看 PPO 如何把 sequence reward 变成 token advantage。

6. trainer/train_ppo.py:186-205
   看 ratio、approx KL、reference KL 和 policy loss。

7. trainer/train_grpo.py:121-143
   对比 GRPO 如何用 group advantage 和同一套 KL/ratio。

8. trainer/train_dpo.py:34-49
   最后用 DPO 对照：reference logprob 也能服务偏好 logratio，而不是只服务 KL penalty。
```

<a id="l25-source-walkthrough"></a>
## 4. MiniMind 源码走读

### 4.1 reward model wrapper：外部模型只负责打分

#### 源码证据：`LMForRewardModel`

文件：`trainer/trainer_utils.py:160-177`

看它是为了理解：MiniMind 把 reward model 封装成 `get_score(messages, response)`，训练脚本只拿到一个标量分数。

代码摘录：

```python
class LMForRewardModel:
    def __init__(self, model_path, device="cuda", dtype=torch.float16):
        self.tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
        self.model = AutoModel.from_pretrained(model_path, torch_dtype=dtype, trust_remote_code=True)
        self.model = self.model.to(device).eval()

    @torch.no_grad()
    def get_score(self, messages, response):
        ...
        score = self.model.get_score(self.tokenizer, eval_messages)
        return max(min(score, 3.0), -3.0)
```

这段代码说明：

- reward model 是 `eval()`，不会被 PPO/GRPO 更新。
- `get_score` 返回一个 Python 标量分数。
- 分数被裁剪到 `[-3.0, 3.0]`，避免极端 reward。
- 这个分数后面会和规则 reward 相加。

### 4.2 PPO/GRPO 的 reward：规则项 + reward model score

#### 源码证据 A：PPO reward

文件：`trainer/train_ppo.py:52-76`

看它是为了理解：PPO 对每条 response 计算一个 sequence reward。

代码摘录：

```python
def calculate_rewards(prompts, responses, reward_model):
    rewards = torch.zeros(len(responses), device=args.device)
    ...
    rewards[i] += 0.5 if 20 <= len(response.strip()) <= 800 else -0.5
    if '</think>' in response:
        rewards[i] += 1.0 if 20 <= len(thinking_content.strip()) <= 300 else -0.5
        rewards[i] += 0.25 if response.count('</think>') == 1 else -0.25
    rewards[i] -= rep_penalty(answer)
    score = reward_model.get_score(messages, answer)
    reward_model_scores.append(score)
    ...
    rewards += reward_model_scores
```

这段代码说明：

- reward 不是纯 reward model 分数。
- MiniMind 还加了长度、thinking 格式、重复惩罚等规则项。
- 返回的 `rewards` shape 是 `[B]`。

#### 源码证据 B：GRPO reward

文件：`trainer/train_grpo.py:37-68`

看它是为了理解：GRPO 的 reward shape 是 `[B * num_generations]`，因为每个 prompt 会生成多条 response。

代码摘录：

```python
for i in range(batch_size):
    for j in range(args.num_generations):
        response_idx = i * args.num_generations + j
        response = responses[response_idx]
        prompt = prompts[i]
        ...
        rewards[response_idx] -= rep_penalty(answer)
        score = reward_model.get_score(messages, answer)
        reward_model_scores.append(score)
...
rewards += reward_model_scores
```

这段代码说明：

- flat 排列顺序是 `prompt_i` 的第 `j` 条生成。
- 这个顺序决定后面可以用 `rewards.view(-1, num_generations)`。
- reward 仍然是 sequence 级，不是 token 级。

### 4.3 old logprob：rollout 阶段保存采样策略概率

#### 源码证据 A：通用 token logprob helper

文件：`trainer/rollout_engine.py:23-36`

看它是为了理解：old logprob 本质也是 `log_softmax + gather`。

代码摘录：

```python
def compute_per_token_logps(model, input_ids, n_keep, attention_mask=None):
    logits = unwrapped(input_ids, attention_mask=attention_mask, logits_to_keep=n_keep + 1).logits[:, :-1, :]
    per_token_logps = []
    for logits_row, ids_row in zip(logits, input_ids[:, -n_keep:]):
        per_token_logps.append(
            torch.gather(logits_row.log_softmax(dim=-1), 1, ids_row.unsqueeze(1)).squeeze(1)
        )
    return torch.stack(per_token_logps)
```

这段代码说明：

- `n_keep` 是 response token 数量。
- `input_ids[:, -n_keep:]` 是要取 logprob 的 completion tokens。
- 返回 shape 是 `[B, R]`。
- 这个 helper 用在 rollout 阶段，所以结果就是 `old_logp`。

#### 源码证据 B：Torch rollout 保存 old logprob

文件：`trainer/rollout_engine.py:71-92`

看它是为了理解：生成 response 后立刻保存 `per_token_logps`。

代码摘录：

```python
output_ids = model.generate(...)
prompt_len = prompt_ids.size(1)
completion_ids = output_ids[:, prompt_len:]
full_mask = (output_ids != self.tokenizer.pad_token_id).long()
per_token_logps = compute_per_token_logps(
    self.policy_model,
    output_ids,
    completion_ids.size(1),
    attention_mask=full_mask
)
return RolloutResult(output_ids, completion_ids, per_token_logps, ...)
```

这段代码说明：

- `output_ids` 是 prompt + response。
- `completion_ids` 是 response 部分。
- `per_token_logps` 是 rollout policy 对 response tokens 的 logprob。
- PPO/GRPO 里后续会把它命名为 `old_resp_logp` 或 `old_per_token_logps`。

#### 源码证据 C：SGLang rollout 也返回 logprob

文件：`trainer/rollout_engine.py:115-173`

看它是为了理解：即使用服务端 rollout，也要拿回同样语义的 `per_token_logps` 和 `completion_mask`。

代码摘录：

```python
payload = {
    "input_ids": all_input_ids,
    "sampling_params": {...},
    "return_logprob": True,
}
...
raw_logprobs = meta.get("output_token_logprobs", [])
...
return RolloutResult(
    output_ids=...,
    completion_ids=...,
    per_token_logps=pad_to_tensor(all_logprobs, max_comp_len, pad_val=0.0),
    completion_mask=torch.tensor([[1] * len(ids) + [0] * ... for ids in all_completion_ids], device=device),
)
```

这段代码说明：

- SGLang 返回的 logprob 也被整理成 `[B * num_generations, R]`。
- SGLang 的 `completion_mask` 能标出 padding。
- Torch rollout 的 mask 可能全 1，SGLang rollout 的 mask 更接近真实变长 completion。

### 4.4 current/ref logprob：训练时对同一批 token 重算

#### 源码证据 A：PPO response 位置和 ref logprob

文件：`trainer/train_ppo.py:117-139`

看它是为了理解：response token 的 logprob 要从 `logits[:, :-1]` 的对应位置取。

代码摘录：

```python
labels = gen_out[:, 1:].clone()
resp_idx = torch.arange(resp_labels.size(1), device=gen_out.device).unsqueeze(0)
logp_pos = prompt_lens.unsqueeze(1) - 1 + resp_idx
...
ref_resp_logp = F.log_softmax(
    ref_model(input_ids=gen_out, attention_mask=full_mask).logits[:, :-1],
    dim=-1
).gather(2, labels.unsqueeze(-1)).squeeze(-1).gather(1, logp_pos)
```

这段代码说明：

- `labels = gen_out[:, 1:]` 表示每个位置预测下一个 token。
- `logp_pos = prompt_len - 1 + response_idx` 是 response token 对应的 logits 位置。
- `ref_resp_logp` shape 是 `[B, R]`。
- reference model 是 frozen，只提供约束基准。

#### 源码证据 B：GRPO current/ref logprob

文件：`trainer/train_grpo.py:97-105`

看它是为了理解：GRPO 的 policy/ref logprob 和 PPO 是同一类 gather 操作。

代码摘录：

```python
res = model_unwrapped(outputs, attention_mask=full_mask)
per_token_logps = F.log_softmax(res.logits[:, :-1, :], dim=-1).gather(
    2, outputs[:, 1:].unsqueeze(-1)
).squeeze(-1).gather(1, logp_pos)

with torch.no_grad():
    ref_per_token_logps = F.log_softmax(
        ref_model(outputs, attention_mask=full_mask).logits[:, :-1, :],
        dim=-1
    ).gather(2, outputs[:, 1:].unsqueeze(-1)).squeeze(-1).gather(1, logp_pos)
```

这段代码说明：

- `per_token_logps` 是 current policy，参与反传。
- `ref_per_token_logps` 在 `torch.no_grad()` 下计算，不参与训练。
- 两者都只取 response token 位置。

### 4.5 reward 怎么变成 advantage

#### 源码证据 A：PPO 把 sequence reward 放到最后有效 token

文件：`trainer/train_ppo.py:140-155`

看它是为了理解：PPO 的 reward 是 sequence 级，但 advantage 是 token 级。

代码摘录：

```python
token_rewards = torch.zeros_like(old_resp_logp)
last_idx = resp_lengths - 1
token_rewards[torch.arange(B, device=args.device)[valid_resp], last_idx[valid_resp]] += rewards[valid_resp]

for t in reversed(range(gen_len)):
    nv = old_resp_values[:, t + 1] if t < gen_len - 1 else 0.0
    delta = token_rewards[:, t] + args.gamma * nv - old_resp_values[:, t]
    lastgaelam = delta + args.gamma * args.lam * lastgaelam
    advs_rev.append(lastgaelam)
advantages = torch.stack(advs_rev[::-1], dim=1)
returns = advantages + old_resp_values
```

这段代码说明：

- 外部 reward 只加在最后一个有效 token。
- Critic 的 `old_resp_values` 参与 GAE。
- 最后得到的 `advantages` shape 是 `[B, R]`。
- padding / EOS 后无效位置后面还会用 mask 去掉。

#### 源码证据 B：GRPO 用组内 reward 标准化

文件：`trainer/train_grpo.py:121-124`

看它是为了理解：GRPO 不需要 critic，直接把 group reward 变成 sequence advantage。

代码摘录：

```python
grouped_rewards = rewards.view(-1, args.num_generations)
mean_r = grouped_rewards.mean(dim=1).repeat_interleave(args.num_generations)
std_r = grouped_rewards.std(dim=1, unbiased=False).repeat_interleave(args.num_generations)
advantages = (rewards - mean_r) / (std_r + 1e-4)
```

这段代码说明：

- `grouped_rewards` shape 是 `[B, num_generations]`。
- 每条 response 的 advantage 是同组相对分数。
- `advantages` shape 是 `[B * num_generations]`，后面通过 `unsqueeze(1)` 广播到 token 维。

### 4.6 KL 的两个名字：approx KL 和 reference KL penalty

#### 源码证据 A：PPO update 中的 ratio 与 KL

文件：`trainer/train_ppo.py:186-205`

看它是为了理解：PPO 同时计算 policy-old 的 approx KL 和 policy-reference 的 KL penalty。

代码摘录：

```python
log_ratio = mb_resp_logp - old_resp_logp[inds]
approx_kl = (0.5 * (log_ratio ** 2) * resp_policy_mask[inds]).sum() / resp_policy_mask[inds].sum().clamp(min=1)
...
ratio = torch.exp(log_ratio)
kl_ref_penalty = ((torch.exp(ref_resp_logp[inds] - mb_resp_logp) - (ref_resp_logp[inds] - mb_resp_logp) - 1.0)
                  * resp_policy_mask[inds]).sum() / resp_policy_mask[inds].sum().clamp(min=1)
policy_loss = (
    masked_ppo_loss
    + args.kl_coef * kl_ref_penalty
)
```

这段代码说明：

- `log_ratio` 比较 current policy 和 rollout old policy。
- `approx_kl` 用于 early stop，不是加到 loss 的 reference KL。
- `kl_ref_penalty` 比较 current policy 和 frozen reference。
- policy loss 里加的是 `kl_coef * kl_ref_penalty`。

#### 源码证据 B：GRPO/CISPO 中的 reference KL

文件：`trainer/train_grpo.py:132-143`

看它是为了理解：GRPO/CISPO 没有 PPO 的 critic/value loss，但仍然使用 reference KL penalty。

代码摘录：

```python
kl_div = ref_per_token_logps - per_token_logps
per_token_kl = torch.exp(kl_div) - kl_div - 1
ratio = torch.exp(per_token_logps - old_per_token_logps)
if args.loss_type == "cispo":
    clamped_ratio = torch.clamp(ratio, max=args.epsilon_high).detach()
    per_token_loss = -(clamped_ratio * advantages.unsqueeze(1) * per_token_logps - args.beta * per_token_kl)
else:
    clipped_ratio = torch.clamp(ratio, 1 - args.epsilon, 1 + args.epsilon)
    per_token_loss1 = ratio * advantages.unsqueeze(1)
    per_token_loss2 = clipped_ratio * advantages.unsqueeze(1)
    per_token_loss = -(torch.min(per_token_loss1, per_token_loss2) - args.beta * per_token_kl)
```

这段代码说明：

- `per_token_kl` 是 token-level reference KL penalty。
- GRPO 和 CISPO 共享 `ratio`、`advantages`、`per_token_kl`。
- 区别只在 policy loss 形式。

### 4.7 DPO 对照：reference logprob 不是只能做 KL

#### 源码证据：DPO logratio

文件：`trainer/train_dpo.py:34-49`

看它是为了理解：DPO 也有 policy/ref logprob，但它不是 token-level KL penalty，而是 chosen/rejected 偏好差。

代码摘录：

```python
ref_log_probs = (ref_log_probs * mask).sum(dim=1)
policy_log_probs = (policy_log_probs * mask).sum(dim=1)
...
pi_logratios = chosen_policy_log_probs - reject_policy_log_probs
ref_logratios = chosen_ref_log_probs - reject_ref_log_probs
logits = pi_logratios - ref_logratios
loss = -F.logsigmoid(beta * logits)
```

这段代码说明：

- DPO 先把 token logprob 沿 sequence 求和。
- `ref_logratios` 是 reference 对 chosen/rejected 的偏好基准。
- DPO 的 reference 用法是“偏好相对提升”，不是 `exp(ref-policy)-...` 这种 KL penalty。

<a id="l25-must-write"></a>
## 5. 本节必须会写 / 暂时不要求

### 5.1 必须会写

这节课不新增一个大的训练模块，但你必须能独立写出这些小公式：

```python
log_probs = F.log_softmax(logits, dim=-1)
token_logps = log_probs.gather(2, labels.unsqueeze(-1)).squeeze(-1)
response_logps = token_logps.gather(1, logp_pos)
```

```python
ratio = torch.exp(policy_logps - old_logps)
```

```python
delta = ref_logps - policy_logps
ref_kl_penalty = torch.exp(delta) - delta - 1.0
```

```python
approx_kl = 0.5 * (policy_logps - old_logps) ** 2
```

```python
masked_mean = (x * mask).sum() / mask.sum().clamp(min=1)
```

以及 GRPO 的 group advantage：

```python
grouped_rewards = rewards.view(-1, num_generations)
mean_r = grouped_rewards.mean(dim=1).repeat_interleave(num_generations)
std_r = grouped_rewards.std(dim=1, unbiased=False).repeat_interleave(num_generations)
advantages = (rewards - mean_r) / (std_r + 1e-4)
```

### 5.2 暂时不要求

本节暂时不要求：

```text
真正加载 reward model 权重；
启动 SGLang server；
写完整 PPO/GRPO 训练脚本；
调 reward shaping 超参；
证明 KL 近似公式的理论来源。
```

现在只要求你能读懂源码中每个张量的身份。

<a id="l25-handwrite"></a>
## 6. 手写模块复盘

本节对应已有手写文件：

```text
course/impl/core/ppo.py
course/impl/core/grpo.py
```

你应该回头确认：

```text
course/impl/core/ppo.py:
  response_logp_positions
  reference_kl_penalty
  ppo_policy_loss

course/impl/core/grpo.py:
  group_relative_advantages
  token_kl_penalty
  grpo_policy_loss
  cispo_policy_loss
```

这些函数共同覆盖了第 25 课的核心：

```text
logprob 位置对齐；
ratio；
mask；
reference KL penalty；
reward -> advantage。
```

验收命令：

```bash
cd /home/sun/minimind
python course/impl/tests/test_ppo.py
python course/impl/tests/test_grpo.py
```

如果 `test_grpo.py` 的 policy loss 失败，优先检查 reduction：

```python
policy_loss = (
    (per_token_loss * completion_mask).sum(dim=1)
    / completion_mask.sum(dim=1).clamp(min=1)
).mean()
```

不要写成全 batch token 直接平均，二者在 response 长度不一致时不同。

<a id="l25-experiment"></a>
## 7. 实验验证

本节新增实验：

```text
course/labs/trace_reward_logprob_kl.py
```

它不加载真实模型，只用固定 logits/token/reward 张量复现 MiniMind 的关键公式。

运行：

```bash
cd /home/sun/minimind
python course/labs/trace_reward_logprob_kl.py
```

它会打印：

```text
[Token logprobs]
old_logps / policy_logps / ref_logps 的 shape

[Reward to advantage]
PPO sequence reward 如何落到最后有效 token
GRPO group rewards 如何变成 group-relative advantages

[KL and ratio]
ratio = exp(policy-old)
approx_kl = 0.5 * (policy-old)^2
reference_kl_penalty = exp(ref-policy) - (ref-policy) - 1

[DPO contrast]
DPO 的 policy/ref logratio 和 reference KL penalty 的区别
```

判断实验正确的重点不是某个固定数值，而是这些关系：

```text
old_logps.shape == policy_logps.shape == ref_logps.shape == response_mask.shape
ratio.shape == policy_logps.shape
reference_kl_penalty.shape == policy_logps.shape
PPO rewards.shape == [B]
GRPO rewards.shape == [B * num_generations]
```

<a id="l25-stage-assembly"></a>
## 8. 阶段组装

到这里，PPO/GRPO 阶段的核心张量已经完整：

```text
rollout_engine
-> output_ids / completion_ids / old_logps / completion_mask
-> reward model / rules
-> PPO: token_rewards + critic values -> GAE advantage
-> GRPO: group rewards -> group advantage
-> current policy logps
-> reference logps
-> ratio / KL penalty
-> policy loss
```

后面组装 `course/impl/train_grpo_impl.py` 时，不需要复刻 MiniMind 的所有工程外围。最小教学版可以只保留：

```text
synthetic prompts 或 tiny RLAIF prompt
-> rollout 生成 response
-> 简化 reward 函数
-> group_relative_advantages
-> token_kl_penalty
-> grpo_policy_loss / cispo_policy_loss
-> optimizer step
```

本节的约束是：

```text
所有 loss 函数必须明确：
  哪个 logprob 来自 old policy；
  哪个 logprob 来自 current policy；
  哪个 logprob 来自 reference model；
  哪个 mask 决定有效 response token；
  reward 是如何进入 advantage 的。
```

<a id="l25-check"></a>
## 9. 本节检查

1. `old_logp` 和 `policy_logp` 的区别是什么？
2. 为什么 PPO/GRPO 不能只保存 response 文本，还要保存 old per-token logprob？
3. `ref_logp - policy_logp` 在 PPO/GRPO 中用来算什么？
4. `approx_kl = 0.5 * log_ratio^2` 比较的是哪两个 policy？
5. PPO 为什么把 sequence reward 加到最后一个有效 response token？
6. GRPO 为什么要求同一个 prompt 的多条 response 在 flat tensor 中连续排列？
7. DPO 的 `ref_logratios` 和 PPO/GRPO 的 reference KL penalty 有什么区别？

<a id="l25-next"></a>
## 10. 下一课

第 26 课进入 Tool Use 与 Agentic RL。

下一课要解决：

```text
train_agent.py 如何组织 tool schema、tool_call、tool_response；
一次多轮 trajectory 如何 rollout；
环境反馈 reward 和普通 reward model score 有什么区别；
Agentic RL 为什么是延迟 reward；
MiniMind 如何把 tool calling 接回 chat template。
```