Lecture 05 · Part 1

Policy Gradients

The first actual RL algorithm of the course. Levine differentiates the RL objective directly — one identity turns an intractable gradient into something you can estimate from rollouts — then spends the rest of the lecture on why the naive estimator barely works and what fixes it: causality and reward-to-go, baselines, importance sampling for off-policy data, the pseudo-loss trick for autodiff, and the natural gradient.

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

The objective, and how to evaluate it from rollouts

Policy gradients are “in some ways kind of the simplest reinforcement learning algorithm”: directly differentiate the RL objective and do gradient ascent on the policy parameters. The policy $\pi_\theta(a_t \mid s_t)$ — a network with weights $\theta$ — induces, together with the unknown dynamics, a trajectory distribution by the chain rule:

$p_\theta(\tau) = p(s_1) \prod_{t=1}^{T} \pi_\theta(a_t \mid s_t)\, p(s_{t+1} \mid s_t, a_t)$

with $\tau = (s_1, a_1, \ldots, s_T, a_T)$ . The objective is $J(\theta) = \mathbb{E}_{\tau \sim p_\theta(\tau)}\left[\sum_t r(s_t, a_t)\right]$ , and this lecture works with the finite-horizon version.

Crucially, model-free RL never assumes we know $p(s_1)$ or the transition probabilities — only that we can run the policy, which samples from them. That is already enough to evaluate $J(\theta)$ : roll out $N$ trajectories and average their total rewards. The whole game is doing the same for the derivative.

One identity makes the gradient estimable

Write $r(\tau)$ for the total reward of a trajectory, so $J(\theta) = \int p_\theta(\tau)\, r(\tau)\, d\tau$ and $\nabla_\theta J(\theta) = \int \nabla_\theta p_\theta(\tau)\, r(\tau)\, d\tau$ — useless as written, since $\nabla_\theta p_\theta(\tau)$ touches the unknown dynamics. Then comes what Levine calls “basically the only piece of mathematical cleverness in this whole derivation”:

The log-gradient identity, and the policy gradient

From $\nabla \log p = \nabla p / p$ , we get $p_\theta(\tau)\, \nabla_\theta \log p_\theta(\tau) = \nabla_\theta p_\theta(\tau)$ . Substituting turns the gradient back into an expectation, and when you expand $\log p_\theta(\tau)$ , the $\log p(s_1)$ and transition terms vanish under $\nabla_\theta$ — they don’t depend on $\theta$ . What survives is exactly the part we can compute:

$\nabla_\theta J(\theta) = \mathbb{E}_{\tau \sim p_\theta(\tau)}\!\left[\left(\sum_{t=1}^{T} \nabla_\theta \log \pi_\theta(a_t \mid s_t)\right)\left(\sum_{t=1}^{T} r(s_t, a_t)\right)\right]$

All the unknowns live only in the sampling distribution — so estimate with rollouts, then step: $\theta \leftarrow \theta + \alpha \nabla_\theta J(\theta)$ . That three-step loop (sample, estimate the gradient, step) is REINFORCE, named by Williams in the 1990s.

What the gradient is actually doing

For a discrete policy, $\nabla_\theta \log \pi_\theta$ is exactly the gradient you’d compute in supervised maximum likelihood — except maximum likelihood raises the log probability of every observed action, while the policy gradient weights each trajectory’s log-likelihood gradient by its reward. High-reward trajectories get their probabilities pushed up; low-reward ones get pushed down.

For continuous actions, take $\pi_\theta(a_t \mid s_t) = \mathcal{N}(f_{\text{NN}}(s_t), \Sigma)$ ; the log-probability is a quadratic in $f(s_t) - a_t$ , and its gradient $-\tfrac{1}{2}\Sigma^{-1}(f(s_t) - a_t)\frac{df}{d\theta}$ backpropagates through the mean network. And a notable freebie: the derivation never used the Markov property, so the same equation holds verbatim in partially observed settings — replace $s_t$ with $o_t$ and run it.

The catch is variance. With three samples and rewards drawn as bars on a number line, where the distribution moves depends alarmingly on things that shouldn’t matter: add a constant to every reward — which provably changes no optimal policy — and a finite-sample gradient estimate moves differently. In expectation the estimator is always correct; for any finite $N$ it is noisy enough that “a lot of advances in policy gradient algorithms basically revolve around different ways to reduce their variance.”

Two variance reducers: causality and baselines

Causality. Distribute the reward sum into the grad-log-pi sum and each action’s term multiplies rewards from all time steps — including the past, which the action provably cannot affect (“this is not an assumption… the only way this would not be true is if you had time travel”). Those past-reward terms cancel in expectation, so drop them:

$\nabla_\theta J(\theta) \approx \frac{1}{N}\sum_{i=1}^{N}\sum_{t=1}^{T} \nabla_\theta \log \pi_\theta(a_{i,t} \mid s_{i,t})\, \hat{Q}_{i,t}, \qquad \hat{Q}_{i,t} = \sum_{t'=t}^{T} r(s_{i,t'}, a_{i,t'})$

$\hat{Q}_{i,t}$ is the reward-to-go — a single-sample estimate of the Q-function, foreshadowing actor-critic. Fewer terms in the sum means smaller numbers means lower variance; the estimator stays unbiased. You always use this.

Baselines. Subtract $b$ from the reward before weighting, so better-than-average trajectories go up and worse-than-average go down regardless of whether raw rewards are all positive:

Baselines are unbiased

$\mathbb{E}\left[\nabla_\theta \log p_\theta(\tau)\, b\right] = \int \nabla_\theta p_\theta(\tau)\, b\, d\tau = b\, \nabla_\theta \int p_\theta(\tau)\, d\tau = b\, \nabla_\theta 1 = 0$ . Subtracting any constant $b$ leaves the expected gradient unchanged — it only moves the variance. The average reward is the “quick and dirty” choice; setting $\frac{d\,\text{Var}}{db} = 0$ gives the optimal per-parameter baseline $b = \mathbb{E}[g^2 r] / \mathbb{E}[g^2]$ (with $g = \nabla_\theta \log p_\theta(\tau)$ ) — expected reward reweighted by squared gradient magnitude, rarely used in practice but instructive to derive.

Going off-policy with importance sampling

Policy gradients are the classic on-policy algorithm: the expectation is under $p_\theta(\tau)$ , so every gradient step demands fresh samples — painful when deep nets need many small steps and samples are expensive. Importance sampling relaxes this. For samples from any $\bar{p}(\tau)$ :

$\mathbb{E}_{x \sim p}[f(x)] = \mathbb{E}_{x \sim q}\!\left[\frac{p(x)}{q(x)} f(x)\right]$

and in the trajectory ratio the unknown initial-state and transition terms cancel, leaving only the product of policy ratios $\prod_t \frac{\pi_{\theta'}(a_t \mid s_t)}{\pi_\theta(a_t \mid s_t)}$ — computable. The trouble: a product of $T$ numbers each below 1 shrinks exponentially, so the importance weights blow variance up “exponentially fast.” The escape hatch is to write the weight over state–action marginals and ignore the state-marginal ratio, keeping only the current action’s ratio. That’s no longer the exact gradient, but the error is bounded when $\theta'$ stays close to $\theta$ — the insight behind practical algorithms like TRPO and PPO, developed properly in the advanced policy gradients lecture.

Implementing it: the pseudo-loss trick

Naively computing $\nabla_\theta \log \pi_\theta$ per state–action pair means ~10,000 million-entry vectors. Instead, give your autodiff package a graph whose gradient is the policy gradient:

Code walkthrough — Weighted maximum likelihood pseudo-loss

Implement maximum likelihood, then weight each sample’s log-likelihood by its reward-to-go:

# maximum likelihood (supervised learning):
logits = policy.predictions(states)
nll = softmax_cross_entropy(logits, actions)
loss = mean(nll)

# policy gradient: same graph, weighted by reward-to-go
weighted_nll = nll * q_values      # q_values: (N*T, 1), computed outside the graph
loss = mean(weighted_nll)
loss.backward()                    # gradient IS the policy gradient

$\tilde{J}(\theta) = \sum_{i,t} \log \pi_\theta(a_{i,t} \mid s_{i,t})\, \hat{Q}_{i,t}$ “is not the reinforcement learning objective — in fact this equation is not anything”; it is chosen so its derivative comes out right, and the package must not know the $\hat{Q}$ values depend on the policy. Compute them in NumPy.

Practice notes, verbatim spirit: the gradient is far noisier than supervised learning, so batch sizes run in the thousands to tens of thousands; learning rates are hard, Adam is “okay-ish” and the right starting point; budget for much more hyperparameter tuning than supervised learning has taught you to expect.

Why step size is broken, and the natural gradient

Gradient ascent is really a constrained problem: maximize the linearized objective within an $\epsilon$ -ball — but that ball lives in parameter space, and parameters affect the policy unevenly. Constrain instead in distribution space with a parameterization- independent divergence: require $D_{KL}(\pi_{\theta'} \| \pi_\theta) \le \epsilon$ . The second-order expansion of the KL is a quadratic form under the Fisher information matrix $\mathbf{F} = \mathbb{E}_{\pi_\theta}\!\left[\nabla \log \pi_\theta\, \nabla \log \pi_\theta^\top\right]$ — estimable from the same samples — and the update becomes

$\theta' = \theta + \alpha\, \mathbf{F}^{-1} \nabla_\theta J(\theta)$

the natural gradient. With the preconditioner in place the vector field points at the optimum. Natural policy gradient picks $\alpha$ ; TRPO picks $\epsilon$ and derives $\alpha$ , using conjugate gradient to get $\mathbf{F}^{-1}\nabla J$ without forming $\mathbf{F}$ .

Back of the chalk tray — what to remember

The log-gradient identity turns $\nabla_\theta J$ into an expectation you can estimate from rollouts; unknown dynamics cancel — that’s REINFORCE.
The policy gradient is weighted maximum likelihood: good trajectories up, bad trajectories down — trial and error as gradient ascent, valid even under partial observability.
The naive estimator’s variance is the central obstacle; causality (reward-to-go) and baselines reduce it without introducing bias.
Importance sampling makes policy gradients off-policy, but full trajectory weights explode exponentially in $T$ — dropping the state-marginal ratio is the bounded-error approximation behind TRPO/PPO.
Implement via the pseudo-loss: weight log-likelihoods by reward-to-go and let autodiff do the rest; use large batches and expect heavy tuning.
Step size in parameter space is the wrong space — the natural gradient $\mathbf{F}^{-1}\nabla J$ constrains steps in distribution space instead.

Check yourself

Derivation-style exercises in the lecture’s spirit — work them on paper.

Problem 5.1 log-gradient derivation

Derive, from scratch, why $\nabla_\theta J(\theta) = \mathbb{E}_{p_\theta(\tau)}[\nabla_\theta \log p_\theta(\tau)\, r(\tau)]$ requires no knowledge of the dynamics, starting from $J(\theta) = \int p_\theta(\tau) r(\tau) d\tau$ . Point to the exact step where the transition probabilities disappear.

Show solution

Push $\nabla_\theta$ inside the integral, then apply $\nabla_\theta p_\theta = p_\theta \nabla_\theta \log p_\theta$ to get $\int p_\theta(\tau)\, \nabla_\theta \log p_\theta(\tau)\, r(\tau)\, d\tau$ , an expectation under $p_\theta$ . Expanding $\log p_\theta(\tau) = \log p(s_1) + \sum_t [\log \pi_\theta(a_t \mid s_t) + \log p(s_{t+1} \mid s_t, a_t)]$ , the dynamics disappear at the differentiation step: $\nabla_\theta \log p(s_1) = 0$ and $\nabla_\theta \log p(s_{t+1} \mid s_t, a_t) = 0$ because neither depends on $\theta$ . Only $\sum_t \nabla_\theta \log \pi_\theta(a_t \mid s_t)$ survives, and the dynamics now appear only in the sampling distribution — handled by rolling out the policy.

Problem 5.2 baseline variance

A one-step problem has two trajectories under the current policy: probability 0.5 each, rewards 101 and 99, with $\nabla_\theta \log \pi_\theta = +g$ for the first and $-g$ for the second. Compute the single-sample gradient estimator’s variance with baseline $b=0$ and with $b = 100$ (the mean reward). What happens with $b = 99$ ?

Show solution

With $b=0$ : the estimator is $+101g$ or $-99g$ with equal probability; mean $= g$ . Variance $= \mathbb{E}[X^2] - g^2 = \tfrac{1}{2}(101^2 + 99^2)g^2 - g^2 = 10000\,g^2$ . With $b=100$ : estimator is $+g$ or $-g$ … mean $= g \cdot \tfrac{1}{2}(1) + ...$ — careful: values are $(101-100)g = g$ and $-(99-100)g = g$ ; the estimator is always $g$ , variance $0$ , mean unchanged at $g$ . With $b = 99$ : values $2g$ and $0$ , mean still $g$ , variance $\tfrac{1}{2}(4g^2) - g^2 = g^2$ . Same expectation every time — exactly the unbiasedness theorem — but variance swings from $10^4 g^2$ to $0$ depending on $b$ .

Problem 5.3 importance-weight blowup

Your new policy agrees with the old one everywhere except it assigns ratio $\pi_{\theta'}/\pi_\theta = 0.9$ at every one of $T = 100$ time steps. What is the trajectory importance weight, and what does it do to the effective sample size? Why does keeping only the per-step action ratio avoid this?

Show solution

The full weight is $0.9^{100} \approx 2.7 \times 10^{-5}$ . With weights that small on every sample (and occasionally enormous ones when ratios exceed 1), the variance of the weighted estimator explodes — effectively almost none of your samples carry information, which is the “variance goes to infinity exponentially fast” failure. Writing the gradient over state–action marginals and dropping the state-marginal ratio leaves a single factor of $0.9$ per term rather than a product over all past steps, so the weight no longer decays exponentially in $T$ . The price is bias — acceptable, with bounded error, only while $\theta'$ stays close to $\theta$ , which is the constraint TRPO/PPO enforce.

Problem 5.4 natural gradient intuition

In the $(k, \sigma)$ example, the vanilla gradient’s $\sigma$ -component dominates as $\sigma \to 0$ . Explain in one or two sentences per item: (a) why this happens, (b) why no single learning rate fixes it, and (c) what property of the KL constraint makes the natural gradient immune.

Show solution

(a) The Gaussian log-density scales as $1/\sigma^2$ , so $\partial/\partial\sigma$ terms grow like $1/\sigma^3$ – $1/\sigma^4$ while $\partial/\partial k$ stays moderate — after normalization the $\sigma$ direction swallows the step. (b) A rate small enough to keep the $\sigma$ updates stable makes progress in $k$ negligible; a rate that moves $k$ overshoots in $\sigma$ — the conditioning problem lives in the geometry, not the scalar. (c) The KL divergence measures change between distributions, independent of parameterization; expanding it gives the Fisher matrix as the local metric, and $\mathbf{F}^{-1}$ rescales each direction by how much it actually changes the policy — big-effect parameters get small steps and vice versa, which is exactly the per-parameter learning rate you wanted.