Policy Gradient Methods

"I never learned what any state was worth. I only learned to want certain actions a little more and others a little less, and to nudge those wants in the direction the reward was pointing. Give me enough episodes and my hesitations harden into habits; give me too few and I commit, confidently, to nonsense."

A Policy Nudging Its Own Probabilities Toward Reward

In Section 25.1 we learned to estimate $Q(s, a)$, the value of taking action $a$ in state $s$, with a deep network, and then acted by choosing the highest-valued action. That is the value-based route to a policy: the policy is implicit, a greedy read-off from a learned value function. This section takes the opposite and equally fundamental route. Instead of learning values and deriving behavior from them, we parameterize the behavior directly, writing the policy as a differentiable function $\pi_\theta(a \mid s)$ with parameters $\theta$ (the weights of a neural network), and we optimize $\theta$ to maximize expected return by gradient ascent. The reinforcement-learning foundations of Chapter 24 gave us the objective, expected return in a Markov decision process (Chapter 22); this section gives us a way to differentiate that objective with respect to the parameters of the policy itself. We use the unified notation of Appendix A: $s_t$ the state, $a_t$ the action, $r_t$ the reward, $\tau$ a trajectory, $\gamma$ the discount, and $\pi_\theta$ the parameterized policy.

Why pursue a separate family of algorithms when value-based methods already produce policies? Because directly parameterizing the policy unlocks settings where the value-based route is awkward or impossible, and because the resulting gradient is so clean that it underlies nearly every modern deep-RL algorithm of practical importance, from A2C and PPO through to the policy-optimization step inside reinforcement learning from human feedback. The continuous-action robotics of Chapter 27 and the stochastic exploration that solves partially observed problems (Chapter 23) both want a policy that emits a probability distribution over actions, something a $Q$-function with a discrete argmax cannot naturally provide. Policy gradients give exactly that.

The four competencies this section installs: to state precisely why and when a directly parameterized policy beats a value-based one; to derive the policy gradient theorem and recognize REINFORCE as its Monte Carlo estimate; to reduce the estimator's variance with a baseline and the advantage function, and to prove the subtraction does not bias the gradient; and to implement, train, and debug REINFORCE with a baseline on a real Gymnasium environment, then recognize the same algorithm inside a library. These are the load-bearing skills for everything in the rest of Chapter 25 and the advanced methods of Chapter 26.

1. Parameterizing the Policy Directly Beginner

A value-based agent answers the question "how good is each action here?" and lets the answer dictate behavior. A policy-based agent answers a different question directly: "what is the probability I should take each action here?" Formally it carries a parameterized conditional distribution $\pi_\theta(a \mid s)$, a function that maps a state to a probability distribution over actions, with the parameters $\theta$ being the weights of a neural network. For a discrete action space the network emits one logit per action and a softmax turns the logits into probabilities; for a continuous action space the network emits the mean (and often the log-standard-deviation) of a Gaussian from which the action is sampled. In both cases the policy is differentiable in $\theta$, which is the entire point: we can compute $\nabla_\theta$ of things that depend on the policy and ascend it.

The objective we maximize is the expected return under the policy, the same return that defines value in Chapter 22. Writing a trajectory $\tau = (s_0, a_0, r_0, s_1, a_1, \dots)$ generated by running $\pi_\theta$ in the environment, and $R(\tau) = \sum_{t \ge 0} \gamma^t r_t$ its discounted return, the objective is

The expectation is taken over trajectories drawn by sampling actions from $\pi_\theta$, which is what makes $\theta$ enter $J$ in two coupled places: it shapes which actions are chosen, and through them which states are visited. That coupling is precisely why differentiating $J$ is not a one-line application of the chain rule, and resolving it is the work of subsection two.

Why is this direct parameterization natural, indeed often necessary, where value-based methods strain? Three reasons recur, and each maps to a setting elsewhere in this book.

