Section 15.2: The policy gradient theorem; REINFORCE

A Careful Control Loop

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.

The key identity is the likelihood-ratio trick:

$$\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.

Worked Example

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.

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.

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. 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.

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; 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))

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.