The policy gradient theorem

This is the load-bearing chapter of Part V. Everything before it was scaffolding; everything after it, REINFORCE, actor-critic, PPO, GRPO, is a variation on the single result I derive here. The policy gradient theorem tells you how to nudge the parameters of a stochastic policy so that its expected reward goes up, and it does so with a formula that, astonishingly, does not require you to differentiate through the environment or the reward. That one property, "environment-gradient-free," is the entire reason you can train a language model against a black-box verifier, a unit-test suite, or a human thumbs-up. If the theorem needed a differentiable reward, RLVR would be impossible, because "did this proof check out" is not a differentiable function of the tokens. It needs only that you can sample and score. Let me earn that claim carefully, twice, by two different routes.

Theory

Parameterized policies

Replace the lookup-table policy of earlier chapters with a differentiable function of parameters $\theta$. For a language model, $\theta$ is the network weights and the policy is the tempered softmax over logits from Chapter 5.3,

$${\pi }_{\theta }(a\mid s)=\frac{\exp (z_{\theta }(s{)}_a)}{{\sum }_b\exp (z_{\theta }(s{)}_b)},\text{\hspace{1em}(1)}$$

where $z_{\theta }(s)$ is the logit vector the network computes for state (context) $s$. The only thing I will use about ${\pi }_{\theta }$ is that it is differentiable in $\theta$ and that it is a proper probability distribution, ${\sum }_a{\pi }_{\theta }(a\mid s)=1$ for every $s$. That normalization identity is quietly the source of the whole theorem, as you will see.

The objective

I want to maximize expected return. Write a full trajectory as $\tau =(s_0,a_0,r_1,s_1,a_1,\dots ,s_T)$ and its return as $R(\tau )={\sum }_{t=0}^{T-1}{\gamma }^t{r}_{t+1}$. The probability of a trajectory under the policy factorizes into the initial-state distribution, the policy, and the environment dynamics:

$$p_{\theta }(\tau )={\rho }_0(s_0)\prod _{t=0}^{T-1}{\pi }_{\theta }(a_t\mid s_t)p(s_{t+1}\mid s_t,a_t).\text{\hspace{1em}(2)}$$

The objective is the expected return over trajectories the policy induces:

$$J(\theta )={\mathbb{E}}_{\tau \sim p_{\theta }}[R(\tau )]=\sum _{\tau }{p}_{\theta }(\tau )R(\tau ).\text{\hspace{1em}(3)}$$

I want ${\nabla }_{\theta }J$. The obstacle is that $\theta$ sits inside $p_{\theta }(\tau )$, which contains the environment term $p(s_{t+1}\mid s_t,a_t)$ that I do not know and cannot differentiate. The log-derivative trick makes the obstacle vanish.

The log-derivative trick

Now push (4) through the objective. This is the trajectory route to the theorem, and it is the one that makes the environment-gradient-free property blindingly obvious.

Why the environment-gradient-free form is the entire reason RLHF exists

Let me hammer this, because it is the thesis in miniature. Suppose the theorem had come out as ${\nabla }_{\theta }J=\mathbb{E}[\dots {\nabla }_{\theta }R(\tau )]$, requiring the gradient of the reward with respect to the parameters. Then you could only train against rewards that are differentiable functions of the model output. But the rewards I care about are things like:

• Does this generated Python program pass its unit tests? (A boolean from a subprocess.)
• Does this math answer equal the reference after symbolic simplification? (A SymPy call.)
• Did an Inspect scorer mark this transcript correct? (A grader, possibly another model.)
• Did a human prefer response A over response B? (A click.)

None of these is a differentiable function of the tokens. Every one of them is a sampling-plus-scoring black box. Equation (5) says that is exactly, and only, what you need. The policy gradient theorem is the mathematical permission slip that lets an eval score, the object of Parts III and IV, stand in as $R(\tau )$ and drive learning. That is why I can call this book "Evals as Rewards" and mean it literally: the theorem is what makes the substitution legal.

The second route: the episodic policy gradient theorem via value functions

Route 1 is the cleanest derivation and the one that matters most for LLMs, but Sutton and Barto reach the same destination from the value-function side, and seeing both routes converge is reassuring and also reveals the $q_{\pi }$ structure that actor-critic methods exploit. I will derive the episodic form, where the objective is the value of the start state, $J(\theta )=v_{{\pi }_{\theta }}(s_0)$.

Equations (5) and (7) are the same theorem seen from two sides. Route 1 gives the trajectory-level form with the full return $R(\tau )$ as the weight; route 2 gives the state-action form with the action value $q_{\pi }(s,a)$ as the weight. They are reconciled by noting that $R(\tau )$ is an unbiased single-sample estimate of $q_{\pi }$, which is the observation the next chapter builds REINFORCE on. The proportionality constant in (6) and (7) is the average episode length in the episodic case and is $1$ in the continuing case with the right normalization; in practice it is absorbed into the learning rate, so I will stop writing $\propto$ and write $=$ from here on.

The score function for a softmax policy

The abstract ${\nabla }_{\theta }\log {\pi }_{\theta }$ becomes concrete and cheap for the softmax policy (1), and since the lab below and every real LLM training step compute exactly this quantity, it is worth deriving once. The result is one of the most-used gradients in all of machine learning, and it is startlingly clean.

Notice that the expected value of (8) under $\pi$ is ${\sum }_a{\pi }_a({\mathbf{e}}_a-\pi )=\pi -\pi =0$, the score-function-integrates-to-zero identity I will need constantly in Chapter 5.5. It falls out of the softmax gradient for free.

On-policy by construction, and the crack that PPO widens

