广义优势估计 GAE

GAE 提出了一种优势函数的一般形式，综合考虑了多步真实采样的累积回报和价值函数的估计，权衡偏差和方差，是目前 PPO 等策略梯度类方法的主流选择。本文将首先梳理策略梯度中 ${\Psi }_t$ 的多种不同选择，理解它们各自的偏差/方差特性，然后介绍 GAE 提出的优势函数的一般形式。

前置：

强化学习策略梯度

强化学习Q-Learning/DQN

基于累积回报的 ${\Psi }_t$

在之前，我们介绍了策略梯度的一般形式：
$${\nabla }_{\theta }J({\pi }_{\theta })={\mathbb{E}}_{\tau \sim {\pi }_{\theta }}\left[\sum _{t=0}^T{\Psi }_t{\nabla }_{\theta }\log {\pi }_{\theta }(a_t|s_t)\right]$$
其中 ${\Psi }_t$ 用来表示状态 $s_t$ 下，动作 $a_t$ 有多 “好”。策略梯度的目标就是要增大 “好” 动作出现的概率，减小 “不好” 动作出现的概率。

在之前策略梯度最简形式的推导中，我们一开始是简单地采用整个 rollout 的累积回报作为 ${\Psi }_t$：
$$\begin{gathered} {\Psi }_t=R(\tau )=\sum _{t=0}^T{r}_t \\ \text{\hspace{1em}(1)} \end{gathered}$$
后来我们对取 ${\Psi }_t=R(\tau )$ 的几个问题进行了分析，得到了两个优化版本的 ${\Psi }_t$。

首先，动作 $a_t$ 有多好，应该与其之前获得的奖励无关，而只与其被采取之后获得的奖励有关。从而得到了 reward-to-go 形式的累积回报：
$$\begin{gathered} {\Psi }_t=R^{\prime }(\tau )=\sum _{t^{\prime }=t}^T{r}_{t^{\prime }} \\ \text{\hspace{1em}(2)} \end{gathered}$$
第二，全正的 ${\Psi }_t$ 可能会导致训练前期恰好没有被采样到的高回报动作一直没有办法被探索到，因此我们进一步减去 baseline $b(s_t)$，即取
$$\begin{gathered} {\Psi }_t=R^{\prime \prime }(\tau )=\sum _{t^{\prime }=t}^T{r}_{t^{\prime }}-b(s_t) \\ \text{\hspace{1em}(3)} \end{gathered}$$
这样就能使得 ${\Psi }_t$ 有正有负，鼓励好的动作，同时打压不够好的动作。EGLP 定理显示，加入只与状态 $s_t$ 有关的 $b(s_t)$，不会影响整体的期望，因此加入 baseline 后整体还是一个无偏估计。

基于价值函数的 ${\Psi }_t$

以上是之前在策略梯度中介绍过的内容，这三种 ${\Psi }_t=R(\tau )$ 的形式都是在 policy 与 environment 的实际互动采样中收集的真实数据。从 bias 的角度来看，它当然是无偏的，但是由于需要进行多步采样，policy 本身有随机性，environment 的状态转移也有随机性，方差在多步采样中逐渐积累，导致最终的方差非常大。无偏但是方差大，如果我们每次都能采样足够多次，那当然没有问题。但实际中我们的采样次数是很有限的，因此方差是一个不得不考虑的问题，大方差的训练会很不稳定。

那么，我们怎么才能降低 ${\Psi }_t$ 的方差呢？在这里，我们想避免多步采样带来的累积方差，那就要直接一步估计出 $R$ 的期望值，而不是逐步采样到最后，这样虽然可能会引入一些 bias（模型的预测可能会不准），但是能够降低方差。这里的 biase/variance 权衡，与我们之前介绍过的 value-based 训练目标中，MC 方法与 TD 方法的差异是一个道理。

“累积回报的期望值”，这听起来非常熟悉，它正是我们之前介绍的 value-based 方法中价值函数的定义。

回顾一下：

状态价值函数，$V^\pi(s)=\mathbb{E}_{\tau\sim\pi}[R(\tau)|s_0=s] \notag \$，是指在状态 $s$，接下来按照特定的策略 $\pi$ 来进行动作，最终累积回报的期望。

