从传统 RL 算法应用到 LLM 定义,推导 RL 算法策略梯度及改进

  • 发布于 2026-01-17
  • 4 次阅读

作者:车中草同学

https://zhuanlan.zhihu.com/p/1893791455769776556*

本篇博客的目的:

陆陆续续学习RL一阵子,看过ppo,grpo,gspo等的论文,但是一直没有建立一个大的领域框架,来兼容最新的改进。

因此本篇博客从RL应用到LLM的初步定义开始,然后推导传统的RL算法策略梯度在LLM领域是如何应用的。然后以此为基础,后续的所有RL算法(PPO,GRPO,REINFORCE++等)改进都是在此基础上来解决策略梯度存在的两个问题(1.信用分配粒度过粗2.绝对优势问题(没有减去baseline导致方差大) 。

写这篇文章我的最大收获是:建立了一个大的框架,来看待各种最新的RL算法改进,希望大家看完也有所收获!

另外,欢迎大家反馈文章中的笔误和概念问题。

前言

在前面我们写过两篇文章,其主要是先入为主的介绍了一些算法(PPO,GRPO,GSPO)的动机和理论。没有涉及其来源,不知道来源会导致无法融会贯通,没有一个大的picture来认识这些算法和后续改进。

[通俗易懂]RLHF第一篇-从PPO->GRPO->GSPO(动机、理论、verl代码、分析)

[通俗易懂]RLHF第二篇-RL中的CLIP机制和重要性采样(动机,原理,最新改进介绍(解决墒崩塌、训推不一致))

因此这篇文章会更细致地介绍传统的RL算法(策略梯度)是如何应用到LLM中的

1.首先明确RL应用到LLM的定义

2.其次在LLM场景下的定义:最大化回报开始推导,直到推导到策略梯度的token级别的loss(surrogate objective)。其实这是PPO, GRPO, REINFORCE++等系列算法的基础,这些最新的算法都是在这一步的基础上做优化

在上述推导的基础上,接下来分析策略梯度存在的两个问题。

  • 绝对优势问题(没有减去baseline导致方差大)
  • 信用分配粒度过粗

3.对于问题1:绝对优势问题(没有减去baseline导致方差大) 。

baseline的实现有很多,我们先介绍下不同的价值评估函数的基本定义,然后我们可以总结对比下PPO(value模型),GRPO(组内平均值),REINFORCE++(batch内平均值)等系列算法,他们是如何针对这一点改进的

4.对于问题2:信用分配粒度过粗

其实目前改进比较少,因为大家目前都在卷好验证的领域(数学,代码),再加上PRM太难

不过我们也进一步了解GAE是如何权衡蒙特卡洛法(看的远),和时序差分法(看的近),来深入理解ppo的信用分配过程。了解这个过程后,我们也介绍下critic模型是如何更新的?

本篇希望回答以下问题?

1.传统的RL中的概念是这么对LLM的RLHF中去的。(给一些详细的定义)

2.将传统的RL中的策略梯度算法应到LLM的RLHF中去的时候,是怎么一步一步展开,并转换为现代深度学习框架的loss的。(使用策略梯度算法,从目标的梯度展开,然后推导到代理目标(surrogate objective),其次,我们也介绍应用了重要性采样后的代理目标(让大家可以对齐paper和代码实现))

3.分析策略梯度的2个问题。1.信用分配粒度过粗2.没有减去baseline导致方差大。后续的研究是如何解决的。

4.介绍下ppo中critic模型的作用,以及在RLHF场景它模型的监督信号是什么?

初步定义

强化学习基础:一般来说,强化学习旨在训练一个智能体,该智能体与外部环境进行多轮交互,通过学习合适的策略进而最大化从外部环境获得的奖励。

状态

​S=\{s_0,s_1,...s_t\}

对于大语言模型而言, ​s_i 代表当前状态就是:用户输入的问题和前i个已经生成的内容(即当前状态)

策略(动作)

​a_t=\pi_{\theta}(.|s_t)

智能体在看到状态 ​s_t 的情况下,其可能执行的动作 ​a_t 服从概率分布 ​\pi_{\theta} .

对于大语言模型而言,策略 ​\pi_{\theta} 就是我们要训练的大语言模型,根据用户输入的问题和已经生成的内容(即当前状态),生成下一个token的概率。

奖励

当智能体采取了某个行动后,状态由 ​s_t 变为 ​s_{t+1} ,那么外部环境会给予智能体一个奖励分数 ​r_t

轨迹

\tau=\{s_0,a_0,r_0,s_1,a_1,r_1,..s_{T-1},a_{T-1},r_{T-1} \}

轨迹是状态state-action-reward chain。代表大语言模型,在训练过程中,​s_o代表输入,​a_0代表生成的第一个词汇,​r_0代表这个词汇的奖励,一直到第T个词。

回报(累积奖励)

回报是针对轨迹而言的,代表这个轨迹的所有累积奖励之和。

​R(\tau)=\sum_{t=0}^{t=T-1}{r_i}

回报包含即时的reward以及未来的reward。。

回报可以来判断两条轨迹(策略)的好坏

累积折扣奖励

R(\tau)=\sum_{t=0}^{t=T-1} \gamma ^tr_i

对于未来的reward,增加一个折损因子,通过这个折损因子可以用来调整对近期或远期奖励的重视程度。

解释

1.LLM 的 Episode 长度是动态的(直到生成 EOS)

2.值得一提的是:对于LLM中的RL

  • 1.对于t如果是中间步骤,那么其通常r=0(或者为了防止模型崩坏加的微小 KL 散度惩罚)。
  • 2.对于t是结尾的token,那么其通常等于奖励模型打出的得分。

策略梯度的推导:

介绍优化目标,推导梯度,引出代理目标(方便在深度学习框架中的计算loss的目标),分析策略梯度loss的问题,引出两个优化(1.减去baseline。2.细粒度的reward分配)

优化目标:

假设参数为 ​\theta 的模型做出决策轨迹的概率为 ​p_{\theta}(\tau) ,该决策轨迹在最终状态能够获得的累积奖励为 ​𝑅(𝜏)。而强化学习的目标就是最大化获得的奖励,即:

\underset{\pi_{\theta}}{arg max}J(\theta)=\underset{\pi_{\theta}}{arx max}\mathbb{E_{\tau\sim \pi_{\theta}}}[R(\tau)]=\underset{\pi_{\theta}}{arx max}\sum\limits_{\tau}R(\tau)p(\tau|\pi_{\theta})

解释:

  1. ​\tau :表示已经生成的轨迹。
  2. ​\pi_{\theta}: 要训练的模型参数。
  3. ​R(\tau) 是轨迹所获得的累积奖励(回报)。
  4. ​p(\tau|\pi_{\theta}) ,在当前模型参数下,模型所生成当前这组轨迹的概率。

这个优化目标的意义是:找到一个策略模型 ​\pi_{\theta} ,使得它生成的一组token,获得回报(累计奖励)期望尽量高。

求上式的梯度:

\nabla J(\theta)= \sum\limits_{\tau}R(\tau) \nabla p(\tau|\pi_{\theta})

我们对​\nabla p(\tau|\pi_{\theta})加一个log进行求导:

\nabla logp(\tau|\pi_{\theta})= \frac{1}{p(\tau|\pi_{\theta})} \nabla p(\tau|\pi_{\theta})

可以得到

\nabla p(\tau|\pi_{\theta})= p(\tau|\pi_{\theta})\nabla logp(\tau|\pi_{\theta})

代入上式,我们得到:

\begin{equation}\begin{split} \nabla J(\theta)&= \sum\limits_{\tau}R(\tau) \nabla p(\tau|\pi_{\theta})\\ &=\sum\limits_{\tau}R(\tau)p(\tau|\pi_{\theta})\nabla logp(\tau|\pi_{\theta})\\ &=\mathbb{E_{\tau\sim\pi_{\theta}}}[R(\tau)\nabla logp(\tau|\pi_{\theta})] \end{split} \end{equation}

上式是轨迹维度的,我们在LLM的RLHF中,还可以继续将其展开。

\begin{equation} \begin{split} \nabla logp(\tau|\pi_{\theta})&=\nabla log(p(s_0)\prod\limits_{t=0}^{T-1}p(s_{t+1}|s_t,a_t)\pi_{\theta}(a_t|s_t) )\\ &=\nabla [log(p(s_0)+\sum\limits_{t=0}^{T-1}logp(s_{t+1}|s_t,a_t)+\sum\limits_{t=0}^{T-1}\pi_{\theta}(a_t|s_t)] \\ &= \sum\limits_{t=0}^{T-1} \nabla log \pi_{\theta}(a_t|s_t) \end{split} \end{equation}

解释:因为我们是对​\theta求导,因此可以把前两项和环境有关系的约掉,最后

\nabla J(\theta)=\mathbb{E_{\tau \sim \pi_{\theta}}}[R(\tau)\sum\limits_{t=0}^{T-1} \nabla log \pi_{\theta}(a_t|s_t))]

这里的期望没办法计算,所以我们用采样的方式采样N个轨迹​\tau,来计算这个梯度,每个轨迹涵盖 ​T_i 步。则优化目标可以被写成:

\nabla J(\theta)=\mathbb{E_{\tau \sim \pi_{\theta}}}[R(\tau)\sum\limits_{t=0}^{T-1} \nabla log \pi_{\theta}(a_t|s_t))] \approx \frac{1}{G} \sum\limits_{i=0}^{G-1} \sum\limits_{t=0}^{T_i-1} R(\tau_{i}) \nabla log \pi_{\theta}(a_{i,t}|s_{i,t}))]

