# Lab 21: a capstone research kit

This lab goes with Chapter 21, "The frontier" (`/rl/learn/21-frontier/`). Earlier labs taught one
algorithm each. This one gives you a small research setup to run your own experiments on
2025–26 questions, at TinyLM scale on a laptop CPU.

It has three parts:

1. **Five short tasks** (`opd.py`, `tools.py`), with tests. These are the pieces the projects
   need: on-policy distillation losses, the pass@k estimator, and MaxRL's advantage.
2. **Project 0, fully worked** (`run.py`): on-policy distillation vs. off-policy distillation
   vs. RL, using your Tasks 1–3.
3. **Projects 1–5**: open questions. Each starter script already runs an experiment end to
   end. You then do the research part: more seeds, a judgement on whether the effect is
   real, and one extension that needs new code.

Everything uses PyTorch on CPU (`torch.set_num_threads(1)`). Each script takes 1–2 minutes.
The first script you run also trains and caches two checkpoints (`teacher.pt` and
`student300.pt`, about 25 s). You can delete them at any time.

## Setup and how to run

```bash
cd 21-frontier
python -m pytest -q                     # your code: fails until you fill in the TODOs
RL_SOLUTIONS=1 python -m pytest -q      # reference solutions: all 16 tests pass (~10 s)
RL_SOLUTIONS=1 python run.py            # Project 0 -> opd_comparison.png (~1.5 min)
python run.py                           # the same, with your code (after Tasks 1-5)
python project5_sharpen_or_expand.py    # any project; most accept --steps and --seed
```

## The cast

| File | What it is | Do you edit it? |
|---|---|---|
| `tinylm.py` | the course's tiny transformer and the "37+48=" task (copied from `_shared`) | no |
| `common.py` | cached teacher/student checkpoints, evaluation (greedy accuracy, pass@k, entropy), a small configurable GRPO-family trainer `rl_train`, difficulty buckets | no (but read `rl_train`) |
| `opd.py` | **Tasks 1–3** | yes |
| `tools.py` | **Tasks 4–5** | yes |
| `run.py` | Project 0, fully worked | read it |
| `project*.py` | Projects 1–5, starter scripts | yes, for the extensions |

- **Teacher**: TinyLM after 1,200 SFT steps, about 100% accurate.
- **Student**: the same architecture after only 300 SFT steps, about 50% greedy and 30%
  sampled accuracy.

The trainer `rl_train` samples `G = 8` responses for each of 16 prompts per step. It scores
them with a reward function (default: exact match), forms group advantages (`"mean"`,
`"std"` or `"maxrl"`), and takes `epochs` clipped-ratio gradient steps with a token-mean
loss and **no KL term**, as most 2025–26 RLVR recipes do (Chapter 16).

## Tasks, in order

