Reinforcement Learning: Policy Gradient 与 Advantage Estimation

Reinforcement learning 中，一个 action 不只带来眼前的 reward，还会改变后续能遇到的 state。这篇从轨迹的期望 return 出发，依次推导 REINFORCE、reward-to-go、baseline、actor-critic 和 GAE，回答如何从采样轨迹构造策略的梯度。

系列导读 · 下一篇

Problem setup

先明确策略在优化什么，再区分几类算法如何表示和改进这个策略。

Definition

我们把 $⟨s_t,a_t⟩$ 打包起来，在 Markov 环境和只依赖当前 state 的策略下，其实它就构成了一个 Markov chain。记一条轨迹为 $\tau =(s_1,a_1,\dots ,s_T,a_T,s_{T+1})$：

这里 ${\pi }_{\theta }(a_t\mid s_t)$ 是策略，$\mathbb{P}$ 是环境转移；初始状态分布与环境转移不依赖 $\theta$。后文也把整条轨迹的分布 $p_{\theta }(\tau )$ 简记为 ${\pi }_{\theta }(\tau )$。

对于 learning objective：

我们定义我现在在 $s_t$，做了 action $a_t$，然后按照 $\pi$ 所得到的期望总 reward 为

然后我们定义 $\mathcal{V}$ 表示现在我在 $s_t$ 的时候遵循 $\pi$ 所得到的期望总 reward：

先考虑有限时域、不打折的 return；有限时域中把剩余时间也视为 state 的一部分，并令终止状态的 value 为 $0$。后面再引入 discount factor。

Types of Algorithms

• Policy Gradients：跟往常一样，就是求导然后去做优化。
• Value-based：去估计最优的 $\mathcal{Q}$ 或 $\mathcal{V}$，然后通过这个推出 $\pi$。
• Actor-critic：通过估计当前策略的 $\hat{\mathcal{Q}}$ 和 $\hat{\mathcal{V}}$，然后用它们指导策略优化。
• Model-based RL：学习或使用环境的转移、reward 模型，再借助模型规划或优化策略。

Policy Gradient Algorithms

Direct policy differentiation

记整条轨迹的总 reward 为 $r(\tau )=\sum _{t=1}^{T}r(s_t,a_t)$，我们想最大化它在当前策略下的期望：

采样得到一条 $\tau$ 后，很容易觉得：既然有它的概率 ${\pi }_{\theta }(\tau )$，又有它的 reward $r(\tau )$，那就直接最大化两者的乘积。这里漏掉了积分：${\pi }_{\theta }(\tau )r(\tau )$ 只是被积函数在这条轨迹上的取值，目标 $\mathcal{J}(\theta )$ 是对所有轨迹的积分。 如果能算出整个积分，直接对 $\mathcal{J}$ 求导当然可以；但只拿到一条采样轨迹时，直接对这个乘积求导，并不能给出整个目标梯度的无偏估计。

实际训练通常无法遍历所有轨迹，只能按当前策略采样。我们希望每条样本给出一个梯度，平均起来就能无偏地估计整个目标的梯度。下面推导的目的，就是把 ${\mathrm{\nabla }}_{\theta }\mathcal{J}(\theta )$ 写成一个可以用采样估计的期望，从而确定每条样本应该贡献什么梯度。

沿用环境与 reward 不依赖 $\theta$ 的假设，在积分与求导可以交换的条件下，利用 ${\mathrm{\nabla }}_{\theta }\log {\pi }_{\theta }(\tau )={\mathrm{\nabla }}_{\theta }{\pi }_{\theta }(\tau )/{\pi }_{\theta }(\tau )$，得到 REINFORCE (Williams, 1992) 的梯度表达式：

这就告诉我们，对 $\tau \sim {\pi }_{\theta }$，$r(\tau ){\mathrm{\nabla }}_{\theta }\log {\pi }_{\theta }(\tau )$ 是所需的无偏梯度估计。为了用反向传播得到它，取下面的 loss (Achiam, 2018)，求导时固定采到的轨迹和 reward：

负号用于把最大化写成梯度下降。使用 $\log {\pi }_{\theta }(\tau )$ 的理由就在这里：这个 loss 的负梯度，在采样所用的策略参数处，是目标梯度的无偏估计。

实现时，由于初始状态分布和环境转移与 $\theta$ 无关，有

因此，对一条轨迹，把实际采到的各步 action 的 log probability 相加，再乘总 reward 即可：

