`core_algos.py` — 这是 PPO 算法的核心算法实现文件¶

文件概述¶

模块路径: verl.trainer.ppo.core_algos

这是 PPO 算法的核心算法实现文件，包含了所有关键的数学计算逻辑。它实现了： - 多种优势估计器（GAE、GRPO、REINFORCE++、RLOO 等） - KL 控制器（固定和自适应） - 策略损失函数注册表

这些函数是"纯计算"的，不依赖任何分布式框架，可以被不同的 trainer 实现（Ray、Megatron 等）调用。

在训练流程中的位置¶

在每个训练步骤中，ray_trainer.py 的 compute_advantage() 函数会调用本文件中的优势估计器。这些计算在驱动进程（driver process）上执行，是轻量级的数值计算。

关键代码讲解¶

1. 注册表模式¶

# 优势估计器注册表
ADV_ESTIMATOR_REGISTRY: dict[str, Any] = {}

def register_adv_est(name_or_enum):
    """装饰器：注册优势估计函数"""
    def decorator(fn):
        name = name_or_enum.value if isinstance(name_or_enum, Enum) else name_or_enum
        ADV_ESTIMATOR_REGISTRY[name] = fn
        return fn
    return decorator

def get_adv_estimator_fn(name_or_enum):
    """根据名称获取优势估计函数"""
    name = name_or_enum.value if isinstance(name_or_enum, Enum) else name_or_enum
    return ADV_ESTIMATOR_REGISTRY[name]

使用注册表模式使得添加新的优势估计器非常方便——只需要写一个函数并用 @register_adv_est 装饰即可。

同样，策略损失函数也有类似的注册表：

POLICY_LOSS_REGISTRY: dict[str, PolicyLossFn] = {}

def register_policy_loss(name):
    def decorator(func):
        POLICY_LOSS_REGISTRY[name] = func
        return func
    return decorator

2. AdvantageEstimator 枚举¶

class AdvantageEstimator(str, Enum):
    GAE = "gae"                                    # 广义优势估计（需要 Critic）
    GRPO = "grpo"                                  # 分组相对策略优化
    REINFORCE_PLUS_PLUS = "reinforce_plus_plus"    # REINFORCE++
    REINFORCE_PLUS_PLUS_BASELINE = "reinforce_plus_plus_baseline"
    REMAX = "remax"                                # ReMax
    RLOO = "rloo"                                  # RLOO
    OPO = "opo"                                    # OPO
    GRPO_PASSK = "grpo_passk"                      # Pass@k 版 GRPO
    GPG = "gpg"                                    # GPG
    RLOO_VECTORIZED = "rloo_vectorized"            # 向量化 RLOO
    GRPO_VECTORIZED = "grpo_vectorized"            # 向量化 GRPO
    OPTIMAL_TOKEN_BASELINE = "optimal_token_baseline"  # 最优 Token 基线
    TIR_OPTIMAL_TOKEN_BASELINE = "tir_optimal_token_baseline"  # 多轮 OTB

3. KL 控制器¶

自适应 KL 控制器¶

class AdaptiveKLController:
    """来自论文 https://arxiv.org/pdf/1909.08593.pdf"""
    def __init__(self, init_kl_coef, target_kl, horizon):
        self.value = init_kl_coef    # 初始 KL 系数
        self.target = target_kl      # 目标 KL 值
        self.horizon = horizon       # 调整速度参数

    def update(self, current_kl, n_steps):
        """根据当前 KL 和目标 KL 调整系数"""
        proportional_error = np.clip(current_kl / self.target - 1, -0.2, 0.2)
        mult = 1 + proportional_error * n_steps / self.horizon
        self.value *= mult

自适应控制器会自动调整 KL 惩罚系数：如果当前 KL 太大（策略偏离太远），就增大系数；如果 KL 太小，就减小系数。

固定 KL 控制器¶

class FixedKLController:
    def __init__(self, kl_coef):
        self.value = kl_coef

    def update(self, current_kl, n_steps):
        pass  # 不做任何调整

4. GAE 优势估计¶

@register_adv_est(AdvantageEstimator.GAE)
def compute_gae_advantage_return(token_level_rewards, values, response_mask, gamma, lam):
    """
    广义优势估计 (Generalized Advantage Estimation)

    公式: \(\hat{A}_t = \delta_t + (\gamma\lambda)\delta_{t+1} + (\gamma\lambda)^2\delta_{t+2} + \cdots\)
    其中: \(\delta_t = r_t + \gamma V(s_{t+1}) - V(s_t)\)
    """
    with torch.no_grad():
        nextvalues = 0
        lastgaelam = 0
        advantages_reversed = []
        gen_len = token_level_rewards.shape[-1]

        for t in reversed(range(gen_len)):
            # TD 误差: δ = r + γV(s') - V(s)
            delta = token_level_rewards[:, t] + gamma * nextvalues - values[:, t]
            # GAE 累积: A_t = δ_t + γλ * A_{t+1}
            lastgaelam_ = delta + gamma * lam * lastgaelam

            # 在非回复 token 上重置（跳过 padding）
            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

GAE 是最经典的优势估计方法： - gamma (\(\gamma\)) 是折扣因子，控制对未来奖励的重视程度 - lam (\(\lambda\)) 控制偏差-方差权衡：\(\lambda=0\) 是纯 TD，\(\lambda=1\) 是蒙特卡洛 - 白化（whiten）使优势的均值为 0、方差为 1，有助于训练稳定

5. GRPO 优势估计¶