• Stochastic policies are first-class. A policy gradient learns a genuine distribution over actions, so a stochastic optimum (mixed strategies in games, exploration under partial observability as in Chapter 23) is representable directly. A greedy value-based policy is deterministic by construction and can only fake stochasticity with an external $\epsilon$-greedy hack.
• Continuous action spaces are immediate. Emitting the parameters of a Gaussian over actions handles a continuous control problem with no change of machinery, whereas a $Q$-function would need to solve $\arg\max_a Q(s, a)$ over a continuous set at every step. This is the entry point to the continuous control of Chapter 27 and is developed further with deterministic and soft variants in Section 25.4.
• The policy can be simpler than the value. Sometimes the optimal behavior is easy to describe ("turn left when you see the wall") even when assigning a precise numeric value to every state is hard. Parameterizing the simpler object directly can be the better-conditioned learning problem.

The cost of these advantages is a different failure mode. Value-based methods can reuse old data (a replay buffer, as in Section 25.1) because a $Q$-value is a property of the environment, not of the current policy. The pure policy gradient we derive next is on-policy: its gradient is an expectation under the very policy being updated, so each batch of experience is valid for exactly one update and is then discarded. That single fact, examined in subsection four, is the engine of policy-gradient sample inefficiency and the reason the field layered critics, trust regions, and importance weighting on top.

2. The Policy Gradient Theorem and REINFORCE Intermediate

We want $\nabla_\theta J(\theta)$, the gradient of expected return with respect to the policy parameters. The obstacle is structural: $J(\theta) = \mathbb{E}_{\tau \sim \pi_\theta}[R(\tau)] = \int p_\theta(\tau)\, R(\tau)\, d\tau$, an integral whose measure $p_\theta(\tau)$ depends on $\theta$. We cannot push the gradient inside the expectation directly, because differentiating an expectation whose distribution moves with the parameter is exactly what the chain rule does not hand you for free. The resolution is the score-function (likelihood-ratio) identity, the single most important trick in this section.

Start by differentiating the integral and applying the identity $\nabla_\theta p_\theta(\tau) = p_\theta(\tau)\,\nabla_\theta \log p_\theta(\tau)$, which is just the derivative of a log rearranged:

The last integral is, by definition, an expectation under $p_\theta$ again, so we have converted a derivative-of-an-integral into an integral-of-a-derivative, now safely an expectation:

Now expand $\log p_\theta(\tau)$. The trajectory probability factors into the environment's transition dynamics and the policy's action choices, $p_\theta(\tau) = p(s_0)\prod_t P(s_{t+1}\mid s_t, a_t)\,\pi_\theta(a_t \mid s_t)$. Taking the log turns the product into a sum, and taking $\nabla_\theta$ annihilates every term that does not contain $\theta$: the initial-state distribution and the transition dynamics are properties of the environment, not the policy, so their gradients are zero. Only the policy terms survive:

This is the decisive simplification: we never needed a model of the environment's dynamics, because its gradient vanished. Substituting back gives the policy gradient theorem in its trajectory form:

A short causality refinement sharpens this into the form we actually implement. An action $a_t$ taken at time $t$ cannot influence rewards earned before $t$, so multiplying $\nabla_\theta \log \pi_\theta(a_t \mid s_t)$ by the full-trajectory return $R(\tau)$ includes past rewards that are, in expectation, just noise. Replacing $R(\tau)$ with the reward-to-go from step $t$ onward, $G_t = \sum_{k \ge t} \gamma^{k-t} r_k$, removes that noise without adding bias, yielding the per-step estimator

Read this equation in words: to increase expected return, push up the log-probability of each action in proportion to the return that followed it. Good actions (high $G_t$) get their probability raised; bad actions (low or negative $G_t$) get theirs lowered. REINFORCE (Williams, 1992) is precisely this expectation estimated by Monte Carlo: run the policy to collect a batch of complete episodes, compute each $G_t$, average the per-step products to form an unbiased gradient estimate, and take a gradient-ascent step. No value function, no dynamics model, no bootstrapping, only sampled returns and the score function. Figure 25.2.1 traces the flow from objective to update.

