Policy gradient theorem

Suppose the policy is a neural network with parameters $\theta$ that outputs a distribution over actions. The objective is the expected return $J(\theta)$ . If you could compute $\nabla_\theta J$ you could do gradient ascent and be done. The difficulty is that $J$ is an expectation over trajectories, and the distribution over trajectories depends on $\theta$ in two ways: through the actions chosen, which is fine, and through which states are visited as a consequence of those actions, which is not fine, because the state visitation distribution is a complicated function of $\theta$ that runs through the unknown dynamics.

The policy gradient theorem says you can ignore the second dependence. The gradient of the expected return is an expectation, under the current policy's own state distribution, of the score $\nabla_\theta \log \pi_\theta(a \mid s)$ weighted by the action value $Q^{\pi_\theta}(s, a)$ . Every piece of that can be sampled by running the policy. The intuition for why the state distribution's gradient disappears is that it is already accounted for: the value $Q^{\pi_\theta}(s, a)$ includes the downstream consequences of the action, so the effect of changing the policy on where you go later is priced into the weight on each action, not carried as a separate term.

The resulting update is easy to read. Increase the probability of actions that led to high value, decrease the probability of actions that led to low value, in proportion to how much better or worse they were than expected.

Setup

Let $\pi_\theta(a \mid s)$ be differentiable in $\theta \in \mathbb{R}^d$ with $\pi_\theta(a \mid s) > 0$ wherever it matters. The objective is the expected discounted return from the start distribution,

J(\theta) = \mathbb{E}_{s_0 \sim \rho_0}\bigl[V^{\pi_\theta}(s_0)\bigr] = \mathbb{E}_{\tau \sim \pi_\theta}\Bigl[\sum_{t=0}^{\infty} \gamma^t r_{t+1}\Bigr].

Define the discounted state visitation distribution

d^{\pi_\theta}(s) = (1 - \gamma) \sum_{t=0}^{\infty} \gamma^t \Pr(s_t = s \mid s_0 \sim \rho_0, \pi_\theta).

The factor $(1 - \gamma)$ normalizes it to a probability distribution, since $\sum_s \Pr(s_t = s) = 1$ for each $t$ and $\sum_t \gamma^t = 1 / (1 - \gamma)$ . It weights states by how soon and how often the policy visits them.

The theorem

Theorem (Sutton, McAllester, Singh, Mansour 1999).

\nabla_\theta J(\theta) = \frac{1}{1 - \gamma}\; \mathbb{E}_{s \sim d^{\pi_\theta},\, a \sim \pi_\theta(\cdot \mid s)}\bigl[\nabla_\theta \log \pi_\theta(a \mid s)\; Q^{\pi_\theta}(s, a)\bigr].

The constant $1 / (1 - \gamma)$ is often absorbed into the step size and the result written as a proportionality; I keep it so that the two proofs below agree exactly.

First proof: unrolling the Bellman equation

This is the original proof. It differentiates the value function's recursive definition and unrolls.

Step 1. Differentiate $V^{\pi_\theta}$ . Start from $V^{\pi_\theta}(s) = \sum_a \pi_\theta(a \mid s)\, Q^{\pi_\theta}(s, a)$ and apply the product rule:

\nabla_\theta V^{\pi_\theta}(s) = \sum_a \Bigl[\nabla_\theta \pi_\theta(a \mid s)\, Q^{\pi_\theta}(s, a) + \pi_\theta(a \mid s)\, \nabla_\theta Q^{\pi_\theta}(s, a)\Bigr].