@register_adv_est(AdvantageEstimator.GRPO)
def compute_grpo_outcome_advantage(token_level_rewards, response_mask, index, epsilon=1e-6,
                                    norm_adv_by_std_in_grpo=True, config=None):
    """
    GRPO: 不需要 Critic！直接用同组回复的奖励分数计算优势。

    \(A_i = \frac{r_i - \text{mean}(r_{\text{group}})}{\text{std}(r_{\text{group}})}\)
    """
    scores = token_level_rewards.sum(dim=-1)  # 每条回复的总分

    with torch.no_grad():
        # 按 prompt（index）分组
        id2score = defaultdict(list)
        for i in range(bsz):
            id2score[index[i]].append(scores[i])

        # 计算每组的均值和标准差
        for idx in id2score:
            if len(id2score[idx]) > 1:
                scores_tensor = torch.stack(id2score[idx])
                id2mean[idx] = torch.mean(scores_tensor)
                id2std[idx] = torch.std(scores_tensor)

        # 归一化：(分数 - 组均值) / 组标准差
        for i in range(bsz):
            if norm_adv_by_std_in_grpo:
                scores[i] = (scores[i] - id2mean[index[i]]) / (id2std[index[i]] + epsilon)
            else:
                scores[i] = scores[i] - id2mean[index[i]]  # Dr.GRPO 变体

        scores = scores.unsqueeze(-1) * response_mask

    return scores, scores

GRPO 的核心思想：不需要 Critic 网络。它对同一个 prompt 的多个回复进行分组，用组内的相对排名作为优势值。这大大简化了训练系统。

6. REINFORCE++ 优势估计¶

@register_adv_est(AdvantageEstimator.REINFORCE_PLUS_PLUS)
def compute_reinforce_plus_plus_outcome_advantage(token_level_rewards, response_mask, config=None):
    """
    REINFORCE++: 使用折扣回报 + 白化

    \(R_t = r_t + \gamma r_{t+1} + \gamma^2 r_{t+2} + \cdots\)
    \(A = \text{whiten}(R)\)
    """
    gamma = config.gamma
    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
            running_return = running_return * response_mask[:, t]  # EOS 后重置

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

    return advantages, returns

7. RLOO 优势估计¶

@register_adv_est(AdvantageEstimator.RLOO)
def compute_rloo_outcome_advantage(token_level_rewards, response_mask, index, ...):
    """
    RLOO (Leave-One-Out): 用组内其他回复的平均分作为基线

    基线 \(= \frac{\text{总分} - r_i}{N-1}\)
    \(A_i = r_i - \text{基线}\)
    """
    for i in range(bsz):
        response_num = len(id2score[index[i]])
        if response_num > 1:
            # 等价于：r_i * N/(N-1) - sum(r_group) / (N-1)
            scores[i] = scores[i] * response_num / (response_num - 1) - \
                         id2mean[index[i]] * response_num / (response_num - 1)

RLOO 是一个无偏估计器，它用"留一法"计算基线：每个样本的基线是其余样本的平均分。

核心类/函数列表¶

名称	类型	作用
`AdvantageEstimator`	Enum	优势估计器类型枚举
`register_adv_est()`	decorator	注册优势估计函数
`get_adv_estimator_fn()`	function	获取优势估计函数
`register_policy_loss()`	decorator	注册策略损失函数
`get_policy_loss_fn()`	function	获取策略损失函数
`AdaptiveKLController`	class	自适应 KL 控制器
`FixedKLController`	class	固定 KL 控制器
`get_kl_controller()`	function	KL 控制器工厂函数
`compute_gae_advantage_return()`	function	GAE 优势估计
`compute_grpo_outcome_advantage()`	function	GRPO 优势估计
`compute_reinforce_plus_plus_outcome_advantage()`	function	REINFORCE++ 优势估计
`compute_rloo_outcome_advantage()`	function	RLOO 优势估计
`compute_opo_outcome_advantage()`	function	OPO 优势估计
`compute_remax_outcome_advantage()`	function	ReMax 优势估计
`compute_gpg_outcome_advantage()`	function	GPG 优势估计
`compute_grpo_passk_outcome_advantage()`	function	Pass@k GRPO 优势估计
`compute_optimal_token_baseline_advantage()`	function	最优 Token 基线优势估计

优势估计器对比¶

估计器	需要 Critic	需要分组	特点
GAE	是	否	最经典，偏差-方差可控
GRPO	否	是	简单高效，无需 Critic
REINFORCE++	否	否	折扣回报 + 白化
RLOO	否	是	无偏留一法
OPO	否	是	长度加权基线
ReMax	否	否	需要额外 greedy 基线
OTB	否	是	逐 token 最优基线

数据流和调用关系¶

ray_trainer.py: compute_advantage()
    |
    +-- AdvantageEstimator.GAE --> compute_gae_advantage_return()
    |                                需要: token_level_rewards, values, response_mask
    |
    +-- AdvantageEstimator.GRPO --> compute_grpo_outcome_advantage()
    |                                需要: token_level_rewards, response_mask, uid(分组)
    |
    +-- 其他估计器 --> get_adv_estimator_fn() --> 对应函数

ray_trainer.py: apply_kl_penalty()
    |
    +-- get_kl_controller() --> AdaptiveKLController / FixedKLController
    +-- kl_penalty()        --> KL 散度计算

小结¶

core_algos.py 是 verl 的算法核心，实现了 RL 训练中最重要的数学计算。它的设计有以下优点：

注册表模式：添加新算法只需一个装饰器
纯计算：不依赖分布式框架，可在任何环境下运行
丰富的算法库：支持 10+ 种优势估计方法
数值稳定：使用 torch.no_grad() 和白化等技术确保计算稳定

core_algos.py — 这是 PPO 算法的核心算法实现文件¶