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

Zhiyuan Yao; Yi-Kai Zhang; Yuxin Chen; Yueqing Sun; Zishan Xu; Yu Yang; Tianhao Hu; Qi Gu; Hui Su; Xunliang Cai

Source-linked AI summary

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

Zhiyuan Yao, Yi-Kai Zhang, Yuxin Chen, Yueqing Sun, Zishan Xu, Yu Yang, Tianhao Hu, Qi Gu, Hui Su, Xunliang Cai

arXiv:2602.03048v3Feb 3, 2026cs.LGcs.AI

TL;DR

CoBA-RL addresses rollout-budget allocation methods that overlook the model’s evolving capability and therefore fail to quantify changing sample training value. It uses a dynamically calibrated Capability-Oriented Value function and heap-based greedy allocation, with experiments reporting stronger performance and exploration-exploitation trade-offs across challenging mathematical benchmarks. The paper concludes that quantifying task training potential and optimizing budget allocation are central to efficient LLM post-training.

Problem
Existing rollout-budget methods rely on uniform, static, or pre-defined value functions that do not explicitly relate individual samples’ training value to the model’s evolving capabilities.
Method
CoBA-RL uses a dynamic Beta-distribution Capability-Oriented Value function and a heap-based greedy strategy to allocate rollout budgets by marginal training value.
Results
Across multiple challenging mathematical benchmarks and models, CoBA-RL significantly outperforms strong baselines and achieves a superior exploration-exploitation trade-off.
Takeaways & Limitations
The results support dynamically quantifying task training potential and optimizing rollout-budget allocation as a route toward more efficient LLM post-training.
Takeaways & Limitations
The allocation formulation assumes that a task’s value diminishes as the current policy masters it, with the saturation rate controlled by a temperature coefficient.

Abstract

from arXiv · show

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.

1. Introduction

CoBA-RL addresses the limitations of uniform or static rollout budgeting by adapting resource allocation to the model’s evolving capability and each task’s training value. It combines a capability-oriented value function with efficient greedy allocation, and experiments report improvements over GRPO and other strategies.

Motivation: GRPO assigns every prompt a uniform rollout budget, while existing adaptive methods often use static difficulty or pass-rate measures that assume harder samples consistently have greater training value.These approaches do not account for the policy model’s changing capabilities during training.
Motivation: As model capabilities evolve, the samples with the highest training value shift, requiring allocation to continually balance exploitation of successful instances with exploration of challenging queries.Exploitation consolidates mastery, whereas exploration samples diverse trajectories to expand the search space for potential solutions.
Approach: CoBA-RL models task training value with a Capability-Oriented Value function conditioned on evolving capability and dynamically calibrated using the current batch’s global failure rate.The function is modeled as a Beta distribution whose high-density regions shift between consolidating knowledge and exploring uncertainty.
Approach: A heap-based greedy strategy solves the constrained allocation problem by iteratively assigning budget to samples with the highest marginal gain, maximizing aggregate batch value.This operationalizes the value distribution into an efficient computational-resource allocation procedure.
Results: Extensive experiments across multiple models and mathematical benchmarks report that CoBA-RL significantly outperforms GRPO, static strategies, and heuristic strategies while improving the exploration-exploitation trade-off.The experiments use Qwen2.5-7B-Base, Qwen2.5-7B-Instruct, and Qwen3-1.7B/4B-Base models.

2. Method

CoBA-RL dynamically allocates rollout budgets according to the model’s evolving capability and each task’s training value. It combines a capability-aware value function with heap-based greedy optimization to balance resource allocation across samples.

