LLM后训练(二)--价值函数

本文是LLM后训练的第二篇，主要介绍强化学习中价值函数以及Advantage Actor-Critic模型(A2C)。

奖励信号的构造方式

在LLM后训练（一）中，我们介绍了强化学习的基本概念和一些常用的算法，并介绍了最原始的奖励损失函数：

$$J(\theta) \approx \frac{1}{N} \sum_{i=1}^{N} \left( R(\tau^i) \sum_{t=0}^{T} \log \pi_\theta(a_t^i|s_t^i) \right)$$

这个函数对整条轨迹使用统一的奖励，因此存在一些问题。那么，有没有对每一步动作使用不同奖励的算法呢？假设这种奖励存在，那么奖励信号就可以写入第二个求和符号中：

$$J(\theta) \approx \frac{1}{N} \sum_{i=1}^{N} \left( \sum_{t=0}^{T} \varPsi_t^i \log \pi_\theta(a_t^i|s_t^i) \right)$$

为了方便表示，假设只针对一条采样轨迹进行分析，忽略下标$i$，将$\varPsi_t^i$简写为$\varPsi_t$。在不同算法中，$\varPsi_t$的定义不同。在介绍A2C算法之前，先总结一下$\varPsi_t$的几种常见构造方式：

$\varPsi_t = \sum_{t=0}^{T}{r_t}$

这种方式等价于最原始的奖励损失函数。

$\varPsi_t =\sum_{t'=t}^{T}{r_{t'}}$

这种方式考虑了从$t$时刻开始的奖励。

$\varPsi_t =\sum_{t'=t}^{T}{r_{t'}} - b(s_t)$

这种方式引入了基线函数$b(s_t)$，用来降低奖励的方差。因为如果只考虑绝对奖励，可能会出现每个奖励都很大或者都很小的情况，从而导致梯度更新不稳定。引入基线后，奖励就变成了某个动作相对于平均水平的“优势”，而训练模型的最终目标正是在候选动作中选择更优动作。$b(s_t)$的定义有很多，这里暂时将其视为采取所有动作时的平均奖励。

$\varPsi_t = Q(s_t, a_t)$

这种方式使用动作价值函数作为奖励信号。
它表示在状态$s_t$和该步对应动作$a_t$的条件下，针对策略$\pi$未来奖励的期望值。由于这里使用的是未来奖励的期望值而非一个具体奖励的采样值，因此它可以看作是一个更稳定的奖励信号。
需要注意的是，动作价值函数中的条件期望是针对特定策略$\pi$定义的。在训练过程中，策略可能会随着参数更新而改变（也可能使用固定采样策略，此时不需要考虑策略改变的问题）。

$\varPsi_t = A(s_t, a_t)$

这种方式使用优势函数作为奖励信号。优势函数的定义为：

$$A(s_t, a_t) = Q(s_t, a_t) - V(s_t)$$

它表示在状态$s_t$和动作$a_t$条件下，动作价值函数与状态价值函数的差值。状态价值函数的定义为

$$V(s_t) = \mathbb{E}_{\tau \sim \pi} \left[ \sum_{t'=t}^{T} \gamma^{t'-t} r_{t'} \mid s_t \right]$$

它表示在状态$s_t$的条件下，针对策略$\pi$未来奖励的期望值。最终，优势函数表示在状态$s_t$和动作$a_t$的条件下，动作相对于平均水平的“优势”，它可以看作是一个更稳定的奖励信号。

$\varPsi_t = \delta_t$

这种方式使用TD误差作为奖励信号。TD误差的定义为

$$\delta_t = r_t + \gamma V(s_{t+1}) - V(s_t)$$

TD误差表示奖励观测值与未来奖励折现值之和，再减去当前状态价值函数值的差值。（它的本质是贝尔曼方程，即当前价值等于即时奖励加上下一个状态折现价值的期望。）

动作价值函数、状态价值函数和TD误差

在上一部分，我们介绍了几种奖励信号的构造方式。本部分着重分析状态价值函数、动作价值函数和TD误差的性质与联系。

动作价值函数的性质

根据动作价值函数和状态价值函数的定义，可以推导：

$$\begin{align} Q(s_t, a_t) &= \mathbb{E}_{\tau \sim \pi} \left[ \sum_{t'=t}^{T} \gamma^{t'-t} r_{t'} \mid s_t, a_t \right] \nonumber \\ &= \mathbb{E}_{\tau \sim \pi} \left[ r_t + \gamma \sum_{t'=t+1}^{T} \gamma^{t'-(t+1)} r_{t'} \mid s_t, a_t \right] \nonumber \\ &= \mathbb{E} \left[r_t\mid s_t, a_t \right] + \mathbb{E}_{\tau \sim \pi} \left[ \gamma\sum_{t'=t+1}^{T} \gamma^{t'-(t+1)} r_{t'} \mid s_t, a_t \right] \nonumber \\ &= \mathbb{E} \left[r_t\mid s_t, a_t \right] + \mathbb{E}_{s_{t+1} \sim P(\cdot|s_t,a_t)} \left[ \gamma\mathbb{E}_{\tau \sim \pi} \left[ \sum_{t'=t+1}^{T} \gamma^{t'-(t+1)} r_{t'} \mid s_{t+1} \right]\right] \nonumber \\ &= \mathbb{E} \left[r_t\mid s_t, a_t \right] + \mathbb{E}_{s_{t+1} \sim P(\cdot|s_t,a_t)} \left[\gamma V(s_{t+1}) \right] \nonumber \\ &= \mathbb{E}_{s_{t+1} \sim P(\cdot|s_t,a_t)} \left[ r_t + \gamma V(s_{t+1}) \right] \nonumber \end{align}$$

上式给出了状态价值函数到动作价值函数的一种变换关系。其中隐含了从$t$到$t+1$的递推关系，这种递推性质在后文会反复用到。

状态价值函数的性质

根据状态价值函数的定义，可以推导：

$$\begin{align} V(s_t) &= \mathbb{E}_{\tau \sim \pi} \left[ \sum_{t'=t}^{T} \gamma^{t'-t} r_{t'} \mid s_t \right] \nonumber \\ &= \mathbb{E}_{a_t \sim \pi(\cdot|s_t)} \left[ \mathbb{E}_{\tau \sim \pi} \left[ \sum_{t'=t}^{T} \gamma^{t'-t} r_{t'} \mid s_t, a_t\right] \right] \nonumber \\ &= \mathbb{E}_{a_t \sim \pi(\cdot|s_t)} \left[ Q(s_t, a_t) \right] \nonumber \end{align}$$

