REINFORCE and variance reduction

The policy gradient theorem of Chapter 5.4 gave me an exact expression for ${\nabla }_{\theta }J$, but it is an expectation, and I cannot evaluate it in closed form for anything real. I have to estimate it by sampling. The moment you estimate a gradient by sampling, you inherit a new adversary that has nothing to do with bias and everything to do with noise: variance. REINFORCE is the most direct Monte Carlo estimator of the policy gradient, and left raw it is so high-variance as to be nearly useless. This chapter derives REINFORCE, then derives the three standard variance-reduction moves, reward-to-go, baselines, and the advantage, proving that each one leaves the gradient unbiased while shrinking its variance. The baseline unbiasedness proof is the theoretical seed of GRPO, so I give it in full. The lab is the payoff of the whole part: a REINFORCE agent on a toy task with a matplotlib figure showing the gradient-variance curve collapse the instant a baseline is switched on.

Theory

REINFORCE: Monte Carlo policy gradient

Start from the trajectory form of the theorem, Chapter 5.4 Equation (5). Sample $N$ trajectories ${\tau }^{(1)},\dots ,{\tau }^{(N)}$ from the current policy and replace the expectation with its sample mean:

$${\nabla }_{\theta }J(\theta )\approx \frac{1}{N}\sum _{i=1}^N\left(\sum _{t=0}^{T_i-1}{\nabla }_{\theta }\log {\pi }_{\theta }(a_t^{(i)}\mid s_t^{(i)})\right)R({\tau }^{(i)}).\text{\hspace{1em}(1)}$$

This is REINFORCE (Williams, 1992). The update rule is gradient ascent, $\theta \leftarrow \theta +\alpha \widehat{{\nabla }_{\theta }J}$. In words: for each sampled trajectory, push up the log-probability of every action it took, scaled by the trajectory's total return. Good trajectories make all their actions more likely; bad trajectories make all their actions less likely. It is beautifully simple and, as stated, badly behaved, for a reason I want to make precise before fixing it.

The trouble is that the single scalar $R(\tau )$ multiplies the log-prob gradient of every action in the trajectory, including actions that had nothing to do with the outcome. If one lucky reward inflates $R(\tau )$, REINFORCE dutifully reinforces every token that happened to be in that sequence, noise and signal alike. The estimator is unbiased, but its variance grows with trajectory length and with the scale of the returns, and for a language model with hundreds of tokens per episode that variance is enormous.

Reward-to-go: the first free variance cut

The first fix follows from causality: an action cannot influence rewards that were already collected before it was taken. So weighting $\nabla \log {\pi }_{\theta }(a_t\mid s_t)$ by the full return $R(\tau )$, which includes rewards from steps $0,\dots ,t-1$, is adding pure noise. Replace the full return with the reward-to-go, the return accumulated from step $t$ onward.

The score function integrates to zero

The identity I just used is the linchpin of every variance-reduction result in RL, so I prove it standalone.

Baselines: subtracting a number for free

Reward-to-go still leaves the returns on an arbitrary scale. If every completion for a prompt scores between $+8$ and $+10$, REINFORCE reinforces all of them strongly (every $G_t$ is large and positive), even though relative to each other some are worse. What matters for learning is not the absolute return but how a return compares to what you expected. So subtract a baseline $b(s_t)$, a number that can depend on the state but not on the action taken:

$${\nabla }_{\theta }J(\theta )=\mathbb{E}\left[\sum _t{\nabla }_{\theta }\log {\pi }_{\theta }(a_t\mid s_t)(G_t-b(s_t))\right].\text{\hspace{1em}(4)}$$

The remarkable fact is that any action-independent baseline leaves the gradient exactly unbiased, so you may choose $b$ purely to minimize variance.

Why a baseline reduces variance, and the optimal choice

Unbiasedness says the baseline does not hurt; now I show a good baseline helps. Write the per-step estimator as $g(a)={\nabla }_{\theta }\log {\pi }_{\theta }(a\mid s)(G-b)$. Its variance, for a scalar parameter to keep the algebra clean, is $\mathrm{Var}[g]=\mathbb{E}[g^2]-(\mathbb{E}[g]{)}^2$, and the second term is fixed (it is the true gradient squared, baseline-independent by the proof above). So minimizing variance means minimizing $\mathbb{E}[g^2]$.