Capability-Oriented Value Function: CoBA-RL quantifies task training value using a Capability-Oriented Value Function conditioned on the policy model’s evolving capability.The function uses global failure-rate statistics to self-calibrate its shape and identify changing sample preferences.
Problem Formulation and Optimization: The allocation problem maximizes aggregate learning potential under total-budget and per-task lower and upper constraints.The value function maps each task and allocated budget to expected learning potential.
Capability-Induced Preference Density: The capability-aware preference density models sampling preferences with a Beta distribution whose parameters evolve with training-step capability metrics.A moving average of global failure rates stabilizes the capability estimate, while a nonlinear transformation increases sensitivity during low-failure stages.
Budget Saturation Factor: The Budget Saturation Factor models diminishing returns as additional rollouts are assigned to a task.The temperature coefficient controls how quickly the value approaches saturation.
Efficient Allocation Optimization via Heap-Based Greedy Strategy: Because marginal value decreases monotonically with allocated budget, a heap-based greedy strategy repeatedly assigns one rollout to the task with the highest current marginal gain.The method updates each task’s marginal gain and reinserts it until the total budget is exhausted or upper bounds are reached.
Efficient Allocation Optimization via Heap-Based Greedy Strategy: The resulting allocation adapts simultaneously to the evolving value-function shape and sample-specific variation within each training batch.This dual-adaptive mechanism supports shifting resource distribution as policy preferences change.

3. Experiments

Experiments evaluate CoBA-RL across models, mathematical benchmarks, allocation strategies, exploration schedules, resource budgets, and computational implementations. CoBA-RL consistently improves accuracy and resource efficiency over GRPO, Knapsack-RL, static or heuristic allocation, while heap-based greedy allocation reduces computational cost.

Main Results: CoBA-RL achieves 46.78% average accuracy on Qwen2.5-7B-Instruct, exceeding GRPO by 4.54 percentage points and Knapsack-RL by 1.39 points.On Qwen2.5-7B-Base, it reaches 47.43% average accuracy; improvements over GRPO are 4.74% and 4.14% for Qwen3-4B-Base and Qwen3-1.7B-Base.
Main Results: AIME25 accuracy rises from 12.71% to 18.33% on Qwen2.5-7B-Instruct, exceeding GRPO by 5.62 percentage points and Knapsack-RL by 3.12 points.On AMC23 with Qwen3-4B-Base, CoBA-RL improves over GRPO by 6.72 percentage points.
Exploration and Exploitation: CoBA-RL’s dynamic allocation reaches 46.78% average accuracy, surpassing Linear Step Decay at 45.39% and the best Static baseline at 45.21%.The results indicate that fixed strategies are less aligned with changing learning dynamics than adaptive allocation.
Budget Constraints: With Btotal = 2048, CoBA-RL attains 45.52% accuracy, exceeding GRPO’s 42.78% with Btotal = 4096.The experiment reports consistent superiority over GRPO and Knapsack-RL across different total-budget constraints.
Computational Efficiency: Heap-based greedy allocation and dynamic programming produce identical allocations theoretically, but dynamic programming has pseudo-polynomial complexity O(M·Btotal·(Bup−Blow)).For M = 512 and Btotal = 8192, the dynamic-programming baseline requires approximately 115.05 seconds.

4. Related Work

Related work spans RLVR optimization, progressive training, budget allocation, and adaptive rollout methods. CoBA-RL differs by adaptively allocating varying rollout budgets according to evolving model capabilities rather than selecting samples or relying on static allocation values.

RLVR Optimization: GRPO is widely adopted for RLVR, while GSPO defines sample importance through sequence likelihood and DAPO introduces four training modifications.These methods extend RL-based reasoning optimization through different objectives or training mechanisms.
Progressive Training: Curriculum methods organize data into difficulty stages, with examples dynamically evaluating later batches or adjusting data distributions using policy-gradient advantages.The cited distinction is that curriculum learning focuses on selecting which samples to train.
Budget Allocation: Budget allocation research applies resource-distribution methods such as Multi-Armed Bandits to domains including advertising, marketing, and LLMs.This literature provides the broader operations-research context for rollout-budget allocation.
Adaptive Rollout Allocation: Dynamic rollout-budget allocation for exploration and exploitation remains challenging because approaches including GVM-RAFT and Knapsack-RL often rely on static or predetermined mechanisms.GVM-RAFT uses rejection sampling to reduce stochastic-gradient variance, whereas Knapsack-RL formulates batch-value maximization as a knapsack problem.

5. Conclusion

CoBA-RL dynamically adapts rollout budgets to evolving LLM capabilities using constrained optimization, a dynamic Beta-based value function, and heap-based greedy allocation. The paper concludes that this approach improves the exploration–exploitation trade-off and that accurately defining training potential and optimizing budget allocation are important directions for efficient LLM post-training.