解释:当​R(\tau_{i})越高,我们就提高这个​a_{i,t}的梯度。然后可以用梯度上升的方式进行优化:

\pi_{\theta} ← \pi_{\theta}+ 𝜂\nabla J(\theta)

在深度学习框架中,我们需要计算一个最后的loss,然后梯度由深度学习框架自己求。

根据上述梯度,我们设计一个surrogate objective (代理目标)

\begin{equation} \begin{split} J(\theta)_{surrogate}&=\mathbb{E_{\tau\sim \pi_{\theta}}}[R(\tau)]\\ &\approx\frac{1}{G}\sum\limits_{i=0}^{G-1}\sum\limits_{t=0}^{T_i-1}R(\tau_i)log\pi_{\theta}(a_{i,t}|s_{i,t}) ) \end{split} \end{equation}

解释

1.这个优化目标的梯度和我们上述推导的梯度一样,因此,在策略梯度中,我们使用这个代理目标来计算loss(这个更容易计算点,大家可以尝试在loss展开到这一步,比较难)。在实际使用深度学习框架计算loss的时候,我们会使用下面的形式:

\begin{equation} \begin{split} L(\theta)&=-J(\theta)_{surrogate} \\&=-\frac{1}{G}\sum\limits_{i=0}^{G-1}\frac{1}{T}\sum\limits_{t=0}^{T_i-1}R(\tau_i)log\pi_{\theta}(a_{i,t}|s_{i,t}) ) \end{split} \end{equation}