动作-状态价值函数，$Q^\pi(s,a)=\mathbb{E}_{\tau\sim\pi}[R(\tau)|s_0=s,a_0=a] \notag \$，是指在状态 $s$，先强制采取动作 $a$（这个动作不一定符合当前策略 $\pi$，是人为采取的），接下来按照特定的策略 $\pi$ 来进行动作，最终累积回报的期望。

状态价值函数 $V$ 是动作-价值函数 $Q$ 在策略 $\pi$ 下的期望：$V^{\pi }(s)={\mathbb{E}}_{a\sim \pi }[Q^{\pi }(s,a)]$。

${\Psi }_t$ 的目标就是要表达出 “在状态 $s_t$ 下，动作 $a_t$ 有多好”，即需要的是 “状态 $s_t$ 下采取动作 $a_t$ 的期望回报”，而这正是动作-状态价值函数 $Q^{\pi }(s_t,a_t)$ 的定义。因此，将 $Q$ 函数作为 ${\Psi }_t$ 是一个合理的选择：
$$\begin{gathered} {\Psi }_t=Q^{\pi }(s_t,a_t) \\ \text{\hspace{1em}(4)} \end{gathered}$$

Actor-Criric 方法：用价值函数来表示 ${\Psi }_t$，同 value-based 方法中一样，我们需要训练一个价值网络来拟合价值函数，这样就将 policy-based 方法（策略梯度）和 value-based 方法（价值网络）结合起来了，这就是所谓的 actor-critic 方法。其中 actor 就是 policy network，即策略网络；critic 就是 value network，是价值网络，用来估计策略对应的价值函数。

更进一步地，如果再将 baseline 考虑在内，${\Psi }_t$ 要表达的就是：“在状态 $s_t$ 下，动作 $a_t$ 相比于该状态下其他动作（的期望）要好多少”，而 “该状态下其他动作的期望”，正是状态价值函数 $V^{\pi }(s_t,a_t)$ 的定义。因此，可以在 $Q$ 的基础上，再将 $V$ 作为 baseline，减去 baseline，即取：
$$\begin{gathered} {\Psi }_t=A^{\pi }(s_t,a_t)=Q^{\pi }(s_t,a_t)-V^{\pi }(s_t) \\ \text{\hspace{1em}(5)} \end{gathered}$$
这里的 $A^{\pi }(s_t,a_t)$ 就是所谓的优势函数。

但是在优势函数 $A$ 中，我们需要估计估计 $Q$ 和 $V$ 两个价值函数，这样就需要训练两个神经网络。这样的训练起来误差会比较大。我们再进行一些变换。首先将 $V$ 写成 TD 残差的形式 $V^{\pi }(s_t)=r_t+V^{\pi }(s_{t+1})$。

我们知道，$Q$ 是 $V$ 的期望：
$$Q^{\pi }(s_t,a_t)=\mathbb{E}[V^{\pi }(s_t)]=\mathbb{E}[r_t+V^{\pi }(s_{t+1})]$$
把期望拿掉，把 $Q$ 写成 TD 残差的形式并代入：
$$A^{\pi }(s_t,a_t)=Q^{\pi }(s_t,a_t)-V^{\pi }(s_t)=r_t+V^{\pi }(s_{t+1})-V^{\pi }(s_t)$$
这样我们就得到了优势函数 $A$ 的 TD 残差形式，它也可以作为 ${\Psi }_t$：
$$\begin{gathered} {\Psi }_t=r_t+V^{\pi }(s_{t+1})-V^{\pi }(s_t) \\ \text{\hspace{1em}(6)} \end{gathered}$$
我们拿掉了期望，会引入一些 bias，但是这样优势函数的估计就只需要训练一个状态价值网络 $V$ 就行了。

广义优势估计 GAE

至此，我们已经介绍了基于累积回报的 ${\Psi }_t$ （式 1-3），它们的特点是无偏但是高方差；基于价值函数函数的 ${\Psi }_t$（式 4-6），它们的特点是有偏但是低方差。这正是 GAE 论文 Introduction 中总结的 6 种 ${\Psi }_t$ 的选择，即下图中的 1-6。

那么 GAE 是做了个什么事儿呢？我们上面看到