trajectory_log_prob = action_log_probs.sum()
loss = -trajectory_return.detach() * trajectory_log_prob

这里 trajectory_return 是总奖励，求导时作为常量；action_log_probs 保留策略参数的计算图。对一个 batch，再平均各条轨迹的 loss。这个估计虽然无偏，但 variance 很高，接下来考虑怎么减小它。

Reduce Variance

仔细观察这个式子，其实这玩意儿是 MLE 那个梯度对 $r(s_t,a_t)$ 加权了：

当前我们是有：${\mathrm{\Psi }}_t=\sum _{t^{\prime }=1}^{T}r(s_{t^{\prime }},a_{t^{\prime }})$

Don't Let the Past Distract You

一种简单的方法来减小 variance，是我们令 ${\mathrm{\Psi }}_t=\sum _{t^{\prime }=t}^{T}r(s_{t^{\prime }},a_{t^{\prime }})$

因为其实对于 $a_t$ 来说，他做啥对于 $t$ 之前的 reward 来说是不具有参考价值的。因此我们主要考虑后面的 reward。这玩意儿直觉上挺清楚的，但数学上想了半天才想明白为啥是对的。主要参考了这篇文章 (Achiam, 2018)。

证明的不造为啥让我想起了 MLE。主要用到的就是一个叫做 EGLP lemma 的东西（其实好像用这个 lemma 需要积分和导数的可交换性，貌似 (Hogg et al., 2013) 里写挺详细的）：

其实跟 MLE 是一样的嘛：

其实我们是要证明嘟是：

也就是要证明当 这个时候：

那么中心思想其实就是咋来区分 捏，我们考虑 是先 reward，再选择：

关键是先固定过去的历史，再对当前 action 取条件期望：给定 $s_t$，有 ${\mathbb{E}}_{a_t\sim {\pi }_{\theta }(\cdot \mid s_t)}[{\mathrm{\nabla }}_{\theta }\log {\pi }_{\theta }(a_t\mid s_t)]=0$。过去的 reward 可以和当前 state 相关，但不会由之后采样的 action 改写。

所以说最终结果是整个期望 $0$。

Introducing Baselines

另一个优化是我们考虑加入 baseline。这个直觉就更对了。就是我们考虑把 $r(s,a)$ 替换成 $r(s,a)-b$。这里减去的是梯度估计器中的 baseline，不是在改环境的 reward；baseline 不依赖当前 action 时，它乘上 score gradient 的期望为零。

不改变期望，不代表方差不变：baseline 的选择会改变梯度估计的波动。先看常数 baseline 的最优取值。

但其实这个 baseline 挺难算的，所以我们通常不会用这个最优的 baseline。而是去找一个相对比较好的。

Actor Critic Methods

前面的 reward-to-go 和 baseline 都在调整每个 action 的梯度权重。接下来用 value function 估计这个权重：先得到 advantage，再讨论 critic 的训练和 GAE。

General Idea

我们回到式子

我们观察这个 ${\mathrm{\Psi }}_t$（先不考虑 baseline）：${\mathrm{\Psi }}_t=\sum _{t^{\prime }=t}^{T}r(s_{t^{\prime }},a_{t^{\prime }})$。

我们发现一件事情，就是它其实是在估计 ${\mathcal{Q}}^{\pi }(s_t,a_t)$。也就是我做了这个 action，之后继续按照 $\pi$，会得到一个什么样的结果。这个和是 ${\mathcal{Q}}^{\pi }(s_t,a_t)$ 的一个无偏估计。

于是我们可以尝试 ${\mathrm{\Psi }}_t={\mathcal{Q}}^{\pi }(s_t,a_t)$。我的理解是，如果能算出这个条件期望，就可以去掉后续轨迹采样带来的部分噪声。

然后我们考虑 baseline。我们大概这么想，就是某个操作比我现在的平均效果好，那么 ${\mathrm{\Psi }}_t$ 要大于 $0$，让它的概率增大。如果比我现在的 $\pi$ 还要垃圾，就减小这个 action 的概率。

所以我们考虑 $b={\mathcal{V}}^{\pi }(s_t)$，于是我们定义：

也就是做了这个操作，相比当前策略的平均表现能好多少。然后

在采样策略就是 $\pi ={\pi }_{\theta }$、advantage 使用真值或条件无偏估计的情况下，这东西是无偏的。那我们考虑怎么算这个 ${\mathcal{A}}^{\pi }(s_t,a_t)$。

我们考虑

