Bayes-Adaptive MDPs for Rollout Allocation to Enhance GRPO Training

Faeze Moradi Kalarde; Armin Toroghi; Nam Ho Koh; Doo Hwan Hwang; Seungjin Kim; Kee-Eung Kim; Scott Sanner

cs.LGOct 27, 2026

Bayes-Adaptive MDPs for Rollout Allocation to Enhance GRPO Training

Authors: Faeze Moradi Kalarde, Armin Toroghi, Nam Ho Koh, Doo Hwan Hwang, Seungjin Kim, Kee-Eung Kim, Scott Sanner

Organizations: University of Toronto · Vector Institute of AI · Korea Advanced Institute of Science & Technology

Abstract

Group Relative Policy Optimization (GRPO) has demonstrated strong performance in LLM post-training. However, it allocates a uniform rollout budget to every training question, ignoring variations in question difficulty. As a result, highly difficult questions may receive too few rollouts to generate informative positive reward signals, limiting their contribution to training, while moderately difficult questions often receive more rollouts than necessary. In this work, we formulate the problem of difficulty-aware rollout budget allocation to maximize training utility under the same total rollout budget as standard GRPO. We propose AdaRoll-GRPO, which reformulates rollout allocation as a sequential decision-making problem and leverages a Bayes-adaptive Markov decision process (MDP) to learn an adaptive allocation policy. The learned policy is computed offline and applied during GRPO training without introducing additional online optimization overhead. Extensive experiments on multiple reasoning datasets and LLM architectures demonstrate that AdaRoll-GRPO consistently improves GRPO training performance under a fixed rollout budget, highlighting the effectiveness of adaptive rollout allocation for GRPO in LLM post-training.

Figures & tables

Figure 1: Question difficulty distributions for the MATH and GSM8K training datasets, measured by the number of rollout generations required to obtain the first correct solution.

Figure 2: Illustration of AdaRoll-GRPO. In vanilla GRPO, for a batch of four questions with a total rollout budget of 16, four rollouts are uniformly allocated to each question. For some questions (e.g., Q1 and Q4), all rollouts may receive identical rewards (all 1s or all 0s), resulting in zero advantage and thus, no contribution to training. AdaRoll-GRPO instead learns a rollout allocation policy,

\phi^{\ast}

, which sequentially selects the next question to sample based on the current state. In this example, the learned policy reduces the rollout budget of Q2 and Q3, which already produce diverse rewards with only two rollouts each, and instead, allocates extra budget to Q1 and Q4 (5 and 7 rollouts, respectively). The additional rollouts increase the likelihood of obtaining diverse rewards for Q1 and Q4, producing non-zero advantages and making these questions effective training samples.

Figure 3: Illustration of the limitation of uniform rollout budget allocation in GRPO. With a fixed budget of

G=8

rollouts per question, the harder question fails to generate any correct rollout, resulting in zero rewards for the entire group and no contribution to the policy update, while the easier question is allocated more rollouts than necessary.

\displaystyle Q(\mathbf{s}_{t},j)

\displaystyle=\hat{p}_{j}^{t}V(\boldsymbol{\alpha}_{t}+\mathbf{e}_{j},\boldsymbol{\beta}_{t})+(1-\hat{p}_{j}^{t})V(\boldsymbol{\alpha}_{t},\boldsymbol{\beta}_{t}+\mathbf{e}_{j}).

Algorithm 1 Rollout Budget Allocation via Value Iteration

Llama 3.1-8B-Instruct
Dataset	$U_{1}$	$U_{2}$	$U_{3}$	$U_{4}$	$U_{5}$	$U_{6}$
MATH	48.60	44.00	47.00	45.00	47.60	46.80
GSM8K	86.05	87.19	86.13	85.06	86.13	85.75
CoLoTa	58.24	54.71	54.71	57.06	60.00	60.59
LogiQA	43.93	43.01	45.16	41.47	43.78	41.78
Llama 3.2-3B-Instruct

Table 1: Final test accuracy (

\%

) of models trained with AdaRoll-GRPO using different training utility functions.

Llama 3.1-8B-Instruct
	Vanilla GRPO		Dr. GRPO		DAPO
Dataset	Def.	+ Ada.	Def.	+ Ada.	Def.	+ Ada.
MATH	42.60	48.60	44.40	48.80	46.60	49.40
GSM8K	83.02	86.05	82.56	84.08	83.32	86.35
CoLoTa	50.00	58.24	46.47	53.52	48.23	54.12
LogiQA	41.17	43.93	42.70	44.85	43.47	46.24

Table 2: Comparison of the final test accuracy (

\%

) of GRPO variants and their AdaRoll-enhanced counterparts.

Figure 4: Normalized relative improvement of AdaRoll-GRPO vs. entropy of the dataset difficulty distributions.

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

