Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Policy Gradient Methods › Variance Reduction with Baselines

Michael BrenndoerferDecember 26, 20257 min read

Language AI HandbookAll chapters

Part 37 of 67Alignment and RLHF

Chapter 277Policy Gradient Methods › Variance Reduction with Baselines

Reading section 8 of 14 · Variance Reduction with Baselines · 7 minSection 8 of 14 · 7 min

Section 8 of 14 · 7 min read

Policy Gradient Methods › Variance Reduction with BaselinesOne section · Sections in this chapter8 / 14 ↓

Sections in this chapter14

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

The main point for reducing variance comes from a useful mathematical property of the policy gradient: subtracting any constant from the reward does not change the expected gradient, but it can substantially reduce variance. This property allows us to center the rewards around zero. This provides clearer learning signals, without introducing any bias into our gradient estimates.

To understand why centering helps, consider the extreme case where all rewards are positive, say ranging from +5 to +15. Without centering, every gradient update pushes to increase the probability of all sampled sequences. The model receives mixed signals: "increase this sequence, increase that sequence too, increase everything." The optimizer cannot easily distinguish which responses are better than others because every response receives a positive gradient push, just of different magnitudes.

With a baseline of 10 (the approximate mean), rewards become -5 to +5. Now the gradient clearly distinguishes: "increase sequences above the baseline, decrease sequences below it." Sequences that scored 14 get positive reinforcement; sequences that scored 6 get negative reinforcement. This clearer signal leads to faster learning and lower variance in gradient estimates.

Consider the gradient estimator with a baseline $b$ :

\nabla_\theta J(\theta) = \mathbb{E}_{y \sim \pi_\theta}[\nabla_\theta \log \pi_\theta(y|x) \cdot (R(x, y) - b)]

where:

$b$ : a baseline value (which must be independent of the action $y$ , though it can depend on the state $x$ )
$(R(x, y) - b)$ : the "centered" reward, which reduces variance by removing the mean level

To verify this does not change the expected gradient, we need to show that the baseline term contributes zero in expectation. We check this by computing the expected value of the baseline term directly:

\mathbb{E}_{y \sim \pi_\theta}[\nabla_\theta \log \pi_\theta(y|x) \cdot b] = b \cdot \mathbb{E}_{y \sim \pi_\theta}[\nabla_\theta \log \pi_\theta(y|x)]

We pulled $b$ outside the expectation because it does not depend on $y$ . Now we need to evaluate the expected value of the score function $\nabla_\theta \log \pi_\theta(y|x)$ . This expectation is always zero:

\begin{aligned} \mathbb{E}_{y \sim \pi_\theta}[\nabla_\theta \log \pi_\theta(y|x)] &= \sum_y \pi_\theta(y|x) \frac{\nabla_\theta \pi_\theta(y|x)}{\pi_\theta(y|x)} && \text{(expand using log derivative)} \\ &= \sum_y \nabla_\theta \pi_\theta(y|x) && \text{(cancel } \pi_\theta \text{)} \\ &= \nabla_\theta \sum_y \pi_\theta(y|x) && \text{(linearity of gradient)} \\ &= \nabla_\theta 1 && \text{(probabilities sum to 1)} \\ &= 0 && \text{(derivative of constant)} \end{aligned}

This result shows that the score function always has zero mean under the policy that generates it. The reason is elegant: the score function measures how changing parameters affects the log-probability of different outcomes. Since increasing the probability of some outcomes necessarily decreases the probability of others (they must sum to one), these effects balance out exactly. Mathematically, the sum of all probabilities is fixed at 1, so its gradient must be zero.

This result implies that we can subtract any baseline without introducing bias. The question becomes: what baseline minimizes variance?

The optimal baseline is the expected reward under the current policy:

b^* = \mathbb{E}_{y \sim \pi_\theta}[R(x, y)]

where:

$b^*$ : the variance-minimizing baseline value
$\mathbb{E}[R(x, y)]$ : the expected reward under the current policy

This is called a value baseline, and in practice we often estimate it with a learned value function $V(x)$ that takes the prompt $x$ as input and predicts the expected reward. With this baseline, the quantity $(R(x, y) - b^*)$ becomes the advantage: how much better (or worse) this specific sequence is compared to what we expect on average. The advantage is positive when a response exceeds expectations and negative when it falls short.

Advantage

The advantage measures how much better a specific response $y$ is compared to the expected reward from the current policy:

A(x, y) = R(x, y) - V(x)

where:

$A(x, y)$ : the advantage of sequence $y$ given prompt $x$
$R(x, y)$ : the actual reward for this sequence
$V(x)$ : the value function estimating the expected reward $\mathbb{E}[R(x, y)]$

Positive advantage means better than expected; negative means worse than expected. The policy gradient with advantage is $\nabla_\theta J(\theta) = \mathbb{E}[\nabla_\theta \log \pi_\theta(y|x) \cdot A(x, y)]$ .

The advantage has a natural interpretation: it measures how surprising a reward is relative to what we expected. A response with advantage +5 did much better than typical; one with advantage -3 did worse. By training on advantages rather than raw rewards, we give the model clearer information about which responses are good relative to the typical performance level. This is the same intuition a manager uses when giving feedback: rather than saying "you scored 85 on this review," it is more informative to say "you scored 15 points above the team average."

In[13]:

Code

def reinforce_with_baseline(log_probs, reward, baseline):
    """
    REINFORCE with baseline for variance reduction.

    Args:
        log_probs: Log probabilities of selected tokens (shape: seq_len)
        reward: Scalar reward for the complete sequence
        baseline: Expected reward estimate from the value network

    Returns:
        Policy loss and value loss as separate tensors
    """
    log_prob_sequence = log_probs.sum()

    # Advantage: how much better or worse than expected
    advantage = reward - baseline

    # Policy gradient uses advantage instead of raw reward.
    # We detach the advantage so gradients do not flow back
    # through the value network when updating the policy.
    policy_loss = -log_prob_sequence * advantage.detach()

    # Value loss: train baseline to predict expected reward
    value_loss = (baseline - reward) ** 2

    return policy_loss, value_loss

In[14]:

Code

import numpy as np

np.random.seed(42)

n_trials = 1000
rewards_trial = np.random.normal(loc=5.0, scale=3.0, size=n_trials)

baselines = [0, 3, 5, 7, 10]
variances = [float(np.mean((rewards_trial - b) ** 2)) for b in baselines]
optimal_baseline = float(np.mean(rewards_trial))

Out[15]:

Visualization

Line plot showing variance minimized at baseline value of 5, with U-shaped curve around it. — Gradient estimate variance as a function of baseline value for rewards drawn from a normal distribution with mean 5.0. The variance is minimized when the baseline equals the expected reward (approximately 5.0, marked by the dashed line), showing a clear U-shaped relationship. Deviating from the optimal baseline increases gradient variance, showing why accurate value estimation is necessary for stable policy gradient training.

Out[16]:

Visualization

Histogram of raw rewards, all positive, with mean marked by a vertical line showing lack of contrast between good and bad sequences. — Distribution of raw rewards when all responses receive positive rewards (e.g., from a reward model with scale 3-7). All gradient updates push in the same direction, giving weak learning signals since the optimizer cannot distinguish good responses from mediocre ones.

Histogram of advantages centered at zero with a vertical baseline marker, showing clear positive and negative regions. — Distribution of advantages after subtracting the mean baseline, centered at zero. Positive advantages (right of zero) identify above-average responses that should be reinforced, while negative advantages identify below-average responses that should be suppressed, giving the optimizer a clear directional signal.