解释:

1.加一个负号,最小化loss,最大化优化目标。

2.把求和变为求平均,这样训练的时候稳定点。

这个loss的形式在加入重要性采样后还会变一下,为了一次性了解这个操作(对齐ppo,grpo论文和代码实现),我们在这里一并介绍。

重要性采样的目标是:为了更充分的利用rollout的轨迹,比如旧策略模型采样出来的轨迹,我们更新多次,但是这样两种期望差异比较大,因此我们引入重要性采样(具体方法不多赘述)来解决这个问题。

对应的梯度变为:

\nabla J(\theta)=\mathbb{E_{\tau \sim \pi_{\theta_{old}}}}[R(\tau)\sum\limits_{t=0}^{T-1} \frac{\pi_{\theta}(a_t|s_t)}{\pi_{\theta_{old}}(a_t|s_t)}\nabla log \pi_{\theta}(a_t|s_t))]

因为有:

\nabla log\pi_{\theta}(a_t|s_t)) \pi_{\theta}(a_t|s_t)=\nabla \pi_{\theta}(a_t|s_t))

带入梯度,得到:

\nabla J(\theta)=\mathbb{E_{\tau \sim \pi_{\theta_{old}}}}[R(\tau)\sum\limits_{t=0}^{T-1} \frac{\nabla \pi_{\theta}(a_t|s_t)}{\pi_{\theta_{old}}(a_t|s_t)}]

对应的loss也可以反出来。

\begin{equation} \begin{split} L(\theta)&=-J(\theta)_{surrogate} \\&=-\frac{1}{G}\sum\limits_{i=0}^{G-1}\frac{1}{T}\sum\limits_{t=0}^{T-1} R(\tau_i)\frac{ \pi_{\theta}(a_{i,t}|s_{i,t})}{\pi_{\theta_{old}}(a_{i,t}|s_{i,t})} \end{split} \end{equation}

解释

1.这里就是ppo,grpo,以及代码实现中常见的形式,不过他们的回报是token级别的,而非序列(轨迹)级别的。

