Contextual Rollout Bandits for Reinforcement Learning with Verifiable Rewards

TL;DR

This paper introduces a neural scheduling framework for reinforcement learning with verifiable rewards, treating rollout selection as a contextual bandit problem to improve efficiency and performance in language model training.

Contribution

It formulates rollout scheduling as a contextual bandit problem and proposes a neural framework for adaptive selection, enhancing RLVR training effectiveness.

Findings

Consistent performance improvements across six reasoning benchmarks.

Enhanced training efficiency with adaptive rollout reuse.

Theoretical guarantees with sublinear regret bounds.

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.