上式给出了动作价值函数到状态价值函数的一种变换关系。

TD误差

根据优势函数的定义，我们可以推导：

$$\begin{align} A(s_t, a_t) &= Q(s_t, a_t) - V(s_t) \nonumber \\ &= \mathbb{E}_{s_{t+1} \sim P(\cdot|s_t,a_t)} \left[ r_t + \gamma V(s_{t+1}) \right] - V(s_t) \nonumber \\ &= \mathbb{E}_{s_{t+1} \sim P(\cdot|s_t,a_t)} \left[ r_t + \gamma V(s_{t+1}) - V(s_t) \right] \nonumber \\ &= \mathbb{E}_{s_{t+1} \sim P(\cdot|s_t,a_t)} \left[ \delta_t \right] \nonumber \end{align}$$

TD误差取期望后等于优势函数，因此在$V_{\pi}$准确的前提下，TD误差可以看作优势函数的无偏估计。然而在实际训练过程中，$V_{\pi}$通常是近似函数，甚至本身也是需要训练的模型，因此TD误差可能存在一定偏差。另外，这里的TD误差只使用一步观测值。实际上也可以使用多步观测值来构造TD误差，这样可以在一定程度上减小偏差，但同时会增大方差。这里默认只使用一步观测值的TD误差，也称为TD(0)误差。

Advantage Actor-Critic模型

上一部分介绍了几种奖励信号的构造方式，并分析了动作价值函数、状态价值函数和TD误差的性质与联系。本部分介绍Actor-Critic模型。相比于原版Actor-Critic模型使用动作价值函数$Q$作为奖励信号，这里使用TD误差（本质上是优势函数）作为奖励信号的模型被称为Advantage Actor-Critic模型（A2C）。如果通过多次采样实际样本，再统计动作价值函数和状态价值函数以得到优势函数，往往会导致优势函数方差较大。而TD误差只使用一步观测值$r_t$以及当前和下一步状态价值函数的估计值$V_{\phi}(s_t)$和$V_{\phi}(s_{t+1})$，因此其方差相对较小，可以加速训练过程。然而，由于TD误差往往是有偏估计（因为估计值未必准确），A2C模型引入了一个额外的状态价值函数模型，即Critic模型，来估计状态价值函数，并与Actor模型（即策略模型）交替优化缓解估计不准的问题。

Actor优化目标

Actor模型的优化目标是最大化下式：

$$J(\theta) \approx \frac{1}{N} \sum_{i=1}^{N} \left( \sum_{t=0}^{T} (r_t^i + \gamma V_{\phi}(s_{t+1}^i) - V_{\phi}(s_t^i)) \log \pi_\theta(a_t^i|s_t^i) \right)$$

在实际优化过程中，可以对$T$中的时间步进行各种采样（包括单步采样），而不需要一次性优化整条轨迹的奖励信号，这样可以加速训练过程。

Critic优化目标

Critic模型的优化目标是最小化下式：

$$\mathcal{L}(\phi) = \mathbb{E}\left[ (r_t + \gamma V_{\phi}(s_{t+1}) - V_{\phi}(s_t))^2 \right] \nonumber$$

其目标是最小化TD误差平方，因此可以看作一个回归问题，目标是让$V_{\phi}$能够更准确地估计状态价值函数。

下图展示了A2C模型的简化优化流程（使用梯度上升最大化目标）：