← All notes

Policy Gradients

July 12, 2026 · Updated July 14, 2026

Policy gradient methods are a class of reinforcement learning algorithms and a sub-class of policy optimization methods. Unlike value-based methods, which learn a value function to derive a policy, policy optimization methods directly learn a policy function $\pi$ that selects actions without consulting a value function. For policy gradient to apply, the policy function $\pi_\theta$ is parameterized by a differentiable parameter $\theta$.¹

Preliminaries: RL and MDPs

Reinforcement learning is about an agent interacting with an environment: at each time step the agent observes the state of the world, takes an action, and the environment responds with a reward and the next state.

The formalism behind this loop is the Markov decision process (MDP): a mathematical model for sequential decision making when outcomes are uncertain. It can be formulated as

a state space $\mathcal{S}$, an action space $\mathcal{A}$, transition probabilities $P(s_{t+1} | s_t, a_t)$, and a reward function $r(s, a)$. The basic objects:

• state $s$: the state of the “world” at time $t$
• action $a$: the decision taken at time $t$
• trajectory $\tau = (s_1, a_1, s_2, a_2, \dots, s_T, a_T)$: sequence of states/observations and actions
• reward function $r(s, a)$: how good is $(s, a)$?
• policy $\pi(a|s)$ (or $\pi(a|o)$ from observations): the behavior, usually what we are trying to learn

For a trajectory $\tau$, the return is the sum of rewards collected along it:

We want policies that generate high-return trajectories.

RL methods split by where their training data comes from:

• offline: using only an existing dataset, no new data from the learned policy
• online: using new data from the learned policy

This note only covers the online case.

Online RL

Online RL is a loop: run the current policy to collect a batch of data, improve the policy using that batch, and repeat. Over iterations, the policy’s trajectories concentrate on the behaviors that earn reward.

Policy Gradient objective

Goal: learn a policy $\pi_\theta$ that maximizes the expected sum of rewards over all possible episodes that this policy could produce:

where the trajectory distribution factorizes as

This factorization is the Markov property: the next state depends only on the current state and action, not on the rest of the history.

The expectation inside the objective is usually named $J(\theta)$, the objective function:

$J(\theta)$ tells us how good the policy is: a higher $J(\theta)$ means a better policy, and $\theta^* = \arg\max_\theta J(\theta)$.

Computing $\mathbb{E}_{\tau \sim p_\theta(\tau)}[\cdot]$ exactly isn’t possible, since there are millions (or infinitely many) possible trajectories due to randomness in the environment’s transitions and the policy’s action choices. (We’ll see how to get around this with Monte Carlo approximation further below.)

Using the return $r(\tau)$ defined above, this is just $J(\theta) = \mathbb{E}_{\tau \sim p_\theta(\tau)}[r(\tau)]$. Writing out the expectation over the trajectory distribution $p_\theta(\tau)$ gives an equivalent, integral form of the objective:

To improve the policy using the batch of data we collect with it, we want the gradient $\nabla_\theta J(\theta)$, because gradient ascent will update

Taking the gradient:

Move the gradient inside the integral:

But $\nabla_\theta p_\theta(\tau)$ is difficult to work with. We don’t know how to directly estimate gradients of trajectory probabilities.²

The log-derivative trick rewrites this into something we can estimate. Since

multiplying both sides by $p_\theta(\tau)$ gives

Substituting this into the gradient, we replace $\nabla_\theta p_\theta(\tau)$ with $p_\theta(\tau) \, \nabla_\theta \log p_\theta(\tau)$:

This now looks like an expectation, since $\mathbb{E}_{x \sim p}[f(x)] = \int p(x) f(x) \, dx$. Therefore

Intuitively, this term reweights the policy update by how good the trajectory turned out to be.³

This can be estimated using Monte Carlo approximation. A fundamental fact from probability: if $X \sim p(x)$, then

for samples $X_1, \dots, X_N$ drawn from $p$. After collecting $N$ trajectories by running the policy,

is an unbiased estimate of $\nabla_\theta J(\theta)$.

Going back to the original gradient integral, $\int \nabla_\theta p_\theta(\tau) \, r(\tau) \, d\tau$: why couldn’t we have estimated that one directly with Monte Carlo, instead of going through the log-derivative trick?⁴