3. Variance Reduction: Baselines and the Advantage Function Intermediate

REINFORCE is unbiased but its variance is brutal. The estimator multiplies each $\nabla_\theta \log\pi$ by a raw Monte Carlo return $G_t$, and returns swing wildly across episodes (one lucky rollout, one disaster), so the gradient direction is correct on average but jittery on any single batch, and learning is slow and unstable. The standard cure is to subtract a baseline $b(s_t)$, a quantity that depends on the state but not on the action taken, from the return before weighting the score:

The remarkable and load-bearing fact is that subtracting any action-independent baseline leaves the gradient exactly unbiased while typically slashing its variance. The unbiasedness follows from a one-line argument. The extra term the baseline introduces is, in expectation over the action drawn from $\pi_\theta$,

The baseline can be pulled out of the action-expectation (it does not depend on $a$), the score-function expectation collapses to the gradient of the total probability mass, and that mass is identically one, whose gradient is zero. So the baseline contributes nothing to the mean and only reshapes the variance. The variance-optimal scalar baseline is a return-weighted average of squared scores, but in practice the near-optimal and standard choice is the state-value function $b(s) = V(s)$, the expected return from $s$ under the current policy. Subtracting it converts the raw return into the advantage,

which measures how much better action $a_t$ was than the policy's average behavior in state $s_t$. The intuition is sharp: weighting the score by the advantage rather than the raw return means we raise the probability of an action only if it did better than expected for that state, and lower it if it did worse, rather than mechanically rewarding every action that happened to occur in a high-reward state. A move that returns $+100$ in a state where the average is $+105$ is, correctly, a discouraged move; the raw-return estimator would have rewarded it. Figure 25.2.2 contrasts the two weightings.

Implementing the baseline costs one extra learned object: a value network $V_\phi(s)$ trained by regression to fit the observed returns $G_t$ (minimizing $\sum_t (V_\phi(s_t) - G_t)^2$), running alongside the policy network. This is the first appearance of the two-network structure (a policy and a value estimate) that, when the value estimate is also used to bootstrap rather than only to center, becomes the actor-critic of Section 25.3. The line between "REINFORCE with a value baseline" and "actor-critic" is exactly whether the value network only subtracts a baseline (this section) or also supplies bootstrapped targets that replace the Monte Carlo return (the next section).

4. On-Policy Inefficiency and the Road to Actor-Critic Advanced

The policy gradient is an expectation under $\pi_\theta$, the current policy. The instant we take a gradient step and change $\theta$, every episode we collected was generated by the old policy and is no longer a valid sample from the new one. So the data must be thrown away after a single update and fresh episodes collected. This on-policy constraint is the structural weakness of pure policy gradients, and it stands in sharp contrast to the value-based methods of Section 25.1, whose replay buffer reused each transition hundreds of times because a $Q$-target is a property of the environment, not of the data-collecting policy.

The consequence is sample inefficiency: REINFORCE needs a great many environment interactions to learn, because every interaction is spent once and discarded. Combined with the Monte Carlo variance of subsection three, even with a baseline, pure REINFORCE is rarely the algorithm one ships. Two complementary fixes define the rest of the chapter, and both are previewed here so the next section reads as their culmination.

• Bootstrap the return with a critic. Instead of waiting for a full episode to compute the Monte Carlo return $G_t$, estimate it from a learned value function using a one-step (or $n$-step) bootstrap, $G_t \approx r_t + \gamma V_\phi(s_{t+1})$. The advantage becomes the temporal-difference error $\delta_t = r_t + \gamma V_\phi(s_{t+1}) - V_\phi(s_t)$, which has far lower variance than a full Monte Carlo return (at the cost of some bias from the imperfect $V_\phi$). A policy updated by this bootstrapped advantage is an actor-critic: the actor is $\pi_\theta$, the critic is $V_\phi$, and the critic's job has grown from "subtract a baseline" to "supply the target." This is the subject of Section 25.3 (A2C and A3C).
• Reuse data with importance weighting and a trust region. If we correct for the mismatch between the old data-collecting policy and the current one with an importance ratio $\pi_\theta / \pi_{\theta_\text{old}}$, we can take several gradient steps on the same batch, recovering some of the sample efficiency that the on-policy constraint cost us. Doing this safely (the ratio must not move too far) is what proximal policy optimization (PPO) and trust-region methods provide, the subject of Section 25.3 and Chapter 26.

