GAE 算法
准备
众所周知, 策略梯度有多种写法, 总的来说, 在保持策略梯度不变的情况下, 策略梯度可以写作
$$
g=\mathbb{E}\left[\sum_{t=0}^{\infty}\Psi_t\nabla_\theta \log\pi_0(a_t\mid s_t)\right]\tag{1}
$$
其中 $\Psi$ 可以是
$$
\begin{aligned}
1.\;&\sum\nolimits_{t=0}^\infty r_t&轨迹的总回报\\\\
2.\;&\sum\nolimits_{t'=t}^\infty r_{t'}&动作后轨迹的总回报\\\\
3.\;&\sum\nolimits_{t'=t}^\infty r_{t'}-b(s_t)&基线形式\\\\
4.\;&Q^{\pi}(s_t,a_t)&状态-动作价值函数\\\\
5.\;&A^{\pi}(s_t,a_t)&优势函数\\\\
6.\;&r_t+V^{\pi}(s_{t+1})-V^\pi(s_t)&\text{TD}\,残差
\end{aligned}
$$
其中
$$
\begin{align}
V^\pi(s_t):=\mathbb{E}_{\substack{s_{t+1:\infty},\\\\a_{t:\infty}}}\left[\sum_{l = 0}^\infty r_{t+l}\right]\qquad Q^\pi(s_t, a_t):=\mathbb{E}_{\substack{s_{t+1:\infty},\\\\a_{t+1:\infty}}}\left[\sum_{l = 0}^\infty r_{t+l}\right]\tag{2}\\\\
A^{\pi}(s_t,a_t)=Q^{\pi}(s_t,a_t)-V^{\pi}(s_t)\tag{3}
\end{align}
$$
这里逗号表示 $a:b$ 指的是 $(a,a+1,\dots,b)$ 这样的序列, $\mathbb{E}$ 的下标枚举了要被积分的变量. 前面 $5$ 项的推导或者推导资料在参数优化中都有说明, 而 $\text{TD}$ 残差其实是优势函数的一种无偏估计.
其中令 $\Psi_t=A^\pi(s_t, a_t)$ (优势函数) 的选择有几乎最小的方差. 这一点可以从策略梯度的角度直观的解释: 策略梯度中的每一步都会增加 “高于平均水平的动作” 的概率, 减少 “低于平均水平的动作” 的概率, 而优势函数 $A^{\pi}(s_t,a_t)=Q^{\pi}(s_t,a_t)-V^{\pi}(s_t)$ 恰好衡量了动作相对平均水平的好坏, 当动作高于平均水平时, 优势函数会取正数, 从而增加其概率; 当动作低于平均水平时, 优势函数会取负数. 从而降低其概率.
我们利用一个参数 $\gamma$ 来降低回报对延迟效应的反应的权重 (即减少未来回报的影响) 来减少方差, 代价是引入偏差. 这个参数相当于有折损的 $\text{MDPs}$ 公式, 但是我们将其当做一个在无折损问题中的一个减少方差的参数. 这些有折损的公式可以表示为
$$
\begin{align}
V^{\pi, \gamma}(s_t):=\mathbb{E}_{\substack{s_{t+1:\infty},\\\\a_{t:\infty}}}\left[\sum_{l = 0}^\infty \gamma^l r_{t+l}\right]\qquad Q^{\pi, \gamma}(s_t, a_t):=\mathbb{E}_{\substack{s_{t+1:\infty},\\\\a_{t+1:\infty}}}\left[\sum_{l = 0}^\infty \gamma^l r_{t+l}\right]\tag{4}\\\\
A^{\pi, \gamma}(s_t,a_t)=Q^{\pi, \gamma}(s_t,a_t)-V^{\pi, \gamma}(s_t)\tag{5}
\end{align}
$$
梯度的有折损近似可以表示为
$$
g^\gamma:=\mathbb{E}_{\substack{s_{0:\infty},\\\\a_{0:\infty}}}\left[\sum_{t=0}^{\infty}A^{\pi, \gamma}(s_t,a_t)\nabla_\theta \log\pi_0(a_t\mid s_t)\right]\tag{6}
$$
由于优势函数是未知的, 所以我们需要对它进行估计. 在实践中, 往往只能学习到 $A^{\pi,\gamma}$ 的一个有偏估计 (但没有那么偏) .
我们引入一个关于 $A^{\pi,\gamma}$ 的一个估计, 并且这个估计是无偏的. 考虑一个与整个轨迹有关的优势函数的估计 $\hat{A}_t(s_{0:\infty},a_{0:\infty})$ .
我们定义: 一个估计 $\hat{A}_t$ 是 $\gamma\text{-just}$ 的当且仅当
$$
\mathbb{E}_{\substack{s_{0:\infty},\\\\a_{0:\infty}}}\left[\hat{A}_t(s_{0:\infty},a_{0:\infty})\nabla_\theta \log\pi_0(a_t\mid s_t)\right]=\mathbb{E}_{\substack{s_{0:\infty},\\\\a_{0:\infty}}}\left[A^{\pi, \gamma}(s_t,a_t)\nabla_\theta \log\pi_0(a_t\mid s_t)\right]\tag{7}
$$
因此如果对于所有的 $t$ 来说 $\hat{A}_t$ 都是 $\gamma\text{-just}$ 的话, 就有
$$
\mathbb{E}_{\substack{s_{0:\infty},\\\\a_{0:\infty}}}\left[\sum_{t=0}^{\infty}\hat{A}_t(s_{0:\infty},a_{0:\infty})\nabla_\theta \log\pi_0(a_t\mid s_t)\right]=g^\gamma\tag{8}
$$
一个令 $\hat{A}_t$ 为 $\gamma\text{-just}$ 的充分条件是 $\hat{A}_t$ 可以被分解为两个函数 $Q_t$ 和 $b_t$ 之间的差. 并且 $Q_t$ 可以依赖于轨迹中的任意变量, 但必须是 $\gamma$-折扣 的 $Q$ 函数 (状态-动作价值函数) 的无偏估计 (如果要求是状态-动作价值函数的估计, 那么由于马尔科夫性, 在时间 $t$ 之前的变量都对状态-动作价值函数的值没有影响, 因此 $Q_t$ 只被时间 $t$ 以及 $t$ 以后的变量决定), 而 $b_t$ 是先于 $a_t$ 被采样的动作和状态的任意函数. 即如果 $\hat{A}_t$ 可以被写成 $\hat{A}_t(s_{0:\infty},a_{0:\infty})=Q_t(s_{t:\infty},a_{t:\infty})-b_t(s_{0:t},a_{0:t-1})$ 的形式, 并且对于任意的 $(s_t,a_t)$ 有 $\mathbb{E}_{s_{t+1:\infty},a_{t+1:\infty}\mid s_t,a_t}[Q_t(s_{t:\infty},a_{t:\infty})]=Q^{\pi, \gamma}(s_t,a_t)$ , 那么 $\hat{A}$ 就是 $\gamma\text{-just}$ 的.
证明:
$$
\begin{align}
&\mathbb{E}_{\substack{s_{0:\infty},\\\\a_{0:\infty}}}\left[\hat{A}_t(s_{0:\infty},a_{0:\infty})\nabla_\theta \log\pi_0(a_t\mid s_t)\right]\\\\
=&\mathbb{E}_{\substack{s_{0:\infty},\\\\a_{0:\infty}}}\left[Q_t(s_{t:\infty},a_{t:\infty})\nabla_\theta \log\pi_0(a_t\mid s_t)\right]-\mathbb{E}_{\substack{s_{0:\infty},\\\\a_{0:\infty}}}\left[b_t(s_{0:t},a_{0:t-1})\nabla_\theta \log\pi_0(a_t\mid s_t)\right]
\end{align}
$$
其中
$$
\begin{align}
&\mathbb{E}_{\substack{s_{0:\infty},\\\\a_{0:\infty}}}\left[Q_t(s_{t:\infty},a_{t:\infty})\nabla_\theta \log\pi_0(a_t\mid s_t)\right]\\\\
=&\mathbb{E}_{\substack{s_{0:t},\\\\a_{0:t}}}\left[\mathbb{E}_{\substack{s_{t+1:\infty},\\\\a_{t+1:\infty}}}\left[Q^{\pi,\gamma}(s_t,a_t)\nabla_\theta \log\pi_0(a_t\mid s_t)\right]\right]\\\\
=&\mathbb{E}_{\substack{s_{0:t},\\\\a_{0:t}}}\left[\mathbb{E}_{\substack{s_{t+1:\infty},\\\\a_{t+1:\infty}}}\left[A^{\pi, \gamma}(s_t,a_t)\nabla_\theta \log\pi_0(a_t\mid s_t)\right]\right]\\\\
=&\mathbb{E}_{\substack{s_{0:\infty},\\\\a_{0:\infty}}}\left[A^{\pi, \gamma}(s_t,a_t)\nabla_\theta \log\pi_0(a_t\mid s_t)\right]
\end{align}
$$
而
$$
\begin{align}
&\mathbb{E}_{\substack{s_{0:\infty},\\\\a_{0:\infty}}}\left[b_t(s_{0:t},a_{0:t-1})\nabla_\theta \log\pi_0(a_t\mid s_t)\right]\\\\
=&\mathbb{E}_{\substack{s_{0:t},\\\\a_{0:t-1}}}\left[\mathbb{E}_{\substack{s_{t+1:\infty},\\\\a_{t:\infty}}}\left[b_t(s_{0:t},a_{0:t-1})\nabla_\theta \log\pi_0(a_t\mid s_t)\right]\right]\\\\
=&\mathbb{E}_{\substack{s_{0:t},\\\\a_{0:t-1}}}\left[\mathbb{E}_{\substack{s_{t+1:\infty},\\\\a_{t:\infty}}}\left[\nabla_\theta \log\pi_0(a_t\mid s_t)\right]b_t(s_{0:t},a_{0:t-1})\right]
\end{align}
$$
根据 $\text{EGLP}$ 定理 (见参数优化) 有
$$
\mathbb{E}_{\substack{s_{t+1:\infty},\\\\a_{t:\infty}}}\left[\nabla_\theta \log\pi_0(a_t\mid s_t)\right]=0
$$
因此
$$
\mathbb{E}_{\substack{s_{0:\infty},\\\\a_{0:\infty}}}\left[b_t(s_{0:t},a_{0:t-1})\nabla_\theta \log\pi_0(a_t\mid s_t)\right]=0
$$
所以有
$$
\mathbb{E}_{\substack{s_{0:\infty},\\\\a_{0:\infty}}}\left[\hat{A}_t(s_{0:\infty},a_{0:\infty})\nabla_\theta \log\pi_0(a_t\mid s_t)\right]=\mathbb{E}_{\substack{s_{0:\infty},\\\\a_{0:\infty}}}\left[A^{\pi, \gamma}(s_t,a_t)\nabla_\theta \log\pi_0(a_t\mid s_t)\right]
$$
也就是说 $\hat{A}$ 是 $\gamma\text{-just}$ 的.
容易验证, 下列估计表达式 $\hat{A}_t$ 也是 $\gamma\text{-just}$ 的:
$$
\begin{align}
&\bullet\sum\nolimits_{l=0}^\infty \gamma^l r_{t+l}&\qquad\qquad\qquad\qquad\bullet& A^{\pi,\gamma}(s_t,a_t)\\\\
&\bullet Q^{\pi,\gamma}(s_t,a_t)&\bullet& r_t+\gamma V^{\pi,\gamma}(s_{t+1})-V^{\pi,\gamma}(s_t)
\end{align}
$$
优势函数的估计
令 $\delta_t^V=r_t+\gamma V(s_{t+1})-V(s_t)$ 为 $V$ 的 $\gamma$ 折扣 $\text{TD}$ 残差. 注意到 $\delta_t^V$ 可以被看做是动作 $a_t$ 的优势. 事实上, 如果我们有确切的价值函数 $V=V^{\pi,\gamma}$ , 那么 $\delta_t^{V^{\pi,\gamma}}$ 会是一个无偏的, $\gamma\text{-just}$ 的优势估计.
$$
\begin{align}
\mathbb{E}_{s_{t+1}}\left[\delta_t^{V^{\pi,\gamma}}\right]&=\mathbb{E}_{s_{t+1}}\left[r_t+\gamma V^{\pi,\gamma}(s_{t+1})-V^{\pi,\gamma}(s_t)\right]\\\\
&=\mathbb{E}_{s_{t+1}}\left[Q^{\pi,\gamma}(s_t,a_t)-V^{\pi,\gamma}(s_t)\right]\\\\
&=A^{\pi,\gamma}(s_t,a_t)
\end{align}
$$
然而, 也只有当 $V=V^{\pi,\gamma}$ 时这个优势估计才是 $\gamma\text{-just}$ 的.
接着, 让我们考虑 $k$ 项这样的 $\delta$ 相加的情况, 我们将其记为 $\hat{A}_t^{(k)}$
$$
\begin{align}
&\hat{A}_t^{(1)}:=\delta_t^V&=&-V(s_t)+r_t+\gamma V(s_{t+1})\tag{11}\\\\
&\hat{A}_t^{(2)}:=\delta_t^V+\gamma\delta_{t+1}^V&=&-V(s_t)+r_t+\gamma r_{t+1}+\gamma^2 V(s_{t+2})\tag{12}\\\\
&\hat{A}_t^{(2)}:=\delta_t^V+\gamma\delta_{t+1}^V+\gamma^2\delta_{t+2}^V&=&-V(s_t)+r_t+\gamma r_{t+1}+\gamma^2 r_{t+2}+\gamma^3 V(s_{t+3})\tag{13}
\end{align}
$$
容易得出当且仅当 $V=V^{\pi,\gamma}$ 时 $\hat{A}_t^{(k)}$ 是 $\gamma\text{-just}$ 的. 注意到当 $k\to \infty$ 时, 由于 $\gamma^kV(s_{t+k})$ 项急剧减小 (当 $V\not= V^{\pi,\gamma}$ 时, 该项会产生偏差), 因此其偏差也会减少, 而 $-V(s_t)$ 项是不会产生偏差的 (因为这是基线). 令 $k\to \infty$ , 我们得到
$$
\hat{A}_t^{(\infty)}=\sum_{l=0}^{\infty}\gamma^l\delta_{t+l}^V=-V(s_t)+\sum_{l=0}^\infty\gamma^lr_{t+l}\tag{15}
$$
仅仅是回报减去价值函数基线.
广义优势估计 $(\text{generalized advantage estimator})$ $\mathrm{GAE}(\gamma,\lambda)$ 被定义为这些 $k$ 步估计量的指数加权平均:
$$
\begin{align}
\hat{A}_t^{\mathrm{GAE}(\gamma,\lambda)}&:=(1-\lambda)\left(\hat{A}_t^{(1)}+\lambda\hat{A}_t^{(2)}+\lambda^2\hat{A}_t^{(3)}+\cdots\right)\\\\
&=(1-\lambda)\left(\delta_t^V+\lambda(\delta_t^V+\gamma\delta_{t+1}^V)+\lambda^2(\delta_t^V+\gamma\delta_{t+1}^V+\gamma^2\delta_{t+2}^V)+\cdots\right)\\\\
&=(1-\lambda)\left(\delta_t^V(1+\lambda+\lambda^2+\cdots)+\gamma\delta_{t+1}^V(\lambda+\lambda^2+\lambda^3+\cdots)\\\\
+\gamma^2\delta_{t+2}^V(\lambda^2+\lambda^3+\lambda^4+\cdots)+\cdots\right)\\\\
&=(1-\lambda)\left(\delta_t^V\left(\frac{1}{1-\lambda}\right)+\gamma\delta_{t+1}^V\left(\frac{\lambda}{1-\lambda}\right)+\gamma^2\delta_{t+2}^V\left(\frac{\lambda^2}{1-\lambda}\right)+\cdots\right)\\\\
&=\sum_{l=0}^\infty(\gamma\lambda)^l\delta_{t+l}^V\tag{16}
\end{align}
$$
这个公式有两个特殊的情况即 $\lambda=0$ 和 $\lambda=1$ ,
$$
\begin{align}
\mathrm{GAE}(\gamma, 0):&\hat{A}_t:=\delta_t=r_t+\gamma V(s_{t+1})-V(s_t)\tag{17}\\\\
\mathrm{GAE}(\gamma,1):&\hat{A}_t:=\sum_{l=0}^\infty \gamma^l\delta_{t+l}=\sum_{l=0}^\infty
\gamma^lr_{t+l}-V(s_t)\tag{18}
\end{align}
$$
$\mathrm{GAE}(\gamma,1)$ 是 $\gamma\text{-just}$ 的, 不论 $V$ 的精度如何, 但由于其表达式中有多项和而导致其有高方差. 而 $\mathrm{GAE}(\gamma,0)$ 是 $\gamma\text{-just}$ 的当且仅当 $V=V^{\pi.\gamma}$ , 如果并非如此, 那么就会包含偏差, 但其往往具有更低的方差. 当 $0<\lambda<1$ 时我们得到更普遍的优势估计, 并且通过调整参数 $\lambda$ , 我们可以在偏差和方差之间取得一个平衡.
我们引入了一个包含两个参数 $\gamma$ 和 $\lambda$ 的优势估计, 这两个参数都会影响偏差和方差之间的平衡. 但是, 它们的作用以及对平衡的影响有本质上的不同. $\gamma$ 决定了价值函数 $V^{\pi,\gamma}$ 的最大值, 而 $\lambda$ 对其并没有影响. 并且无论价值函数是否是正确的, 只要 $\gamma<1$ 那么对于梯度 $g$ 来说就一定会有偏差 (因为已经使用了有折损回报, 而我们要优化的是无折损回报, 而 $\gamma$ 可以看做是影响方差与偏差的一个参数) . 而当 $\lambda<1$ 时只有错误的价值函数才会带来偏差. 因此我们往往会发现 $\lambda$ 的最优值往往比 $\gamma$ 的最优值要低得多, 这有可能是因为当获取到一个足够准确的价值函数后 $\lambda$ 引入的误差要比 $\gamma$ 小得多.
使用广义优势估计, 我们可以构建一个 $g^\gamma$ 的有偏估计
$$
g^\gamma\approx\mathbb{E}\left[\sum_{t=0}^{\infty}\nabla_\theta\log\pi_{\theta}(a_t\mid s_t)\hat{A}_t^{\mathrm{GAE(\gamma,\lambda)}}\right]=\mathbb{E}\left[\sum_{t=0}^{\infty}\nabla_\theta\log\pi_{\theta}(a_t\mid s_t)\sum_{l=0}^\infty (\gamma\lambda)^l\delta_{t+l}^V\right]\tag{19}
$$
当 $\lambda=1$ 时就是无偏的了.
对回报变形后的解释
令 $\Phi \colon\mathcal{S}\to \mathbb{R}$ 为在状态空间上的任意函数考虑一个变形后的回报 $\tilde{r}$
$$
\tilde{r}(s,a,s')=r(s,a,s')+\gamma\Phi(s')-\Phi(s)\tag{20}
$$
容易证明, 对于变形后的回报, 其有折损和为
$$
\sum_{l=0}^\infty \gamma^l \tilde{r}(s_{t+l},a_t,a_{t+l+1})=\sum_{l=0}^\infty \gamma^l r(s_{t+l},a_t,a_{t+l+1})-\Phi(s_t)\tag{21}
$$
类似的, 我们定义
$$
\begin{align}
&\tilde{Q}^{\pi,\gamma}(s,a)=Q^{\pi,\gamma}(s,a)-\Phi(s)\tag{22}\\\\
&\tilde{V}^{\pi,\gamma}(s)=V^{\pi,\gamma}(s)-\Phi(s)\tag{23}\\\\
&\tilde{A}^{\pi,\gamma}(s,a)=(Q^{\pi,\gamma}(s,a)-\Phi(s)\tag{24})-(V^{\pi,\gamma}(s)-\Phi(s))=A^{\pi,\gamma}(s,a)
\end{align}
$$
分别为变形后的回报, 价值, 优势函数.
注意到如果 $\Phi$ 如果恰好等于价值函数 $V^{\pi,\gamma}$ , 那么 $\tilde{V}^{\pi,\gamma}(s)$ 对于任何状态都等于零.
同样注意到使用变形后的回报对最大化有折损回报 $\sum_{t=0}^\infty \gamma^tr(s_t,a_t,s_{t+1})$ 的策略梯度是没有影响的. 但本文的目的其实是要最大化无折损回报, 而 $\gamma$ 是一个用来降低方差的参数.
引出回报变形后, 让我们考虑将其使用在梯度估计上. 注意到 $\tilde{r}$ 的形式与 $\delta^V$ 完全相同. 如果我们令 $\Phi=V$ , 那么就有
$$
\sum_{l=0}^\infty(\gamma\lambda)^l\tilde{r}(s_{t+l},a_t,s_{t+l+1})=\sum_{l=0}^\infty(\gamma\lambda)^l\delta_{t+l}^V=\hat{A}_t^{\mathrm{GAE}(\gamma,\lambda)}\tag{25}
$$
因此可以将 $\delta_{t}^V$ 理解为变形后的回报, 而 $\gamma\lambda$ 其实就是折损.
还有另一种解释, 记响应函数 $\chi$ 为
$$
\chi(l;s,a)=\mathbb{E}[r_{t+l}\mid s_t,a_t]-\mathbb{E}[r_{t+l}\mid s_t]\tag{26}
$$
注意到 $A^{\pi,\gamma}(s,a)=\sum_{l=0}^\infty \gamma^l\chi(l;s,a)$ . 然后重新回忆我们用折扣因子 $\gamma$ 和 $A^{\pi,\gamma}$ 来估计的策略梯度 (式 $(6)$) 中有如下形式
$$
\nabla_\theta \log\pi_0(a_t\mid s_t)A^{\pi, \gamma}(s_t,a_t)=\nabla_\theta \log\pi_0(a_t\mid s_t)\sum_{l=0}^\infty \gamma^l\chi(l;s,a)\tag{27}
$$
使用折扣 $\gamma<1$ 相当于急剧减少当 $l\gg 1/(1-\gamma)$ 的项 (因为 $\gamma^l$ 会变得很小) . 因此这些项带来的偏差会随着 $l$ 的增加而变少, 也就是说, 在大约 $1/(1-\gamma)$ 步后行动对回报的影响看起来像是被 “遗忘” 了.
如果我们使用变形后的回报 $\tilde{r}$ 并且有 $\Phi=V^{\pi,\gamma}$ , 不难得到当 $l>0$ 时有 $\mathbb{E}[\tilde{r}_{t+l}\mid s_t,a_t]=\mathbb{E}[\tilde{r}_{t+l}\mid s_t]=0$ (由于 $t$ 之后时刻采取的动作的概率完全依赖于那时候的状态 $s$ , 而 $t$ 时刻的动作已经给定为 $a_t$) . 因此只有 $l=0$ 是项才是非零的. 因此变形回报将依赖长时间的梯度估计变成了瞬时的. 当然我们往往难以得到非常好的估计, 但这已经能够给出式 $(16)$ 的一个很好的解释: 变形后的回报可以通过拟合 $V^{\pi,\gamma}$ 来减少响应函数起作用的范围, 通过引入一个更陡的折损 $\gamma\lambda$ 来减少由回报延迟 (即梯度估计中有很多项回报的和) 引起的噪声 (这是产生方差的原因), 也就是当 $l\gg 1/(1-\gamma\lambda)$ 时忽略项 $\nabla_\theta \log\pi_0(a_t\mid s_t)\delta^{V}_{t+l}$ .
后记
这篇文章可能可以算是论文 High-Dimensional Continuous Control Using Generalized Advantage Estimation 的一个翻译, 但是其中增添了一些我自己的理解并且有些东西换了一些表述. 这篇文章只有论文前半部分的内容, 后面则是讲如何拟合价值函数以及一些探究, 有兴趣的可以看原论文. 其实原论文后面这段解释我看着怪怪的… 虽然说不出什么问题吧… 但是总感觉没什么必要… 因为前面的介绍以及推导已经解释了 $\text{GAE}$ 是如何起作用的… 还有我本人英语非常菜, 借助了翻译才一点点看下来… 如果哪个地方理解不对请指出, 还有大佬们请轻喷.