One more property of (5) and (7) deserves flagging now, because it is the seam along which every later algorithm is built. The expectation is taken over trajectories sampled from ${\pi }_{\theta }$, the same policy whose gradient you are computing. The estimator is on-policy: it is only valid for data generated by the current parameters. The instant you take a gradient step, $\theta$ changes, and strictly speaking your previous rollouts are stale, drawn from a policy that no longer exists.

For a single ascent step per batch of rollouts this is fine, and it is exactly what REINFORCE does. But generating rollouts is the expensive part on one GPU, all that vLLM decoding, and you would love to take several optimizer steps on one batch of samples to amortize the cost. Doing so violates the on-policy assumption. The repair is importance sampling: reweight each sample by the ratio of the new policy's probability to the old one's,

$$r_t(\theta )=\frac{{\pi }_{\theta }(a_t\mid s_t)}{{\pi }_{{\theta }_{\text{old}}}(a_t\mid s_t)},\text{\hspace{1em}(9)}$$

which corrects the expectation back to being over ${\pi }_{{\theta }_{\text{old}}}$. That ratio (9) is the protagonist of the PPO and GRPO chapters: PPO clips it to keep the reused updates from wandering too far from the sampling policy, and GRPO inherits the same ratio. So the humble on-policy caveat here is the reason trust-region methods exist at all. Keep (9) in mind; it is the first thing that changes when you move past pure REINFORCE.

Tooling

The tool is autograd plus a sanity check. The single most valuable habit when implementing any policy-gradient method is to verify the estimator (5) against a brute-force finite-difference gradient on a tiny problem, because a sign error or a missing log is otherwise invisible: the loss still goes down sometimes, just for the wrong reasons. I reuse the uv project from Chapter 5.1 and add nothing; NumPy is enough to make the point without dragging in a deep-learning framework.

Lab

The artifact numerically confirms the policy gradient theorem. I build a one-step (bandit) softmax policy over three actions with a fixed reward per action, compute the analytic policy gradient estimator from (5) by sampling, and compare it against a finite-difference estimate of the true ${\nabla }_{\theta }J$. Agreement to a couple of decimals is the proof, in code, that "sample the log-prob gradient and weight by reward" really does climb $J$.

"""Numerically verify the policy gradient theorem on a 3-armed softmax
bandit. Compare the score-function (log-derivative) estimator, eq (5),
against a finite-difference estimate of the true gradient of J(theta).

If they agree, the theorem's sampling estimator is correctly implemented.

Run:  uv run python labs/policy_gradient_check.py
"""
from __future__ import annotations
import numpy as np

REWARDS = np.array([1.0, -0.5, 2.0])     # fixed reward per action (the "verifier")
N_ACTIONS = len(REWARDS)


def softmax(theta):
    z = theta - theta.max()
    e = np.exp(z)
    return e / e.sum()


def J(theta):
    """True objective: expected reward = sum_a pi(a) R(a).
    Here the environment is trivial so we can compute J in closed form,
    which is what lets us finite-difference it as ground truth."""
    return float(softmax(theta) @ REWARDS)


def finite_difference_grad(theta, eps=1e-6):
    """Ground-truth gradient of J via central differences."""
    g = np.zeros_like(theta)
    for i in range(len(theta)):
        d = np.zeros_like(theta)
        d[i] = eps
        g[i] = (J(theta + d) - J(theta - d)) / (2 * eps)
    return g


def grad_log_pi(theta, a):
    """d/dtheta log pi(a): for a softmax, this is (onehot(a) - pi)."""
    pi = softmax(theta)
    onehot = np.zeros_like(theta)
    onehot[a] = 1.0
    return onehot - pi


def score_function_grad(theta, n_samples, rng):
    """Monte Carlo estimate of nabla J via eq (5): average of
    R(a) * grad_log_pi(a) over sampled actions."""
    pi = softmax(theta)
    est = np.zeros_like(theta)
    for _ in range(n_samples):
        a = rng.choice(N_ACTIONS, p=pi)
        reward = REWARDS[a]                       # scalar, detached
        est += reward * grad_log_pi(theta, a)     # weight log-prob grad by reward
    return est / n_samples


if __name__ == "__main__":
    rng = np.random.default_rng(0)
    theta = np.array([0.3, -0.2, 0.1])

    fd = finite_difference_grad(theta)
    print(f"finite-difference grad (ground truth): {fd}")

    for n in (100, 1_000, 100_000):
        sf = score_function_grad(theta, n, rng)
        err = np.linalg.norm(sf - fd)
        print(f"score-function grad, n={n:>7}: {sf}   ||err||={err:.4f}")

    # Take a few gradient-ascent steps and watch J climb.
    print("\ngradient ascent on J using the score-function estimator:")
    lr = 0.2
    for step in range(6):
        g = score_function_grad(theta, 20_000, rng)
        theta = theta + lr * g
        print(f"step {step}: J(theta)={J(theta):.4f}  pi={softmax(theta).round(3)}")

What you should see. The score-function estimator (5) converges to the finite-difference ground truth as the sample count grows: the error norm shrinks from order $10^{-1}$ at 100 samples to order $10^{-2}$ or better at $10^5$ samples, the classic Monte Carlo $1/\sqrt{n}$ shrinkage. That convergence is the policy gradient theorem verified: two completely different computations, one differentiating the reward-weighted log-prob and the other perturbing $\theta$ and measuring $J$, agree. The final block does gradient ascent and you watch $J(\theta )$ climb toward $2.0$ while the policy concentrates on action index 2, the highest-reward arm. Notice what the estimator never touched: the reward vector was used only as a scalar multiplier, never differentiated. Swap REWARDS for a call to a unit-test runner or an Inspect scorer and the exact same code trains a language model. That substitutability is the whole game.