Where the variance actually comes from

It helps to name the two sources of noise in the REINFORCE estimator separately, because the baseline attacks one of them and reward-to-go attacks the other. The law of total variance gives the decomposition cleanly.

Advantage as the general form

Put reward-to-go and a value baseline together. With $b(s_t)=v_{\pi }(s_t)$ and recalling that $\mathbb{E}[G_t\mid s_t,a_t]=q_{\pi }(s_t,a_t)$, the weight on the log-prob gradient becomes, in expectation,

$$G_t-v_{\pi }(s_t)⟶q_{\pi }(s_t,a_t)-v_{\pi }(s_t)=A_{\pi }(s_t,a_t),\text{\hspace{1em}(7)}$$

the advantage from Chapter 5.2. The advantage is the general form of the policy gradient weight, and every method in the rest of Part V is a different way of estimating $A_{\pi }$:

$${\nabla }_{\theta }J(\theta )=\mathbb{E}\left[\sum _t{\nabla }_{\theta }\log {\pi }_{\theta }(a_t\mid s_t)A_{\pi }(s_t,a_t)\right].\text{\hspace{1em}(8)}$$

• REINFORCE with baseline estimates $A$ by $G_t-b$ with $b$ a running average or the value function.
• Actor-critic and GAE (5.6) estimate $A$ with a learned critic and bootstrapping.
• PPO (5.7) uses the same $A$ but clips the update to stay in a trust region.
• GRPO (5.8) estimates $A$ by $(R_k-\bar{R})/\mathrm{std}(R)$ across a group, no critic at all.

Seen this way, the whole progression from REINFORCE to GRPO is one question asked repeatedly: what is the cheapest low-variance unbiased estimate of the advantage that fits in 16 GB?

Tooling

The tooling is matplotlib for the artifact plus NumPy for the agent, both already in the uv project from Chapter 5.1. I deliberately avoid a deep-learning framework here so that every line of the gradient estimator is visible and the variance measurement is unambiguous. The one methodological tool worth naming is measuring gradient variance directly: draw many independent gradient estimates at a fixed $\theta$ and take the variance across them, rather than trusting the loss curve to tell you whether your estimator is noisy. That measurement is the figure.

Lab

The task is a contextual bandit that stands in for one-step LLM generation: a "prompt" (context) is one of a few categories, the "action" is one of several tokens, and a verifier rewards the correct token for each context. I train a softmax policy with REINFORCE, once with no baseline and once with a running-mean baseline, and at each step I estimate the variance of the gradient across a fresh minibatch. The artifact is a matplotlib figure with two panels: reward-over-time (both variants learn) and gradient-variance-over-time (the baseline variant is dramatically lower). That variance gap, at equal final reward, is the entire lesson of the chapter made visual.

"""REINFORCE on a contextual bandit (a stand-in for one-step LLM
generation), with and without a baseline. Measure the variance of the
gradient estimator at each step and plot it. Artifact: reinforce_variance.png

Context = "prompt" category. Action = "token". Reward = 1 if the action
is the correct token for that context, else 0 (a toy verifier).

Run:  uv run python labs/reinforce_variance.py
"""
from __future__ import annotations
import numpy as np
import matplotlib
matplotlib.use("Agg")                     # headless: write a file, no display
import matplotlib.pyplot as plt

N_CONTEXTS = 4
N_ACTIONS = 5
CORRECT = np.array([0, 2, 4, 1])          # the "right token" per context
STEPS = 400
BATCH = 32                                 # completions sampled per step
LR = 0.5
SEED = 0


def softmax_rows(logits):
    z = logits - logits.max(axis=1, keepdims=True)
    e = np.exp(z)
    return e / e.sum(axis=1, keepdims=True)


def reward(context, action):
    return 1.0 if action == CORRECT[context] else 0.0


def grad_log_pi(probs_row, action):
    """d/dlogits log pi(action) for a softmax = onehot(action) - probs."""
    g = -probs_row.copy()
    g[action] += 1.0
    return g


