# Lab 16: one loss for the whole critic-free zoo

The chapter argues that REINFORCE, RLOO, GRPO, Dr. GRPO, DAPO, GSPO, CISPO and REINFORCE++ are one estimator with six knobs. In this lab you make that literal: you write **one** configurable `pg_loss`, and every algorithm becomes one line of configuration in `common.PRESETS`. The tests check each preset against its paper's formula, written out with explicit loops. Then you race the presets on the TinyLM addition task, and watch GRPO's length bias appear on a toy built to expose it.

Chapter: <https://learn.xhelio.ai/rl/learn/16-policy-gradient-zoo/>

## Files

| File | What it is |
|---|---|
| `zoo.py` | **Your file.** Five tasks marked `TODO`. |
| `common.py` | Given: `PGConfig` (the knobs), `PRESETS` (the zoo), the TinyLM training loop `train`, and `LengthToy` / `train_length_toy`. Read the `PGConfig` docstring and `PRESETS` first. |
| `tinylm.py` | Given: the tiny transformer and the "37+48=" task (same as Labs 11–19). |
| `test_zoo.py` | Tests. |
| `run.py` | The experiment: saves `zoo_comparison.png`. |
| `solutions/zoo.py` | Reference solution. Try not to peek. |

## Tasks

Shapes: `N` responses (prompts × G), `T` the longest response. `logp`, `old_logp`, `ref_logp`, `mask` are `(N, T)`; `rewards` and `group_ids` are `(N,)`.

1. **`group_advantages`** (knobs 1–2). Baseline: none, batch mean, group mean (GRPO), leave-one-out (RLOO). Scale: none, group std (GRPO), batch std (REINFORCE++). The docstring has the chapter's 8-response example worked out. One test checks that group-mean advantages are exactly (G−1)/G times leave-one-out ones (Exercise 16.2).
2. **`aggregate`** (knob 3). Sequence-mean, token-mean, constant. One test checks the length bias directly: under sequence-mean each token of response i gets weight 1/(N·|y_i|).
3. **`token_surrogate`** (knob 4). Token ratio or GSPO's sequence (geometric-mean) ratio; no clip, PPO clip with separate ε_low/ε_high, or CISPO's stop-gradient clipped weight. The tests check the gradient of every case in the chapter's PPO-vs-CISPO table.
4. **`kl_per_token`** (knob 5). k1, k2, k3.
5. **`pg_loss`** (everything). KL in the reward or in the loss, then advantages, surrogate, aggregation. The tests compare loss *and gradient* for seven presets against the papers' formulas, check that an all-correct group moves nothing, check both KL placements, and train on the LengthToy for a few seconds.

```bash
python -m pytest -q                   # all tests; they fail until you fill in the TODOs
python -m pytest -q -k task1          # just Task 1 (also task2 ... task5)
RL_SOLUTIONS=1 python -m pytest -q    # run the tests against the reference solution
```

All tests together take under 10 seconds.

## Then run the experiment

```bash
python run.py            # about 2 minutes; saves zoo_comparison.png
```

The first run also fine-tunes and caches the starting TinyLM (`sft_tinylm.pt`, about 15 seconds). Then:

- **Top row, TinyLM addition.** REINFORCE, RLOO, GRPO and Dr. GRPO train for 120 steps from the same SFT checkpoint (about 50% sampled accuracy), 16 prompts × G = 8 per step, two gradient updates per batch.
- **Bottom row, LengthToy.** A tabular "language model": a response is one answer token, then filler tokens, then `<eos>`. Filler never changes the reward, but the policy can learn to write more of it after some answers than after others. GRPO, DAPO and Dr. GRPO train for 300 on-policy steps.

**What you should see** (reference solution; your numbers can differ a little):

- **TinyLM accuracy**: every preset improves, from about 0.51 to roughly 0.63–0.70, and the order changes from one evaluation to the next. With 500 evaluation prompts the noise is about ±2 points, so most differences are noise. That matches the bigger picture in the chapter: on a plain task, these estimators perform similarly.
- **TinyLM length** stays flat: the answer's length is fixed by the sum, so there is nothing for a length bias to act on.
- **LengthToy**: under GRPO (sequence-mean), wrong answers roughly double in length (about 4 to about 8 tokens, with noisy spikes past 10, since few wrong answers remain late in training) and right answers shrink from about 3.5 to about 2.2. Under DAPO (token-mean) and Dr. GRPO (constant normalizer) both stay flat. All three solve the task. This is the modified reward (r − b)/|y| of the chapter's "Knob 3" section, made visible.

## Things to try

- `python run.py --toy-presets grpo gspo`: the two curves coincide. Why? (Exercise 16.10.)
- `python run.py --presets grpo dapo cispo gspo`: the clip-related presets. Watch the printed "clip/outside-band frac" (the fraction of tokens whose ratio left the clip band on the last update): GSPO's tiny band clips almost every response on the second update.
- Make the clip matter: call `train(..., updates_per_batch=8, lr=1e-3)` from a Python prompt and compare `grpo` with `cispo`.
- Add your own preset with `common.preset("grpo", scale="batch_std")` or invent one (MaxRL's "divide by the number of successes" is a good challenge: it needs a new baseline/scale option in `group_advantages`).