【AI Infra】【RLHF框架】四、VeRL中PPO、GRPO、REINFORCE++、RLOO实现源码解析

系列文章：

【AI Infra】【RLHF框架】一、VeRL中基于Ray的执行流程源码解析
【AI Infra】【RLHF框架】二、VeRL中colocate实现解析
【AI Infra】【RLHF框架】三、VeRL中的Rollout实现源码解析
【AI Infra】【RLHF框架】四、VeRL中PPO、GRPO、REINFORCE++、RLOO实现源码解析

​ 相比于前三篇博客，本偏博客涉及的公式较多，但对于理解RL框架这也是不可避免的。

​ 本篇博客尽量采用统一的符号来介绍PPO、GRPO、REINFORCE++和RLOO的理论，这样便于理解各类算法之间的关联。此外，很多时候实现代码为了简洁或者数值稳定等原因，不一定完全按照公式进行实现。本文也尽量补齐理论和实现的差距。

希望这篇博客对你有所帮助，如有错误，欢迎指正。

一、LLM场景下的PPO

1. PPO原理

​ 自回归语言模型${\pi }_{\theta }$将提示$x$作为输入，然后通过自回归的方式逐步产生输出$y$。若采用贪心解码的方式，则生成$y$中第$t$个token的过程为
$$\begin{gathered} y_t=\arg \underset{y}{\max }{\pi }_{\theta }(y|x,y_{<t}) \\ \text{\hspace{1em}(1)} \end{gathered}$$
若将这个过程建模为马尔科夫决策过程(MDP)，则$s_t=x\oplus y_{<t}$，$a_t$表示进行第$t$个token的选择。单个样本的生成过程构成了一条完整的轨迹$\tau =(s_0,a_t,s_1,a_1,...)$。

​ 基于上面的定义，LLM场景下PPO损失函数定义为
$$\begin{gathered} L_{\text{PPO}}(\theta )=-{\text{E}}_{x\sim p,y\sim \pi (\cdot |x)}\frac{1}{|y|}\sum _{t=1}^{|y|}\min [r_t(\theta )A_t,\text{clip}(r_t(\theta ),1-\epsilon ,1+\epsilon )A_t] \\ \text{\hspace{1em}(2)} \end{gathered}$$

• $A_t$是优势函数，表示采取动作$a_t$的预期收益，该值的计算会在后续GAE小节中详细介绍。
• $r_t(\theta )=\frac{{\pi }_{\theta }(a_t|s_t)}{{\pi }_{{\theta }_{\text{old}}}(a_t|s_t)}$是重要性采样，通过重要性采样能够调整旧策略${\pi }_{{\theta }_{\text{old}}}$样本的权重来适应新策略${\pi }_{\theta }$的分布。例如，旧策略以高概率选择了某个动作$a$，但新策略以低概率选择该动作，那么$r_t(\theta )$会较小，从而降低该样本对整个梯度的影响。反之，权重会放大其贡献。
• $\text{clip}$是截断函数，用于防止出现大幅度更新，导致模型退化。

​ 在verl中该损失函数的计算如下：

# verl/trainer/ppo/core_algos.py
def compute_policy_loss(old_log_prob, log_prob, advantages, eos_mask, cliprange):
    """
    old_log_prob: (bs, response_length)
    log_prob: (bs, response_length)
    advantages: (bs, response_length)
    eos_mask: (bs, response_length)
    """
    negative_approx_kl = log_prob - old_log_prob
    ratio = torch.exp(negative_approx_kl)
    ppo_kl = verl_F.masked_mean(-negative_approx_kl, eos_mask)

    pg_losses = -advantages * ratio
    pg_losses2 = -advantages * torch.clamp(ratio, 1.0 - cliprange, 1.0 + cliprange)

    pg_loss = verl_F.masked_mean(torch.max(pg_losses, pg_losses2), eos_mask)
    pg_clipfrac = verl_F.masked_mean(torch.gt(pg_losses2, pg_losses).float(), eos_mask)
    return pg_loss, pg_clipfrac, ppo_kl

