Reproduce “Adaptive Rollout Allocation for Online Reinforcement Learning with Verifiable Rewards” on OpenResearch

About this paper

Hieu Trung Nguyen, Bao Nguyen, Wenao Ma, Yuzhi Zhao, Ruifeng She, Viet Anh Nguyen

The Variance-Informed Predictive (VIP) allocation strategy optimizes online reinforcement learning for large language models (LLMs) with verifiable rewards by adaptively distributing computational rollouts based on dynamically predicted prompt success probabilities. This method consistently boosts performance in mathematical and tool-augmented reasoning tasks, yielding improvements such as +12.3 Pass@32 on AIME benchmarks for Qwen2.5-Math-1.5B with a minimal training overhead of under 1.2%.

Key insights

• The gradient variance for on-policy components in Dr. GRPO and RLOO is proportional to p(1-p), where p is the prompt success probability, justifying adaptive allocation based on prompt informativeness.
• A Gaussian Process effectively models the non-stationary nature of prompt success probabilities during RL training, providing robust and adaptively calibrated uncertainty estimates for resource allocation.
• Minimizing total gradient variance via an adaptive, variance-informed rollout allocation strategy leads to enhanced sampling efficiency and superior performance compared to uniform or heuristic methods.

Summary by alphaXiv.