Chapter 15 · Section 15.3 · 60 min

Policy gradients and REINFORCE

A pretrained language model already has a policy, and a good one. It writes fluent text, follows most instructions, and gets some fraction of math problems right. What we want from RL is to make the responses that earn high reward a little more likely and the ones that earn low reward a little less likely, without throwing away everything the model already knows. Q-learning, from §15.2, would start over, learning a value for every (prefix, token) pair from scratch and acting greedily on it. That is the wrong tool.

The right tool treats the policy's parameters $\theta$ as the thing to optimize and climbs the expected reward by gradient ascent, exactly as supervised learning descends a loss. Two obstacles seem to stand in the way. The reward is often not differentiable at all (a unit test passes or fails), and the expectation is taken over samples that the policy itself generates, so changing $\theta$ changes the distribution we are averaging over. One identity, the log-derivative trick from §0.4, removes both obstacles. This section derives the policy gradient from it, proves the policy gradient theorem, builds REINFORCE, runs it on a bandit and on a small environment where only a stochastic policy is optimal, and then confronts its great weakness, variance.

The objective

Let $\pi_\theta(a \mid s)$ be a policy with parameters $\theta$ (for us, a softmax over logits produced by a network). Running it in the environment for at most $T$ steps produces a random trajectory $\tau = (s_0, a_0, r_0, \dots, s_{T-1}, a_{T-1}, r_{T-1}, s_T)$ , where as in §15.1 $r_t$ is the reward for taking $a_t$ in $s_t$ . Its probability factors into the start distribution $p_0$ , the policy's choices, and the environment's transitions $P$ ,

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

Write $R(\tau) = \sum_{t=0}^{T-1} \gamma^t r_t$ for the discounted return of the whole trajectory. The objective is the expected return,

J(\theta) = \mathbb{E}_{\tau \sim p_\theta}\big[R(\tau)\big] = \sum_\tau p_\theta(\tau)\, R(\tau),

where the sum runs over all possible trajectories (an integral, if states are continuous). In plain words, $J(\theta)$ is the average total reward the policy earns, and we want to make it larger.

We cannot differentiate $R(\tau)$ with respect to $\theta$ , and we do not need to, because $R(\tau)$ does not depend on $\theta$ at all. A given trajectory has a given return. What depends on $\theta$ is how likely each trajectory is.

The log-derivative trick

Differentiate $J$ and move the gradient inside the sum (fine for finite sums, and for integrals under mild regularity conditions),

\nabla_\theta J(\theta) = \sum_\tau R(\tau)\, \nabla_\theta p_\theta(\tau).

This is not yet an expectation, so we cannot estimate it by sampling. The trick is the chain rule for the logarithm, $\nabla_\theta \log p_\theta = \nabla_\theta p_\theta / p_\theta$ , rearranged as $\nabla_\theta p_\theta = p_\theta \nabla_\theta \log p_\theta$ . Substituting,

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

In plain words, the gradient of the expected return is the average, over trajectories the policy actually produces, of "how good was this trajectory" times "which way in parameter space makes this trajectory more likely". Every sampled trajectory pulls the parameters toward making itself more likely, with a strength set by its return, so high-return trajectories win the tug of war (and a negative return pushes away instead).

Why not differentiate through the sampling instead, as the reparameterization trick of §0.4 does for Gaussians? Because that trick needs the sample to be a differentiable function of $\theta$ and the reward to be differentiable in the sample. A discrete action (a token) is neither, and a unit test is not differentiable at all. The score-function estimator needs neither, which is why it is the workhorse of RL on language models, at the price of the variance we will fight for the rest of the chapter.

The environment drops out

Now take the log of the trajectory probability. The product becomes a sum,

\log p_\theta(\tau) = \log p_0(s_0) + \sum_{t=0}^{T-1} \log \pi_\theta(a_t \mid s_t) + \sum_{t=0}^{T-1} \log P(s_{t+1}\mid s_t, a_t),

and only the middle sum depends on $\theta$ . The start distribution and the dynamics vanish when we take the gradient,