Conclusion: CoBA-RL adapts rollout budgets to evolving LLM capabilities by formulating allocation as a constrained optimization problem.The method uses a dynamic Beta-distribution value function and heap-based greedy optimization of marginal gains.
Conclusion: Experiments show that CoBA-RL achieves a superior exploration–exploitation trade-off and significantly outperforms static and heuristic baselines.The analysis reports a typical transition from exploitation toward exploration during training.
Conclusion: Accurately defining task training potential and optimizing budget allocation are identified as critical directions for advancing efficient LLM post-training.This statement is presented as a future-oriented implication of the paper’s findings.

A. Proof of Proposition 3.2

The proof expresses marginal gain as the value increase from one additional rollout and establishes that these gains decrease strictly with allocated budget. This diminishing-return property supports efficient greedy allocation.

Under p_i ∈ (0, 1) and τ > 0, the value function can be rewritten using positive constants C and k.The derivation states that C > 0 and k > 0 in this non-trivial case.
Marginal gain ΔV(B_i, p_i) is defined as V(B_i + 1, πθ, p_i) − V(B_i, πθ, p_i).
The proof compares consecutive marginal gains to establish strict monotonic decrease as the allocated budget B_i increases.It examines the ratio between ΔV(B_i, p_i) and ΔV(B_i + 1, p_i).
The resulting marginal gain is a strictly decreasing geometric sequence with respect to B_i.

B. Implementation Details of Main Training Loop

The main training loop integrates CoBA-RL’s BudgetAllocator into GRPO by computing per-prompt rollout counts and resampling the batch before generation. The remaining GRPO loss calculation and update proceed afterward.

The BudgetAllocator is executed inside the GRPO training loop for each batch from the dataloader.
The allocator returns a dictionary mapping each prompt index to its assigned rollout count.
The batch is expanded by repeating each index according to its allocation count, while indices with zero counts are excluded.
The resampled batch is prepared before the generation phase, after which the GRPO loss calculation and update are performed.

C. Experiment Details

Experiments use multiple Qwen model scales within a Verl-based pipeline with SGLang inference and AdamW optimization. Training batch size and duration vary by model scale, while rollout methods use Clip-higher.

The training pipeline uses Verl with SGLang as the inference engine and AdamW with learning rate 1 × 10^-6.
Models smaller than 7B use global batch size M = 512 and approximately 500 training steps, whereas the 7B model uses M = 256 and nearly 1000 steps.
Table 5 provides the hyperparameters used for evaluation rollout.
Clip-higher is adopted for both Knapsack-RL and CoBA-RL following DAPO recommendations.

D.1. Analysis of Task Difficulty Transition

CoBA-RL produces the highest task-transition conversion rates across difficulty levels, including medium and extremely-hard tasks, while retaining more extremely-easy instances. The transition matrix reports percentages from initial to final status.

71.2% of medium tasks convert under CoBA-RL, versus 46.8% for GRPO and 50.0% for Knapsack-RL.
The reported results characterize CoBA-RL as aligning resources with model capability across the full difficulty spectrum.
Each Figure 7 cell gives the percentage of samples transitioning from an initial status to a final status.

D.2. Sensitivity Analysis of κ

The analysis examines how κ, the constant sum of Beta-distribution shape parameters, affects capability-oriented value sharpness and training performance. Performance is robust across κ settings, while baseline visualizations contrast static exploitation, static exploration, and rigid heuristic allocation.

κ = α + β controls Beta-distribution variance and the sharpness of the capability-oriented value function.
Accuracy ranges from 45.40% at κ = 7 to 46.61% at κ = 11, indicating minimal performance fluctuation across κ ∈ {7, 11, 15, 21}.
The authors interpret this stability as evidence that improvement stems from dynamic allocation rather than specific hyperparameter tuning.
κ = 11 is selected as the default configuration for the main experiments because it achieves the highest observed performance.
Static baselines use Beta distributions skewed toward high success rates for exploitation or low success rates for exploration, while Linear Step Decay anneals α from 10 to 1 using a rigid schedule.