Baselines for Variance Reduction

The REINFORCE algorithm uses the full Monte Carlo return $G_t$ to scale the gradient $\nabla_\theta \log \pi(a_t|s_t; \theta)$. While this provides an unbiased estimate of the policy gradient, the high variance of $G_t$ often leads to noisy updates and slow convergence. Imagine an episode where all rewards are large and positive; $G_t$ will be large for all timesteps, pushing the probabilities of all actions taken higher, even if some actions were much better than others within that successful episode. Conversely, a poor episode might unduly penalize all actions taken.

To mitigate this, we can introduce a baseline function $b(s)$ that depends only on the state $s_t$. We subtract this baseline from the return $G_t$ in the policy gradient update:

$$\nabla_\theta J(\theta) \approx \frac{1}{N} \sum_{i=1}^N \sum_{t=0}^{T_i-1} \nabla_\theta \log \pi(a_{i,t}|s_{i,t}; \theta) (G_{i,t} - b(s_{i,t}))$$

Why is this valid? Crucially, subtracting a state-dependent baseline does not introduce bias into the gradient estimate. The expected value of the subtracted term is zero:

$$\mathbb{E}_{\pi_\theta} [\nabla_\theta \log \pi(a_t|s_t; \theta) b(s_t)] = \sum_{s} d^{\pi_\theta}(s) \sum_{a} \pi(a|s; \theta) [\nabla_\theta \log \pi(a|s; \theta) b(s)]$$ $$= \sum_{s} d^{\pi_\theta}(s) \sum_{a} \nabla_\theta \pi(a|s; \theta) b(s)$$ $$= \sum_{s} d^{\pi_\theta}(s) b(s) \nabla_\theta \sum_{a} \pi(a|s; \theta)$$ $$= \sum_{s} d^{\pi_\theta}(s) b(s) \nabla_\theta 1 = 0$$

Here, $d^{\pi_\theta}(s)$ is the state distribution under policy $\pi_\theta$. Since the expectation of the term we subtract is zero, the overall expectation of the gradient remains unchanged.

While it doesn't change the expectation, a well-chosen baseline can significantly reduce the variance of the gradient estimate. The intuition is to center the returns around a reference point. Instead of simply scaling the gradient by the total return $G_t$, we scale it by how much better or worse that return was compared to what we typically expect from state $s_t$. If $G_t > b(s_t)$, the action $a_t$ led to a better-than-expected outcome, and its probability should be increased. If $G_t < b(s_t)$, the action led to a worse-than-expected outcome, and its probability should be decreased (or increased less strongly).

Choosing the Baseline

What makes a good baseline? Ideally, $b(s_t)$ should be an estimate of the expected return from state $s_t$.

1. Simple Constant Baseline: One could use a running average of all returns encountered so far across all episodes. This is very simple but doesn't adapt to the value of specific states.
2. State-Value Function $V(s_t)$: A much more effective and common choice for the baseline is the state-value function $V(s_t) = \mathbb{E}_{\pi_\theta}[G_t | S_t = s_t]$. This represents the expected return if we follow the current policy $\pi_\theta$ starting from state $s_t$.

Using $V(s_t)$ as the baseline gives rise to the term $G_t - V(s_t)$. Recall that the action-value function is $Q(s_t, a_t) = \mathbb{E}_{\pi_\theta}[G_t | S_t = s_t, A_t = a_t]$, and the advantage function is defined as $A(s_t, a_t) = Q(s_t, a_t) - V(s_t)$. Since $G_t$ is a Monte Carlo sample of the return starting from $(s_t, a_t)$, the term $G_t - V(s_t)$ is actually an estimate of the advantage function $A(s_t, a_t)$.

Using the advantage estimate centers the scaling factor: actions leading to returns better than the average for that state get positive reinforcement, while actions leading to returns worse than average get negative reinforcement. This often results in a significantly lower variance compared to using the raw return $G_t$.

Example showing how subtracting a baseline ($V(s_t) \approx 10$ in this simplified case) centers the scaling factor around zero and reduces its variance compared to using raw returns $G_t$.

Learning the Value Function Baseline

Of course, we typically don't know the true $V(s_t)$ under the current policy. Therefore, we need to learn it concurrently with the policy. This is usually done using function approximation, often with another neural network. We introduce a value function approximator $V(s; w)$ with parameters $w$.

This value network $V(s; w)$ is trained to predict the expected return from a given state. Since REINFORCE uses Monte Carlo returns $G_t$, a common approach is to train the value network by minimizing the Mean Squared Error (MSE) between its predictions $V(s_t; w)$ and the observed returns $G_t$ from the collected trajectories:

$$L(w) = \mathbb{E}_{\pi_\theta} [(G_t - V(s_t; w))^2]$$

So, during training, we perform two main updates in each iteration (after collecting trajectories):

1. Update Policy Parameters $\theta$: Use the policy gradient with the learned baseline:

   $$\Delta \theta = \alpha_\theta \nabla_\theta \log \pi(a_t|s_t; \theta) (G_t - V(s_t; w))$$

   Note that we treat $V(s_t; w)$ as a fixed baseline value when calculating the policy gradient; we don't backpropagate the policy loss through the value network.

2. Update Value Function Parameters $w$: Minimize the MSE loss between predicted values and actual returns:

   $$\Delta w = \alpha_w \nabla_w (G_t - V(s_t; w))^2 = \alpha_w \cdot 2 (G_t - V(s_t; w)) (-\nabla_w V(s_t; w))$$

Here, $\alpha_\theta$ and $\alpha_w$ are the learning rates for the policy and value function, respectively.

This approach of using a learned value function $V(s; w)$ as a baseline for a policy gradient method forms the basis of Actor-Critic algorithms, which we will discuss in detail in the next chapter. The policy network $\pi(a|s; \theta)$ is the "Actor," deciding which actions to take, and the value network $V(s; w)$ is the "Critic," evaluating the states visited by the Actor and providing the baseline to reduce variance in the Actor's learning signal.