So the narrative arc of this chapter is set by the two weaknesses this section exposes. Variance, addressed by the baseline of subsection three and pushed further by the critic's low-variance bootstrap. Sample inefficiency, addressed by the critic's reuse of single transitions and by PPO's importance-weighted multi-step updates. Every advanced policy-gradient algorithm you will meet is REINFORCE plus a particular answer to "how do I estimate the advantage with less variance?" and "how do I reuse data without breaking the on-policy assumption?"

5. Worked Example: REINFORCE With a Baseline From Scratch, Then a Library Advanced

We now make the whole pipeline executable on a real environment. The plan: implement REINFORCE with a value baseline entirely from scratch in PyTorch on the classic CartPole-v1 Gymnasium task (balance a pole by pushing a cart left or right), train it until it solves the task, and plot the return curve climbing. Then we reproduce a comparable policy-gradient agent in a few library lines and state the line-count reduction. Code 25.2.1 is the from-scratch agent: a policy network, a value-baseline network, episode rollout, return computation, and the policy-plus-value update of subsections two and three.

import gymnasium as gym
import torch, torch.nn as nn
import numpy as np

torch.manual_seed(0); np.random.seed(0)
env = gym.make("CartPole-v1")
obs_dim = env.observation_space.shape[0]   # 4 state features
n_act   = env.action_space.n               # 2 discrete actions (left, right)

policy = nn.Sequential(nn.Linear(obs_dim, 128), nn.Tanh(), nn.Linear(128, n_act))
value  = nn.Sequential(nn.Linear(obs_dim, 128), nn.Tanh(), nn.Linear(128, 1))
opt_p = torch.optim.Adam(policy.parameters(), lr=3e-3)
opt_v = torch.optim.Adam(value.parameters(),  lr=1e-2)
gamma = 0.99

def run_episode():
    """Roll the current policy out for one full episode; cache log-probs, values, rewards."""
    s, _ = env.reset()
    log_probs, values, rewards = [], [], []
    done = False
    while not done:
        st = torch.as_tensor(s, dtype=torch.float32)
        logits = policy(st)
        dist = torch.distributions.Categorical(logits=logits)  # softmax policy pi(a|s)
        a = dist.sample()                                      # sample an action ON-policy
        log_probs.append(dist.log_prob(a))                    # log pi(a_t | s_t)
        values.append(value(st).squeeze())                    # baseline V(s_t)
        s, r, term, trunc, _ = env.step(a.item())
        rewards.append(r); done = term or trunc
    return log_probs, values, rewards

def returns_to_go(rewards):
    """G_t = r_t + gamma G_{t+1}, computed backward over the episode."""
    G, out = 0.0, []
    for r in reversed(rewards):
        G = r + gamma * G
        out.append(G)
    return torch.tensor(out[::-1], dtype=torch.float32)

history = []
for episode in range(600):
    log_probs, values, rewards = run_episode()
    G = returns_to_go(rewards)                      # Monte Carlo returns-to-go
    V = torch.stack(values)
    adv = G - V.detach()                            # advantage A_t = G_t - V(s_t); detach: critic is a baseline only
    logp = torch.stack(log_probs)
    policy_loss = -(logp * adv).sum()               # ascend grad-log-pi * advantage  => minimize its negative
    value_loss  = ((V - G) ** 2).mean()             # regress the baseline V_phi(s) -> G_t
    opt_p.zero_grad(); policy_loss.backward(); opt_p.step()
    opt_v.zero_grad(); value_loss.backward();  opt_v.step()
    history.append(sum(rewards))
    if episode % 100 == 0:
        print(f"ep {episode:3d}  return {np.mean(history[-20:]):6.1f}")