\nabla_\theta \log p_\theta(\tau) = \sum_{t=0}^{T-1} \nabla_\theta \log \pi_\theta(a_t \mid s_t), \qquad\text{so}\qquad \nabla_\theta J(\theta) = \mathbb{E}_{\tau}\Big[ R(\tau) \sum_{t=0}^{T-1} \nabla_\theta \log \pi_\theta(a_t \mid s_t)\Big].

This is the first form of the policy gradient. Read it slowly, because it is remarkable. To estimate the gradient we need only samples of trajectories, their returns, and the gradient of the log-probability of the actions the policy took. We never need $P$ , we never need its derivative, and the reward can come from anywhere, a human, a unit test, a reward model, or a game engine. This is why policy gradient methods are called model-free. The vector $\nabla_\theta \log \pi_\theta(a\mid s)$ is called the score of the action.

The score of a softmax policy

For a softmax policy with logits $\mathbf{z} \in \R^{K}$ over $K$ actions, $\log\pi(a) = z_a - \log\sum_{b} e^{z_b}$ . The gradient of the log-sum-exp is the softmax itself (§0.2), so

\nabla_{\mathbf{z}} \log \pi(a) = \mathbf{e}_a - \boldsymbol{\pi},

where $\mathbf{e}_a$ is the one-hot vector of action $a$ and $\boldsymbol{\pi} = \operatorname{softmax}(\mathbf{z})$ . In plain words, to make action $a$ more likely, raise its logit and lower every logit in proportion to its current probability. This is the same $\mathbf{p} - \mathbf{y}$ as the gradient of cross-entropy in that section, with the sign flipped, which is no coincidence. $\log\pi(a)$ is the negative cross-entropy loss with target $a$ . For a network, the chain rule takes it from the logits to every weight, and autograd does that for us.

Warm-up

The score of a softmax policy

Implement softmax_score(logits, actions) where logits has shape (N, K) and actions is an int array of shape (N,). Return the array of shape (N, K) whose row $i$ is $\nabla_{\mathbf{z}}\log\operatorname{softmax}(\mathbf{z}_i)_{a_i} = \mathbf{e}_{a_i} - \operatorname{softmax}(\mathbf{z}_i)$ . It must be numerically stable for large logits.

import numpy as np

def softmax_score(logits, actions):
    raise NotImplementedError

Hidden tests, hints, and AI help, from $1.99 a month.

A worked example on a bandit

Take a two-armed bandit (a one-step MDP) with deterministic rewards $\mathbf{r} = (1, 0)$ and a softmax policy with logits $\mathbf{z} = (0, 0)$ , so $\boldsymbol{\pi} = (0.5, 0.5)$ . The objective is $J(\mathbf{z}) = \sum_a \pi_a r_a = \pi_1$ . Its exact gradient follows from the softmax Jacobian of §4.3, $\partial \pi_a/\partial z_j = \pi_a(\mathbb{1}[a=j] - \pi_j)$ , giving $\nabla_{\mathbf{z}} J = \boldsymbol{\pi}\odot(\mathbf{r} - J\mathbf{1}) = (0.25, -0.25)$ .

The policy gradient estimator sees one pull at a time. If it pulls arm 1 (probability 0.5), it reports $r_1(\mathbf{e}_1 - \boldsymbol\pi) = (0.5, -0.5)$ . If it pulls arm 2, it reports $r_2(\mathbf{e}_2 - \boldsymbol\pi) = (0, 0)$ . The average of the two reports, weighted by their probabilities, is $(0.25, -0.25)$ , exactly the true gradient. Every single report is wrong, though, either twice too large or zero. The estimator is unbiased and noisy, and that combination is the story of the rest of this chapter.

In general, for a $K$ -armed bandit the expectation of the estimator is $\sum_a \pi_a r_a(\mathbf{e}_a - \boldsymbol\pi) = \boldsymbol\pi\odot\mathbf{r} - J\boldsymbol\pi = \boldsymbol\pi\odot(\mathbf{r} - J\mathbf{1})$ , which matches the exact gradient term by term. The next exercise asks for this exact gradient, which is useful as a test oracle for any sampled implementation.

Core

The exact policy gradient of a softmax bandit

