""" ============================================================================= 第 16 讲 · 共享模块：一个"玩具推理模型"——可验证奖励的强化学习（RLVR） ============================================================================= 对应讲义：lectures/lecture16_frontiers_review.md ----------------------------------------------------------------------------- 【设定】 ----------------------------------------------------------------------------- 真正的推理模型（比如用 RL 训练做数学题的大语言模型）有几十亿参数，生成的回答是很长的文本。 这里把它简化为（与第 9 讲的"玩具语言模型"同样的思路）： - 有 P 道题（提示词 p = 0, 1, ..., P-1）； - 每道题，模型从 M 个候选回答中选一个： 回答 0 ：正确答案 回答 1 ："作弊"的回答（只在 02 文件中有意义：它能骗过一个有漏洞的验证器） 回答 2..M-1 ：各种错误答案 - 模型 = 每道题上 M 个回答的概率分布 π(y | p) = softmax(logits[p] + 共享的作弊倾向 h · [y 是作弊回答])； - **参考模型**（预训练 + 监督微调之后、RL 之前的模型）：不同的题难度不同—— 参考模型答对的概率从 1% 到 60% 不等；作弊回答的概率很小（0.5%）。 - **可验证的奖励**：答对得 1，否则得 0。不需要奖励模型（第 9 讲），直接检查答案—— 数学题比对最终答案、代码题运行单元测试。 为什么还要一个"共享的作弊倾向" h？ 每道题各自的 logits 只影响这一道题；而 h 是**所有题共享**的参数—— 就像真实的大模型中，"学会了钻验证器的空子"这种行为会泛化到所有的题上。 ----------------------------------------------------------------------------- 【本模块提供】 ----------------------------------------------------------------------------- make_task(P, M, seed) : 生成题目难度与参考模型 Policy : 策略（logits 表 + 共享的作弊倾向 h），可以抽样、计算概率 advantages(r, method, ...) : 各种优势估计：无基线、学到的基线、组内均值（RLOO）、GRPO train(...) : 通用的训练循环（每步：抽一批题，每题抽 G 个回答，计算奖励与优势，策略梯度更新） """ import unicodedata import numpy as np def pad(text, width): """在右边补空格，使字符串在终端里的显示宽度为 width（中文字符算 2 格），用来对齐表格。""" w = sum(2 if unicodedata.east_asian_width(ch) in "WF" else 1 for ch in text) return text + " " * max(0, width - w) # ============================================================================= # 第 1 部分：任务 # ============================================================================= DIFFICULTY_LEVELS = np.array([0.01, 0.05, 0.2, 0.6]) # 参考模型答对的概率：很难、难、中等、容易 LEVEL_NAMES = ["很难(1%)", "难(5%)", "中等(20%)", "容易(60%)"] def make_task(P=200, M=16, seed=0, hack_prob=0.005): """ 生成 P 道题。每道题随机分配一个难度等级，参考模型在这道题上： P(答对) = 该等级的概率，P(作弊回答) = hack_prob，其余概率在 M-2 个错误答案之间随机分配。 【返回】字典： "ref_logits" : 形状 (P, M)，参考模型的 logits（log 概率） "level" : 形状 (P,)，每道题的难度等级（0～3） 【边界】各概率之和必须小于 1：DIFFICULTY_LEVELS 的最大值 + hack_prob < 1。 """ rng = np.random.default_rng(seed) level = rng.integers(len(DIFFICULTY_LEVELS), size=P) probs = np.zeros((P, M)) probs[:, 0] = DIFFICULTY_LEVELS[level] probs[:, 1] = hack_prob rest = 1.0 - probs[:, 0] - probs[:, 1] w = rng.dirichlet(np.ones(M - 2), size=P) probs[:, 2:] = rest[:, None] * w return {"ref_logits": np.log(probs), "level": level, "P": P, "M": M} # ============================================================================= # 第 2 部分：策略 # ============================================================================= class Policy: """ π(y | p) = softmax(logits[p] + h · onehot(作弊回答)) 【参数】 logits : 形状 (P, M)，每道题自己的参数（初始化为参考模型） h : 标量，所有题共享的"作弊倾向"（初始为 0） """ def __init__(self, task): self.logits = task["ref_logits"].copy() self.h = 0.0 def probs(self, prompts=None): """返回 π(· | p)，形状 (题数, M)。prompts = None 表示所有题。""" lg = self.logits if prompts is None else self.logits[prompts] lg = lg.copy() lg[:, 1] += self.h lg -= lg.max(axis=1, keepdims=True) e = np.exp(lg) return e / e.sum(axis=1, keepdims=True) def sample(self, prompts, G, rng): """对每道题抽 G 个回答。返回形状 (题数, G) 的回答编号。""" pr = self.probs(prompts) cum = pr.cumsum(axis=1) u = rng.random((len(prompts), G, 1)) return (u > cum[:, None, :]).sum(axis=2) # 【Python 小知识】按累积概率抽样的向量化写法：随机数 u 超过了几个累积概率，就落在第几个回答上 # ============================================================================= # 第 3 部分：优势估计 # ============================================================================= def advantages(r, method, baseline=None, eps=1e-6): """ 根据一组回答的奖励 r（形状 (题数, G)）计算每个回答的优势 A（同样的形状）。 method 取值： "none" : A = r （REINFORCE，无基线） "learned" : A = r - V(p) （学到的基线：V 由调用者维护，传入 baseline，形状 (题数,)） "rloo" : A_i = r_i - mean_{j≠i} r_j （留一法组内基线，RLOO） "group" : A_i = r_i - mean(r) （组内均值基线，即不除以标准差的 GRPO，"Dr. GRPO"） "grpo" : A_i = (r_i - mean(r)) / (std(r) + eps) （GRPO，DeepSeekMath, 2024） 【直觉】GRPO 不训练价值网络（第 7 讲 PPO 需要一个和策略一样大的 critic），而是对同一道题抽 G 个回答， 用这组回答的平均奖励当基线：比同组的平均水平好，就加强；差，就减弱。 【边界】一组回答的奖励全部相同（全对或全错）时，组内均值基线和 GRPO 的优势都是 0——这组数据没有任何学习信号。 """ G = r.shape[1] if method == "none": return r.copy() if method == "learned": return r - baseline[:, None] if method == "rloo": return r - (r.sum(axis=1, keepdims=True) - r) / (G - 1) mean = r.mean(axis=1, keepdims=True) if method == "group": return r - mean if method == "grpo": return (r - mean) / (r.std(axis=1, keepdims=True) + eps) raise ValueError(method) # ============================================================================= # 第 4 部分：训练循环 # ============================================================================= def train(task, method, steps, batch, G, lr, seed, hack_reward=0.0, kl_coef=0.0, lr_h=None, eval_every=None, reward_offset=0.0): """ 用策略梯度训练策略。 【每一步】 1. 随机抽 batch 道题，每道题抽 G 个回答； 2. 奖励：回答 0（正确）得 1；回答 1（作弊）得 hack_reward（验证器有漏洞时为 1，否则为 0）；其余为 0； 3. 计算优势 A（见 advantages）； 4. 策略梯度：对每道题 p，logits[p] += lr · (1/G) Σ_i A_i · (onehot(y_i) - π_p) （这就是第 6 讲的得分函数 ∇ log π(y) 对 softmax logits 的形式）； 共享的作弊倾向 h 的梯度是所有题的梯度之和：Σ_p (1/G) Σ_i A_i · ([y_i 是作弊] - π_p(作弊))； 5. 如果 kl_coef > 0，再减去 kl_coef · ∇KL(π_p ‖ π_ref)，让策略不要离参考模型太远（第 9 讲）。 【参数】 method : 优势估计方法（见 advantages） steps : 训练步数 batch, G : 每步的题数、每题的回答数（每步共 batch × G 个样本） lr, lr_h : logits 和共享参数 h 的学习率（lr_h 默认等于 lr / batch，使 h 的更新幅度与单道题相当） hack_reward : 作弊回答得到的奖励（02 文件中用来模拟有漏洞的验证器）。 也可以是一个函数 step → 奖励，用来模拟"训练到一半修好了验证器" reward_offset : 给**所有**回答都加上的常数奖励（比如"格式正确就给分"），用来观察基线的作用 kl_coef : KL 惩罚系数 eval_every : 每隔多少步记录一次评估结果 【返回】字典： "history" : 列表 [(步数, 真实准确率, 作弊概率, 验证器给的平均奖励, 按难度的准确率), ...] "zero_frac" : 所有训练步中"组内奖励全部相同、没有学习信号"的组所占的比例 "policy" : 训练好的策略 【真实准确率】= 所有题上 π(正确答案) 的平均，即期望的 pass@1。这里直接用概率计算，不需要抽样。 """ rng = np.random.default_rng(seed) P, M = task["P"], task["M"] pol = Policy(task) ref_logp = task["ref_logits"] - np.log(np.exp(task["ref_logits"]).sum(axis=1, keepdims=True)) V = np.zeros(P) # "learned" 方法的基线：每道题的平均奖励估计 lr_h = lr / batch if lr_h is None else lr_h history, zero_groups, total_groups = [], 0, 0 eval_every = eval_every or max(1, steps // 10) def hack_r(step): return hack_reward(step) if callable(hack_reward) else hack_reward def evaluate(step): pr = pol.probs() acc = pr[:, 0].mean() by_level = [pr[task["level"] == k, 0].mean() for k in range(len(DIFFICULTY_LEVELS))] history.append((step, float(acc), float(pr[:, 1].mean()), float((pr[:, 0] + hack_r(step) * pr[:, 1]).mean()), by_level)) evaluate(0) for step in range(1, steps + 1): prompts = rng.choice(P, size=batch, replace=False) y = pol.sample(prompts, G, rng) # (batch, G) r = (y == 0).astype(float) + hack_r(step) * (y == 1) + reward_offset A = advantages(r, method, baseline=V[prompts]) same = np.all(r == r[:, :1], axis=1) zero_groups += int(same.sum()) total_groups += batch pr = pol.probs(prompts) # (batch, M) onehot = np.zeros((batch, G, M)) np.put_along_axis(onehot, y[:, :, None], 1.0, axis=2) grad = (A[:, :, None] * (onehot - pr[:, None, :])).mean(axis=1) # (batch, M) grad_h = grad[:, 1].sum() # h 只出现在作弊回答的 logit 上 if kl_coef > 0: logp = np.log(pr + 1e-300) diff = logp - ref_logp[prompts] kl = (pr * diff).sum(axis=1, keepdims=True) g_kl = pr * (diff - kl) # ∇_logits KL(π ‖ π_ref) grad -= kl_coef * g_kl grad_h -= kl_coef * g_kl[:, 1].sum() pol.logits[prompts] += lr * grad pol.h += lr_h * grad_h if method == "learned": V[prompts] += 0.2 * (r.mean(axis=1) - V[prompts]) if step % eval_every == 0: evaluate(step) return {"history": history, "zero_frac": zero_groups / max(total_groups, 1), "policy": pol} def pass_at_k(p, k): """一道题上答对概率为 p 时，独立抽 k 次至少答对一次的概率：1 - (1 - p)^k。""" return 1.0 - (1.0 - p) ** k if __name__ == "__main__": task = make_task() pol = Policy(task) pr = pol.probs() print("参考模型的平均准确率（pass@1）：", round(pr[:, 0].mean(), 3)) for k, name in enumerate(LEVEL_NAMES): print(f" {name}：{int((task['level'] == k).sum())} 道题")