We don’t need the environment dynamics

$\nabla_\theta \log p_\theta(\tau)$ can be expanded further. Using the trajectory factorization,

and differentiating with respect to $\theta$:

since the initial state distribution $p(s_1)$ and the environment dynamics $p(s_{t+1} \mid s_t, a_t)$ don’t depend on $\theta$, only the policy does, so their gradients vanish. This is a huge simplification: computing $\nabla_\theta J(\theta)$ no longer requires knowing the environment’s transition probabilities at all, only the policy’s own log-probabilities, which we have direct, analytical access to. The gradient and its Monte Carlo estimate become

Full algorithm

Putting it all together, this is the online RL loop from earlier, made concrete. This is the “REINFORCE algorithm,” the vanilla policy gradient:

Then repeat from step 1 with the updated policy.

This algorithm works, but it is noisy and has very high variance. A single lucky episode with a huge reward can dominate the gradient estimate.⁶

Reward-to-go: fixing credit assignment

One natural question comes up: why should an action at time $t$ get credit (or blame) for rewards that happened before it? Policy behavior at time $t$ can’t affect rewards at $t' < t$; future actions cannot change the past.⁷

Causality says we should only weight $\nabla_\theta \log \pi_\theta(a_t \mid s_t)$ by the rewards that came after it, not the whole trajectory’s return. So we replace

with the reward-to-go,

giving a new estimator:

Before, every action in a trajectory was weighted by the same total return. Now, each action is weighted only by the rewards that came after it, the only ones it could have influenced. Removing this irrelevant past-reward noise decreases variance, while the expected gradient stays correct.⁸

Baselines: centering the reward

Even with reward-to-go, there’s still a problem.⁹ If every trajectory in a batch gets a similar, uniformly positive reward, REINFORCE increases the probability of all of them, even the relatively worst one, since it only looks at each trajectory’s raw reward, not how it compares to the rest of the batch.

To fix this, we introduce a baseline $b$ and subtract it from the reward:

A natural choice is the batch’s average reward,

so instead of weighting by $r(\tau)$, we weight by how much better or worse than average each trajectory did.¹⁰ This also reduces variance: rewards that used to be large, similar numbers become small numbers centered at zero, so the gradient magnitude no longer swings wildly from batch to batch.¹¹

Subtracting a constant baseline doesn’t bias the gradient. In expectation,

so the estimator stays unbiased. (The full proof isn’t given here; see Stanford CS224R’s policy gradient slides.) The average reward is a pretty good baseline: unbiased, and lower variance.

Limits of baselines: sparse rewards

Even with a baseline, policy gradient still struggles when the reward is sparse, given only once at the very end of a trajectory. The baseline centers a trajectory’s total return around the batch average, but it can’t see inside the trajectory: two trajectories that happen to end with the same final reward get exactly the same weight, even if one made real progress along the way and the other made none at all.¹²

Because the reward is sparse, the gradient doesn’t know which specific action along the timeline was good or bad; it just scales the entire sequence of actions by the same final number. Policy gradient is still noisy and high-variance in this regime: even with a baseline, it struggles to tell “close failures” from “total failures” apart. Fixing this needs either dense rewards (small intermediate rewards for sub-steps toward the goal) or large batches, so the noise averages out over enough samples.

Implementing this efficiently: the surrogate objective

Putting reward-to-go and the baseline together, the full estimator is

Computing this naively means calling backward() once per state-action pair: differentiate $\log \pi_\theta(a_t \mid s_t)$, multiply by its own $(G_t - b)$, and sum. That’s a lot of individual backward passes.¹³

Deep learning frameworks are built around a single scalar loss and one backward() call, not thousands of gradients accumulated by hand. So instead we construct a scalar surrogate objective $\tilde{J}(\theta)$: not the true objective $J(\theta) = \mathbb{E}[r]$, but a quantity whose gradient happens to equal the policy gradient we actually want:

$(G_{i,t} - b)$ comes from the environment and the batch of collected data, not from the network, so autograd treats it as a constant with no dependence on $\theta$. Differentiating $\tilde{J}(\theta)$ therefore only differentiates $\log \pi_\theta(a_{i,t} \mid s_{i,t})$, carrying $(G_{i,t} - b)$ straight through:

