26 KiB

Raw Blame History

第 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. 本节主线

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

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

一句话：

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

这三个量的角色不能混：

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

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

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

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

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：

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

2. 完整原理：三个量分别管什么

2.1 Reward：回答好不好

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


R_i \in \mathbb{R}

这个 reward 可以来自：

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

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

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

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
}

所以：

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 都是通过同一个基本操作拿到的：

logits -> log_softmax -> gather 目标 token id

即：


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

注意 logprob 的对齐关系：

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

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

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)

它回答的问题是：

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

如果：

ratio > 1

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

如果：

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

源码里对应：

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

它的作用是：

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 是否更新过猛：

如果 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

3. 源码阅读顺序图

建议按这个顺序读：

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。

4. MiniMind 源码走读

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

源码证据：`LMForRewardModel`

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

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

代码摘录：

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。

代码摘录：

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。

代码摘录：

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。

代码摘录：

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。

代码摘录：

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。

代码摘录：

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] 的对应位置取。

代码摘录：

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 操作。

代码摘录：

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 级。

代码摘录：

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。

代码摘录：

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。

代码摘录：

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。

代码摘录：

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 偏好差。

代码摘录：

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。

5. 本节必须会写 / 暂时不要求

5.1 必须会写

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

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)

ratio = torch.exp(policy_logps - old_logps)

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

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

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

以及 GRPO 的 group advantage：

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 暂时不要求

本节暂时不要求：

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

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

6. 手写模块复盘

本节对应已有手写文件：

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

你应该回头确认：

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 课的核心：

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

验收命令：

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

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

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

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

7. 实验验证

本节新增实验：

course/labs/trace_reward_logprob_kl.py

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

运行：

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

它会打印：

[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 的区别

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

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]

8. 阶段组装

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

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 的所有工程外围。最小教学版可以只保留：

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

本节的约束是：

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

9. 本节检查

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

10. 下一课

第 26 课进入 Tool Use 与 Agentic RL。

下一课要解决：

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

26 KiB Raw Blame History Unescape Escape

第 25 课：Reward、Logprob 与 KL Penalty

目录

0. 本节主线

1. 本节要懂的 7 个原理

2. 完整原理：三个量分别管什么

2.1 Reward：回答好不好

2.2 Logprob：policy 更新的入口

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

2.4 Reference KL penalty：别离 SFT reference 太远

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

3. 源码阅读顺序图

4. MiniMind 源码走读

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

源码证据：LMForRewardModel

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

源码证据 A：PPO reward

源码证据 B：GRPO reward

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

源码证据 A：通用 token logprob helper

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

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

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

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

源码证据 B：GRPO current/ref logprob

4.5 reward 怎么变成 advantage

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

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

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

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

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

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

源码证据：DPO logratio

5. 本节必须会写 / 暂时不要求

5.1 必须会写

5.2 暂时不要求

6. 手写模块复盘

7. 实验验证

8. 阶段组装

9. 本节检查

10. 下一课

26 KiB

Raw Blame History

源码证据：`LMForRewardModel`