In Part 1, we saw that the asymptotic variance of the vanilla REINFORCE estimator scales cubically ($O(T^3)$) with trajectory length $T$.

In this second post, we explore how baselines systematically tame this variance. Rather than introducing baselines as an ad-hoc trick, we formulate them through the classical Monte Carlo technique of control variates, exploring variance-reducing baselines and showing how centering the return drops the estimator variance upper bound down to quadratic ($O(T^2)$).

The $O(T^3)$ Variance of REINFORCE

As established for vanilla REINFORCE in Part 1, even when replacing total reward with Reward-to-Go, the variance upper bound of the gradient estimator continues to grow cubically with trajectory length $T$: \begin{equation}\label{eq:variance_decomp} \text{Var}(\hat{g}) \le \underbrace{\left( \sum_{t=0}^{T-1} \EE\left[ \|\nabla_\theta \log {\color{colorpolicy}\pi_\theta(a_t \mid s_t)}\|^2 \right] \right)}_{O(T)} \cdot \underbrace{\vphantom{\sum_{t=0}^{T-1}} \EE\left[ |{\color{colorreturn}R(\tau)}|^2 \right]}_{O(T^2)} = O(T^3)\,. \end{equation} A cubic dependence on the horizon length is bad news. Doubling the episode length increases variance by a factor of eight ($2^3 = 8$), while a ten-fold increase magnifies it by a factor of a thousand ($10^3 = 1{,}000$). Because the batch size required to estimate gradients with fixed precision scales directly with variance, long-horizon reinforcement learning tasks quickly become intractable without variance reduction.

Interlude: Control Variates

How can we systematically reduce the variance of a Monte Carlo estimator? A standard technique in statistics and machine learning is the method of control variates.

Suppose we want to estimate a vector $\mu = \EE[Y]$ using a random sample $Y$. If we have access to another random variable $Z$ with known mean $\EE[Z] = \mathbf{0}$, we can construct a modified estimator for any scalar constant $b$: \begin{equation}\label{eq:generic_cv} \hat{Y}(b) \defas Y - b Z\,. \end{equation} Because $\EE[Z] = \mathbf{0}$, the estimator $\hat{Y}(b)$ remains strictly unbiased: $\EE[\hat{Y}(b)] = \EE[Y] - b \EE[Z] = \mu$. For any random vector $X$, the total variance is defined as $\text{Var}(X) \defas \EE[\|X - \EE[X]\|^2]$. Expanding the variance of $\hat{Y}(b)$ yields: \begin{equation}\label{eq:cv_variance} \text{Var}(\hat{Y}(b)) = \text{Var}(Y) + {\color{colorvariance}\underbrace{b^2 \, \text{Var}(Z) - 2b \, \text{Cov}(Y, Z)}_{\text{variance change}}}\,, \end{equation} where $\text{Cov}(U, V) \defas \EE[\langle U - \EE[U], V - \EE[V] \rangle]$. To achieve maximum variance reduction, we want the “variance change” term to be as negative as possible. Minimizing this convex quadratic function with respect to $b$ yields the optimal coefficient: \begin{equation}\label{eq:optimal_cv_scalar} b^* = \frac{\text{Cov}(Y, Z)}{\text{Var}(Z)} = \frac{\EE[\langle Y, Z \rangle]}{\EE[\|Z\|^2]}\,. \end{equation} Plugging $b^*$ back into the variance reveals the optimal variance reduction: \begin{equation}\label{eq:cv_variance_reduction} \text{Var}(\hat{Y}(b^*)) = \text{Var}(Y) {\color{colorvariance}\underbrace{\left( 1 - \rho_{Y, Z}^2 \right)}_{\text{variance reduction}}}\,, \end{equation} where $\rho_{Y, Z} = \frac{\text{Cov}(Y, Z)}{\sqrt{\text{Var}(Y)\text{Var}(Z)}}$ is the correlation coefficient between $Y$ and $Z$. Thus, for effective variance reduction, $Z$ must be strongly correlated with $Y$: the higher the correlation $|\rho_{Y, Z}|$, the larger the reduction in variance.

Baselines as Control Variates