| Task | File | Function | Chapter idea |
|---|---|---|---|
| 1 | `opd.py` | `reverse_kl_rewards` | OPD's per-token reward $r_t = \log\pi_T(y_t\mid s_t) - \log\pi_\theta(y_t\mid s_t)$ |
| 2 | `opd.py` | `opd_pg_loss` | that reward inside a REINFORCE loss with discount 0 |
| 3 | `opd.py` | `full_vocab_reverse_kl` | the exact KL over the whole vocabulary at each visited state (DeepSeek-V4's choice) |
| 4 | `tools.py` | `pass_at_k` | the unbiased pass@k estimator $1 - \binom{n-c}{k}/\binom{n}{k}$ |
| 5 | `tools.py` | `maxrl_advantages` | MaxRL's advantage $(r_i - \hat p)/\hat p$ |

Run one task's tests with, e.g., `python -m pytest -q test_opd.py -k full_vocab`.

---

## Project 0 (worked example): on-policy distillation

**Question.** You have a strong teacher and a weak student. Is it better to (a) fine-tune the
student on the teacher's samples (off-policy distillation, i.e. SFT), (b) let the student
sample and have the teacher grade every token (on-policy distillation, sampled-token or
full-vocabulary), or (c) run RL with the 0/1 verifier?

**Setup** (`run.py`). There are four arms, each run for 200 steps with 64 sequences per step,
AdamW at learning rate 3e-4, and the same student checkpoint. The script measures sampled
accuracy (pass@1), pass@8, and the per-token KL(student ‖ teacher) on the student's own
samples.

**What we observed** (seed 0; your numbers will differ a little):

| arm | pass@1, start → end | KL to teacher, start → end |
|---|---|---|
| opd-full | 0.31 → 0.67 | 1.98 → 0.70 |
| sft-teacher | 0.31 → 0.64 | 1.98 → 0.75 |
| opd-sampled | 0.31 → 0.54 | 1.98 → 1.01 |
| rl | 0.31 → 0.54 | 1.98 → 1.07 |

**How to read it.**

- *Dense beats sparse.* Both full-vocabulary OPD and SFT on teacher samples give a learning
  signal at every token. The RL arm learns from a single bit per response.
- *Full-vocabulary beats sampled-token.* With the same samples, the sampled-token estimate
  of the KL is much noisier. DeepSeek-V4 gives exactly this reason for choosing
  full-vocabulary OPD.
- *Why doesn't on-policy beat SFT by more?* The answers here are only 3–4 tokens long, so the
  student rarely drifts into states the teacher never visits. The compounding-error argument
  from Chapter 10 needs long sequences before it bites.

**Extensions.** (i) Make sequences longer: change the task to output the sum *twice*, or to
add three numbers. Does OPD's lead over SFT grow? (ii) Give the student a different
architecture from the teacher (`TinyLM(d=48)`). Does sampled-token OPD still work?

---

## Project 1: does clip-higher slow entropy collapse?

**Hypothesis.** DAPO's decoupled clip (ε_low = 0.2, ε_high = 0.28) keeps policy entropy
higher than the symmetric clip [0.8, 1.2], and this helps pass@k at large k (Chapter 16,
Chapter 19).

**Starter.** `project1_clip_higher.py` runs symmetric, clip-higher and no-clip for 150 steps,
with 4 gradient epochs per batch so that the ratio actually moves away from 1.

**Measure.** Entropy, pass@1 and pass@16 against step.

**What we observed (one seed).** Entropy collapsed in every arm, from 0.54 to about 0.22–0.26
nats per token. Clip-higher ended at 0.235 and symmetric at 0.222, a difference well inside
seed noise. No clip ended with the highest entropy but the lowest accuracy (0.50 vs 0.57).
**On this tiny problem we found no clear clip-higher effect.**

**Your job.** Run seeds 0–4 for each arm and report the mean ± spread of final entropy.
Then try these extensions:

- Increase `epochs` to 8 and ε_high to 0.4.
- Log the **clip fraction** (the share of tokens whose ratio hit either bound). If the
  upper clip almost never binds, clip-higher cannot matter. That could explain a null
  result.

**Expected.** Entropy collapse itself is robust. Whether clip-higher helps depends on how
often the upper clip binds, so it may only show up with more off-policy epochs.

---

## Project 2: rubric rewards, and how a rubric gets hacked

**Hypotheses.**

1. A partial-credit rubric learns at least as fast as a binary verifier.
2. A rubric item that is *not* the goal ("looks thorough": the answer has 3+ digits) gets
   exploited once its weight is comparable to correctness (Goodhart; Chapters 14 and 19).

**Starter.** `project2_rubric.py` runs three reward functions (binary, partial credit, and
"stylish" with weight `--w`, default 2) for 200 steps.

**Measure.**

- Gold accuracy (exact match).
- Proxy reward (the training reward).
- The fraction of 3-digit answers to problems whose sum is below 100.

**What we observed (one seed):**

| reward | gold pass@1 | proxy | 3-digit answers to small sums |
|---|---|---|---|
| binary | 0.29 → 0.55 | 0.56 | 0.02 |
| partial credit | 0.29 → 0.54 | 0.74 | 0.01 |
| stylish, w = 2 | 0.29 → **0.29** | 0.61 | **0.98** |

Partial credit did not speed up learning here: its items are strongly correlated with
exactness. The stylish rubric was hacked almost completely. The proxy went up, the gold
accuracy went nowhere, and the model now pads small sums to 3 digits. With `--w 1` the
hack was much weaker (0.08 of small sums padded, gold 0.50).

**Your job.**

- Find the threshold weight $w^*$ where the hack takes over (try w = 1, 1.25, 1.5, 2).
  Explain $w^*$ with a back-of-envelope argument. Compare the expected reward of "pad
  everything" with that of "answer honestly" at the current accuracy.
- Then add a **pitfall** item, a negative criterion as in Rubrics as Rewards, that
  removes the hack without hurting gold accuracy.

---

## Project 3: a self-play-style proposer at the learner's frontier

**Hypothesis.** Absolute Zero's learnability reward ($0$ if the solver never succeeds,
else $1-\hat p$) makes a proposer avoid both unsolvable and already-solved problems.
Training only on a target the student cannot solve gives (almost) no gradient.