ep   0  return   21.0
ep 100  return  102.4
ep 200  return  187.6
ep 300  return  421.8
ep 400  return  498.5
ep 500  return  500.0

The numbers tell the story of subsection three made empirical: the return rises monotonically once the baseline starts centering the signal, and it reaches the environment's 500-step cap. Code 25.2.2 turns the recorded history into the learning curve, a return-versus-episode plot with a smoothed trend so the climb is legible through the Monte Carlo noise that subsection three warned about.

import matplotlib.pyplot as plt

h = np.array(history)
smooth = np.convolve(h, np.ones(20) / 20, mode="valid")   # 20-episode moving average
plt.figure(figsize=(7, 4))
plt.plot(h, alpha=0.3, label="per-episode return")        # raw, noisy (high variance)
plt.plot(range(19, len(h)), smooth, lw=2, label="20-ep moving average")
plt.xlabel("episode"); plt.ylabel("return"); plt.title("REINFORCE + baseline on CartPole-v1")
plt.legend(); plt.tight_layout(); plt.savefig("reinforce_curve.png", dpi=120)

Now the library equivalent. The from-scratch agent above is roughly 45 lines of rollout, return, advantage, and dual-optimizer bookkeeping. A modern RL library collapses the same on-policy policy-gradient training, network construction, rollout collection, advantage estimation, and the update loop, into a handful of lines. Code 25.2.3 trains a comparable agent with Stable-Baselines3, whose A2C is a synchronous actor-critic built directly on the advantage-weighted policy gradient of this section.

from stable_baselines3 import A2C

# Construct, train, and evaluate an advantage-actor-critic policy-gradient agent.
model = A2C("MlpPolicy", "CartPole-v1", gamma=0.99, seed=0, verbose=0)
model.learn(total_timesteps=100_000)               # rollout + advantage + policy/value update, all internal

# Roll out the trained policy and report mean return.
env = gym.make("CartPole-v1")
returns = []
for _ in range(20):
    s, _ = env.reset(); done = False; total = 0.0
    while not done:
        a, _ = model.predict(s, deterministic=True) # greedy action from the learned policy
        s, r, term, trunc, _ = env.step(int(a)); total += r; done = term or trunc
    returns.append(total)
print(f"library A2C mean return over 20 eval episodes: {np.mean(returns):.1f}")

library A2C mean return over 20 eval episodes: 500.0

Read the three code blocks together. Code 25.2.1 implemented the policy gradient theorem and the value baseline by hand and trained CartPole to solved; Code 25.2.2 plotted the climb that subsection three's variance reduction made smooth enough to learn; Code 25.2.3 reproduced the result in three lines, with every piece, rollout, advantage, dual update, handled by the library. The pedagogical payoff is that the policy gradient is not a black box: you can derive it, code it on a real environment, watch it learn, and then recognize the identical algorithm inside the production tools of the temporal-AI toolbox.

from stable_baselines3 import A2C
model = A2C("MlpPolicy", "CartPole-v1", gamma=0.99)   # policy + value networks built for you
model.learn(total_timesteps=100_000)                  # rollout, advantage, update: all internal

We have parameterized a policy directly, derived the gradient of its expected return with the score-function trick, tamed that gradient's variance with a provably unbiased baseline, and trained a working agent on a real environment, by hand and then in three library lines. But two weaknesses remain visible in the return curve and the code: the gradient still carries Monte Carlo variance even with the baseline, and every episode is spent on a single update. The next section attacks both at once. Section 25.3 promotes the value baseline into a full critic that bootstraps the return, giving the low-variance temporal-difference advantage, and then builds A2C and A3C on it, before PPO and SAC layer on the trust-region and data-reuse machinery that make policy optimization the workhorse of modern deep reinforcement learning. The advantage you just learned to subtract becomes the quantity the critic learns to predict.