Bayes-Adaptive MDPs for Rollout Allocation to Enhance GRPO Training
GRPO spends the same rollout budget on every question, so hard questions stay under-sampled. A Bayes-adaptive policy allocates rollouts offline, with no extra online overhead.
The claim
A Bayes-adaptive policy spends rollouts where they teach GRPO the most, at no extra cost.
The problem
GRPO gives every training question the same rollout budget.
Question difficulty varies widely, so hard questions get too few rollouts.
Moderately difficult questions often receive more rollouts than they need.
How it works
The authors cast rollout allocation as a sequential decision problem.
A Bayes-adaptive Markov decision process learns the allocation policy.
That policy is computed offline before GRPO training begins.
GRPO then spends its fixed rollout budget adaptively per question.
No additional online optimization overhead is introduced during training.
Figures
What changes
Training performance improves under the same total rollout budget as standard GRPO.
Hard questions receive rollouts that can carry positive reward signal.
The gains hold across multiple reasoning datasets and model architectures.
Limits
The abstract reports no accuracy figures for the gain.
The allocation policy must be trained offline before use.
The total rollout budget stays fixed, so compute per run holds.
Source · arXiv