Section 13

Policy gradients & REINFORCE

How a reward becomes a weight update

So far we have a reward model that scores a response, and we want the language model to produce responses it scores highly. The obvious idea (just backpropagate the reward into the model, the way we train on next-token loss) doesn’t work. Generating a response means sampling tokens one at a time, and you can’t take a gradient through a random choice of which token got picked. (And the reward might be a black box we can’t differentiate at all.) So we need a way to improve the model from nothing but the single number the reward hands back at the end. That tool is the policy gradient, and this chapter builds it from scratch.

The LLM as a policy

In reinforcement learning, a policy is a rule for choosing actions. We write it $\pi_\theta(a \mid s)$ : given a state $s$ , it gives a probability distribution over actions $a$ , controlled by parameters $\theta$ . The whole vocabulary of RL maps cleanly onto a language model:

The state is the prompt-so-far: the input $x$ plus whatever tokens have already been generated.
An action is choosing the next token.
The policy is the model itself. At each step it emits logits, the softmax turns them into a distribution over the vocabulary, and we sample.

Because the model generates one token at a time and feeds each choice back in, the probability of a complete response $y = (y_1, \ldots, y_T)$ given prompt $x$ factorizes:

\pi_\theta(y \mid x) = \prod_{t=1}^{T} \pi_\theta(y_t \mid x, y_{<t})

This is the same autoregressive factorization you already know from next-token prediction. Nothing about the network has changed. What changes is how we score the output and what gradient we push through it.

A single act of generating a full response is a rollout. The sequence of states and actions it produces (prompt, token, new state, token, …) is a trajectory. At the end, a scalar reward $r(y)$ (here, the reward-model score) judges the whole thing. Summed over the trajectory, that’s the return; for a single terminal reward at the end of a response, return and reward coincide, and we’ll just write $r(y)$ .

To make all four words concrete, fix one prompt and sample the model four times. Say the prompt is $x$ = “Explain in one sentence why you see lightning before you hear thunder.” Each sample below is one rollout; sampling is random, so the same prompt yields different responses, some good, some not:

Rollout	Sampled response $y$	Reward $r(y)$	Why the reward model scores it that way
1	”Light travels much faster than sound, so the flash reaches your eyes almost instantly while the rumble takes several seconds to reach your ears.”	$+2.1$	Correct, complete, follows the one-sentence instruction.
2	”Because light is faster than sound.”	$+0.3$	Correct but unhelpfully terse; states the fact without explaining the experience.
3	”Because the lightning happens first, and the thunder only happens a few seconds later.”	$-1.8$	Fluent and confident, but a misconception: they happen together; only their signals arrive apart.
4	”You see the flash first because the light arrives first, the light arrives first, and the light arrives…”	$-0.9$	Starts correctly, then degenerates into repetition.

Each row’s trajectory is the full state-action alternation behind that response. For rollout 2 it looks like: state = the prompt $x$ ; action = sample the token “Because”; new state = $x$ + “Because”; action = sample ” light”; and so on until the model emits its end-of-response token, seven actions in total. Rollout 1’s trajectory is the same kind of object, just ~30 actions long. The reward is one number for the whole trajectory, handed out only at the end: nothing scored “Because” on its own, and no token was labeled right or wrong. That’s the whole learning problem in miniature: from these four (response, score) pairs alone, the update we build next should make responses like rollout 1 more likely and responses like rollouts 3 and 4 less likely.

The objective: maximize expected reward

We want a policy whose rollouts tend to score high. Formally, maximize the expected reward over the responses the policy itself generates:

J(\theta) = \mathbb{E}_{y \sim \pi_\theta(\cdot \mid x)}\big[\, r(y) \,\big]

Read it carefully: the expectation is taken over $y$ drawn from the current policy. That is the crux of the difficulty. The thing we are differentiating with respect to, $\theta$ (the model parameters), also controls the distribution we are averaging over. Change $\theta$ and you don’t just change $r(y)$ : you change which $y$ ‘s show up at all. Worse, $r$ may be a non-differentiable black box: a reward model, a unit-test pass/fail, a human thumbs-up. We can’t simply backpropagate through $r$ .

The log-derivative trick

The way out is a small, beautiful identity. Write the expectation as a sum (over all possible responses) and differentiate:

\nabla_\theta J(\theta) = \nabla_\theta \sum_y \pi_\theta(y) \, r(y) = \sum_y r(y) \, \nabla_\theta \pi_\theta(y)

Two symbols are new here. $\nabla_\theta$ is the gradient with respect to the parameters: the vector of partial derivatives, one per weight, pointing in the direction in weight-space that most increases whatever follows it. And $\sum_y$ is a sum over every possible response $y$ the model could produce, an astronomically large set, which is exactly what makes this expression intractable until we fix it. ( $J(\theta)$ , $\pi_\theta(y)$ , and $r(y)$ are the objective, policy probability, and reward from above.)

The reward $r(y)$ doesn’t depend on $\theta$ , so it slides outside the gradient. We’re left with $\nabla_\theta \pi_\theta(y)$ , which is still awkward: it’s not an expectation we can sample. The trick is to multiply and divide by $\pi_\theta(y)$ :

\nabla_\theta \pi_\theta(y) = \pi_\theta(y) \, \frac{\nabla_\theta \pi_\theta(y)}{\pi_\theta(y)} = \pi_\theta(y) \, \nabla_\theta \log \pi_\theta(y)

The one new object is $\nabla_\theta \log \pi_\theta(y)$ : the gradient of the log-probability the policy assigns to response $y$ . Read it as a direction in weight-space: the way to nudge every parameter so that this exact response becomes more likely.

That last step uses the identity $\nabla \log f = \nabla f / f$ . Substituting back, the $\pi_\theta(y)$ out front turns the sum back into an expectation:

\boxed{\,\nabla_\theta J(\theta) = \mathbb{E}_{y \sim \pi_\theta}\big[\, r(y) \, \nabla_\theta \log \pi_\theta(y) \,\big]\,}

Every symbol here is now one you’ve met, so read the finished formula term by term. $\mathbb{E}_{y \sim \pi_\theta}[\cdot]$ averages over responses sampled from the current policy; for each sampled response, $\nabla_\theta \log \pi_\theta(y)$ is the direction that makes that response more likely, and $r(y)$ is a scalar weight on it. Put together, it says: sample responses, and step toward making each one more likely in proportion to the reward it earned. Push up what scored well, push down what scored badly.

This is the policy gradient, also called the score-function estimator (the quantity $\nabla_\theta \log \pi_\theta$ is the “score” in statistics). It is the entire foundation of RL for LLMs.

Why the log-derivative trick is the whole game

We wanted the gradient of an average, but the average is over a distribution that itself depends on $\theta$ . Naively that’s intractable. The identity $\nabla_\theta \pi_\theta = \pi_\theta \, \nabla_\theta \log \pi_\theta$ rewrites the gradient as an expectation of $r(y)\,\nabla_\theta \log \pi_\theta(y)$ , and an expectation is exactly what we can estimate by sampling.

So in practice we draw $N$ rollouts $y^{(1)}, \ldots, y^{(N)}$ from the current policy and average:

\nabla_\theta J \approx \frac{1}{N} \sum_{i=1}^{N} r(y^{(i)}) \, \nabla_\theta \log \pi_\theta(y^{(i)})

Crucially, $r(y)$ appears only as a scalar multiplier. We never differentiate through it. That is precisely why the reward is allowed to be a black box.

REINFORCE: push up what worked

Turning that estimator into a learning rule gives the REINFORCE algorithm (Williams, 1992). One step is:

Sample a batch of rollouts from the current policy.
Score each one with the reward function to get $r(y^{(i)})$ .
Update by ascending the estimated gradient:

\theta \leftarrow \theta + \eta \, \frac{1}{N} \sum_i r(y^{(i)}) \, \nabla_\theta \log \pi_\theta(y^{(i)})

The intuition is exactly what you’d hope. $\nabla_\theta \log \pi_\theta(y)$ is the direction in parameter space that makes response $y$ more likely. Multiply it by the reward and sum: rollouts with high reward get pushed up in probability, and (if rewards can be negative) low-reward rollouts get pushed down. The policy reshapes itself to put more mass on the responses that scored well. It is supervised learning where the model writes its own training examples and the reward decides how hard to learn from each one: weighted maximum likelihood, with the weights handed down by the reward.

Because the factorization is autoregressive, $\nabla_\theta \log \pi_\theta(y) = \sum_t \nabla_\theta \log \pi_\theta(y_t \mid x, y_{<t})$ , a sum of per-token log-prob gradients, every one of which you already compute during an ordinary forward/backward pass.

On-policy vs. off-policy

One subtlety is hiding in the expectation $\mathbb{E}_{y \sim \pi_\theta}$ . The data must come from the policy we are currently updating. An algorithm that learns only from its own fresh rollouts is on-policy; REINFORCE is the canonical example. The moment you take a gradient step, $\theta$ moves, your old samples are stale, and strictly speaking you must collect new rollouts before the next step.

The alternative is off-policy learning: reusing data generated by a different (older, or entirely separate) policy. That’s more sample-efficient (you don’t throw rollouts away after a single step) but it requires a correction, because you’re now averaging over the wrong distribution. The fix, importance sampling, is exactly the bridge that leads from REINFORCE to PPO in chapter 17.

Try it

Below is a tiny world: a bandit with a handful of actions, each hiding its own average reward. The policy is a single softmax over the actions. Watch pure REINFORCE in action (sample, reward, update) and see the action distribution migrate toward whatever earns reward.

Policy gradient sandbox (REINFORCE bandit)

Click an arm to sample it. Each pull collects a noisy reward and applies one update: logit ← logit + η·reward. Watch the policy shift.

steps 0avg reward 0.000

Policy π(arm) — click a bar to sample

Last sampled arm

—

Reward received

—

Every reward here is positive, so every pull pushes its own arm up — the only question is by how much. That means early lucky pulls on a mediocre arm inflate it, and the policy can wander for a long time before arm C's edge wins out. Reset and repeat a few times: some runs lock on quickly, others get stuck chasing a lucky arm. That run-to-run scatter is the variance problem, and the next chapter's baseline is the fix.

Play with it and one thing jumps out: the updates are noisy. With a small number of samples, a couple of lucky high-reward draws can yank the distribution around (sometimes in the wrong direction) before later samples correct it. That noise is not a quirk of the toy. It is the central weakness of REINFORCE.

The catch: variance

The policy-gradient estimator is unbiased (on average it points the right way), but it has punishing variance. Two reasons. First, we estimate an expectation from only a handful of samples. Second, and more insidiously, the raw reward $r(y)$ scales every update. If all your rewards happen to be large and positive (say every response scores between $+8$ and $+10$ ) then REINFORCE pushes up the log-prob of every rollout, hard, and only weakly distinguishes the $+10$ from the $+8$ . The gradient is dominated by the level of the reward rather than by which responses were better than average. With billions of parameters and long sequences, that variance makes training slow and brittle.

The good news: there is a clean fix that costs nothing in bias. We can subtract a reference value from the reward (a baseline) so that what multiplies the log-prob becomes “how much better than expected was this rollout?” rather than the raw score. That single idea is the subject of the next chapter — and the value functions and advantage estimates built on top of it fill the two after that.