Baselines and Variance Reduction

Recall the REINFORCE gradient estimator for a trajectory-based objective $J(\theta) = \mathbb{E}_\pi[\sum_t \gamma^t r_t]$. The policy gradient theorem gives

where $G_t = \sum_{k \ge t} \gamma^{k-t} r_k$ is the return from time $t$ onward. This estimator is unbiased: its expectation equals the true gradient. But $G_t$ is a Monte Carlo sample of a random variable that can have huge variance, especially in long-horizon or sparse-reward tasks, and this variance propagates directly into the gradient estimate. A baseline is any function $b(s_t)$ that depends on the state but not on the action. The baselined gradient estimator replaces $G_t$ with $(G_t - b(s_t))$:

The key structural requirement is that $b$ must not depend on the action $a_t$ (it may depend on the state, the time step, or even be a constant); if this is violated, the subtraction can introduce bias. Subject to that requirement, the identity above holds for any choice of $b$, including the zero function, a running average of past returns, or — most usefully — an estimate of the state-value function $V^\pi(s_t)$. The terminology to fix: the quantity $G_t - b(s_t)$, when $b(s_t) \approx V^\pi(s_t)$, is an estimator of the advantage function $A^\pi(s_t, a_t) = Q^\pi(s_t, a_t) - V^\pi(s_t)$, which measures how much better action $a_t$ is than the average action in state $s_t$. This is the natural conceptual endpoint of the baseline idea and the direct precursor to Actor-Critic Methods and The Advantage Function, where a learned critic network provides $b(s_t)$ online rather than via Monte Carlo averaging.

Step 1 — Why subtracting a baseline is unbiased. The crux fact is a single lemma: for any function $b(s)$ that does not depend on $a$,

Proof: pull $b(s)$ out of the expectation (it is constant with respect to $a$), and use the "score function" identity $\mathbb{E}_a[\nabla_\theta \log \pi_\theta(a|s)] = \sum_a \pi_\theta(a|s) \frac{\nabla_\theta \pi_\theta(a|s)}{\pi_\theta(a|s)} = \nabla_\theta \sum_a \pi_\theta(a|s) = \nabla_\theta (1) = 0$. Since $b(s)$ is a constant multiplier, this whole expectation vanishes. This means adding or subtracting $b(s) \nabla_\theta \log \pi_\theta(a|s)$ to the REINFORCE gradient contributes exactly zero in expectation, so the estimator's mean is preserved for any choice of $b$.

Step 2 — Why it reduces variance. Write $X = \nabla_\theta \log \pi_\theta(a|s) \, G$ (the un-baselined per-step gradient contribution, treated as a vector-valued random variable; consider one coordinate for simplicity). The baselined version is $X - Y$ where $Y = \nabla_\theta \log \pi_\theta(a|s)\, b(s)$. Since $\mathbb{E}[Y] = 0$ (Step 1), we have $\mathbb{E}[X - Y] = \mathbb{E}[X]$, and

If $b(s)$ is chosen close to $\mathbb{E}[G \mid s] = V^\pi(s)$, then $Y$ tracks the state-dependent part of the noise in $X$ closely, making $\mathrm{Cov}(X,Y)$ large and positive, and the subtraction removes far more variance than it adds. In the ideal limit $b(s) = V^\pi(s)$, the baselined target becomes $G_t - V^\pi(s_t)$, which has zero conditional mean given $s_t$ and thus the smallest possible variance attributable to the "which state did we land in" component of randomness — only the action-selection and future-trajectory noise conditioned on the state remains.

Step 3 — Choosing a practical baseline. Three common choices, in increasing order of effectiveness: 1. Constant baseline: subtract the running average return $\bar{G}$ over recent episodes. Trivial to implement, removes only the overall scale/offset of returns. 2. State-dependent average: bucket states (or use a simple regression) to estimate average return per state. Better, but crude for large or continuous state spaces. 3. Learned value function baseline: train a separate function approximator $\hat{v}(s; w)$ (via regression toward observed returns) and use $b(s_t) = \hat{v}(s_t; w)$. This is the standard choice and is exactly the "critic" of actor-critic methods.

