GAE 推导

3 minute read

Generalized Advantage Estimation（GAE）全推导

RL 里面的公式只能说是常学常新，最近看 PPO 的代码，其中最难理解的其实就是 GAE，纯看代码的话硬理顺其实也能理顺，但是你会完全不知道为什么要这样操作：

def gae(r, v, gamma=0.99, lam=0.95):  # r: [B, T], v: [B, T]
    B, T = r.shape
    adv = torch.zeros_like(r)
    last_gae = torch.zeros(B, device=r.device)
    last_v = torch.zeros(B, device=r.device)
    for t in reversed(range(T)):
        delta = r[:, t] + gamma * last_v - v[:, t]       # δ_t
        last_gae = delta + gamma * lam * last_gae        # A_t = δ_t + γλ A_{t+1}
        adv[:, t] = last_gae
        last_v = v[:, t]
    ret = adv + v                                        # R_t = A_t + V(s_t)
    adv = (adv - adv.mean()) / (adv.std() + 1e-8)
    return adv, ret

他里面实现的其实是一个递推公式：

\[A_t = \delta_t + \gamma\lambda A_{t+1}.\]

但是如果我们去翻原始论文 High-Dimensional Continuous Control Using Generalized Advantage Estimation，更熟悉的是以下这种形式：

\[\hat{A}_t^{\text{GAE}(\gamma,\lambda)} = \sum_{l=0}^{\infty} (\gamma\lambda)^l \,\delta_{t+l}, \qquad \delta_t = r_t + \gamma V^{\pi}(s_{t+1}) - V^{\pi}(s_t).\]

这个式子很漂亮，但要理解为什么 Advantage 能化成未来一串 $\delta$ 的几何加权和，需要一些推导上的直觉和耐心。

本文的目的就是把这两种形式进行完整的推导记录。

1. 从策略梯度和 Advantage 出发

给定策略梯度的公式（推导过程本身可以单独开一篇文章，这里暂且当作已知）：

\[\nabla_\theta J(\theta) = \mathbb{E}_{t \sim \pi_\theta}\Big[\,\nabla_\theta \log \pi_\theta(a_t|s_t)\, A^{\pi}(s_t, a_t)\,\Big].\]

其中 $\pi_\theta$ 是带参数 $\theta$ 的策略，$A^\pi(s_t,a_t)$ 是 Advantage function。

Advantage 的含义：

\[A^{\pi}(s_t,a_t) = Q^{\pi}(s_t,a_t) - V^{\pi}(s_t).\]

它刻画了在状态 $s_t$ 下，选择动作 $a_t$ 相比于“平均水平”（即该状态下策略的期望回报 $V^\pi(s_t)$）到底是好还是坏。

2. 一步 TD 残差：Advantage 的原子形式

Bellman 方程给出动作价值函数：

\[Q^{\pi}(s_t,a_t) = \mathbb{E}\big[\,r_t + \gamma V^{\pi}(s_{t+1}) \,\big|\, s_t,a_t\big].\]

于是：

\[A^\pi(s_t,a_t) = Q^{\pi}(s_t,a_t) - V^{\pi}(s_t) \;\approx\; r_t + \gamma V^{\pi}(s_{t+1}) - V^{\pi}(s_t).\]

定义 TD 残差：

\[\delta_t = r_t + \gamma V^{\pi}(s_{t+1}) - V^{\pi}(s_t).\]

$\delta_t$ 可以理解为“一步的意外”，即实际即时回报加上下一个状态的估值，与当前预测的差值。

3. k-step Advantage：为什么要它？

只用一步 $\delta_t$ 容易有偏差；完全展开（Monte Carlo return）方差又很大。为了在两者之间折中，引入 $k$-step return：

\[G_t^{(k)} = r_t + \gamma r_{t+1} + \cdots + \gamma^{k-1} r_{t+k-1} + \gamma^k V^{\pi}(s_{t+k}).\]

并定义：

\[A_t^{(k)} = G_t^{(k)} - V^{\pi}(s_t).\]

这里第一次出现 $G_t^{(k)}$，它表示“前 $k$ 步的实际奖励，加上第 $k$ 步的 bootstrap 估值”。而 $A_t^{(k)}$ 则是这个截断回报相较于当前预测值的差异。

问题驱动：为了最终能推出论文里的 $\sum (\gamma\lambda)^l \delta_{t+l}$ 形式，我们需要一个桥梁。直觉告诉我们，既然一步 Advantage 可以写成 $\delta_t$，那 $k$-step Advantage 是否也能写成若干 $\delta$ 的累积？如果能，就很容易和几何加权结合。下面严格展开。

4. 关键一步：把 $k$-step Advantage 写成 $\delta$ 累积

目标：

\[A_t^{(k)} \stackrel{?}{=} \sum_{l=0}^{k-1} \gamma^l \delta_{t+l}.\]

推导：

$k=1$：

\[A_t^{(1)} = r_t + \gamma V(s_{t+1}) - V(s_t) = \delta_t.\]

$k=2$：

