Lecture 17: Alignment - RL 2

What this lecture covers

This is the final lecture of the course, and it goes deeper into the mechanics introduced in the previous lecture rather than covering new material. It reframes reinforcement learning for language models explicitly: the state is the prompt plus tokens generated so far, each action is a token, and (for this course) the reward is computed only once the full response is done. The lecture argues that language modeling's RL setting is unusual compared to robotics, because any sequence of tokens is reachable, so the hard part is making the model's own generated "scratchpad" lead to a correct answer, not reaching some particular state.

The core of the lecture works from the plain policy gradient objective through the idea of baselines and advantage functions, using a small two-state numeric example to make variance reduction concrete. It then walks through real Python code for a toy sorting task, implementing GRPO's reward shaping, clipped loss, and KL penalty from scratch, and shows training runs on this toy task. After watching, you should be able to derive why a baseline doesn't change the optimization problem, and read GRPO-style training code with an understanding of what each term does.

Key ideas

• State, action and reward for LMs: the state is the prompt plus generated tokens so far, an action is the next token, and this course focuses on outcome rewards that are deterministic and verifiable, computed once per full response.
• Naive policy gradient is SFT weighted by reward: it looks just like supervised fine-tuning on the model's own generations, except each example is scaled by its reward.
• Sparse rewards stall training: if a policy is bad enough that it rarely gets reward 1, most gradient updates are zero, and the model does not improve.
• Baselines don't bias the objective: subtracting any function B(s) that doesn't depend on the action a leaves the expected value of the policy gradient unchanged, because it integrates to zero.
• Baselines reduce variance: in a worked two-state example, subtracting a well-chosen baseline cuts gradient variance roughly fivefold without changing what the optimal policy is.
• The advantage function: A(s,a) = Q(s,a) - V(s), and subtracting the expected reward given the state as a baseline is equivalent to optimizing the advantage.
• GRPO's group structure: prompting a language model to generate several responses per prompt gives a natural comparison set, so the group mean becomes the baseline and the group standard deviation normalizes the scale.
• Freezing constants correctly: when computing a ratio like p / p_old for the same parameters, you must stop gradients through p_old, or the gradient of the ratio becomes zero.

Walkthrough

Defining the RL problem for language models (0:05)

The lecture sets up state, action, reward, transition dynamics, and policy for language modeling explicitly. Because appending a token is the entire transition function, planning at test time is straightforward compared to robotics, but the "state" itself is just whatever tokens the model chooses to write, giving it a lot of freedom without any guarantee that freedom leads to a correct answer.

Naive policy gradient and its variance problem (5:10)

Taking the gradient of expected reward gives the standard policy gradient form, which the lecture compares directly to SFT: same update, but every example is weighted by its reward. With binary correct/incorrect rewards, this reduces to fine-tuning only on responses that happened to be correct. The lecture explains that reinforcement learning has much higher variance than supervised learning, and that with sparse rewards a weak initial policy can produce mostly zero gradients, stalling progress. It also explains why the effective training set changes over time: after each update, the policy that generates the next batch of responses is different.

Baselines and variance reduction (15:23)

A two-state, two-action toy example (with rewards 11, 9, 0, and 2) shows how naive policy gradient can get stuck reinforcing a suboptimal action just because its reward is locally larger than another state's best action. The lecture proves that subtracting any baseline B(s) independent of the action leaves the expectation unchanged, then shows numerically that choosing a good baseline cuts gradient variance from about 5.3 to about 1.1. It connects this to the advantage function and notes that GRPO-style algorithms are a specific, simple choice of estimator within this general framework.

Building a toy sorting task and model (33:33)

To make the algorithm concrete without a full transformer, the lecture defines a task of sorting n numbers and a small non-autoregressive model that decodes each output position independently. It compares reward functions: a sparse all-or-nothing sorted check, a reward counting correctly placed positions, and a denser reward that adds credit for each input token reproduced and each adjacent correctly-ordered pair, noting this richer reward has an exploitable loophole left as an exercise.

Computing deltas, log-probs, and the GRPO loss in code (43:43)

The lecture walks through compute_deltas, showing raw rewards, centered rewards (subtracting the group mean), and normalized rewards (also dividing by group standard deviation), and discusses what happens when all responses in a group get the same reward: centering produces zero updates, since there's no useful comparison signal. It then covers a subtlety with stop-gradients when computing ratios like p / p_old, walks through the GRPO clipped loss and its similarity to the naive weighted loss, and derives the low-variance KL divergence estimator used in the KL penalty term.

Full training loop and toy results (59:04)

The complete algorithm is assembled: an outer loop that samples a batch of responses, then an inner loop that takes several gradient steps on those same responses (since generating rollouts is expensive relative to gradient steps), optionally regularizing toward a frozen reference model via a KL penalty. Running this on the sorting task, the lecture shows centered and normalized rewards behaving as expected, including batches where all responses tie and no update happens. It closes by pointing out that GRPO training loss curves are not a reliable progress signal, since the underlying data distribution shifts every iteration; reward is the more meaningful metric to watch.

Before you watch

• Watch the previous lecture on RL from verifiable rewards, PPO, and GRPO, since this lecture builds directly on those definitions.
• Be comfortable with expectations, gradients, and basic probability notation, since the lecture works through several derivations on the board.
• Familiarity with PyTorch-style tensor operations will help with the code walkthrough.

Check your understanding

1. Why is a naive policy gradient update, in the binary-reward case, mathematically equivalent to supervised fine-tuning on only the correct responses?
2. Why does subtracting a baseline B(s) not change what the optimal policy is, even though it changes the variance of the gradient estimate?
3. In the toy sorting task, what happens to the update when every sampled response for a prompt receives the same reward, and why?
4. Why must p_old be wrapped in a no-gradient operation when computing the PPO/GRPO ratio p / p_old?
5. Why does the lecture argue that the training loss curve is not a trustworthy signal of progress in this RL setting, even though the reward is increasing?