---

2. 广义优势估计(GAE)

​ 在PPO原理介绍中提及了优势函数$A_t$，本小节介绍估计$A_t$的常用方式GAE。

2.1 优化函数的定义

$$\begin{gathered} A_t=Q(s_t,a_t)-V(s_t) \\ \text{\hspace{1em}(3)} \end{gathered}$$

其含义是在当前状态$s_t$下采取动作$a_t$的未来预期收益$Q(s_t,a_t)$与状态$s_t$下平均动作预期收益的差值，即衡量动作$a_t$比平均动作好多少。在LLM的RL场景中，$A_t$表示在第$t$个token选择能获得的收益比平均选择所有token高多少。

2.2 TD估计

​ 单步TD估计。由于$Q$和$V$这两函数都是未知的，理论上需要两个深度模型来分别拟合。但是，通过贝尔曼方程将$Q$转换为$V$，即
$$A_t^1=r_t+\gamma V(s_{t+1})-V(s_t)$$
其中$\gamma$是折扣因子，$r_t$是第$t$步的实际奖励。这样就能使用单个神经网络来估计优势值，其中拟合$V$的这个网络就是Critic模型。这种包含了一步真实奖励$r_t$的估计称为优势函数的单步TD估计。但是，单步TD存在低方差但偏差高的问题。因为，仅依赖一步真实奖励的方差较低，而更多使用Critic模型的预测值则可能导致高偏差(模型预测值往往波动较小，但可能存在整体预测偏高或偏低的问题)。