For a bandit with expected rewards r of shape (K,) and a softmax policy with logits z of shape (K,), implement bandit_policy_gradient(z, r) returning $\nabla_{\mathbf{z}} J$ where $J(\mathbf{z}) = \sum_a \operatorname{softmax}(\mathbf{z})_a\, r_a$ . Also implement sampled_policy_gradient(z, r, n, seed), which draws n actions with np.random.default_rng(seed).choice(K, size=n, p=pi) and returns the average of the single-sample estimates $r_{a_i}(\mathbf{e}_{a_i} - \boldsymbol\pi)$ .

import numpy as np

def bandit_policy_gradient(z, r):
    raise NotImplementedError

def sampled_policy_gradient(z, r, n, seed):
    raise NotImplementedError

Hidden tests, hints, and AI help, from $1.99 a month.

The policy gradient theorem

The trajectory form multiplies every action's score by the return of the entire trajectory, including rewards that arrived before the action was taken. That cannot be the most efficient use of the data. An action at step 10 cannot have caused the reward at step 3. The policy gradient theorem (Sutton, McAllester, Singh, and Mansour 2000) gives a sharper form, in which each action is weighted by its own action value.

Theorem. For a discounted MDP with start distribution $p_0$ and objective $J(\theta) = \sum_s p_0(s)V^{\pi_\theta}(s)$ ,

\nabla_\theta J(\theta) = \sum_{s} \rho^{\pi}(s) \sum_a \nabla_\theta\pi_\theta(a\mid s)\, Q^{\pi}(s, a) = \mathbb{E}_{\tau}\Big[\sum_{t=0}^{T-1} \gamma^t\, Q^{\pi}(s_t, a_t)\, \nabla_\theta\log\pi_\theta(a_t\mid s_t)\Big],

where $\pi = \pi_\theta$ and $\rho^\pi(s) = \sum_{t\ge0}\gamma^t\,\mathbb{P}(s_t = s)$ is the discounted state visitation, how often (with discounting) the policy finds itself in $s$ .

What is surprising is what is missing. Changing $\theta$ changes which states the agent visits, so you might expect a term involving the derivative of $\rho^\pi$ , which would be hopeless to estimate. There is no such term. The gradient only needs the policy's own score, weighted by how good each action is.

∂ Proof of the policy gradient theorem

Fix a state $s$ . The value is $V^\pi(s) = \sum_a \pi_\theta(a\mid s)Q^\pi(s,a)$ , and both factors depend on $\theta$ . By the product rule,

\nabla V^\pi(s) = \sum_a \Big[\nabla\pi_\theta(a\mid s)\, Q^\pi(s,a) + \pi_\theta(a\mid s)\,\nabla Q^\pi(s,a)\Big].

The reward $r(s,a)$ and the transitions $P$ do not depend on $\theta$ , so differentiating $Q^\pi(s,a) = r(s,a) + \gamma\sum_{s'}P(s'\mid s,a)V^\pi(s')$ gives $\nabla Q^\pi(s,a) = \gamma\sum_{s'}P(s'\mid s,a)\nabla V^\pi(s')$ . Write $\mathbf{u}(s) = \sum_a\nabla\pi_\theta(a\mid s)Q^\pi(s,a)$ for the first part (a vector with the shape of $\theta$ ), and $\mathbb{P}(s\to x, k)$ for the probability of being in state $x$ after exactly $k$ steps of $\pi$ from $s$ . Then

\nabla V^\pi(s) = \mathbf{u}(s) + \gamma\sum_{x}\mathbb{P}(s\to x, 1)\,\nabla V^\pi(x).

The same identity holds for $\nabla V^\pi(x)$ , so substitute it in, and again for the states after that,

\nabla V^\pi(s) = \mathbf{u}(s) + \gamma\sum_x\mathbb{P}(s\to x,1)\mathbf{u}(x) + \gamma^2\sum_x\mathbb{P}(s\to x, 2)\nabla V^\pi(x) = \dots = \sum_x\Big(\sum_{k=0}^{\infty}\gamma^k\,\mathbb{P}(s\to x, k)\Big)\mathbf{u}(x).