这个loss存在以下的问题。

1.信用分配问题:当前我们使用的回报的力度太粗了,一个轨迹(序列)的整体的回报很高,不能说明他的每一个动作都好,我们对轨迹中所有的动作都使用相同的回报是否合理?我们希望对一个动作能有两种方法衡量其价值:1)评估动作当下的影响(单步回报);2)又能体现动作对后续策略轨迹的长远效益。(轨迹整体回报)

2.绝对优势问题:不同轨迹的 ​R(\tau_i) 是不同的,而且我们是采样了一些轨迹,如果我们采样的轨迹​ R(\tau_i) 都是正的,把轨迹中对应的动作都提高,但一些我们没有采样到,是正确答案的轨迹对应的动作就得降低,因为,我们不希望奖励总是正的,而是要减去一个基线(平均回报),表示相对优势。只有一个动作相比相对优势是正的,我们才提高其动作的概率,如果相比相对优势是负向的,我们要降低其动作的概率。

针对分析出来的这两个问题,我们一个一个去解决,但在解决之前,我们需要引入一些基础概念。

价值评估函数

更细粒度的评估:我们先引入状态价值函数和动作价值函数,把对整个轨迹的回报细粒度到对1.当前策略达到的状态预期能获得的总回报。2.当前策略达到的状态下,执行的动作a,预期能获得的总回报。为我们解决信用分配打好基础。

1.累积折扣奖励:

G_t=\sum_{t=0}^{t=T-1} \gamma ^tr_i

2.t时刻后的累积折扣奖励:

由于MDP假设,t时刻发生之前的事情和t时刻无关,而t时刻发生后的事情才受到t时刻的影响。

G_t=\sum_{t=t'}^{t=T-1} \gamma ^{t'}r_i

我们在之后的论述中,都是使用t时刻后的累积折扣奖励。

3.状态价值函数:

agent在当前状态下 ​s_t 开始,一直根据策略 ​\pi 行动,预期能获得的总回报(累计折扣奖励)的期望值:

V_{\pi}(s_t)=E_{\pi}[G_t|s_t]

解释

1.其依赖当前状态 ​s_t ,也依赖当前策略模型 ​\pi (是策略模型使得其到达这个状态)。

2.直观数值可以反应:当前策略模型目前到达的这个状态好不好。

在RLHF的LLM里举例来说:其可以表示策略模型当前生成的上文好不好。

1.比如:面对 “法国的首都是” 这个上文,当前的 LLM 平均表现如何? 如果当前的 LLM 已经训练得很好了(Knowledgeable),它大概率后面会接 “巴黎”。那么这个状态 的价值就很高,因为预期能拿到高分。

2.实现:这个值是critic模型来预测得到的,在PPO的训练过程中,critic模型一直在学习如何更好的预测这个目标。

4.动作价值函数:(动作-状态价值函数)

agent在当前状态下 ​s_t 开始,执行动作 ​a_t ,一直根据策略 ​\pi 行动,预期能获得的总回报(累计折扣奖励)的期望值:

Q_{\pi}(s_t,a_t)=E_{\pi}[G_t|s_t,a_t]

解释:

1.其依赖当前状态 ​s_t ,也依赖当前策略模型 ​\pi 和动作 ​a_t

2.其数值可以反应,在当前状态下,执行这个动作好不好?

在RLHF的LLM举例来说:其可以表示策略模型面对已经生成的上文(状态)下(“法国的首都是”)选择的token好不好?

1.比如:面对 “法国的首都是” 这个上文,

  • 选择token a:巴黎。

    Q_{\pi}(s_t,“巴黎”)=0.99
  • 选择token b:伦敦。

    Q_{\pi}(s_t,“北京”)=0.01

但是大家注意,在实际llm训练中,我们往往不会直接训练一个动作价值函数模型来预测,因为词表太大了,计算量无法接受,具体的实现在下面介绍:

5.状态价值函数和动作价值函数的关系

V_{\pi}(s_t)=E_{\pi}[\pi (a_t|s_t)Q(s_t,a_t)]=\sum\limits_{a_t\in A}\pi (a_t|s_t)Q(s_t,a_t)

解释:

1.动作-状态价值函数的期望是状态价值函数。

其次,我们刚才也说了,计算 ​Q_{\pi}(s_t,a_t) 不现实,因此,我们将其进行一个转换:

Q_{\pi}(s_t,a_t)=E_{\pi}[G_t|s_t,a_t]