​ n步TD估计。为了方便表示，令${\delta }_t=r_t+\gamma V(s_{t+1})-V(s_t)$。那么单步TD估计表示为
$$A_t^1={\delta }_t$$
为了降低单步TD估计中的偏差，可以考虑估计中使用更多的真实奖励，那么通过贝尔曼方程将单步TD中的$V(s_{t+1})$展开，就得到2步TD估计
$$\begin{array}{rl}{A}_t^2& =r_t+\gamma (r_{t+1}+\gamma 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{array}$$
以此类推，n步的TD估计为
$$\begin{gathered} A_t^n=\sum _{l=0}^{n-1}{\gamma }^l{\delta }_{t+l} \\ \text{\hspace{1em}(4)} \end{gathered}$$

2.3 广义优势估计(GAE)原理

​ n步TD估计中n越大则使用的真实奖励越多，对应的方差就越高，偏差越低。由于不同任务对应的最优n是不同的，很难精确选择合适的n，GAE中则选择融合所有可能的步长($n=1,2,\dots ,\infty$)。此外，引入一个衰减因子$\lambda \in [0,1]$来对远期步长施加衰减，即
$$\begin{gathered} A_t^{\text{GAE}}=\sum _{l=0}^{\infty }(\gamma \lambda {)}^l{\delta }_{t+1} \\ \text{\hspace{1em}(5)} \end{gathered}$$
当$\lambda =0$时，GAE估计退化为${\delta }_t$；当$\lambda =1$时，GAE等价于蒙特卡洛估计(全部使用真实奖励)。所以，$\lambda$可以当做是调节方差和偏差的超参数，$\lambda$越大则真实奖励越多，方差越大，偏差越小。

2.4 GAE的计算

​ GAE的计算。在实际计算中，通常需要计算所有步骤的GAE，为了避免重复计算可以采用迭代的方式从轨迹末端开始计算。
$$\begin{array}{rl}{A}_t^{\text{GAE}}& ={\delta }_t+(\gamma \lambda ){\delta }_{t+1}+(\gamma \lambda {)}^2{\delta }_{t+2}+...\\ & ={\delta }_t+\gamma \lambda ({\delta }_{t+1}+(\gamma \lambda ){\delta }_{t+2}+\dots )\\ & ={\delta }_t+\gamma \lambda A_{t+1}^{\text{GAE}}\end{array}$$
通过上面的公式可以方便的迭代计算GAE。下面是verl中的实现代码：

# verl/trainer/ppo/core_algos.py
def compute_gae_advantage_return(token_level_rewards: torch.Tensor, values: torch.Tensor, eos_mask: torch.Tensor, gamma: torch.Tensor, lam: torch.Tensor):
    """
    token_level_rewards: (bs, response_length)
    values: (bs, response_length)
    """
    with torch.no_grad():
        lastgaelam = 0
        advantages_reversed = []
        gen_len = token_level_rewards.shape[-1]

        for t in reversed(range(gen_len)):
            nextvalues = values[:, t + 1] if t < gen_len - 1 else 0.0
            delta = token_level_rewards[:, t] + gamma * nextvalues - values[:, t]
            lastgaelam = delta + gamma * lam * lastgaelam
            advantages_reversed.append(lastgaelam)
        advantages = torch.stack(advantages_reversed[::-1], dim=1)

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

​ 代码中优势函数的计算与上面迭代公式完全一致。此外，这里还计算了回报returns。回忆一下优势的定义为$A_t=Q(s_t,a_t)-V(s_t)$，那么$Q(s_t,a_t)=A_t+V(s_t)$。$Q(s_t,a_t)$的定义就是状态$s_t$下动作$a_t$的折扣回报。

---

3. KL约束以及近似计算

3.1 KL约束

​ 损失函数$L_{\text{PPO}}(\theta )$中主要是最大化预期收益，但是在LLM场景中仅最大化预期收益可能导致生成的response不符合自然语言的表达。为了解决这个问题引入了KL约束，使得优化后的模型不要太偏离ref(sft)模型。那么损失函数变为
$$\begin{gathered} L_{\text{PPO}}(\theta )=-{\text{E}}_{x\sim p,y\sim \pi (\cdot |x)}\frac{1}{|y|}\{\sum _{t=1}^{|y|}\min [r_t(\theta )A_t,\text{clip}(r_t(\theta ),1-\epsilon ,1+\epsilon )A_t] \\ -\beta \text{KL}[{\pi }_{{\theta }_{\text{old}}}(a_t|s_t),{\pi }_{{\theta }_{\text{ref}}}(a_t|s_t)]\} \\ \text{\hspace{1em}(6)} \end{gathered}$$
其中
$$\begin{gathered} \text{KL}[q,p]=\sum _{x}q(x)\log \frac{q(x)}{p(x)}=E_{x\sim q}\left[\log \frac{q(x)}{p(x)}\right] \\ \text{\hspace{1em}(7)} \end{gathered}$$

3.2 近似计算

​ 在计算KL散度时需要计算两个模型在整个词表上的概率，为了简化计算过程，在实现时通常采用近似计算。

​ k1估计器。一种朴素的近似方式是使用无偏估计
$$\begin{gathered} k1=\log \frac{q(x)}{p(x)} \\ \text{\hspace{1em}(8)} \end{gathered}$$
但这个估计器是高方差的。因为只要$q(x)<p(x)$，$k1$就是负值，但KL散度应该是非负的。

​ k2估计器。为了解决$k1$的高方差问题，可以使用$k2$估计器
$$\begin{gathered} k2=\frac{1}{2}(\log \frac{p(x)}{q(x)}{)}^2 \\ \text{\hspace{1em}(9)} \end{gathered}$$
显然，$k2$估计器是有偏的，但在实际使用中偏差较小且能有效降低方差。

​ k3估计器。因为k1是无偏的，那么降低方差可以通过添加一个期望为0但与k1负相关的项。$\frac{p(x)}{q(x)}-1$的期望为0，因为${\text{E}}_{x\sim q}[\frac{p(x)}{q(x)}-1]={\text{E}}_{x\sim q}\left[\frac{p(x)}{q(x)}\right]-1=0$。因此，k3估计器定义为
$$\begin{gathered} k3=\frac{p(x)}{q(x)}-1-\log \frac{p(x)}{q(x)} \\ \text{\hspace{1em}(10)} \end{gathered}$$

# verl/trainer/ppo/core_algos.py
def kl_penalty(logprob: torch.FloatTensor, ref_logprob: torch.FloatTensor, kl_penalty) -> torch.FloatTensor:
    if kl_penalty == "kl":
        return logprob - ref_logprob

    if kl_penalty == "abs":
        return (logprob - ref_logprob).abs()

    if kl_penalty == "mse":
        return 0.5 * (logprob - ref_logprob).square()

    if kl_penalty == 'low_var_kl':
        kl = ref_logprob - logprob
        ratio = torch.exp(kl)
        kld = (ratio - kl - 1).contiguous()
        return torch.clamp(kld, min=-10, max=10)

​ kl、mse、low_var_kl分别对应$k1$、$k2$和$k3$，abs则是通过绝对值的方式解决$k1$为负的问题。

---

4. 熵正则

​ 在PPO训练的过程，有可能过早将概率集中在某个特定的token上。这意味着过早放弃了探索，容易陷入局部最优。熵正则就是通过在损失函数中添加概率分布${\pi }_{\theta }(a_t|s_t)$的信息熵，从而鼓励学习过程中保持一定的随机性。
$$\begin{gathered} L_{\text{PPO}}(\theta )=-{\text{E}}_{x\sim p,y\sim \pi (\cdot |x)}\frac{1}{|y|}\sum _{t=1}^{|y|}\{\min [r_t(\theta )A_t,\text{clip}(r_t(\theta ),1-\epsilon ,1+\epsilon )A_t] \\ -{\beta }_1\text{KL}[{\pi }_{{\theta }_{\text{old}}}(a_t|s_t),{\pi }_{{\theta }_{\text{ref}}}(a_t|s_t)]-{\beta }_2\sum _{a_t}{\pi }_{\theta }(a_t|s_t)\ln {\pi }_{\theta }(a_t|s_t)\} \\ \text{\hspace{1em}(11)} \end{gathered}$$

4.1 熵正则的计算

​ 令$l=[l_1,...,l_{|v|}]$表示未经过softmax的logit。那么词表中第$i$个token的概率为
$$p_i=\frac{e^{l_i}}{{\sum }_{j=1}^{|v|}{e}^{l_j}}$$
两边取对数，有
$$\log p_i=l_i-\log \sum _{j=1}^{|v|}{e}^{l_j}$$
那么信息熵为
$$\begin{gathered} -\sum _{i=1}^{|v|}{p}_i\log p_i=-\sum _{i=1}^{|v|}{p}_i(l_i-\log \sum _{j=1}^{|v|}{e}^{l_j})=\log \sum _{j=1}^{|v|}{e}^{l_j}-\sum _{i=1}^{|v|}{p}_i{l}_i \\ \text{\hspace{1em}(12)} \end{gathered}$$
verl中的熵正则计算代码为

# verl/utils/torch_functional.py
def entropy_from_logits(logits: torch.Tensor):
    """logits: [bs, response_length, vocab_size]"""
    pd = torch.nn.functional.softmax(logits, dim=-1)
    entropy = torch.logsumexp(logits, dim=-1) - torch.sum(pd * logits, dim=-1)
    return entropy

---

5. Critic模型loss的计算

​ Critic模型的作用是预测状态价值函数$V$的取值。如果有明确的标签，则使用MSE计算回归loss即可。

​ 但是实际计算中并没有这个标签。因此，在Actor-Critic范式中使用包含了真实奖励的折扣回报作为标签。令$V_t$是Critic模型直接针对第$t$个token预测的价值，$R_t$则是通过优势计算出包含真实奖励的折扣回报，那么损失函数为
$$\begin{gathered} L_{\text{PPO}}^{\text{Critic}}(\theta )={\text{E}}_{x\sim p,y\sim \pi (\cdot |x)}\frac{1}{|y|}\sum _{t=1}^{|y|}\frac{1}{2}\max ((V_t-R_t{)}^2,(\text{clip}(V_t)-R_t{)}^2) \\ \text{\hspace{1em}(13)} \end{gathered}$$
这里也对$V_t$进行了clip。下面是verl中的实现

# verl/trainer/ppo/core_algos.py
def compute_value_loss(vpreds, returns, values, eos_mask, cliprange_value):
    """
    vpreds:(`batch_size`, `response_length`)
    values:(`batch_size`, `response_length`)
    returns:(`batch_size`, `response_length`)
    """
    vpredclipped = verl_F.clip_by_value(vpreds, values - cliprange_value, values + cliprange_value)
    vf_losses1 = (vpreds - returns)**2
    vf_losses2 = (vpredclipped - returns)**2
    vf_loss = 0.5 * verl_F.masked_mean(torch.max(vf_losses1, vf_losses2), eos_mask)
    vf_clipfrac = verl_F.masked_mean(torch.gt(vf_losses2, vf_losses1).float(), eos_mask)
    return vf_loss, vf_clipfrac

---

6. verl中PPO的整体流程

​ 通过代码来看PPO的计算流程并不直观，为了更直观的理解绘制了下图。绿色的节点代表模型，红色的菱形代码处理逻辑，虚线的椭圆形代表输出。

---

二、GRPO

论文：DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models

1. 原理

​ PPO中需要额外训练一个值函数(Critic模型)，这会带来很大的显存需求和计算负担。GRPO的目标就是去掉Critic模型。回忆一下，PPO中Critic模型主要是用于估计优势函数中的$V$，从而进一步计算优势$A_t=Q(s_t,a_t)-V(s_t)$。

​ 在优势计算中之所以减去$V(s_t)$，是因为直接利用真实奖励蒙特卡洛估计$Q(s_t,a_t)$会导致高方差，而$V(s_t)$作为$Q(s_t,a_t)$的期望，充当了baseline。那么，除了减去baseline实现降低方差外，也可以通过多次采样求均值的方式来减少方差。因此GRPO相较于PPO的主要改进为：(1) 单个样本多次采取求均值减少方差；(2) 优势计算中去掉了对值函数的依赖。

​ 具体来说，损失函数更新为
L GRPO ( θ ) = − E x ∼ p , { y } i = 1 G ∼ π ( ⋅ ∣ x ) 1 G ∑ i = 1 G 1 ∣ y i ∣ ∑ t = 1 ∣ y i ∣ { min ⁡ [ r i , t ( θ ) A i , t , clip ( r i , t ( θ ) , 1 − ε , 1 + ε ) A i , t ] − β KL [ π θ old ( a i , t ∣ s i , t ) , π θ ref ( a i , t ∣ s i , t ) ] } (14) L_{\text{GRPO}}(\theta)=-\text{E}_{x\sim p,\{y\}_{i=1}^G\sim\pi(\cdot|x)}\frac{1}{G}\sum_{i=1}^G\frac{1}{|y_i|}\sum_{t=1}^{|y_i|}\Big\{\min[r_{i,t}(\theta)A_{i,t},\text{clip}(r_{i,t}(\theta),1-\varepsilon,1+\varepsilon)A_{i,t}]\\ -\beta\text{KL}[\pi_{\theta_{\text{old}}}(a_{i,t}|s_{i,t}),\pi_{\theta_{\text{ref}}}(a_{i,t}|s_{i,t})]\Big\} \tag{14}\\ LGRPO​(θ)=−Ex∼p,{y}i=1G​∼π(⋅∣x)​G1​i=1∑G​∣yi​∣1​t=1∑∣yi​∣​{min[ri,t​(θ)Ai,t​,clip(ri,t​(θ),1−ε,1+ε)Ai,t​]−βKL[πθold​​(ai,t​∣si,t​),πθref​​(ai,t​∣si,t​)]}(14)
其中重要性采样为$r_{i,t}(\theta )=\frac{{\pi }_{\theta }(a_{i,t}|s_{i,t})}{{\pi }_{{\theta }_{\text{old}}}(a_{i,t}|s_{i,t})}$。在GRPO的论文中，KL散度指定使用$k3$估计器。

​ 此外，由于单个样本$x$会进行G次采样，则会得到一组输出 { o 1 , o 2 , … , o G } \{o_1,o_2,\dots,o_G\} {o1​,o2​,…,oG​}，对应的得到一组奖励 r = { r 1 , r 2 , … , r G } r=\{r_1,r_2,\dots,r_G\} r={r1​,r2​,…,rG​}。GRPO将normalize后的奖励作为优势函数的取值
$$\begin{gathered} A_{i,t}={\tilde{r}}_i=\frac{r_i-\text{mean}(r)}{\text{std}(r)} \\ \text{\hspace{1em}(15)} \end{gathered}$$

---

2. 实现

​ 优势计算的代码如下：

# verl/trainer/ppo/core_algos.py
def compute_grpo_outcome_advantage(token_level_rewards: torch.Tensor,
                                   eos_mask: torch.Tensor,
                                   index: torch.Tensor,
                                   epsilon: float = 1e-6):
    """
    token_level_rewards: (bs, response_length)
    eos_mask:  (bs, response_length)
    """
    response_length = token_level_rewards.shape[-1]
    scores = token_level_rewards.sum(dim=-1)
    
    # id2score的value是一个列表，其包含了G个奖励值，key则是为每个输入x生成的id
    id2score = defaultdict(list)
    # id2mean和id2std分别存储每组的均值和标准差
    id2mean = {}
    id2std = {}

    with torch.no_grad():
        bsz = scores.shape[0]
        for i in range(bsz):
            id2score[index[i]].append(scores[i])
        for idx in id2score:
            if len(id2score[idx]) == 1:
                id2mean[idx] = torch.tensor(0.0)
                id2std[idx] = torch.tensor(1.0)
            elif len(id2score[idx]) > 1:
                id2mean[idx] = torch.mean(torch.tensor(id2score[idx]))
                id2std[idx] = torch.std(torch.tensor([id2score[idx]]))
            else:
                raise ValueError(f"no score in prompt index: {idx}")
        breakpoint()
        for i in range(bsz):
            # normalize
            scores[i] = (scores[i] - id2mean[index[i]]) / (id2std[index[i]] + epsilon)
        scores = scores.unsqueeze(-1).tile([1, response_length]) * eos_mask

    return scores, scores

​ 除了优势计算需要重写外，整个loss的计算直接复用PPO的即可，只要指定每个输入$x$的重复采样次数actor_rollout_ref.rollout.n即可。

---

三、REINFORCE++

论文：REINFORCE++: A Simple and Efficient Approach for Aligning Large Language Models

1. 原理

​ REINFORCE++的动机仍然是去掉PPO中的Critic模型。相比于GRPO采用组内平均的方式减少方差，REINFORCE++则更加直接，其在保留了PPO绝大数trick的情况下直接去掉了baseline。

​ 简单来说，去掉优势函数中的baseline $V(s_t)$，那么优势为$A_t=Q(s_t,a_t)$。采用蒙特卡洛估计来计算$Q(s_t,a_t)$，即使用折扣回报$G_t={\sum }_{k=t+1}^T{\gamma }^{k-t}{r}_k$作为优势$A_t$的估计值。此外，为了保持稳定的梯度，再对计算出来的优势进行normalize
$$\begin{gathered} {\tilde{A}}_t=\frac{A_t-{\mu }_A}{{\sigma }_A} \\ \text{\hspace{1em}(16)} \end{gathered}$$
其中，${\mu }_A$和${\sigma }_A$是整个batch优势的均值和标准差。

---

2. 实现

​ verl中优势和回报的计算

# verl/trainer/ppo/core_algos.py
def compute_reinforce_plus_plus_outcome_advantage(
    token_level_rewards: torch.Tensor, 
    eos_mask: torch.Tensor,
    gamma: torch.Tensor):
    """
    token_level_rewards:(bs, response_length)
    eos_mask:(bs, response_length)
    """
    with torch.no_grad():
        returns = torch.zeros_like(token_level_rewards)
        running_return = 0

        for t in reversed(range(token_level_rewards.shape[1])):
            running_return = token_level_rewards[:, t] + gamma * running_return
            returns[:, t] = running_return
            # Reset after EOS
            running_return = running_return * eos_mask[:, t]

        advantages = verl_F.masked_whiten(returns, eos_mask)
        advantages = advantages * eos_mask

    return advantages, returns

​ 下面是normalize的代码，其中mean和var都是标量值

def masked_whiten(values, mask, shift_mean=True):
    # mean和var都是标量值
    mean, var = masked_mean(values, mask), masked_var(values, mask)
    whitened = (values - mean) * torch.rsqrt(var + 1e-8)
    if not shift_mean:
        whitened += mean
    return whitened

---

四、RLOO

论文：Back to Basics: Revisiting REINFORCE Style Optimization for Learning from Human Feedback in LLMs

1. 原理

​ RLOO同样是要去掉Critic模型，但采取了不同于GRPO的baseline估计方式。类似于GRPO，通常是单个样本进行多次采样。假设对样本$x$进行G次采样，则会得到一组输出 { o 1 , o 2 , … , o G } \{o_1,o_2,\dots,o_G\} {o1​,o2​,…,oG​}，对应的得到一组奖励 r = { r 1 , r 2 , … , r G } r=\{r_1,r_2,\dots,r_G\} r={r1​,r2​,…,rG​}。那么$[x;o_i]$的优势函数为
$$\begin{gathered} A_i=\frac{1}{G}\sum _{i=1}^G[r_i-\frac{1}{G-1}\sum _{j\ne i}{r}_j] \\ \text{\hspace{1em}(17)} \end{gathered}$$
这样通过采样一组输出后，将除自身以外的奖励均值作为baseline。

---

2. 实现

# verl/trainer/ppo/core_algos.py
def compute_rloo_outcome_advantage(token_level_rewards: torch.Tensor,
                                   eos_mask: torch.Tensor,
                                   index: torch.Tensor,
                                   epsilon: float = 1e-6):
    """
    token_level_rewards:(bs, response_length)
    eos_mask:(bs, response_length)
    """
    response_length = token_level_rewards.shape[-1]
    # (bs,)，聚合token_level_rewards形成每个轨迹的奖励
    scores = token_level_rewards.sum(dim=-1)

    id2score = defaultdict(list)
    id2mean = {} # 组内均值

    with torch.no_grad():
        bsz = scores.shape[0]
        for i in range(bsz):
            id2score[index[i]].append(scores[i])
        for idx in id2score:
            if len(id2score[idx]) == 1:
                id2mean[idx] = torch.tensor(0.0)
            elif len(id2score[idx]) > 1:
                id2mean[idx] = torch.mean(torch.tensor(id2score[idx]))
            else:
                raise ValueError(f"no score in prompt index: {idx}")
        breakpoint()
        for i in range(bsz):
            response_num = len(id2score[index[i]])
            if response_num > 1:
                scores[i] = scores[i] * response_num / (response_num - 1) - id2mean[index[i]] * response_num / (response_num - 1)
        scores = scores.unsqueeze(-1).tile([1, response_length]) * eos_mask

    return scores, scores

​ 乍看起来，这里的实现与论文中并不一致。这里简单推导一下实现和论文的关系
$$\begin{array}{rl}{A}_i& =\frac{G}{G-1}\cdot r_i-\frac{G}{G-1}\cdot \frac{1}{G}\sum _{j=1}^G{r}_j\\ & =\frac{G}{G-1}\cdot r_i-\frac{1}{G-1}\sum _{j=1}^G{r}_j\\ & =\frac{G-1}{G-1}\cdot r_i+\frac{1}{G-1}\cdot r_i-\frac{1}{G-1}(r_i+\sum _{j\ne i}{r}_j)\\ & =r_i-\frac{1}{G-1}\sum _{j\ne i}{r}_j\end{array}$$

参考资料

http://joschu.net/blog/kl-approx.html

https://arxiv.org/html/2501.03262v1

https://arxiv.org/pdf/2402.03300

https://arxiv.org/pdf/2402.14740

https://zhuanlan.zhihu.com/p/25208314999

https://zhuanlan.zhihu.com/p/675309680

https://zhuanlan.zhihu.com/p/675348061