Policy gradients and REINFORCE

Here is a puzzle. You have a policy with a knob $\theta$. Turn the knob and the policy’s action probabilities change. You run the policy, it rolls a die to pick an action, the world responds, and at the end you get a score. You want to turn the knob in the direction that raises the average score. So you want a gradient. But the score depends on $\theta$ only through a dice roll, and the world in between is a black box you cannot differentiate. It might be a physics simulator, a human clicking thumbs-up, or a unit test that passes or fails. How do you take a derivative through a dice roll?

The answer is one line of calculus, the log-derivative trick, and it is the foundation of almost everything from here to the end of the course. In Chapter 6 you learned a value function $Q$ and acted by taking the argmax. In this chapter you skip the middleman: you write down a policy $\pi_\theta$ and climb the expected return directly. The algorithm you will build, REINFORCE, is more than thirty years old. It is also, with a better baseline, the core of the methods used to train today’s reasoning models.

Why learn the policy directly?

DQN learns how good each action is, and the policy is implied: pick the action with the largest $Q$. That works well for Atari. It has three limitations that matter.

Continuous actions. A robot arm chooses torques, real numbers. To act greedily you would have to solve $\max_a Q(s, a)$ over a continuum at every step, which is itself an optimization problem. A policy that outputs “a Gaussian with this mean and width” just samples.

Discontinuity. The greedy policy jumps. If $Q(s, \text{left})$ and $Q(s, \text{right})$ are close, a tiny change in the network flips the action from always-left to always-right. Small parameter changes cause large behavior changes, which makes learning jittery. A parameterized stochastic policy changes smoothly: nudge $\theta$ and the probabilities move a little.

Sometimes the best policy is random. In rock-paper-scissors against an opponent who adapts, any deterministic policy loses; the best you can do is play each move with probability 1/3. The same thing happens when the agent cannot tell two states apart. In Sutton and Barto’s “short corridor” example, two states look identical to the agent but the actions have opposite effects in them. Every deterministic policy loops forever, and the best policy goes right with probability about 0.59. A greedy-over-$Q$ agent can only express deterministic policies (plus ε of noise it does not choose).

There is a fourth reason, the one that matters most for this course: a language model is already a policy. It outputs a probability distribution over the next token, it is differentiable in its parameters, and it is very good before RL even starts. Nobody trains a $Q$-function over 100,000 tokens and takes an argmax. You start from the policy you have and improve it.

Policies with knobs

Softmax policies for discrete actions

Give each action $a$ a real-valued score $\theta_a$, called a logit (later, a neural network will compute the logits from the state). Turn the logits into probabilities with the softmax:

$$\pi_\theta(a) = \frac{e^{\theta_a}}{\sum_{b} e^{\theta_b}}.$$

Physicists will recognize the Boltzmann distribution: action $a$ is a state with energy $-\theta_a$ at temperature 1. Dividing all logits by a temperature $\tau$ makes the policy greedier as $\tau \to 0$ and more uniform as $\tau \to \infty$. That is exactly how sampling temperature works in a language model (Chapter 11). One property will matter later: adding the same constant to every logit changes nothing, since it cancels in the ratio.

Gaussian policies for continuous actions

For a real-valued action, let the policy be a normal distribution whose mean $\mu_\theta(s)$ (and maybe width $\sigma_\theta(s)$) comes from a network:

$$\pi_\theta(a \mid s) = \frac{1}{\sqrt{2\pi}\,\sigma} \exp\!\left(-\frac{(a - \mu_\theta(s))^2}{2\sigma^2}\right).$$

Everything in this chapter works for both kinds. We will mostly use the softmax because you can check every number by hand.

The objective

The quantity to maximize is the expected return of the policy,

$$J(\theta) = \mathbb{E}_{\tau \sim \pi_\theta}\big[G_0\big],$$

where $\tau = (s_0, a_0, r_0, s_1, \dots)$ is an episode generated by running $\pi_\theta$ in the environment, and $G_0 = \sum_t \gamma^t r_t$ is its return. In a bandit (one step, no state) this is just the average reward: $J(\theta) = \sum_a \pi_\theta(a)\, q_*(a)$, where $q_*(a)$ is arm $a$‘s true mean reward, as in Chapter 2.

A worked example you will see all chapter

Take a 3-armed bandit with mean rewards $q_*(1) = 1$, $q_*(2) = 2$, $q_*(3) = 4$, and start from logits $\theta = (0, 0, 0)$. The policy is uniform, $\pi = (\tfrac13, \tfrac13, \tfrac13)$, and

$$J = \tfrac13 \cdot 1 + \tfrac13 \cdot 2 + \tfrac13 \cdot 4 = \tfrac73 \approx 2.33.$$

Because you know $q_*$ here, you can differentiate $J$ directly. Using $\partial \pi(a) / \partial \theta_k = \pi(a)\,(\mathbb{1}[a = k] - \pi(k))$ (differentiate the softmax quotient; the next section derives the same fact in log form, and Exercise 7.2 does it in full), you get

$$\frac{\partial J}{\partial \theta_k} = \sum_a q_*(a)\, \pi(a)\big(\mathbb{1}[a=k] - \pi(k)\big) = \pi(k)\,\big(q_*(k) - J\big).$$

Read it in words: raise the logit of every arm that is better than the current average, in proportion to how often you play it. At the uniform policy:

$$\nabla_\theta J = \tfrac13\,(1 - \tfrac73,\; 2 - \tfrac73,\; 4 - \tfrac73) = \left(-\tfrac49,\; -\tfrac19,\; \tfrac59\right) \approx (-0.44,\; -0.11,\; 0.56).$$

Arm 3 goes up, arm 1 goes down, and arm 2 (slightly below average) goes down a little. Keep these three numbers in mind.

The trouble is that a real agent does not know $q_*$. It can only pull arms and see noisy rewards. So the question becomes: can you get $(-0.44, -0.11, 0.56)$, on average, from samples alone?

The log-derivative trick

The identity

Write the objective generically as an expectation of some score $f(x)$ over samples $x$ from a distribution $p_\theta$ that you control:

$$J(\theta) = \mathbb{E}_{x \sim p_\theta}[f(x)] = \sum_x p_\theta(x)\, f(x).$$

The parameters sit inside the distribution, not inside $f$. Differentiate, and use $\nabla p = p\, \nabla \log p$ (that is just the chain rule for $\log$, read backwards):

$$\nabla_\theta J = \sum_x f(x)\, \nabla_\theta p_\theta(x) = \sum_x p_\theta(x)\, f(x)\, \nabla_\theta \log p_\theta(x) = \mathbb{E}_{x \sim p_\theta}\big[f(x)\, \nabla_\theta \log p_\theta(x)\big].$$

That last line is the whole trick. The gradient of an expectation becomes an expectation of something you can compute for each sample. Draw $x$ from your own policy, compute $f(x)\,\nabla_\theta \log p_\theta(x)$, and average. You never differentiate $f$, so $f$ can be a black box, discontinuous, or a 0/1 pass/fail signal.

The vector $\nabla_\theta \log p_\theta(x)$ is called the score function (statisticians’ name, from maximum likelihood), and this estimator is called the score-function or likelihood-ratio estimator. In RL it goes by REINFORCE.

The score of a softmax policy

For the softmax, $\log \pi_\theta(a) = \theta_a - \log \sum_b e^{\theta_b}$. Differentiating with respect to $\theta_k$:

$$\frac{\partial \log \pi_\theta(a)}{\partial \theta_k} = \mathbb{1}[a = k] - \pi_\theta(k), \qquad\text{that is,}\qquad \nabla_\theta \log \pi_\theta(a) = e_a - \pi_\theta,$$

where $e_a$ is the one-hot vector for action $a$. There is a nice picture here. Both $e_a$ and $\pi_\theta$ are points of the probability simplex (the triangle of all distributions over 3 arms). So the score of action $a$ is the arrow from where your policy is now to the corner “always play $a$”. The demo below draws these three arrows as dashed lines in the triangle. (The components of $e_a - \pi$ sum to zero, so the score, and every gradient in this chapter, lives in a two-dimensional plane. That is what lets the demo draw it flat.)