which is exactly the policy gradient term above. A single backward() call on $\tilde{J}(\theta)$ therefore gives the same gradient as summing every individual term by hand.¹⁴

$\log \pi_\theta(a_t \mid s_t)$ is also exactly the term that shows up in maximum likelihood training: for classification, cross-entropy loss is $-\log p(y \mid x)$. The policy gradient loss, $-\log \pi_\theta(a_t \mid s_t) \, (G_t - b)$, is the same thing weighted by $(G_t - b)$, a weighted maximum likelihood that pushes up the log-probability of actions in proportion to how much better than the baseline they did.

In PyTorch, this is:

log_probs = policy.log_prob(actions)
loss = -(log_probs * advantages).mean()

optimizer.zero_grad()
loss.backward()
optimizer.step()

where advantages holds $(G_t - b)$ per timestep, and the sign is flipped since optimizers minimize by convention.

On-policy vs. off-policy

Every variant of the estimator covered so far, from vanilla REINFORCE through reward-to-go, the baseline, and the surrogate objective, shares one property: the expectation $\mathbb{E}_{\tau \sim p_\theta(\tau)}[\cdot]$ is over trajectories sampled from the current policy $\pi_\theta$. But step 3 of the full algorithm changes $\theta$. Once that happens, the trajectories used to compute that gradient no longer come from the policy we’re now trying to improve, so we have to go back to step 1 and collect a fresh batch under the new $\theta$ before taking another gradient step.

Every estimator above is on-policy: each gradient step needs its own brand-new batch of trajectories, and none of the old data can be reused once $\theta$ changes. That makes them sample-inefficient.

Importance sampling: reusing old trajectories

Being on-policy means all the old data is, technically, invalid the moment the policy changes: we need fresh rollouts under the new $\theta$ at every single gradient step, which is expensive. Importance sampling gives a way to reuse trajectories collected under an older policy instead.

The general trick: suppose we want an expectation under a distribution $p$,

but only have samples from a different distribution $q(x)$. Multiplying the integrand by $\frac{q(x)}{q(x)} = 1$,

so an expectation under $p$ can be estimated from samples drawn from $q$ instead, as long as each sample is reweighted by the importance weight $\frac{p(x)}{q(x)}$.

Applying this to the policy gradient objective: we want $J(\theta) = \mathbb{E}_{\tau \sim p_\theta(\tau)}[r(\tau)]$, the expectation under the current policy $\pi_\theta$, but only have trajectories $\tau \sim \bar{p}(\tau)$ collected under an older policy $\bar\pi$. Importance sampling rewrites the objective as

so the old trajectories can be reused, each one reweighted by the correction factor $\frac{p_\theta(\tau)}{\bar{p}(\tau)}$: how much more (or less) likely that trajectory would be under the new policy than under the old one.

To compute this correction factor, recall the trajectory factorization $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)$, and likewise for the old policy, $\bar{p}(\tau) = p(s_1) \prod_{t=1}^{T} \bar\pi(a_t \mid s_t) \, p(s_{t+1} \mid s_t, a_t)$. Both trajectories were generated in the same environment, so the initial state distribution $p(s_1)$ and the dynamics $p(s_{t+1} \mid s_t, a_t)$ are identical top and bottom, and cancel out of the ratio:¹⁵

leaving only a ratio of policy probabilities. We never need to know the environment dynamics, only the old and new policies’ probabilities of the actions actually taken.¹⁶

Off-policy policy gradient

Combining reward-to-go, the baseline, and importance sampling gives a full off-policy policy gradient: update the latest policy $\pi_\theta$ using trajectories sampled from an old policy $\bar\pi$, instead of $\pi_\theta$ itself.

Start from the on-policy reward-to-go-with-baseline estimator, written as an exact expectation:

The problem: this expectation is over trajectories $\tau \sim p_\theta(\tau)$, sampled from the very policy we’re trying to update, exactly the data we don’t have.

Using importance sampling, rewrite the expectation over the old policy’s trajectory distribution $\bar{p}(\tau)$ instead, correcting each sample with the importance weight $\frac{p_\theta(\tau)}{\bar{p}(\tau)}$:

and substituting in the per-timestep form of that ratio,