Dataset	Sample Question	Answer
GSM8K	Five food companies sponsored a local food bank. Foster Farms donated 45 dressed chickens; American Summits donated twice the number of bottled water than the number of dressed chicken donated by Foster Farms; Hormel donated three times the number of dressed chickens that Foster Farms donated; Boudin Butchers donated one-third of the number of dressed chickens that Hormel donated; Del Monte Foods donated 30 fewer bottles of water than American Summits. How many food items did the companies donate in total?	375
MATH	A parabola has vertex $(4,2)$ and passes through $(2,0).$ Enter the align of the parabola in the form “ $y=ax^{2}+bx+c$ ”.	$y=-\frac{1}{2}x^{2}+4x-6$
LogiQA	Article: The personal right of the author, also known as the moral right of the author, refers to the rights of the copyright owner to his work that are related to the person or are inseparable without direct property content. Reputation, prestige, and the right to maintain the integrity of the work. Copyright property rights, also known as copyright economic rights, are the symmetry of the author’s personal rights and refer to the right of the copyright owner to use or authorize others to use the work in a certain way to obtain material benefits. Question: According to the above definition, Which of the following statements are correct: Options: A. A well-known online novel writer asks the publisher to publish his work with his real name, then the writer exercises his personal rights. B. The painter Wang sold his painting to Jiang at a high price, and Jiang obtained the copyright of the painting. C. The dance created by a well-known dancer was loved by the masses, and many people imitated and performed commercial performances. The dancer asked to stop this infringement, but she exercised her personal rights. D. A TV drama created by a playwright is remake because it is loved by the audience. The playwright asks the TV show producer not to modify the ending of the story, then he is exercising the property right of copyright.	A
CoLoTa	Could someone in Vaughan take a taxi to the Jamjuree Art Gallery?	No

Appendix

Table 3: Example questions and their answers from the studied evaluation datasets.

Panel of Figure 5 — Figure 5: Comparison of the test-set reward achieved by different AdaRoll-GRPO variants and vanilla GRPO throughout training the Llama-3.1-8B-Instruct model . Left: training on the GSM8K dataset. Right: training on the MATH dataset.

Panel of Figure 6 — Figure 6: Comparison of the test-set reward achieved by different AdaRoll-GRPO variants and vanilla GRPO throughout training the Llama-3.2-3B-Instruct model . Left: training on the GSM8K dataset. Right: training on the MATH dataset.

Figure 7: Comparison of rollout allocations across prompts within a batch for different AdaRoll-GRPO variants at a randomly selected training step of Llama-3.1-8B-Instruct on GSM8K , together with the corresponding mean reward of the generated rollouts.

Figure 8: Comparison of rollout allocations across prompts within a batch for different AdaRoll-GRPO variants at a randomly selected training step of Llama-3.1-8B-Instruct on MATH , together with the corresponding mean reward of the generated rollouts.

Hyperparameter	Value
Optimizer	AdamW
Scheduler	Cosine
Fine-tuning	Full-parameter
Loss type	GRPO / DR-GRPO / DAPO
Learning rate (Llama-3.1-8B-Instruct, Llama-3.2-3B-Instruct)	$1\times 10^{-6}$
Learning rate (Gemma3-4B-IT)	$5\times 10^{-6}$

Appendix

Table 4: Selected training hyperparameter values.

RL Post-Training

Reinforcement learning post-training adapts pretrained models with reward or feedback signals to improve reasoning, alignment, and interactive behavior.

See all 426 papers

Cross-Epoch Adaptive Rollout Optimization for RL Post-Training
Jun 4, 2026Yiming Zong, Yige Wang, Jiashuo JiangSequential Decision Making RL for Language Model Reasoning
Department of Industrial Engineering & Decision Analytics, Hong Kong University of Science and Technology
The Easy, the Hard, and the Learnable: Confidence and Difficulty-Adaptive Policy Optimization for LLM Reasoning
Jun 6, 2026Zhanke Zhou, Xiangyu Lu, Chentao Cao +4RL for Language Model Reasoning Policy Optimization
1TMLR Group, Department of Computer Science, Hong Kong Baptist University · 2Stanford University · 3Sydney AI Centre, The University of Sydney
Rollout-Level Advantage-Prioritized Experience Replay for GRPO
Jun 3, 2026Gyeongtae Yoo, Sanghyeok Park, Soohyuk Jang +2Sample-Efficient RL Group Relative Policy Optimization
Department of Electrical and Computer Engineering, Seoul National University · Interdisciplinary Program in AI, Seoul National University · AIIS, ASRI, INMC, and ISRC, Seoul National University

Jun 4, 2026J/K move · Enter open · S save

Cross-Epoch Adaptive Rollout Optimization for RL Post-Training

Yiming Zong, Yige Wang, Jiashuo Jiang

Department of Industrial Engineering & Decision Analytics, Hong Kong University of Science and Technology

Sequential Decision Making RL for Language Model Reasoning RL Post-Training Online Resource Allocation Sample-Efficient RL

Open paper AskPDF arXiv

LLM post-training often relies on reinforcement learning methods that sample multiple rollouts per prompt, yet most existing approaches use a fixed rollout budget for every prompt, despite large differences in the training signal different prompts provide. In this paper, we study adaptive rollout allocation under a fixed global budget and formulate the problem as online resource allocation with prompt-level diminishing returns. Our method, CERO, maintains a Beta posterior over each prompt's success probability and uses the posterior expected Bernoulli variance as a Bayesian estimate of the value of additional rollouts. We use this estimate to construct a concave, saturating utility over cumulative allocations, yielding an objective in which decisions across prompts and epochs are coupled by the global budget. Since the resulting objective is temporally nonseparable, we derive a Fenchel-dual reformulation and update both prompt-level and budget-level dual variables via projected online gradient descent. Under fixed prompt utilities, we prove an

O(\sqrt{K})

regret bound against the offline allocation benchmark. Experiments on mathematical-reasoning problems show that CERO consistently outperforms GRPO across multiple open-weight LLMs and benchmarks, demonstrating that adaptive rollout budgeting can improve sample efficiency.

Abstract

RL Post-Training

Explore similar work

Cross-Epoch Adaptive Rollout Optimization for RL Post-Training

The Easy, the Hard, and the Learnable: Confidence and Difficulty-Adaptive Policy Optimization for LLM Reasoning

Rollout-Level Advantage-Prioritized Experience Replay for GRPO