There is a strange little subtraction in policy-gradient code:

gradient estimate = (return - baseline) * grad_log_probability

It looks like bookkeeping. It is not.

The baseline is one of the few places where reinforcement learning lets you change the learning signal without changing the expected gradient. Pick it well and the estimator calms down. Pick it lazily and the update is still unbiased, but you may spend most of your samples paying for avoidable variance.

The usual explanation is:

subtract the value function so actions get advantage, not return

That is a good operational slogan. It is not the primitive reason the trick works. The primitive reason is smaller:

the score has expectation zero.

That sentence is the hinge between REINFORCE, score-function estimators, control variates, actor-critic methods, and a surprising fact from Greensmith, Bartlett, and Baxter: the average-reward baseline can be suboptimal.1

The Zero-Mean Handle

Let \(X \sim p_\theta\) and suppose we want the gradient of an expectation

\[J(\theta) = E_{p_\theta}[f(X)].\]

The score-function identity writes this as

\[\nabla_\theta J(\theta) = E_{p_\theta}\left[f(X)\nabla_\theta \log p_\theta(X)\right].\]

This estimator is old enough to have several names: likelihood-ratio gradient, score-function estimator, and, in reinforcement learning, REINFORCE.2 Sutton, McAllester, Singh, and Mansour later put the policy itself at the center of the reinforcement-learning objective, updating explicit stochastic policies in the direction of expected reward.3

Now use the one fact that makes the baseline legal:

\[E_{p_\theta}\left[\nabla_\theta \log p_\theta(X)\right] = \nabla_\theta \int p_\theta(x)\,dx = \nabla_\theta 1 = 0.\]

So for any constant \(b\) that does not depend on the sampled action,

\[E\left[(f(X)-b)\nabla_\theta \log p_\theta(X)\right] = \nabla_\theta J(\theta).\]

The baseline is not a correction to the target. It is a zero-mean term inserted to reduce variance. Mohamed, Rosca, Figurnov, and Mnih’s review makes the same point in general Monte Carlo-gradient language: a baseline keeps the estimator unbiased because the added score term has zero expectation, and is preferred because it can reduce variance.4

The control-variate question is not:

Which baseline sounds like value?

It is:

Which zero-mean subtraction is most correlated with the noise I want gone?

A One-Logit Bandit

Here is the smallest world where the distinction bites.

There are two actions. Action 0 pays mean \(\mu_0\). Action 1 pays mean \(\mu_1\). A one-parameter stochastic policy chooses action 1 with probability

\[p = \sigma(\theta).\]

For a sampled action \(A \in \{0,1\}\), the score is

\[\nabla_\theta \log \pi_\theta(A) = A - p.\]

The expected reward is

\[J(\theta) = (1-p)\mu_0 + p\mu_1,\]

and the exact gradient is

\[\nabla_\theta J(\theta) = p(1-p)(\mu_1-\mu_0).\]

The REINFORCE estimator with baseline \(b\) is

\[g_b = (R-b)(A-p).\]

Every constant \(b\) gives the same expectation. Not every constant gives the same variance.

For this one-parameter bandit, the variance-minimizing constant baseline is

\[b^* = \frac{E[R(A-p)^2]}{E[(A-p)^2]} = (1-p)\mu_1 + p\mu_0.\]

That is not the average reward

\[J = (1-p)\mu_0 + p\mu_1.\]

The weights are flipped.

Why? Because the rare action carries a larger score when it appears. If the better action has probability 12%, then a better-action sample has score 0.88, while a common-action sample has score -0.12. Variance cares about score squared, so the rare action gets much more leverage in the optimal baseline.

The value-like baseline is still unbiased. It is just not the variance optimum.

Baseline Lab

The lab below uses that one-logit bandit. The default policy chooses the better action with probability 12%. The common action has mean reward 6.00; the better action has mean reward 8.00; reward noise has standard deviation 0.60; each displayed gradient estimate averages a batch of 80 samples.

