Lecture 09 · Part 2

Advanced Policy Gradients

Why does policy gradient actually work? Levine reinterprets it as a softened policy iteration, proves that maximizing the old policy's advantage under the new policy's trajectories maximizes the true RL objective, and shows that the substitution making this tractable is safe exactly when the new policy stays close to the old one. That turns policy gradient into a constrained optimization problem, enforced with a KL penalty (the route PPO takes) or approximated in closed form by the natural gradient (the route TRPO takes).

▶ Lecture videos (4 parts) Slides (current revision, PDF) ↗

Policy gradient is softened policy iteration

This is “one of the more technically nuanced lectures in the course,” and it answers a question the earlier lectures skipped: why should we expect policy gradients to improve the policy? The obvious answer — it’s a gradient, and gradient ascent works — hides most of the structure, and the analysis that reveals it carries over to actor-critic methods and RL algorithm design generally.

Recap: REINFORCE samples trajectories from the current policy, computes each step’s reward-to-go (Monte Carlo, or a learned value function from the actor-critic lecture), and steps along $\nabla_\theta J(\theta) \approx \frac{1}{N}\sum_i \sum_t \nabla_\theta \log \pi_\theta(a_{i,t} \mid s_{i,t})\, \hat{A}_{i,t}$ . Conceptually, though, every policy gradient method alternates two steps: estimate the advantage $\hat{A}^\pi(s_t, a_t)$ of the current policy, then use the estimate to get an improved policy $\pi'$ . That structure is exactly policy iteration from the value-based methods week — except policy iteration jumps straight to the $\arg\max$ , putting probability 1 on the highest-advantage action, while policy gradient “makes a much gentler update”: grad-log-pi times $\hat{A}$ nudges probability toward large-advantage actions without going all the way to one and zero.

The improvement lemma

The technical heart of the lecture is one identity. Work with discounted returns where the discount starts at step one — not the convention used in practice, but it makes the calculation clean.

Policy improvement as expected advantage

For any parameters $\theta$ (old) and $\theta'$ (new),

$J(\theta') - J(\theta) = \mathbb{E}_{\tau \sim p_{\theta'}(\tau)}\!\left[\sum_{t=0}^{\infty} \gamma^t A^{\pi_\theta}(s_t, a_t)\right]$

— the improvement in the RL objective equals the expected discounted advantage of the old policy under the new policy’s trajectory distribution. Since $J(\theta)$ is constant in $\theta'$ , maximizing the right-hand side maximizes $J(\theta')$ itself — which is precisely what policy iteration does, so proving the identity proves that policy iteration “does the right thing.”

Proof sketch. Write $J(\theta) = \mathbb{E}_{s_0 \sim p(s_0)}[V^{\pi_\theta}(s_0)]$ . Every policy induces the same initial-state marginal, so the expectation may be taken under $p_{\theta'}(\tau)$ instead. Then telescope: $V^{\pi_\theta}(s_0) = \sum_{t=0}^{\infty} \gamma^t V^{\pi_\theta}(s_t) - \sum_{t=1}^{\infty} \gamma^t V^{\pi_\theta}(s_t)$ . Flipping the sign and pairing terms off by one index gives $-J(\theta) = \mathbb{E}_{\tau \sim p_{\theta'}(\tau)}\!\left[\sum_{t=0}^{\infty} \gamma^t \left(\gamma V^{\pi_\theta}(s_{t+1}) - V^{\pi_\theta}(s_t)\right)\right]$ ; adding $J(\theta') = \mathbb{E}_{\tau \sim p_{\theta'}}[\sum_t \gamma^t r(s_t, a_t)]$ and grouping, the bracket becomes $r(s_t, a_t) + \gamma V^{\pi_\theta}(s_{t+1}) - V^{\pi_\theta}(s_t)$ — exactly $A^{\pi_\theta}(s_t, a_t)$ . $\square$

Now unpack the right-hand side into per-time-step state–action marginals and importance-sample the action expectation back to the old policy:

$\sum_t \mathbb{E}_{s_t \sim p_{\theta'}(s_t)}\!\left[\mathbb{E}_{a_t \sim \pi_\theta(a_t \mid s_t)}\!\left[\frac{\pi_{\theta'}(a_t \mid s_t)}{\pi_\theta(a_t \mid s_t)}\, \gamma^t A^{\pi_\theta}(s_t, a_t)\right]\right]$

