最近正在入门大模型 RL 微调,阅读 VeRL 框架代码。在这篇博客中,我们从 VeRL 源码入手,解读一下 ppo 和 grpo 算法的 advantage 计算方法:
# 来自 verl/trainer/ppo/core_algos.py@register_adv_est(AdvantageEstimator.GAE)# or simply: @register_adv_est("gae")defcompute_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) """withtorch.no_grad():nextvalues=0lastgaelam=0advantages_reversed=[]gen_len=token_level_rewards.shape[-1]fortinreversed(range(gen_len)):delta=token_level_rewards[:,t]+gamma*nextvalues-values[:,t]lastgaelam_=delta+gamma*lam*lastgaelam# skip values and TD-error on observation tokensnextvalues=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)returnadvantages,returns# NOTE(sgm): this implementation only consider outcome supervision, where the reward is a scalar.@register_adv_est(AdvantageEstimator.GRPO)# or simply: @register_adv_est("grpo")defcompute_grpo_outcome_advantage(token_level_rewards:torch.Tensor,response_mask:torch.Tensor,index:np.ndarray,epsilon:float=1e-6,norm_adv_by_std_in_grpo:bool=True,config:Optional[AlgoConfig]=None,)->tuple[torch.Tensor,torch.Tensor]:""" Compute advantage for GRPO, operating only on Outcome reward (with only one scalar reward for each response). Args: token_level_rewards: `(torch.Tensor)` shape is (bs, response_length) response_mask: `(torch.Tensor)` shape is (bs, response_length) index: `(np.ndarray)` index array for grouping epsilon: `(float)` small value to avoid division by zero norm_adv_by_std_in_grpo: `(bool)` whether to scale the GRPO advantage config: `(Optional[AlgoConfig])` algorithm configuration object Note: If norm_adv_by_std_in_grpo is True, the advantage is scaled by the std, as in the original GRPO. If False, the advantage is not scaled, as in Dr.GRPO (https://arxiv.org/abs/2503.20783). Returns: advantages: `(torch.Tensor)` shape is (bs, response_length) Returns: `(torch.Tensor)` shape is (bs, response_length) """scores=token_level_rewards.sum(dim=-1)id2score=defaultdict(list)id2mean={}id2std={}withtorch.no_grad():bsz=scores.shape[0]foriinrange(bsz):id2score[index[i]].append(scores[i])foridxinid2score:iflen(id2score[idx])==1:id2mean[idx]=torch.tensor(0.0)id2std[idx]=torch.tensor(1.0)eliflen(id2score[idx])>1:scores_tensor=torch.stack(id2score[idx])id2mean[idx]=torch.mean(scores_tensor)id2std[idx]=torch.std(scores_tensor)else:raiseValueError(f"no score in prompt index:{idx}")foriinrange(bsz):ifnorm_adv_by_std_in_grpo:scores[i]=(scores[i]-id2mean[index[i]])/(id2std[index[i]]+epsilon)else:scores[i]=scores[i]-id2mean[index[i]]scores=scores.unsqueeze(-1)*response_maskreturnscores,scores以上的两段代码来自 VeRL 框架,分别实现了PPO 中的 GAE(广义优势估计)和GRPO(组相对策略优化)的 advantage 计算方法。下面,我们逐一解读,并对比它们的设计思路。
文章目录
- 1 PPO + GAE 优势计算 (`compute_gae_advantage_return`)
- 1.1 函数签名
- 1.2 核心逻辑(GAE 原理)
- 2 GRPO 优势计算 (`compute_grpo_outcome_advantage`)
- 2.1 函数签名
- 2.2 核心逻辑(GRPO 原理)
- 3 PPO(GAE) 与 GRPO 优势计算对比
1 PPO + GAE 优势计算 (compute_gae_advantage_return)
1.1 函数签名
defcompute_gae_advantage_return(token_level_rewards:torch.Tensor,values:torch.Tensor,response_mask:torch.Tensor,gamma:torch.Tensor,lam:torch.Tensor,):- token_level_rewards: 形状
(batch_size, response_length),每个 token 的即时奖励,通常由奖励模型或规则给出。如果是 LLM 做数学题,整个回答的 reward 可能只会给在最后一个 token。 - values: 形状与 reward 相同,Critic 网络输出的每个 token 的状态价值估计。
- response_mask: 形状相同,标记哪些 token 属于有效响应([EOS] 之后的 token 被 mask 掉,值为0)。
- gamma: 折扣因子,控制未来奖励的重要性。
- lam: GAE 中的 λ 参数,平衡偏差与方差。
1.2 核心逻辑(GAE 原理)
GAE 的递推公式为:
δ t = r t + γ ⋅ V ( s t + 1 ) − V ( s t ) , A t = δ t + γ ⋅ λ ⋅ A t + 1 δ_t = r_t + γ·V(s_{t+1}) - V(s_t), \quad A_t = δ_t + γ·λ·A_{t+1}δt=rt+γ⋅V(st+1)−V(st),At=δt+γ⋅λ⋅At+1
其中A t A_tAt是优势,V VV是价值函数。δ t = Q ( s t , a t ) − V ( s t ) \delta_t = Q(s_t, a_t) - V(s_t)δt=Q(st,at)−V(st),而 advantage 的这种计算方法,相当于不仅考虑当前时刻的 advantage,还以 γ 的 discount factor 考虑未来的 advantage。
代码中,从最后一个 token 反向计算,避免了显式循环的多次张量操作,同时利用response_mask处理变长序列。
nextvalues,lastgaelam=0,0advantages_reversed=[]gen_len=token_level_rewards.shape[-1]fortinreversed(range(gen_len)):# 计算 TD 误差 deltadelta=token_level_rewards[:,t]+gamma*nextvalues-values[:,t]# 递推计算优势lastgaelam_=delta+gamma*lam*lastgaelam# 用 mask 控制哪些位置的 value 和 advantage 要传播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# 计算 returns (用于价值网络更新)advantages=verl_F.masked_whiten(advantages,response_mask)# 对有效 token 做白化(标准化)returnadvantages,returns关键点:
- 使用
response_mask确保填充位置不参与计算,也不传播梯度。 masked_whiten对每个 batch 的有效 token 进行标准化(均值 0,方差 1),有助于稳定训练。- 最终返回的
advantages和returns形状均为(batch_size, response_length)。
2 GRPO 优势计算 (compute_grpo_outcome_advantage)
2.1 函数签名
defcompute_grpo_outcome_advantage(token_level_rewards:torch.Tensor,response_mask:torch.Tensor,index:np.ndarray,epsilon:float=1e-6,norm_adv_by_std_in_grpo:bool=True,config:Optional[AlgoConfig]=None,):- token_level_rewards: 形状
(batch_size, response_length),每个 token 的奖励。这里假设是结果监督(outcome supervision),即一个响应只对应一个标量奖励(如正确/错误),因此所有 token 的奖励和等于该响应的最终得分。 - response_mask: 同上,标记有效 token。
- index: 形状
(batch_size,),每个样本所属的组 ID(例如同一个 prompt 的不同采样)。GRPO 需要在组内进行标准化。 - norm_adv_by_std_in_grpo: 如果为 True,则按原始 GRPO 方式除以标准差;如果为 False,则只减均值(Dr.GRPO 风格)。
- epsilon: 防止除零。
2.2 核心逻辑(GRPO 原理)
GRPO 不使用 Critic 网络,而是用同一 prompt 下多个采样输出的平均奖励作为 baseline。advantage 就是组内标准化的奖励:
- 将 token 奖励求和,得到每个 response 的标量得分:
scores=token_level_rewards.sum(dim=-1)# shape: (batch_size,)- 按 index 分组,计算每个 group 的均值和标准差:
id2score=defaultdict(list)foriinrange(bsz):id2score[index[i]].append(scores[i])foridxinid2score:scores_tensor=torch.stack(id2score[idx])id2mean[idx]=torch.mean(scores_tensor)id2std[idx]=torch.std(scores_tensor)如果组内只有一个样本,均值为 0,标准差为 1(避免除零)。
- 根据均值和标准差,对每个得分进行标准化:
foriinrange(bsz):ifnorm_adv_by_std_in_grpo:scores[i]=(scores[i]-id2mean[index[i]])/(id2std[index[i]]+epsilon)else:scores[i]=scores[i]-id2mean[index[i]]标准化后,每个组内的平均优势为 0(可能方差为 1)。
- 每个 token 都被赋值为相同的标量 advantage(非响应位置仍为 0):
scores=scores.unsqueeze(-1)*response_mask# shape: (bs, response_length)returnscores,scores# 这里 returns 和 advantages 相同(GRPO 不更新价值网络)关键点:
- GRPO 完全抛弃了价值网络,baseline 来自同组的其他样本,避免了训练 Critic 的开销和不稳定性。
- 标准化操作可以在组内消除奖励尺度差异,使策略更新更稳健。
- 返回的两个张量相同,因为 GRPO 没有独立的 returns 概念(一般直接使用优势更新策略)。
3 PPO(GAE) 与 GRPO 优势计算对比
| 维度 | PPO + GAE | GRPO |
|---|---|---|
| baseline 来源 | Critic 网络估计的状态值V ( s t ) V(s_t)V(st) | 同一 prompt 下多个样本的平均奖励 |
| 优势形式 | 每个时间步t tt都有一个 advantageA t A_tAt | 整个响应共用一个标量 advantage(广播到每个 token) |
| 是否需要价值网络 | 是 | 否 |
| 处理变长序列 | 通过 mask 控制,逐 token 计算 | 汇总为标量后广播,mask 仅用于置零 |
| 计算复杂度 | 较高(需维护价值网络并做反向传播) | 较低(仅需组内统计) |
| 适用场景 | 一般 RLHF,需要细粒度信用分配 | 结果监督任务(如数学推理、代码生成),可节省资源 |
总结:
- PPO-GAE代码是标准的 GAE 实现,考虑了 mask 和最终的 whitening,适用于需要 Critic 的在线策略优化。
- GRPO代码则利用了结果监督和分组归一化的思想,简化了优势计算,适合大模型微调中每个响应只有一个标量奖励的场景(如 RLHF 中的“正确/错误”信号)。
这两段代码清晰体现了两种算法在设计哲学上的差异:PPO 追求细粒度时序信用分配,而 GRPO 通过组内比较降低方差并简化实现。