Now let's apply this principle to policy gradients. From Part 1, the Reward-to-Go policy gradient estimator is given by \begin{equation}\label{eq:pg_raw_estimator} \hat{g}_{\text{RTG}}(\tau) = \sum_{t=0}^{T-1} \nabla_\theta \log {\color{colorpolicy}\pi_\theta(a_t \mid s_t)} \, {\color{colorreturn}G_t}\,, \quad \text{where } {\color{colorreturn}G_t} \defas \sum_{t'=t}^{T-1} {\color{colorreturn}r_{t'}}\,. \end{equation} What should we choose as our zero-mean control variate $Z$? Remember that the individual per-step scores $\nabla_\theta \log {\color{colorpolicy}\pi_\theta(a_t \mid s_t)}$ have zero conditional mean given the current state $s_t$: \begin{align}\label{eq:score_zero_mean_conditional} \EE_{a_t \sim {\color{colorpolicy}\pi_\theta(\cdot \mid s_t)}}\left[ \nabla_\theta \log {\color{colorpolicy}\pi_\theta(a_t \mid s_t)} \ \Big|\ s_t \right] &= \sum_{a_t} \nabla_\theta {\color{colorpolicy}\pi_\theta(a_t \mid s_t)} \nonumber \\ &= \nabla_\theta \underbrace{\sum_{a_t} {\color{colorpolicy}\pi_\theta(a_t \mid s_t)}}_{=\, 1} = 0\,. \end{align} Because this zero-mean property holds conditionally at every state $s_t$, we can scale each step's score by any arbitrary function $b(s_t)$ that depends solely on the state $s_t$ (and not on the action $a_t$ or future rewards) without introducing bias into the expected gradient: \begin{equation} \EE\left[ \nabla_\theta \log {\color{colorpolicy}\pi_\theta(a_t \mid s_t)} \, b(s_t) \ \Big|\ s_t \right] = b(s_t) \EE\left[ \nabla_\theta \log {\color{colorpolicy}\pi_\theta(a_t \mid s_t)} \ \Big|\ s_t \right] = 0\,. \end{equation} Subtracting this localized zero-mean control variate at each step yields the general policy gradient baseline estimator: \begin{equation}\tag{BASE}\label{eq:baseline_estimator} \bbox[10px, border: 1.5px solid #81A1C1]{ \hat{g}_b(\tau) = \sum_{t=0}^{T-1} \nabla_\theta \log {\color{colorpolicy}\pi_\theta(a_t \mid s_t)} \left( {\color{colorreturn}G_t} - b(s_t) \right) }\,. \end{equation}

Expected Return Baselines

What is a principled choice for the baseline function $b(s_t)$ in $\eqref{eq:baseline_estimator}$? Because subtracting a state-dependent baseline leaves the expected policy gradient unchanged, we want to choose $b(s_t)$ to minimize the variance of each step's gradient estimator, $\hat{g}_{b, t} \defas \nabla_\theta \log {\color{colorpolicy}\pi_\theta(a_t \mid s_t)} \left({\color{colorreturn}G_t} - b(s_t)\right)$.

Following the same decomposition used in $\eqref{eq:variance_decomp}$, we can bound the conditional variance of this localized gradient estimator by separating the expectation into the product of the expected squared score norm and the expected squared return: \begin{equation}\label{eq:per_step_variance_bound} \text{Var}(\hat{g}_{b, t} \mid s_t) \le \EE\left[ \left\|\nabla_\theta \log {\color{colorpolicy}\pi_\theta(a_t \mid s_t)}\right\|^2 \ \Big|\ s_t \right] \cdot \EE\left[ \left( {\color{colorreturn}G_t} - b(s_t) \right)^2 \ \Big|\ s_t \right]\,. \end{equation} The first factor—the expected squared score norm—depends purely on the policy parametrization and is completely invariant to the choice of baseline $b(s_t)$. Therefore, minimizing this variance upper bound reduces entirely to minimizing the second factor: the conditional mean squared error $\EE\left[ \left( {\color{colorreturn}G_t} - b(s_t) \right)^2 \ \Big|\ s_t \right]$.

For any random variable, the constant that minimizes its mean squared deviation is always its expected value! Setting the baseline to the expected reward-to-go conditional on state $s_t$ gives: \begin{equation}\label{eq:expected_return_baseline}\tag{RET} \bbox[10px, border: 1.5px solid #81A1C1]{ b(s_t) = \EE_{\tau \sim {\color{colorpolicy}\pi_\theta}}\left[ {\color{colorreturn}G_t} \ \Big|\ s_t \right] }\,. \end{equation} This is commonly known as the state value function, often denoted $V(s_t)$ or $V^{\pi_\theta}(s_t)$. While minimizing the variance upper bound leads to $V(s_t)$, Weaver and Tao (2001) show that directly minimizing the exact parameter variance weights the return by the score norm: $b^*(s_t) = \frac{\EE[\|\nabla_\theta \log \pi_\theta(a_t \mid s_t)\|^2 G_t \mid s_t]}{\EE[\|\nabla_\theta \log \pi_\theta(a_t \mid s_t)\|^2 \mid s_t]}$. Because estimating score-weighted expectations is computationally difficult and $V(s_t)$ performs comparably in practice, algorithms almost universally adopt the state value function. By optimizing the structural variance bound from $\eqref{eq:variance_decomp}$, the choice of $V(s_t)$ emerges directly as a consequence of variance reduction.

To see the effect of subtracting a baseline during optimization, consider a toy 1D Gaussian policy trained via gradient ascent. At each step, gradient updates are driven by the slope of a sample surrogate objective constructed from a randomly sampled action. As illustrated in the animation below, without a baseline (left), uncentered rewards produce steep, high-variance surrogate curves that repeatedly deflect the optimization trajectory in unstable directions. By contrast, centering around a baseline (right) replaces the absolute reward with a low-variance centered return ${\color{colorreturn}G_t} - b(s_t)$, yielding calm, stable surrogate curves that guide parameters smoothly to the optimum.

The simulation script to generate this animation and comparative trajectory landscape is available at gen_objective_landscape_1d.py.

Quantifying the Quadratic Variance Bound

How much does the state-dependent baseline $b(s_t) = \EE[{\color{colorreturn}G_t} \mid s_t]$ reduce the overall policy gradient variance?

With $b(s_t) = \EE[{\color{colorreturn}G_t} \mid s_t]$, the centered reward-to-go ${\color{colorreturn}G_{b,t}} \defas {\color{colorreturn}G_t} - b(s_t)$ has zero conditional mean ($\EE[{\color{colorreturn}G_{b,t}} \mid s_t] = 0$). Its conditional second moment therefore coincides with its variance: \begin{equation}\label{eq:centered_moment_is_variance} \EE\left[ |{\color{colorreturn}G_{b,t}}|^2 \ \Big|\ s_t \right] = \text{Var}\left( {\color{colorreturn}G_t} \ \Big|\ s_t \right) \le (T - t) K_r = O(T - t)\,. \end{equation} Eliminating the squared mean drift $\left(\EE[{\color{colorreturn}G_t} \mid s_t]\right)^2 = \Theta((T-t)^2)$ from each step drops the per-step return second moment from quadratic $\Theta((T-t)^2)$ down to linear $O(T-t)$. Summing over all $T$ time steps reduces the total variance upper bound from cubic $O(T^3)$ down to quadratic $O(T^2)$. Below is the formal statement.

Assume bounded rewards, bounded per-step score norms, and geometrically decaying conditional reward covariances across time steps: \begin{align} |{\color{colorreturn}r_t}| &\le {\color{colorreturn}r_{\max}}\,, \label{eq:bounded_rewards} \\ \|\nabla_\theta \log {\color{colorpolicy}\pi_\theta(a_t \mid s_t)}\|^2 &\le C\,, \label{eq:bounded_scores} \\ |\text{Cov}({\color{colorreturn}r_i}, {\color{colorreturn}r_j} \mid s_t)| &\le M \rho^{|i - j|} \quad \text{for some } M \ge 0 \text{ and } \rho \in [0, 1)\,. \label{eq:mixing_assumption} \end{align} Then the variance of the baseline estimator $\hat{g}_b(\tau)$ with $b(s_t) = \EE[{\color{colorreturn}G_t} \mid s_t]$ satisfies: \begin{equation}\label{eq:baseline_variance_bound} \text{Var}(\hat{g}_b(\tau)) \le \boxed{T^2} ~K_r C = O(T^2)\,, \end{equation} where $K_r \defas {\color{colorreturn}r_{\max}}^2 + \frac{2 M \rho}{1 - \rho}$ is a constant independent of $T$.
The baseline estimator hence reduces the asymptotic variance upper bound from cubic $O(T^3)$ down to quadratic $O(T^2)$.

To verify this theoretical reduction, we simulated the baseline estimator across varying trajectory lengths $T \in \{5, 10, \ldots, 640\}$. As shown below, centering the return drops the empirical log-log slope down to $1.98$, confirming that in general environments this quadratic bound is tight (exhibiting $\Theta(T^2)$ empirical scaling) and yielding orders-of-magnitude lower variance as the horizon grows.

The complete simulation script to reproduce this figure is available at policy_gradient_variance.py.

In Practice. Learned state-dependent baselines are a cornerstone of modern reinforcement learning. For example, in AlphaGo, rewards are strictly episodic and terminal ($+1$ for a win, $-1$ for a loss), yet a deep neural network estimates the win probability $\EE[{\color{colorreturn}R(\tau)} \mid s_t]$ from any intermediate board state. Subtracting this state-dependent baseline centers each move's update around its current position value rather than the global game average, significantly reducing policy gradient variance and guiding tree search.

Similarly, in large language model instruction tuning with Reinforcement Learning from Human Feedback (RLHF), baselines mitigate the extreme variance of token-level policy gradients. While traditional actor-critic setups rely on a secondary parameterized value model, recent methods like Group Relative Policy Optimization (GRPO) eliminate the need for an auxiliary critic entirely. By sampling a group of response trajectories for each prompt and subtracting the group average reward, GRPO provides an inexpensive empirical baseline that stabilizes policy updates without the computational overhead of training and serving a separate value network.

Citing

If you find this blog post useful, please consider citing as:

Policy Gradients Part 2: Baselines, Fabian Pedregosa, 2026

with bibtex entry:

@misc{pedregosa2026policygradientbaselines,
  author = {Pedregosa, Fabian},
  title = {Policy Gradients Part 2: Baselines},
  year = {2026},
  howpublished = {\url{http://fa.bianp.net/blog/2026/policy-gradient-baselines/}},
}
  

References