For the second term, $Q^{\pi_\theta}(s, a) = R(s, a) + \gamma \sum_{s'} P(s' \mid s, a) V^{\pi_\theta}(s')$ , and neither $R$ nor $P$ depends on $\theta$ , so

\nabla_\theta Q^{\pi_\theta}(s, a) = \gamma \sum_{s'} P(s' \mid s, a)\, \nabla_\theta V^{\pi_\theta}(s').

Step 2. Write it as a recursion. Let $\Phi(s) = \sum_a \nabla_\theta \pi_\theta(a \mid s)\, Q^{\pi_\theta}(s, a)$ and let $P_\theta(s' \mid s) = \sum_a \pi_\theta(a \mid s) P(s' \mid s, a)$ be the policy induced transition kernel. Then

\nabla_\theta V^{\pi_\theta}(s) = \Phi(s) + \gamma \sum_{s'} P_\theta(s' \mid s)\, \nabla_\theta V^{\pi_\theta}(s').

This has the same shape as a Bellman equation with $\Phi$ playing the role of the reward. Note that the gradient of the transition kernel $P_\theta$ never appears: the recursion was derived from the value function's structure, and the only place $\theta$ enters explicitly is through $\pi_\theta$ inside $\Phi$ .

Step 3. Unroll. Substitute the recursion into itself $k$ times:

\nabla_\theta V^{\pi_\theta}(s) = \sum_{t=0}^{k-1} \gamma^t \sum_{s'} \Pr(s_t = s' \mid s_0 = s, \pi_\theta)\, \Phi(s') + \gamma^k \sum_{s'} \Pr(s_k = s' \mid s_0 = s, \pi_\theta)\, \nabla_\theta V^{\pi_\theta}(s').

The tail term vanishes as $k \to \infty$ : $\nabla_\theta V^{\pi_\theta}$ is bounded (rewards are bounded and the score is assumed bounded), and $\gamma^k \to 0$ . So

\nabla_\theta V^{\pi_\theta}(s) = \sum_{t=0}^{\infty} \gamma^t \sum_{s'} \Pr(s_t = s' \mid s_0 = s, \pi_\theta)\, \Phi(s').

Step 4. Average over the start state. Taking $\mathbb{E}_{s_0 \sim \rho_0}$ on both sides,

\nabla_\theta J(\theta) = \sum_{s'} \Bigl[\sum_{t=0}^{\infty} \gamma^t \Pr(s_t = s' \mid \rho_0, \pi_\theta)\Bigr] \Phi(s') = \frac{1}{1 - \gamma} \sum_{s'} d^{\pi_\theta}(s')\, \Phi(s').

Step 5. The log derivative trick. Since $\nabla_\theta \pi_\theta(a \mid s) = \pi_\theta(a \mid s)\, \nabla_\theta \log \pi_\theta(a \mid s)$ wherever $\pi_\theta > 0$ ,

\Phi(s) = \sum_a \pi_\theta(a \mid s)\, \nabla_\theta \log \pi_\theta(a \mid s)\, Q^{\pi_\theta}(s, a) = \mathbb{E}_{a \sim \pi_\theta(\cdot \mid s)}\bigl[\nabla_\theta \log \pi_\theta(a \mid s)\, Q^{\pi_\theta}(s, a)\bigr].

Substituting into Step 4 gives the theorem. $\blacksquare$

The moment to notice is Step 2. The state distribution depends on $\theta$ , and yet $\nabla_\theta P_\theta$ never had to be computed, because the derivative was taken of the value function's recursive definition rather than of the trajectory distribution. The unrolling then reconstructs the state distribution from the policy induced kernel without ever differentiating it.

Second proof: the likelihood ratio on trajectories

The same result from a different direction, which is the one most implementations are actually built on. Work with finite trajectories of length $T$ for clarity; the infinite horizon case follows by letting $T \to \infty$ with the discount.

A trajectory $\tau = (s_0, a_0, s_1, a_1, \dots, s_T)$ has probability

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),

and return $G(\tau) = \sum_{t=0}^{T-1} \gamma^t r_{t+1}$ . The objective is $J(\theta) = \sum_\tau p_\theta(\tau)\, G(\tau)$ .

Step 1. Differentiate under the sum and apply the log derivative trick to the whole trajectory.

\nabla_\theta J(\theta) = \sum_\tau \nabla_\theta p_\theta(\tau)\, G(\tau) = \sum_\tau p_\theta(\tau)\, \nabla_\theta \log p_\theta(\tau)\, G(\tau) = \mathbb{E}_{\tau \sim \pi_\theta}\bigl[\nabla_\theta \log p_\theta(\tau)\, G(\tau)\bigr].

Step 2. The dynamics drop out of the log. Taking the log of $p_\theta(\tau)$ turns the product into a sum, and only the policy factors depend on $\theta$ :

\nabla_\theta \log p_\theta(\tau) = \nabla_\theta \log \rho_0(s_0) + \sum_{t=0}^{T-1} \Bigl[\nabla_\theta \log \pi_\theta(a_t \mid s_t) + \nabla_\theta \log P(s_{t+1} \mid s_t, a_t)\Bigr] = \sum_{t=0}^{T-1} \nabla_\theta \log \pi_\theta(a_t \mid s_t).

This is the second proof's version of the same miracle. The transition probabilities are inside the trajectory likelihood, but their gradient with respect to $\theta$ is zero, so they vanish from the score of the trajectory. You never need to know $P$ to compute the gradient. Hence

\nabla_\theta J(\theta) = \mathbb{E}_{\tau \sim \pi_\theta}\Bigl[\Bigl(\sum_{t=0}^{T-1} \nabla_\theta \log \pi_\theta(a_t \mid s_t)\Bigr) G(\tau)\Bigr].

Step 3. Causality: an action cannot affect rewards that came before it. Expand $G(\tau) = \sum_{k=0}^{T-1} \gamma^k r_{k+1}$ and consider a single cross term $\mathbb{E}[\nabla_\theta \log \pi_\theta(a_t \mid s_t)\, r_{k+1}]$ with $k < t$ . Condition on everything up to and including $s_t$ ; the reward $r_{k+1}$ is then already determined, and

\mathbb{E}\bigl[\nabla_\theta \log \pi_\theta(a_t \mid s_t) \mid s_t\bigr] = \sum_a \pi_\theta(a \mid s_t)\, \nabla_\theta \log \pi_\theta(a \mid s_t) = \nabla_\theta \sum_a \pi_\theta(a \mid s_t) = \nabla_\theta 1 = 0.

So every term pairing an action with an earlier reward has expectation zero and can be dropped:

\nabla_\theta J(\theta) = \mathbb{E}_{\tau \sim \pi_\theta}\Bigl[\sum_{t=0}^{T-1} \nabla_\theta \log \pi_\theta(a_t \mid s_t) \sum_{k=t}^{T-1} \gamma^k r_{k+1}\Bigr] = \mathbb{E}_{\tau \sim \pi_\theta}\Bigl[\sum_{t=0}^{T-1} \gamma^t\, \nabla_\theta \log \pi_\theta(a_t \mid s_t)\, G_t\Bigr],

where $G_t = \sum_{k=t}^{T-1} \gamma^{k-t} r_{k+1}$ is the reward to go from time $t$ .

Step 4. Replace the return by its conditional expectation. Conditioning on $(s_t, a_t)$ , $\mathbb{E}[G_t \mid s_t, a_t] = Q^{\pi_\theta}(s_t, a_t)$ by definition, so

\nabla_\theta J(\theta) = \mathbb{E}_{\tau \sim \pi_\theta}\Bigl[\sum_{t=0}^{T-1} \gamma^t\, \nabla_\theta \log \pi_\theta(a_t \mid s_t)\, Q^{\pi_\theta}(s_t, a_t)\Bigr].

Finally, the sum over $t$ weighted by $\gamma^t$ of an expectation over $s_t$ is $1 / (1 - \gamma)$ times an expectation over $s \sim d^{\pi_\theta}$ , by the definition of the discounted visitation distribution, which recovers the theorem. $\blacksquare$

The two proofs establish the same identity by different routes, and each makes a different feature visible. The first shows that the gradient of the state distribution is never needed. The second shows that the gradient of the dynamics is never needed, and gives the trajectory form that is actually implemented.

REINFORCE

The trajectory form in Step 3 of the second proof is an expectation over trajectories generated by the policy, so it can be estimated by sampling. That estimator is REINFORCE (Williams 1992).

Algorithm. Repeat:

Run the policy for one episode, collecting $(s_0, a_0, r_1, \dots, s_{T-1}, a_{T-1}, r_T)$ .
For each $t$ , compute the reward to go $G_t = \sum_{k=t}^{T-1} \gamma^{k-t} r_{k+1}$ .
Update

\theta \leftarrow \theta + \alpha \sum_{t=0}^{T-1} \gamma^t\, \nabla_\theta \log \pi_\theta(a_t \mid s_t)\, G_t.

Proposition. The REINFORCE update direction is an unbiased estimate of $\nabla_\theta J(\theta)$ .

Proof. The update is a single sample of the random variable inside the expectation in Step 3, and Step 3 established that the expectation of that random variable is $\nabla_\theta J(\theta)$ . $\square$

That is the whole justification, and it is worth appreciating how little it requires: a differentiable policy, the ability to run it, and nothing about the environment. The cost is variance. $G_t$ is a sum over the entire remaining trajectory and fluctuates with every future action and transition, so the estimator is unbiased but extremely noisy, and REINFORCE in its raw form needs many episodes per useful step.

Two remarks on the algorithm as written. In practice the $\gamma^t$ factor in front of the score is usually dropped, which makes the estimator biased toward the undiscounted objective but reduces variance for later time steps; this is one of the places theory and practice quietly part ways. And the update is a gradient ascent, since $J$ is to be maximized.

Baselines

The most important variance reduction costs nothing in bias.

Lemma. For any function $b(s)$ of the state alone,

\mathbb{E}_{a \sim \pi_\theta(\cdot \mid s)}\bigl[\nabla_\theta \log \pi_\theta(a \mid s)\, b(s)\bigr] = 0.

Proof. The baseline does not depend on $a$ , so it factors out of the expectation:

b(s) \sum_a \pi_\theta(a \mid 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. \qquad \square

This is the same computation as the causality step, and it means $Q^{\pi_\theta}(s, a)$ can be replaced by $Q^{\pi_\theta}(s, a) - b(s)$ in the theorem, or $G_t$ by $G_t - b(s_t)$ in REINFORCE, without changing the expectation.

Why it reduces variance. The gradient estimator at a state is $\nabla_\theta \log \pi_\theta(a \mid s)$ times a scalar weight. If the weight is large and positive for every action, as it is when all returns are positive, every sampled action gets pushed up, and the useful signal, that some actions are better than others, is a small difference between large numbers. Subtracting a baseline centers the weights so that better than average actions are pushed up and worse than average actions are pushed down. Formally, the variance of $\nabla \log \pi \cdot (Q - b)$ is minimized over constant $b$ at a weighted average of $Q$ , and the state value $V^{\pi_\theta}(s)$ is close to that optimum and has a clear meaning.

With $b(s) = V^{\pi_\theta}(s)$ the weight becomes the advantage,

A^{\pi_\theta}(s, a) = Q^{\pi_\theta}(s, a) - V^{\pi_\theta}(s),

how much better action $a$ is than the policy's average behavior in $s$ . The advantage form of the theorem,

\nabla_\theta J(\theta) = \frac{1}{1 - \gamma}\; \mathbb{E}_{s \sim d^{\pi_\theta},\, a \sim \pi_\theta}\bigl[\nabla_\theta \log \pi_\theta(a \mid s)\, A^{\pi_\theta}(s, a)\bigr],

is the one every modern method starts from.

Actor critic

The baseline $V^{\pi_\theta}$ is not known, so estimate it. An actor critic method keeps two function approximators: the actor $\pi_\theta$ , updated by the policy gradient, and the critic $V_\phi$ , updated by Temporal difference learning. The advantage is then estimated by bootstrapping,

\hat{A}_t = r_{t+1} + \gamma V_\phi(s_{t+1}) - V_\phi(s_t),

which is exactly the TD error of the critic. This replaces the Monte Carlo return $G_t$ with a one step bootstrapped estimate, trading the high variance of REINFORCE for the bias of an imperfect critic. Generalized advantage estimation (Schulman et al. 2016) interpolates between the two with a parameter $\lambda$ , in the same way TD( $\lambda$ ) interpolates between TD(0) and Monte Carlo.

The trust region and proximal methods, TRPO and PPO, keep the same gradient and constrain how far each update moves the policy, since the theorem is a statement about the local gradient and a large step can leave the region where it is informative. PPO is what I used for the verifiable rewards project in the lab, and its clipped objective is a direct descendant of the advantage form above.

Where this breaks

The theorem is exact, but every estimator of it is noisy, and variance rather than bias is the practical limit. Baselines, bootstrapped critics and large batches are all responses to the same problem.
The discounted state distribution $d^{\pi_\theta}$ is what the theory asks for, and almost no implementation samples from it. Implementations average over the states actually visited, which is the undiscounted distribution, and drop the $\gamma^t$ weight. The gradient being followed is therefore not quite the gradient of $J$ . In practice this is fine; in theory it is a known gap (Nota and Thomas 2020).
The score function must be well behaved. Policies that put probability near zero on actions produce huge scores, which is one reason entropy regularization and clipping are common.
Nothing here guarantees convergence to a global optimum. Gradient ascent on a nonconvex $J$ finds a stationary point, and with function approximation there are results on convergence to local optima under step size conditions (Sutton et al. 1999; Konda and Tsitsiklis 2000) but not more.

References

Williams. Simple Statistical Gradient-Following Algorithms for Connectionist Reinforcement Learning. Machine Learning, 1992.
Sutton, McAllester, Singh, Mansour. Policy Gradient Methods for Reinforcement Learning with Function Approximation. NeurIPS, 1999.
Konda, Tsitsiklis. Actor-Critic Algorithms. NeurIPS, 2000.
Schulman, Levine, Moritz, Jordan, Abbeel. Trust Region Policy Optimization. ICML, 2015.
Schulman, Moritz, Levine, Jordan, Abbeel. High-Dimensional Continuous Control Using Generalized Advantage Estimation. ICLR, 2016.
Schulman, Wolski, Dhariwal, Radford, Klimov. Proximal Policy Optimization Algorithms. arXiv, 2017.
Nota, Thomas. Is the Policy Gradient a Gradient? AAMAS, 2020.
Sutton, Barto. Reinforcement Learning: An Introduction, 2nd ed. MIT Press, 2018, chapter 13.