Adaptive Rollout Allocation for Online Reinforcement Learning with Verifiable Rewards

AI summary

Overview

Research area: reinforcement learning with verifiable rewards (RLVR), group-based policy optimization, and sampling efficiency for language model post-training.
Technical level: Advanced.
Scope: The paper introduces VIP, a variance-informed predictive rollout allocation strategy that distributes a fixed rollout budget across prompts to minimize expected policy-gradient variance under a hard compute constraint.

What This Paper Is About

Group-based reinforcement learning methods such as GRPO, Dr. GRPO, and RLOO typically assign the same number of rollouts to every training prompt. This uniform allocation treats all prompts as equally informative, which can waste computation and slow training. The paper’s goal is to allocate rollouts adaptively across prompts in each mini-batch so that the policy update has lower gradient variance and the available sampling budget is used more efficiently.

Key Contributions

1. Gradient variance analysis: The paper derives the relationship between per-prompt gradient variance and the prompt’s success probability for prominent group-based RL methods, including Dr. GRPO and RLOO.
2. Variance prediction: VIP uses a Gaussian process over prompt embeddings to predict the probability that the current model solves each prompt, with recursive Bayesian updates that leverage past rollout outcomes and prompt similarity.
3. Variance-minimizing rollout allocation: Given predicted variances, VIP formulates a convex optimization problem that minimizes gradient variance subject to a rollout budget constraint. It provides an exact solution to the continuous relaxation and a greedy rounding heuristic for feasible integer allocations.
4. Empirical validation: VIP is integrated on top of Dr. GRPO and RLOO and evaluated on mathematical reasoning and tool-augmented reasoning tasks, showing improvements over uniform or heuristic allocation strategies.

Main Findings

• Consistent improvements in mathematical reasoning: Adding VIP yields consistent improvements on Pass@32 and Mean@32 across all three base models and both rollout budgets tested.
• Large gain for a smaller model: On Qwen2.5-Math-1.5B at 8×Q, RLOO+VIP improves Pass@32 by +12.3 and Mean@32 by +6.3 points over RLOO.
• Stronger relative benefit for weaker backbones: The relative performance gain from VIP is larger for the 1.5B and 3B models than for the 7B model, suggesting VIP’s budget-aware variance reduction particularly helps weaker backbones that otherwise underutilize the rollout budget.
• Tool-augmented reasoning improvements: On Bamboogle and MuSiQue, VIP improves over base methods. For example, Dr. GRPO+VIP reaches EM 23.2 on Bamboogle versus Dr. GRPO’s 20, and EM 10.5 on MuSiQue versus Dr. GRPO’s 6. RLOO+VIP reaches EM 17.6 on Bamboogle versus RLOO’s 10.4, and EM 11 on MuSiQue versus RLOO’s 8.5.
• Ablation study indicated: The paper includes an ablation study on AIME24 and AIME25, but the specific numerical values are not included in the provided content.

Methodology in Plain English

At each training iteration, VIP takes a mini-batch of prompts. It embeds each prompt into a 384-dimensional vector using all-MiniLM-L6-v2. A Gaussian process with an RBF kernel and a sigmoid link function predicts each prompt’s success probability. After rollouts are generated, the observed rewards are clipped and converted into latent observations, and the Gaussian process posterior is updated recursively for the next iteration.

The predicted success probabilities are translated into per-prompt variance estimates using the theoretical formulas for Dr. GRPO and RLOO. VIP then solves a convex optimization problem to decide how many rollouts each prompt should receive, subject to a total rollout budget C and lower and upper bounds L and U on per-prompt rollouts. The continuous solution is found efficiently by bisection, and a greedy heuristic rounds the allocation to integers. VIP is integrated with RLOO and Dr. GRPO. Models are trained for two epochs under two total rollout budgets: 8×Q and 16×Q, where Q is the dataset size. Mathematical reasoning experiments train on DAPO-MATH-17k and evaluate on AIME2024 and AIME2025 using Pass@32, Mean@32, and Maj@32. Tool-augmented experiments use Bamboogle and MuSiQue with EM, F1@5, and Precision@5. The implementation uses VERL for math experiments and follows prior work for tool-augmented tasks.

Why This Matters

Impact on research: The paper provides a principled framework for improving sampling efficiency in RLVR, linking success probability to gradient variance and enabling adaptive budget allocation instead of fixed rollout counts.

Real-world applications:

• Efficient post-training of language models for mathematical reasoning.
• Tool-augmented reasoning and retrieval-augmented question answering, as tested on Bamboogle and MuSiQue.
• Compute-constrained RL training pipelines where rollout generation is a bottleneck.
• Improving smaller or weaker language models that otherwise underutilize rollout budgets.

Industry relevance: VIP can reduce the computational cost of RL-based post-training by allocating rollouts where they most reduce gradient variance. This is directly relevant to organizations training large language models under memory and compute limits.

Future Directions

• Extend the allocation framework to additional group-based RL methods and variants beyond Dr. GRPO and RLOO.
• Relax the paper’s assumptions, including the zero KL regularization setting, uncorrelated reward and gradient assumptions, and equal variance of the projected gradient across prompts.
• Improve or scale the Gaussian process predictor, for example through alternative probability models or more efficient approximations for large prompt sets.
• Evaluate VIP on a broader range of reasoning tasks, model sizes, and budget regimes, and study how the lower and upper rollout bounds should be set.

Target Audience

This paper benefits machine learning researchers working on reinforcement learning for language models, LLM post-training practitioners, and engineers optimizing compute budgets for RLVR. It is also suitable for advanced students familiar with policy optimization, Gaussian processes, and convex optimization.

Authors’ 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 VIP, 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, VIP 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 VIP consistently improves sampling efficiency and achieves higher performance than uniform or heuristic allocation strategies in multiple benchmarks.