为什么 A2C 中减去 Baseline 函数可以减小方差

在 PPO × Family 课程第一讲中，已经介绍了策略梯度（policy gradient）的基本原理。但是如果直接使用最朴素的策略梯度方法，会发现实际训练中梯度的方差很大，进而导致训练效果不佳。为了解决这一问题，需要引入基线（baseline）函数来尽可能减小梯度的方差，从而提升算法性能。首先回顾一下 policy gradient 中标准的优化公式：

$$\begin{array}{c}{\mathrm{\nabla }}_{\theta }{R}_{\theta }={\mathbb{E}}_{\tau }[\sum _{t=0}^{T-1}{G}_t(\tau ){\mathrm{\nabla }}_{\theta }\log p_{\theta }(a_t\mid s_t)],\\ \mathrm{\Delta }\theta =\eta {\mathrm{\nabla }}_{\theta }{R}_{\theta }.\end{array}$$

在算法实践中，我们使用采样来获得 ${\mathrm{\nabla }}_{\theta }{R}_{\theta }$ 的无偏估计 ${\mathrm{\nabla }}_{\theta }{\overline{R}}_{\theta }$：

$${\mathrm{\nabla }}_{\theta }{\overline{R}}_{\theta }=\frac{1}{N}\sum _{n=1}^N\sum _{t=0}^{T_n-1}{G}_t({\tau }^n){\mathrm{\nabla }}_{\theta }\log p_{\theta }(a_t^n\mid s_t^n).$$

可以看出，上式将期望改成了对 $N$ 组轨迹进行采样和平均的形式。

在实践中，网络的训练会因为 ${\mathrm{\nabla }}_{\theta }{\overline{R}}_{\theta }$ 的方差而出现不稳定：虽然均值相同，但每次采样数据的 ${\mathrm{\nabla }}_{\theta }{\overline{R}}_{\theta }$ 与真实的期望值有较大偏离，而环境本身奖励函数的随机性也会进一步加重这个问题。因此我们希望降低上述优化形式带来的梯度方差。

Baseline 函数正是为了解决这个问题而提出的。具体来说，将上述公式改写为包含 baseline 的版本：

$${\mathrm{\nabla }}_{\theta }{\overline{R}}_{\theta }=\frac{1}{N}\sum _{n=1}^N\sum _{t=0}^{T_n-1}(G_t({\tau }^n)-b(s_t^n)){\mathrm{\nabla }}_{\theta }\log p_{\theta }(a_t^n\mid s_t^n).$$

其中 $b(s_t)$ 是关于 $t$ 时刻状态的一个函数，被称作基线（baseline）函数。这里要注意：其实我们并没有限制 baseline 函数必须是某一特定的形式，只要它是 $s_t$ 的函数就可以保证更新公式的正确性。但是一般认为 $b(s_t)=V(s_t)$ 时是最优的，因为此时能使得 ${\mathrm{\nabla }}_{\theta }{\overline{R}}_{\theta }$ 的方差最小。

接下来我们要论证：

1. 为什么添加了 baseline 函数之后，仍然能够保持估计的无偏性；
2. 为什么添加了合适的 baseline 函数之后，可以减小方差。

1. 为什么添加 baseline 函数之后仍然能够保持梯度估计的无偏性

要证明这一点，本质上是要证明：

$${\mathbb{E}}_{\tau }[\sum _{t=0}^{T-1}{G}_t(\tau ){\mathrm{\nabla }}_{\theta }\log p_{\theta }(a_t\mid s_t)]={\mathbb{E}}_{\tau }[\sum _{t=0}^{T-1}(G_t(\tau )-b(s_t)){\mathrm{\nabla }}_{\theta }\log p_{\theta }(a_t\mid s_t)].$$

移项化简，即证：

$${\mathbb{E}}_{\tau }[\sum _{t=0}^{T-1}b(s_t){\mathrm{\nabla }}_{\theta }\log p_{\theta }(a_t\mid s_t)]=0.$$

考虑到对于任意 $t$ 时刻的 transition，都有：

$$\begin{array}{rl}{\mathbb{E}}_{a_t}[b(s_t){\mathrm{\nabla }}_{\theta }\log p_{\theta }(a_t\mid s_t)]& =\int \frac{{\mathrm{\nabla }}_{\theta }{p}_{\theta }(a_t\mid s_t)}{p_{\theta }(a_t\mid s_t)}{p}_{\theta }(a_t\mid s_t)b(s_t)\mathrm{d}{a}_t\\ & =b(s_t){\mathrm{\nabla }}_{\theta }\int p_{\theta }(a_t\mid s_t)\mathrm{d}{a}_t\\ & =b(s_t){\mathrm{\nabla }}_{\theta }1\\ & =0.\end{array}$$

因此对于整体轨迹而言，自然对每个时刻 $t$ 都成立，也就有：

$${\mathbb{E}}_{\tau }[\sum _{t=0}^{T-1}b(s_t){\mathrm{\nabla }}_{\theta }\log p_{\theta }(a_t\mid s_t)]=0.$$

至此，就证明了添加 baseline 函数之后，梯度的无偏性仍然能够得以保持。

2. 为什么添加 baseline 函数之后梯度估计的方差会降低

严格来说，并不是添加任何一种形式的 baseline 函数都能够减小方差；相反，不合适的 baseline 函数甚至会增大方差。因此这里要说明的是：

• 为什么添加合适的 baseline 函数可以减小方差；
• 为什么通常选用 $b(s_t)=V(s_t)$ 作为实践中的 baseline 函数。

为简化分析，考虑单个时间步 $t$ 的梯度估计项：

$$g_t=(G_t(\tau )-b(s_t)){\mathrm{\nabla }}_{\theta }\log p_{\theta }(a_t\mid s_t).$$

由于 baseline 的引入不改变期望，即 $\mathbb{E}[g_t]$ 保持不变，因此最小化 $\mathrm{Var}(g_t)$ 等价于最小化 $\mathbb{E}[\parallel g_t{\parallel }^2]$。在常见假设下（如 $G_t(\tau )-b(s_t)$ 与 ${\mathrm{\nabla }}_{\theta }\log p_{\theta }(a_t\mid s_t)$ 在给定 $s_t$ 时不相关），最优的 baseline 满足：

$$b^{\ast }(s_t)={\mathbb{E}}_{\tau }[G_t(\tau )\mid s_t]=V(s_t).$$

直观上，因为 $V(s_t)$ 是在状态 $s_t$ 下按照当前策略能够获得的期望回报，用实际回报 $G_t(\tau )$ 减去它之后，剩下的部分 ${\delta }_t=G_t(\tau )-V(s_t)$ 衡量的是当前轨迹相对于平均水平的“优势”。这相当于把奖励中由状态本身决定的部分去掉，只保留动作带来的相对好坏信息，从而降低了采样的绝对数值大小，减小了梯度估计的方差。

因此，选取合适的 baseline 函数可以减小对梯度的估计方差；而价值函数 $V(s_t)$ 正是在常见假设下使方差最小的最优 baseline，这也是实践中一般使用价值函数作为 baseline 的理论依据。不过值得注意的是，如果在不满足上述假设的一些决策环境里，这种方法的效果可能会大打折扣。