# Lab 17: process rewards and test-time compute

Goes with Chapter 17 (`/rl/learn/17-process-rewards/`). numpy, matplotlib and the
standard library only (no torch). Tests take about a second; `run.py` about 5 s.

You will:

1. compute how much extra samples can help (pass@N, and majority vote's exact accuracy);
2. write the three ways of picking one answer from N samples (majority vote,
   best-of-N, weighted vote) and watch an imperfect verifier start to *hurt*;
3. label every step of a solution by rolling out from it (Math-Shepherd's
   trick), and see that these labels are Monte Carlo value estimates whose
   differences are step-level advantages;
4. steer a step-level beam search with a process reward model (PRM), and
   compare it with best-of-N at the same compute.

## How to run

```bash
cd 17-process-rewards
python -m pytest -q                       # your code: fails until you fill in the TODOs
python -m pytest -q test_ttc.py -k pass   # just one task's tests
RL_SOLUTIONS=1 python -m pytest -q        # reference solutions: all 26 pass
python run.py                             # plots -> process_rewards.png
```

## The two toy worlds (in `chain.py`, complete)

**Answer simulator** (Tasks 1–3), the same model as the chapter's first demo.
Each sample is correct with probability p. Answer `0` is correct, `1` is a
"popular" wrong answer (a common mistake), `2..21` are rare wrong answers. A
verifier scores each sample: correct answers get a latent score around `d`
(the verifier's *sharpness*), ordinary wrong answers around 0, and a fraction
`f` of wrong answers are *convincing*, scoring around `d + 1`, higher than real
solutions.

**Chain task** (Tasks 4–5), the chapter's worked example and second demo.
"Start with 4; add 7, 3, 9, 5, 8, 2; keep only the last digit." A solution
writes the running digit after each step (1, 4, 3, 8, 6, 8). The policy makes
a mistake at step t with probability `0.04 + 0.03 * op_t`, then carries on
from its own wrong digit, so one mistake usually ruins the answer. A later
mistake can land back on the right track by luck (a "right answer for the
wrong reason"). `ChainTask.exact_value(problem, prefix)` returns the true
probability of finishing correctly from a prefix; you never need it in your
code, but the tests and plots compare your Monte Carlo estimates to it.
`NoisyPRM(sigma)` is a stand-in for a learned PRM: the true value plus
log-odds noise of size `sigma`, fixed per prefix.

## Tasks, in order

| Task | File | Function | What it is |
|---|---|---|---|
| 1 | `ttc.py` | `pass_at_n`, `pass_at_k` | 1 − (1 − p)^N; the unbiased estimator 1 − C(n−c, k)/C(n, k) |
| 2 | `ttc.py` | `majority_vote_accuracy` | exact binomial sum for right/wrong answers, coin-flip ties |
| 3 | `ttc.py` | `majority_vote`, `best_of_n`, `weighted_vote` | the three selection rules (ties: first occurrence) |
| 4a | `process.py` | `mc_step_labels` | soft and hard labels from k rollouts per prefix |
| 4b | `process.py` | `process_rewards` | r_j = V_j − V_(j−1): step rewards as advantages |
| 5 | `process.py` | `beam_search` | step-level beam search with a PRM; count generated steps |

## What you should see

`python run.py` prints a one-line summary and saves `process_rewards.png`:

- **(a)** Majority vote with right/wrong answers: accuracy climbs to 1 when
  p > 1/2 and sinks to 0 when p < 1/2 (the Condorcet jury theorem). The
  circles (simulation using your `majority_vote`) sit on your exact formula.
- **(b)** A benchmark of 300 problems of mixed difficulty. pass@N (the oracle)
  heads to 1. Majority vote **plateaus** at about 0.5: the fraction of
  problems where the right answer is the most common one (dotted line). Best-of-N with an
  honest verifier keeps climbing; with a verifier that is fooled by 1% of the
  wrong answers it **levels off near 0.84 (from N ≈ 16 to 256) and then declines** (to about 0.76 at N = 1024). Weighted
  voting is not hurt by the convincing wrong answers, but it plateaus too.
- **(c)** Monte Carlo labels for a solution whose first mistake is step 3.
  The true value drops from about 0.39 to 0.05 at step 3. Soft labels with
  k = 4 are very noisy; with k = 32 they track the truth. Hard labels say "1"
  for every good prefix, and often for bad ones too (lucky recoveries).
- **(d)** The step rewards V_j − V_(j−1): a large negative reward lands on
  step 3, while the outcome-only advantage puts the same −0.28 on every step.
- **(e)** With a good PRM, step-level beam search reaches about 0.8 accuracy with
  only ~12 generated steps per problem, where best-of-N with the same compute
  gets about 0.5. Pruning bad prefixes early saves compute.
- **(f)** With a poor PRM, beam search loses most of its edge. Beyond the
  smallest budget it trails weighted voting and is roughly level with
  best-of-N, because the search now optimizes against the PRM's mistakes.
  Majority vote ignores the PRM entirely, and at the largest budget it is as
  good as anything else.

Your exact numbers will match these if your functions use the random number
generator in the documented order. Otherwise they will be close.

## Going further (optional)

- In `run.py`, change `fool_rate` in panel (b) to 0.05 and to 0.001. Where does
  best-of-N peak? Compare with the chapter's estimate N* ≈ 1/(f(1 − p)).
- Make the chain longer (`ChainTask(n_steps=12)`): does the gap between beam
  search and best-of-N grow? (It should: longer solutions give the PRM more
  chances to prune.)
- Replace the soft labels in panel (d) with hard labels. What happens to the
  reward at the mistake?
- Real-LM version: take `tinylm.py` from the Lab 15 download, treat each
  answer digit as a step, and label digits by rolling out from each prefix
  (sample completions of the prompt "37+48=7", and so on). On which digit does
  the value usually collapse when the model gets a sum wrong?