(The remainder after $k$ unrollings is $\gamma^k$ times a bounded quantity, so it vanishes as $k\to\infty$ when $\gamma < 1$ , or when every trajectory ends in finite time.) Average over the start state $s\sim p_0$ . The inner sum becomes $\rho^\pi(x)$ , and we get the first form of the theorem.

For the second form, use $\nabla\pi_\theta(a\mid s) = \pi_\theta(a\mid s)\nabla\log\pi_\theta(a\mid s)$ inside $\mathbf{u}$ , so that $\sum_x\rho^\pi(x)\mathbf{u}(x) = \sum_{t}\gamma^t\sum_x\mathbb{P}(s_t = x)\sum_a\pi_\theta(a\mid x)Q^\pi(x,a)\nabla\log\pi_\theta(a\mid x)$ , which is $\mathbb{E}_\tau[\sum_t\gamma^tQ^\pi(s_t,a_t)\nabla\log\pi_\theta(a_t\mid s_t)]$ written out. $\square$

You should not take a theorem like this on faith, and in a gridworld we do not have to. For a tabular softmax policy (one logit per state and action, $\theta \in \R^{|\mathcal{S}|\times|\mathcal{A}|}$ ), the gradient with respect to the logits of state $s$ is $\rho^\pi(s)\,\boldsymbol\pi(\cdot\mid s)\odot(Q^\pi(s,\cdot) - V^\pi(s))$ , by the same softmax algebra as the bandit. Everything on the right can be computed exactly with linear solves, including $\rho^\pi = (\mathbf{I} - \gamma\mathbf{P}_\pi^\top)^{-1}\mathbf{p}_0$ , and the left side can be checked by finite differences of the exact $J$ .

Python · cell 1

The theorem's gradient agrees with brute-force finite differences to about ten decimal places, with no derivative of the state distribution anywhere. (The visitation sums to $1/(1-\gamma) = 10$ , because the absorbing terminal states keep collecting visitation after the episode ends. They contribute nothing to the gradient, since all their actions have the same value.)

Research

The policy gradient theorem, exactly, for a tabular policy

For a finite MDP with P of shape (S, A, S), R of shape (S, A), start distribution p0 of shape (S,), discount $\gamma < 1$ , and a tabular softmax policy with logits theta of shape (S, A), implement exact_policy_gradient(P, R, p0, gamma, theta) that returns $\nabla_\theta J$ , shape (S, A), using the policy gradient theorem (no finite differences). Then implement discounted_visitation(P, p0, gamma, pi) returning $\rho^\pi$ , shape (S,).

import numpy as np

def discounted_visitation(P, p0, gamma, pi):
    raise NotImplementedError

def exact_policy_gradient(P, R, p0, gamma, theta):
    raise NotImplementedError

Hidden tests, hints, and AI help, from $1.99 a month.

REINFORCE

We rarely know $Q^\pi$ , but we can sample it. Given $(s_t, a_t)$ , the reward-to-go $G_t = \sum_{k=t}^{T-1}\gamma^{k-t}r_k$ , the discounted return from step $t$ onward, has conditional expectation exactly $Q^\pi(s_t, a_t)$ , by the definition of $Q^\pi$ . Substituting it into the theorem gives an unbiased estimator that needs nothing but sampled episodes. Williams (1992) called the resulting algorithm REINFORCE.

Run the current policy $\pi_\theta$ to collect $N$ episodes.
For every step of every episode, compute the reward-to-go $G_t^{(i)}$ with one backward pass (Exercise 15.1a).
Estimate the gradient and take a step of gradient ascent,

\hat{\mathbf{g}} = \frac{1}{N}\sum_{i=1}^{N}\sum_{t=0}^{T_i - 1}\gamma^t\, G_t^{(i)}\,\nabla_\theta\log\pi_\theta\big(a_t^{(i)}\mid s_t^{(i)}\big), \qquad \theta \leftarrow \theta + \eta\,\hat{\mathbf{g}}.

Throw the episodes away and repeat. They were sampled from the old policy, and the estimator is only unbiased for the policy that generated the data. This is what on-policy means, and it is the waste that PPO (§15.6) will try to reduce.

In plain words, after each batch of episodes, raise the log-probability of every action in proportion to how much reward followed it.

