Section 15.2: The policy gradient theorem; REINFORCE

"Weight actions you took by how well things turned out afterward. The theorem says this is enough to climb the hill; the variance says the hill is steep."

Section 15.2

Picture a robot that slips on gravel at timestep 200 of a 500-step walk: how does the learning signal know whether to blame the foot placement at step 1 or the recovery flail at step 199? REINFORCE answers with one rule, illustrated in Figure 15.2A. Link every sampled action to the delayed return through likelihood-ratio credit assignment. This lets the policy learn from outcomes it can score even when the world between action and return is not differentiable. The policy objective is still expected return, \(J(\theta)=\mathbb{E}_{\tau\sim\pi_\theta}[G(\tau)]\). The obstacle is that the trajectory distribution contains environment dynamics, contact physics, resets, sensor noise, and reward delays. REINFORCE works because it differentiates the probability of the sampled trajectory with respect to the policy, while treating the environment as a source of samples.

From objective to estimator

The key identity is the likelihood-ratio trick, an identity that rewrites the gradient of an expectation over trajectories as an expectation of a score function (the gradient of a log probability) times the outcome, so the world only needs to be sampled, not differentiated:

$$\nabla_\theta J(\theta)=\mathbb{E}_{\tau\sim\pi_\theta}\left[G(\tau)\nabla_\theta \log p_\theta(\tau)\right].$$

Because the environment transition probabilities do not depend on \(\theta\), the policy-dependent part becomes a sum of action log probabilities:

$$\nabla_\theta J(\theta)=\mathbb{E}_{\tau\sim\pi_\theta}\left[\sum_t \nabla_\theta \log \pi_\theta(a_t\mid s_t)G_t\right].$$

Theory

For a sampled step, \(\nabla_\theta \log \pi_\theta(a_t\mid s_t)\) points in the direction that would make the sampled action more likely. Multiplying by \(G_t\) says how strongly to move in that direction. Positive high return increases the action probability; low or negative return pushes it down.

A baseline \(b(s_t)\) can be subtracted from the return without biasing the expected gradient:

$$\mathbb{E}_{a\sim\pi_\theta}\left[\nabla_\theta \log \pi_\theta(a\mid s)b(s)\right]=b(s)\nabla_\theta\sum_a \pi_\theta(a\mid s)=0.$$

This result is the reason value functions can serve as variance-reduction baselines. They change the noise of the update, not its target direction in expectation. To feel why that matters: without a baseline, the estimator is like trying to detect a 1-gram improvement in a 10-kilogram package by weighing the whole thing on a bathroom scale; subtract the known package weight first (the baseline) and that same 1-gram difference becomes the only thing the scale is measuring.

Worked Example

Monte Carlo returns, meaning returns computed by summing actual observed rewards along a full sampled trajectory rather than estimating them from a model, are what REINFORCE uses in place of a value prediction. Code Fragment 1 computes a tiny REINFORCE-style update for a two-action policy. The numbers show why a good return attached to a low-probability action produces a large learning signal.

# Compute REINFORCE loss terms for sampled actions.
# The loss uses frozen log probabilities and advantage estimates from rollout data.
import math

actions = ["left", "right", "right"]
action_probs = {"left": 0.25, "right": 0.75}
returns = [1.2, 0.4, -0.3]
baseline = 0.2

for action, return_value in zip(actions, returns):
    advantage = return_value - baseline
    log_prob = math.log(action_probs[action])
    loss_term = -log_prob * advantage
    print(action, "advantage:", round(advantage, 2), "loss term:", round(loss_term, 3))

The exact gradient depends on the policy parameterization, but the sign and scale are already visible. The first action was unlikely and better than baseline, so it receives a strong positive update. The final action was likely but worse than baseline, so its probability should decrease.

Practical Recipe

1. Save the old log probability for every sampled action during rollout collection.
2. Compute returns from rewards collected after the action, not from a reward prediction made before the action.
3. Subtract a baseline or value estimate to reduce variance before multiplying by the log-probability gradient.
4. Normalize advantages within a batch when reward scale changes across tasks or resets.
5. Debug with a one-state bandit before applying the estimator to a contact-rich simulator.

Both self-check questions hinge on the same algebraic fact, so it is worth unpacking exactly why the world drops out of the derivative. The theorem is often written compactly, but the cancellation is the teaching point. A trajectory probability factors into an initial-state term, transition terms, and policy terms. The initial-state and transition terms may decide which data you see, but they do not contain \(\theta\) if the policy parameters do not change the simulator or world dynamics directly. This is the score-function trick, and it is the reason gradient-based learning reaches through non-differentiable physics.

For embodied agents this matters practically. Contact forces, joint limits, and ground reactions are not differentiable functions of policy parameters, and a walking robot's foot strike is a discontinuous event through which a gradient is undefined. The score-function trick sidesteps that discontinuity, so one update rule serves both a smooth simulation and a physical robot with real friction and latency.

The mechanism is algebraic cancellation. When you write \(\nabla_\theta \log p_\theta(\tau)\), the transition probabilities \(p(s_{t+1}\mid s_t,a_t)\) appear in \(p_\theta(\tau)\) but are constant in \(\theta\); their derivatives are zero. Only \(\log \pi_\theta(a_t\mid s_t)\) survives differentiation. The world contributes only sampled data, not gradients, so no model of the world is needed.

That is why the gradient can be estimated from sampled rollouts. The price is variance: one unlucky slip can assign a poor return to several reasonable earlier actions. To make this concrete, consider bare REINFORCE on a 10-second locomotion task (500 timesteps at 50 Hz). It typically needs around 5,000 episodes to show consistent improvement. A linear value baseline drops that figure to roughly 500. Actor-critic methods and GAE in the next section keep the same log-probability mechanism while replacing raw Monte Carlo returns with more local advantage estimates. A gradient that cannot pass through the world still reaches the policy because the world's only job here is to hand back a number.

Code Fragment 2 makes the baseline identity concrete with a two-action policy. The expected score contribution of the baseline sums to zero because the probability derivatives across all actions cancel.

1. Verify that action probabilities sum to one before computing log probabilities.
2. Compute returns and baselines from the same reward convention and discount factor.
3. Check that the mean advantage is close to zero after baseline subtraction on a stable batch.
4. Track gradient norm because Monte Carlo returns can produce rare but very large updates.
5. Keep rollout data on-policy for REINFORCE (meaning the trajectories used for the update were sampled from the current policy parameters, not an older version); old trajectories require importance correction or a different algorithm.

# Verify that a state baseline has zero expected score contribution.
# This is why baselines reduce variance without changing the policy-gradient target.
prob_left = 0.25
prob_right = 0.75
baseline = 2.0

score_left = 1.0 - prob_left
score_right = 0.0 - prob_left

expected_baseline_term = (
    prob_left * score_left * baseline
    + prob_right * score_right * baseline
)
print(round(expected_baseline_term, 6))

Because the baseline only reshapes the noise of an already-unbiased estimator, when training stalls the right question is not whether the policy network is wrong but which property of the estimator has broken. When REINFORCE fails, classify the failure by estimator pathology before changing the policy network. High variance suggests better baselines, shorter-horizon shaping, or advantage normalization. Biased updates suggest stale log probabilities, off-policy data, action post-processing, or a reward convention mismatch.