Back to the worked example

At the uniform policy the three possible scores are

$$e_1 - \pi = \left(\tfrac23, -\tfrac13, -\tfrac13\right),\quad e_2 - \pi = \left(-\tfrac13, \tfrac23, -\tfrac13\right),\quad e_3 - \pi = \left(-\tfrac13, -\tfrac13, \tfrac23\right).$$

Suppose rewards are noiseless for now, so arm $a$ always pays $q_*(a)$. A single pull gives one of three possible gradient estimates $\hat g = R \cdot (e_a - \pi)$:

You pull	Probability	Reward $R$	Estimate $\hat g = R\,(e_a - \pi)$
arm 1	1/3	1	$(\tfrac23, -\tfrac13, -\tfrac13)$
arm 2	1/3	2	$(-\tfrac23, \tfrac43, -\tfrac23)$
arm 3	1/3	4	$(-\tfrac43, -\tfrac43, \tfrac83)$

Every single estimate is wrong; none of them is $(-\tfrac49, -\tfrac19, \tfrac59)$. Average them with weights 1/3:

$$\tfrac13\left[\left(\tfrac23 - \tfrac23 - \tfrac43\right),\; \left(-\tfrac13 + \tfrac43 - \tfrac43\right),\; \left(-\tfrac13 - \tfrac23 + \tfrac83\right)\right] = \left(-\tfrac49,\; -\tfrac19,\; \tfrac59\right).$$

Exactly the true gradient. The estimator is unbiased: right on average, noisy every time. Look at how noisy: the true gradient has length about 0.72, while a single estimate has length between 0.82 and 3.3. How to tame that noise is the second half of this chapter.

Notice also something odd about the arm-1 row. Arm 1 is the worst arm, yet pulling it gives an estimate that raises arm 1’s logit, because its reward (1) is positive. Every pull pushes the arm you pulled up; good arms just get pushed harder. Hold that thought until the section on baselines.

The score of a Gaussian policy

For $\pi(a) = \mathcal{N}(a; \mu, \sigma^2)$ with parameter $\mu$,

$$\nabla_\mu \log \pi(a) = \frac{a - \mu}{\sigma^2}.$$

Sample $a = \mu + 0.5$ and get a good reward, and the update moves $\mu$ toward $a$. Get a bad reward, and $\mu$ moves away, provided “good” and “bad” are measured relative to some reference level, which is again the baseline story. Exercise 7.3 does the width $\sigma$ as well.

From one action to whole episodes: the policy gradient theorem

A bandit has a single action. An episode has many, and between them the environment moves the state around according to dynamics you do not know. Apply the same trick with $x = \tau$, the whole trajectory, and $f(\tau) = G_0$:

$$\nabla_\theta J = \mathbb{E}_{\tau \sim \pi_\theta}\big[G_0\, \nabla_\theta \log p_\theta(\tau)\big].$$

What is the probability of a trajectory? Start state, then alternately the policy picks an action and the environment picks the next state:

$$p_\theta(\tau) = p(s_0)\prod_{t=0}^{T-1} \pi_\theta(a_t \mid s_t)\; P(s_{t+1} \mid s_t, a_t).$$

Take the log, and the product becomes a sum:

$$\log p_\theta(\tau) = \log p(s_0) + \sum_{t} \log \pi_\theta(a_t \mid s_t) + \sum_t \log P(s_{t+1} \mid s_t, a_t).$$

Now take $\nabla_\theta$. The start distribution and the transition probabilities do not depend on $\theta$, so their gradients are zero. Only the policy terms survive:

$$\nabla_\theta \log p_\theta(\tau) = \sum_{t} \nabla_\theta \log \pi_\theta(a_t \mid s_t).$$

Putting it together gives the policy gradient theorem in its trajectory form:

$$\nabla_\theta J(\theta) = \mathbb{E}_{\tau \sim \pi_\theta}\!\left[\,G_0 \sum_{t=0}^{T-1} \nabla_\theta \log \pi_\theta(a_t \mid s_t)\right].$$

Stop and appreciate what just happened. The environment’s dynamics dropped out. You need to sample from $P$ (run the episode), but you never need to know $P$, write it down, or differentiate it. The physics simulator, the opponent, the human rater, the code interpreter: all black boxes. The only things you differentiate are your own log-probabilities, which you control. This is why policy gradients apply to almost anything you can run, and why they work for language models whose “environment” is a reward model or a verifier.

∂ The policy gradient theorem, carefully

Setup: episodes of length at most $T$ (or discounted, so returns are bounded), a policy $\pi_\theta(a \mid s)$ differentiable in $\theta$ for every $(s, a)$, and rewards $r_t$ drawn from some distribution depending on $(s_t, a_t, s_{t+1})$ but not on $\theta$. Write $R(\tau) = G_0 = \sum_t \gamma^t r_t$.

1. The objective is an integral over trajectories (sums for discrete parts):

$$J(\theta) = \int p_\theta(\tau)\, R(\tau)\, d\tau .$$

2. Under the conditions above (bounded $R$, finite horizon, smooth $\pi_\theta$), you can move the gradient inside the integral (dominated convergence):

$$\nabla_\theta J = \int \nabla_\theta p_\theta(\tau)\, R(\tau)\, d\tau .$$

Note that $R(\tau)$ itself has no $\theta$ in it: once the trajectory is fixed, its return is fixed.

3. Wherever $p_\theta(\tau) > 0$, $\nabla_\theta p_\theta(\tau) = p_\theta(\tau)\,\nabla_\theta \log p_\theta(\tau)$. Where $p_\theta(\tau) = 0$ in a neighbourhood, the gradient is zero too. So

$$\nabla_\theta J = \int p_\theta(\tau)\, \nabla_\theta \log p_\theta(\tau)\, R(\tau)\, d\tau = \mathbb{E}_{\tau \sim \pi_\theta}\big[R(\tau)\,\nabla_\theta \log p_\theta(\tau)\big].$$

4. The trajectory density factorizes. Including the reward densities $p(r_t \mid s_t, a_t, s_{t+1})$:

$$p_\theta(\tau) = p(s_0) \prod_{t} \pi_\theta(a_t \mid s_t)\, P(s_{t+1} \mid s_t, a_t)\, p(r_t \mid s_t, a_t, s_{t+1}).$$

5. In $\log p_\theta(\tau)$, every factor except $\pi_\theta$ is constant in $\theta$, so

$$\nabla_\theta \log p_\theta(\tau) = \sum_t \nabla_\theta \log \pi_\theta(a_t \mid s_t).$$

6. Substituting into step 3 gives the theorem. $\blacksquare$

What the proof needs: you can sample from the environment, the policy is differentiable and gives positive probability to the actions it takes, and the environment does not depend on $\theta$ except through your actions. What it does not need: a model, differentiable rewards, or differentiable dynamics.

REINFORCE

Replace the expectation by an average over $N$ episodes and you have an algorithm. Williams (1992) named it REINFORCE (an acronym: “REward Increment = Nonnegative Factor × Offset Reinforcement × Characteristic Eligibility”; nobody remembers that, everybody remembers the name).

REINFORCE (vanilla, with the whole-episode return)

initialise policy parameters θ
repeat:
    1. run the current policy π_θ for N episodes; record (s_t, a_t, r_t) for each
    2. for each episode, compute its return G_0 = Σ_t γ^t r_t
    3. ĝ = (1/N) Σ_episodes Σ_t  G_0 · ∇_θ log π_θ(a_t | s_t)
    4. θ ← θ + α ĝ
    5. throw the episodes away (they came from the old θ)

Step 5 is not a detail. The expectation in the theorem is over trajectories from the current $\pi_\theta$. After one update, the old episodes come from a different policy, and reusing them naively gives a biased gradient. REINFORCE is on-policy: every update needs fresh samples. Reusing old samples correctly needs importance ratios, which is the subject of Chapter 9.

How you actually code it: a surrogate loss

In PyTorch you don’t assemble $\hat g$ by hand. You build a scalar whose gradient is $\hat g$ and call backward():