所以说

也就是说，$r(s_t,a_t)+{\mathcal{V}}^{\pi }(s_{t+1})-{\mathcal{V}}^{\pi }(s_t)$ 是对 ${\mathcal{A}}^{\pi }(s_t,a_t)$ 的一个条件无偏估计。这里需要真实的 ${\mathcal{V}}^{\pi }$；换成学出来的近似值后，一般会引入 bias。

也就是我们现在的问题来到了怎么搞这个 ${\mathcal{V}}^{\pi }$，最直接的方法是大力去做，然后求平均。

当然，更成熟的想法是我们可以训一个模型 $\varphi$ 去预测这个 ${\mathcal{V}}^{\pi }$。我们考虑 loss，一种直接的方法是：

但是我们考虑有没有方差更低一些的方法，就是用 $r(s_t,a_t)+{\hat{\mathcal{V}}}_{\varphi }^{\pi }(s_{t+1})$ 来代替整个求和。这就变成了 TD error 的平方：

这里 $\mathrm{sg}$ 表示 stop-gradient：当前更新把 bootstrap target 当作固定目标。它降低了对完整 Monte Carlo return 的依赖，但 target 本身的估计误差也会传过来。

Introducing Discount Factors

有时候我们会考虑 $T\to \mathrm{\infty }$ 的情况。为了在 reward 有界时让总 return 也是有穷的，我们可以取 ，增加一个 discount factor。这个其实是比较直觉的，因为现在给你一块钱和以后给你一块钱，肯定选马上要。

那么我们就需要稍稍改一改式子。此时 ${\mathcal{V}}^{\pi ,\gamma }$ 从当前时刻开始计算折扣 return：

这时候其实我们在估计导数的时候，要小心折扣的位置。直接用整条轨迹的 return，可以写成：

利用 reward-to-go，也可以使用下面的估计量：

外层的 ${\gamma }^{t-1}$ 与 return 内的 ${\gamma }^{t^{\prime }-t}$ 作用不同，不能直接漏掉。两种估计量的期望相同，同一批轨迹上算出的数值不必相同。Thomas (Thomas, 2014) 专门讨论了 discount 与 policy gradient bias 的问题，有空填坑。

Implementation Details

在真正写代码的时候，我们需要构造一个能产生 policy gradient 的 surrogate。和前面 REINFORCE 的 loss 一样，考虑一个看似无意义的函数：

求导时，采到的 state、action 和 ${\hat{\mathcal{A}}}_{\varphi }$ 都作为常量，所以这个 ${\mathrm{\nabla }}_{\theta }\tilde{\mathcal{J}}(\theta )$ 就是相应的 actor gradient 估计。也就是说，我们其实是借助自动求导程序去计算前面推出来的更新。它的数值不是 $\mathcal{J}(\theta )$；critic 若有误差，也不能直接声称这个梯度无偏。对于前面从初始状态定义的 discounted objective，还要给第 $t$ 项乘上 ${\gamma }^{t-1}$。

也就是说我们要训练两个网络 $\theta$ 和 $\varphi$。用当前策略采到的数据集 ${\mathcal{D}}_k$，一轮更新里的两个目标可以写成：

两者都用梯度下降；actor loss 前面的负号对应最大化 surrogate。这里沿用不额外写外层时间权重的实现记法，具体采样与折扣约定仍要和目标一致。

Generalized Advantage Estimation

Actor-critic 用 advantage 判断一个 action 比当前策略的平均表现好多少。实际训练时，真实的 $\mathcal{Q}$ 和 $\mathcal{V}$ 都未知，我们需要从采到的轨迹和 critic 的预测中构造 ${\hat{\mathcal{A}}}_t$。

我们需要决定观察多少步真实 reward，再让 critic 预测剩下的收益。只观察一步，估计很依赖 critic；一直观察到终点，又会混入后续 action 和环境带来的随机性。Schulman 等人提出的 Generalized Advantage Estimation（GAE） (Schulman et al., 2015) 把不同步数的估计结合起来，用 $\lambda$ 调节两者的取舍。

From one step to multiple steps

固定采样策略 $\pi$ 和 discount factor $\gamma$，记 $r_t=r(s_t,a_t)$，把 critic 对 ${\mathcal{V}}^{\pi ,\gamma }(s_t)$ 的预测简记为 $v_t={\hat{\mathcal{V}}}_{\varphi }(s_t)$。在同一条轨迹上，我们可以在不同位置交给 critic 接手：

