""" ============================================================================= 第 9 讲 · 01：agentic 后训练的两个问题 —— 奖励怎样设计？RL 从哪里起步？ ============================================================================= 对应讲义：lectures/lecture09_post_training_evolution.md 第 5、6 节 ----------------------------------------------------------------------------- 【设定】 ----------------------------------------------------------------------------- 第 1 讲的算术题，10B 模型（它的算术能力固定不变）。模型每做一步之前，要决定： 自己算，还是调用计算器？ · 计算器一定算对（但有 5% 的概率调用失败，失败时只能自己算）； · 每调用一次有成本 λ（真实的 agent 里：延迟、API 费用、上下文长度）。 策略：对每种"情境"（运算 × 难度，和第 6 讲相同），有一个参数 θ，调用计算器的概率 = σ(θ)。 奖励 = [最终答案正确] − λ × 调用次数。用第 6 讲的组内相对优势训练 θ（每题采样 8 次）。 三个演示： 演示 1 成本 λ 怎样塑造学到的工具使用策略：λ = 0 / 0.02 / 0.2（优势 = 奖励 − 组内其他样本的平均，即 RLOO）。 演示 2 同样是 λ = 0.02，换成 GRPO 的优势（再除以组内标准差）会怎样？ 演示 3 冷启动（cold start）：先用 50 条"专家轨迹"（只在乘两位数时调用计算器）做 SFT，再做 RL； 和直接从"几乎从不调用工具"的初始策略开始 RL 比较。训练步数不同时，谁更好？ 【对应的真实问题】训练 agent 时，奖励不只是"做对没有"，还要考虑成本、安全、格式； 而 RL 很难学会初始策略几乎从不尝试的行为（第 4 讲 RLEF 的"探索不足"），所以通常先用示范数据冷启动。 运行时间：约 3 秒（只用标准库）。 """ import math import os import random import sys _HERE = os.path.dirname(os.path.abspath(__file__)) sys.path.insert(0, os.path.join(_HERE, "..", "lecture06")) from trainable_llm import (HARD_CONTEXTS, PROCEDURES, TrainableLLM, apply_op, context_key, # noqa: E402 make_problem_set, mean, pad) from toy_llm import CARELESS_OFFSETS, apply_bug, sample_index # noqa: E402 TOOL_FAIL = 0.05 THETA0 = math.log(0.05 / 0.95) # 初始策略：只有 5% 的概率调用计算器 GROUP, BATCH, STEPS, LR = 8, 16, 60, 20.0 def sigmoid(z): return 1 / (1 + math.exp(-z)) if z >= 0 else math.exp(z) / (1 + math.exp(z)) def own_step(model, step, x, i, rng): """模型自己算一步（按第 1 讲的三种做法采样，temperature = 1）。""" probs = model.step_probs(step.op, x, step.operand, i, 1.0) proc = PROCEDURES[sample_index(probs, rng)] right = apply_op(step.op, x, step.operand) if proc == "correct": return right if proc == "bug": return apply_bug(step.op, x, step.operand) return right + CARELESS_OFFSETS[rng.randrange(len(CARELESS_OFFSETS))] def rollout(model, theta, problem, rng): """ 跑一个回合：每一步按 σ(θ[情境]) 决定是否调用计算器。 【返回】(是否做对, 调用次数, 决策记录 [(情境, 是否调用)]) 【思路】决策记录只包含**模型自己做的决定**（调不调用）。计算器返回的结果不是模型的动作 —— 训练时不应该对它求梯度。真实的 agentic RL 里，工具返回的内容（网页、执行输出）会被从损失里**屏蔽**掉 （loss masking），道理相同。 """ x, calls, decisions = problem.start, 0, [] for i, step in enumerate(problem.steps): key = context_key(step.op, x, step.operand) call = rng.random() < sigmoid(theta.get(key, THETA0)) decisions.append((key, call)) if call: calls += 1 if rng.random() >= TOOL_FAIL: x = apply_op(step.op, x, step.operand) continue x = own_step(model, step, x, i, rng) return x == problem.answer(), calls, decisions def train_policy(model, theta, train_set, lam, rng, steps=STEPS, adv="rloo"): """ 训练工具使用策略 θ（原地修改），返回训练后的 θ。 【思路】伯努利决策的对数概率梯度：∂ log π / ∂θ = [调用] − σ(θ)。奖励 = [做对] − λ × 调用次数。 每道题采样 8 次，两种优势： adv="rloo"：A = 奖励 − 组内其他 7 个样本的平均奖励（第 6 讲的 RLOO） adv="grpo"：A = (奖励 − 组内均值) / (组内标准差 + 1e−6)（第 6 讲的 GRPO） 【边界】GRPO 的陷阱：一组 8 个全错（或全对）时，组内奖励的差别**只剩下调用成本**（−λ × 调用次数）， 标准差很小，除以它之后，这点微小的差别被放大成和"做对 / 做错"一样大的信号 —— 模型被强烈地推向"少调用"。 """ for _ in range(steps): grads = {} for p in rng.sample(train_set, BATCH): outs = [rollout(model, theta, p, rng) for _ in range(GROUP)] rewards = [(1.0 if ok else 0.0) - lam * calls for ok, calls, _ in outs] mu = mean(rewards) sd = math.sqrt(mean((r - mu) ** 2 for r in rewards)) for (ok, calls, decisions), r in zip(outs, rewards): a = (r - mu) / (sd + 1e-6) if adv == "grpo" else (r - mu) * GROUP / (GROUP - 1) for key, call in decisions: p_call = sigmoid(theta.get(key, THETA0)) grads[key] = grads.get(key, 0.0) + a * ((1.0 if call else 0.0) - p_call) / (BATCH * GROUP) for key, g in grads.items(): theta[key] = theta.get(key, THETA0) + LR * g return theta def evaluate(model, theta, problems, rng, repeats=3): """返回 (准确率, 平均每题调用次数)。""" outs = [rollout(model, theta, p, rng) for p in problems for _ in range(repeats)] return mean(1.0 if ok else 0.0 for ok, _, _ in outs), mean(calls for _, calls, _ in outs) def report_row(name, model, theta, test, rng): acc, calls = evaluate(model, theta, test, rng) probs = "".join(pad(f"{sigmoid(theta.get(key, THETA0)):.2f}", 14, "right") for _, key in HARD_CONTEXTS) return pad(name, 24) + pad(f"{acc:.3f}", 9, "right") + pad(f"{calls:.2f}", 10, "right") + probs SHORT_NAMES = ("加法无进位", "乘一位数", "乘两位数", "乘两位·大数") # HARD_CONTEXTS 的简称 def header(): return (pad("", 24) + pad("准确率", 9, "right") + pad("调用/题", 10, "right") + "".join(pad(name, 14, "right") for name in SHORT_NAMES)) def main(): model = TrainableLLM(1e10) train_set = make_problem_set(random.Random(210), n_problems=300, min_steps=3, max_steps=8) test_set = make_problem_set(random.Random(211), n_problems=200, min_steps=3, max_steps=8) print("=" * 78) print("演示 1：工具调用的成本 λ 怎样塑造学到的策略（RLOO 式的优势，训练 60 步）") print("=" * 78) print("后几列：几个情境下调用计算器的概率（" + "、".join(name for name, _ in HARD_CONTEXTS) + "）") print(header()) print(report_row("训练前（几乎不调用）", model, {}, test_set, random.Random(212))) for lam in (0.0, 0.02, 0.2): theta = train_policy(model, {}, train_set, lam, random.Random(213)) print(report_row(f"RL，λ = {lam}", model, theta, test_set, random.Random(214))) print("\n观察：") print(" · λ = 0（调用免费）：RL 学会了在乘法上几乎总是调用计算器，准确率从 0.39 升到 0.88。") print(" 无进位的加法本来就很少出错，调不调用差别很小，梯度也小，所以那里的调用概率没怎么变。") print(" · λ = 0.02：RL 学会了\"按需调用\"：乘法上仍然几乎总是调用，其余情境很少调用 ——") print(" 调用次数少了三分之一，准确率仍然远高于训练前。一个小小的成本项，就塑造出了合理的工具使用习惯。") print(" · λ = 0.2：调用太贵了（一次调用的成本抵得上不少正确率），RL 学会了几乎不调用，准确率甚至略低于训练前。") print(" · 结论：**奖励写什么，模型就学什么**。agent 的奖励不只是\"做对没有\"，还包括成本、延迟、安全 ——") print(" 怎样在它们之间权衡，是后训练的设计者要做的价值判断。\n") print("=" * 78) print("演示 2：同样是 λ = 0.02，换成 GRPO 的优势（除以组内标准差）") print("=" * 78) rng = random.Random(219) groups = [[rollout(model, {}, p, rng)[0] for _ in range(GROUP)] for p in train_set[:200]] same = mean(1.0 if (all(g) or not any(g)) else 0.0 for g in groups) print(f"训练开始时，8 个样本全对或全错的组占 {same:.1%}。在这些组里，奖励的差别只剩下调用成本。\n") print(header()) for name, lam, adv in (("RLOO，λ = 0.02", 0.02, "rloo"), ("GRPO，λ = 0.02", 0.02, "grpo"), ("GRPO，λ = 0（对照）", 0.0, "grpo")): theta = train_policy(model, {}, train_set, lam, random.Random(213), adv=adv) print(report_row(name, model, theta, test_set, random.Random(214))) print("\n观察：") print(" · GRPO 在 λ = 0.02 时学会了**完全不调用**计算器，准确率回到训练前的水平 —— 和 RLOO 的结果截然相反。") print(" · 原因：在全对或全错的组里，组内奖励的差别只是 −0.02 × 调用次数，标准差很小；除以它之后，") print(" \"少调用一次\"得到了和\"做对一道题\"一样大的优势。这些组的信号压倒了其余的组。") print(" · λ = 0 时（没有成本），全对或全错的组优势全为 0，GRPO 没有这个问题。") print(" · 这正是第 6 讲提到的 Dr. GRPO 等工作对标准差归一化的批评：它会放大\"几乎没有差别\"的组里的噪声。") print(" 在 agent 的训练里，奖励往往由多项组成（正确性 + 成本 + 格式 + 安全），这个陷阱更容易出现。\n") print("=" * 78) print("演示 3：冷启动 —— 先用 50 条专家轨迹做 SFT，再做 RL（RLOO，λ = 0.02）") print("=" * 78) # 专家轨迹：只在乘两位数时调用计算器。SFT = 对专家的每个决策做最大似然。 expert_rng = random.Random(215) theta_sft = {} for _ in range(30): # 30 个 epoch 的梯度上升 grads = {} for p in train_set[:50]: x = p.start for i, step in enumerate(p.steps): key = context_key(step.op, x, step.operand) expert_call = step.op == "mul" and step.operand >= 10 p_call = sigmoid(theta_sft.get(key, THETA0)) grads[key] = grads.get(key, 0.0) + ((1.0 if expert_call else 0.0) - p_call) x = apply_op(step.op, x, step.operand) if expert_call else own_step(model, step, x, i, expert_rng) for key, g in grads.items(): theta_sft[key] = theta_sft.get(key, THETA0) + 0.05 * g print("后几列：几个情境下调用计算器的概率") print(pad("训练步数", 10) + header()[10:]) print(report_row("只做 SFT（冷启动）", model, theta_sft, test_set, random.Random(216))) for steps in (10, 30, 60): t_rl = train_policy(model, {}, train_set, 0.02, random.Random(217), steps=steps) t_both = train_policy(model, dict(theta_sft), train_set, 0.02, random.Random(217), steps=steps) print(report_row(f"只做 RL，{steps} 步", model, t_rl, test_set, random.Random(218))) print(report_row(f"SFT + RL，{steps} 步", model, t_both, test_set, random.Random(218))) print("\n观察：") print(" · 只做 RL 时，一开始模型只有 5% 的概率调用计算器，\"调用\"这个行为很少出现，RL 的信号很弱，") print(" 前 10 步几乎没有进步 —— 和第 4 讲 RLEF 的\"探索不足\"是同一个问题。") print(" · 50 条专家轨迹的 SFT 就让模型学会了\"乘两位数时调用计算器\"，起点高得多；在此基础上做 RL，") print(" 10 步就超过了只做 RL 30 步的水平。这是现代后训练流水线先 SFT（冷启动）再 RL 的主要理由。") print(" · 但训练久了，只做 RL 反而更好（60 步时 0.765 vs 0.647）：SFT 让模型在**专家没调用过的情境**上") print(" 变得非常确定地\"不调用\"（概率 0.01 左右），这些行为几乎再也不会被尝试，RL 也就学不到它们的价值；") print(" 只做 RL 的策略保留了更多的探索，在更多情境上学会了调用。") print(" · 冷启动的数据会**固化**示范者的习惯：它加快了早期的学习，也限制了之后能探索到的范围。") print(" DeepSeek-R1 用几千条长 CoT 数据冷启动，主要是为了可读性和训练的稳定；数据选什么、选多少，") print(" 决定了 RL 从哪里出发、还能走多远。\n") if __name__ == "__main__": main()