其中的累积折扣奖励可以拆解。

G_t= r_t+\gamma r_{t+1}+\gamma^2 r_{t+2}+...\gamma^{T_n-1-t} r_{T_n-1}=r_t+\gamma G_{t+1}

代入上式可以得到:

\begin{equation} \begin{split} Q_{\pi}(s_t,a_t)&=E_{\pi}[G_t|s_t,a_t]\\&=E_{\pi}[r_t+\gamma G_{t+1}|s_t,a_t]\\&=r_t+\gamma E_{\pi}[ G_{t+1}|s_t,a_t]\\&=r_t+\gamma V_{\pi}(s_{t+1}|s_t) \end{split} \end{equation}

解释:

1.​Q_{\pi}(s_t,a_t) 是当前状态下执行当前动作,未来总回报的期望。我们通过上面的式子将其展开代表,除了从不好计算的原因,也可以理解为我们多观测一步当前动作的奖励 ​r_t ,来逼近真实的未来总回报的期望。

2.从第二行到第三行,我们为什么可以把期望去掉呢?

原始的论文做实验的时候发现这样比较好,后续大家就沿用了。

绝对优势问题

在上面的分析中提到,我们不希望奖励总是正的,而是要减去一个基线(平均回报),表示相对优势。进一步的,在更细力度的评估问题下,我们也引入优势价值函数(让动作价值函数-状态价值函数)。

优势价值函数:(动作状态价值函数-状态价值函数)

A_{\pi}(s_t,a_t) = Q_{\pi}(s_t,a_t)-V_{\pi}(s_t)

解释

1.这个优势价值函数,就可以解决我们上面提出的绝对优势问题,通过减去一个和当前动作无关的baseline( ​V_{\pi}(s_t) ),来成为相对优势函数,可以稳定训练。

下面我们介绍下最新的算法是怎么计算这个值的。

PPO

ppo的优势价值函数如下面公式所示的,是动作-状态价值函数-价值函数,不过它引入了参数GAE来平衡方差和偏差(后面在信用分配环节详细介绍)。

A^{GAE}_t = \sum\limits_{l=0}^{l=+\infty }(γλ)^l * δ_{t+l}
δ_t = Q_{\pi}(s_t,a_t)-V_{\pi}(s_t)= r_t + γV_{\pi}(s_{t+1}) - V_{\pi}(s_t)

GRPO

A_{i,t}=\frac{r_i-mean(r)}{std(r)}

对于一个prompt,它使用采样出来一组轨迹的回报的平均值作为baseline,然后除一个标准差,让 ​A_{i,t} 归一化下,作为当前轨迹的优势价值函数。

REINFORCE++

对于一个prompt,它使用采样出来batch内所有轨迹的回报的平均值作为baseline,然后除一个batch内所有轨迹的标准差,让 ​A_{i,t} 归一化下,作为当前轨迹的优势价值函数。

作者认为:使用一个更大、更全局的样本集合(整个训练批次)来计算优势函数的均值和标准差,可以得到一个比局部(提示级别)标准化更稳定、更准确(有效无偏)的优势估计,从而在提升无 Critic 方法效率的同时,保证甚至超越传统 PPO 的训练稳定性和最终性能。

公式4:是参考RLHF场景下的ppo reward格式(token级别的奖励减去一个当前模型和参考模型的kl散度(防止reward hacking))

公式5:是我们上面介绍的优势价值函数,和grpo的区别是其来自batch。

ReMax

