PPO&GRPO
详解PPO算法和GRPO算法。
强化学习的概念
- State: 状态，表示当前的局面，如果是LLM那就是一个句子（可以是不完整的句子），用\(s\)表示。
- Action: 动作，在当前局面下作出的反应，如果是LLM那就是生成一个token，用\(a\)表示。
- Action Space: 动作空间，可选择的动作，在LLM中就是可以输出的所有token。
- Policy: 策略函数，输入State，输出Action的概率分布，用\(\pi\)表示。
- Trajectory: 轨迹，一连串状态和动作的序列。表示为\(\tau = s_0, a_0, s_1, a_1, ....\)。
- Reward: 当前时间点的奖励。
- Return: 回报，从当前时间点到结束的Reward的累计和。
期望计算方式：
- 每个可能结果的概率与其结果值的乘积之和。
- 采样多次取平均值。（约等于，采样次数趋于无穷时取等号）
\[E(x)_{x \sim p(x)}=\sum_{x}x * p(x) \approx \frac{1}{n}\sum_{i=1}^{n}{x_i}\]
强化学习的目标：训练一个Policy（神经网络\(\pi\)），在所有的状态s下，给出相应的Action，使得到的Return期望最大。
PPO
将上面强化学习的目标变为公式，一步一步计算。
Return的期望表示为：
\[E(R(\tau))_{\tau \sim P_{\theta}(\tau)}=\sum_{\tau}R(\tau)P_{\theta}(\tau)\]
因为要对上面的期望求最大值，所以对上式求梯度：
\[ \begin{align} \nabla E(R(\tau))_{\tau \sim P_{\theta}(\tau)}&=\nabla \sum_{\tau}R(\tau)P_{\theta}(\tau) \\ &=\sum_{\tau}R(\tau) \nabla P_{\theta}(\tau) \\ &=\sum_{\tau}R(\tau) \nabla P_{\theta}(\tau) \frac{P_{\theta}(\tau)}{P_{\theta}(\tau)} \\ &=\sum_{\tau}P_{\theta}(\tau)R(\tau) \frac{\nabla P_{\theta}(\tau)}{P_{\theta}(\tau)} \end{align} \]
这里将\(R(\tau) \frac{\nabla P_{\theta}(\tau)}{P_{\theta}(\tau)}\)看作一个整体，则概率\(P_{\theta}(\tau)\)和这个整体的乘积可表示为期望，用期望计算方式二来变换上式得到：
数学技巧： \[\nabla logf(x)=\frac{\nabla f(x)}{f(x)}\]
\[ \begin{align} \nabla E(R(\tau))_{\tau \sim P_{\theta}(\tau)}&=\sum_{\tau}P_{\theta}(\tau)R(\tau) \frac{\nabla P_{\theta}(\tau)}{P_{\theta}(\tau)} \\ & \approx \frac{1}{N}\sum_{n=1}^{N}R(\tau^n) \frac{\nabla P_{\theta}(\tau^n)}{P_{\theta}(\tau^n)} \\ &=\frac{1}{N}\sum_{n=1}^{N}R(\tau^n) \nabla logP_{\theta}(\tau^n) \end{align} \]
更进一步，我们将上式中的Trajectory拆分细化：
\[ \begin{align} \nabla E(R(\tau))_{\tau \sim P_{\theta}(\tau)}&=\frac{1}{N}\sum_{n=1}^{N}R(\tau^n) \nabla logP_{\theta}(\tau^n) \\ &=\frac{1}{N}\sum_{n=1}^{N}R(\tau^n) \nabla log \prod_{t=1}^{T_n}P_{\theta}(a_n^t|s_n^t) \\ &=\frac{1}{N}\sum_{n=1}^{N}R(\tau^n) \sum_{t=1}^{T_n} \nabla logP_{\theta}(a_n^t|s_n^t) \\ &=\frac{1}{N}\sum_{n=1}^{N}\sum_{t=1}^{T_n} R(\tau^n) \nabla logP_{\theta}(a_n^t|s_n^t) \end{align} \]
这样，我们需要最大化的期望就可以写作最小化的Loss函数形式，如下所示：
\[ \begin{align} Loss=-\frac{1}{N}\sum_{n=1}^{N}\sum_{t=1}^{T_n} R(\tau^n)logP_{\theta}(a_n^t|s_n^t) \end{align} \]
我们继续对上式中的\(R(\tau^n)\)进行拆解。
最寻常的Return就是从当前时间点到结束的Reward的累计和，如下所示：
\[R(\tau^n) = \sum_{t=1}^{T_n}r_t^n\]
而PPO算法中引入了两点改进：
- 当前时刻所作出的动作应该只会影响之后的reward，和之前步骤的reward应该无关。---》\(R(\tau^n)\)之计算当前时刻之后的累积。
- 一个动作对接下来的几步影响是大的，但对接下来很远步骤的影响力应该减小。---》\(R(\tau^n)\)引入衰减因子\(\gamma\)，越远影响越小。
则上式优化为：
\[R(\tau^n) = \sum_{t^{\prime}=t}^{T_n} \gamma^{t^{\prime}-t}r_{t^{\prime}}^n=R_t^n\]
这样，我们的Loss函数可以写成：
\[ \begin{align} Loss=-\frac{1}{N}\sum_{n=1}^{N}\sum_{t=1}^{T_n} R_t^nlogP_{\theta}(a_n^t|s_n^t) \end{align} \]
到这里，我们就可以根据上面的Loss函数训练我们的强化学习模型了。其中\(T_n\)可视做LLM中句子的长度，\(N\)可视做batch的大小。
但是，仅通过上述Loss函数做优化，还存在许多缺点，我们继续一步一步添加PPO算法中的优化对上述Loss函数进行改进。
奖励减去baseline
在好的局势下，无论做什么动作都能得到正的Reward；而在坏的局势下，无论做什么动作都会得到负的Reward。如下图所示：
这是不合适的，我们需要的是当前这个动作相对于其他动作的优势，即只要这个动作相对于其他的动作好的话，则给他正向的奖励，反之给负向的奖励。所以这就需要我们对奖励减去一个baseline。
\[ \begin{align} Loss=-\frac{1}{N}\sum_{n=1}^{N}\sum_{t=1}^{T_n} (R_t^n - B(s_n^t))logP_{\theta}(a_n^t|s_n^t) \end{align} \]
我们把\((R_t^n -
B(s_n^t))\)称为一个动作对其他动作的优势。
这样，无论在好的局势下还是坏的局势下，我们得到的奖励都是有正有负的，更利于强化学习的训练，如下图所示：
引入优势函数
计算Return的式子：
\[ \begin{align} R_t^n - B(s_n^t) \\ R_t^n = \sum_{t^{\prime}=t}^{T_n} \gamma^{t^{\prime}-t}r_{t^{\prime}}^n \end{align} \]
可以看到，\(R_t^n\)是一次随机采样，方差很大，只有采样次数足够多的时候才会比较准确。那么是不是可以通过一个神经网络来估计\(R_t^n - B(s_n^t)\)这个优势呢？PPO算法在这里就引入了优势函数。
几个概念：
- Action-Value Function（动作价值函数）： \(Q_{\theta}(s, a)\)，在状态s下，做出动作a的期望回报。
- State-Value Function（状态价值函数）：\(V_{\theta}(s)\)，在状态s下的期望回报。
- Advantage Function（优势函数）：\(A_{\theta}(s, a)=Q_{\theta}(s, a)-V_{\theta}(s)\)，在状态s下，做出动作a，相比于其他动作能带来多少优势。
这样，Loss函数就可以更新为：
\[ \begin{align} Loss=-\frac{1}{N}\sum_{n=1}^{N}\sum_{t=1}^{T_n} A_{\theta}(s_n^t, a_n^t)logP_{\theta}(a_n^t|s_n^t) \end{align} \]
下面进一步拆解优势函数\(A_{\theta}(s, a)=Q_{\theta}(s, a)-V_{\theta}(s)\)。
首先\(Q_{\theta}(s_t, a)\)可以理解为当前状态\(s_t\)下，作出动作\(a\)的期望回报。同时，它也等效于当前状态\(s_t\)下，作出动作\(a\)的奖励加上下一个状态\(s_{t+1}\)的期望回报，写作下式：
\[Q_{\theta}(s_t, a) = r_t + \gamma * V_{\theta}(s_{t+1})\]
然后，把\(Q_{\theta}(s_t, a)\)代入优势函数可以得到：
\[A_{\theta}(s_t, a)=r_t + \gamma * V_{\theta}(s_{t+1})-V_{\theta}(s_t)\]
这样我们就可以通过一个预测状态价值\(V_{\theta}(s)\)的神经网络去计算得到优势函数，其中\(V_{\theta}(s)\)可以和要训练的强化学习网络结构相同，如下图所示：
GAE
对于状态价值\(V_{\theta}(s_t)\)的计算，我们可以直接在当前时间直接通过神经网络预测，即直接得到:
\[V_{\theta}(s_t)\]
我们也可以先现在当前时间随机采样一个动作\(a\)，得到奖励\(r_t\)，然后再通过神经网络预测下一个状态的价值，即得到:
\[V_{\theta}(s_t) \approx r_t + \gamma * V_{\theta}(s_{t+1})\]
更进一步，我们可以采样两步，然后再通过神经网络预测，即得到
\[V_{\theta}(s_t) \approx r_t + \gamma * r_{t+1} + \gamma^2 * V_{\theta}(s_{t+2})\]
以此类推，我们可以向后一直采样，得到采样多个步骤的优势函数如下：
\[ \begin{align} A^1_{\theta}(s_t, a) &= r_t + \gamma * V_{\theta}(s_{t+1}) - V_{\theta}(s_t) \\ A^2_{\theta}(s_t, a) &= r_t + \gamma * r_{t+1} + \gamma^2 * V_{\theta}(s_{t+2}) - V_{\theta}(s_t) \\ A^3_{\theta}(s_t, a) &= ... \end{align} \]
定义中间变量：
\[ \begin{align} \delta_t^V &= r_t + \gamma * V_{\theta}(s_{t+1}) - V_{\theta}(s_t) \\ \delta_{t+1}^V &= r_{t+1} + \gamma * V_{\theta}(s_{t+2}) - V_{\theta}(s_{t+1}) \\ \delta_{t+2}^V &= ... \end{align} \]
采样多个步骤的优势函数可简化为：
\[ \begin{align} A^1_{\theta}(s_t, a) &= \delta_t^V \\ A^2_{\theta}(s_t, a) &= \delta_t^V + \gamma \delta_{t+1}^V \\ A^3_{\theta}(s_t, a) &= ... \end{align} \]
那么在优势函数计算的时候，应该使用采样了几步的优势函数呢，GAE的做法是全都要：
\[ \begin{align} A_{\theta}^{GAE}&=(1 - \lambda)(A_{\theta}^1 + \lambda A_{\theta}^2 + \lambda^2A_{\theta}^3) \\ &=\sum_{b=0}^{\infty}(\gamma \lambda)^b\delta_{t+b}^V \end{align} \]
这样，经过PPO算法中GAE步骤的改进，我们的损失函数变为如下形式：
\[ \begin{align} \delta_t^V &= r_t + \gamma * V_{\theta}(s_{t+1}) - V_{\theta}(s_t) \\ A_{\theta}^{GAE}(s_t, a)&=\sum_{b=0}^{\infty}(\gamma \lambda)^b\delta_{t+b}^V \\ Loss&=-\frac{1}{N}\sum_{n=1}^{N}\sum_{t=1}^{T_n} A_{\theta}^{GAE}(s_n^t, a_n^t)logP_{\theta}(a_n^t|s_n^t) \end{align} \]