The policy gradient theorem: a full derivation, checked numerically

The policy gradient theorem gives the gradient of a policy's expected return with respect to its parameters as an expectation you can sample: ${\nabla }_{\theta }J(\theta )={\mathbb{E}}_{\pi }[{\nabla }_{\theta }\log {\pi }_{\theta }(a\mid s)Q^{\pi }(s,a)]$. It needs no derivative of the environment's dynamics, only the log-probability of the actions the agent actually took and an estimate of how good they turned out to be. Every policy gradient method, from REINFORCE to actor-critic and PPO, is a way of estimating this one expectation.

This page derives the theorem from scratch, starting with the log-derivative trick, then checks every step numerically on a two-state problem small enough to solve exactly. You will see the exact gradient, the sampled estimate, how noisy that estimate is, and two mistakes that quietly turn it into something that is not a gradient at all. All numbers come from seeded numpy runs we did for this article.

A two-state example to keep in mind#

Take the smallest problem where a policy gradient is interesting. There are two states, A and B, and two actions in each. In A, the action stay pays +1 and keeps you in A; the action go pays nothing and moves you to B. In B, stay pays +2 and keeps you in B; go pays nothing and sends you back to A. The agent starts in A, and future rewards are discounted by $\gamma =0.9$. (If the vocabulary of states, actions and discounting is new, start with the Markov decision process page.)

The policy has two parameters, one logit per state: ${\pi }_{\theta }(\text{go}\mid A)=\sigma ({\theta }_A)$ and ${\pi }_{\theta }(\text{stay}\mid B)=\sigma ({\theta }_B)$, where $\sigma$ is the logistic function. At $\theta =(0,0)$ the agent flips a fair coin everywhere, and its expected return from A is $J=7.25$. The best policy goes to B and stays there forever, for a return of $2/(1-\gamma )=20$. A second, worse policy stays in A forever for a return of 10, and it matters later: it is a local optimum.

Which way should $\theta$ move? Leaving A gives up a sure +1 now for a chance at +2 per step later, and whether that is worth it depends on how the agent behaves in B. The two parameters interact. The policy gradient answers the question precisely: at $\theta =(0,0)$, it is $(-0.756,2.756)$. Increase the logit of staying in B a lot, and for now slightly decrease the probability of leaving A, because leaving A only pays once the agent reliably stays in B. We will derive where those numbers come from.

What is the policy gradient theorem?#

Write a trajectory as $\tau =(s_0,a_0,r_0,s_1,a_1,r_1,\dots )$ and the discounted return as $R(\tau )={\sum }_{t\ge 0}{\gamma }^t{r}_t$. The objective is the expected return of the policy from the start-state distribution:

Here ${\rho }_0$ is the start-state distribution, ${\pi }_{\theta }$ the policy, and $P$ the environment's transition probabilities, which we do not know and cannot differentiate. The theorem, in its state-action form, says

where $d_{\gamma }^{\pi }(s)={\sum }_t{\gamma }^t\mathrm{Pr}(s_t=s)$ is the discounted visitation of state $s$, and $Q^{\pi }(s,a)$ is the expected discounted return after taking $a$ in $s$ and following $\pi$ afterwards (see the Bellman equation page for how $Q^{\pi }$ is defined and computed). In words: push up the probability of each action in proportion to how good it is, weighted by how often the policy visits the state where it is taken. The striking part is what is missing. Changing $\theta$ changes which states the agent visits, yet no derivative of $d_{\gamma }^{\pi }$ appears.

Deriving it step by step#

Step 1: the log-derivative trick#

Differentiate $J$ and move the gradient inside the sum over trajectories:

This is not yet useful: ${\sum }_{\tau }{\nabla }_{\theta }{p}_{\theta }(\tau )R(\tau )$ is not an expectation under anything we can sample. The trick is the identity ${\nabla }_{\theta }p=p{\nabla }_{\theta }\log p$, which is just the chain rule for the logarithm read backwards:

Now the gradient is an expectation under the policy's own trajectory distribution, so averaging over sampled trajectories estimates it without bias. Statisticians call ${\nabla }_{\theta }\log p_{\theta }$ the score function, and this construction is the score-function or likelihood-ratio gradient estimator; Glynn studied it for general stochastic systems [4], and Williams introduced it to reinforcement learning as REINFORCE [3]. Mohamed and colleagues survey it alongside the alternatives [6].

