Understanding Baselines for Variance Reduction in Policy Gradient Methods

Introduction

Deep reinforcement learning has achieved remarkable success in recent years, enabling agents to learn complex behaviors in high-dimensional environments. Policy gradient methods are a key class of deep RL algorithms that directly optimize the policy to maximize expected reward. While powerful and widely applicable, policy gradient methods often suffer from poor sample efficiency and unstable learning, largely due to the high variance of the gradient estimator.

Fortunately, we can mitigate this issue through the use of a baseline function. In this post, we‘ll take a deep dive into what baselines are, why they help, and how to choose an effective one. Whether you‘re an aspiring deep RL researcher or practitioner, understanding baselines is crucial for successfully applying policy gradients to real-world problems. Let‘s jump in!

The Trouble with Raw Policy Gradients

To understand the need for baselines, let‘s briefly review the standard policy gradient algorithm. The goal is to find a policy $\pi_\theta(a|s)$, parameterized by $\theta$, that maximizes the expected discounted return:

$$J(\theta) = \mathbb{E}{\tau \sim p\theta(\tau)}[R(\tau)]$$

where $\tau = (s_0, a_0, r_0, s_1, a_1, r1, …)$ is a trajectory, $p\theta(\tau)$ is the probability of $\tau$ under policy $\pi\theta$, and $R(\tau) = \sum{t=0}^T \gamma^t r_t$ is the discounted return.

The policy gradient theorem tells us that the gradient of $J(\theta)$ with respect to the policy parameters is:

$$\nabla\theta J(\theta) = \mathbb{E}{\tau \sim p\theta(\tau)} [\sum{t=0}^T \nabla\theta \log \pi\theta(a_t|s_t) R(\tau)]$$

We can estimate this expectation by collecting a batch of trajectories ${\tau^{(i)}}_{i=1}^N$ and computing the Monte Carlo estimate:

$$\nabla\theta J(\theta) \approx \frac{1}{N} \sum{i=1}^N \sum{t=0}^T \nabla\theta \log \pi_\theta(a_t^{(i)}|s_t^{(i)}) R(\tau^{(i)})$$

Seems straightforward, right? The problem is that this gradient estimator has very high variance. Even with a large batch size $N$, the estimates can vary wildly from update to update, leading to unstable learning and poor sample efficiency. This is where baselines come to the rescue.

Baselines to the Rescue

The key idea behind a baseline is to subtract a quantity $b(s_t)$ from the return $R(\tau)$ at each timestep $t$. Importantly, $b(s_t)$ should be a function only of the state $s_t$, not the action $a_t$. With a baseline, our policy gradient estimate becomes:

$$\nabla\theta J(\theta) \approx \frac{1}{N} \sum{i=1}^N \sum{t=0}^T \nabla\theta \log \pi_\theta(a_t^{(i)}|s_t^{(i)}) (R(\tau^{(i)}) – b(s_t^{(i)}))$$

Why does this help? Let‘s look at it from two angles: an intuitive view and a mathematical derivation of the optimal baseline.

Intuition: Reducing Gradient Noise

To build intuition, consider a simple bandit problem with only a single state. Our policy $\pi_\theta(a)$ specifies a distribution over the $k$ possible actions, and the environment gives a random reward $r \sim R(a)$ for the chosen action.

Without a baseline, the policy gradient estimate is proportional to $r \nabla\theta \log \pi\theta(a)$ for each sample. Since $r$ can be positive or negative, this estimate will be very noisy, pulling $\theta$ in different directions.

Now let‘s introduce a baseline $b$. The gradient is now proportional to $(r – b) \nabla\theta \log \pi\theta(a)$. If we choose $b$ close to the average reward, then $r – b$ will be positive for better-than-average actions and negative for worse-than-average actions. This pushes the policy to favor good actions and avoid bad ones, rather than just getting pushed around by the noise in the rewards. The closer $b$ is to $\mathbb{E}[r]$, the more it will reduce variance while keeping the expected value of the gradient unchanged.

The Optimal Baseline

More formally, we can derive the optimal state-dependent baseline $b^*(s)$ that minimizes the variance of the policy gradient estimate, while keeping its expectation unchanged. The variance of an estimator $\hat{x}$ is defined as:

$$\text{Var}[\hat{x}] = \mathbb{E}\big[(\hat{x} – \mathbb{E}[\hat{x}])^2\big]$$

Applying this to the policy gradient estimator $\hat{g} = \frac{1}{N} \sum{i=1}^N \sum{t=0}^T (R(\tau^{(i)}) – b(st^{(i)})) \nabla\theta \log \pi_\theta (a_t^{(i)}|s_t^{(i)})$:

$$\begin{aligned}
\text{Var}[\hat{g}] &= \mathbb{E} \Big[ \big(\hat{g} – \mathbb{E}[\hat{g}]\big)^2 \Big] \
&= \mathbb{E} \Big[ \big(\hat{g} – \nabla\theta J(\theta) \big)^2 \Big] \
&= \mathbb{E} \bigg[ \Big( \frac{1}{N} \sum{i=1}^N \sum_{t=0}^T (R(\tau^{(i)}) – b(st^{(i)}) – Q^{\pi\theta}(s_t^{(i)},at^{(i)})) \nabla\theta \log \pi_\theta (a_t^{(i)}|s_t^{(i)}) \Big)^2 \bigg] \end{aligned}$$

where $Q^{\pi\theta}(s,a) = \mathbb{E}{\tau \sim p_\theta(\tau|s_0=s, a0=a)}[R(\tau)]$ is the state-action value function. The last step follows from the fact that $\nabla\theta J(\theta) = \mathbb{E}{s \sim d^{\pi\theta}, a \sim \pi\theta}[Q^{\pi\theta}(s,a) \nabla\theta \log \pi\theta(a|s)]$.