Default exact facts:

  • true gradient: 0.2112;
  • average-reward baseline: 6.24;
  • score-weighted optimal baseline: 7.76;
  • no-baseline batch variance: 0.079962;
  • average-baseline batch variance: 0.003525;
  • optimal-baseline batch variance: 0.000475;
  • optimal baseline cuts batch variance by 99.4% versus no baseline;
  • the average-reward baseline is 7.4x noisier than the optimal constant baseline in this default setting;
  • audit grid: 8 / 8 exact configurations passed.

The left panel samples many batch gradients. The middle panel sweeps the baseline value and plots exact estimator variance. The right panel runs a small policy-gradient loop so you can see how estimator noise turns into policy motion.

no baseline average reward score-weighted baseline true gradient exact variance curve

Deterministic one-logit Bernoulli bandit. The audit checks finite-difference gradients, exact unbiasedness for several baselines, optimal-baseline variance minimization, reward-offset invariance after optimal centering, and finite learning trajectories across 8 parameter settings.

Move Reward offset upward. The no-baseline estimator gets worse because it multiplies the score by a large constant reward. The optimal centered estimator barely cares. A constant added to all rewards should not change the decision problem, and a good baseline makes the gradient estimator act like it knows that.

Move Better action probability toward 50%. The average baseline and score-weighted baseline meet. At 50%, both actions have the same score magnitude, so ordinary reward averaging and score-squared weighting agree.

Move Better action probability toward an edge. The optimal baseline drifts toward the rare action’s reward, because rare actions carry large score magnitude when they appear.

Why This Matters More Than the Toy

The bandit formula is scalar, but the warning generalizes.

In a real policy-gradient system, the term multiplying the return is not one number. It is a vector of log-probability derivatives across actions, states, tokens, tools, or continuous control dimensions. The variance-minimizing baseline is therefore not simply “the expected return” in the abstract. It is the part of the return that best cancels the score-weighted noise you are actually injecting into the parameter update.

Greensmith, Bartlett, and Baxter analyze this in the broader GPOMDP setting and show that both baselines and actor-critic methods can be understood as additive control variates. They also show something that should stay mildly uncomfortable: using the average reward as a baseline can be suboptimal, and using the true value function as a critic can be suboptimal for variance in general.1

That does not mean value functions are bad. It means their job in policy gradients is not metaphysical truth. Their job is variance reduction without biasing the gradient you meant to estimate.

The practical questions become sharper:

  1. Is the baseline independent of the sampled action at the point where it is subtracted?
  2. Does it remove reward offsets that should not affect learning?
  3. Does it track the score-weighted variance of the actual estimator, not only the mean return?
  4. Is the variance reduction measured, or merely hoped for?

The last one is the habit I like most. When using a control variate, log the gradient variance you would have had without it. A baseline that looks elegant but does not reduce update noise is just an unbiased ornament.

The Mental Model

A policy-gradient sample has two pieces:

how surprising was the action under the policy?
how good was the outcome?

The score answers the first question. The return answers the second.

The baseline is allowed because it only subtracts something from the second piece that the first piece will average back to zero. It is useful when that subtraction removes outcome variation that was never informative about the direction of improvement.

That is the small, durable idea:

do not make the score carry reward scale it cannot learn from.

REINFORCE is often introduced as the simple, high-variance ancestor of modern policy optimization. Fair. But the baseline trick is not an old footnote. It is the cleanest place to see what policy-gradient learning is doing: building a usable signal by subtracting everything the score can legally ignore.

  1. Evan Greensmith, Peter L. Bartlett, and Jonathan Baxter, “Variance Reduction Techniques for Gradient Estimates in Reinforcement Learning”, Journal of Machine Learning Research 5:1471-1530, 2004. ↩ ↩2

  2. Ronald J. Williams, “Simple statistical gradient-following algorithms for connectionist reinforcement learning”, Machine Learning 8:229-256, 1992. ↩

  3. Richard S. Sutton, David McAllester, Satinder Singh, and Yishay Mansour, “Policy Gradient Methods for Reinforcement Learning with Function Approximation”, NeurIPS 1999. ↩

  4. Shakir Mohamed, Mihaela Rosca, Michael Figurnov, and Andriy Mnih, “Monte Carlo Gradient Estimation in Machine Learning”, Journal of Machine Learning Research 21(132):1-62, 2020. ↩