Step 2: the dynamics drop out#

Take the log of the trajectory probability. The product becomes a sum:

Only the policy terms depend on $\theta$, so the start distribution and every transition term vanish under the gradient:

This is the whole reason policy gradients work in model-free reinforcement learning. The environment can be a simulator, a physical robot or a human rater; as long as you can run the policy and compute $\log {\pi }_{\theta }(a_t\mid s_t)$ for the actions it took, you can estimate the gradient. The same property is what lets reinforcement learning for language models treat a reward model or a unit test as a black box.

Step 3: rewards cannot be caused by later actions#

The estimator above multiplies every action's score by the whole return, including rewards that arrived before the action was taken. Those terms average to zero. For $k<t$, the reward $r_k$ is fixed once we condition on the history up to time $t$, and the expected score of a distribution is zero:

This one-line fact does a lot of work in the rest of this cluster: it is also why subtracting a baseline leaves the gradient unbiased. Dropping the past rewards leaves the reward-to-go form

where $G_t$ is the discounted return from time $t$ onward and the factor ${\gamma }^t$ is what remains of ${\gamma }^k$ after factoring out the rewards from $t$ on. Note that ${\gamma }^t$ stays. We come back to it below because dropping it is the most common error in implementations.

Step 4: from trajectories to states#

Conditioning on $(s_t,a_t)$, the expected reward-to-go is exactly $Q^{\pi }(s_t,a_t)$, so $G_t$ can be replaced by $Q^{\pi }(s_t,a_t)$ inside the expectation. Grouping the sum over time by which state is visited turns ${\sum }_t{\gamma }^t\mathrm{Pr}(s_t=s)$ into $d_{\gamma }^{\pi }(s)$, and that gives the state-action form stated above. Sutton and colleagues reached the same result by a different route [2]: differentiate the Bellman equation $V^{\pi }(s)={\sum }_a{\pi }_{\theta }(a\mid s)Q^{\pi }(s,a)$ with the product rule, giving $\nabla V^{\pi }(s)={\sum }_a[\nabla {\pi }_{\theta }(a\mid s)Q^{\pi }(s,a)+{\pi }_{\theta }(a\mid s)\gamma {\sum }_{s^{\prime }}P(s^{\prime }\mid s,a)\nabla V^{\pi }(s^{\prime })]$, then unroll the recursion. Each unrolling step moves one transition further along the trajectory and multiplies by $\gamma$, and summing the series yields $d_{\gamma }^{\pi }$. The two proofs agree because they compute the same quantity; Sutton and Barto's textbook walks through the episodic case in this style [1].

Two remarks. First, $d_{\gamma }^{\pi }$ is not a probability distribution; it sums to $1/(1-\gamma )$. In our example the visitation from A at $\theta =0$ is $d(A)=5.5$ and $d(B)=4.5$, which sum to 10. Many texts normalize it and absorb the constant into the step size. Second, the theorem holds for any differentiable stochastic policy, tabular, linear or neural. Silver and colleagues later proved an analogous theorem for deterministic policies in continuous action spaces, which underlies DDPG-style methods [7].

Checking the theorem numerically#

A derivation is easier to trust when you can watch it hold. For the two-state problem everything is computable exactly. $V^{\pi }$ solves a 2-by-2 linear system, $Q^{\pi }(s,a)=r(s,a)+\gamma V^{\pi }(s^{\prime })$, and $d_{\gamma }^{\pi }$ solves another. With a logistic policy, ${\nabla }_{{\theta }_s}\pi (\text{go or stay}\mid s)=\pm \sigma ({\theta }_s)(1-\sigma ({\theta }_s))$, so the theorem reduces to one product per state.

Setup (simulation). At $\theta =(0,0)$ we computed (a) the gradient from the policy gradient theorem, (b) a central finite difference of the exact $J$, and (c) sampled score-function estimates from 200,000 simulated episodes, each truncated at 150 steps ($0.9^{150}\approx 10^{-7}$, so truncation is negligible).

The exact theorem and the finite difference agree to every printed digit, and both unbiased estimators land within their standard errors (shown after ±) of the truth. The reward-to-go estimator, which simply deletes terms whose mean is zero, has a per-episode standard deviation of 6.8 against 32.5 for the full-return version: about 23 times less variance from a free change.