(That $G_t$ may replace the full $R(\tau)$ of the trajectory form, discarding the rewards before step $t$ , also has a direct one-line proof, which we give in §15.4 along with the baseline, because the same lemma proves both.)

The surrogate loss

In code, nobody writes down $\nabla_\theta\log\pi_\theta$ by hand. We build a scalar whose gradient is the estimator and let autograd differentiate it. Treating the returns as constants (detached from the graph), define

\mathcal{L}^{\text{PG}}(\theta) = -\frac{1}{N}\sum_{i}\sum_{t} \gamma^t\, G_t^{(i)}\, \log \pi_\theta\big(a_t^{(i)} \mid s_t^{(i)}\big).

Its gradient is $-\hat{\mathbf{g}}$ , so minimizing it with any optimizer performs REINFORCE. Look at its shape. It is a cross-entropy loss on the policy's own samples, with each sample's log-likelihood weighted by its return. REINFORCE is weighted maximum likelihood on your own behavior. With a reward that is 1 for success and 0 for failure and no discounting, the loss is exactly the supervised fine-tuning loss on the successful samples, divided by $N$ , and the failures contribute nothing. That is a useful bridge to rejection-sampling fine-tuning (§18.2), which does exactly that.

Two warnings about the surrogate. Its value means nothing. It is not the expected return, it can go up while the policy improves, and plotting it as a training curve misleads everyone who reads the plot. Track the average return instead. And the normalization matters more than it looks. Dividing by the number of episodes $N$ gives the estimator above. Dividing by the total number of steps in the batch only rescales it by a batch-dependent number, which ties the step size to how long the episodes happened to be. Dividing each episode's sum by its own length before averaging, as the original GRPO loss does, is different in kind. It weights every step of a long episode less than every step of a short one, which changes the objective, a subtle bias that GRPO's variants argue about at length (§18.4).

Interview

The REINFORCE loss and its gradient

Implement reinforce_loss(logits, actions, weights) for a batch of $N$ decisions, where logits has shape (N, K), actions is an int array of shape (N,), and weights of shape (N,) holds each decision's weight (for example $\gamma^tG_t$ ), treated as constants. Return (loss, dlogits) with

$\text{loss} = -\frac{1}{N}\sum_{i=1}^N w_i\,\log\operatorname{softmax}(\mathbf{z}_i)_{a_i}$

and dlogits $= \partial\,\text{loss}/\partial\mathbf{z}$ , of shape (N, K). Compute the log-softmax stably.

import numpy as np

def reinforce_loss(logits, actions, weights):
    raise NotImplementedError

Hidden tests, hints, and AI help, from $1.99 a month.

REINFORCE on a bandit

Start with the simplest possible case, a three-armed Bernoulli bandit with success probabilities $0.2$ , $0.5$ , and $0.8$ , a softmax policy over three logits, and small batches of 4 pulls per update, over 30 independent runs. For a bandit every episode is one step, so the reward-to-go is just the reward.

Python · cell 2

With the small learning rate every one of the 30 runs ends with most of its probability on the best arm. With a learning rate of 4, two runs end certain of the wrong arm, and with 16, a third of them do. A few lucky pulls of a mediocre arm drove its probability so high that the best arm was almost never tried again. A softmax policy that has collapsed onto one action stops exploring, and since the gradient of a nearly deterministic softmax is nearly zero ( $\boldsymbol\pi\odot(\mathbf{r} - J)$ vanishes when $\boldsymbol\pi$ is one-hot), it cannot climb back out. Keeping the policy from collapsing too early is the job of the entropy bonus in §15.6 and the KL penalty in §15.7.

REINFORCE where only a random policy is optimal

Here is a small environment that shows something value-based methods cannot do, from Sutton and Barto's Example 13.1, the short corridor with switched actions. There are three non-terminal cells in a row and a goal at the right end. Each step costs $1$ . Actions are "left" and "right", and they do what they say in the first and third cells, but in the middle cell they are reversed. The catch is that the agent cannot tell the cells apart. Every state looks the same, so the policy is a single probability $p$ of choosing "right", used everywhere.