def train(use_baseline: bool, rng):
    # Policy parameters: one logit vector per context.
    theta = np.zeros((N_CONTEXTS, N_ACTIONS))
    baseline = np.zeros(N_CONTEXTS)        # running mean reward per context
    rewards_hist, gradvar_hist = [], []

    for step in range(STEPS):
        probs = softmax_rows(theta)
        # Sample a batch of (context, action, reward).
        contexts = rng.integers(N_CONTEXTS, size=BATCH)
        batch_reward = 0.0
        # Accumulate per-sample gradients so we can measure their spread.
        per_sample_grads = np.zeros((BATCH, N_CONTEXTS * N_ACTIONS))
        grad_sum = np.zeros_like(theta)

        for i, c in enumerate(contexts):
            a = rng.choice(N_ACTIONS, p=probs[c])
            r = reward(c, a)
            batch_reward += r
            advantage = r - (baseline[c] if use_baseline else 0.0)
            g = np.zeros_like(theta)
            g[c] = advantage * grad_log_pi(probs[c], a)
            grad_sum += g
            per_sample_grads[i] = g.ravel()
            if use_baseline:
                # Update running-mean baseline for this context (eq: b ~ E[G]).
                baseline[c] += 0.05 * (r - baseline[c])

        rewards_hist.append(batch_reward / BATCH)
        # Variance of the gradient estimator: trace of the covariance of
        # the per-sample gradients (total variance across all components).
        gradvar_hist.append(float(per_sample_grads.var(axis=0).sum()))
        # Gradient ASCENT on J.
        theta += LR * grad_sum / BATCH

    return np.array(rewards_hist), np.array(gradvar_hist)


def smooth(x, k=15):
    kernel = np.ones(k) / k
    return np.convolve(x, kernel, mode="valid")


if __name__ == "__main__":
    rng = np.random.default_rng(SEED)
    r_no, v_no = train(use_baseline=False, rng=rng)
    rng = np.random.default_rng(SEED)        # same seed: fair comparison
    r_yes, v_yes = train(use_baseline=True, rng=rng)

    print(f"final reward  (no baseline): {r_no[-50:].mean():.3f}")
    print(f"final reward  (baseline)   : {r_yes[-50:].mean():.3f}")
    print(f"mean grad-var (no baseline): {v_no.mean():.4f}")
    print(f"mean grad-var (baseline)   : {v_yes.mean():.4f}")
    print(f"variance reduction factor  : {v_no.mean() / max(v_yes.mean(), 1e-9):.2f}x")

    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(11, 4.2))
    ax1.plot(r_no, alpha=0.35, color="tab:red")
    ax1.plot(r_yes, alpha=0.35, color="tab:blue")
    ax1.plot(range(len(smooth(r_no))), smooth(r_no), color="tab:red",
             label="no baseline")
    ax1.plot(range(len(smooth(r_yes))), smooth(r_yes), color="tab:blue",
             label="with baseline")
    ax1.set_title("Reward over training (both learn)")
    ax1.set_xlabel("step"); ax1.set_ylabel("mean batch reward"); ax1.legend()

    ax2.plot(v_no, color="tab:red", alpha=0.8, label="no baseline")
    ax2.plot(v_yes, color="tab:blue", alpha=0.8, label="with baseline")
    ax2.set_title("Gradient-estimator variance (lower is better)")
    ax2.set_xlabel("step"); ax2.set_ylabel("trace of grad covariance")
    ax2.set_yscale("log"); ax2.legend()

    fig.tight_layout()
    fig.savefig("reinforce_variance.png", dpi=130)
    print("\nwrote reinforce_variance.png")

What you should see. The printed summary reports that both variants reach essentially the same final reward (both solve the toy task, climbing toward a mean batch reward near $1.0$), while the mean gradient variance with a baseline is several times lower than without, a variance-reduction factor typically in the range of roughly 2x to 5x on this problem with these seeds. The saved figure reinforce_variance.png makes it unmistakable: the left panel shows two reward curves rising to the same ceiling, so the baseline costs nothing in final performance, and the right panel, on a log scale, shows the blue (baseline) variance curve sitting well below the red (no-baseline) one for the whole run. That separation at equal reward is the empirical face of the unbiasedness proof: the baseline moved variance, not bias. This is the exact mechanism GRPO scales up, swapping my per-context running-mean baseline for a per-prompt group-mean baseline, and it is why a critic-free method can train a reasoning model on a single RTX 5080 16GB (when you run the real GRPO loop in Part VII, log this same gradient-variance and entropy pair: measured on the baseline machine — record value, date, driver).