One obstacle remains: the states are still distributed according to $p_{\theta'}(s_t)$ , which we cannot sample from — $\theta'$ is the thing we’re solving for. If we could replace $p_{\theta'}(s_t)$ with $p_\theta(s_t)$ — call the result $\bar{A}(\theta')$ — then differentiating at $\theta' = \theta$ recovers exactly the policy gradient, and policy gradient stands revealed as approximate policy iteration. Everything hinges on one claim: $p_\theta(s_t)$ is close to $p_{\theta'}(s_t)$ when $\pi_\theta$ is close to $\pi_{\theta'}$ . That sounds obvious; proving it non-vacuously is the next section’s work.

Bounding the distribution mismatch

Start with a deterministic $\pi_\theta$ and say $\pi_{\theta'}$ is close to it if it assigns probability at most $\epsilon$ to any action $\pi_\theta$ would not have taken. Then the new policy’s state marginal splits into two terms:

$p_{\theta'}(s_t) = (1 - \epsilon)^t\, p_\theta(s_t) + \left(1 - (1 - \epsilon)^t\right) p_{\text{mistake}}(s_t)$

If no mistake was made for $t$ steps — probability $(1-\epsilon)^t$ — the state distribution is exactly the old one; otherwise we land in some $p_{\text{mistake}}$ about which we assume nothing, like the tightrope walker from the imitation learning lecture who falls and can never recover. This is exactly the behavioral cloning analysis from lecture 2, and it yields the same bound: any total variation divergence is at most 2, and $(1-\epsilon)^t \ge 1 - \epsilon t$ for $\epsilon \in [0, 1]$ , so

$\left| p_{\theta'}(s_t) - p_\theta(s_t) \right| \le 2\left(1 - (1-\epsilon)^t\right) \le 2\,\epsilon\, t$

“It’s not a great bound, but it is a bound”: the marginals converge as $\epsilon \to 0$ .

For the general stochastic case — following the trust region policy optimization (TRPO) paper — define closeness by total variation: $|\pi_{\theta'}(a_t \mid s_t) - \pi_\theta(a_t \mid s_t)| \le \epsilon$ for all $s_t$ (an in-expectation version also works). A lemma from prior work does the heavy lifting: if two distributions have total variation divergence $\epsilon$ , there exists a joint distribution with those marginals under which the two variables agree with probability $1 - \epsilon$ . Geometrically: overlay the two distributions as bar graphs; the slivers where the bars differ have total volume $\epsilon$ , so the probability of acting differently is at most $\epsilon$ . That recovers the “probability of a mistake” structure of the deterministic case, and the same $2 \epsilon t$ bound.

To turn a bound on distributions into a bound on the objective, bound the expectation of any function $f$ :

$\mathbb{E}_{s_t \sim p_{\theta'}}[f(s_t)] \ge \mathbb{E}_{s_t \sim p_\theta}[f(s_t)] - 2\,\epsilon\, t \cdot \max_{s_t} f(s_t)$

(write $p_{\theta'} = p_\theta + (p_{\theta'} - p_\theta)$ and be maximally pessimistic on the difference — Levine’s advice: pause and derive this yourself). Taking $f$ to be the bracketed action expectation above, the tractable objective lower-bounds the true one up to an error $2 \epsilon t C$ , with $C$ the largest value the bracket can take. The importance weights and expectations are over probabilities that sum to one, and an advantage is a sum of rewards over time, so $C$ is of order $T\, r_{\max}$ for a finite horizon and $r_{\max} / (1 - \gamma)$ for an infinite one.

The conclusion: maximizing $\bar{A}(\theta')$ — whose derivative is exactly the policy gradient — is guaranteed to improve the true objective, provided $\theta'$ stays close to $\theta$ . Policy gradient is sound, but only as a constrained procedure.

Enforcing the constraint: KL penalties

Total variation constraints are awkward to impose (the absolute value isn’t differentiable everywhere). The KL divergence — “basically the most widely used divergence measure between distributions,” with tractable expressions and closed forms for common continuous distributions — bounds it:

$\left| \pi_{\theta'}(a_t \mid s_t) - \pi_\theta(a_t \mid s_t) \right| \le \sqrt{\tfrac{1}{2} D_{\mathrm{KL}}\!\left(\pi_{\theta'}(a_t \mid s_t)\, \|\, \pi_\theta(a_t \mid s_t)\right)}$

so constraining $D_{\mathrm{KL}}(\pi_{\theta'} \| \pi_\theta) \le \epsilon$ preserves every bound above. The practical problem becomes: maximize the importance-sampled objective $\bar{A}(\theta')$ subject to the KL constraint. One principled way in is the Lagrangian,

$\mathcal{L}(\theta', \lambda) = \bar{A}(\theta') - \lambda \left( D_{\mathrm{KL}}\!\left(\pi_{\theta'} \| \pi_\theta\right) - \epsilon \right)$

optimized by dual gradient descent: alternate maximizing $\mathcal{L}$ with respect to $\theta'$ (even incompletely, for a few gradient steps) with a gradient step on $\lambda$ — raise $\lambda$ when the constraint is violated, lower it when it’s slack. A more heuristic variant just fixes $\lambda$ by hand and treats the KL term as a regularizer penalizing deviation from the old policy. Both work well in practice; this family includes the original guided policy search method and PPO (proximal policy optimization).

The natural gradient

The last part trades exactness for simplicity: approximate the constrained problem so aggressively that a single linear-algebra step solves it. First, Taylor-expand the objective to first order around $\theta$ , giving the trust-region picture: optimize the linear approximation (the green line), but only inside the region where you trust it (the red box). The problem becomes

$\theta' \leftarrow \arg\max_{\theta'}\; \nabla_\theta \bar{A}(\theta)^\top (\theta' - \theta) \quad \text{s.t.} \quad D_{\mathrm{KL}}\!\left(\pi_{\theta'} \| \pi_\theta\right) \le \epsilon$

and the payoff of linearizing is that at $\theta' = \theta$ the importance weights cancel, so $\nabla_\theta \bar{A}(\theta) = \nabla_\theta J(\theta)$ — the standard policy gradient, with any advantage estimator you like plugged in.

Could plain gradient ascent handle the constraint? It does solve a constrained problem — but the wrong one: $\arg\max \nabla J^\top (\theta' - \theta)$ subject to $\|\theta' - \theta\|^2 \le \epsilon$ , a circle in parameter space, with the learning rate $\alpha = \sqrt{\epsilon / \|\nabla_\theta J(\theta)\|^2}$ emerging as its Lagrange multiplier. But the constraint we want lives in distribution space: some parameters barely move the policy’s probabilities while others move them enormously, so a small ball in $\theta$ can violate any KL budget. We want an ellipse — squished along directions where the distribution is sensitive, elongated where it isn’t.

So do to the constraint what we did to the objective: Taylor-expand. The KL’s first derivative at $\theta' = \theta$ is zero (it’s flat at its minimum), so go to second order:

$D_{\mathrm{KL}}\!\left(\pi_{\theta'} \| \pi_\theta\right) \approx \tfrac{1}{2} (\theta' - \theta)^\top \mathbf{F}\, (\theta' - \theta)$

With a linear objective and a quadratic constraint, the Lagrangian is a second-order polynomial with a closed-form maximizer:

$\theta' = \theta + \alpha\, \mathbf{F}^{-1} \nabla_\theta J(\theta), \qquad \alpha = \sqrt{\frac{2\epsilon}{\nabla_\theta J(\theta)^\top \mathbf{F}^{-1} \nabla_\theta J(\theta)}}$

This is the natural gradient: the regular gradient premultiplied by the inverse Fisher matrix, with $\alpha$ chosen (if you wish) to satisfy the KL budget $\epsilon$ exactly — plug the update into the quadratic constraint and check the algebra.

Worked example — The (k, sigma) pathology, resolved

Recall the example from the policy gradient lecture: 1D state and action, reward $-s^2 - a^2$ , policy $\pi_\theta(a \mid s) = \mathcal{N}(k s, \sigma)$ . The vanilla gradient always points in an ascent direction but is “extremely ill-conditioned”: as $\sigma$ shrinks it affects the probabilities far more than $k$ does, its gradient component dwarfs $k$ ‘s, and you never reach the optimum $k = -1$ — you’re too busy shrinking $\sigma$ . A fixed ball in $\theta$ space, “even in this very simple two-dimensional space, is just not enough.” Multiplying by $\mathbf{F}^{-1}$ accounts for $\sigma$ ‘s outsized influence, boosts the $k$ -component, and points the field at the optimum — a real practical fix, especially for continuous Gaussian policies.

Practical notes, as Levine gives them: classical natural policy gradient estimates $\mathbf{F}$ from samples and picks $\alpha$ by hand — a good stabilizer in general (see Peters and Schaal, “Reinforcement learning of motor skills with policy gradients”). TRPO instead picks $\epsilon$ , solves for $\alpha$ , and uses the conjugate gradient method to compute $\mathbf{F}^{-1} \nabla_\theta J$ without ever forming the full matrix. Or skip the approximation entirely and optimize the importance-sampled objective for many gradient steps with a KL regularizer — the path that leads to PPO.

Back of the chalk tray — what to remember

Policy gradient is softened policy iteration: estimate the current policy’s advantage, then gently shift probability toward high-advantage actions instead of jumping to the $\arg\max$ .
The improvement lemma $J(\theta') - J(\theta) = \mathbb{E}_{\tau \sim p_{\theta'}}[\sum_t \gamma^t A^{\pi_\theta}(s_t, a_t)]$ — proved by a telescoping sum on values — says maximizing the old policy’s advantage under the new policy maximizes the true objective.
The only intractable piece is sampling states from $p_{\theta'}$ ; substituting $p_\theta$ costs an error of at most $2\epsilon t C$ when the policies are within total variation $\epsilon$ — the same bound, via the same argument, as behavioral cloning’s.
Policy gradient is therefore only sound as a constrained procedure: improvement is guaranteed for small enough policy divergence, in KL form $D_{\mathrm{KL}}(\pi_{\theta'} \| \pi_\theta) \le \epsilon$ .
Enforce the constraint with a Lagrangian and dual gradient descent (or a fixed KL penalty) — the PPO route — or linearize the objective and quadratify the constraint to get the closed-form natural gradient $\theta' = \theta + \alpha \mathbf{F}^{-1} \nabla_\theta J(\theta)$ — the TRPO route.
Vanilla gradient ascent also solves a constrained problem, but with a ball in parameter space instead of distribution space — the wrong constraint, and the source of the $(k, \sigma)$ ill-conditioning.

Check yourself

Derivation-heavy exercises in the lecture’s spirit — pen and paper recommended.

Problem 9.1 telescoping improvement lemma

Prove the improvement lemma $J(\theta') - J(\theta) = \mathbb{E}_{\tau \sim p_{\theta'}(\tau)}\left[\sum_{t=0}^{\infty} \gamma^t A^{\pi_\theta}(s_t, a_t)\right]$ . Identify (a) the step where the expectation switches from $p(s_0)$ to $p_{\theta'}(\tau)$ and why it is legal, and (b) where the telescoping happens.

Show solution

Start from $J(\theta) = \mathbb{E}_{s_0 \sim p(s_0)}[V^{\pi_\theta}(s_0)]$ . (a) The quantity inside depends only on $s_0$ , and every policy’s trajectory distribution has the same initial-state marginal $p(s_0)$ , so the expectation may be taken under $p_{\theta'}(\tau)$ without changing its value — this is the step that plants the new policy’s distribution into an expression about the old policy’s value. (b) Replace $V^{\pi_\theta}(s_0)$ with $\sum_{t=0}^\infty \gamma^t V^{\pi_\theta}(s_t) - \sum_{t=1}^\infty \gamma^t V^{\pi_\theta}(s_t)$ : all terms with $t \ge 1$ cancel, so the two infinite sums equal $V^{\pi_\theta}(s_0)$ exactly. Flip the sign of $J(\theta)$ , pair the term at $t+1$ (which carries an extra $\gamma$ ) with the term at $t$ , and pull out $\gamma^t$ to get $-J(\theta) = \mathbb{E}_{p_{\theta'}}[\sum_t \gamma^t(\gamma V^{\pi_\theta}(s_{t+1}) - V^{\pi_\theta}(s_t))]$ . Adding $J(\theta') = \mathbb{E}_{p_{\theta'}}[\sum_t \gamma^t r(s_t, a_t)]$ and merging the sums leaves $\mathbb{E}_{p_{\theta'}}[\sum_t \gamma^t (r(s_t, a_t) + \gamma V^{\pi_\theta}(s_{t+1}) - V^{\pi_\theta}(s_t))]$ , and the bracket is the definition of $A^{\pi_\theta}(s_t, a_t)$ .

Problem 9.2 how loose is the bound

Policies $\pi_\theta$ and $\pi_{\theta'}$ satisfy the total variation closeness condition with $\epsilon = 0.001$ at every state, over a horizon $T = 100$ with $r_{\max} = 1$ . (a) Bound the total variation between the state marginals at the final step. (b) Bound the error term $2 \epsilon t C$ of the surrogate objective at $t = T$ , using $C = T\, r_{\max}$ . (c) Compare (b) to the largest possible return and comment on how the bound scales with horizon.

Show solution

(a) $2 \epsilon t = 2 \times 0.001 \times 100 = 0.2$ . (b) $2 \epsilon t C = 2 \times 0.001 \times 100 \times (100 \times 1) = 20$ . (c) The largest possible return is $T r_{\max} = 100$ , so the worst-case error is 20% of the maximum achievable return even with a policy that deviates with probability only one in a thousand per step. The culprit is the $t \cdot C \sim T^2 r_{\max}$ scaling: one factor of $T$ from compounding state-distribution drift (the behavioral cloning argument), another from the advantage being a sum of up to $T$ rewards. The bound is loose but honest — it vanishes linearly as $\epsilon \to 0$ , which is all the constrained-update argument needs, and it explains why trust regions must shrink as horizons grow.

Problem 9.3 step sizes satisfy constraints

(a) Show that the gradient ascent step $\theta' = \theta + \alpha \nabla_\theta J(\theta)$ with $\alpha = \sqrt{\epsilon / \|\nabla_\theta J(\theta)\|^2}$ satisfies $\|\theta' - \theta\|^2 \le \epsilon$ with equality. (b) Show that the natural gradient step $\theta' = \theta + \alpha \mathbf{F}^{-1} \nabla_\theta J(\theta)$ with $\alpha = \sqrt{2\epsilon / (\nabla_\theta J^\top \mathbf{F}^{-1} \nabla_\theta J)}$ satisfies the quadratized KL constraint $\tfrac{1}{2}(\theta' - \theta)^\top \mathbf{F} (\theta' - \theta) \le \epsilon$ with equality.

Show solution

(a) $\|\theta' - \theta\|^2 = \alpha^2 \|\nabla J\|^2 = \frac{\epsilon}{\|\nabla J\|^2} \|\nabla J\|^2 = \epsilon$ . (b) Write $g = \nabla_\theta J$ . Then $\theta' - \theta = \alpha \mathbf{F}^{-1} g$ and $\tfrac{1}{2}(\alpha \mathbf{F}^{-1} g)^\top \mathbf{F} (\alpha \mathbf{F}^{-1} g) = \tfrac{1}{2} \alpha^2\, g^\top \mathbf{F}^{-1} \mathbf{F} \mathbf{F}^{-1} g = \tfrac{1}{2} \alpha^2\, g^\top \mathbf{F}^{-1} g$ (using symmetry of $\mathbf{F}$ ). Substituting $\alpha^2 = 2\epsilon / (g^\top \mathbf{F}^{-1} g)$ gives $\tfrac{1}{2} \cdot \frac{2\epsilon}{g^\top \mathbf{F}^{-1} g} \cdot g^\top \mathbf{F}^{-1} g = \epsilon$ . In both cases the maximizer of a linear objective over a convex constraint region sits on the boundary, which is why equality is the right answer — the step uses the entire budget $\epsilon$ in the steepest useful direction.

Problem 9.4 Fisher preconditioning by hand

A two-parameter policy has, at the current $\theta$ , Fisher matrix $\mathbf{F} = \mathrm{diag}(100, 1)$ and policy gradient $\nabla_\theta J = (1, 1)^\top$ . (a) Compare the directions of the vanilla gradient step and the natural gradient step. (b) Interpret what $\mathbf{F}_{11} = 100$ says about parameter 1, and explain how this mirrors the $(k, \sigma)$ example.

Show solution

(a) The vanilla step moves along $(1, 1)$ — equal effort in both coordinates. The natural step moves along $\mathbf{F}^{-1} \nabla J = (\tfrac{1}{100} \cdot 1,\ 1 \cdot 1)^\top = (0.01, 1)^\top$ — almost entirely along parameter 2. (b) A large diagonal Fisher entry means small changes in that parameter produce large changes in the policy’s action probabilities (the expected squared grad-log-prob in that direction is big). The natural gradient therefore takes small steps in high-sensitivity directions and large steps in low-sensitivity ones — per-direction step sizes chosen by distribution geometry rather than parameter geometry. In the $(k, \sigma)$ example, $\sigma$ plays the role of parameter 1: as $\sigma \to 0$ it dominates the probabilities, the vanilla gradient pours all its effort into shrinking it further, and $k$ never reaches $-1$ ; premultiplying by $\mathbf{F}^{-1}$ suppresses the $\sigma$ direction and boosts $k$ , pointing the update at the optimum.