• 基于累积回报的 ${\Psi }_t$ 是一直采样到最后，获取真实的累积奖励，是 MC 方法，无偏但是多次采样方差很大；
• 基于价值函数的 ${\Psi }_t$ 是用一个价值网络估计动作的价值（或优势），是 TD 方法，TD 残差形式下只有一步采样，方差小，但是毕竟网络估计可能不准，因此是有偏的估计

现在要解决的问题就是 偏差/方差 权衡的问题。

先补充说明一件事，我们之前都是用的无折扣的累积回报，将未来所有回报都直接加到 $a_t$ 的累积回报中。但是实际情况下，应当是更近获得的奖励受到当前动作 $a_t$ 的影响较大，更远获得的奖励则受到 $a_t$ 影响较小，因此我们一般会对未来获得的奖励累乘一个小于 1 的折扣因子 $\gamma$，因此上面的几种 ${\Psi }_t$ 形式都要改写。举例来说，reward-to-go 累积回报（式 2）应该写为：
$$\begin{gathered} R^{\prime }(\tau )=\sum _{l=0}^{\infty }{\gamma }^l{r}_t \\ \text{\hspace{1em}(7)} \end{gathered}$$
而优势函数的 TD 残差形式（式 6） 应该写为：
$$\begin{gathered} r_t+\gamma V^{\pi }(s_{t+1})-V^{\pi }(s_t)\stackrel{记作}{=}{\delta }_t \\ \text{\hspace{1em}(8)} \end{gathered}$$
我们现在将其记作 ${\delta }_t$。

GAE 怎么权衡偏差和方差呢？

首先将 $k$ 个这样的 $\delta$ 项加起来，记为 ${\hat{A}}_t^{(k)}$：
$$\begin{array}{rlrl}{\hat{A}}_t^{(1)}& :={\delta }_t& & =-V(s_t)+r_t+\gamma V(s_{t+1})\\ {\hat{A}}_t^{(2)}& :={\delta }_t+\gamma {\delta }_{t+1}& & =-V(s_t)+r_t+\gamma r_{t+1}+{\gamma }^{2}V(s_{t+2})\\ {\hat{A}}_t^{(3)}& :={\delta }_t+\gamma {\delta }_{t+1}+{\gamma }^2{\delta }_{t+2}& & =-V(s_t)+r_t+\gamma r_{t+1}+{\gamma }^2{r}_{t+2}+{\gamma }^{3}V(s_{t+3})\\ ...& & ...& \\ {\hat{A}}_t^{(k)}& :=\sum _{l=0}^{k-1}{\gamma }^l{\delta }_{t+l}& & =-V(s_t)+r_t+\gamma r_{t+1}+{\gamma }^{k-1}{r}_{t+k-1}+{\gamma }^{k}V(s_{t+k})\end{array}$$
可以看到，${\hat{A}}_t^{(k)}$ 相当于是先采样 $k$ 步，得到 $k$ 步的真实回报，然后估计剩余的价值函数 $\gamma V(s_{t+k})$，再减去 baseline $V(s_t)$。从 ${\hat{A}}_t^{(1)}={\delta }_t$ 推广（${\delta }_t$ 是优势函数的 TD 残差形式的估计），我们可以将所有 ${\hat{A}}_t^{(k)}$ 都看作是对优势函数的一种估计，

• $k$ 越大，真实采样的步数越多，对价值函数估计的依赖越小，偏差越小，方差越大；
• $k$ 越小，真实采样的步数越少，对价值函数估计的依赖越大，偏差越大，方差越小；

特别的，当 $k\to \infty$ 时，有
$${\hat{A}}_t^{(\infty )}=\sum _{l=0}^{\infty }{\gamma }^l{\delta }_{t+l}=-V(s_t)+\sum _{l=0}^{\infty }{\gamma }^l{r}_{t+l}$$
这样就是一直采样到结束，相当于又成了基于累积回报的 ${\Psi }_t$（式 7），只是 baseline 采用了价值函数的估计 $V(s_t)$。