This is now genuinely off-policy: it can be computed entirely from trajectories collected under the old policy $\bar\pi$, with no fresh rollouts under $\pi_\theta$ required.

But the product term, $\prod_{t=1}^{T} \frac{\pi_\theta(a_t \mid s_t)}{\bar\pi(a_t \mid s_t)}$, is the catch: it multiplies $T$ per-timestep ratios together, and for longer horizons this product can become vanishingly small or explosively large, since even a small, consistent per-step difference between the two policies compounds multiplicatively over many timesteps.¹⁷

Reducing variance: per-timestep importance sampling

Because of this high-variance trajectory-level product, it helps to stop thinking about importance-sampling an entire trajectory at once, and instead importance-sample one timestep, one state-action pair, at a time.¹⁸ Instead of

consider

Write each timestep’s importance ratio as $\rho_t = \frac{\pi_\theta(s_t, a_t)}{\bar\pi(s_t, a_t)}$, and its gradient contribution as $g_t = \nabla_\theta \log \pi_\theta(a_t \mid s_t) \, (G_t - b)$. The old, trajectory-level estimator, for a single trajectory, was

one giant product of ratios, multiplied by one giant sum of gradient contributions: every $g_t$ in the sum gets scaled by the same trajectory-wide weight. The new, per-timestep estimator is instead

Now every timestep carries its own weight. Instead of

we have

• step 1 → its own weight $\rho_1$
• step 2 → its own weight $\rho_2$
• step 3 → its own weight $\rho_3$
• …

Each $g_t$ is scaled only by its own timestep’s ratio, not by a product of every ratio in the trajectory. This is exactly why per-timestep importance sampling reduces variance so dramatically.

But there’s a new problem: a policy doesn’t output $\pi_\theta(s_t, a_t)$ directly, only $\pi_\theta(a_t \mid s_t)$. Expanding $\pi(s, a)$ with the chain rule of probability, $\pi(s, a) = \pi(a \mid s) \, \pi(s)$, splits $\rho_t$ into two factors:

The first factor, the ratio of action probabilities, is easy: both policies hand it to us directly. The second factor, the ratio of state marginals, is not. $\pi(s_t)$ isn’t something either policy outputs: it’s the probability that the agent reaches state $s_t$ at all, which depends on the policy, every action taken before $t$, and the environment’s transition dynamics, none of which we have direct access to.

In practice, this state-marginal ratio is just approximated as $1$. If $\pi_\theta$ is only a small update away from $\bar\pi$, the two policies tend to visit nearly the same states, so $\pi_\theta(s_t) \approx \bar\pi(s_t)$ and

This approximation only holds when the two policies are close together.

Substituting it back in gives a fully practical estimator:

The huge trajectory-length product is gone. What’s left is a single probability ratio per timestep, far more stable than multiplying together $T$ of them.

When does our off-policy approximation stop working?

The per-timestep estimator above leans on approximating $\frac{\pi_\theta(s_t)}{\bar\pi(s_t)} \approx 1$, which quietly assumes $\pi_\theta$ hasn’t drifted far from $\bar\pi$.¹⁹ Once the new policy diverges enough, this breaks: the two policies start visiting different states, so the trajectories collected under $\bar\pi$ no longer describe what $\pi_\theta$ would actually encounter.²⁰

To keep the approximation valid, we can constrain how far a single update is allowed to move the policy: require the new policy to stay close to the old one, measured by KL divergence,

KL divergence measures how different two probability distributions are. If $\pi_\theta$ stays within $\delta$ of $\bar\pi$ at every state, the old data stays relevant, the state-marginal approximation $\pi_\theta(s) \approx \bar\pi(s)$ stays reasonable, and the resulting gradient estimate stays trustworthy.

This constraint shouldn’t be checked at a single state; it should hold on average, over the states the old policy actually visited:

Pick a small threshold $\delta$, and reject (or shrink) any update that would push the policy further than that. This constrained-optimization idea, keeping every update inside a “trust region” around the old policy, is exactly what TRPO is built around.

This topic is what convinced me to start writing these notes properly and publishing them here, so I can come back later and find the detailed examples waiting in the footnotes instead of re-deriving them from scratch. Made while going through Stanford’s CS224R and talking it through with GPT-5.5.