How many episodes does an estimate need? We repeated the reward-to-go estimate 2,000 times at each batch size and measured the error relative to the true gradient's length. With a single episode the median relative error is 2.56, so a typical one-episode estimate is mostly noise, and only 63% of single-episode estimates even point uphill (positive inner product with the true gradient). With 10 episodes the median error is 0.89 and 94% point uphill; with 100 it is 0.28, and with 1,000 it is 0.083. The error falls as $1/\sqrt{N}$, exactly as an average of independent samples should.

Explore the gradient field#

The figure below draws the exact return $J(\theta )$ over both logits as a heat map, with the direction of the policy gradient at a grid of points. Click or tap anywhere (or use the sliders) to move the policy. The thick arrow is the exact gradient from the theorem; the thin red arrows are 20 independent sampled estimates, each averaging $N$ episodes, drawn at the same scale. The dashed path follows the exact gradient uphill. The estimator menu lets you switch between the three estimators in the table.

Three things are worth trying. Set $N=1$ and redraw a few times: single-episode estimates scatter in every direction, which is the variance problem the REINFORCE algorithm lives with. Switch to the dropped-${\gamma }^t$ estimator: the arrows grow roughly fifteenfold, because each visit is now weighted as if it happened at time zero. And click in the lower left quadrant: the exact gradient leads to the "stay in A" policy, whose return is 10, not 20. The policy gradient theorem gives you a true gradient, and gradient ascent on a non-concave objective finds local optima. In this problem, starting from $(0,-3)$ or $(1,-3)$ ends at return 10, while $(-1,-1)$ reaches 18.0 and climbing.

Policy gradient vs Q-learning#

Value-based methods such as Q-learning learn $Q^{\ast }$ and act greedily on it; the policy is implicit. Policy gradient methods parameterize the policy directly and climb $J$. The theorem explains when the second approach is the better tool:

• Continuous or huge action spaces. A greedy policy needs $\arg {\max }_{a}Q(s,a)$, which is an optimization problem of its own when actions are continuous torques or entire token sequences. A policy gradient only needs to sample $a\sim {\pi }_{\theta }$ and evaluate $\log {\pi }_{\theta }(a\mid s)$.
• Stochastic optimal policies. In partially observed problems and games, the best policy can be random. Greedy value-based policies are deterministic; a softmax or Gaussian policy can represent the right mixture.
• Smooth improvement. A small change in $\theta$ changes action probabilities a little. A small change in $Q$ can flip an argmax and change behaviour abruptly, which is one reason value-based methods with function approximation can diverge.
• The cost: variance and sample efficiency. The estimator above is on-policy and noisy. Value-based methods reuse old data freely; vanilla policy gradients throw each batch away after one step. Most of the methods built on the theorem exist to reduce that cost: baselines and advantages, learned critics in actor-critic methods and GAE, and the trust regions of TRPO and PPO that let one batch support several update steps.

The two families are less separate than they look. A learned critic that estimates $Q^{\pi }$ or $V^{\pi }$ is exactly a value function, and Schulman, Chen and Abbeel showed that entropy-regularized (soft) Q-learning and policy gradients are closely related, in some settings equivalent [8]. For the broader map of where these methods sit, see what is reinforcement learning.

Common mistakes when implementing the policy gradient#

Dropping the ${\gamma }^t$ factor#

Most implementations multiply $\nabla \log {\pi }_{\theta }(a_t\mid s_t)$ by the discounted reward-to-go $G_t$ but omit the outer ${\gamma }^t$. That update weights a state visited at step 50 as if it were visited at the start. Nota and Thomas showed that the resulting direction is, in general, not the gradient of any function [5]. In our example it happens to point almost the right way (cosine similarity 0.999 with the true gradient), but its components are 14 and 17 times too large, so even the relative weighting of the two states is off. In larger problems with states that are only reached late, the direction can be badly wrong. The usual justification is that $\gamma$ is a variance-reduction device and the real objective is undiscounted; that is a reasonable engineering choice, but it should be a choice, not an accident.

Using data from a different policy#

The expectation is over the current policy's trajectories. Reusing old samples without correction estimates the gradient of a different objective. Importance weights ${\pi }_{\theta }(a\mid s)/{\pi }_{\text{old}}(a\mid s)$ fix the action mismatch, and the trust-region view in the TRPO and PPO article explains when the state-distribution mismatch can be ignored. Off-policy data from logs is the territory of offline reinforcement learning.