​A_{i,t}=r_i-r_{o'}

​o' 是一个greedy decoding的回复(可以认为从分布中取众数)

作者的发表论文的心路历程讲的很生动:ReMax背后的故事

信用分配(credit assignment)问题

再把评估粒度变细,以及解决了绝对优势问题之后,我们来看如何更好的做信用分配,有两种典型的方法

1. 蒙特卡洛法(Monte Carlo, MC):

这是最直观的方法。它的核心思想是 “等到这一句生成完了,甚至整个对话结束了,再回头算账”。

因此它的优势函数是:

A_{\pi}(s_t,a_t)=G_t-V_{\pi}(s_t)

这种方法的特点是:

1.无偏差:使用的是真实发生的奖励总和,数学期望上是准确的。

2.方差高(High Variance):这是致命伤。因为一个 Token 的选择会影响后面几十上百个 Token 的生成,任何一步的随机性都会累积。长文本生成中,这种随机性导致的方差极其巨大,会导致训练难以收敛。

2.时序差分法(Temporal Difference, TD):

TD 的核心思想是 “只看眼前一步,剩下的靠猜(利用 Value Model)”。这里主要指一步 TD(1-step TD)。

A_{\pi}(s_t,a_t) =r_t+V_{\pi}(s_{t+1})-V_{\pi}(s_t)

使用当前时刻的奖励+下一个状态的状态价值函数来作为 ​Q_{\pi}(s_t,a_t) 的估计值。

特点:

1.方差低(Low Variance):只依赖当前一步的随机性,受后续长序列生成的干扰小。

2.有偏差(Biased):它高度依赖 Critic Model 的准确性。在训练初期,Critic Model 也是瞎猜的,这会导致估计出来的优势也是错的,

这两种方法都各有其缺点,合理的方法应该取一个折中。

3.广义优势估计器(generalized advantage estimator,gae)

在这里我们考虑时序差分法,不仅考虑一步,也多考虑下一步,不过乘一个折扣因子γ。

1-step TD:

A_t^{(1)}= δ_t = r_t + γV_{\pi}(s_{t+1}) - V_{\pi}(s_t)

2-step TD:(在前面基础上加上折扣因子γ乘多看一步的优势。)

A_t^{(2)}= δ_t+γδ_{t+1} = r_t + γV_{\pi}(s_{t+1}) - V_{\pi}(s_t)+γ( r_{t+1} + γV_{\pi}(s_{t+2}) - V_{\pi}(s_{t+1}))=-V_{\pi}(s_t)+r_t +γr_{t+1}+γ^2V_{\pi}(s_{t+2})

k-step TD:

A_t^{(k)} = \sum\limits_{l=0}^{l=k-1}γ^lδ_{t+l}=-v(s_t)+r_t+γr_{t+1}+γ^2r_{t+2}+...γ^{k-1}r_{t+k-1}+γ^kv(s_{t+k})

+∞-step TD:

A_t^{(+∞)} = δ_t+γδ_{t+1}+... = -v(s_t)+\sum\limits_{l=0}^{l=+∞}γ^lr_{t+l}

解释:

i.这里省略了一个+∞的状态价值函数。(这个是终止状态,一般设置为0,或者 ​γ 一般小于1,则l->+∞的时候, ​γ^l 这一项为0)。

ii.计算 GAE:对于每一个时间步t,其广义优势估计Â_t是未来所有 TD 残差的指数加权折扣和。

动机

1.我们如果采用1-step TD,偏差大;如果使用+∞-step TD,方差大。

2.那么GAE利用指数移动平均 (EMA, Exponential Moving Average)来把1,2,3...+∞-step的TD都加进来。

A^{GAE}_t = (1-\lambda) (A^{(1)}_t + \lambda A^{(2)}_t + \lambda^2 A^{(3)}_t + \dots)

乘 (​1-\lambda) 的原因是求这么多项的平均。在 ​\lambda 小于1的前提下(一般ppo设置)。

1+\lambda+\lambda^2+...\lambda^{+∞}=1/(1-\lambda)

经过化简(无穷级数求和),最后是:

A^{GAE}_t = \sum\limits_{l=0}^{l=+\infty }(γλ)^l * δ_{t+l}

展开:

A^{GAE}_t = δ_t+γλ * δ_{t+1}+(γλ)^2 * δ_{t+2} +...(γλ)^{l} * δ_{t+l}

其中

δ_t = r_t+γV_{\pi}(s_{t+1})-V_{\pi}(s_t)

直观理解:当前的优势 = 现在的优势 + (打折的) 下一步的优势 + (打折更多) 下下步的优势...

1.​\lambda (Lambda)是 GAE 的平衡因子。它决定了我们在计算优势时,多大程度上信任 Critic 网络的估值,多大程度上信任真实的轨迹

分情况讨论:

​\lambda=0 的时候:

A^{GAE}_t =r_t+V_{\pi}(s_{t+1})-V_{\pi}(s_t)

就是时序差分法。

​\lambda=1 的时候:

A^{GAE}_t = \sum \gamma^l \delta_{t+l}=-v(s_t)+\sum\limits_{l=0}^{l=+∞}γ^lr_{t+l}

就是蒙特卡洛。

​\lambda 在0-1之间,就是这两者之间。

2.γ (Gamma):是时间折扣因子,计算回报的时候,多考虑当前动作的奖励,而对后续动作的奖励,都打一个折扣,。

总结下GAE:

1.从本质来看:解决在评估当前动作的优势时,解决时序差分法的高偏差问题,以及蒙特卡洛法的高方差问题,在这两者之间来进行折中,来计算当前动作的优势。

2.从实现上看:在评估当前动作的优势时,使用从1到k step TD的指数加权折扣和(推导后,等价于当前动作TD残差+打折后的下一个动作TD残差+...(再打折后)到最后一个动作的残差)。

verl中计算的代码:

https://github.com/volcengine/verl/blob/d82d39d17d08eb3e32ed8e0e23002f12a4550663/verl/trainer/ppo/core_algos.py#L213

使用递归法来求:当前时刻的优势,因此从后往前求:

A^{GAE}_t = δ_t+γλ * δ_{t+1}+(γλ)^2 * δ_{t+2} +...(γλ)^{l} * δ_{t+l}= δ_t+γλ\sum\limits_{l=0}^{l=+\infty }(γλ)^{l} * δ_{t+1+l}

另t=t+1,则有:

A_t^{GAE} = δ_t + γλ·A_{t+1}^{GAE}

用这个公式可以从后往前递归的求GAE.

@register_adv_est(AdvantageEstimator.GAE)  # or simply: @register_adv_est("gae")
def compute_gae_advantage_return(
    token_level_rewards: torch.Tensor,
    values: torch.Tensor,
    response_mask: torch.Tensor,
    gamma: torch.Tensor,
    lam: torch.Tensor,
):
    """Adapted from https://github.com/huggingface/trl/blob/main/trl/trainer/ppo_trainer.py

    Args:
        token_level_rewards: `(torch.Tensor)`
            shape is (bs, response_length)
        values: `(torch.Tensor)`
            shape is (bs, response_length)
        response_mask: `(torch.Tensor)`
            shape is (bs, response_length). [EOS] mask. The token after [EOS] have mask zero.
        gamma is `(float)`
            discounted factor used in RL
        lam: `(float)`
            lambda value when computing Generalized Advantage Estimation (https://arxiv.org/abs/1506.02438)

    Returns:
        advantages: `(torch.Tensor)`
            shape: (bs, response_length)
        Returns: `(torch.Tensor)`
            shape: (bs, response_length)

    """
    with torch.no_grad():
        nextvalues = 0
        lastgaelam = 0
        advantages_reversed = []
        gen_len = token_level_rewards.shape[-1]

        for t in reversed(range(gen_len)):
            delta = token_level_rewards[:, t] + gamma * nextvalues - values[:, t]
            lastgaelam_ = delta + gamma * lam * lastgaelam

            # skip values and TD-error on observation tokens
            nextvalues = values[:, t] * response_mask[:, t] + (1 - response_mask[:, t]) * nextvalues
            lastgaelam = lastgaelam_ * response_mask[:, t] + (1 - response_mask[:, t]) * lastgaelam

            advantages_reversed.append(lastgaelam)
        advantages = torch.stack(advantages_reversed[::-1], dim=1)

        returns = advantages + values
        advantages = verl_F.masked_whiten(advantages, response_mask)
    return advantages, returns

critic模型的作用以及是如何训练的?

A_{\pi}(s_t,a_t) = Q_{\pi}(s_t,a_t)-V_{\pi}(s_t)

1.critic模型的作用是来计算 ​v(s_t) 的,评估当前状态可以获得的累积折扣回报的期望,作为一个baseline。

输入: 状态(用户输入的问题和前i个已经生成的内容)。输出:基于当前已经生成的内容,未来最终生成的回复得到的预期分数(预期回报)

举例子来理解。

如果大模型生成的上文:“法国的首都是”,价值模型(critic)预测这句话很好,最后能拿到不错的回报:0.5分。

下一个token大模型预测采样到两个词

token a:巴黎。

Q_{\pi}(s_t,“巴黎”)=0.99

那么

A_{\pi}(s_t,"巴黎")=0.99-0.5=0.44

token b:伦敦。

Q_{\pi}(s_t,“北京”)=0.01

那么

A_{\pi}(s_t,"北京")=0.01-0.5=-0.49

我们就鼓励提升token a的概率,降低token b的概率。

从这个例子: ​V_{\pi}(s_t) 的值和 ​Q_{\pi}(s_t,a_t) 的值有一定关系,可以理解为前者是后者的平均值。

在当前状态下,采样到的所有动作最后的回报都很低,那么 ​V_{\pi}(s_t) 一定很低,证明这个状态(这句话)不是一个好的话。

2.critic模型可以提供可靠的token级别的中间价值估计。

在RLHF场景,因为PRM比较难,那么存在下面的计算reward的方式:

  • 1.对于t如果是中间tiken,那么其通常r=0(或者为了防止模型崩坏加的微小 KL 散度惩罚)。
  • 2.对于t是结尾的token,那么其通常等于奖励模型打出的得分。

critic模型的另外一个作用是:可以将最后的 Reward 分配到中间步骤(Token-level)。他是如何做到的呢?

1.我们先介绍已经有了一个训练好的critic模型,他可以判断当前已经生成的句子的好坏程度。

对于中间token的优势,上文提到我们有GAE来计算:假设简略版本的GAE是这

A_t \approx r_t+γV_{\pi}(s_{t+1})-V_{\pi}(s_t)

由于中间token的奖励是0(或者微小的负的kl散度), ​γ =1, 0<\lambda<1 。那么

A_t \approx V_{\pi}(s_{t+1})-V_{\pi}(s_t)

这意味着:当前动作的 “优势 / 得分”,等于执行动作后,Critic 认为局势变好了多少。

因此一个好的critic模型可以:可以将最后的 Reward 分配到中间步骤(Token-level)

2.假设我们没有一个训练好的critic模型,他是如何训练的?他的监督信号是怎么来的?

答案是:可以在强化学习训练的过程中训练,用mse loss,学习的目标就是在当前状态 ​s_t 下,执行动作 ​a_t ,所能得到的回报 ​Q_{\pi}(s_t,a_t) 。因为会采样到很多个轨迹,因此critic模型最终预测得到的是一个平均回报,也就是 ​V_{\pi}(s_t) 的定义。

critic 的mse loss公式如下

loss=mse(V_{new}(s_t), A_t^{\text{GAE}} + V_{old}(s_t))

解释

1.当前状态下,critic模型的预测值是 ​V_{new} , ​V_{old}(s_t) 是老的critic模型对当前状态的预测值 , ​A_t^{\text{GAE}} 是当前动作在时序差分法和蒙特卡洛之间的一个折中的优势价值函数。

2.假设生成的一批样本,ppo更新两次。

  • 1.第一个step里,这两个值是相同的。(虽然两者剪掉了,但是梯度仍然会回传)
  • 2.在第二个step,两个值不同。
V_{target} = A_t^{\text{GAE}} + V_{old}(s_t)

代表了当前状态下执行特定动作后的修正回报(也就是 Q 值的一个估计样本)。

我们之所以用它作为 Critic 的 Target,是因为:

  • 我们没有上帝视角,拿不到真实的 ​V(s),只能拿到采样的 ​Q(s, a)
  • 利用 MSE 的性质,让 ​V(s) 通过拟合这些 ​Q 样本,最终学会预测期望值。
  • 利用 GAE,是为了让这个 Target 值比纯粹的蒙特卡洛回报方差更小,让 Critic 学得更稳。

考虑一种RLHF场景下常见的情况,展开GAE,当lamba和gemma==1的时候。

A_t^{\text{GAE}}(s_t,a_t) = r_t+V_{\pi}(s_{t+1})-V_{\pi}(s_t)+r_{t+1}+V_{\pi}(s_{t+2})-V_{\pi}(s_{t+1}))+.... \\ =\sum_{l=0}^{l=\infty }r_{t+l}-V_{\pi}(s_t)
这个时候:
loss= (V_{new}(s_t)-(\sum_{l=0}^{l=\infty }r_{t+l}-v_{old}(s_t)+v_{old}(s_t)))^2\\ =(V_{new}(s_t)-\sum_{l=0}^{l=\infty }r_{t+l})^2

可以看出,就是在学习回报(RLHF场景下,包含reward值-生成的中间token的kl散度)。

参考文献

1.High-Dimensional Continuous Control Using Generalized Advantage Estimation(GAE)
https://arxiv.org/abs/1506.02438
2. 蘑菇书EasyRL
https://datawhalechina.github.io/easy-rl/#/chapter4/chapter4
3. Mathematical Foundations of Reinforcement Learning。
4. https://arxiv.org/pdf/2310.10505
5. https://arxiv.org/pdf/2501.03262
6. GAE(Generalized Advantage Estimation)理解及推导
https://zhuanlan.zhihu.com/p/675309680