To minimize this variance with respect to $b$, we take the derivative, set it to zero, and solve for the optimal baseline $b^*(s)$:

$$\begin{aligned}
0 &= \frac{\partial}{\partial b(s)} \text{Var}[\hat{g}] \
&= \mathbb{E}{a \sim \pi\theta(\cdot|s)} \big[ (Q^{\pi\theta}(s,a) – b^*(s)) \nabla\theta \log \pi\theta(a|s) \big] \
b^*(s) &= \frac{\mathbb{E}{a \sim \pi\theta(\cdot|s)} \big[ Q^{\pi\theta}(s,a) \nabla\theta \log \pi\theta(a|s) \big]}{\mathbb{E}{a \sim \pi\theta(\cdot|s)}[\nabla\theta \log \pi\theta(a|s)]}
\end{aligned}$$

Intuitively, $b^*(s)$ is a weighted average of the Q-values $Q^{\pi_\theta}(s,a)$, with weights proportional to how much each action contributes to the policy gradient. This makes sense: we want to subtract off the average Q-value according to the current policy to center the gradient estimates.

Baseline Functions in Practice

Let‘s look at some common choices of baselines used in practice:

Constant Baseline

The simplest choice is a constant baseline $b(s) = c$ for all states. The most common constant is the average return across the batch:

$$b = \frac{1}{N} \sum_{i=1}^N R(\tau^{(i)})$$

This can be interpreted as setting a "reference point" representing how well the policy is doing overall. Actions leading to higher-than-average return push the policy in that direction, while those with lower return push the policy away.

State-Dependent Baseline

A natural extension is to allow the baseline to vary with the state: $b_\phi(s) \approx V^{\pi_\theta}(s)$, where $\phi$ are learned parameters. The state-value function $V^{\pi}(s)$ measures the expected return when starting in state $s$ and following policy $\pi$. It serves as a perfect baseline since $Q^{\pi}(s,a) – V^{\pi}(s)$ has zero mean for all $s$.

In practice, $V^{\pi}$ is not known and must be estimated. A common approach is to train a function approximator $V_\phi(s)$ to minimize the mean squared error:

$$\mathcal{L}(\phi) = \frac{1}{N} \sum{i=1}^N \sum{t=0}^T \big(V_\phi(s_t^{(i)}) – R_t^{(i)}\big)^2$$

where $Rt = \sum{t‘=t}^T \gamma^{t‘-t} r{t‘}$ is the return from step $t$ onwards. With a sufficiently expressive $V\phi(s)$ and enough data, this will approach the true $V^{\pi_\theta}(s)$ as $\phi$ is optimized.

State-Action-Dependent Baseline

The optimal baseline derived earlier depended on both state $s$ and action $a$, but this is tricky to estimate directly, since the expectation over $a \sim \pi(\cdot|s)$ requires knowledge of the current policy.

One solution is to use an estimator $Q\phi(s,a) \approx Q^{\pi\theta}(s,a)$ for the state-action value function, learned similarly to $V_\phi(s)$:

$$\mathcal{L}(\phi) = \frac{1}{N} \sum{i=1}^N \sum{t=0}^T \big(Q_\phi(s_t^{(i)}, a_t^{(i)}) – R_t^{(i)}\big)^2$$

Then, we can form a baseline using the expectation of $Q_\phi$ under the current policy:

$$b(s) = \mathbb{E}{a \sim \pi\theta(\cdot|s)} \big[ Q_\phi(s,a) \big] \approx \frac{1}{Ns} \sum{i=1}^{Ns} Q\phi(s, a^{(i)}) \text{ where } a^{(i)} \sim \pi_\theta(\cdot|s)$$

This combines the benefits of a state-dependent baseline with the optimal weighting by the policy gradient contributions. The downside is that it requires more computation and potentially more samples to train $Q_\phi$ well.

Wrapping Up

Baselines are a key technique for stabilizing policy gradient learning in deep reinforcement learning. By subtracting a state-dependent function from the returns, we can reduce gradient variance without introducing bias. The optimal baseline is a weighted average of state-action values, but in practice simpler baselines like a constant average return or learned state-value function are very effective.

When applying policy gradients to a new problem, consider starting with a simple constant baseline and then moving to a learned state-value function if needed. For the most challenging problems, a state-action-dependent baseline can squeeze out even more efficiency, at the cost of some added complexity.

I hope this post has demystified baselines and helped you understand this important concept in policy gradients. While we focused on basic algorithms here, baselines have been extended to fancier policy gradient variants like PPO and TRPO – but that‘s a topic for another time. For now, go forth and conquer RL problems with the power of baselines!

References and Further Reading

• Williams, R. J. (1992). Simple statistical gradient-following algorithms for connectionist reinforcement learning. Machine learning.
• Sutton, R. S., McAllester, D. A., Singh, S. P., & Mansour, Y. (1999). Policy gradient methods for reinforcement learning with function approximation. NeurIPS.
• Greensmith, E., Bartlett, P. L., & Baxter, J. (2004). Variance reduction techniques for gradient estimates in reinforcement learning. JMLR.
• Wu, C., Rajeswaran, A., Duan, Y., Kumar, V., Bayen, A. M., Kakade, S., … & Abbeel, P. (2018). Variance reduction for policy gradient with action-dependent factorized baselines. ICLR.
• Tucker, G., Bhupatiraju, S., Gu, S., Turner, R. E., Ghahramani, Z., & Levine, S. (2018). The mirage of action-dependent baselines in reinforcement learning. ICML.