**Starter.** `project3_learnability.py` has five problem buckets. The student's success rate
is about 2% on 1d+1d (rare in its SFT data), 30–50% on the 2-digit buckets, and ~0% on
3d+1d (never seen). It compares three ways of choosing problems:

- **target-only**: train only on the 3d+1d bucket;
- **uniform**: pick buckets uniformly;
- **learnability**: a softmax bandit (Chapter 2) over each bucket's running learnability
  reward.

**What we observed (one seed, 200 steps).**

- *Target-only* learned nothing: 3d+1d stayed at 0.00, and the other buckets drifted down.
- *Uniform* and *learnability* both improved the 2-digit buckets. Learnability ended
  slightly ahead on sum ≥ 100 (0.62 vs 0.50).
- Learnability spent only 3% of its samples on the unsolvable bucket, against 22% for
  uniform.
- Neither arm fixed 1d+1d. At p ≈ 0.02, most groups of 8 contain no success.

**Your job.**

- Run 3 seeds.
- Plot bucket usage over time.
- Try `G = 32` for the 1d+1d bucket only. Bigger groups make a success in a rare-success
  group likely, as in BroRL's "more rollouts" argument.
- Hardest extension: let the *model itself* propose problems. Prompt it with a new token,
  parse its output as "a+b=", and reward it with the learnability reward. That is real
  self-play.

---

## Project 4: MaxRL vs. the mean baseline vs. GRPO's std

**Hypothesis.** MaxRL's advantage $(r-\hat p)/\hat p$ up-weights hard prompts. Its
population weight is $w(p)\to 1/p$, the maximum-likelihood weight, while plain RL has
$w(p) = 1$ and GRPO has $1/\sqrt{p(1-p)}$. So it should improve pass@k at large k more than
pass@1 (Tajwar et al. 2026).

**Starter.** `project4_maxrl.py` runs 3 arms for 200 steps from the same student.

**What we observed (one seed).**

| advantage | final pass@1 | final pass@16 |
|---|---|---|
| mean | 0.55 | 0.89 |
| std | 0.51 | 0.88 |
| maxrl | 0.51 | 0.91 |

The direction matches the hypothesis, but the gaps are about the size of the evaluation
noise (100 prompts × 16 samples).

**Your job.**

- Increase `n_prompts` in `evaluate` to shrink the noise.
- Run 3 seeds.
- Start from a weaker student (`--sft-steps 200`), where more prompts have small p. Does
  MaxRL's advantage grow?

---

## Project 5: does RL sharpen or expand?

**Hypothesis.** RL raises pass@1 a lot and pass@k at large k little or not at all (Yue et al.
2025). The base model may even win at large k.

**Starter.** `project5_sharpen_or_expand.py` trains the student with RL for 300 steps. It then
plots pass@k for k = 1…64 before and after, on in-distribution 2-digit problems and on
never-seen 3d+1d problems.

**What we observed (one seed).** In distribution, pass@1 rose from 0.31 to 0.67, while
pass@64 *fell* from 0.97 to 0.92. The curves cross between k = 16 and k = 32. On 3d+1d both
models are near zero at every k: RL did not create the skill. This is the "sharpening"
picture in miniature.

**Your job.**

- Train longer (`--steps 1000`) and check whether the crossing point moves. ProRL (Liu et al.
  2025) reports that prolonged RL with KL control and reference resets expands the
  boundary.
- Add a KL penalty to `rl_train` and see if it protects pass@64.

---

## Writing it up

For each project, write half a page with:

1. the hypothesis;
2. the setup, including everything you held fixed;
3. a plot with seeds shown;
4. the result, in one sentence;
5. the confounders you could not rule out.

The chapter's section "How to read (and write) an RL-for-LLM paper" has the checklist: base
model, baselines, KL and reference, evaluation sampling, contamination, seeds and compute.
Apply it to your own work first.

## Files

```
tinylm.py, common.py           helpers (complete)
opd.py, tools.py               Tasks 1-5 (yours)
test_opd.py, test_tools.py     16 tests
run.py                         Project 0 (worked)
project1_clip_higher.py ... project5_sharpen_or_expand.py
solutions/                     reference code for Tasks 1-5
```