Policy Gradients from First Principles: REINFORCE, Variance and Baselines

Policy-gradient methods are the corner of reinforcement learning where the mathematics fits on one page but most bugs come from skipping a line of it. This post derives REINFORCE from the definition of the objective, shows where the variance comes from, proves that a baseline does not bias the estimate, and then checks each claim on a toy environment written in NumPy. No framework is involved; the whole experiment is a few dozen lines.

The loop every policy-gradient method lives in. The agent only touches the environment through sampled actions, which is why the gradient has to be estimated from samples.

The objective and the trick that makes it tractable

An episode is a trajectory $τ = (s_{0}, a_{0}, r_{1}, s_{1}, a_{1}, r_{2}, \dots, s_{T})$ generated by a start distribution $μ (s_{0})$ , a stochastic policy $π_{θ} (a ∣ s)$ and transition dynamics $p (s^{'} ∣ s, a)$ that we do not know. The objective is the expected discounted return:

J (θ) = E_{τ \sim p_{θ}} [R (τ)], R (τ) = \sum_{t = 0}^{T - 1} γ^{t} r_{t + 1} .

We cannot differentiate through the environment, but we can differentiate the probability of a trajectory. The identity $\nabla_{θ} p_{θ} = p_{θ} \nabla_{θ} \log p_{θ}$ turns the gradient of an expectation into the expectation of a gradient:

\nabla_{θ} J (θ) = \int R (τ) \nabla_{θ} p_{θ} (τ) d τ = E_{τ \sim p_{θ}} [R (τ) \nabla_{θ} \log p_{θ} (τ)] .

Now expand the trajectory probability. The dynamics and the start distribution do not depend on $θ$ , so they vanish under the gradient:

\begin{aligned} \log p_{θ} (τ) & = \log μ (s_{0}) + \sum_{t = 0}^{T - 1} [\log π_{θ} (a_{t} ∣ s_{t}) + \log p (s_{t + 1} ∣ s_{t}, a_{t})], \\ \nabla_{θ} \log p_{θ} (τ) & = \sum_{t = 0}^{T - 1} \nabla_{θ} \log π_{θ} (a_{t} ∣ s_{t}) . \end{aligned}

This is the whole reason policy gradients are model-free: the estimator needs the policy's own log-probabilities and nothing about the environment beyond samples.

Causality and the return-to-go

The product $R (τ) \sum_{t} \nabla \log π_{θ} (a_{t} ∣ s_{t})$ credits the action at time $t$ with rewards that arrived before $t$ . Those terms have zero mean (a reward at time $k < t$ cannot depend on an action chosen later) but non-zero variance, so dropping them is free. What remains is the return-to-go:

\nabla_{θ} J (θ) = E [\sum_{t = 0}^{T - 1} \nabla_{θ} \log π_{θ} (a_{t} ∣ s_{t}) G_{t}], G_{t} = \sum_{k \geq 0} γ^{k} r_{t + k + 1} .

Replacing the expectation by a single sampled episode gives REINFORCE (Williams, 1992). It is unbiased, and that is its main virtue: the variance can be enormous because $G_{t}$ is a sum of many noisy rewards, and the sign of each update depends on whether that particular return happened to land above or below what the policy usually achieves.

Baselines: subtracting a number without changing the expectation

Pick any function $b (s)$ that does not depend on the action. Then

E_{a \sim π_{θ} (\cdot ∣ s)} [b (s) \nabla_{θ} \log π_{θ} (a ∣ s)] = b (s) \sum_{a} \nabla_{θ} π_{θ} (a ∣ s) = b (s) \nabla_{θ} \sum_{a} π_{θ} (a ∣ s) = b (s) \nabla_{θ} 1 = 0.

So $\sum_{t} \nabla \log π_{θ} (a_{t} ∣ s_{t}) (G_{t} - b (s_{t}))$ has exactly the same expectation as the raw estimator for every choice of $b$ . Its variance does change with $b$ : for each term it is a quadratic in $b$ , minimised near the expected return from that state. The common choice $b (s) = V^{π} (s)$ is not the optimum but is close, and $G_{t} - V^{π} (s_{t})$ is an estimate of the advantage $A^{π} (s_{t}, a_{t})$ , which is where actor–critic methods come from.

A tiny environment and a tabular softmax policy

The environment is a line of 11 cells with the goal in the middle. An episode starts in a random non-goal cell; each step moves one cell left or right and costs $- 1$ ; reaching the goal pays $+ 10$ and ends the episode, which is also cut off after 20 steps. The optimal policy walks straight to the goal, so the optimal expected return is $8.0$ (the start distance is uniform over $1, \dots, 5$ and every extra step costs one point).

import numpy as np

N_POS, GOAL, MAX_T = 11, 5, 20

def reset(rng):
    s = GOAL
    while s == GOAL:
        s = int(rng.integers(0, N_POS))
    return s

def step(s, a):                       # a = 0: left, a = 1: right
    s2 = int(np.clip(s + (1 if a == 1 else -1), 0, N_POS - 1))
    done = s2 == GOAL
    return s2, (10.0 if done else -1.0), done

def policy(theta, s):                 # tabular softmax, theta has shape (N_POS, 2)
    z = theta[s] - theta[s].max()
    p = np.exp(z)
    return p / p.sum()

def grad_logp(theta, s, a):           # d log pi(a|s) / d theta = e_a - pi(.|s), in row s only
    g = np.zeros_like(theta)
    g[s] = -policy(theta, s)
    g[s, a] += 1.0
    return g

For a softmax over a table, $\nabla_{θ_{s, \cdot}} \log π_{θ} (a ∣ s) = e_{a} - π_{θ} (\cdot ∣ s)$ , which is the only piece of calculus the implementation needs. The two update rules differ by a single term:

def returns_to_go(rewards, gamma=0.99):
    G, g = np.zeros(len(rewards)), 0.0
    for t in reversed(range(len(rewards))):
        g = rewards[t] + gamma * g
        G[t] = g
    return G

def update(theta, S, A, Rw, lr=0.02):
    G = returns_to_go(Rw)
    grad = np.zeros_like(theta)
    for t, (s, a) in enumerate(zip(S, A)):
        grad += G[t] * grad_logp(theta, s, a)
    theta += lr * grad

def update(theta, V, S, A, Rw, lr=0.02, lr_v=0.1):
    G = returns_to_go(Rw)
    grad = np.zeros_like(theta)
    for t, (s, a) in enumerate(zip(S, A)):
        grad += (G[t] - V[s]) * grad_logp(theta, s, a)   # advantage-like weight
    for t, s in enumerate(S):                             # Monte Carlo value update
        V[s] += lr_v * (G[t] - V[s])
    theta += lr * grad

Both versions were run for 500 episodes with learning rate 0.02 and five random seeds.

Episodic return of REINFORCE with and without a baseline over 500 episodes — Episodic return (20-episode moving average), mean and one standard deviation over five seeds. Both estimators solve the task; the baseline mainly tightens the spread between seeds late in training.

Method	Mean return, last 50 episodes	Std across seeds
REINFORCE	7.59	0.26
REINFORCE + $V (s)$ baseline	7.49	0.14
optimal policy	8.00	—

On a problem this small the two learning curves are hard to tell apart, which is itself a useful lesson: a baseline reduces the variance of the gradient estimate, and whether that shows up in the return curve depends on how much variance the problem had to begin with.

Measuring the variance directly

The cleaner experiment is to freeze the policy (uniform, $θ = 0$ ), sample 3000 episodes, compute the single-episode gradient estimate for each one, and measure its total variance $E ‖ \hat{g} - E \hat{g} ‖^{2}$ as a function of a constant baseline $b$ . The state-value baseline $V^{π} (s)$ was estimated by Monte Carlo from a separate sample of 20,000 episodes.

Variance of the gradient estimate as a function of the baseline value — Variance of the single-episode gradient estimate at the uniform policy. The constant-baseline curve is a parabola, as the algebra predicts, with its minimum close to the mean return-to-go; the state-dependent baseline does slightly better than the best constant.

Estimator	Total variance	Relative
plain REINFORCE ( $b = 0$ )	550.0	1.00×
best constant baseline ( $b = - 6$ )	344.7	0.63×
state-value baseline $V^{π} (s)$	336.8	0.61×

The sample means of the two estimators differ by at most 0.19 in any coordinate over the 3000 episodes, consistent with the unbiasedness argument given the measured spread. The parabola is the real content of the figure: the baseline is a free parameter of the estimator, and a bad baseline (far to either side) is strictly worse than none at all.

WARNING

The derivation assumes the episodes were sampled from the current policy. Reusing old episodes after a parameter update makes the estimate biased unless you correct it with importance weights, which is exactly what off-policy and PPO-style methods do. A surprising share of "my REINFORCE does not learn" bugs are stale data.

Why the optimal constant baseline is a weighted mean of returns

For a single term $\hat{g} = g (G - b)$ with $g = \nabla \log π_{θ} (a ∣ s)$ , the variance is $E [g^{2} (G - b)^{2}] - (E [g G])^{2}$ , and the second part does not depend on $b$ . Expanding the first part and setting its derivative in $b$ to zero gives

b^{⋆} = \frac{E [g^{2} G]}{E [g^{2}]},

a mean of the returns weighted by the squared score. When the magnitude of the score is uncorrelated with the return this collapses to $E [G]$ , which is why the minimum in the figure sits close to the mean return-to-go. (Here $g^{2}$ is read coordinate-wise; the total-variance version sums over coordinates.)

Where this leads

Everything after REINFORCE is variance reduction plus trust: actor–critic replaces the Monte Carlo return by a bootstrapped estimate, generalized advantage estimation interpolates between the two with a $λ$ parameter, and PPO limits how far one batch of samples may move the policy. The derivation above is the common root of all of them, and the parabola is the picture to keep in mind whenever someone adds a baseline.

Policy Gradients from First Principles: REINFORCE, Variance and Baselines

The objective and the trick that makes it tractable

Causality and the return-to-go

Baselines: subtracting a number without changing the expectation

A tiny environment and a tabular softmax policy

Measuring the variance directly

WARNING

Where this leads

Further reading

Cite this post

Keep reading