Policy gradient baselines and the advantage function: why they reduce variance without bias

A policy gradient baseline is a quantity $b(s)$ subtracted from the return before it multiplies ${\nabla }_{\theta }\log {\pi }_{\theta }(a\mid s)$. Because the expected score of any policy is zero, subtracting anything that does not depend on the action leaves the expected gradient unchanged while it can cut the variance enormously. When the baseline is the state value $V^{\pi }(s)$, the weight becomes the advantage function $A^{\pi }(s,a)=Q^{\pi }(s,a)-V^{\pi }(s)$: how much better an action is than the policy's average in that state.

This page proves the unbiasedness claim in two lines, derives the variance-optimal baseline, and measures what baselines buy on two small problems. Shifting every reward by +20 leaves the problem unchanged but multiplies the variance of the plain gradient estimate by 150; a baseline removes that entirely. On a two-state MDP, a state-value baseline cut variance 17-fold, while knowing the exact $Q$ values without a baseline helped almost not at all. All numbers come from seeded numpy runs we did for this article.

Why the raw return is a bad learning signal#

Start with a four-armed bandit, the simplest setting where a policy gradient makes sense (the multi-armed bandits page covers the value-based approaches). The arms pay Gaussian rewards with means 1.0, 1.5, 2.0 and 0.5 and standard deviation 1. The policy is a softmax over four logits, currently $(0,0.5,1,-0.5)$, so it pulls the arms with probabilities 0.167, 0.276, 0.455 and 0.102. The REINFORCE estimate from one pull of arm $a$ with reward $R$ is

where $e_a$ is the one-hot vector for the pulled arm and $\pi$ is the vector of arm probabilities; $e_a-\pi$ is the score ${\nabla }_{\theta }\log {\pi }_{\theta }(a)$ of a softmax. The true gradient for the best arm's logit is ${\pi }_3({\mu }_3-\bar{\mu })=0.208$, where $\bar{\mu }=1.542$ is the policy's expected reward.

Now add a constant $c$ to every reward. Nothing about the problem changes: the same arm is best, the same policy is optimal, and the true gradient is identical. But the estimator is not invariant. Pulling any arm with a large positive reward pushes up that arm's logit and pushes down all others, whether the arm was good or not; only the small differences between arms carry information, buried under a large common term. We measured the total variance of the one-pull estimate (summed over the four logits) from 400,000 pulls at each offset.

Even at $c=0$ the baseline helps: 2.04 against 0.83, a 2.5-fold reduction, because rewards near 1.5 are still a large common term compared with the gaps between arms. In practice this is the normal situation. Rewards in real tasks have arbitrary offsets (a step penalty of −1, a reward model whose scores sit around +3), and the REINFORCE article shows how the all-negative returns of a step-penalty task limit the usable step size.

Why a baseline does not bias the policy gradient#

The whole argument rests on one identity: the expected score of a probability distribution is zero. For any state $s$,

The probabilities sum to one for every $\theta$, so their gradients sum to zero. Now take a baseline $b(s_t)$ that depends on the state, or on anything that happened before the action, but not on the action $a_t$ itself. Conditioning on $s_t$,

So the term you subtract has mean zero, and the expected gradient is unchanged. Williams proved this for his REINFORCE class in 1992, with the condition that the baseline be conditionally independent of the action [1]; Sutton and Barto present the episodic version as REINFORCE with baseline, using a learned state value [2]. Two consequences are worth spelling out:

• The baseline can be wrong and still unbiased. A poor baseline cannot bias the gradient; it can only fail to reduce variance, or increase it. Our bandit with $b=0$ and $c=20$ is a poor baseline in exactly this sense.
• The baseline must not see the action. If $b$ depends on $a_t$, the inner expectation no longer factors and the estimate is biased unless you add a correction term. Action-dependent baselines with such corrections exist [7], but Tucker and colleagues found that in standard benchmarks their apparent gains mostly came from other implementation differences, because the action-dependent part removes little variance in practice [6].

The same identity also justifies reward-to-go: rewards earned before $a_t$ act like a baseline for $a_t$ and can be dropped. The derivation is on the policy gradient theorem page.

What is the advantage function?#

The natural choice of state-dependent baseline is the state value $V^{\pi }(s)={\mathbb{E}}_{a\sim \pi }[Q^{\pi }(s,a)]$. Subtracting it from $Q^{\pi }$ gives the advantage

where $Q^{\pi }(s,a)$ is the expected return after taking $a$ in $s$ and $V^{\pi }(s)$ is the expected return from $s$ under the policy. The advantage is positive for actions better than the policy's average and negative for worse ones, so the policy gradient in advantage form pushes probability toward better-than-average actions and away from worse ones, which is what we wanted all along:

In practice neither $Q^{\pi }$ nor $V^{\pi }$ is known. The single-sample estimate $G_t-\hat{V}(s_t)$ is an unbiased estimate of the advantage up to the baseline error, which, as shown above, does not bias the gradient. Replacing $G_t$ by bootstrapped estimates such as $r_t+\gamma \hat{V}(s_{t+1})$ trades that unbiasedness for lower variance; the actor-critic and GAE article is about exactly that trade, which Schulman and colleagues made tunable with generalized advantage estimation [5]. The Bellman equation page covers how $V^{\pi }$ and $Q^{\pi }$ relate.

The optimal baseline#

Is $V^{\pi }$ the best baseline? Not quite. Take one gradient component with score $s=\partial \log {\pi }_{\theta }(a)/\partial {\theta }_i$ and return $R$. The estimate is $s(R-b)$, its mean does not depend on $b$, so minimizing variance means minimizing the second moment:

The optimal baseline is the expected return weighted by the squared score: outcomes that produce large gradient magnitudes count more. Weaver and Tao derived the optimal constant baseline for policy gradient methods [4], and Greensmith, Bartlett and Baxter analysed baselines as control variates in general, showing that the excess variance of any baseline is a weighted squared distance from the optimal one [3]. If one scalar must serve all components, replace $s^2$ by the squared norm of the score vector. The excess variance of a suboptimal baseline is

a parabola in $b$, so being slightly off costs almost nothing, while being far off (a baseline of 0 when returns sit around 20) is catastrophic.

How much does the optimal baseline add?#

On the bandit, the expected reward is 1.542. The optimal single scalar is 1.349, and the optimal per-logit values are 1.109, 1.516, 1.618 and 0.618. Measured total variance per pull: 0.827 with the expected reward, 0.802 with the optimal scalar, and 0.722 with per-logit optimal baselines. The first, simple baseline captures most of the possible gain: it removes 59% of the $c=0$ variance, and the fully optimal version removes 65%. That pattern is common, which is why almost every implementation uses an estimate of $V^{\pi }$ rather than $b^{\ast }$: it is easy to learn, it gives an interpretable advantage, and the remaining gap is small next to the noise in learning the baseline itself.

Baselines in sequential problems#

Bandits hide one thing that matters in multi-step problems, the kind formalized as a Markov decision process: the baseline can vary by state. We returned to the two-state MDP from the policy gradient theorem page (stay in A for +1, or move to B where staying pays +2, $\gamma =0.9$) at the uniform policy, where $V(A)=7.25$ and $V(B)=7.75$, and measured the variance of five per-episode estimators of the same gradient over 200,000 episodes. All five are unbiased; their means agreed with the exact gradient to within sampling error.

Three lessons sit in that figure. First, most of the variance came from the common offset of returns around 7, not from the randomness of future rewards: replacing the sampled return $G_t$ with its exact expectation $Q(s_t,a_t)$ removed only 11% of the variance, while subtracting a constant removed 93%. Second, in this problem the two states have similar values, so a state-dependent baseline adds little over a constant one (5.16 against 5.84); in problems where values differ a lot between states, the difference is much larger. Third, what remains after the baseline is mostly the noise in $G_t$ itself. Removing that requires replacing $G_t$ with a learned estimate, which is the step from baselines to critics.

Explore: gradient estimates with and without a baseline#

The figure draws 2,000 independent batches on the bandit and plots the histogram of the gradient estimate for the best arm's logit, without a baseline (top) and with the baseline you choose (bottom). The vertical line is the true gradient and the red triangle the mean of the estimates. Move the batch size and the reward offset.

Two things to try. Slide $c$ from −10 to +20 with the baseline fixed: the bottom histogram does not move, while the top one widens. Then pick the last baseline option, the batch mean including each sample's own reward, and set $N=2$. The histogram is narrow, but its mean sits at half the true gradient. That bias is the subject of the next section.

Batch-mean baselines and a hidden bias#

When no value function is learned, a cheap baseline is the average return of the other samples in the batch, for example the other responses to the same prompt in RL fine-tuning of language models. The subtle point is whether the sample's own return is included in the average. If it is, the baseline depends on the sample's own action and the unbiasedness argument fails. Since the mean of $N$ returns includes $R_i/N$, the estimate is exactly $(N-1)/N$ times the unbiased one in expectation:

We confirmed this on the bandit with 200,000 batches per size. With $N=2$ the mean estimate was 0.50 of the true gradient for every logit; with $N=4$, 0.75; with $N=16$, 0.94. The leave-one-out average, which excludes the sample's own return, was unbiased at every $N$, and at $N=16$ it had less than half the variance of no baseline (0.056 against 0.127). Within a single group the shrinkage only rescales the update, which a step size absorbs, but it is a reason to prefer leave-one-out when group sizes vary, and it is the estimator Ahmadian and colleagues used for REINFORCE-style RLHF [8]. Normalizing advantages by the group's standard deviation, as some group-relative methods do, is a further change: it reweights prompts by how consistent their rewards are and is no longer a pure baseline. For where these methods are used, see reinforcement learning with verifiable rewards and RLHF.

Code: baselines in a few lines#

import numpy as np

def advantages(returns, values=None, groups=None):
    """returns: (N,) returns-to-go. values: (N,) learned V(s_t) estimates, or None.
    groups: (N,) group ids (e.g. prompt ids) for a leave-one-out baseline, or None."""
    if values is not None:                      # state-value baseline -> advantage estimate
        return returns - values
    if groups is not None:                      # leave-one-out mean within each group
        adv = np.empty_like(returns, dtype=float)
        for g in np.unique(groups):
            m = groups == g
            n = m.sum()
            adv[m] = returns[m] - (returns[m].sum() - returns[m]) / max(n - 1, 1)
        return adv
    return returns - returns.mean()             # batch mean: biased by (N-1)/N, fine for large N

def optimal_constant_baseline(scores, returns):
    """scores: (N, d) score vectors grad log pi(a|s); returns: (N,).
    Variance-optimal scalar baseline E[|s|^2 R] / E[|s|^2]."""
    w = (scores ** 2).sum(axis=1)
    return (w * returns).sum() / w.sum()

Choosing a baseline in practice#

Failure modes#

• Baseline that peeks at the action. Including the current sample in a batch mean, or using $\hat{Q}(s,a)$ as the baseline without a correction, biases the gradient.
• Stale or untrained value estimates. A value network that lags far behind the policy is still unbiased but reduces little variance, and can increase it. Warm up the critic or use a simple running mean early.
• Reading advantage normalization as a baseline. Subtracting the batch mean of advantages is a baseline; dividing by their standard deviation rescales the step size per batch. Both are common, and only the first is covered by the unbiasedness argument.
• Judging a baseline from one run. Variance reductions show up as faster, steadier learning curves, and single runs are noisy; compare several seeds with intervals, as described in how to evaluate RL agents.
• Expecting a baseline to fix long horizons. A baseline removes offsets, not the accumulated randomness of future rewards. For that you need bootstrapped critics, at the price of bias.

A baseline is the cheapest improvement in policy gradient methods: it costs no bias, little computation, and on our two small problems it reduced gradient variance between 2.5-fold and 150-fold. The next step, replacing sampled returns with a critic's estimates, is covered in actor-critic methods and GAE, and the trust-region step-size control that sits on top of it in TRPO and PPO.

Frequently asked questions#

Sources#

1. R. J. Williams. Simple statistical gradient-following algorithms for connectionist reinforcement learning, 1992. doi.org
2. R. S. Sutton, A. G. Barto. Reinforcement Learning: An Introduction (2nd ed.), 2018. incompleteideas.net
3. E. Greensmith, P. L. Bartlett, J. Baxter. Variance Reduction Techniques for Gradient Estimates in Reinforcement Learning, 2004. jmlr.org
4. L. Weaver, N. Tao. The Optimal Reward Baseline for Gradient-Based Reinforcement Learning, 2001. arxiv.org
5. J. Schulman, P. Moritz, S. Levine, M. Jordan, P. Abbeel. High-Dimensional Continuous Control Using Generalized Advantage Estimation, 2015. arxiv.org
6. G. Tucker, S. Bhupatiraju, S. Gu, R. E. Turner, Z. Ghahramani, S. Levine. The Mirage of Action-Dependent Baselines in Reinforcement Learning, 2018. arxiv.org
7. C. Wu, A. Rajeswaran, Y. Duan, V. Kumar, A. M. Bayen, S. Kakade, I. Mordatch, P. Abbeel. Variance Reduction for Policy Gradient with Action-Dependent Factorized Baselines, 2018. arxiv.org
8. A. Ahmadian, C. Cremer, M. Gallé, M. Fadaee, J. Kreutzer, O. Pietquin, A. Üstün, S. Hooker. Back to Basics: Revisiting REINFORCE Style Optimization for Learning from Human Feedback in LLMs, 2024. arxiv.org