Adaptive Rollout Allocation for Online Reinforcement Learning with Verifiable Rewards

TL;DR

The paper tackles sampling efficiency in online reinforcement learning with verifiable rewards by introducing VIP, a variance-informed predictive rollout allocator. VIP combines a Gaussian-process-based predictor of per-prompt success with a convex optimization-based budget allocator to minimize the minibatch gradient variance under a fixed rollout budget. Theoretical analysis derives per-prompt gradient-variance expressions for GRPO and RLOO, linking variance to prompt success and rollout counts, which VIP exploits for adaptive budgeting. Empirical results across mathematical reasoning and tool-augmented tasks show consistent improvements over uniform or heuristic allocations, with low computational overhead, highlighting VIP as a practical, principled approach to resource-efficient RL training for language models.

Abstract

Sampling efficiency is a key bottleneck in reinforcement learning with verifiable rewards. Existing group-based policy optimization methods, such as GRPO, allocate a fixed number of rollouts for all training prompts. This uniform allocation implicitly treats all prompts as equally informative, and could lead to inefficient computational budget usage and impede training progress. We introduce \Ours, a Variance-Informed Predictive allocation strategy that allocates a given rollout budget to the prompts in the incumbent batch to minimize the expected gradient variance of the policy update. At each iteration, \Ours~uses a lightweight Gaussian process model to predict per-prompt success probabilities based on recent rollouts. These probability predictions are translated into variance estimates, which are then fed into a convex optimization problem to determine the optimal rollout allocations under a hard compute budget constraint. Empirical results show that \Ours~consistently improves sampling efficiency and achieves higher performance than uniform or heuristic allocation strategies in multiple benchmarks. Our code will be available at https://github.com/HieuNT91/VIP.

Figures (7)

• Figure 1: The process starts with an initial belief over prompt success probabilities. At each step $t$, a mini-batch $\mathcal{B}_t$ is selected, and the belief function $m_t(\cdot)$ predicts the success probabilities of the prompts in $\mathcal{B}_t$. A budget allocation module assigns rollout budgets $\{n_q\}$, rollouts are generated, and the resulting data updates the model and beliefs. Repeated for $T$ steps, this yields a fine-tuned model $\pi_{\theta_{T + 1}}$ with improved performance and efficient rollout usage.
• Figure 2: Prediction mean absolute error (MAE) over training steps for two model scales. Our GPR-based predictor achieves consistently lower MAE than moving average and Ridge Regression baselines for both the 1.5B and 7B models.
• Figure 3: Comparison of optimal rollout allocations produced by different heuristics versus our proposed variance-aware allocation strategy. The figure plots the optimal number of rollouts $n_i^\star$ against prompt difficulty $p_i$, highlighting how our method allocates budget differently from inverse-accuracy and inverse-variance baselines.
• Figure 4: Prompt template for mathematical reasoning
• Figure 5: Prompt template for tool augmented reasoning

Theorems & Definitions (10)

• Proposition 4.2: Dr. GRPO gradient variance
• Proposition 4.3: RLOO gradient variance
• Theorem 5.1: Continuous allocation, Dr. GRPO
• Theorem 5.2: Continuous allocation, RLOO
• proof : Proof of Proposition \ref{['prop:drgrpo_var']}
• proof : Proof of Proposition \ref{['prop:rloo_var']}
• proof : Proof of Theorem \ref{['thm:opt-drgrpo']}
• proof : Proof of Theorem \ref{['thm:opt-rloo']}
• Proposition E.1: Dr. GRPO gradient variance, continuous reward
• Proposition E.2: RLOO gradient variance, continuous reward