# Lab 11: the language model as a policy

Goes with Chapter 11 (`/rl/learn/11-llm-as-policy/`). PyTorch and numpy, CPU only.
The tests take a few seconds; `run.py` takes 15–40 s.

You build the plumbing that every LLM RL algorithm in Chapters 12–21 relies on:
log-probabilities of a response (with the shift and the mask done right),
temperature sampling, the three KL estimators, and the KL-regularized optimum
π* ∝ π_ref·exp(r/β), computed both for whole responses and one token at a time.

## Files

| File | What it is |
|---|---|
| `policy.py` | **Your file.** Five tasks marked `TODO`. |
| `tinylm.py` | A ~100k-parameter transformer and the "37+48=" addition task (copied from the course's shared code; complete). |
| `toylm.py` | An enumerable toy LM: every response is 2 tokens, so π_ref, π* and all KLs can be computed exactly. Same scenarios as the chapter's demo. Complete. |
| `test_policy.py` | The tests. |
| `run.py` | Plots your functions at work (saves `llm_as_policy.png`). |
| `solutions/policy.py` | Reference solution. Try first! |

## How to run

```bash
cd 11-llm-as-policy
python -m pytest -q                  # your code: fails until you fill in the TODOs
python -m pytest -q -k task1         # just one task's tests
RL_SOLUTIONS=1 python -m pytest -q   # reference solution: all 24 pass
python run.py                        # plots -> llm_as_policy.png (first run trains and caches sft_tinylm.pt)
```

## Tasks, in order

| Task | Function | What it is | Chapter section |
|---|---|---|---|
| 1a | `response_logprobs` | log_softmax, then pick the log-prob of the token that actually came next. **The logits at column t belong to the token at column t + 1.** Return only the response part. | "Log-probs in practice: the shift and the mask" |
| 1b | `sequence_logprob` | Sum the per-token log-probs over real response tokens. Use `torch.where`, not `* mask`: 0 × (−∞) = nan. | same |
| 2 | `sample_with_temperature` | Sample from softmax(z / T); T = 0 means argmax. `generate` (provided) uses it to run the token-level MDP. | "Temperature" |
| 3 | `kl_estimators` | With ℓ = log π − log π_ref: k1 = ℓ, k2 = ℓ²/2, k3 = e^(−ℓ) − 1 + ℓ. | "Estimating the KL from samples" |
| 4 | `tilted_optimum`, `kl_objective` | π* = π_ref·e^(r/β)/Z over all responses (in log space), and J(π) = E_π[r] − β·KL(π ‖ π_ref). | "The KL-regularized objective" |
| 5 | `token_level_optimum` | The same π*, built backwards with soft values V(i) = β log Σ_j π_ref(j∣i) e^(r_ij/β). | "The optimum, one token at a time" |

Tips:
- Task 1a is three lines, but it is the line people get wrong in real code. Draw
  the 4-column example of `test_task1_shift_on_a_hand_example` (or the 11-column
  TinyLM row in the chapter) on paper first: which logits column predicts which token?
- Task 4: `np.log(0)` is `-inf` and `exp(-inf)` is 0, so zero-probability
  responses take care of themselves, as long as you silence numpy's warning
  (`np.errstate(divide="ignore")`) and subtract the max of the finite entries.
- Task 5: `np.logaddexp.reduce(x, axis=...)` is a stable log-sum-exp.

## What you should see

`pytest`: 24 tests pass. They check hand-computed values from the chapter's
worked examples (π* = 0.172, 0.312, 0.062, 0.454 at β = 1; V(a) = 0.523,
V(b) = 1.434), that the prompt gets no gradient, that k1 and k3 are unbiased
and k2 is not, and that no random policy beats π* (the gap is exactly
β·KL(q ‖ π*), the Gibbs identity from Chapter 10).

`run.py` prints a one-line summary and saves three panels:

1. **Temperature trades accuracy for diversity.** The SFT TinyLM is right about
   80% of the time with greedy decoding, about 55% at T = 1 and below 30% at
   T = 2, while the number of distinct answers among 8 samples climbs from 1 to
   about 5. (Exact numbers depend on your machine's PyTorch version.)
2. **How far temperature moves the policy**, per response, estimated with k1 and
   k3 on the sampled tokens, against the exact value (dashed; computed by
   summing over the whole vocabulary at every position, which you can afford for
   a 15-token vocabulary). At T > 1 both estimators agree with the truth. At
   T = 0.25 **k3 is far too low**, and its error bars don't even cover the
   truth. That's not a bug: k3 is unbiased, but when the sampling policy is much
   *sharper* than the reference, k3's large values sit on tokens the sampler
   almost never draws (Exercise 11.4 works through a two-token example). RL
   makes policies sharper, so this matters in practice.
3. **The reward–KL frontier of the toy LM.** The solid curves are π* as β
   sweeps from large to small. Temperature (dotted) reaches a different and
   always lower curve: in the "worked" scenario, lowering T even *reduces*
   reward, because the most likely responses there are not the best ones. In
   the "hack" scenario the dashed curve is the *true* reward of π*: it rises,
   peaks around KL ≈ 0.5 nats, and then falls as π* moves onto the response
   the reward model scores wrongly. That's why the leash exists.

## Going further (optional)

- In `run.py`, also estimate the KL with k2. Where is its bias largest?
- Suppose the tokens were sampled from an older policy π_old, not from π.
  Multiply k3 by the importance weight π/π_old (this is what the DeepSeek-V3.2
  report does) and check numerically, on the 3-token example of the tests, that
  the weighted estimator is unbiased for KL(π ‖ π_ref) while plain k3 is not.
- Take the "hack" scenario and find the β that maximizes the *true* reward.
  How much proxy reward does π* have there? (Chapter 19 turns this into a
  scaling law.)