观察 $n$ 步后，前半段用实际 reward，后半段用 $v_{t+n}$ 补足。这种用估计值预测尚未观察到的收益的做法叫做 bootstrap。再减去当前状态的 baseline $v_t$，得到 $n$-step advantage estimate：

$n=1$ 时，它就是 TD residual：

增加 $n$，就是用更多实际发生的 reward 替换 critic 对未来的预测。这样通常能减少 bootstrap 误差的影响，但也纳入了更多后续轨迹的随机性。

Mixing the horizons

GAE 为这些不同长度的估计分配指数衰减的权重。先考虑 ：

这些权重之和为 $1$。较小的 $\lambda$ 把权重集中在短步数估计上；较大的 $\lambda$ 让长步数估计参与得更多。对于已经终止的 episode，可以把终止后的 reward 和 value 都延拓为 $0$。

把各个 $n$-step estimate 展开并合并，GAE 就变成了 TD residual 的折扣和：

前一个式子说明 GAE 如何结合不同步数的估计，后一个式子更方便计算。$\lambda =1$ 时使用后一个表达式，或取前一个表达式在 $\lambda \to 1$ 时的极限。

What λ controls

Bootstrap 可能把 critic 的预测误差带入 policy gradient。若 critic 恰好等于真实的 ${\mathcal{V}}^{\pi ,\gamma }$，一步 TD 就已经是 advantage 的条件无偏估计；在同样的采样与边界条件下，其他 $\lambda$ 也不会额外引入这种 bias。

$\gamma$ 和 $\lambda$ 的作用不同：$\gamma$ 定义本节要估计的 discounted return；固定 $\gamma$ 后，$\lambda$ 调节估计这个量时观察多长的轨迹。改变 $\lambda$ 不会改变真实 advantage 的定义。

Computing GAE on a rollout

实际收集到的轨迹是有限的。设最后一个 action 是 $a_T$，下一状态为 $s_{T+1}$。从后往前计算即可：

这里的 ${\hat{\mathcal{A}}}_{T+1}=0$ 表示没有更多已采样的 TD residual；未来收益仍通过最后一个 ${\delta }_T$ 中的 $v_{T+1}$ 进入估计。边界 value 按轨迹停止的原因确定：

• 真正终止：未来没有 reward，令 $v_{T+1}=0$。
• rollout 收集结束或外部 time limit 截断：任务本可继续，使用停止前最后一个状态的 critic 预测作为 $v_{T+1}$。

Spinning Up 的 PPOBuffer 实现 就通过最后一个 value 处理这两种情况。Gymnasium (Farama Foundation, n.d.) 用 terminated 和 truncated 区分真正终止与外部截断；任务本身定义的有限时域终点属于前者。

有限 rollout 下，即使 $\lambda =1$，仍有

采到真正终止时，边界项为零，才得到纯 Monte Carlo return 减 baseline；截断时则保留 bootstrap。

下面用普通列表实现单段 rollout 的计算。rewards 长度为 $T$，values 长度为 $T+1$，最后一个元素已经按上面的规则设置。它们都是采样时记录并固定下来的数值。

def gae(rewards, values, gamma, lam):
    assert len(values) == len(rewards) + 1
    advantages = [0.0] * len(rewards)
    carry = 0.0
    for t in reversed(range(len(rewards))):
        delta = rewards[t] + gamma * values[t + 1] - values[t]
        carry = delta + gamma * lam * carry
        advantages[t] = carry
    return advantages

多条 episode 应分别计算；截断后 reset 得到的新初始状态也不属于上一段轨迹，不能把它的 value 或 advantage 接到上一段。

例如，一条三步后终止的轨迹只有最后一步得到 reward。取 $\gamma =0.9$、$\lambda =0.8$，因此 $\gamma \lambda =0.72$，并设终止状态 $v_4=0$：

从最后一步开始，先得到 ${\hat{\mathcal{A}}}_3=0.4$，再算 ${\hat{\mathcal{A}}}_2=0.04+0.72\times 0.4=0.328$，最后得到 ${\hat{\mathcal{A}}}_1=0.05+0.72\times 0.328=0.28616$。终点的信息通过这次反向递推传到更早的 action。

训练 actor 时，把算好的 advantage 当作固定权重，代入前面的 policy-gradient loss 或 PPO objective。GAE 负责构造这个权重；下一篇的 TRPO 和 PPO 则控制策略如何利用它更新。

系列导读 · 下一篇

References