""" ============================================================================= 第 8 讲 · 代码 3：REINFORCE + KL 惩罚 —— 策略梯度、基线，以及它收敛到哪里 ============================================================================= 对应讲义：lectures/lecture08_rlhf.md 第 3、4 节 ----------------------------------------------------------------------------- 【这个文件要解决什么问题？】 ----------------------------------------------------------------------------- RLHF 的第三步：用强化学习优化 J(π) = E_{y~π}[ r(y) ] − β · KL(π ‖ π_ref) r 是代码 02 的奖励模型（线性 + 众包）。这个文件回答三个问题： · 策略梯度是怎样从采样估计出来的？为什么要减"基线"？ · REINFORCE 真的能优化 J 吗？—— 玩具世界里 J 的最优解 π* ∝ π_ref · exp(r/β) 可以精确算出来， 所以能直接检查 RL 是否到了那里。 · KL 系数 β 控制什么？ 演示 1：梯度估计的质量：不减基线 / 减平均 / 留一法（RLOO）/ 组内标准化（GRPO）， 和**精确的梯度**（穷举全部回答算出来的）比较。 演示 2：训练曲线：三种做法谁收敛得快、稳？离解析最优解还有多远？ 演示 3：不同的 β：RL 找到的策略 vs 解析最优解。 ----------------------------------------------------------------------------- 【运行方式】 ----------------------------------------------------------------------------- python3 03_reinforce_kl.py # 约 25 秒（CPU） 需要 numpy 和 PyTorch。依赖本目录的 toy_lm.py。 """ import numpy as np import torch from toy_lm import (ALL_ENC, INDEX, advantages, all_logprobs, clone_policy, encode, exact_stats, kl, make_reference, make_reward_model, optimal_logprobs, pad, print_header, rl_train, sample, sequence_logprob) def flat_grad(policy): """把策略所有参数的梯度拼成一个长向量。""" return torch.cat([p.grad.reshape(-1) for p in policy.parameters()]).clone() def exact_gradient(policy, ref_lp_t, reward_t, beta): """ 目标 J(θ) = Σ_y π_θ(y) [r(y) − β(log π_θ(y) − log π_ref(y))] 的**精确**梯度（穷举全部 19531 个回答）。 【思路】J 是全部回答上的一个有限和，直接让 PyTorch 对它自动求导。真实的语言模型做不到（回答有无穷多个）， 只能采样估计 —— 这正是策略梯度方法要解决的问题。 """ policy.zero_grad() lp = sequence_logprob(policy, ALL_ENC) J = (lp.exp() * (reward_t - beta * (lp - ref_lp_t))).sum() J.backward() return flat_grad(policy), J.item() def sampled_gradient(policy, ref, reward_t, beta, batch, gen, method): """ 用 batch 个采样回答估计策略梯度（梯度**上升**的方向）。 【思路】策略梯度定理：∇J = E_{y~π}[ (R(y) − b) ∇log π(y) ]，R = r − β(log π − log π_ref)，b 是基线。 这里构造"代理损失" mean(A · log π(y))，对它求导就得到这个估计（A 当作常数，不求导）。 """ ys = sample(policy, batch, gen) enc = encode(ys) policy.zero_grad() lp = sequence_logprob(policy, enc) with torch.no_grad(): lp_ref = sequence_logprob(ref, enc) r = reward_t[torch.tensor([INDEX[y] for y in ys])] A = advantages(r - beta * (lp.detach() - lp_ref), method) (A * lp).mean().backward() return flat_grad(policy) # ============================================================================= # 演示 1 # ============================================================================= def demo_gradient_quality(ref, reward, beta=1.0, trials=200, batch=32): print_header("演示 1：梯度估计的质量 —— 为什么要减基线") reward_t = torch.tensor(reward, dtype=torch.float32) ref_lp_t = torch.tensor(all_logprobs(ref), dtype=torch.float32) policy = clone_policy(ref) g_true, J = exact_gradient(policy, ref_lp_t, reward_t, beta) print(f"在 SFT 策略处（β = {beta}），精确的目标值 J = {J:.3f}，精确梯度的长度 {g_true.norm():.3f}。") print(f"每种做法各估计 {trials} 次（每次 {batch} 个回答），和精确梯度比较：\n") print(f" {pad('做法', 26)} {pad('平均余弦相似度', 15, 'right')} {pad('相对均方误差', 13, 'right')} {pad('平均长度 / 真实长度', 20, 'right')}") names = {"none": "不减基线（原始 REINFORCE）", "mean": "减这一批的平均", "rloo": "留一法基线（RLOO）", "grpo": "组内标准化（GRPO）"} for method in ["none", "mean", "rloo", "grpo"]: gen = torch.Generator().manual_seed(11) gs = torch.stack([sampled_gradient(policy, ref, reward_t, beta, batch, gen, method) for _ in range(trials)]) cos = torch.nn.functional.cosine_similarity(gs, g_true.unsqueeze(0), dim=1).mean().item() rel_mse = ((gs - g_true) ** 2).sum(1).mean().item() / float(g_true.norm() ** 2) ratio = gs.mean(0).norm().item() / g_true.norm().item() print(f" {pad(names[method], 26)} {cos:>15.3f} {rel_mse:>13.2f} {ratio:>20.2f}") print(""" 观察： · 不减基线：每个采到的回答都按它的回报 R 被“往上推”（在 SFT 策略处 R 平均约 +1.4，大多是正的）， 推的力度只有细微差别。这些推力大部分互相抵消，有用的信号混在大量噪声里：均方误差约为梯度长度平方的 75%。 · 减去基线后（只告诉模型“这个回答比平常好还是差”），均方误差降到约一半，方向也更准，期望不变（E[∇log π] = 0）。 留一法基线（RLOO）不包含自己，是严格无偏的（最后一列 = 1.00）；“减这一批的平均”包含自己，梯度被缩小了 1/B。 · GRPO 再除以组内标准差：方向和 RLOO 一样准，但梯度的整体大小被改变了（最后一列 0.63），所以它的均方误差 和别的不能直接比。除以标准差相当于自动调整步长：奖励差异小的提示词，梯度被放大。 真实的 GRPO 是“同一个提示词采一组回答”，在组内标准化。 · 回报的平均值离 0 越远，不减基线的噪声越大（试试给奖励整体加 10）。奖励模型的分数本来就没有固定的零点 （第 1 讲：BT 只确定差值），所以不减基线的 REINFORCE 结果会随一个任意的常数变化 —— 减基线让它不受这个常数影响。 · PPO 用一个学出来的价值函数 V(状态) 做基线（“演员—评论家”），还能给每个词单独的优势； 代价是要多训练一个和策略差不多大的网络。""") # ============================================================================= # 演示 2 # ============================================================================= def demo_training(ref, reward, beta=1.0, steps=400): print_header("演示 2：训练曲线 —— 离解析最优解还有多远") ref_lp = all_logprobs(ref) opt_lp = optimal_logprobs(ref_lp, reward, beta) st = exact_stats(opt_lp, ref_lp, reward) J_opt = st["proxy"] - beta * st["kl"] J_ref = exact_stats(ref_lp, ref_lp, reward)["proxy"] print(f"β = {beta}。SFT 策略的 J = {J_ref:.3f}；解析最优解 π* ∝ π_ref·exp(r/β) 的 J = {J_opt:.3f}" f"（KL = {st['kl']:.3f}，真实质量 {st['gold']:+.3f}）。") print(f"每步 128 个回答，Adam 学习率 1e-3，共 {steps} 步。表中是精确的 J 和 KL(π_θ ‖ π*)：\n") checkpoints = [25, 50, 100, 200, 400] names = {"none": "不减基线", "rloo": "RLOO", "grpo": "GRPO"} rows = {} for method in ["none", "rloo", "grpo"]: rec = {} def cb(step, pol, rec=rec): lp = all_logprobs(pol) s = exact_stats(lp, ref_lp, reward) rec[step] = (s["proxy"] - beta * s["kl"], kl(lp, opt_lp)) rl_train(clone_policy(ref), ref, reward, beta, steps, method=method, seed=3, log_every=25, callback=cb) rows[method] = rec print(f" {pad('步数', 6)}" + "".join(pad(f"{names[m]}：J / 离最优", 24, "right") for m in rows)) for c in checkpoints: print(f" {c:>6d}" + "".join(pad(f"{rows[m][c][0]:.3f} / {rows[m][c][1]:.4f}", 24, "right") for m in rows)) print(""" 观察： · RLOO 和 GRPO 在一两百步内就几乎到了解析最优解（KL(π_θ ‖ π*) 接近 0）：REINFORCE 真的在优化 J， 而且找到的就是第 1 讲练习 4 推出来的那个分布。 · 不减基线的版本也在往上走，但一直在最优解附近抖动，到不了 —— 梯度噪声太大（演示 1）。 · 这个结论很重要：**RLHF 的目标有一个明确的答案 π* ∝ π_ref·exp(r/β)**。RL 算法（REINFORCE、PPO、GRPO） 只是“怎样到达那里”的不同方法。所以“RLHF 学到了什么”主要由奖励模型 r 和 β 决定，而不是由 RL 算法决定。 第 9 讲的 DPO 正是利用这一点，跳过 RL 直接求 π*。""") # ============================================================================= # 演示 3 # ============================================================================= def demo_beta(ref, reward, steps=500): print_header("演示 3：KL 系数 β —— RL 找到的策略 vs 解析最优解") ref_lp = all_logprobs(ref) s0 = exact_stats(ref_lp, ref_lp, reward) print(f"SFT 策略：奖励 {s0['proxy']:+.3f}，真实质量 {s0['gold']:+.3f}，平均长度 {s0['length']:.2f}。RLOO，{steps} 步。\n") print(f" {pad('β', 5)} {pad('', 6)} {pad('KL(π ‖ π_ref)', 14, 'right')} {pad('奖励', 7, 'right')} {pad('真实质量', 9, 'right')} " f"{pad('平均长度', 9, 'right')} {pad('华丽', 6, 'right')} {pad('要点', 6, 'right')} {pad('KL(π ‖ π*)', 11, 'right')}") for beta in [3.0, 1.0, 0.3]: opt_lp = optimal_logprobs(ref_lp, reward, beta) pol = rl_train(clone_policy(ref), ref, reward, beta, steps, method="rloo", seed=4) lp = all_logprobs(pol) for name, l in [("RL", lp), ("解析", opt_lp)]: s = exact_stats(l, ref_lp, reward) gap = f"{kl(lp, opt_lp):.4f}" if name == "RL" else "" print(f" {beta:>5.1f} {pad(name, 6)} {s['kl']:>14.3f} {s['proxy']:>+7.3f} {s['gold']:>+9.3f} " f"{s['length']:>9.2f} {s['fancy']:>6.2f} {s['points']:>6.2f} {gap:>11s}") gen = torch.Generator().manual_seed(8) pol = rl_train(clone_policy(ref), ref, reward, 0.3, steps, method="rloo", seed=4) from toy_lm import show, gold_reward print("\n β = 0.3 的 RL 策略生成的 5 个回答：") for s in sample(pol, 5, gen): print(f" {pad(show(s), 34)} 奖励 {reward[INDEX[s]]:+.2f} 真实质量 {gold_reward(s):+.2f}") print(""" 观察： · 每个 β 下，RL 训练出的策略都和解析最优解几乎一样（最后一列很小）。 · β 大（3.0）：策略只离开 SFT 一点点（KL 小），奖励和真实质量都提高了一些。 · β = 1.0：走得更远，真实质量提高得更多 —— 这时奖励模型还在“正确的方向”上。 · β 小（0.3）：策略离 SFT 很远，奖励继续上升，可**真实质量反而下降了**：回答越来越长、要点越来越多、 华丽也多了 —— 正是代码 02 里奖励模型“最喜欢”的方向。这就是**奖励过度优化**，代码 04 会系统地画出这条曲线。""") def main(): ref = make_reference() _, reward, _ = make_reward_model(ref, kind="linear", labeler="crowd") demo_gradient_quality(ref, reward) demo_training(ref, reward) demo_beta(ref, reward) if __name__ == "__main__": main()