Getting the sign or the stop-gradient wrong#

In autodiff frameworks the theorem is implemented as a surrogate loss $-\frac{1}{N}\sum \log {\pi }_{\theta }(a_t\mid s_t){\hat{A}}_t$, where ${\hat{A}}_t$ must be treated as a constant. If the return or advantage estimate depends on $\theta$ (through a shared critic, for example) and gradients flow through it, you are no longer computing the policy gradient. The surrogate's value is also meaningless as a progress metric: it is not $J$, and it can go up while $J$ goes down.

Reading a single noisy estimate as signal#

With one episode per update, our estimate pointed downhill 37% of the time. Learning still works, because steps in the wrong direction cancel over many updates if the step size is small, but it is slow, and it makes run-to-run comparisons noisy. If you compare policy gradient variants, use several seeds and report intervals; the evaluating RL agents page covers how.

A minimal estimator in numpy#

The function below computes the reward-to-go policy gradient for a tabular softmax policy from a batch of episodes. It is the estimator used for Figure 2, written for any number of states and actions.

import numpy as np

def policy_gradient(theta, episodes, gamma):
    """theta: (S, A) logits of a tabular softmax policy.
    episodes: list of (states, actions, rewards) integer/float arrays.
    Returns the reward-to-go estimate of grad J, shape (S, A)."""
    pi = np.exp(theta - theta.max(1, keepdims=True))
    pi /= pi.sum(1, keepdims=True)
    grad = np.zeros_like(theta)
    for s, a, r in episodes:
        T = len(r)
        G = np.zeros(T)                      # discounted reward-to-go
        running = 0.0
        for t in range(T - 1, -1, -1):
            running = r[t] + gamma * running
            G[t] = running
        w = gamma ** np.arange(T) * G        # keep the gamma^t factor
        score = -pi[s]                       # d log pi(a|s) / d theta[s, :]
        score[np.arange(T), a] += 1.0
        np.add.at(grad, s, w[:, None] * score)
    return grad / len(episodes)

For the score of a softmax policy, $\partial \log \pi (a\mid s)/\partial {\theta }_{s,b}=1[b=a]-\pi (b\mid s)$, which is what the two lines building score compute. Replace the softmax with a neural network and the score with an autodiff call, and this is the core of every policy gradient implementation.

What to take into practice#

• Use reward-to-go, not total return. It is unbiased and, in our example, cut variance 23-fold for free.
• Decide deliberately about ${\gamma }^t$. Keep it if you want the gradient of the discounted objective. Drop it only if you mean to optimize something closer to the undiscounted return and accept that the update is not an exact gradient.
• Budget samples by $1/\sqrt{N}$. Halving gradient noise costs four times as many episodes, which is why variance reduction pays better than brute force.
• Expect local optima. The theorem guarantees a gradient, not a global answer. Initialization, entropy bonuses and exploration decide which basin you end in.
• Treat the surrogate loss as plumbing. Monitor the return, not the loss, and keep advantages out of the computation graph.

Frequently asked questions#

Sources#

1. R. S. Sutton, A. G. Barto. Reinforcement Learning: An Introduction (2nd ed.), 2018. incompleteideas.net
2. R. S. Sutton, D. McAllester, S. Singh, Y. Mansour. Policy Gradient Methods for Reinforcement Learning with Function Approximation, 1999. proceedings.neurips.cc
3. R. J. Williams. Simple statistical gradient-following algorithms for connectionist reinforcement learning, 1992. doi.org
4. P. W. Glynn. Likelihood ratio gradient estimation for stochastic systems, 1990. doi.org
5. C. Nota, P. S. Thomas. Is the Policy Gradient a Gradient?, 2019. arxiv.org
6. S. Mohamed, M. Rosca, M. Figurnov, A. Mnih. Monte Carlo Gradient Estimation in Machine Learning, 2019. arxiv.org
7. D. Silver, G. Lever, N. Heess, T. Degris, D. Wierstra, M. Riedmiller. Deterministic Policy Gradient Algorithms, 2014. proceedings.mlr.press
8. J. Schulman, X. Chen, P. Abbeel. Equivalence Between Policy Gradients and Soft Q-Learning, 2017. arxiv.org