CoBA-RL: Capability-Oriented Budget Allocation for Reinforcement Learning in LLMs

TL;DR

This paper tackles the inefficiency of uniform rollout budgets in reinforcement learning for large language models by introducing CoBA-RL, a dynamic budgeting framework that adapts to the model’s evolving capabilities. It defines a Capability-Oriented Value Function, leveraging a Beta distribution whose shape parameters respond to global failure rate, and couples it with a Budget Saturation Factor to model diminishing returns, enabling principled budget allocation. The allocation is carried out with a heap-based greedy strategy that maximizes marginal gain, providing a near-optimal solution with scalable runtime. Experiments across multiple Qwen models and math benchmarks show CoBA-RL consistently outperforms static and heuristic baselines, demonstrating improved sample efficiency and exploration-exploitation balance with substantial practical impact for efficient LLM post-training.

Abstract

Reinforcement Learning with Verifiable Rewards (RLVR) has emerged as a key approach for enhancing LLM reasoning.However, standard frameworks like Group Relative Policy Optimization (GRPO) typically employ a uniform rollout budget, leading to resource inefficiency. Moreover, existing adaptive methods often rely on instance-level metrics, such as task pass rates, failing to capture the model's dynamic learning state. To address these limitations, we propose CoBA-RL, a reinforcement learning algorithm designed to adaptively allocate rollout budgets based on the model's evolving capability. Specifically, CoBA-RL utilizes a Capability-Oriented Value function to map tasks to their potential training gains and employs a heap-based greedy strategy to efficiently self-calibrate the distribution of computational resources to samples with high training value. Extensive experiments demonstrate that our approach effectively orchestrates the trade-off between exploration and exploitation, delivering consistent generalization improvements across multiple challenging benchmarks. These findings underscore that quantifying sample training value and optimizing budget allocation are pivotal for advancing LLM post-training efficiency.

Figures (9)

• Figure 1: Comparison between GRPO-based methods and CoBA-RL. (a) GRPO employs a uniform strategy independent of training progress. (b) CoBA-RL dynamically self-calibrates the allocation strategy throughout the training process. It autonomously directs the rollout budget toward instances with high training value, aligned with the model's evolving capability. In this visualization, $p_i$ denotes the pass rate corresponding to the task instance $x_i$.
• Figure 2: Overview of CoBA-RL. (a) The training pipeline of CoBA-RL. (b) The adaptive value function influenced by model capability. (c) Illustration of the heap-based greedy allocation.
• Figure 3: Performance comparison of different models and methods on the Olympiad benchmark (avg@16). The curves track the validation accuracy over training steps for GRPO, Knapsack, and CoBA-RL across varying model scales.
• Figure 4: Visualization of different budget allocation across varying model capabilities on Qwen2.5-7B-instruct.
• Figure 5: Evolution of $\alpha_t$ on Qwen2.5-7B-Instruct. Left: The "Exploit $\to$ Explore" strategy, where $\alpha_t$ exhibits a fluctuating downward trend. Right: The "Explore $\to$ Exploit" strategy, where $\alpha_t$ shows a fluctuating upward trend.

Theorems & Definitions (2)

• Definition 2.1: Global Capability
• Proposition 2.2