A deterministic policy is a disaster here. Always going right gets stuck bouncing between the first two cells, and always going left never leaves the first cell. Let us compute the value of a general $p$ exactly. Write $V_1, V_2, V_3$ for the values of the three cells (all with $\gamma = 1$ ). In cell 1, "right" (probability $p$ ) moves to cell 2 and "left" stays put. In cell 2, the reversed actions mean "right" moves back to cell 1 and "left" moves on to cell 3. In cell 3, "right" reaches the goal. So

V_1 = -1 + (1-p)V_1 + pV_2, \qquad V_2 = -1 + pV_1 + (1-p)V_3, \qquad V_3 = -1 + (1-p)V_2 .

The first equation gives $V_1 = V_2 - 1/p$ . Substitute it and the third equation into the second, $V_2 = -1 + p(V_2 - 1/p) + (1-p)(-1 + (1-p)V_2)$ , which simplifies to $V_2\,p(1-p) = p - 3$ . Then

J(p) = V_1 = \frac{p-3}{p(1-p)} - \frac{1}{p} = -\frac{2(2 - p)}{p(1-p)} .

Setting the derivative to zero gives $p^2 - 4p + 2 = 0$ , so the best policy goes right with probability $p^\star = 2 - \sqrt 2 \approx 0.586$ and earns $J(p^\star) = -(6 + 4\sqrt 2) \approx -11.66$ . For comparison, an $\varepsilon$ -greedy agent with $\varepsilon = 0.1$ , which goes right with probability $0.95$ or $0.05$ , earns $-44.2$ or $-82.1$ . An agent that learns values and acts (nearly) greedily on them cannot represent the optimum. A policy gradient method can, because it searches over $p$ directly.

Python · cell 3

Starting from the $\varepsilon$ -greedy-like policy $p = 0.95$ , REINFORCE walks $p$ down toward the optimum $0.586$ , and the expected return climbs from $-44$ toward $-11.7$ . It needed hundreds of episodes for a problem with one parameter, though, and the step size had to be tiny ( $2^{-12}$ ). Every return is negative, often large (from $-1$ near the goal to $-100$ or worse after a long bounce), so every update pushes down on the probability of every action taken, and the useful signal is only the small difference between how hard good and bad actions get pushed. That observation is the subject of the next part of this section.

Why REINFORCE has high variance

An unbiased gradient is only useful if its noise is small enough that a reasonable number of samples averages it out. REINFORCE's noise is large, for four reasons that are worth separating.

It scales with the size of the returns, not their differences. Add a constant $c$ to every reward of a bandit. The expected gradient does not change, since $\mathbb{E}_{a\sim\pi}[c\,\nabla\log\pi(a)] = c\,\nabla\sum_a\pi(a) = c\,\nabla 1 = 0$ . But each single-sample estimate gains a term $c\,\nabla\log\pi(a)$ , whose variance grows like $c^2$ . The corridor's returns are all large and negative, so most of each update is this useless term.
It sums over the horizon. The trajectory estimator multiplies a sum of $T$ random scores by a sum of $T$ random rewards, so its variance tends to grow with the episode length. A 1,000-token response is a 1,000-step episode.
One number takes the credit for everything. With a single reward at the end, every action in the episode gets the same weight, whether it helped, hurt, or did nothing.
Rare, large returns dominate. A trajectory that is unlikely under the policy but has a huge return contributes a huge estimate on the rare occasions it is sampled.

The first effect is the easiest to see and to fix. The widget below uses a two-armed bandit with rewards $r_1 > r_2$ and a single parameter $\theta$ , the logit of arm 1, so $\pi_1 = \sigma(\theta)$ and the true gradient is $\pi_1\pi_2(r_1 - r_2)$ . Shift both rewards by the same offset and watch the spread of the single-sample estimates, and the learning curves, change while the true gradient stays put.

Subtracting a number from the rewards before multiplying by the score is called using a baseline, and the widget's toggle is the simplest one, the average reward. §15.4 proves that a baseline that does not depend on the action never introduces bias, finds the baseline that minimizes the variance, and arrives at the advantage function, actor-critic methods, and the form of the policy gradient used in every modern algorithm.

