Baselines for Variance Reduction

The REINFORCE algorithm is a fundamental Monte Carlo policy gradient method. Understanding its update rule for the policy parameters $\theta$ is essential for effectively implementing baselines for variance reduction.

$$\theta \leftarrow \theta + \alpha G_t \nabla_\theta \ln \pi_\theta(a_t|s_t)$$

Here, $G_t$ is the total discounted return starting from time step $t$. While this update rule correctly pushes the parameters in the direction of increased expected return, it often suffers from a significant drawback: high variance.

The Problem of High Variance

The return $G_t$ is calculated based on a single, complete trajectory sampled according to the current policy $\pi_\theta$. The rewards received along this trajectory can vary substantially due to the stochastic nature of the environment and the policy itself. A slight change in an early action might lead to a completely different sequence of states and rewards, resulting in a very different $G_t$.

This high variance in $G_t$ translates directly into high variance in the gradient estimates $\nabla_\theta \ln \pi_\theta(a_t|s_t)$. Imagine two different episodes starting from the same state $s_t$. In one episode, the agent gets lucky and achieves a high return $G_t = 10$. In another, perhaps due to an unlucky transition, the return is much lower, $G_t = -5$. The resulting gradient updates for the policy parameters $\theta$ will point in very different directions, even if the action $a_t$ taken was the same.

This noisy gradient signal makes learning inefficient and unstable. The parameter updates can oscillate wildly, requiring many samples (episodes) to average out the noise and make consistent progress. Consequently, the learning process can be extremely slow.

Introducing Baselines

To address this high variance, we can modify the policy gradient update by subtracting a baseline $b(s_t)$ from the return $G_t$. The important requirement for the baseline is that it must not depend on the action $a_t$; it should only depend on the state $s_t$.

The modified update rule becomes:

$$\theta \leftarrow \theta + \alpha (G_t - b(s_t)) \nabla_\theta \ln \pi_\theta(a_t|s_t)$$

Why does this help? Intuitively, we are no longer just increasing the probability of an action based on the raw return $G_t$. Instead, we are comparing the obtained return $G_t$ to an expected or average return $b(s_t)$ for being in state $s_t$.

• If $G_t > b(s_t)$, the term $(G_t - b(s_t))$ is positive. This means the action $a_t$ led to a better-than-expected outcome for state $s_t$, so we increase its probability.
• If $G_t < b(s_t)$, the term $(G_t - b(s_t))$ is negative. The action $a_t$ led to a worse-than-expected outcome, so we decrease its probability.

This comparison focuses the update on the relative quality of the action $a_t$ given the state $s_t$, rather than the absolute return $G_t$, which can fluctuate greatly. By centering the returns around a baseline, we reduce the magnitude and variability of the update term, leading to lower variance in the gradient estimates and more stable learning.

Maintaining an Unbiased Gradient

A critical property of the baseline $b(s_t)$ is that subtracting it does not change the expected value of the gradient update. In other words, it reduces variance without introducing bias. We can show this mathematically. The expected value of the term we subtracted, scaled by the gradient log-probability, is zero:

$$E_{\pi_\theta} [b(s_t) \nabla_\theta \ln \pi_\theta(a_t|s_t)] = \sum_s d(s) \sum_a \pi_\theta(a|s) b(s) \nabla_\theta \ln \pi_\theta(a|s)$$

Using the identity $\nabla_\theta \ln \pi_\theta(a|s) = \frac{\nabla_\theta \pi_\theta(a|s)}{\pi_\theta(a|s)}$, we get:

$$= \sum_s d(s) b(s) \sum_a \pi_\theta(a|s) \frac{\nabla_\theta \pi_\theta(a|s)}{\pi_\theta(a|s)}$$ $$= \sum_s d(s) b(s) \sum_a \nabla_\theta \pi_\theta(a|s)$$

Since the sum of probabilities $\sum_a \pi_\theta(a|s)$ must equal 1 for any state $s$, its gradient with respect to $\theta$ is zero:

$$= \sum_s d(s) b(s) \nabla_\theta \sum_a \pi_\theta(a|s)$$ $$= \sum_s d(s) b(s) \nabla_\theta (1) = 0$$

Because the expected value of the subtracted term is zero, the expected value of the overall gradient remains unchanged. The baseline successfully reduces variance without altering the direction the policy parameters should move in, on average.

Choosing the Baseline: The State-Value Function

What makes a good baseline $b(s_t)$? An effective and widely used choice is an estimate of the state-value function, $V(s_t)$. The state-value function $V(s_t)$ represents the expected return starting from state $s_t$ and following the current policy $\pi_\theta$. This is precisely the kind of "expected outcome" we want to compare $G_t$ against.

So, we use an approximation $\hat{v}(s_t, w)$, parameterized by weights $w$, as our baseline:

$$b(s_t) = \hat{v}(s_t, w)$$

The policy update then becomes:

$$\theta \leftarrow \theta + \alpha (G_t - \hat{v}(s_t, w)) \nabla_\theta \ln \pi_\theta(a_t|s_t)$$

The term $(G_t - \hat{v}(s_t, w))$ is now an estimate of the advantage of taking action $a_t$ in state $s_t$. Recall that the advantage function is defined as $A(s, a) = Q(s, a) - V(s)$. Since $G_t$ is a Monte Carlo sample (an estimate) of $Q^{\pi_\theta}(s_t, a_t)$ (under policy $\pi_\theta$ for the specific action $a_t$ taken) and $\hat{v}(s_t, w)$ estimates $V^{\pi_\theta}(s_t)$, their difference approximates the advantage $A^{\pi_\theta}(s_t, a_t)$. Using the advantage estimate generally leads to lower variance than using the raw return $G_t$.

How do we get the value function estimate $\hat{v}(s_t, w)$? We need to learn it! This value function can be learned concurrently with the policy using methods we've already encountered, such as Monte Carlo prediction or Temporal-Difference learning (like TD(0)), often using techniques like function approximation (linear or neural networks) if the state space is large. This involves updating the value function parameters $w$ based on the observed returns or TD errors.

This approach, where we learn both a policy (the "actor") and a value function (the "critic" which provides the baseline), forms the foundation for Actor-Critic methods, which we will introduce in the next section. By having the critic evaluate the actor's actions relative to the expected value of the state, we can significantly stabilize and accelerate policy gradient learning.