Spend Your Rollouts Where It Counts: Rollout Allocation for Group-Based RL Post-Training

📄 arXiv: 2605.26606v1 📥 PDF

作者: Woojeong Kim, Ziyi Yang, Jing Nathan Yan, Jialu Liu

分类: cs.LG, cs.AI

发布日期: 2026-05-26


💡 一句话要点

提出Pilot-Commit框架,通过预算感知的rollout分配,加速基于群组的RL后训练。

🎯 匹配领域: 支柱二:RL算法与架构 (RL & Architecture) 支柱九:具身大模型 (Embodied Foundation Models)

关键词: 强化学习 语言模型 后训练 rollout分配 预算感知

📋 核心要点

  1. 在线强化学习中,rollout生成是训练大型语言模型的主要计算瓶颈,尤其是在基于群组的策略优化方法中。
  2. Pilot-Commit框架通过pilot阶段评估prompt的信息量,然后将rollout预算分配给高价值的prompt,从而优化rollout分配。
  3. 实验表明,Pilot-Commit在数学推理任务上,能以更低的采样成本达到与基线相当的准确率,显著加速训练过程。

📝 摘要(中文)

强化学习(RL)是后训练大型语言模型的主流范式。然而,在在线、on-policy设置中,rollout生成占据了训练的主要计算成本。基于群组的策略优化方法计算每个prompt的多个rollout的优势,但它们不加区分地将预算分配给奖励分布崩溃的prompt,从而在可忽略的学习信号上浪费了昂贵的rollout。我们证明了基于群组的更新在奖励方差高的状态下最有效。由于策略在整个训练过程中不断演变,prompt的信息量必须在线估计,而不是预先计算,但穷尽地评估每个prompt在计算上是禁止的。我们引入了Pilot-Commit,这是一个用于基于群组的RL后训练的预算感知rollout分配框架。Pilot-Commit将prompt评估与利用分离:pilot阶段使用一部分预算来估计每个prompt的信息量,剩余的rollout被分配给高杠杆的prompt,而低信号的prompt被跳过。在多个数学推理基准测试和从1.5B到14B参数的模型规模上,Pilot-Commit以显著降低的采样成本匹配了基线的准确性,在累积rollout中,达到目标准确率比GRPO快1.9倍,比DAPO快4.0倍。

🔬 方法详解

问题定义:论文旨在解决基于群组的强化学习后训练中,rollout预算分配不合理的问题。现有方法如GRPO和DAPO,对所有prompt平均分配rollout,导致大量计算资源浪费在信息量低的prompt上,降低了训练效率。这些方法没有充分考虑prompt的奖励方差,而奖励方差是衡量prompt信息量的关键指标。

核心思路:论文的核心思路是根据prompt的信息量动态分配rollout预算。通过一个pilot阶段,使用少量rollout评估每个prompt的奖励方差,从而估计prompt的信息量。然后,将剩余的rollout预算优先分配给信息量高的prompt,跳过信息量低的prompt。这种预算感知的分配策略可以更有效地利用计算资源,加速训练过程。

技术框架:Pilot-Commit框架包含两个主要阶段:Pilot阶段和Commit阶段。在Pilot阶段,对每个prompt生成少量rollout,计算其奖励方差,并根据方差大小对prompt进行排序。在Commit阶段,根据Pilot阶段的排序结果,将剩余的rollout预算分配给排名靠前的prompt。对于排名靠后的prompt,则直接跳过,不进行rollout。整个框架是budget-aware的,可以根据总预算和每个阶段的预算比例进行调整。

关键创新:Pilot-Commit的关键创新在于将prompt评估与rollout利用解耦。传统方法在rollout过程中同时进行prompt评估和策略更新,效率较低。Pilot-Commit通过独立的Pilot阶段进行prompt评估,可以更准确地估计prompt的信息量,从而更有效地分配rollout预算。此外,Pilot-Commit是一种在线方法,可以根据策略的演变动态调整rollout分配策略。

关键设计:Pilot阶段的关键设计在于如何有效地估计prompt的奖励方差。论文采用了一种简单的方差计算方法,即计算每个prompt的rollout奖励的方差。Commit阶段的关键设计在于如何根据Pilot阶段的排序结果分配rollout预算。论文采用了一种简单的比例分配方法,即根据prompt的排名比例分配rollout预算。此外,论文还考虑了总预算和每个阶段的预算比例的设置,这些参数需要根据具体任务进行调整。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明,在多个数学推理基准测试和从1.5B到14B参数的模型规模上,Pilot-Commit以显著降低的采样成本匹配了基线的准确性,在累积rollout中,达到目标准确率比GRPO快1.9倍,比DAPO快4.0倍。这表明Pilot-Commit能够显著提高训练效率,降低计算成本。

🎯 应用场景

该研究成果可广泛应用于大型语言模型的后训练,尤其是在计算资源有限的情况下。通过更有效地利用rollout预算,可以加速模型的训练过程,提高模型的性能。此外,该方法还可以应用于其他强化学习任务,例如机器人控制和游戏AI。

📄 摘要(原文)

Reinforcement learning (RL) is the dominant paradigm for post-training large language models. However, in the online, on-policy setting, rollout generation dominates the computational cost of training. Group-based policy optimization methods compute advantages from multiple rollouts per prompt, yet they indiscriminately allocate budget to prompts with collapsed reward distributions, wasting expensive rollouts on negligible learning signals. We demonstrate that group-based updates are most effective in regimes of high reward variance. Since the policy evolves throughout training, prompt informativeness must be estimated online rather than precomputed, but exhaustively evaluating every prompt is computationally prohibitive. We introduce Pilot-Commit, a budget-aware rollout allocation framework for group-based RL post-training. Pilot-Commit decouples prompt evaluation from exploitation: a pilot stage estimates per-prompt informativeness using a fraction of the budget, and the remaining rollouts are allocated to high-leverage prompts while low-signal prompts are skipped. Across multiple math reasoning benchmarks and model scales from 1.5B to 14B parameters, Pilot-Commit matches baseline accuracy with significantly lower sampling costs, reaching target accuracy up to $1.9\times$ faster than GRPO and $4.0\times$ faster than DAPO in cumulative rollouts.