Do not conclude from the corridor that REINFORCE is only a teaching device. With a baseline it is a serious method for language models, whose pretrained policy starts so close to good behavior that small, noisy steps are enough. Ahmadian et al. (2024) fine-tuned models of about 7 billion parameters (Pythia-6.9B and Llama-7B) from human preferences on the TL;DR summarization and Anthropic helpful-and-harmless datasets, and found that REINFORCE-style updates matched or beat PPO there, as well as DPO and rejection-sampling fine-tuning. Their best variant used a leave-one-out baseline over several responses to the same prompt, which §15.4 derives and §18.4 turns into RLOO and GRPO.

In PyTorch

The production version of REINFORCE is three lines inside the training loop. Compute the log-probabilities of the taken actions, multiply by the detached returns, and call backward on the negative mean. Here we check that autograd's gradient of that surrogate loss on the gridworld equals the REINFORCE estimator computed by hand with the score formula.

🔥 In PyTorchfor a machine with PyTorch

🎯 Interview questions on this section

Derive the policy gradient

\nabla_\theta\mathbb{E}_{\tau\sim\pi_\theta}[R(\tau)]

Write $J = \sum_\tau p_\theta(\tau)R(\tau)$ and differentiate, $\nabla J = \sum_\tau R(\tau)\nabla p_\theta(\tau)$ . Use $\nabla p_\theta = p_\theta\nabla\log p_\theta$ to turn it back into an expectation, $\nabla J = \mathbb{E}_\tau[R(\tau)\nabla\log p_\theta(\tau)]$ . Since $\log p_\theta(\tau) = \log p_0(s_0) + \sum_t\log\pi_\theta(a_t\mid s_t) + \sum_t\log P(s_{t+1}\mid s_t,a_t)$ and only the policy terms depend on $\theta$ , $\nabla J = \mathbb{E}_\tau[R(\tau)\sum_t\nabla\log\pi_\theta(a_t\mid s_t)]$ . The dynamics never need to be known or differentiated, and $R$ can be any black-box score.

State the policy gradient theorem. What is surprising about it?