GAE，就是对这一系列不同的对优势函数的估计 ${\hat{A}}_t^{(k)}$ 的指数加权平均：
$$\begin{array}{rl}{\hat{A}}_t^{\text{GAE}}& :=(1-\lambda )({\hat{A}}_t^{(1)}+\lambda {\hat{A}}_t^{(2)}+{\lambda }^2{\hat{A}}_t^{(3)}+...)\\ & =(1-\lambda )({\delta }_t+\lambda ({\delta }_t+\gamma {\delta }_{t+1})+{\lambda }^2({\delta }_{t+1}+\gamma {\delta }_{t+1}+{\gamma }^2{\delta }_{t+2})+...)\\ & =(1-\lambda )({\delta }_t(1+\lambda +{\lambda }^2+...)+\gamma {\delta }_{t+1}(\lambda +{\lambda }^2+{\lambda }^3+...)+{\gamma }^2{\delta }_{t+2}({\lambda }^2+{\lambda }^3+{\lambda }^4+...)+...)\\ & =(1-\lambda )({\delta }_t\frac{1}{1-\lambda }+\gamma {\delta }_{t+1}\frac{\lambda }{1-\lambda }+{\gamma }^2{\delta }_{t+2}\frac{{\lambda }^2}{1-\lambda }+...)\\ & =\sum _{l=0}^{\infty }(\gamma \lambda {)}^l{\delta }_{t+l}\end{array}$$
（这里乘上个 $(1-\lambda )$ 是为了化简形式？）可以看到，我们引入一个超参数 $\lambda$，对所有的 ${\hat{A}}_t^{(k)}$ 进行加权表示，得到了一个一般形式的对优势函数的估计，也就是 “广义优势估计 Generalized Advantage Estimation”，一般直接简记作 ${\hat{A}}_t$。这样，我们就可以通过调节超参数 $\lambda$，来对偏差和方差进行 trade-off。

• $\lambda$ 越小，$k$ 值高的项 ${\hat{A}}_t^{(k)}$ 衰减得越厉害，价值网络越主导，偏差越大，方差越小；
• $\lambda$ 越大，$k$ 值高的项 ${\hat{A}}_t^{(k)}$ 衰减得越少，真实采样越主导，偏差越小，方差越大；

特别地，当 $\lambda =0$ 或 $\lambda =1$ 时，有：
$$\begin{array}{rl}\lambda =0,{\hat{A}}_t& :={\delta }_t=r_t+\gamma V^{\pi }(s_{t+1})-V^{\pi }(s_t)\\ \lambda =1,{\hat{A}}_t& :=\sum _{l=0}^{\infty }{\gamma }^l{\delta }_{t+l}=\sum _{l=0}^{\infty }{\gamma }^l{r}_{t+l}-V(s_t)\end{array}$$
分别对应了只一次真实采样，偏差最大、方差最小的情况，和全都真实采样，方差最大、偏差最小的情况。

顺便提一句，GAE 这种构造形式和 $\text{TD}(\lambda )$ 很像，但是 $\text{TD}(\lambda )$ 是对价值函数的估计，而 GAE 是对优势函数的估计。

最终，我们在策略梯度中，就可以将 GAE ${\hat{A}}_t$ 作为 ${\Psi }_t$；

$$\begin{gathered} {\nabla }_{\theta }J({\pi }_{\theta })={\mathbb{E}}_{\tau \sim {\pi }_{\theta }}\left[\sum _{t=0}^T{\hat{A}}_t{\nabla }_{\theta }\log {\pi }_{\theta }(a_t|s_t)\right] \\ {\hat{A}}_t=\sum _{l=0}^{\infty }(\gamma \lambda {)}^l{\delta }_{t+l},{\delta }_t=r_t+\gamma V^{\pi }(s_{t+1})-V^{\pi }(s_t) \end{gathered}$$

总结

本文中，我们首先回顾了之前介绍过的基于累积回报的 ${\Psi }_t$ 选择，这类选择虽然是无偏的，但是由于需要多次采样，会造成累积方差很大。为了减小方差，我们考虑了基于价值函数的 ${\Psi }_t$ 选择，但是在价值网络估计的价值函数不准时，这类选择的偏差又比较大。GAE 提出了一种一般形式，将多项对优势函数的估计值进行加权表示，通过调节超参数 $\lambda$，实现了偏差和方差之间进行权衡，目前已经成为策略梯度类算法中优势函数形式的主流选择。