\[\begin{aligned} A_t^{(2)} &= r_t + \gamma r_{t+1} + \gamma^2 V(s_{t+2}) - V(s_t) \\ &= (r_t + \gamma V(s_{t+1}) - V(s_t)) + \gamma(r_{t+1} + \gamma V(s_{t+2}) - V(s_{t+1})) \\ &= \delta_t + \gamma \delta_{t+1}. \end{aligned}\]

$k=3$：同理

\[A_t^{(3)} = \delta_t + \gamma \delta_{t+1} + \gamma^2 \delta_{t+2}.\]

一般情况（归纳）：

\[A_t^{(k+1)} = A_t^{(k)} + \gamma^k \delta_{t+k}.\]

因此归纳得证：

\[A_t^{(k)} = \sum_{l=0}^{k-1} \gamma^l \delta_{t+l}.\]

解释：$\delta_{t+l}$ 是“第 $t+l$ 步的意外”，折算到当前时刻需要再乘上 $\gamma^l$。所以 $k$-step Advantage 就是一段时间内所有意外的折扣和。

5. Generalized Advantage Estimation (GAE)

终于到 GAE 出场了。GAE 的想法是：不要固定某个 $k$，而是对所有 $k$-step Advantage 取一个几何加权平均：

\[A_t^{\text{GAE}(\gamma,\lambda)} = (1-\lambda) \sum_{k=1}^{\infty} \lambda^{k-1} A_t^{(k)}.\]

这里 $(1-\lambda)\lambda^{k-1}$ 是几何分布的权重。

微妙之处：

每个 $A_t^{(k)}$ 内部已经是对 $\delta$ 的 $\gamma$-加权和；
GAE 又在外层对这些 $A_t^{(k)}$ 做了 $\lambda$-几何加权。

也就是说，GAE 本质上是“双层加权”：内层的 $\gamma$ 折扣未来的 $\delta$，外层的 $\lambda$ 平滑不同长度的截断 Advantage。最终这两层权重会合并成一个优雅的 $(\gamma\lambda)^l$。

6. 化简：交换求和次序

代入 $A_t^{(k)} = \sum_{l=0}^{k-1}\gamma^l \delta_{t+l}$：

\[A_t^{\text{GAE}} = (1-\lambda)\sum_{k=1}^{\infty} \lambda^{k-1}\Big(\sum_{l=0}^{k-1}\gamma^l \delta_{t+l}\Big).\]

这是一个双重求和。关键是交换求和次序：

\[= \sum_{l=0}^\infty \gamma^l \delta_{t+l} \cdot (1-\lambda)\sum_{k=l+1}^\infty \lambda^{k-1}.\]

几何级数计算：

\[\sum_{k=l+1}^\infty \lambda^{k-1} = \lambda^l \cdot \frac{1}{1-\lambda}.\]

于是：

\[A_t^{\text{GAE}} = \sum_{l=0}^\infty (\gamma\lambda)^l \delta_{t+l}.\]

这就是论文里最常见的形式。注意到这里 $\gamma\lambda$ 看起来像放在了一起，但其实 $\gamma$ 在 $\delta$ 里面本身已经出现过一次，它既控制未来回报的折扣，也在最后和 $\lambda$ 一起出现在系数里。推导的优雅恰恰在于这种双重作用。

7. 递推形式：问题驱动的引入

现在我们有了一个无限和的形式。问题是：实际实现中我们不能真的求无限和，怎么计算？

答案就是把它写成递推。把第一项取出来：

\[A_t = \delta_t + \sum_{l=1}^\infty (\gamma\lambda)^l \delta_{t+l}.\]

令 $j = l-1$：

\[A_t = \delta_t + \gamma\lambda \sum_{j=0}^\infty (\gamma\lambda)^j \delta_{t+1+j}.\]

注意括号里的和正好是 $A_{t+1}$。于是得到递推：

\[A_t = \delta_t + \gamma\lambda A_{t+1}.\]

问题解决：我们不需要真的算无限和，只要反向循环一次就能算出所有的 $A_t$。

此处也正好对应代码里的实现方式，呼应开头的那段：

def gae(r, v, gamma=0.99, lam=0.95):  # r: [B, T], v: [B, T]
    B, T = r.shape
    adv = torch.zeros_like(r)
    last_gae = torch.zeros(B, device=r.device)
    last_v = torch.zeros(B, device=r.device)
    for t in reversed(range(T)):
        delta = r[:, t] + gamma * last_v - v[:, t]       # δ_t
        last_gae = delta + gamma * lam * last_gae        # A_t = δ_t + γλ A_{t+1}
        adv[:, t] = last_gae
        last_v = v[:, t]
    ret = adv + v                                        # R_t = A_t + V(s_t)
    adv = (adv - adv.mean()) / (adv.std() + 1e-8)
    return adv, ret

delta 对应 $\delta_t$；
last_gae 的更新正是递推公式；
ret = adv + v 得到的是回报目标 $\hat{R}_t = A_t + V(s_t)$。

8. 总结

本文详细推导了 GAE，其核心思路是首先得到 $\delta_t$ 的形式，然后推出来 k-step Advantage 的形式（是 $\delta$ 的 $\gamma$-加权和），然后对所有 k-step Advantage 进行几何加权平均，得到 GAE 的形式。

最后，我们得到了 GAE 的递推形式，并将其与代码实现进行了对应。

Share on

Twitter Facebook LinkedIn