When a baseline makes an agent too sure of itself

Imagine an agent learning to press one of two buttons. The blue button gives it a treat every time; the red button gives it nothing. The agent does not know this yet. It has to try buttons and gradually adjust its chances of pressing each one.

A policy-gradient learner updates itself using the reward from the button it pressed, compared with a reference point called a baseline. The signal is simply reward minus baseline. Baselines are usually introduced as a way to make these updates less noisy. That story is true, but it misses what happens along one particular learning journey.

Two reference points, two reactions

Blue pays 1; red pays 0. After a red press, a low baseline gives red a positive signal (0 − (−0.4) = +0.4). A higher baseline gives red a negative one (0 − 0.4 = −0.4). What does that mean over a whole learning run, when each next press comes from the policy the agent has learned so far?

With Adam and baseline minus 0.4, mean expected reward reaches 0.88 at step 60. — Faint lines show 100 example runs; thick lines average 1,000 runs per baseline. With Adam on the raw score-function gradient, the mean expected reward at step 60 is 0.88 for the low baseline versus 0.96 for the higher baseline. No sampled run is below 0.1 at step 60 in this setting: here the low baseline mainly slows improvement rather than producing a visible failure. Both runs start at 50% blue and use paired on-policy draws.

With Adam and baseline plus 0.4, mean expected reward reaches 0.96 at step 60. — Faint lines show 100 example runs; thick lines average 1,000 runs per baseline. With Adam on the raw score-function gradient, the mean expected reward at step 60 is 0.88 for the low baseline versus 0.96 for the higher baseline. No sampled run is below 0.1 at step 60 in this setting: here the low baseline mainly slows improvement rather than producing a visible failure. Both runs start at 50% blue and use paired on-policy draws.

With the low baseline, even a disappointing red press gets reinforced. If the agent happens to press red a few times early on, red becomes more likely each time. Blue becomes rarer, so the agent gets fewer chances to discover that blue was better all along. Its own decisions determine what evidence it will see next.

With the higher baseline, a red press moves the agent away from red. It leaves the door open for more blue presses. The difference is less about one update than about the feedback loop: try a button, change the policy, then collect the next experience using that changed policy.

But isn't the average update the same?

Yes. At a fixed policy, subtracting either baseline leaves the expected policy gradient unchanged. Averaged over possible button presses, both point toward blue. But a learner lives through one sequence of presses, not the average of every possible sequence. An unlucky early streak can make the better button so unlikely that the corrective evidence almost never arrives.

In our ICML 2021 paper, we call quick reinforcement of the sampled action committal behaviour. A less committal update keeps alternatives in play for longer. That is why the baseline is tied to exploration and exploitation: it affects how quickly the learner acts as though its early experience has settled the question.

Why variance alone does not settle it

The two-button story is only the beginning. Give the agent three choices: blue pays 1, orange pays 0.7, and gray pays 0. The baseline that minimizes a natural-gradient update's variance can still let the agent settle on orange. The value baseline—the current policy's expected reward—adapts as the policy changes and helps it keep exploring.

Watch one learning run

Each point is a policy: the closer it is to a corner, the more often it chooses that arm. The triangle shows one selected random seed, with one trajectory per baseline. Choose whether learning begins uniformly or already favors orange; each learner samples its next action from its own current probabilities.

Starts uniformly (33% per arm) · policy gradient + Adam step size 0.04 · β₁ = 0.9 · β₂ = 0.999 · ε = 10⁻⁸

Blue 1.0 Orange 0.7 Gray 0

Minimum variance Value baseline

Minimum-variance baselineb* = 0.57

Blue 33%Orange 33%Gray 33%

Value baselineV = 0.57

Blue 33%Orange 33%Gray 33%

Variance of the sampled gradient

Exact one-step variance of the raw policy gradient fed into Adam, before its moment updates; log scale.

Minimum variance: 0.059Value baseline: 0.059

Each curve is measured at that learner's own current policy. “Minimum variance” only promises the smallest variance among baselines at the same policy, so its curve can sit above the value-baseline curve later.

This triangle follows one selected on-policy seed, with one trajectory for each baseline. With Adam and a uniform start, both illustrated policies favor blue (about 92–93%) by step 120. Among 1,000 simulated runs with these Adam settings, neither baseline put over 95% probability on orange at step 120. This is an empirical comparison, not the NPG convergence theorem.

Changing the optimizer changes the story. Adam does not show the same orange outcome from a uniform start in this short simulation. Starting already close to orange makes it a harder test: with the settings above, the minimum-variance baseline can keep Adam near the second-best arm much longer than the value baseline. That is an observed delay in these runs, not a theorem about Adam's eventual destination.

How do we derive the minimum-variance baseline?

For arm i, write its reward as r_i, its current probability as π_i, and its gradient direction as x_i. The sampled gradient is G_i(b) = (r_i − b) x_i.

The expected sampled gradient does not depend on b, because Σ_i π_ix_i = 0. So minimizing its variance means minimizing the part that does depend on the baseline:

Σ_i π_i(r_i − b)² ‖x_i‖².

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

b* = [Σ_i π_ir_i ‖x_i‖²] / [Σ_i π_i ‖x_i‖²].

Natural policy gradient

For the three-arm softmax policy, the minimum-norm natural-gradient direction has ‖x_i‖² = 2 / (3π_i²). The common factors cancel:

b*(π) = [1/π_blue + 0.7/π_orange + 0/π_gray] / [1/π_blue + 1/π_orange + 1/π_gray].

The value baseline is simply V(π) = π_blue + 0.7π_orange. At the uniform start, both baselines equal about 0.57, so both initially reinforce orange after choosing it. As learning unfolds, the minimum-variance baseline can fall while the value baseline rises with the policy's expected reward. This changes how strongly the learner commits to orange.

Policy gradient with Adam

In the Adam view, the sampled raw policy gradient uses x_i = e_i − π rather than the natural-gradient direction. The same differentiation gives a different minimum-variance baseline:

b*_PG(π) = [Σ_i π_ir_i ‖e_i − π‖²] / [Σ_i π_i ‖e_i − π‖²].

Adam then averages and rescales these gradients using its past first and second moments. The chart reports variance of the raw gradient before Adam; this formula does not claim to minimize variance of Adam's history-dependent parameter step.

What does the variance chart measure?

Var[G(b)] = Σ_i π_i‖G_i(b)‖² − ‖Σ_i π_iG_i(b)‖².

We calculate this exactly from all three possible arms at each displayed policy. The log scale keeps both small and large values visible.

This is not a rule to always choose a larger baseline. A learner that never commits has its own problems, and the precise results depend on the algorithm and setting. The useful shift in perspective is to ask what a baseline makes the agent do after each experience, and how that changes the experiences available later.

Our NeurIPS 2022 follow-up develops the story further, including why the value function is a particularly useful reference point under its assumptions. The broad lesson remains simple: in learning from your own actions, a baseline can change more than the noise. It can change when you decide you have seen enough.