""" ============================================================================= 第 5 讲 · 01：树搜索 —— LATS（MCTS + 价值估计 + 反思）与 AB-MCTS（往宽走还是往深走） ============================================================================= 对应讲义：lectures/lecture05_reasoning_planning.md 第 2、4 节（Zhou et al. 2023；Inoue et al. 2025） ----------------------------------------------------------------------------- 【共同的设定】 ----------------------------------------------------------------------------- 第 1 讲的算术题，10B 生成器，temperature = 0.7。环境能告诉我们一个完整的解答"通过测试没有"（像代码的 unit test）。 预算按**生成的步数**计算（相当于 token 数）：从头写一份 k 步的解答花 k，从第 i 步续写花 k − i。 预算用"完整解答的份数"表示：预算 B = B × k 步。在预算内做对就算成功。 ----------------------------------------------------------------------------- 【演示 1：只有"通过 / 不通过"的反馈时】 ----------------------------------------------------------------------------- 重复采样（repeated sampling）：一份一份地从头写，直到通过或预算用完。 Reflexion 式的顺序修改：不通过就让模型自己复查、修改（第 2 讲的 revise）。 LATS：把部分解答组织成一棵树，用 MCTS 搜索： 选择（selection）：从根往下，按 PUCT 选子节点，直到一个叶子 —— 平均价值高的、先验高的、访问少的优先； 扩展（expansion）：在叶子上让模型写几次"下一步"，生成几个子节点； 评估（evaluation）：用一个模型给每个新的一步打分（LM 的价值估计 / 先验）； 模拟（simulation）：从先验最高的子节点续写到底，交给环境测试； 反向传播（backpropagation）：把结果沿路径往上更新； 反思（reflection）：失败的完整解答记进"记忆"，之后扩展时避开已知失败的那一步。 关键的省钱之处：**共享前缀**。树上已经写好的前缀不用重写，新的尝试只从分叉的地方开始。 ----------------------------------------------------------------------------- 【演示 2：有"部分分数"的反馈时（AB-MCTS 的设定）】 ----------------------------------------------------------------------------- 反馈是一个分数：从第一步开始连续做对了几步 / 总步数（像"通过了多少个测试用例"）。于是知道第一个错误在哪里。 往宽走（wider / GEN）：从头采样一份新的解答。 往深走（deeper / REFINE）：拿当前最好的解答，在第一个错误的那一步排除错误值重算，往后续写（第 4 讲的 repair）。 只往宽、只往深、两者交替，以及 AB-MCTS 式的**自适应**：用 Thompson 采样决定每一次往宽还是往深 —— 为两个动作各维护一个 Beta 分布，记录"这个动作有没有带来更好的分数"，每次从两个分布里各抽一个数，选大的那个。 运行时间：约 1 秒（只用标准库）。 """ import math import os import random import sys _HERE = os.path.dirname(os.path.abspath(__file__)) sys.path.insert(0, os.path.join(_HERE, "..", "lecture03")) from verifier_tools import (CARELESS_OFFSETS, PROCEDURES, ToyLLM, apply_bug, apply_op, extend_trace, # noqa: E402 make_problem_set, mean, pad, step_prob) from tts_tools import revise, sample_index # noqa: E402 （verifier_tools 已经把第 2 讲的目录加进了 sys.path） TEMPERATURE = 0.7 BUDGETS = (2, 4, 8, 16) # ============================================================================= # 共用的小工具 # ============================================================================= def sample_step_excluding(model, step, x, i, rng, avoid): """ 让模型写第 i 步，但避开 avoid 里的值（"我记得这个值试过，失败了"）。 【思路】三种做法里，一定会写出 avoid 里某个值的做法概率置 0，剩下的重新归一化；careless 的偏移也避开。 LATS 的反思（reflection）和 AB-MCTS 的往深走（refine）都用它：**记住失败，不再重复**。 【参数】avoid：要避开的值的集合（可以为空） 【返回】(值, 做法) 【边界】如果三种做法全被排除（极少见），退回不排除的普通采样。 """ probs = model.step_probs(step.op, x, step.operand, i, TEMPERATURE) right, bug = apply_op(step.op, x, step.operand), apply_bug(step.op, x, step.operand) allowed = [0.0 if (proc == "correct" and right in avoid) or (proc == "bug" and bug in avoid) else p for proc, p in zip(PROCEDURES, probs)] if sum(allowed) == 0: allowed = probs total = sum(allowed) proc = PROCEDURES[sample_index([a / total for a in allowed], rng)] if proc == "correct": return right, proc if proc == "bug": return bug, proc offsets = [d for d in CARELESS_OFFSETS if right + d not in avoid] or list(CARELESS_OFFSETS) return right + offsets[rng.randrange(len(offsets))], proc def beta_sample(a, b, rng): """ 从 Beta(a, b) 抽一个数（a、b 是正整数）。 【思路】一个漂亮的事实：取 a + b − 1 个 (0, 1) 上的均匀随机数，从小到大排好，第 a 个就服从 Beta(a, b)。 为什么不用 random.betavariate？它内部的算法在不同 Python 版本里不保证相同；自己写可以保证输出逐字相同。 【参数】a、b：正整数（这里是"成功次数 + 1"和"失败次数 + 1"） """ us = sorted(rng.random() for _ in range(a + b - 1)) return us[a - 1] def leading_correct(problem, sample): """部分分数：从第一步开始连续做对了几步（"通过了多少个测试用例"）。""" n = 0 for v, truth in zip(sample.values, problem.correct_values()): if v != truth: break n += 1 return n # ============================================================================= # 演示 1 的三种方法 # ============================================================================= def repeated_sampling(gen, problem, rng, budget): """一份一份地从头写，直到通过或预算用完。返回是否成功。""" spent = 0 while spent + problem.n_steps <= budget: spent += problem.n_steps if gen.generate(problem, rng, TEMPERATURE).correct: return True return False def reflexion(gen, problem, rng, budget): """ Reflexion 式的顺序修改：不通过时，模型自己复查（低温重算每一步），在第一处不一致的地方改写。 【思路】反馈只有"不通过"，模型必须自己找错在哪里 —— 它对自己的误解是"盲"的（第 2 讲 revise 的性质）。 每次修改的成本记为一份完整解答（复查 + 续写）。 """ spent = problem.n_steps s = gen.generate(problem, rng, TEMPERATURE) while not s.correct and spent + problem.n_steps <= budget: spent += problem.n_steps s = revise(gen, problem, s, rng, TEMPERATURE) return s.correct class Node: """ 搜索树的一个节点 = 一个部分解答（前缀）。 【字段】values：前缀里各步的值；parent：父节点；children：子节点列表； N：访问次数；W：价值之和（平均价值 = W / N）；prior：评估模型给"这一步"的分数； dead：完整的解答且测试失败（不会再被选中）；failed_next：在这个节点上写过、后来被证明失败的"下一步"的值（反思的记忆） """ def __init__(self, values, parent, prior): self.values, self.parent, self.prior = values, parent, prior self.children, self.N, self.W = [], 0, 0.0 self.dead, self.failed_next = False, set() def lats(gen, evaluator, problem, rng, budget, n_expand=3, c_puct=1.0): """ LATS 的简化版：在部分解答的树上做 MCTS。返回是否在预算内做对。 【思路】每一轮： 1. 选择：从根往下，在子节点里按 PUCT 选一个，直到走到一个还没有子节点的叶子： PUCT = 平均价值 + c × 先验 × √(父节点访问次数 + 1) / (1 + 子节点访问次数) 失败过的子节点平均价值是 0，于是"先验高、还没试过"的兄弟节点会被优先尝试。 2. 扩展：在叶子上让生成器写 n_expand 次"下一步"（避开反思记忆里失败过的值），每次花 1 步的预算； 不同的值成为子节点，先验 = 评估模型给这一步的分数（step_prob）。 如果新的子节点已经是完整的解答，直接交给测试：通过就成功，失败就标记为 dead。 3. 模拟：从先验最高的、还活着的子节点续写到底（花剩余步数的预算），交给测试。通过就成功。 4. 反思：失败的完整解答记进父节点的 failed_next，以后扩展时不再写这个值； 一个节点的所有子节点都失败了，它自己也标记为失败，并记进它父节点的记忆。 5. 反向传播：沿路径往上，访问次数 N += 1（失败的模拟价值为 0）。 【参数】gen：生成器；evaluator：给步骤打先验分数的模型（自己 = self-evaluation，或更强的模型）； n_expand：每次扩展写几个候选（1 = 不分叉）；c_puct：探索的强度 【边界】 · 共享前缀：从深度 d 的节点续写，只需要付 k − d 步的成本。但每次扩展也要多付 n_expand 步， 短题上这部分开销不可忽略。 · 树会把预算集中在"评估者相信"的分支上。评估者有盲点时，树会在错误的分支上反复打转； 预算很大时，这种"集中"反而不如独立采样的多样性（见演示 1 的结果）。 """ k = problem.n_steps root = Node([], None, 1.0) spent = 0 while spent < budget: node = root while node.children: alive = [ch for ch in node.children if not ch.dead] if not alive: # 所有子节点都失败了：这个节点也失败 node.dead = True if node.parent is None: return False node.parent.failed_next.add(node.values[-1]) node = root continue parent = node node = max(alive, key=lambda ch: ch.W / max(1, ch.N) + c_puct * ch.prior * math.sqrt(parent.N + 1) / (1 + ch.N)) i = len(node.values) x = node.values[-1] if node.values else problem.start for _ in range(n_expand): v, _proc = sample_step_excluding(gen, problem.steps[i], x, i, rng, node.failed_next) spent += 1 if any(ch.values[-1] == v for ch in node.children): continue # 写出了已有的值：不新增子节点 child = Node(node.values + [v], node, step_prob(evaluator, problem.steps[i], x, v, i)) if i == k - 1: # 已经是完整的解答：直接测试 if v == problem.answer(): return True child.dead = True node.failed_next.add(v) node.children.append(child) alive = [ch for ch in node.children if not ch.dead] if not alive: continue child = max(alive, key=lambda ch: ch.prior) rollout = extend_trace(gen, problem, child.values, ["?"] * len(child.values), rng, TEMPERATURE) spent += k - len(child.values) if rollout.correct: return True n = child while n is not None: n.N += 1 n = n.parent return False def demo_lats(problems, gen, big, rng): print("=" * 78) print(f"演示 1：只有\"通过 / 不通过\"的反馈 —— 预算内做对的比例（{gen.name} 生成器，{len(problems)} 道 5~8 步题 × 2 次）") print("=" * 78) methods = (("重复采样", lambda p, b: repeated_sampling(gen, p, rng, b)), ("Reflexion 式修改", lambda p, b: reflexion(gen, p, rng, b)), (f"LATS（先验 = {gen.name} 自评）", lambda p, b: lats(gen, gen, p, rng, b)), (f"LATS（先验 = {big.name}）", lambda p, b: lats(gen, big, p, rng, b))) print(pad("预算（份）", 12) + "".join(pad(name, 24, "right") for name, _ in methods)) for b in BUDGETS: row = pad(str(b), 12) for _, method in methods: row += pad(f"{mean(method(p, b * p.n_steps) for p in problems for _ in range(2)):.3f}", 24, "right") print(row) print("\n观察：") print(" · 预算小到中等时（2~8 份），LATS 好于重复采样：**共享前缀**让每次新的尝试只从分叉处开始，") print(" 同样的预算能做更多次尝试；反思又让它不再重复已知失败的那一步。") print(" · 预算大时（16 份），重复采样追了上来并略微反超：树把预算集中在少数分支附近，多样性不如独立采样。") print(" 这和第 2 讲 Snell et al. 对 beam search 的发现一致：搜索在小预算下占优，大预算下 best-of-N 追上来。") print(" · 评估模型用 10B 自评还是 100B，在这里几乎没有差别：先验只用来决定先试哪个子节点，") print(" 而每次扩展的 3 个候选里多半已经包含了正确的一步，真正起作用的是\"失败了就换一个兄弟节点\"。") print(" · Reflexion 式的修改在小预算下很强（低温复查能修好手滑，第 2 讲的结论），但大预算下停滞：") print(" 它修不了模型自己的误解，又不像采样那样能\"碰\"出别的做法。") print(" · LATS 论文里的反思是一段文字（\"上次失败是因为……\"），作为提示词的一部分影响之后的生成；") print(" 这里用\"记住失败的那一步、不再写它\"来近似。\n") # ============================================================================= # 演示 2：AB-MCTS —— 往宽走还是往深走 # ============================================================================= def refine(gen, problem, sample, rng): """往深走：在第一个错误的那一步排除错误值重算，往后续写。返回 (新解答, 成本)。""" i = leading_correct(problem, sample) x = problem.start if i == 0 else sample.values[i - 1] v, proc = sample_step_excluding(gen, problem.steps[i], x, i, rng, {sample.values[i]}) new = extend_trace(gen, problem, list(sample.values[:i]) + [v], list(sample.procs[:i]) + [proc], rng, TEMPERATURE) return new, problem.n_steps - i def search_with_policy(gen, problem, rng, budget, policy): """ 在预算内交替地"往宽走 / 往深走"，返回是否成功。 【参数】policy："wider"（只采样）/ "deeper"（只修改当前最好的）/ "alternate"（交替）/ "adaptive"（Thompson 采样） 【思路】adaptive：为 GEN、REFINE 两个动作各维护 Beta(成功 + 1, 失败 + 1)，"成功"= 这个动作带来了比当前最好 还高的部分分数。每次从两个 Beta 分布各抽一个数，做数值大的那个动作 —— 越常带来进步的动作越常被选中， 但另一个动作也保留着被尝试的机会（探索）。这是 AB-MCTS 在每个节点上做的决定的简化版（只在根部决策）。 """ k = problem.n_steps best = gen.generate(problem, rng, TEMPERATURE) spent, best_score = k, leading_correct(problem, best) stats = {"GEN": [1, 1], "REFINE": [1, 1]} turn = 0 while not best.correct: if policy == "wider": action = "GEN" elif policy == "deeper": action = "REFINE" elif policy == "alternate": action = "GEN" if turn % 2 == 0 else "REFINE" else: draws = {a: beta_sample(s, f, rng) for a, (s, f) in stats.items()} action = max(draws, key=lambda a: draws[a]) cost = k if action == "GEN" else k - leading_correct(problem, best) if spent + cost > budget: break if action == "GEN": cand = gen.generate(problem, rng, TEMPERATURE) else: cand, cost = refine(gen, problem, best, rng) spent += cost score = leading_correct(problem, cand) improved = score > best_score or cand.correct stats[action][0 if improved else 1] += 1 if improved: best, best_score = cand, score turn += 1 return best.correct def demo_wider_deeper(problems_by_len, gen, rng): print("=" * 78) print("演示 2：有部分分数的反馈 —— 往宽走、往深走、交替、自适应（预算 = 4 份完整解答）") print("=" * 78) policies = (("只往宽（重复采样）", "wider"), ("只往深（修改最好的）", "deeper"), ("交替", "alternate"), ("自适应（Thompson）", "adaptive")) print(pad("题目", 16) + "".join(pad(name, 22, "right") for name, _ in policies)) for label, probs in problems_by_len: row = pad(label, 16) for _, pol in policies: row += pad(f"{mean(search_with_policy(gen, p, rng, 4 * p.n_steps, pol) for p in probs for _ in range(2)):.3f}", 22, "right") print(row) label, long_probs = problems_by_len[-1] print(f"\n{label}的题，不同预算下：") print(pad("预算（份）", 16) + "".join(pad(name, 22, "right") for name, _ in policies)) for budget in (4, 8, 16): row = pad(str(budget), 16) for _, pol in policies: row += pad(f"{mean(search_with_policy(gen, p, rng, budget * p.n_steps, pol) for p in long_probs):.3f}", 22, "right") print(row) print("\n观察：") print(" · 有了\"第一个错误在哪里\"的信息，往深走（修改）非常有效：它保留已经做对的前缀，") print(" 并且排除了已知错误的值，模型就跳出了自己的误解（第 4 讲 RLEF 的机制）。") print(" · 只往宽在短题上不差，题目越长越吃亏：每次从头写，前面做对的部分全部浪费。") print(" · 在这个 toy world 里，往深走在每种题长、每种预算下都是最好的。自适应策略不知道这一点，") print(" 要从\"哪个动作带来了进步\"里慢慢学：预算小时只能做几次决策，试错的代价占比很大，它落在两者之间；") print(" 预算增大，它逐渐逼近往深走。") print(" · 为什么这里往深走总是赢？因为反馈**精确地**指出了第一个错误，而且修好这一步就能继续往前。") print(" 真实任务里常有另一种情况：整个思路就是错的（例如选错了算法），在错误的思路上修修补补没用，") print(" 必须重新采样一个新的思路 —— 这时往宽走更好。哪种情况更多，事先往往不知道。") print(" AB-MCTS 的价值正在于此：它在每个节点上自适应地选择，论文报告它在代码和工程类任务上") print(" 稳定地胜过重复采样和标准的 MCTS。这个 toy world 没有\"思路错了\"的情况，所以体现不出它的这一面。\n") def main(): gen, big = ToyLLM(1e10), ToyLLM(1e11) problems = make_problem_set(random.Random(120), n_problems=150, min_steps=5, max_steps=8) demo_lats(problems, gen, big, random.Random(121)) by_len = [("3~4 步", make_problem_set(random.Random(122), 150, 3, 4)), ("5~6 步", make_problem_set(random.Random(123), 150, 5, 6)), ("7~9 步", make_problem_set(random.Random(124), 150, 7, 9))] demo_wider_deeper(by_len, gen, random.Random(125)) if __name__ == "__main__": main()