$\nabla J(\theta) = \sum_s\rho^\pi(s)\sum_a\nabla\pi_\theta(a\mid s)Q^\pi(s,a) = \mathbb{E}_\tau[\sum_t\gamma^tQ^\pi(s_t,a_t)\nabla\log\pi_\theta(a_t\mid s_t)]$ , with $\rho^\pi$ the discounted state visitation. Changing $\theta$ changes the distribution of visited states, yet the gradient contains no derivative of $\rho^\pi$ . The proof unrolls $\nabla V = \sum_a[\nabla\pi\,Q + \pi\nabla Q]$ with $\nabla Q = \gamma\sum_{s'}P\nabla V$ , and the state-distribution effects are absorbed into the visitation weights.

What is the REINFORCE surrogate loss, and why is its value meaningless?

$\mathcal{L} = -\frac1N\sum_i\sum_t G_t^{(i)}\log\pi_\theta(a_t^{(i)}\mid s_t^{(i)})$ with the returns detached. Its gradient equals minus the policy gradient estimate, so minimizing it does REINFORCE. Its value depends on the arbitrary scale and offset of the returns and on the current batch, and it is not the expected return, so it can rise while the policy improves. Monitor average return, and KL to the previous policy, instead.

Why does REINFORCE have high variance? Give at least three reasons.

Each single-sample estimate is $G\nabla\log\pi$ , so its size scales with the raw returns, and a constant offset in the rewards adds a zero-mean term with variance proportional to the offset squared. Over long horizons it sums many random scores and many random rewards. A single terminal reward gives every action the same credit regardless of its effect. And rare high-return trajectories make occasional huge contributions. Baselines, reward-to-go, critics, GAE, and larger batches are the standard remedies.

How is REINFORCE related to supervised fine-tuning?

The REINFORCE loss is a log-likelihood loss on the model's own samples, weighted by return. With binary rewards and $\gamma = 1$ , it is exactly the SFT loss on the successful samples (scaled by $1/N$ ), with failures ignored. That is why rejection-sampling fine-tuning, which generates samples, filters the good ones, and fine-tunes on them, behaves like a crude, off-policy policy gradient. Real policy gradient methods add negative weights for bad samples (through baselines), on-policy sampling, and a constraint on how far the policy moves.

When can a stochastic policy beat every deterministic one, and what does that mean for value-based methods?

When the agent cannot distinguish states, as in the short corridor where the middle cell's actions are reversed but all cells look alike. The best policy there goes right with probability $2 - \sqrt2 \approx 0.586$ and earns about $-11.66$ , while $\varepsilon$ -greedy policies earn $-44$ or $-82$ . Value-based methods act (nearly) greedily on learned values, so they cannot represent such an optimum. Policy gradient methods search over the probabilities directly, so they can.

Most implementations drop the

\gamma^t

factor in the policy gradient. Is that a problem?

Without $\gamma^t$ , steps late in an episode are weighted as much as early ones, which corresponds to weighting states by their undiscounted rather than discounted visitation. The update is then not the gradient of the discounted objective, nor, as Nota and Thomas (2020) showed, of any objective. It still works well in practice, and one common reading is that it uses $\gamma$ as a variance-reduction device on the returns rather than as part of the objective. For language model RL with $\gamma = 1$ the question does not arise.

🧭 What to remember

The log-derivative trick, $\nabla p_\theta = p_\theta\nabla\log p_\theta$ , turns the gradient of an expected reward into an expectation that can be estimated from samples, even when the reward is not differentiable.
The environment's dynamics drop out of $\nabla\log p_\theta(\tau)$ , leaving $\nabla J = \mathbb{E}[R(\tau)\sum_t\nabla\log\pi_\theta(a_t\mid s_t)]$ . For a softmax policy the score is $\mathbf{e}_a - \boldsymbol\pi$ .
The policy gradient theorem sharpens this to $\mathbb{E}[\sum_t\gamma^tQ^\pi(s_t,a_t)\nabla\log\pi_\theta(a_t\mid s_t)]$ , with no derivative of the state distribution, and it can be checked exactly in a gridworld.
REINFORCE replaces $Q^\pi$ by the sampled reward-to-go. In code it is a return-weighted log-likelihood loss on the policy's own samples, whose value is meaningless and whose normalization matters.
Policy gradient methods can represent optimal stochastic policies that value-based methods cannot, but plain REINFORCE is slow and noisy, mostly because its estimates scale with the raw returns.
Next, §15.4 proves that subtracting a baseline removes much of that noise without adding bias, and builds the advantage function and actor-critic methods from it.

Williams (1992). Simple Statistical Gradient-Following Algorithms for Connectionist Reinforcement Learning. The REINFORCE paper, including the observation that a baseline leaves the gradient unbiased.
Sutton, McAllester, Singh & Mansour (2000). Policy Gradient Methods for Reinforcement Learning with Function Approximation. The policy gradient theorem.
OpenAI Spinning Up. Part 3: Intro to Policy Optimization. A clean derivation of the simplest policy gradient, with the reward-to-go argument and a minimal implementation.
Sutton & Barto (2018). Reinforcement Learning: An Introduction, 2nd edition. Chapter 13 covers policy gradient methods, including the short-corridor example and the proof we followed.
Ahmadian et al. (2024). Back to Basics: Revisiting REINFORCE Style Optimization for Learning from Human Feedback in LLMs. REINFORCE and its leave-one-out variant against PPO and DPO on 7-billion-parameter models.
Nota & Thomas (2020). Is the Policy Gradient a Gradient? What dropping the $\gamma^t$ factor really does.
Levine. CS 285, Lecture 5, Policy Gradients. The lecture slides derive the same estimator with the same trajectory notation.

Zero to Frontier is made for larger screens

Policy gradients and REINFORCE

The objective

The log-derivative trick

The environment drops out

The score of a softmax policy

The score of a softmax policy

A worked example on a bandit

The exact policy gradient of a softmax bandit

The policy gradient theorem

The policy gradient theorem, exactly, for a tabular policy

REINFORCE

The surrogate loss

The REINFORCE loss and its gradient

REINFORCE on a bandit

REINFORCE where only a random policy is optimal

Why REINFORCE has high variance

In PyTorch

Read more

Watch next