"""Helper code for the Chapter 17 lab (complete: you don't need to change anything here). Two toy worlds, both small enough to understand completely. 1. ANSWER SIMULATOR (Tasks 1-3): what you get when you sample N answers to one problem and score them with an imperfect verifier. This is the same model as the chapter's "test-time scaling" demo. * Each sample is correct with probability p. Answers are integers: 0 the correct answer 1 the "popular" wrong answer (a common mistake) 2 .. 21 rare wrong answers (scattered mistakes) A wrong sample gives the popular wrong answer with probability `wrong_top_share`, otherwise one of the 20 rare ones uniformly. * The verifier looks at a sample and produces a latent score z: correct answer z ~ Normal(d, 1) ordinary wrong answer z ~ Normal(0, 1) "convincing" wrong answer z ~ Normal(d + 1, 1) A wrong answer is convincing with probability `fool_rate`. The verifier reports v = sigmoid(d * (z - d / 2)), a probability-like score in (0, 1). (For f = 0 and p = 1/2 this is exactly the posterior P(correct | z).) d is the verifier's *sharpness* (signal-detection d'); d = 0 is useless. 2. CHAIN TASK (Tasks 4-5): a toy multi-step "reasoning" problem. Problem: "start with 4; add 7, 3, 9, 5, 8, 2; keep only the last digit". A solution writes the running digit after every step: 1, 4, 3, 8, 6, 8. The final answer is the last digit written. The policy (a stand-in for an LLM) works from whatever it wrote last. At step t it writes (previous + op_t) mod 10 with probability 1 - err_t, and otherwise a uniformly random *wrong* digit (one of the other 9). Bigger numbers to add are harder: err_t = 0.04 + 0.03 * op_t (7% for +1, 31% for +9). Because it continues from its own (possibly wrong) digit, one mistake usually ruins the answer. But a later mistake can, by luck, land back on the right track (probability err/9 per step): "right answer for the wrong reasons". `exact_value(problem, prefix)` gives the true probability that the policy finishes with the right answer from a given prefix: the value V(prefix) that Monte Carlo step labels estimate. It is a two-state Markov chain (on track / off track), solved backwards. """ import math import random from typing import NamedTuple import numpy as np # ====================================================================== 1. answers N_RARE = 20 # rare wrong answers are 2 .. 21 def sample_answers(p, n, wrong_top_share, rng): """n sampled answers (int array) for one problem. 0 is the correct answer.""" u = rng.random(n) correct = u < p popular = rng.random(n) < wrong_top_share rare = 2 + rng.integers(0, N_RARE, n) return np.where(correct, 0, np.where(popular, 1, rare)) def verifier_scores(answers, d, fool_rate, rng, bonus=1.0): """Verifier scores v in (0, 1) for an array of answers (0 = correct).""" answers = np.asarray(answers) correct = answers == 0 convincing = (~correct) & (rng.random(answers.shape) < fool_rate) mean = np.where(correct, d, np.where(convincing, d + bonus, 0.0)) z = mean + rng.standard_normal(answers.shape) return 1.0 / (1.0 + np.exp(-d * (z - d / 2))) # ====================================================================== 2. chain task class Problem(NamedTuple): start: int ops: tuple # numbers to add, one per step errs: tuple # the policy's error rate at each step def step_error(op): """Policy error rate for adding `op` (harder for bigger numbers).""" return 0.04 + 0.03 * op def make_problem(start, ops): return Problem(int(start), tuple(int(o) for o in ops), tuple(step_error(o) for o in ops)) EXAMPLE = make_problem(4, (7, 3, 9, 5, 8, 2)) # the chapter's worked example class ChainTask: """The toy multi-step task. All randomness comes from the `rng` you pass in (a `random.Random`), so results are reproducible.""" def __init__(self, n_steps=6): self.n_steps = n_steps # -------------------------------------------------------------- problems def sample_problem(self, rng): return make_problem(rng.randrange(10), [rng.randint(1, 9) for _ in range(self.n_steps)]) @staticmethod def correct_steps(problem): """The correct running digits, e.g. [1, 4, 3, 8, 6, 8] for EXAMPLE.""" out, s = [], problem.start for op in problem.ops: s = (s + op) % 10 out.append(s) return out @staticmethod def answer(problem): return ChainTask.correct_steps(problem)[-1] # -------------------------------------------------------------- the policy @staticmethod def policy_step(problem, prefix, rng): """Sample the next digit given the steps written so far (a list of ints).""" t = len(prefix) prev = prefix[-1] if prefix else problem.start right = (prev + problem.ops[t]) % 10 if rng.random() < problem.errs[t]: wrong = rng.randrange(9) # one of the 9 other digits, uniformly return wrong if wrong < right else wrong + 1 return right def rollout(self, problem, prefix, rng): """Complete `prefix` with the policy. Returns the full list of steps (a new list; `prefix` is not modified).""" steps = list(prefix) while len(steps) < len(problem.ops): steps.append(self.policy_step(problem, steps, rng)) return steps def sample_solution(self, problem, rng): return self.rollout(problem, [], rng) # -------------------------------------------------------------- outcome verifier @staticmethod def is_correct(problem, steps): """Outcome check: is the final digit right? (Ignores how it got there.)""" return len(steps) == len(problem.ops) and steps[-1] == ChainTask.answer(problem) # -------------------------------------------------------------- ground truth @staticmethod def on_track(problem, prefix): if not prefix: return True return prefix[-1] == ChainTask.correct_steps(problem)[len(prefix) - 1] @staticmethod def exact_value(problem, prefix): """True P(final answer correct | prefix) under the policy. Two states: on track (current digit equals the correct running digit) or off track. Per remaining step with error rate e: on -> off with prob e off -> on with prob e / 9 (a lucky mistake) """ v_on, v_off = 1.0, 0.0 # value after the last step for e in reversed(problem.errs[len(prefix):]): v_on, v_off = (1 - e) * v_on + e * v_off, (e / 9) * v_on + (1 - e / 9) * v_off return v_on if ChainTask.on_track(problem, prefix) else v_off @staticmethod def first_error(problem, steps): """Index (1-based step number) of the first wrong step, or None.""" good = ChainTask.correct_steps(problem) for i, (s, g) in enumerate(zip(steps, good)): if s != g: return i + 1 return None class NoisyPRM: """A stand-in for a learned process reward model. Score of a prefix = sigmoid( logit(V) + sigma * noise ), where V is the true value (clipped to [0.05, 0.95]) and the noise is a fixed standard-normal number for each (problem, prefix): like a real PRM, it gives the same score every time you ask about the same prefix. sigma = 0 is a perfect PRM, sigma = 1 is a good one, and sigma = 3 (used in run.py) makes many mistakes. """ def __init__(self, sigma, seed=0): self.sigma = sigma self.seed = seed self.calls = 0 def __call__(self, problem, prefix): self.calls += 1 v = min(max(ChainTask.exact_value(problem, prefix), 0.05), 0.95) key = (self.seed, problem.start, problem.ops, tuple(prefix)) noise = random.Random(hash(key)).gauss(0.0, 1.0) # hash of int tuples is stable logit = math.log(v / (1 - v)) + self.sigma * noise return 1.0 / (1.0 + math.exp(-logit))