CoBA-RL：面向能力的强化学习 Rollout 预算分配 CoBA-RL: Capability-Oriented Budget Allocation for Reinforcement Learning in LLMs
用 Beta 分布动态建模模型能力，自适应分配 RL rollout 预算以优化探索-利用权衡
前置知识
GRPO (Group Relative Policy Optimization)
GRPO 是 DeepSeek 提出的一种强化学习算法，用于提升大语言模型的推理能力。其核心思想是：对每个 prompt 生成一组（group）G 个 rollout 轨迹，利用组内相对优势（group-relative advantage）来更新策略，从而避免了传统 PPO 中需要单独训练一个 critic/value network 的开销。每个 prompt 获得固定数量 G 的 rollout 预算，通过二值奖励函数（正确/错误）计算 pass rate 和 advantage。
CoBA-RL 正是在 GRPO 框架上进行改进的，理解 GRPO 的均匀预算分配机制是理解本文动机和方法的前提。
Beta 分布
Beta 分布是定义在 [0,1] 区间上的连续概率分布，由两个形状参数 α 和 β 控制。当 α > β 时分布向右偏（偏好高值），当 α < β 时向左偏（偏好低值），α = β 时对称。Beta 分布常用于建模概率的不确定性，在贝叶斯推断中是伯努利分布的共轭先验。本文用 Beta 分布来建模不同 pass rate 区间的偏好密度。
CoBA-RL 的核心创新之一就是用 Beta 分布来建模 Capability-Induced Preference Density，理解其形状参数如何影响分布形态对理解方法至关重要。
探索-利用权衡 (Exploration-Exploitation Trade-off)
这是强化学习中的经典问题：exploitation 指利用已知的最优策略获取最大奖励，exploration 指尝试新的、未知的策略以发现可能更好的选择。在 LLM 后训练中，exploitation 意味着巩固模型已掌握的简单题，exploration 意味着让模型尝试更多难题以拓展解题空间。过早探索会导致基础不牢，过晚探索则会错过提升机会。
本文的核心论点就是 GRPO 的均匀分配没有优化这个权衡，而 CoBA-RL 通过动态调整偏好密度来自动平衡探索和利用。
边际收益递减 (Diminishing Marginal Returns)
经济学和优化中的基本原理：随着某一资源投入量的增加，每额外增加一单位资源所带来的收益逐渐减少。在本文中，给一个任务实例增加 rollout 预算的边际收益是递减的——第128个 rollout 带来的训练价值远小于第1个。这个性质使得贪婪策略能够获得最优解。
作者严格证明了 CoBA-RL 的价值函数具有边际收益递减性质（Proposition 2.2），这是使用堆贪心策略获得最优分配解的理论基础。
Max-Heap（最大堆）
最大堆是一种特殊的完全二叉树数据结构，父节点的值始终大于等于子节点。它支持 O(log n) 的插入和删除最大值操作。在预算分配中，最大堆用于维护所有任务的边际增益，每次从堆顶取出增益最大的任务分配预算，再更新后重新插入，从而实现高效的贪心分配。
堆贪心策略是 CoBA-RL 高效实现的关键，相比动态规划基线实现了约 928 倍的加速。
研究动机
在使用强化学习（如 GRPO）提升大语言模型推理能力的训练过程中，一个关键但被忽视的问题是 rollout 预算的分配方式。标准 GRPO 对每个 prompt 都分配完全相同的 G 个 rollout，这种均匀分配策略存在严重的资源浪费。具体而言，当模型已经能够轻松解决某个简单题目时，分配 16 个 rollout 来计算 group-relative advantage 是冗余的；相反，对于模型完全不会的极难题目，16 个 rollout 可能不足以探索到正确的解题路径。现有的自适应方法（如 Knapsack-RL）虽然开始利用实例级指标（如历史 pass rate）来调整预算，但它们依赖的是静态价值函数——假设难样本始终比简单样本具有更高的训练价值，且这种关系在整个训练过程中不变。然而现实是：随着模型能力的提升，最有训练价值的样本集合是不断变化的。训练初期简单题最有价值（快速获得训练信号），训练后期难题更有价值（拓展解题空间）。静态方法无法捕捉这种动态变化，导致在错误的时间点将大量资源投入到错误难度的样本上。
本文的目标是本文的目标是设计一种能够根据模型实时能力动态调整 rollout 预算分配的强化学习算法。具体而言，CoBA-RL 要解决三个子目标：第一，定义一个能够量化每个任务实例在当前训练步骤下的训练价值的价值函数；第二，该价值函数的形状必须能够随着模型能力的演化而自动调整，而非预先固定；第三，在总预算约束下高效求解最优分配方案，且求解过程的计算开销必须足够小，能够无缝集成到在线 RL 训练循环中。最终目标是在多个数学推理基准上实现一致的性能提升，同时消耗更少或相同的计算资源。
与已有工作不同的是，本文的独特切入点在于一个关键观察：样本的训练价值不仅取决于样本本身的难度，更取决于模型当前的能力状态。这就像一个学生做练习题——同一道题在学生不同学习阶段的价值是不同的。已有的工作要么使用固定的难度-价值映射（如 Knapsack-RL 的背包问题公式化），要么使用预定义的退火策略（如线性衰减），都无法捕捉这种模型能力与样本价值之间的动态耦合关系。CoBA-RL 的独特视角是：通过监控整个训练 batch 的全局失败率来量化模型的全局能力，并将这种能力信号引入价值函数的参数化中，使得分配策略能够自校准。这种从样本视角到模型视角的转变——即从关注这个样本有多难转向关注这个样本对当前模型有多大训练价值——是本文相对于已有工作的根本性差异。
核心方法
CoBA-RL 的核心思想可以用一个直觉来理解：想象一个老师在给学生布置练习题，老师会根据学生当前的水平来动态调整题目分配——基础薄弱时多练简单题巩固基础，基础扎实后多安排难题拓展能力。技术路线上，CoBA-RL 在 GRPO 框架上增加了两个核心组件：第一，一个 Capability-Oriented Value Function（面向能力的价值函数），它将每个任务实例映射到一个标量训练价值，该价值同时取决于任务的 pass rate 和模型的当前能力；第二，一个 Heap-Based Greedy Budget Allocation（基于堆的贪心预算分配），它在总预算约束下高效求解使 batch 总价值最大化的分配方案。整个流程是：在每个训练步骤中，首先计算当前 batch 的全局失败率来估计模型能力，然后据此调整 Beta 分布的形状参数，接着计算每个实例的边际增益，最后通过堆贪心策略迭代分配预算。这个过程是自校准的——随着训练推进，模型能力提升，价值函数的偏好区域自动从高 pass rate 样本转移到低 pass rate 样本。
CoBA-RL 最核心的创新在于用一个参数化的能力感知价值函数取代了已有方法中的静态价值函数。具体而言，价值函数由两部分相乘组成：第一部分是 Budget Saturation Factor（预算饱和因子）$\eta(B_i, p_i) = 1 - e^{-\frac{B_i}{\tau p_i(1-p_i)}}$，它建模了增加预算的收益递减效应，其中 $p_i(1-p_i)$ 项确保 pass rate 接近 0 或 1 的样本（极端简单或极端困难）的价值增长缓慢；第二部分是 Capability-Induced Preference Density（能力诱导偏好密度），它是一个 Beta 分布 $p_i^{\alpha_t-1}(1-p_i)^{\beta_t-1}/B(\alpha_t, \beta_t)$，其参数 $\alpha_t$ 和 $\beta_t$ 由模型的全局失败率 $F_t$ 动态决定。当 $F_t$ 高时（模型能力弱），$\alpha_t$ 大、$\beta_t$ 小，偏好密度集中在高 pass rate 区间，即优先利用简单样本；当 $F_t$ 低时（模型能力强），$\alpha_t$ 小、$\beta_t$ 大，偏好密度转移到低 pass rate 区间，即转向探索难题。这种机制与已有方法的本质区别在于：Knapsack-RL 使用固定的 pass rate 作为价值指标，而 CoBA-RL 的什么样本有价值这一判断本身是随训练动态变化的。
方法步骤详情
CoBA-RL 的训练流程包含以下步骤：(1) 在每个训练步骤 t，从数据加载器中取出一批任务 $X_t = \{x_1, ..., x_M\}$；(2) 计算全局能力指标：首先对每个任务 $x_i$ 计算其 pass rate $p_i$，然后计算全局成功率 $S_t = \frac{1}{M}\sum_i p_i$ 和全局失败率 $F_t = 1 - S_t$；(3) 计算移动平均失败率 $\bar{F}_t$（过去 k 步的平均值以获得稳定估计），再通过非线性变换 $\Psi$ 得到 $\tilde{F}_t$：当 $\bar{F}_t > 0.5$ 时直接使用 $\bar{F}_t$，否则用 Sigmoid 函数 $\sigma(\gamma(\bar{F}_t - 0.5))$ 增强低失败率阶段的灵敏度，其中 $\gamma = 10$；(4) 根据 $\tilde{F}_t$ 线性映射得到 Beta 分布参数：$\alpha_t = \text{clip}(\alpha_{\min} + \lambda \cdot \tilde{F}_t, \alpha_{\min}, \alpha_{\max})$，$\beta_t = \kappa - \alpha_t$，其中 $\kappa = \alpha + \beta$ 是固定常数（论文中默认取 11）；(5) 初始化每个实例的预算为下界 $B_i = B_{\text{low}}$（默认为 2），计算剩余预算 $R = B_{\text{total}} - \sum_i B_{\text{low}}$；(6) 构建最大堆，对每个实例计算边际增益 $\Delta V_i = V(B_i+1, p_i) - V(B_i, p_i)$ 并压入堆中；(7) 迭代地从堆顶弹出边际增益最大的实例，为其分配一个单位预算，更新其边际增益后重新压入堆，直到剩余预算耗尽或所有实例达到上界 $B_{\text{up}}$（默认为 128）；(8) 根据分配结果重采样 batch，生成轨迹并计算 GRPO loss 更新策略。
技术新颖性
CoBA-RL 的技术新颖性体现在三个层面。第一，价值函数设计的自适应性：已有工作（Knapsack-RL）使用固定的 pass rate 或预定义的难度-价值映射，而 CoBA-RL 的 Beta 分布参数是模型能力的函数，使得价值函数的拓扑结构随训练动态变化。这种能力条件化（capability-conditioned）的价值建模在 LLM RL 领域是首次提出。第二，边际收益递减的理论保证：作者严格证明了价值函数的边际增益随预算递减（Proposition 2.2），证明过程简洁优雅——将价值函数分解为与 $B_i$ 相关和无关的部分后，边际增益简化为 $A \cdot e^{-kB_i}$ 的几何递减序列。这一性质保证了贪心策略的最优性，避免了使用计算昂贵的动态规划。第三，分配策略的高效性：通过堆数据结构，算法复杂度从动态规划的 $O(M \cdot B_{\text{total}} \cdot (B_{\text{up}} - B_{\text{low}}))$ 降低到 $O(B_{\text{total}} \log M)$，在 $M=512, B_{\text{total}}=8192$ 的设置下从 115 秒加速到 0.124 秒，加速比约 928 倍。这种效率使得方法能够无缝集成到在线训练循环中。
实验结果
CoBA-RL 在四个模型（Qwen2.5-7B-Instruct、Qwen2.5-7B-Base、Qwen3-4B-Base、Qwen3-1.7B-Base）和五个数学推理基准（AIME24、AIME25、AMC23、MATH500、OLYMPIAD）上进行了全面评估，所有结果均使用 avg@16 指标报告。在 Qwen2.5-7B-Instruct 上，CoBA-RL 的整体平均准确率达到 46.78%，相比 GRPO 基线的 42.24% 提升了 4.54 个百分点，也优于 Knapsack-RL 的 45.39%。最引人注目的提升出现在 AIME25 基准上，准确率从 GRPO 的 12.71% 提升到 18.33%（+5.62%），从 Knapsack-RL 的 15.21% 提升了 3.12%。在 Qwen2.5-7B-Base 上，CoBA-RL 在 OLYMPIAD 基准上达到 43.11%，超过 Knapsack-RL 的 41.33%（+1.78%）。在较小的 Qwen3-4B-Base 上，CoBA-RL 在 AMC23 数据集上取得了最显著的提升，从 GRPO 的 65.62% 提升到 72.34%（+6.72%）。在最小的 Qwen3-1.7B-Base 上，AIME24 从 8.96% 提升到 16.25%（+7.29%），提升幅度最大。数据效率分析显示，CoBA-RL 在 Btotal=2048 时达到 45.52% 准确率，不仅超过同预算下所有基线，甚至超过了 GRPO 在双倍预算 Btotal=4096 时的 42.78%，证明了能力感知分配的高效性。探索-利用顺序实验表明，先利用后探索策略（平均 46.78%）显著优于先探索后利用策略（平均 44.50%），在 AIME25 上差距尤为明显（18.33% vs 10.41%）。任务难度转换分析显示，CoBA-RL 在中等难度任务上的转换率达到 71.2%，远超 GRPO 的 46.8% 和 Knapsack-RL 的 50.0%；在困难任务上达到 36.7%，近乎翻倍于 GRPO 的 17.3%。与静态策略和启发式基线的对比进一步验证了动态分配的优势：固定参数策略（α=10.5, β=1.5）仅达 44.22%，线性步衰减策略达 45.39%，均低于 CoBA-RL 的 46.78%。对超参数 κ 的敏感性分析表明方法在 κ 取值 7、11、15、21 范围内表现稳定，准确率在 45.40% 到 46.61% 之间波动。计算效率方面，堆贪心策略仅需 0.124 秒完成分配，相比动态规划的 115.05 秒实现约 928 倍加速。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| AIME25 (Qwen2.5-7B-Instruct) | Accuracy avg@16 | 18.33% | GRPO 12.71% | +5.62% |
| AIME24 (Qwen3-1.7B-Base) | Accuracy avg@16 | 16.25% | GRPO 8.96% | +7.29% |
| AMC23 (Qwen3-4B-Base) | Accuracy avg@16 | 72.34% | GRPO 65.62% | +6.72% |
| OLYMPIAD (Qwen2.5-7B-Instruct) | Accuracy avg@16 | 43.19% | GRPO 37.68% | +5.51% |
| Overall Avg (Qwen2.5-7B-Instruct) | Average Accuracy | 46.78% | GRPO 42.24% | +4.54% |
| Overall Avg (Qwen3-4B-Base) | Average Accuracy | 49.46% | GRPO 44.72% | +4.74% |
| Overall Avg (Qwen3-1.7B-Base) | Average Accuracy | 35.75% | GRPO 31.61% | +4.14% |
| Data Efficiency (Qwen2.5-7B-Instruct) | Accuracy @ Btotal=2048 | 45.52% | GRPO @ Btotal=4096: 42.78% | +2.74% (half budget) |
局限与改进
尽管 CoBA-RL 取得了显著的实验结果，但本文存在几个值得关注的局限性。首先，实验仅在数学推理任务上进行验证，使用的数据集为 DAPO-Math-17K，评估基准也全部是数学竞赛题（AIME、AMC、MATH500、OLYMPIAD）。这使得方法在代码生成、agent 任务、通用推理等其他 RLVR 应用场景中的有效性尚不清楚。其次，模型规模的验证范围有限，最大的模型仅为 7B 参数，尚未在 70B+ 的大规模模型上进行实验，而这些模型恰恰是计算资源最紧张、预算分配优化价值最大的场景。第三，方法依赖于二值奖励函数 R(x, o) ∈ {0, 1}，对于需要过程奖励（process reward）或连续奖励信号的场景，价值函数的设计可能需要重新考虑。第四，全局失败率作为能力指标的假设在 batch 组成不均匀时可能存在偏差——如果某个 batch 恰好包含大量同质难度的任务，$\tilde{F}_t$ 的估计可能不够准确。最后，作者提到 Knapsack-RL 的官方实现未开源，因此是在 Verl 框架中重新实现的，这可能引入不公平比较的风险。
独立分析的弱点
从独立分析的角度看，CoBA-RL 存在以下几个可改进的弱点。第一，价值函数中 Beta 分布的参数化方式相对简单——通过线性映射将全局失败率映射到 α 参数。这种线性关系可能过于粗糙，尤其是在模型能力快速变化的训练阶段，线性映射可能无法精确捕捉最优的偏好密度迁移轨迹。一个可能的改进方向是学习一个神经网络来参数化这个映射关系，或者引入更丰富的局部能力信号（如不同难度区间的子集失败率）来替代单一的全局失败率。第二，方法假设每个训练 step 的 batch 能够准确反映模型的全局能力，但实际训练中 batch 组成是随机采样的，可能与整体数据分布存在偏差。可以考虑引入滑动窗口或多 batch 采样来获得更稳定的能力估计。第三，Budget Saturation Factor 中的温度系数 τ 是一个固定的超参数，而不同难度区间可能需要不同的饱和速度。例如，对于极难题目可能需要更慢的饱和（更多探索），对于简单题则可以快速饱和。第四，堆贪心策略虽然高效，但它是一种逐单位分配的方法，无法考虑任务之间的协同效应——某些任务的组合训练可能比单独训练的总和更有价值。
未来方向
基于 CoBA-RL 的成果，未来研究可以从以下几个方向延伸。第一，将方法推广到代码生成、agent 任务、多模态推理等更广泛的 RLVR 场景中，验证能力感知预算分配的通用性。第二，探索将 CoBA-RL 与 curriculum learning 方法结合——不仅动态分配 rollout 预算，还动态调整训练数据的采样分布，实现数据选择和资源分配的联合优化。第三，将能力感知的价值函数思想应用到其他需要动态资源分配的场景，如推理时的 token 预算分配（与 ROI-Reasoning 相关）、多任务学习中的计算资源分配等。第四，作者在结论中指出准确定义任务的训练价值并优化预算分配是推进 LLM 后训练效率的关键方向，这暗示了更精细的价值建模（如考虑轨迹多样性、奖励分布形状等）可能带来进一步提升。第五，探索在大规模模型（70B+）上的应用，这些模型的训练成本极高，能力感知的资源分配带来的效率提升可能更加显著。
复现评估
在复现性方面，CoBA-RL 的表现相对友好。代码已在 GitHub 开源（https://github.com/Within-yao/CoBA-RL），训练框架基于 Verl 和 SGLang，这两个都是活跃维护的开源项目。训练数据使用 DAPO-Math-17K，这是一个广泛使用的公开数学推理数据集。实验设置相对清晰：AdamW 优化器，学习率 1e-6，7B 模型 batch size 256 训练约 1000 步，小模型 batch size 512 训练约 500 步，最大响应长度 4096 tokens。关键超参数包括：每个实例预算范围 [2, 128]，总预算 Btotal 默认为 batch_size × default_group_size，κ 默认为 11，评估时温度 1.0、top-p 0.9、采样数 16。方法的计算开销很小（堆贪心策略仅需 0.124 秒），不会显著增加训练时间。需要注意的是，作者没有使用 KL 散度惩罚（βKL = 0），这与某些标准 GRPO 设置不同。整体而言，具备 RL 训练基础设施（多 GPU）的研究者应该能够在合理时间内复现主要结果。
论文图表
