← All posts

When a baseline makes an agent too sure of itself

Subtracting a baseline leaves the expected gradient unchanged. Yet its choice can determine whether an agent discovers the best action or commits to a worse one.

Consider an agent choosing between two buttons. Blue always gives a reward of 1; red always gives a reward of 0. With deterministic rewards and no credit-assignment problem, learning which button to press should be straightforward.

Can a policy-gradient algorithm nevertheless learn to press the red button forever?

For on-policy natural policy gradient, the answer is yes. Whether this happens can depend on the baseline, a choice usually motivated by variance reduction and thought to affect learning speed rather than the final solution.

In this post, we show that the baseline plays a central role in exploration and exploitation, beyond its effect on gradient variance.

Two buttons, two different learning outcomes

The agent samples from its current policy. Follow the individual runs: some learn blue, while others become increasingly likely to choose red.

With Adam, 15% initial probability on blue, learning rate 0.05 and baseline minus 0.4, mean expected reward reaches 0.931 at step 500. With Adam, 15% initial probability on blue, learning rate 0.05 and baseline plus 0.4, mean expected reward reaches 0.998 at step 500.
Faint lines show 100 runs; thick lines average 1,000 runs per baseline. Each action is sampled on-policy. The sliders reuse the same paired random draws.

For a sampled action a, ordinary policy gradient uses the ascent direction:

G(a) = (r(a) − b) ∇θ log πθ(a).

The score gradient points toward making a more likely. Reward minus baseline determines whether we follow or oppose that direction:

Observed rewardGradient contribution
r(a) > bReinforce the sampled action, encouraging repetition
r(a) < bDiscourage the sampled action, shifting probability toward alternatives
r(a) = bNo contribution

Viewed this way, the baseline sets a standard for what counts as a satisfactory reward. If that standard is too low, even a poor action can exceed it: with a negative baseline, red’s zero reward becomes a positive learning signal. Raising the baseline between the two rewards makes red disappointing while blue remains worth reinforcing. With two buttons, discouraging one favors the other; with more actions, the increased probability of alternatives need not be shared across all of them. These signs describe the sampled gradient; Adam also depends on accumulated moments.

The baseline can therefore shift the balance between exploitation and exploration. A low standard encourages repetition of current choices; a higher one more often discourages them, creating opportunities to try alternatives. With a low baseline, a streak of red samples can reinforce itself: red becomes more likely, blue is sampled less, and opportunities to correct the policy become rarer. We call this committal behavior in our ICML 2021 paper.

An unbiased gradient can still lead to the wrong policy

But how is this compatible with the expected gradient being unchanged?

For a fixed policy and an action-independent baseline b, treated as a constant when differentiating, the expected gradient is the same with or without the baseline:

Ea ∼ πθ[(r(a) − b) ∇θ log πθ(a)] = Ea ∼ πθ[r(a) ∇θ log πθ(a)].
Why does the baseline leave the expected gradient unchanged?

Expanding the expectation for an action-independent baseline b separates the reward and baseline terms:

E[(r(a) − b) ∇θ log πθ(a)]
= Σa πθ(a) r(a) ∇θ log πθ(a) − b Σa πθ(a) ∇θ log πθ(a).

Using πθ(a) ∇θ log πθ(a) = ∇θ πθ(a), the second sum becomes:

Σa πθ(a) ∇θ log πθ(a)
= ∇θ Σa πθ(a) = ∇θ 1 = 0.

The first sum is the gradient of expected reward. Therefore:

E[(r(a) − b) ∇θ log πθ(a)]
= ∇θ Σa πθ(a) r(a) = ∇θ J(θ).

In an MDP, the cancellation applies conditionally at each state, allowing a baseline b(s). A baseline that depends on policy parameters must still be held constant inside this estimator.

The identity compares estimators at the same policy; applying the same natural-gradient preconditioner preserves that equality. But after different sampled updates, the learners can have different policies and therefore collect different data. Equal expected directions at a given policy do not imply equal learning trajectories.

With natural policy gradient, a rare blue sample can produce a large corrective update, keeping the expected direction favorable. But if blue becomes sufficiently rare, a particular run may never receive that correction.

