Calibrated rollout allocation improves training efficiency in language models
Measure First, Then Allocate: Calibrated Rollout Budgets for MCP Agents in Reinforcement Learning
Machine LearningComputation and Language
Summary
Training language models that use tools can be very costly because they need many tries to learn from rewards. The authors show that when rewards are the same in a group of attempts, it doesn’t help the model learn which version is better. They develop a method called Discord to focus only on cases where the rewards disagree, making comparisons more meaningful. Based on this, the BRACE system predicts how useful new attempts will be and allocates resources more efficiently, improving training while using fewer new attempts overall.
What this means in practice
- •For machine learning engineers: Allocate interaction budget more efficiently when training language models that use tools in reinforcement learning settings.
- •For ai platform operators: Improve resource management and cost control when running large-scale model training involving feedback with correlated rewards.
Authors
Hieu Tran, Zonghai Yao, Sunjae Kwon, Hong Yu
Abstract
Training tool-using language models requires costly interaction, yet identical rewards within a rollout group provide no group-relative learning signal. We use reward disagreement to guide harness evaluation and the allocation of rollouts for policy training in Model Context Protocol (MCP) environments. We first introduce Discord, which tests harness edits using exact paired inference and decision-preserving early stopping: only disagreeing paired outcomes favor one harness over another. This insight motivates prospective allocation in BRACE (Bank-calibrated Rollout Allocation under Correlated Episodes). BRACE combines historical task evidence, uncertainty in task success rates, and a correlated reward model to forecast mixed-reward groups before allocating generation budget. Controlled harness experiments favor Discord, while live MCP experiments expose a retention-cost trade-off. At three Qwen3-VL model scales, BRACE exceeds the strongest plotted baseline peaks with fewer new training rollouts; historical-bank collection is counted separately. It also leads the compared methods on BFCL and NESTFUL without further training. Ablations probe the bank, estimator, and allocation score; prospective evaluation checks forecast calibration. These results support forecasting reward contrast before generation.