The Complete Guide to Policy Gradients

What You'll Master

RL Foundations & Notation

Reinforcement learning is about an agent interacting with an environment over time, trying to discover — through trial and error — which actions lead to the most total reward. Think of a robot learning to walk: nobody programs exact motor commands; the robot tries stuff, gets a score, and adjusts.

The Trajectory Distribution

Three ingredients:

• p(s₁) — initial state distribution. Given by environment.

• πθ(at|st) — our policy's choice. Only part we control via θ.

• p(st+1|st,at) — environment dynamics. Unknown in model-free RL.

On-Policy vs. Off-Policy vs. Offline

The RL Objective

Everything reduces to one problem. Find θ that maximizes expected total reward:

• maxθ — search over all parameter vectors for the best.

• J(θ) — scalar: expected total reward. Make this as big as possible.

• 𝔼τ ~ pθ(τ) — expectation over trajectories. Both policy and env are stochastic → maximize the average.

• Σt r(st,at) — total reward of one trajectory.

With shorthand r(τ) = Σt r(st,at):

We integrate over all possible trajectories. Can't compute this exactly — estimate with samples (actual rollouts).

From Imitation to Policy Gradients

Previously: behavioral cloning — given expert (s,a) pairs, maximize likelihood:

Limitation: can never outperform the demonstrator. No improvement from practice.

The bridge: keep the same gradient structure, but weight each action by its reward:

Deriving the Policy Gradient

The most important derivation in this course. Starting point: J(θ). Goal: ∇_θJ(θ).

The REINFORCE Algorithm

Collect → compute → update → repeat. After updating θ, old data is invalid → recollect. This is on-policy.

Improving the Gradient: Causality

Our gradient weights every action by the total trajectory reward. But action at step 5 can't affect reward at step 3.

Replace total reward Σt'=1T with future-only Σt'=tT.

Formally: for t' < t, the expectation of ∇_θ log π_θ(a_t|s_t) · r(s_t',a_t') = 0. Past reward is independent of current action given the state. Removing these terms removes noise without changing the expected gradient.

Improving the Gradient: Baselines

The Baseline Trick

Subtract constant b:

Reward > b → positive weight. Reward < b → negative. Clear signal.

Best simple baseline: average reward b = (1/N) Σ_i r(τ_i).

Implementation: Surrogate Objective

Can't plug gradient formula into autodiff efficiently (N×T backward passes). Instead, build a surrogate objective whose autodiff gradient = our policy gradient:

Why it works: autodiff differentiates only log πθ(a|s). Reward weights are detached constants. So ∇_θ J̃ = our gradient.

Discrete actions: cross-entropy loss (actions as labels). Continuous (Gaussian):

But Why Does the Surrogate Work?

Let’s slow down. You’ve seen the formula. But many people memorize it without understanding why it exists and why it’s fast. The answer comes from understanding what a backward pass actually does.

When you call .backward() on a scalar in PyTorch, it walks backward through the computation graph and fills in ∂(that scalar)/∂θ for every parameter. That walk is one backward pass. It doesn’t matter if the scalar came from summing 3 things or 3 million things — backward is one sweep through the graph.

In supervised learning, you do this all the time without thinking:

python
loss = (1/N) * sum(cross_entropy(net(x_i), y_i) for i in range(N))
loss.backward()  # one backward pass, done

Mathematically ∇L = (1/N) Σ_i ∇ℓ_i, but you never compute each ∇ℓ_i separately. You build the sum, call backward, done. You trust that ∇(sum) = sum(∇).

The Policy Gradient Formula Has ∇ Inside the Sum

Now look at the policy gradient formula on paper:

Notice the ∇ is inside the double sum. In supervised learning, you never saw a formula with ∇ written inside — you just saw the loss and knew to take the gradient. Here, the derivation hands you a formula that already has ∇ in it.

So the question is: do you implement this literally, or rearrange first?

The Naive (Literal) Translation: N×T Backward Passes

If you read Σ_i Σ_t ∇(something) as literal instructions, you get:

python
# SLOW: one backward pass per (i,t) pair
total = 0
for i in range(N):
    for t in range(T):
        logp = log_pi(a[i,t], s[i,t])   # tiny scalar
        g = autograd.grad(logp, theta)   # backward pass!
        total += g * w[i,t]              # accumulate

Each iteration builds a tiny scalar logp, calls backward to get a gradient vector, multiplies by w_i,t, and accumulates. That’s N×T backward passes — each one a full sweep through the neural network. With N=100 trajectories and T=200 steps, that’s 20,000 backward passes per gradient update.

The One-Line Algebra Fix

Since w_i,t doesn’t depend on θ (it’s just a number computed from rewards), it slides inside the gradient:

And since ∇(sum) = sum(∇):

Look at the right side. What’s inside the brackets? A single scalar — compute log π_θ for each (i,t), multiply by w_i,t, sum them up. One number. And the ∇ is outside. Call this scalar J̃. Now you’re in the exact same situation as supervised learning: you have a scalar, you want its gradient, you call backward. One backward pass.

The Surrogate in Three Lines

python
logps = policy(s).log_prob(a)     # [N, T] — one forward pass
surrogate = (logps * w).mean()    # sum into one scalar
surrogate.backward()              # one backward pass

The sums over N and T still happen — but in the forward pass (adding numbers is cheap). The expensive part is the backward sweep through the neural network, and that happens exactly once. Internally, autodiff computes the same thing as N×T separate backward passes summed together — but in one sweep through one graph.

Why .detach() Matters

The only reason we could move ∇ outside the sum was that w_i,t doesn’t depend on θ. If w depended on θ, you’d get extra terms and the rearrangement would be invalid. That’s why you .detach() the rewards and baseline in code — to tell autodiff “don’t differentiate through this, it’s a constant.”

What Is the Surrogate, Really?

The real objective J(θ) is expected return. You can’t build J as a scalar that autodiff can differentiate — sampling trajectories isn’t differentiable. So you build a different scalar J̃ whose value is meaningless, but whose gradient happens to equal ∇J. You never care about J̃’s value. You only care that .backward() on it gives you the right gradient.

Off-Policy: Importance Sampling

Here’s the frustrating thing about REINFORCE: you collect a batch of trajectories, compute one gradient, update θ, and throw the entire batch away. Every single gradient step requires fresh data from the current policy. This is called on-policy learning, and it’s enormously wasteful.

Why can’t we reuse old data? Because our gradient formula assumes samples from π_θ. After we update θ, the old data came from a different policy. Using it directly gives a biased gradient. But what if we could mathematically correct for the mismatch?

Applying IS to Trajectories

We want to evaluate the policy gradient of π_θ′ (our new policy), but using trajectories collected from π_θ (our old policy). The IS correction requires the ratio of trajectory probabilities:

Something beautiful happens here: the initial state distribution p(s₁) and all the dynamics terms p(s_t+1|s_t,a_t) appear in both numerator and denominator — they cancel completely:

The Curse of Horizon: Exploding Products

There’s a devastating problem. The IS weight is a product of T individual ratios. Even if each ratio is close to 1, the product can be catastrophic:

Fix: Per-Timestep Importance Weights

The solution: instead of one IS weight per trajectory (product of T ratios), use one IS weight per timestep (a single ratio). This requires switching from a trajectory-level expectation to a timestep-level one.

The full off-policy gradient with per-timestep IS becomes:

The Hidden Approximation

There’s a subtle but important approximation buried in this formula. The full per-timestep IS weight should include a marginal state distribution ratio: π_θ′(s_t) / π_θ(s_t). This captures the fact that different policies visit different states.

The problem: π_θ′(s_t) is the probability of being in state s_t at time t under the new policy. Computing this requires rolling out the new policy or solving for the stationary distribution — which is exactly what we’re trying to avoid. So in practice, this ratio is approximated as 1: we pretend the state distribution doesn’t change between policies.

The Off-Policy Algorithm