For the two-arm natural-gradient setting, we prove that a fixed negative baseline gives a strictly positive probability of choosing red forever. The Adam animations illustrate finite-time behavior; they do not establish this asymptotic claim.

The limits of variance reduction

Perhaps the low baseline simply produces too much noise. Would minimizing the variance of the gradient estimator prevent the agent from making these bad updates?

During my internship at Google Brain Montréal, I noticed that symmetric perturbations of the minimum-variance baseline increased variance by the same amount, but had asymmetric effects on performance. The direction of the perturbation appeared to matter.

Let b* minimize the total variance of the sampled gradient estimator G(b). Since variance is quadratic in the baseline, symmetric perturbations around its minimum give:

Var[G(b* − ε)] = Var[G(b* + ε)].

This equality holds at a fixed policy, using the same gradient estimator and norm. Once the learning trajectories diverge, their variances can differ too.

How do we derive the minimum-variance baseline?

For ordinary policy gradient, sampling arm i gives the gradient estimate:

Gi(b) = (ri − b) ∇θ log πθ(i).

As derived above, the expected gradient does not depend on b. The mean term in the variance is therefore constant with respect to the baseline, and we only need to minimize the second moment:

Σi πθ(i) (ri − b)2 ‖∇θ log πθ(i)‖2.

Differentiate with respect to b and set the result to zero:

−2 Σi πθ(i) (ri − b*) ‖∇θ log πθ(i)‖2 = 0.

Solving gives a weighted average of rewards, with more weight on actions whose score vectors have larger squared norms:

b* = [Σi πθ(i) ri ‖∇θ log πθ(i)‖2] / [Σi πθ(i) ‖∇θ log πθ(i)‖2].

This is why the variance-minimizing baseline generally differs from the value baseline: expected reward weights actions only by their probability. For the natural-gradient view, replace the score vector with the natural-gradient direction. In this three-arm softmax example, the resulting weights are proportional to 1/πi. “Minimum variance” therefore depends on which gradient and parameter-space norm we measure. In the Adam view, it refers to the raw gradient before Adam's moment updates.

Nevertheless, these symmetric perturbations produced qualitatively different learning paths. Matching the expected gradient and total variance did not capture the tendency to commit.

Although minimizing variance solves the two-arm example, it can fail with three arms. We construct a bandit where natural policy gradient with the minimum-variance baseline itself can converge to a suboptimal policy, while another baseline with larger variance guarantees convergence to the optimal one.

What makes the value baseline useful?

How should we choose this standard when we do not yet know which actions are good?

In our toy examples, we can place a baseline between the rewards because we already know which actions are better. In a realistic task, the agent must discover what constitutes a good outcome through experience, often while learning its baseline from that same experience. Exceeding a low standard does not establish that an action is better than alternatives the agent has not yet explored.

A common choice is the value function. In a bandit, this is the policy’s expected reward, Vπ = Σi π(i) ri, generally different from b*. Actions above that average receive a positive signal; those below it receive a negative one. This reference point changes as the policy improves.

The next animation compares the minimum-variance and value baselines, whose variances need not match. Orange is better than gray, but still suboptimal. If the baseline stays below orange’s reward, sampling it can reinforce an early commitment. Can the value baseline help the agent move beyond this second-best action?

Watch five learning runs

Each point is a policy; corners correspond to choosing one arm exclusively. Faint paths show five runs per baseline; bold paths show their mean, which need not follow any individual run.

Starts favoring orange (10% blue, 80% orange, 10% gray) · policy gradient + Adam step size 0.04 · β₁ = 0.9 · β₂ = 0.999 · ε = 10⁻⁸

Blue 1.0 Orange 0.7 Gray 0
Five three-arm learning trajectories per baseline and their averages Faint curves show five runs per baseline; bold curves average their policies at each step. Press Play to trace learning. Blue · 1.0 Orange · 0.7 Gray · 0
Minimum variance Value baseline
Minimum-variance baselineMean b* = 0.53
Blue 10%Orange 80%Gray 10%
Value baselineMean V = 0.66
Blue 10%Orange 80%Gray 10%

Seeds concentrated on the suboptimal arm