Step 4 — Putting it together in an algorithm. The following is REINFORCE with a learned baseline (sometimes called "REINFORCE with a value-function baseline," a stepping stone to full actor-critic): ``` Initialize policy parameters θ (actor) Initialize value-function parameters w (baseline), e.g. via a neural network v̂(s; w) Set learning rates α_θ (policy) and α_w (baseline) for each episode: generate trajectory s_0, a_0, r_1, s_1, a_1, r_2, ..., s_{T-1}, a_{T-1}, r_T using policy π_θ for t = 0 to T-1: G_t ← sum_{k=t}^{T-1} γ^{k-t} r_{k+1} # Monte Carlo return from step t for t = 0 to T-1: δ_t ← G_t - v̂(s_t; w) # baselined advantage estimate w ← w + α_w · δ_t · ∇_w v̂(s_t; w) # regress baseline toward G_t θ ← θ + α_θ · γ^t · δ_t · ∇_θ log π_θ(a_t | s_t) # policy gradient step end for end for ``` Notes on the pseudocode:

• The return $G_t$ is computed exactly as in plain REINFORCE (no change to the target being estimated).
• $\delta_t = G_t - \hat v(s_t; w)$ plays a dual role: it is the baselined return used to scale the policy gradient, and it is also the regression error used to update the baseline itself toward $G_t$ (this is why $w$'s update uses the same $\delta_t$, scaled by $\nabla_w \hat v$, i.e., gradient descent on squared error $\tfrac{1}{2}(G_t - \hat v(s_t;w))^2$).
• The $\gamma^t$ factor in the policy update accounts for discounting from the start of the episode, consistent with the discounted objective; it is sometimes dropped in practice as a design choice that slightly changes what's being optimized.
• Because $w$ and $\theta$ are updated using data from the same trajectory, updating $w$ first, before using $\delta_t$ for the $\theta$ update in the same step, is a matter of convention — both orders are used in practice and remain unbiased since we use the same $\delta_t$ computed before either update. The version shown updates $w$ then uses the same $\delta_t$ for $\theta$, which is standard.

Time complexity. The baseline adds only a constant-factor overhead to REINFORCE. Per episode of length $T$:

• Computing all returns $G_t$: $O(T)$ (one backward pass accumulating the discounted sum).
• Evaluating $\hat v(s_t; w)$ at each of the $T$ states: $O(T \cdot C_v)$ where $C_v$ is the cost of one forward pass of the value network (e.g., $O(d)$ for a linear model in feature dimension $d$, or the cost of a neural network forward pass).
• Updating $w$ via gradient descent on the squared error: $O(T \cdot C_v)$ for the backward pass, same order as forward.
• The policy gradient update itself: $O(T \cdot C_\pi)$, unchanged from plain REINFORCE.

So overall the algorithm remains $O(T \cdot (C_\pi + C_v))$ per episode — the same asymptotic order as plain REINFORCE, since a baseline is just one more function approximator evaluated alongside the policy. Space complexity is $O(T)$ to store the trajectory (states, actions, rewards, returns) plus $O(|\theta| + |w|)$ for the two parameter vectors — again, only an additive constant beyond REINFORCE's $O(T + |\theta|)$.

Correctness (unbiasedness). This is the central theoretical claim, and it deserves a full derivation, not just an assertion. Start from the exact policy gradient (single time step, to keep notation light; summing over $t$ recovers the full trajectory case):

Now consider the baselined estimator's expectation:

The second term vanishes because, as shown above, $\mathbb{E}_a[\nabla_\theta \log \pi_\theta(a|s)] = \nabla_\theta \sum_a \pi_\theta(a|s) = \nabla_\theta 1 = 0$, using the fact that $\pi_\theta(\cdot|s)$ is a probability distribution summing to $1$ for every $\theta$ (this differentiate-under-the-sum step is valid under standard regularity conditions on $\pi_\theta$, e.g., smoothness and dominated convergence, which hold for the softmax and Gaussian policies used in practice). Hence the baselined expectation equals the unbaselined one exactly — no approximation, no asymptotic argument, an exact algebraic cancellation for any function $b(s)$ that does not depend on $a$. This is why the "must not depend on the action" restriction is not a minor technicality: it is precisely the condition that makes the cancellation term equal to $b(s) \cdot 0$ rather than something nonzero.

Correctness (variance reduction, and why it isn't automatic). Variance reduction is not guaranteed for every baseline — it is guaranteed only that the estimator stays unbiased. To see when variance actually goes down, expand:

Variance decreases exactly when $2\,\mathrm{Cov}[\cdot,\cdot] > \mathrm{Var}[\nabla_\theta \log \pi_\theta(a|s)\,b(s)]$ — i.e., when $b(s)$ is correlated enough with $G$ (through the shared randomness in $s$) relative to the extra noise $b(s)$ itself injects via $\nabla_\theta \log \pi_\theta$. Since $b(s)$ is constant given $s$, choosing $b(s)$ close to $\mathbb{E}[G \mid s] = V^\pi(s)$ makes the covariance term large, because it directly cancels the state-dependent mean shift in $G$ that has nothing to do with which action was chosen. A provably optimal (minimum-variance) baseline, per state, can even be derived as a weighted average of returns weighted by squared score-function norm, but in practice $V^\pi(s)$ is used because it is simple to estimate and empirically works extremely well.

Tradeoffs. The baseline is free of bias but not free of cost: a second function approximator $\hat v(s;w)$ must be trained, adding hyperparameters (its own learning rate, architecture, regularization) and a second source of approximation error. If $\hat v$ is badly miscalibrated (e.g., early in training), the baseline can be a poor variance reducer — it never introduces bias into the gradient, but a bad baseline may fail to reduce variance much, or even (in finite samples, due to using the same returns to fit $w$ and to compute $\delta_t$) introduce subtle correlations that complicate finite-sample analysis, though the population-level unbiasedness of the gradient is unaffected.

Consider a tiny two-state MDP: state $s_A$ (start) and state $s_B$ (leads to episode end), with two actions "up" and "down" available in $s_A$. Suppose across a batch of 4 episodes we observe:

Without a baseline, REINFORCE would scale the gradient for the "up" action in $s_A$ by $+105$ and for "down" in $s_A$ by $+96$ — both positive, so both actions get reinforced, even though "down" was clearly the worse choice in that state. Worse, the raw returns in $s_A$ (around 100) and in $s_B$ (around 10) live on completely different scales, purely because $s_A$ happens to be a state with high value regardless of action. This scale difference is the "nuisance variance" the baseline exists to remove. With the learned baseline $\hat v(s_A) = 100$ and $\hat v(s_B) = 10$ (estimates of each state's average return), the baselined signal $G_t - \hat v(s_t)$ becomes $+5$ for "up" and $-4$ for "down" in $s_A$ — now correctly signed to increase the probability of "up" and decrease the probability of "down," at a much smaller, more comparable magnitude to the signal in $s_B$ ($+2$ and $-3$). The direction of the policy update is unchanged in expectation (both the raw and baselined estimators point the same way on average, by the unbiasedness proof above), but the per-sample gradient magnitudes are now centered near zero rather than swinging between roughly 7 and 105 — a reduction in scale of over an order of magnitude in this toy case, illustrating concretely why gradient noise (and hence required learning-rate tuning and number of samples) shrinks dramatically with a good baseline. To make the variance argument numeric: the sample variance of $\{105, 96, 12, 7\}$ (crudely as "the scale multiplying the score function") is large — spanning nearly 100 in range — while the variance of $\{5, -4, 2, -3\}$ spans only 9. Even in this four-point toy batch, the reduction in spread mirrors the general phenomenon: a state-appropriate baseline removes the part of the return's variance attributable to "which state am I in," leaving only the part attributable to "was this action better or worse than the state's average" — precisely the informative signal for policy improvement.

The baseline trick sits at the intersection of statistics (control variates, a decades-old variance-reduction technique in Monte Carlo estimation) and the specific machinery of policy gradient reinforcement learning developed in the late 20th century.

Baseline subtraction is not a niche trick; it is load-bearing infrastructure in nearly every deployed policy-gradient system. Robotics control policies trained with policy gradients rely on learned value-function baselines to make training tractable within realistic sample budgets, since raw REINFORCE variance would otherwise require infeasibly many rollouts on physical or simulated hardware. Large-scale reinforcement learning from human feedback (RLHF) pipelines used to align language models employ a value-function baseline (often called the "critic" or "value head") precisely to stabilize policy updates over noisy, high-variance reward signals coming from a learned reward model. Game-playing agents, from classic Atari-playing deep RL agents to more recent large-scale game AI systems, use actor-critic architectures whose critic is functionally a baseline in the sense developed here — this article's advantage estimate $G_t - \hat v(s_t)$ is the direct ancestor of the advantage estimates used inside PPO and A3C. Simulation-based scientific and engineering optimization (e.g., chip placement, molecule design, resource scheduling formulated as sequential decision problems) similarly leans on baseline-subtracted policy gradients to make Monte Carlo policy optimization practical at all. Understanding this idea is therefore a prerequisite for understanding Actor-Critic Methods and The Advantage Function, both of which are simply more sophisticated, more tightly integrated versions of "learn a good baseline and subtract it."

Ask the tutor: Can you show why the minimum-variance baseline is weighted by the squared score-function norm, and how far that is from just using $V^\pi(s)$?