logits = policy(states)                                  # (N, n_actions)
logp = torch.log_softmax(logits, -1)                     # log π(· | s_t)
logp_taken = logp.gather(1, actions[:, None]).squeeze(1) # log π(a_t | s_t)
loss = -(weights.detach() * logp_taken).mean()           # weights = G_0, G_t, or G_t - b
optimizer.zero_grad(); loss.backward(); optimizer.step()

The gradient of $-\frac1N\sum_t w_t \log \pi_\theta(a_t \mid s_t)$ with the weights held constant is $-\frac1N \sum_t w_t \nabla_\theta\log\pi_\theta(a_t \mid s_t)$, exactly minus the REINFORCE estimate. Gradient descent on this loss is gradient ascent on $J$.

Reward-to-go: don’t blame an action for the past

Vanilla REINFORCE multiplies the score of every action in an episode by the same number $G_0$. Think about what that means. An action at step 50 gets credit for rewards collected at steps 0 to 49, before it was taken. It could not possibly have caused them. Those rewards are pure noise as far as the step-50 action is concerned: on average they contribute nothing, but in any one episode they jostle the estimate.

The fix is to weight each action only by the rewards that come after it, its reward-to-go:

$$\nabla_\theta J = \mathbb{E}_{\tau \sim \pi_\theta}\!\left[\sum_{t} \nabla_\theta \log \pi_\theta(a_t \mid s_t)\, \sum_{t' \ge t} \gamma^{t'} r_{t'}\right].$$

This is still exactly the gradient: no approximation, just fewer noise terms. The reason is the fact you met in the first quiz. Conditioned on everything up to $s_t$, the score of $a_t$ has mean zero, so multiplying it by anything already determined (like an earlier reward) gives zero on average.

∂ Why dropping past rewards keeps the estimator unbiased

Write $\psi_t = \nabla_\theta \log \pi_\theta(a_t \mid s_t)$ and $h_t = (s_0, a_0, r_0, \dots, s_t)$, the history up to state $s_t$ but not including $a_t$.

1. The score has zero mean given the state:

$$\mathbb{E}[\psi_t \mid h_t] = \sum_a \pi_\theta(a \mid s_t)\, \nabla_\theta \log \pi_\theta(a \mid s_t) = \sum_a \nabla_\theta \pi_\theta(a \mid s_t) = \nabla_\theta \sum_a \pi_\theta(a \mid s_t) = \nabla_\theta 1 = 0.$$

(This uses the Markov property: given $h_t$, the action is drawn from $\pi_\theta(\cdot \mid s_t)$.)

2. Take a reward from the past, $r_{t'}$ with $t' < t$. It is a function of $h_t$. By the tower rule (the law of total expectation, $\mathbb{E}[XY] = \mathbb{E}[X\,\mathbb{E}[Y \mid h]]$ when $X$ is fixed by $h$),

$$\mathbb{E}[r_{t'}\, \psi_t] = \mathbb{E}\big[\, r_{t'}\, \mathbb{E}[\psi_t \mid h_t] \,\big] = \mathbb{E}[r_{t'} \cdot 0] = 0.$$

3. Expand the vanilla estimator $G_0 \sum_t \psi_t = \sum_t \sum_{t'} \gamma^{t'} r_{t'}\, \psi_t$. All terms with $t' < t$ have expectation zero by step 2. Dropping them leaves

$$\nabla_\theta J = \mathbb{E}\Big[\sum_t \psi_t \sum_{t' \ge t} \gamma^{t'} r_{t'}\Big] = \mathbb{E}\Big[\sum_t \gamma^{t}\, G_t\, \psi_t\Big], \qquad G_t = \sum_{k \ge 0} \gamma^{k} r_{t+k}. \qquad \blacksquare$$

The terms we dropped have zero mean but nonzero variance. Removing a zero-mean term $Y$ from $X + Y$ changes the variance by $-\mathrm{Var}(Y) - 2\,\mathrm{Cov}(X, Y)$, which is a reduction unless $Y$ is negatively correlated with the kept terms strongly enough, so this is not a theorem. In practice, reward-to-go almost always reduces the variance a lot.

So the improved estimator is

$$\widehat{\nabla J} = \frac1N \sum_{\text{episodes}} \sum_{t} G_t\, \nabla_\theta \log \pi_\theta(a_t \mid s_t).$$

On CartPole, which you will train in the lab, this one change takes REINFORCE from barely learning to balancing the pole for hundreds of steps in most runs, with the same number of episodes.

Baselines: same gradient, much less noise

The problem: everything gets pushed up

Go back to the bandit and shift every reward by $+10$, so the arms pay 11, 12 and 14. Nothing about the problem has changed: arm 3 is still best by the same margins, and $\nabla J$ is still $(-\tfrac49, -\tfrac19, \tfrac59)$ (the offset cancels in $q_*(k) - J$). But now every single estimate $R\,(e_a - \pi)$ has $R \ge 11$. Whichever arm you pull, its logit gets a big push up. The true gradient emerges only as a small difference between large, nearly cancelling pushes. In CartPole every reward is $+1$, so every action in every episode gets pushed up; the good actions are only distinguished by being pushed up slightly more.

You can measure the damage with the total variance of one estimate, $\mathbb{E}\lVert \hat g - \nabla J \rVert^2$ (the trace of its covariance matrix: the mean squared distance of a dot from the true gradient). The demo below calls it the spread. For the uniform policy, $\lVert e_a - \pi \rVert^2 = \tfrac49 + \tfrac19 + \tfrac19 = \tfrac23$ for every arm, so

$$\mathbb{E}\lVert \hat g \rVert^2 = \tfrac23\, \mathbb{E}[R^2], \qquad \mathrm{Var} = \tfrac23\,\mathbb{E}[R^2] - \lVert \nabla J \rVert^2, \qquad \lVert \nabla J \rVert^2 = \tfrac{16 + 1 + 25}{81} \approx 0.52.$$

With rewards 1, 2, 4: $\mathbb{E}[R^2] = 7$, so the variance is $\tfrac{14}{3} - 0.52 \approx 4.15$. With rewards 11, 12, 14: $\mathbb{E}[R^2] = \tfrac{461}{3}$, and the variance is about 102. Same gradient, 25 times the noise.

The fix: subtract a baseline

Measure each reward relative to a reference level $b$, the baseline:

$$\widehat{\nabla J} = \frac1N \sum_{\text{episodes}} \sum_t \big(G_t - b(s_t)\big)\, \nabla_\theta \log \pi_\theta(a_t \mid s_t).$$

Now an action is pushed up only if it did better than $b$, and pushed down if it did worse. Is this still the right gradient? Yes, by the same zero-mean fact, as long as $b$ does not depend on the action $a_t$:

$$\mathbb{E}\big[b(s_t)\, \nabla_\theta \log \pi_\theta(a_t \mid s_t)\big] = \mathbb{E}\Big[b(s_t)\, \underbrace{\mathbb{E}\big[\nabla_\theta \log \pi_\theta(a_t \mid s_t) \mid s_t\big]}_{=\,0}\Big] = 0.$$

The baseline can depend on the state, on the time step, on anything from the past, even on other independent episodes. It must not depend on the action whose score it multiplies; Exercise 7.4 shows what goes wrong if it does.

Back to the worked example, with $b = J = \tfrac73$ (rewards 1, 2, 4):

You pull	$R - b$	Estimate $(R - b)(e_a - \pi)$
arm 1	$-\tfrac43$	$(-\tfrac89, \tfrac49, \tfrac49)$
arm 2	$-\tfrac13$	$(\tfrac19, -\tfrac29, \tfrac19)$
arm 3	$+\tfrac53$	$(-\tfrac59, -\tfrac59, \tfrac{10}9)$

The average is still $(-\tfrac49, -\tfrac19, \tfrac59)$; check it. But now pulling arm 1 pushes arm 1 down, as it should. The variance is $\tfrac23 \mathrm{Var}(R) - 0.52 = \tfrac23 \cdot \tfrac{14}{9} - 0.52 \approx 0.52$, down from 4.15, an 8-fold reduction. With the $+10$ offset, $b = J = 12.33$ gives the same $0.52$, down from 102. The baseline makes the estimator blind to the offset, which is what you want, because the offset never carried information.

Which baseline is best?

A natural choice is the average return, $b \approx J$, or, when there are states, $b(s_t) \approx V^\pi(s_t)$, the expected return from $s_t$. Then $G_t - V^\pi(s_t)$ is an estimate of the advantage $A^\pi(s_t, a_t) = Q^\pi(s_t, a_t) - V^\pi(s_t)$: how much better this action turned out than the policy’s average from that state. Learning $V$ well enough to use as a baseline (and then using it for more than that) is what Chapter 8 is about.

$J$ is close to the best constant, but not exactly the best. Minimizing the variance over $b$ gives a score-weighted average of the rewards.

∂ The minimum-variance constant baseline

Consider a single-sample estimate $\hat g(b) = (R - b)\,\psi$, where $\psi = \nabla_\theta \log \pi_\theta(a)$ and $R$ is the reward (or return). We want the constant $b$ minimizing the total variance $\mathbb{E}\lVert \hat g(b) - \nabla J \rVert^2$.

1. Since $\mathbb{E}[\hat g(b)] = \nabla J$ for every $b$, the variance is

$$\mathbb{E}\lVert \hat g(b)\rVert^2 - \lVert \nabla J\rVert^2 = \mathbb{E}\big[(R - b)^2 \lVert\psi\rVert^2\big] - \lVert \nabla J\rVert^2 .$$

2. This is a quadratic in $b$. Differentiate and set to zero:

$$\frac{d}{db}\,\mathbb{E}\big[(R - b)^2 \lVert\psi\rVert^2\big] = -2\,\mathbb{E}\big[(R - b)\lVert\psi\rVert^2\big] = 0 .$$

3. Solve:

$$b^* = \frac{\mathbb{E}\big[R\,\lVert \psi \rVert^2\big]}{\mathbb{E}\big[\lVert \psi \rVert^2\big]} .$$

The second derivative is $2\,\mathbb{E}\lVert\psi\rVert^2 > 0$, so this is a minimum.

So $b^*$ is an average of the rewards, weighted by how large the score was on each sample. If all scores have the same size, $b^* = \mathbb{E}[R] = J$. That happens at the uniform softmax policy (every $\lVert e_a - \pi\rVert^2 = \tfrac23$), so in the worked example $b^* = J = \tfrac73$ exactly. Away from uniform, rarely chosen actions have larger scores and count more; Exercise 7.6 works out a two-arm case where the answer is surprising.

If you allow a different baseline for each parameter component $k$, the same computation gives $b_k^* = \mathbb{E}[R\,\psi_k^2] / \mathbb{E}[\psi_k^2]$. In practice nobody uses these exactly: a learned $V(s)$ is close enough, and it adapts to the state.

See it

The demo below is the worked example, live. The first panel is the probability triangle. The large dot is the current policy (tap or drag to move it), the dashed lines run from it to the three corners (the three possible scores $e_a - \pi$), and the orange arrow shows which way a small step along the true gradient moves the policy. The second panel shows single-sample gradient estimates as dots, colored by the arm that was pulled, together with the true gradient (solid orange arrow) and the mean of the dots so far (dashed arrow). It is drawn in the same orientation as the triangle, so “raise arm 3” points up. Each dot lies on the dashed spoke of the arm that was pulled, parallel to that arm’s dashed line in the triangle, at distance $\lvert R - b\rvert\,\lVert e_a - \pi\rVert$ from the centre. The strip plot below the baseline slider shows the exact spread of one estimate as a function of $b$ (log scale), with markers at $J$ and $b^*$; the orange dot is your current $b$.

Policy π (tap or drag to move it)

Gradient estimates (one dot per sample)

π J(θ)

true ∇J

mean of samples

spread (total variance) of one estimate exact , measured

Baseline b =

b = 0b = J(θ)b = b* (least spread)

Add c to every reward: c = Reward noise σ =

+1 sample+25ClearZoom outReset

What to try: (1) Press +25 a few times and watch the dashed mean arrow settle onto the solid true gradient: unbiased, but noisy. (2) Press b = J(θ): every dot slides along its own spoke toward the centre. The true gradient does not move; the spread drops from 4.15 to 0.52. (3) Set c = 10 and b = 0. The dots fly off the chart (press Zoom out to find them), and the spread is about 100. Now press b = J(θ). (4) Drag the policy near a corner and compare the b = J and b = b* buttons. They differ; Exercise 7.6 explains why.

REINFORCE in action, and why it is not enough

Here is the full algorithm as practitioners usually run it.

REINFORCE with reward-to-go and a baseline

initialise policy π_θ and a baseline (a running average, or a small value network V_φ)
repeat:
    1. run π_θ for a batch of episodes
    2. for every step t of every episode:
           G_t = r_t + γ G_{t+1}                    (backward pass, G_T = 0)
           Â_t = G_t − b(s_t)                       (advantage estimate)
    3. loss = −mean_t [ Â_t · log π_θ(a_t | s_t) ]   (Â_t held constant)
    4. one gradient step on the loss
    5. update the baseline toward the new returns (e.g. regress V_φ(s_t) on G_t)

The demo below runs it on the 3-armed bandit, one pull per update: 12 runs without a baseline (orange) and 12 runs with a running-average baseline (blue), from the same start and with the same random seeds. The paths show how each run’s policy moves across the probability triangle.

Policy paths (tap to choose the start)

Expected reward Σ π(a) q(a) (offset c not included)

no baseline: weight Rrunning-average baseline: weight R − b

step

no baseline mean J runs not yet committed to arm 3 (π₃ < 0.5):

with baseline mean J runs not yet committed to arm 3:

Add c to every reward: c = (rewards ≈ , plus noise σ = 1)Step size α =

Play+10 stepsResetSpeed 1×

Each step pulls one arm and updates θ by α(R − b)(e_a − π), with reward noise σ = 1. Orange runs use b = 0; blue runs use b = the running average of the rewards seen so far (it starts at 0). What to try: press Play at c = 5 (rewards about 6, 7, 9). Then try c = 0, and c = 15. Tap near the arm-2 corner to start from a bad policy. Raise α. Count the orange runs that end up stuck in the wrong corner.

Two lessons are hiding in that demo, and both foreshadow the rest of Part II.

Unbiased is not the same as safe. Without a baseline, whichever arm you happen to pull gets reinforced, so early luck compounds: the policy drifts toward an arbitrary corner. Once it is nearly deterministic on arm 2, it almost never tries arm 3 again, and the true gradient there is nearly zero anyway ($\pi(k)(q_*(k) - J)$ is tiny when $\pi(3)$ is tiny). The run is stuck. The estimator was correct on average at every single step; the finite step size plus high variance did the damage. In the demo with $c = 10$ and $\alpha = 0.1$, 5 of the 12 no-baseline runs end up committed to a wrong arm after 400 steps (other random seeds give 3 to 5), and none of the 12 baseline runs do.

The fundamental costs of REINFORCE are:

• High variance. Even with reward-to-go and a baseline, one episode’s return is a noisy measure of how good each action was. Noise forces small learning rates and many samples. A learned critic cuts it further (Chapter 8).
• Sample inefficiency. On-policy: every batch is used for one gradient step, then thrown away. Squeezing several steps out of one batch safely needs importance ratios and a trust region (Chapter 9).
• Fragile steps. One overly large step can make the policy nearly deterministic, and then it stops exploring and can’t collect the data it would need to recover. The demo’s stuck runs are a small version of this. In LLM training the same thing shows up as entropy collapse (Chapter 19).
• Credit assignment. Every action in a long episode shares the noise of the whole return-to-go. The longer the episode, the worse.

Two practical tricks you will see in every codebase. Batching: average over several episodes per update. Return normalization: subtract the batch mean of the weights and divide by their standard deviation. Subtracting the mean is a baseline (estimated from the same batch, so it is very slightly biased); dividing by the standard deviation just rescales the step size per batch.

From CartPole to language models

Everything above carries over to a language model with almost no change. Take a prompt $x$. The model samples a response $y = (y_1, \dots, y_{|y|})$ one token at a time. The probability of the whole response is the product of the next-token probabilities, so its log-probability is a sum of token log-probabilities, and so is its gradient:

$$\log \pi_\theta(y \mid x) = \sum_{t=1}^{|y|} \log \pi_\theta(y_t \mid x, y_{\lt t}), \qquad \nabla_\theta \log \pi_\theta(y \mid x) = \sum_{t} \nabla_\theta \log \pi_\theta(y_t \mid x, y_{\lt t}).$$

That is the trajectory form of the policy gradient theorem, with state $(x, y_{\lt t})$ and action $y_t$. The “environment” is trivial (appending a token is deterministic), and usually the only reward is a single number $r(x, y)$ at the end: a reward model’s score, or 1 if the final answer is correct. With $\gamma = 1$ and one terminal reward, every token’s reward-to-go is just $r(x, y)$. So REINFORCE for a language model is

$$\widehat{\nabla J} = \frac1N \sum_{i=1}^N \big(r(x, y^{(i)}) - b\big) \sum_{t} \nabla_\theta \log \pi_\theta\big(y^{(i)}_t \mid x, y^{(i)}_{\lt t}\big).$$

Every token in a good response is pushed up by the same amount; every token in a bad response is pushed down. And the variance problem is acute: one scalar reward is spread over hundreds or thousands of tokens.

So everything hinges on the baseline. Sample a group of $G = 4$ responses to the same prompt and suppose they get rewards $1, 0, 0, 1$. Using the group mean as the baseline, $b = 0.5$, the weights are $+0.5, -0.5, -0.5, +0.5$: correct answers pushed up, wrong ones pushed down. Using the leave-one-out mean (the average of the other three responses) for each response instead, the weights are $1 - \tfrac13 = +\tfrac23$ for the correct ones and $0 - \tfrac23 = -\tfrac23$ for the wrong ones. That is RLOO. If you divide the group-mean version by the group’s standard deviation (0.5 here), you get $\pm 1$, and that is the advantage used by GRPO. (Some implementations use the sample standard deviation, $0.58$, which gives $\pm 0.87$; Chapter 15 discusses such details.)

Summary

• A policy gradient method parameterizes the policy $\pi_\theta$ directly and does gradient ascent on the expected return $J(\theta)$. It handles continuous actions, changes smoothly, can represent stochastic optimal policies, and fits language models, which are already policies.
• The log-derivative trick, $\nabla_\theta \mathbb{E}_{p_\theta}[f] = \mathbb{E}_{p_\theta}[f \,\nabla_\theta \log p_\theta]$, turns the gradient of an expectation into an expectation you can sample. $f$ need not be differentiable.
• For trajectories the dynamics drop out: $\nabla J = \mathbb{E}[\sum_t \nabla_\theta \log \pi_\theta(a_t \mid s_t)\cdot(\text{return})]$. You need samples from the environment, never its model.
• For a softmax policy the score is $e_a - \pi_\theta$: the arrow from the current policy toward always playing $a$.
• REINFORCE is weighted maximum likelihood on your own samples. In code it is the surrogate loss $-\text{mean}(w_t \log \pi_\theta(a_t \mid s_t))$ with the weights detached. It is on-policy.
• Reward-to-go and baselines remove terms with zero mean and nonzero variance. Both keep the estimator unbiased because the score has zero mean given the state. A baseline must not depend on the action.
• A good baseline is about the expected return ($V(s)$, leading to the advantage); the exact minimum-variance constant is a score-weighted average of returns. A baseline is a control variate.
• REINFORCE’s weaknesses (variance, one step per batch, fragile steps) motivate actor-critic (Chapter 8) and PPO (Chapter 9). Its strengths (no critic, black-box rewards) are why RLOO and GRPO are REINFORCE at heart.

Exercises

Exercise 7.1Where argmax strugglesConcept●○○

For each problem, say what makes “learn $Q$, act greedily” awkward or wrong, and what a parameterized stochastic policy does instead.

1. A robot arm that outputs three joint torques, each a real number in $[-1, 1]$.
2. Rock-paper-scissors against an opponent that learns to exploit any pattern in your play.
3. A language model choosing the next token from a vocabulary of 100,000.

Solution

1. Acting greedily means solving $\max_{a \in [-1,1]^3} Q(s, a)$ at every step, a nonconvex optimization over a continuum, inside the training loop. A Gaussian policy outputs a mean and width and samples an action in one forward pass, and its log-probability is differentiable in $\theta$.
2. Every deterministic policy, and every greedy policy, is exploitable: once the opponent learns what you will play, you lose every round. The optimal (Nash) strategy is uniformly random, which a softmax policy represents exactly with equal logits, and a greedy-over-$Q$ policy cannot represent at all (ε-greedy adds randomness, but not a chosen amount).
3. The action space is huge and the “state” (the text so far) is never seen twice, so learning accurate $Q$-values for every token is hopeless. But the model already is a softmax policy over tokens, pretrained to be very good. Policy gradients improve it in place.

I've done this one

Exercise 7.2The score of a softmax, with temperaturePen & paper●●○

Let $\pi_\theta(a) = e^{\theta_a/\tau} / \sum_b e^{\theta_b/\tau}$ with temperature $\tau > 0$.

1. Show that $\partial \log \pi_\theta(a) / \partial \theta_k = (\mathbb{1}[a = k] - \pi_\theta(k))/\tau$.
2. Show that the components of the score sum to zero, and explain why that had to be true.
3. Show that $\sum_a \pi_\theta(a)\, \nabla_\theta \log \pi_\theta(a) = 0$.
4. At $\tau = 1$ and $\theta = (0, 0, 0)$, compute the score of action 3. What happens to its size as $\tau \to 0$, and what does that suggest about learning with a very low temperature?

Hint

Write $\log \pi_\theta(a) = \theta_a/\tau - \log Z$ with $Z = \sum_b e^{\theta_b/\tau}$, and remember that $\partial \log Z / \partial \theta_k$ is a probability.

Solution

1. $\log \pi_\theta(a) = \theta_a/\tau - \log Z$. Then $\partial(\theta_a/\tau)/\partial\theta_k = \mathbb{1}[a=k]/\tau$, and $\partial \log Z / \partial \theta_k = \frac{1}{Z} e^{\theta_k/\tau}\cdot\frac1\tau = \pi_\theta(k)/\tau$. Subtract.

2. $\sum_k (\mathbb{1}[a=k] - \pi_\theta(k))/\tau = (1 - 1)/\tau = 0$. It had to be: adding the same constant $c$ to every logit leaves $\pi$ unchanged, so the directional derivative of $\log \pi_\theta(a)$ along $(1, 1, \dots, 1)$ is zero, and that directional derivative is the sum of the components.

3. $\sum_a \pi(a)(\mathbb{1}[a=k] - \pi(k))/\tau = (\pi(k) - \pi(k))/\tau = 0$ for every $k$. (This is the general fact $\mathbb{E}[\nabla\log\pi] = \nabla\sum\pi = 0$.)

4. At $\tau = 1$: $e_3 - \pi = (-\tfrac13, -\tfrac13, \tfrac23)$, with squared length $\tfrac23$. At $\theta = 0$ the policy is uniform for every $\tau$, so the score $(e_3 - \pi)/\tau$ simply grows like $1/\tau$: a lower temperature amplifies every update, like a larger step size. With unequal logits it gets worse. As $\tau \to 0$ the policy becomes nearly deterministic. The likely action’s score goes to zero (it is already certain), while an unlikely action’s score has size of order $1/\tau$, so on the rare occasions it is sampled it produces a huge estimate. Learning at very low temperature combines vanishing gradients for likely actions with rare, enormous ones for unlikely actions: very high variance. This is one reason RL for LLMs usually samples at temperature about 1.

I've done this one

Exercise 7.3The Gaussian scorePen & paper●●○

Let $\pi(a) = \mathcal{N}(a; \mu, \sigma^2)$ with parameters $\mu$ and $\eta = \log \sigma$ (parameterizing by $\log\sigma$ keeps $\sigma$ positive).

1. Compute $\partial \log \pi(a)/\partial \mu$ and $\partial \log \pi(a)/\partial \eta$.
2. Check that both have mean zero under $a \sim \pi$.
3. With $\mu = 0$, $\sigma = 1$, you sample $a = 1.5$ and get return $G = 3$. The baseline is $b = 1$, and the step size is $\alpha = 0.1$. Compute the REINFORCE update of $\mu$ and $\eta$, and say in words what the policy learned.

Hint

$\log \pi(a) = -\frac{(a-\mu)^2}{2}e^{-2\eta} - \eta - \tfrac12\log 2\pi$. For the mean-zero check you need $\mathbb{E}[(a-\mu)^2] = \sigma^2$.

Solution

1. $\partial \log\pi/\partial\mu = (a - \mu)/\sigma^2$. For $\eta$: $\partial/\partial\eta\left[-\tfrac12(a-\mu)^2 e^{-2\eta}\right] = (a-\mu)^2 e^{-2\eta}$, and $\partial(-\eta)/\partial\eta = -1$, so $\partial \log\pi/\partial\eta = (a-\mu)^2/\sigma^2 - 1$.

2. $\mathbb{E}[(a-\mu)/\sigma^2] = 0$ because $\mathbb{E}[a] = \mu$. $\mathbb{E}[(a-\mu)^2/\sigma^2 - 1] = \sigma^2/\sigma^2 - 1 = 0$.

3. $G - b = 2$. Scores: $\partial_\mu = 1.5/1 = 1.5$, and $\partial_\eta = 1.5^2 - 1 = 1.25$. Updates: $\Delta\mu = 0.1 \cdot 2 \cdot 1.5 = 0.3$, and $\Delta\eta = 0.1 \cdot 2 \cdot 1.25 = 0.25$, so $\sigma$ becomes $e^{0.25} \approx 1.28$. The mean moves toward the action that beat the baseline. The width grows, because a better-than-expected outcome came from far out in the tail (more than one standard deviation away), which suggests exploring more widely. A good outcome from close to the mean ($|a - \mu| < \sigma$) would have shrunk $\sigma$ instead.

I've done this one

Exercise 7.4Baselines may depend on the state, not on the actionPen & paper●●○

1. Prove that for any function $b(s)$, $\mathbb{E}_{a \sim \pi_\theta(\cdot \mid s)}\big[b(s)\,\nabla_\theta \log \pi_\theta(a \mid s)\big] = 0$.
2. Now cheat: in the 3-armed worked example (uniform policy, noiseless rewards 1, 2, 4), use an action-dependent “baseline” $b(a) = q_*(a)$. Compute the expected estimate $\mathbb{E}[(R - b(a))\,(e_a - \pi)]$. Compare with $\nabla J$.
3. Use $b(a) = 2$ for arms 1 and 2 and $b(a) = 0$ for arm 3. Compute the expected estimate and show it is biased.

Hint

In part 1, $b(s)$ is a constant for the sum over $a$, and $\pi\,\nabla\log\pi = \nabla\pi$. In parts 2 and 3, write out the three possible estimates and average them with weight 1/3.

Solution

1. $\sum_a \pi_\theta(a\mid s)\, b(s)\, \nabla_\theta \log\pi_\theta(a \mid s) = b(s)\sum_a \nabla_\theta \pi_\theta(a\mid s) = b(s)\,\nabla_\theta \sum_a \pi_\theta(a \mid s) = b(s)\,\nabla_\theta 1 = 0$. Pulling $b(s)$ out of the sum is the step that needs $b$ not to depend on $a$.

2. With noiseless rewards $R = q_*(a)$, so $R - b(a) = 0$ for every arm: every estimate is exactly zero, and so is its expectation. The true gradient is $(-\tfrac49, -\tfrac19, \tfrac59)$. The “baseline” removed all of the signal, because the signal is the difference between arms. (With an action-dependent baseline you can correct the bias by adding back $\mathbb{E}_a[b(a)\nabla\log\pi(a)]$ analytically; that is the idea behind some “action-dependent control variate” methods. Without the correction, it is just wrong.)

3. The estimates are arm 1: $(1 - 2)(\tfrac23, -\tfrac13, -\tfrac13) = (-\tfrac23, \tfrac13, \tfrac13)$; arm 2: $(2 - 2)(\cdots) = 0$; arm 3: $4\,(-\tfrac13, -\tfrac13, \tfrac23) = (-\tfrac43, -\tfrac43, \tfrac83)$. The average is $\tfrac13(-2, -1, 3) = (-\tfrac23, -\tfrac13, 1)$, not $(-\tfrac49, -\tfrac19, \tfrac59)$. The bias is $\mathbb{E}[-b(a)(e_a - \pi)] = -\tfrac13[2(e_1 - \pi) + 2(e_2 - \pi)] = (-\tfrac29, -\tfrac29, \tfrac49)$. Here it happens to point roughly the right way, but in general a biased estimator can converge to the wrong policy.

I've done this one

Exercise 7.5A REINFORCE update by handPen & paper●●○

A tiny MDP has two states, A and B, each with actions left and right. The policy is tabular softmax: each state has its own two logits, all initially 0. One episode is: in A take right, receive $r_0 = 1$, move to B; in B take left, receive $r_1 = 0$, episode ends. Use $\gamma = 0.9$ and step size $\alpha = 0.5$.

1. Compute the update to all four logits using vanilla REINFORCE (every action weighted by $G_0$).
2. Repeat with reward-to-go weights $G_t$.
3. Repeat with reward-to-go and a constant baseline $b = 0.5$.
4. In each case, what is the new probability of left in state B? Which update is the most sensible, and why?

Hint

At a uniform 2-action softmax, the score of right with respect to that state’s logits (left, right) is $(-\tfrac12, \tfrac12)$, and the score of left is $(\tfrac12, -\tfrac12)$. The other state’s logits get zero.

Solution

Returns: $G_0 = r_0 + \gamma r_1 = 1$, and $G_1 = r_1 = 0$.

1. Both actions are weighted by $G_0 = 1$. State A (took right): $\Delta\theta_A = 0.5 \cdot 1 \cdot (-\tfrac12, \tfrac12) = (-0.25, +0.25)$. State B (took left): $\Delta\theta_B = 0.5 \cdot 1 \cdot (\tfrac12, -\tfrac12) = (+0.25, -0.25)$.

2. Step 0 gets $G_0 = 1$, the same as before: $\Delta\theta_A = (-0.25, +0.25)$. Step 1 gets $G_1 = 0$: $\Delta\theta_B = (0, 0)$.

3. Weights $G_0 - b = 0.5$ and $G_1 - b = -0.5$. $\Delta\theta_A = 0.5\cdot 0.5\cdot(-\tfrac12, \tfrac12) = (-0.125, +0.125)$. $\Delta\theta_B = 0.5\cdot(-0.5)\cdot(\tfrac12, -\tfrac12) = (-0.125, +0.125)$.

4. $\pi(\text{left} \mid B) = \sigma(\theta_{B,\text{left}} - \theta_{B,\text{right}})$: (1) $\sigma(0.5) \approx 0.62$; (2) $\sigma(0) = 0.5$; (3) $\sigma(-0.25) \approx 0.44$.

Vanilla REINFORCE made left in B more likely, even though it earned nothing: it got credit for the reward collected in A before it was taken. Reward-to-go removes that false credit. With a baseline, earning 0 is below the expected level, so left in B becomes less likely, while right in A, which earned more than expected, is still reinforced. All three are unbiased estimates of the same gradient; the third assigns credit most sensibly. (With the exact $\gamma^t$ factor, step 1’s weight would be $\gamma\,(G_1 - b)$ rather than $G_1 - b$; with $G_1 = 0$ this changes only part 3, to $\Delta\theta_B = (-0.1125, +0.1125)$.)

I've done this one

Exercise 7.6The optimal baseline can be surprisingPen & paper●●●

Take a 2-armed bandit with a softmax over two logits and policy $\pi = (p, 1-p)$. Rewards are noiseless: arm 1 pays $q_1$, arm 2 pays $q_2$.

1. Compute the squared norm of the score $\lVert e_a - \pi\rVert^2$ for each arm.
2. Use the minimum-variance formula $b^* = \mathbb{E}[R\lVert\psi\rVert^2] / \mathbb{E}[\lVert\psi\rVert^2]$ to show that $b^* = (1-p)\,q_1 + p\,q_2$.
3. Compare with $J = p\,q_1 + (1-p)\,q_2$. When are they equal? Explain in words why the optimal baseline weights each arm by the probability of the other arm.
4. With $p = 0.9$, $q_1 = 1$, $q_2 = 0$: compute $J$, $b^*$, and the variance of a single estimate at $b = 0$, $b = J$ and $b = b^*$.

Hint

The score of arm 1 is $(1-p, -(1-p))$ and that of arm 2 is $(-p, p)$. For part 4, use $\mathrm{Var}(b) = \mathbb{E}[(R-b)^2\lVert\psi\rVert^2] - \lVert\nabla J\rVert^2$, and note that $\nabla J = \pi \odot (q - J)$.

Solution

1. $\lVert e_1 - \pi\rVert^2 = 2(1-p)^2$ and $\lVert e_2 - \pi\rVert^2 = 2p^2$.

2. Numerator: $p \cdot q_1 \cdot 2(1-p)^2 + (1-p)\cdot q_2 \cdot 2p^2 = 2p(1-p)\,[(1-p)q_1 + p\,q_2]$. Denominator: $p\cdot 2(1-p)^2 + (1-p)\cdot 2p^2 = 2p(1-p)\,[(1-p) + p] = 2p(1-p)$. The ratio is $b^* = (1-p)\,q_1 + p\,q_2$.

3. They are equal when $p = \tfrac12$ (or $q_1 = q_2$). The variance comes from $(R - b)^2\lVert\psi\rVert^2$. The rare arm is sampled rarely, but when it is, its score is large ($\lVert\psi\rVert^2 = 2p^2$ for arm 2 is near 2 when $p$ is near 1), so its term dominates the variance. The best baseline therefore sits close to the rare arm’s reward, to shrink its huge contributions, even though that makes the common arm’s (small) contributions slightly larger. This is why the demo’s b* button differs from J away from the uniform policy.

4. $J = 0.9$, $b^* = 0.1$. The gradient: $\nabla J = (0.9(1 - 0.9), 0.1(0 - 0.9)) = (0.09, -0.09)$, so $\lVert\nabla J\rVert^2 = 0.0162$. Score norms: arm 1 $2(0.1)^2 = 0.02$, arm 2 $2(0.9)^2 = 1.62$.

   • $b = 0$: $\mathbb{E}[(R-b)^2\lVert\psi\rVert^2] = 0.9\cdot 1\cdot 0.02 + 0.1\cdot 0\cdot 1.62 = 0.018$, so the variance is $0.018 - 0.0162 = 0.0018$.
   • $b = J = 0.9$: $0.9 \cdot 0.01 \cdot 0.02 + 0.1\cdot 0.81\cdot 1.62 = 0.00018 + 0.1312 = 0.1314$; variance $\approx 0.115$.
   • $b = b^* = 0.1$: $0.9\cdot 0.81 \cdot 0.02 + 0.1 \cdot 0.01 \cdot 1.62 = 0.01458 + 0.00162 = 0.0162$; variance $= 0$.

   With two arms and noiseless rewards, $b^*$ makes every estimate exactly equal to $\nabla J$, so the variance is zero. More striking: $b = J$ is 64 times worse than no baseline at all here. “Use the average reward” is a good rule of thumb, not a theorem.

I've done this one

Exercise 7.7What the environment has to provideConcept●●○

The policy gradient theorem never uses $P(s' \mid s, a)$ or $\partial r/\partial a$.

1. What does REINFORCE need from the environment? List everything.
2. Suppose you had a perfect differentiable simulator: you know $P$, and the reward is a smooth function of the state. Describe a different way to get $\nabla_\theta J$, and why it might have lower variance.
3. A coding model is rewarded 1 if its program passes all unit tests and 0 otherwise. Explain why the method in part 2 is unavailable here and REINFORCE is not.

Solution

1. The ability to (a) reset to a start state drawn from $p(s_0)$, (b) execute an action and observe the next state and reward, many times, under the current policy. Plus the ability to compute $\nabla_\theta \log\pi_\theta(a\mid s)$ for your own policy. No model, no reward derivatives, no knowledge of what the dynamics are.

2. With a differentiable simulator and a reparameterized policy (actions $a_t = \mu_\theta(s_t) + \sigma\varepsilon_t$), you can backpropagate the return through the whole rollout: through the reward, through the dynamics $s_{t+1} = f(s_t, a_t)$, into $\theta$. This pathwise gradient uses the derivative of the reward with respect to actions, information the score-function estimator ignores, and so often has much lower variance. (It has its own problems: chaotic dynamics give exploding gradients over long horizons.)

3. The reward is a step function of a discrete output: tokens are discrete, so there is no $\partial y/\partial\theta$ through sampling, and “pass/fail” has zero derivative almost everywhere. There is nothing to backpropagate through. REINFORCE only needs the reward’s value for each sample and the gradient of the model’s own log-probabilities, both of which are available. This is why RL with verifiable rewards (Chapter 15) is built on policy gradients.

I've done this one

Exercise 7.8REINFORCE with 0/1 rewards is fine-tuning on your successesPen & paper●●○

You sample $N$ responses $y^{(1)}, \dots, y^{(N)}$ to a prompt; a verifier gives $r^{(i)} \in \lbrace 0, 1\rbrace$, and $k$ of them are correct.

1. Show that the no-baseline REINFORCE estimate equals $\tfrac{k}{N}$ times the average gradient of $\log\pi_\theta$ over the $k$ correct responses, which is the gradient of supervised fine-tuning on your own correct samples, scaled by the success rate $\hat p = k/N$.
2. Now use the baseline $b = \hat p$. Write the weight on a correct response and on an incorrect one. What does the baseline add that filtered fine-tuning lacks?
3. For $G = 4$ responses with rewards 1, 0, 0, 1, write the four weights with and without the baseline. What happens when all four rewards are 1?

Hint

Split the sum over $i$ into correct and incorrect responses; the incorrect ones have $r^{(i)} = 0$.

Solution

1. $\frac1N\sum_i r^{(i)}\nabla\log\pi_\theta(y^{(i)}) = \frac1N\sum_{i:\,r^{(i)}=1}\nabla\log\pi_\theta(y^{(i)}) = \frac{k}{N}\cdot\frac1k\sum_{i:\,r^{(i)}=1}\nabla\log\pi_\theta(y^{(i)})$. The last factor is the maximum-likelihood gradient on the correct samples. Fine-tuning on your own filtered successes, known as rejection-sampling fine-tuning or expert iteration, and used in STaR (Zelikman et al., 2022), is therefore REINFORCE with 0/1 rewards and no baseline, up to step size.

2. Correct responses get weight $1 - \hat p$; incorrect ones get $-\hat p$. The baseline adds negative examples: wrong answers are actively pushed down, and the push is strongest on prompts the model usually gets right (when $\hat p$ is near 1, a failure is surprising and informative). It also makes easy prompts contribute little positive push, since $1 - \hat p$ is small there: no need to reinforce what is already reliable.

3. No baseline: $(1, 0, 0, 1)$ (divided by $N = 4$). With $b = 0.5$: $(+0.5, -0.5, -0.5, +0.5)$. If all four are correct, $\hat p = 1$ and every weight is $1 - 1 = 0$: the group gives no gradient at all. Without a baseline it would still push all four up. That “all-same group gives zero signal” property is central to GRPO (Chapter 15).

I've done this one

Exercise 7.9Explore the demosExplore the demo●○○

Use the two demos in this chapter.

1. In the gradient-cloud demo at the uniform policy with $c = 0$ and noise $\sigma = 0$, how many distinct dot positions are there, and why? Turn the noise up to 1: what shape does each arm’s set of dots take, and why is it a line segment rather than a blob?
2. Still at the uniform policy, with $\sigma = 0$, read off the spread at $b = 0$, $b = J$ and $b = 10$. Which is the largest?
3. Drag the policy to about $\pi = (0.8, 0.1, 0.1)$. Is the true gradient bigger or smaller than at the uniform policy? Where do $J$ and $b^*$ sit now?
4. In the race demo, with $c = 0$, both methods usually work. Explain why, using what you learned about the sign of $R - b$.
5. In the race demo, set $c = 15$ and tap near the arm-2 corner to start there. Play. What fraction of no-baseline runs escape to arm 3?

Solution

1. Three: with noiseless rewards each arm always gives the same estimate $q_*(a)(e_a - \pi)$. With noise, $R$ varies but the direction $e_a - \pi$ does not, so each arm’s dots spread along a line through the origin, its spoke. A single-sample estimate can only ever point along one of three directions; only averaging produces other directions.

2. About 4.15, 0.52, and $\tfrac23(\mathbb{E}[(R-10)^2]) - 0.52 = \tfrac23\cdot\tfrac{81+64+36}{3} - 0.52 \approx 39.7$. A badly chosen baseline (far above every reward) is as harmful as none: now every pull pushes its own arm down.

3. Smaller: $\nabla J = \pi \odot (q_* - J)$ is small when most probability sits on one arm. Here $J = 0.8 + 0.2 + 0.4 = 1.4$ and $\nabla J = (0.8 \cdot (-0.4),\ 0.1 \cdot 0.6,\ 0.1 \cdot 2.6) = (-0.32, 0.06, 0.26)$, of length about 0.42 instead of 0.72. For $b^*$: the squared score norms are $\lVert e_1 - \pi\rVert^2 = 0.04 + 0.01 + 0.01 = 0.06$ and $\lVert e_2 - \pi\rVert^2 = \lVert e_3 - \pi\rVert^2 = 0.64 + 0.81 + 0.01 = 1.46$, so $b^* = (0.8 \cdot 1 \cdot 0.06 + 0.1 \cdot 2 \cdot 1.46 + 0.1 \cdot 4 \cdot 1.46)/(0.8 \cdot 0.06 + 0.2 \cdot 1.46) = 0.924/0.34 \approx 2.72$. It is well above $J$, pulled toward the rewards of the rarely chosen arms 2 and 3, whose scores are large (as in Exercise 7.6).

4. With rewards 1, 2, 4, every reward is already positive, so without a baseline every pull still pushes its own arm up, but arm 3’s pushes are four times as large as arm 1’s, so the drift toward arm 3 dominates before the policy can collapse. The relative spread of the rewards is large. With $c = 15$, rewards 16, 17, 19 differ by only about 15%, and the “push whatever you pulled” component dominates the signal.

5. Few: starting at about $\pi = (0.05, 0.9, 0.05)$ with the default $\alpha = 0.1$, only about 2 of the 12 no-baseline runs reach arm 3 within 400 steps (the readout then shows about 10 of 12 not committed to arm 3). Arm 3 is rarely sampled, and every pull of arm 2 (reward about 17) reinforces arm 2 strongly. All 12 baseline runs escape, because pulling arm 2 gives $R - b \approx 0$ once $b$ has adapted, while a rare pull of arm 3 gives $R - b \approx +2$. (Exact counts depend on where you tap.)

I've done this one

Exercise 7.10Build REINFORCE: the labCode●●○

Download the lab below and implement, in pg.py:

• Task 1 softmax_score(theta, actions): the score $e_a - \pi$.
• Task 2 exact_gradient(theta, q): $\pi \odot (q - J)$. The test checks the worked example’s $(-\tfrac49, -\tfrac19, \tfrac59)$.
• Task 3 reinforce_estimates(theta, actions, rewards, baseline): $(R - b)(e_a - \pi)$ per sample. Tests check unbiasedness by Monte Carlo and the variance cut from a baseline.
• Task 4 rewards_to_go(rewards, gamma).
• Task 5 reinforce_loss(logits, actions, weights): the surrogate loss with detached weights.

Then run python run.py. It plots the bandit estimator’s variance against the baseline, and CartPole learning curves for three weightings: whole return $G_0$, reward-to-go $G_t$, and $G_t - V(s_t)$ with a small learned value baseline. Before you run it, predict the order of the three curves.

Hint

Task 4 is one backward loop, running = r[t] + gamma * running. In Task 5, torch.log_softmax(logits, -1).gather(1, actions[:, None]) picks out the taken actions’ log-probabilities; don’t forget the minus sign and .detach() on the weights.

Solution

Reference code is in solutions/pg.py. The key lines:

# Task 1: one-hot rows minus pi
return np.eye(len(pi))[actions] - pi[None, :]

# Task 3: scale each row by (R - b)
return (rewards - baseline)[:, None] * softmax_score(theta, actions)

# Task 4: one backward pass
running = 0.0
for t in reversed(range(len(rewards))):
    running = rewards[t] + gamma * running
    out[t] = running

# Task 5: minus the weighted log-likelihood of the taken actions
logp = torch.log_softmax(logits, -1).gather(1, actions[:, None]).squeeze(1)
return -(weights.detach() * logp).mean()

What you should see: the variance curve is a valley with its minimum near $b = J = 12.33$ (about 1 there, versus about 100 at $b = 0$). The minimum is about 1.2 rather than the chapter’s 0.52 because the lab’s rewards have noise $\sigma = 1$, which adds $\tfrac23 \sigma^2 \approx 0.67$ at every $b$ (the noise term $\mathbb{E}[\varepsilon^2]\,\lVert e_a - \pi\rVert^2$). On CartPole, typical final mean returns over 4 seeds are roughly 40 for the whole return, 300 for reward-to-go (with a wide spread across seeds), and close to 500 for the learned baseline. The algorithm, learning rate and number of episodes are the same in all three. Only the variance of the weights differs.

I've done this one

Further reading

• Sutton & Barto (2018), Reinforcement Learning: An Introduction, ch. 13. The textbook treatment: the policy gradient theorem for continuing and episodic tasks, REINFORCE with baseline, and the short-corridor example. Read 13.1 to 13.4.
• OpenAI Spinning Up, Part 3: Intro to Policy Optimization. A short, careful derivation with code, including reward-to-go and the “expected grad-log-prob” lemma used throughout this chapter.
• Williams (1992), Simple statistical gradient-following algorithms for connectionist reinforcement learning. The original REINFORCE paper, including the idea of comparing the reward against a reference level (the baseline).
• Sutton, McAllester, Singh & Mansour (1999), Policy Gradient Methods for Reinforcement Learning with Function Approximation. The state-distribution form of the policy gradient theorem, and the bridge to actor-critic (Chapter 8).
• Greensmith, Bartlett & Baxter (2004), Variance Reduction Techniques for Gradient Estimates in Reinforcement Learning. Baselines and critics as control variates, with optimal baselines. Dense, but the source for the control-variate view.
• Mohamed, Rosca, Figurnov & Mnih (2020), Monte Carlo Gradient Estimation in Machine Learning. A survey comparing score-function, pathwise and measure-valued gradient estimators; read it to see where REINFORCE sits among the alternatives.
• Ahmadian et al. (2024), Back to Basics: Revisiting REINFORCE Style Optimization for Learning from Human Feedback in LLMs. Argues that much of PPO’s machinery is unnecessary for RLHF, and reports that REINFORCE with a leave-one-out baseline (RLOO) does as well or better. The direct link from this chapter to Chapters 15 and 16.
• Nota & Thomas (2020), Is the Policy Gradient a Gradient? The careful story behind the dropped $\gamma^t$ factor, for when you want to know exactly what your code is optimizing.

Next, in Chapter 8, you replace the noisy Monte Carlo return with a learned critic and get the advantage estimates that PPO and RLHF are built on.