Contextual Rollout Bandits for Reinforcement Learning with Verifiable Rewards
Abstract
Reinforcement Learning with Verifiable Rewards (RLVR) is an effective paradigm for improving the reasoning capabilities of large language models. However, existing RLVR methods utilize rollouts in an indiscriminate and short-horizon manner: responses of heterogeneous quality within each prompt are treated uniformly, and historical rollouts are discarded after a single use. This leads to noisy supervision, poor sample efficiency, and suboptimal policy updates. We address these issues by formulating rollout scheduling in RLVR as a contextual bandit problem and proposing a unified neural scheduling framework that adaptively selects high-value rollouts throughout training. Each rollout is treated as an arm whose reward is defined by the induced performance gain between consecutive optimization steps. The resulting scheduler supports both noise-aware intra-group selection and adaptive global reuse of historical rollouts within a single principled framework. We provide theoretical justification by deriving sublinear regret bounds and showing that enlarging the rollout buffer improves the achievable performance upper bound. Experiments on six mathematical reasoning benchmarks demonstrate consistent gains in performance and training efficiency across multiple RLVR optimization methods.
Lay Summary
Large language models can improve their reasoning ability by practicing on problems whose answers can be automatically checked, such as math questions. However, during this training process, the model often generates many answers of uneven quality: some are helpful, some are noisy, and some may even be misleading. Existing methods usually use these generated answers almost indiscriminately and discard them after one training step, which wastes computation and useful learning signals.We propose CBS, a lightweight training scheduler that learns which generated answers are most useful for improving the model. CBS treats each generated answer as a candidate to choose from, scores it using simple training signals such as reward, uncertainty, length, and reuse history, and then selects high-value answers for training. It can both filter noisy answers within the current batch and reuse helpful answers from previous batches. Experiments on mathematical reasoning tasks show that CBS improves both model accuracy and training efficiency across several popular reinforcement learning methods. This makes training reasoning models more data-efficient and less wasteful.