Share of 1,000 seeds with over 95% probability on orange at the displayed step.

Proportion of seeds concentrated on orange over time Both baselines start with zero seeds assigning over 95% probability to orange. Play reveals the proportion at each step. 100% 50% 0% 0 250 500 steps
Minimum variance: 0.0% (0/1,000)Value baseline: 0.0% (0/1,000)

Measured at each step, so the curve can fall when runs recover. Crossing the threshold does not establish permanent failure.

Mean variance of the sampled gradient

Raw-gradient variance before Adam; log scale.

Gradient variance over learning steps The two baselines start with 80% probability on orange. The minimum-variance baseline minimizes raw-gradient variance at each policy. Press Play to reveal each variance curve. 0 250 500 steps
Minimum variance: 0.059Value baseline: 0.059

Each curve averages variance at the learners’ own policies. Minimum variance is a comparison at a fixed policy; its curve can rise above the value-baseline curve once their policies differ.

Paired seeds 109–113. The orange-concentration chart uses a separate cohort of 1,000 runs.
How does the optimizer affect these examples?

Optimizer choice matters. From a uniform start, both Adam baselines learn the best arm reliably in these runs. Favoring the second-best arm exposes a much longer delay with the minimum-variance baseline.

Adam rescales gradients using accumulated moments. Equality of the expected raw gradients therefore does not imply equality of the expected parameter updates.

What does an “optimal” baseline optimize?

Suppose we could compute the exact minimum-variance baseline at every step. Would that prevent premature commitment?

The three-arm result says no. Even this oracle baseline can reinforce the second-best action until sampling the best action becomes too rare to correct the policy. The failure does not require an inaccurate critic: it occurs with the exact solution to the variance-minimization problem. Under the natural-policy-gradient assumptions of our theorem, a different baseline with higher variance guarantees convergence to the best action.

The distinction is between making the current gradient estimate less noisy and ensuring that future samples remain informative. Variance minimization considers the current policy. But each sampled update changes that policy and, with on-policy sampling, the data the agent will collect next. If the variance-minimizing baseline lies below orange’s reward, orange receives positive feedback. Repeated orange samples can then make the agent increasingly confident without providing new evidence that orange is the best choice.

Why might the value baseline help? As performance improves, it raises the standard an action must exceed to be reinforced. In our softmax examples, it also tends to sit above the minimum-variance baseline as the policy concentrates on good actions. This can make a practical difference: an action that receives positive feedback under the minimum-variance baseline may receive negative feedback under the value baseline, leaving more probability for alternatives.

Why can the value baseline be higher?

For ordinary policy gradient, write w(a) = ‖∇θ log πθ(a)‖². The minimum-variance baseline weights rewards by squared score norms, giving b* = E[rw] / E[w]. Subtracting Vπ = E[r] yields:

b* − Vπ = Cov(r, w) / E[w].

All expectations are under the current policy. As a softmax policy concentrates on a good action, that action’s score gradient becomes small. Rare, worse alternatives have larger score gradients and receive disproportionate weight in the variance calculation. A negative reward–squared-score covariance therefore places b* below Vπ.

This ordering is not universal: near a confidently wrong policy, rare actions may be better and reverse the comparison. The minimum-variance baseline also depends on the gradient estimator and parameter-space norm. The ordering helps explain the tendency to commit; it is not by itself a convergence guarantee.

Revisiting the usual explanation

What does this change about the usual justification for value baselines?

The value function is often motivated as an approximation to the minimum-variance baseline. Yet departing from that baseline can be part of what makes it useful. A baseline shapes both the noise in the update and the agent’s willingness to keep sampling alternatives; optimizing the first does not guarantee the second.

Our NeurIPS 2022 follow-up establishes convergence guarantees for on-policy natural policy gradient with the true value baseline in tabular settings. The baseline prevents the optimal action’s probability from vanishing, maintaining exploration even though gradient variance can remain unbounded.

Could an inaccurate critic reinforce a suboptimal action until alternatives become too rare to correct it?

For actor–critic methods, tracking advantage signs and individual policy trajectories can reveal behavior that average return and gradient variance conceal. The bandit analysis motivates these diagnostics; their implications depend on the actual updates used.