The key win: K gradient steps per batch instead of 1. With K=10, you get 10× more optimization per unit of environment interaction. But there’s a catch: as θ′ drifts further from θ with each step, the IS weights become more extreme, the state distribution approximation breaks down, and the gradient becomes unreliable. After too many steps, you’re optimizing garbage.

KL Divergence Constraints

We’ve established the dilemma: off-policy policy gradients let us take multiple gradient steps per batch (great for sample efficiency), but each step moves θ′ further from the data-collecting policy θ, degrading the IS correction and eventually ruining the gradient. We need a principled way to say “stop updating before the gradient becomes unreliable.”

Measuring Policy Distance with KL Divergence

First, we need a way to measure how much π_θ′ has changed from π_θ. Raw parameter distance ||θ′ − θ|| doesn’t work because a small change in parameters can cause a large change in the policy (or vice versa). We need to measure distance in distribution space.

Why KL, Not Just ||θ′ − θ||?

Consider two policies that differ by Δθ = 0.001 in one parameter. If that parameter controls a softmax temperature, this tiny change might completely flip the action distribution. Conversely, Δθ = 100 in a dead parameter changes nothing. KL divergence measures the functional change in the policy — how different the action distributions actually are, not how different the numbers in memory are.

The KL-Constrained Optimization

We want to maximize the off-policy objective, but constrain how far the policy can move:

Why This Fixes Everything

Inside the trust region:

• IS weights near 1: If D_KL(π_θ′ ∥ π_θ) is small, the policies assign similar probabilities, so each ratio π_θ′(a|s)/π_θ(a|s) stays near 1. No extreme weights.
• State distribution valid: Similar policies visit similar states, so the approximation π_θ′(s) ≈ π_θ(s) holds.
• Multiple steps safe: You can take as many gradient steps as you want within the trust region. The constraint, not the number of steps, determines safety.

The Road to PPO: Three Levels of Approximation

The KL-constrained optimization above is the right objective, but how do you solve it?

What PPO’s Clipped Objective Looks Like

For reference, PPO’s objective (covered in depth in our PPO Veanors lesson):

Where r_t(θ) = π_θ′(a_t|s_t) / π_θ(a_t|s_t) is the IS ratio, and Â_t is the advantage estimate. The clip caps the ratio at 1 ± ε, and the min selects the more conservative estimate. When the ratio drifts outside the clipping range, the gradient vanishes — the optimizer can’t keep pushing the policy further.

Policy Architectures

The derivation above is deliberately architecture-agnostic. The policy π_θ(a|s) can be any differentiable function that outputs a probability distribution over actions. The entire derivation only requires one thing: that you can compute ∇_θ log π_θ(a|s). Any architecture that satisfies that works.

In practice, what people actually use depends on the input type:

Full Summary & Cheat Sheet

The Evolution

Derivation Cheat Sheet

Strengths vs. Weaknesses

What Comes Next

Actor-Critic: replace noisy reward-to-go with a learned value function (the "critic"). Dramatically reduces variance. Combined with KL constraints → PPO, the workhorse of modern RL.

import torch

def compute_reinforce_loss(policy, value_net, states, actions, rewards, gamma=0.99):
    T = len(rewards)

    # Step 1: Compute discounted reward-to-go
    returns = torch.zeros(T)
    G = 0
    for t in reversed(range(T)):
        G = rewards[t] + gamma * G
        returns[t] = G

    # Step 2: Compute advantages
    values = value_net(states).squeeze(-1)  # [T]
    advantages = returns - values.detach()  # don't backprop through V here

    # Step 3: Get log-probabilities of taken actions
    dist = policy(states)  # e.g., Normal(mu, sigma)
    log_probs = dist.log_prob(actions).sum(-1)  # [T], sum over action dims

    # Step 4: Policy surrogate loss (negate for gradient ascent)
    policy_loss = -(log_probs * advantages).mean()

    # Step 5: Value function loss (regress toward actual returns)
    value_loss = ((values - returns.detach()) ** 2).mean()

    return policy_loss, value_loss