PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2610.37429v1cs.LGcs.AI

TL;DR

GRPO allocates the same rollout count to every question despite variation in question difficulty, which can leave some questions without informative advantages. AdaRoll-GRPO uses a Bayes-adaptive MDP to learn a fixed-budget allocation policy. Experiments report consistent training improvements across reasoning benchmarks, LLM architectures, and GRPO variants.

  • Problem

    Uniform GRPO rollout allocation ignores question difficulty, although identical rewards can yield zero advantages and no learning signal for some questions.

    Source: §1 · p. 1

  • Method

    AdaRoll-GRPO models rollout allocation as a sequential decision process and uses a Bayes-adaptive MDP to learn an adaptive policy for a fixed total budget.

    Source: §1 · p. 2

  • Results

    AdaRoll-GRPO consistently improves training performance across multiple reasoning benchmarks, LLM architectures, and GRPO variants under the same rollout budget.

    Source: §8 · p. 12

  • Takeaways & Limitations

    The learned allocation policy is computed offline and applied during GRPO training without additional online optimization overhead.

    Source: §8 · p. 12

    The BAMDP uses question-specific rollout success probabilities and represents them initially with uniform [0, 1] priors.

    Source: §5.1 · p. 7, §5.2 · p. 8

Abstract

from arXiv · show

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.

1 Introduction

GRPO uses uniform rollout budgets despite substantial variation in question difficulty, which can leave some questions without informative reward signals. AdaRoll-GRPO reallocates a fixed budget through a learned policy, and experiments report improved performance across GRPO variants, datasets, and LLMs.

1 Introduction

  • Uniform rollout allocation can leave easy or very difficult questions with identical rewards and zero advantages, while moderately difficult questions may need fewer rollouts.Questions contribute to the policy update through advantages computed from rollout rewards.Source: §1 · p. 1
  • More than 28% of MATH training questions require at least eight rollouts before the first correct solution, illustrating variation in question difficulty.The passage also reports that many MATH questions are solved with one rollout.Source: §1 · p. 2
  • The method models allocation as sequential question selection in a finite-horizon MDP, then uses a BAMDP belief state to handle unknown rollout success probabilities.Value iteration yields a closed-form Bellman update for the value function.Source: §1 · p. 2
  • The allocation policy is computed offline and applied during GRPO training without additional online optimization overhead.AdaRoll-GRPO dynamically reallocates rollouts while maintaining the predefined total budget.Source: §1 · p. 2
  • AdaRoll-GRPO improves training performance across different GRPO variants, datasets, and LLMs.The rollout budget remains the same as the predefined total.Source: §1 · pp. 2–3

2 Related Work

Related work improves GRPO by changing its optimization objective or its advantage estimation. AdaRoll-GRPO instead changes how rollout budgets are distributed, making it complementary to those approaches.

2 Related Work

  • Prior work on GRPO improvement falls into objective modification and advantage-estimation improvement.The paper positions rollout allocation as a distinct perspective.Source: §2 · p. 3

2.1 Objective Modifications

  • DAPO, Dr. GRPO, GSPO, and TIC-GRPO modify clipping, normalization, or token-versus-sequence importance ratios in the optimization objective.The paper describes rollout budget allocation as orthogonal to these objective changes and combinable with them.Source: §2.1 · p. 3

2.2 Advantage Estimation Improvements

  • Dr. GRPO, MC-GRPO, KRPO, and MS-GRPO alter reward normalization, centering, filtering, or subgroup aggregation to improve advantage estimation.AdaRoll-GRPO allocates rollouts across questions rather than modifying advantage estimation, so the approaches can be combined.Source: §2.2 · p. 4

3 Background on GRPO

GRPO estimates advantages from relative rewards among responses sampled for the same prompt, avoiding the separate value model used by PPO. Its equal-sized rollout groups can nevertheless be poorly matched to question difficulty.

3 Background on GRPO

  • GRPO avoids PPO’s separate value model by estimating advantages from relative rewards among responses sampled for the same prompt.This reduces the computational and memory overhead associated with training and maintaining a value model.Source: §3 · p. 4
  • Each rollout is a token sequence, and GRPO applies a clipped surrogate objective using the current policy’s token probabilities and a reference policy.The cited passages identify the current policy as the one being optimized and the reference policy as the one used to generate rollouts.Source: §3 · p. 4
  • Within each prompt’s group, responses with higher rewards than their peers receive positive advantages, while lower-reward responses receive negative advantages.The work considers binary rewards indicating whether the rollout’s final answer is correct.Source: §3 · pp. 4–5
  • With a uniform budget of eight rollouts, a harder question can receive no correct rollout and contribute no policy update while an easier question receives unnecessary rollouts.This illustrates why equal rollout counts may fail to match question difficulty.Source: Figure 3 · p. 5

4 Rollout Budget Allocation Problem Formulation

Uniform allocation can underserve difficult questions and overspend on questions already yielding informative rewards. The formulation allocates a fixed total rollout budget to maximize expected utility computed from stochastic rollout rewards.

4.1 Illustrative Example

  • A question solved only after ten rollouts can receive no correct sample under uniform allocation of eight, while a question solved after two receives more rollouts than needed.The example proposes shifting budget from the moderately difficult question to the harder one.Source: §4.1 · p. 5

4.2 Problem Formulation

  • The allocation vector assigns a rollout count to each question in the batch, with total allocations constrained to the vanilla GRPO budget.The objective maximizes expected training utility because rollout rewards are stochastic and observed only after generation.Source: §4.2 · p. 6
  • The minimum-reward-variance utility seeks diverse rewards for every question, since reward diversity yields more non-zero advantages and more effective training signals.It uses the minimum reward variance across questions to encourage all questions to provide diverse signals.Source: §4.2 · p. 6
  • The minimum-advantage-contribution utility uses the smallest question-level sum of squared advantages to encourage every question to contribute to training.The paper connects larger advantage magnitudes to greater contributions to the batch gradient.Source: §4.2 · p. 6
  • The summed-maximum-rewards utility counts correct-answer generation by summing each question’s maximum reward across the batch.A maximum reward of one indicates at least one correct trajectory for that question.Source: §4.2 · p. 7
  • The choice of utility determines the optimization objective and therefore the resulting rollout allocation strategy.The paper also lists additional objectives: minimum maximum rewards, summed reward variances, and summed advantage contributions.Source: §4.2 · pp. 6–7

5 Adaptive Rollout Budget Allocation

The rollout-allocation objective depends on rewards observed only after rollouts are generated, making it an online decision problem. AdaRoll-GRPO models allocation sequentially and uses beliefs over unknown question-solving probabilities to choose allocations under a finite rollout horizon.

5 Adaptive Rollout Budget Allocation

  • The rollout-allocation objective is an online integer program because utility depends on rewards observed only after allocation and generation.Source: §5 · p. 7

5.1 MDP Formalization

  • With horizon N, each MDP step assigns one rollout to a question, and the state tracks rollout totals and correct-answer counts needed by the utility functions.Source: §5.1 · p. 7
  • Intermediate decisions receive zero reward; the MDP assigns utility only after all N rollouts, using the final allocation vector G.Source: §5.1 · p. 8

5.2 Bayes-Adaptive MDP (BAMDP) Formalization

  • Because GRPO updates change each question’s success probability, the MDP’s transition probabilities are unknown during training and are represented with beliefs in a BAMDP.Source: §5.2 · p. 8
  • A Beta prior yields a Beta posterior after rollout outcomes, and its success and failure parameters capture the original MDP state’s rollout counts.Source: §5.2 · pp. 8–9

6 Optimal Rollout Allocation Policy

The paper solves the BAMDP with value iteration, assigning terminal value from the final reward and recursively evaluating intermediate allocation choices. It computes the policy offline and applies it as a lookup during GRPO training.

6 Optimal Rollout Allocation Policy

  • Value iteration evaluates intermediate states with the Bellman optimality equation, comparing successful and unsuccessful rollout outcomes to obtain the optimal allocation policy.Source: §6 · p. 9
  • The policy is computed once offline and stored as a belief-state lookup table, so rollout allocation adds no online computational cost during training.Source: §6 · p. 9
  • For fixed batch size, offline value iteration is polynomial in rollout budget N and does not affect GRPO’s per-step computational cost.Source: §6 · p. 10

7 Experiments

Experiments compare AdaRoll-GRPO with GRPO variants across reasoning datasets and LLMs, and examine utility objectives and dataset difficulty heterogeneity. Adaptive allocation consistently improves performance over its original counterparts, with larger gains associated with higher dataset difficulty-distribution entropy.

7 Experiments

  • Across four benchmark datasets and four LLM architectures, evaluation spans mathematical, logical, and commonsense reasoning.The datasets are GSM8K, MATH, LogiQA, and CoLoTa; the models are Llama 3.1-8B-Instruct, Llama 3.2-3B-Instruct, Gemma 3-4B-IT, and Qwen 2-MATH-1.5B.Source: §7 · p. 10, §7.1 · p. 10

7.2 Baselines and Evaluation

  • The rollout-allocation strategy leaves the underlying GRPO objective unchanged and is evaluated with vanilla GRPO, Dr.GRPO, and DAPO.Each variant is compared with its counterpart using the proposed allocation strategy, and models are evaluated by exact-match accuracy on the corresponding test set.Source: §7.2 · p. 10
  • The minimum reward variance objective, U1, achieves the best results in 6 of 14 benchmark-model combinations, although no objective consistently wins.The authors adopt U1 for the remaining experiments.Source: §7.2 · p. 10

RQ 3. How does training data heterogeneity impact the effectiveness of AdaRoll-GRPO?

  • Figure 4 relates normalized relative improvement over vanilla GRPO to dataset difficulty-distribution entropy; performance gain increases with entropy.Higher entropy indicates greater heterogeneity in question difficulty.Source: Figure 4 · p. 11, RQ 3. How does training data… · p. 12

7.2 Baselines and Evaluation

  • AdaRoll-GRPO consistently improves performance over vanilla GRPO, Dr.GRPO, and DAPO counterparts.Table 2 compares final test accuracy for each variant and its AdaRoll-enhanced counterpart.Source: §7.2 · p. 11, Table 2 · p. 12

8 Conclusion

AdaRoll-GRPO uses a Bayes-adaptive MDP to learn difficulty-aware allocation of a fixed rollout budget without additional online optimization. Experiments across reasoning benchmarks, LLM architectures, and GRPO variants show consistent training-performance improvements under that budget.

8 Conclusion

  • AdaRoll-GRPO allocates a fixed rollout budget across questions according to estimated difficulty, without additional online optimization overhead.The allocation policy is learned by formulating rollout allocation as a BAMDP.Source: §8 · p. 12
  • Experiments across multiple reasoning benchmarks, LLM architectures, and GRPO variants demonstrate consistent training-performance improvements under the same rollout budget.The authors characterize the framework as an efficient and effective approach to improving GRPO-based LLM post-training through adaptive rollout allocation.Source: §8 · p. 12

A Alternative Training Utility Functions

The appendix presents alternative training utility objectives and describes a Bayesian posterior update for rollout success probabilities.

A Alternative Training Utility Functions

  • U4 prioritizes the weakest question’s maximum reward, U5 maximizes total reward variance, and U6 maximizes total advantage contribution without explicitly balancing questions.U4 encourages at least one successful trajectory per question whenever possible; U5 promotes batch reward diversity without explicitly balancing question-level variance.Source: A Alternative Training Utility Functions · p. 15

B Posterior Update

  • Bayes’ rule combines the prior with the Bernoulli likelihood of an observed binary reward to produce a posterior with a Beta-distribution kernel.The passage describes the posterior density update for the probability of a correct reasoning trajectory.Source: B Posterior Update · pp. 15–16

C Information on Tasks and Datasets

The appendix describes the reasoning datasets and evaluation splits, then compares training curves for two Llama models. It also provides example questions and answers from the studied datasets.

C.1 Mathematical Reasoning

  • Mathematical reasoning is evaluated on GSM8K and MATH; training uses the full GSM8K split and the first 4,000 MATH samples, with MATH500 evaluation.GSM8K covers multi-step arithmetic word problems; MATH contains advanced competition-level problems.Source: §C.1 · p. 16

C.2 Logical Reasoning

  • LogiQA evaluates deductive reasoning over language-based premises, using the first 5,000 training samples and its official test split.Its questions test formal reasoning skills including implication, contradiction, and quantification under complex contextual constraints.Source: §C.2 · p. 16

C.3 Commonsense Reasoning

  • CoLoTa assesses commonsense reasoning over long-tail entities in an out-of-distribution setting, with a random 90/10 train-test split.The dataset has no official train/test split.Source: §C.3 · p. 16

C Information on Tasks and Datasets

D.1.1 Llama-3.1-8B-Instruct model

  • Across GSM8K and MATH, all AdaRoll-GRPO variants consistently achieve higher test rewards than vanilla GRPO throughout training Llama-3.1-8B-Instruct, and attain higher final rewards.Source: §D.1.1 · p. 17, Figure 5 · p. 18

D.1.2 Llama-3.2-3B-Instruct model

  • On Llama-3.2-3B-Instruct, most AdaRoll-GRPO variants outperform vanilla GRPO on GSM8K, while all variants do so on the more challenging, heterogeneous MATH dataset.The reported comparison is of test rewards throughout training.Source: Figure 6 · p. 18, §D.1.2 · p. 18

D.2 Analysis of Learned Rollout Allocation Policies

The learned policies allocate rollouts differently according to their utility objectives. In the analyzed batch, sum-based and minimum-based objectives produce distinct allocation patterns.

D.2 Analysis of Learned Rollout Allocation Policies

  • The allocation visualization compares per-question rollout counts with mean rewards; its dashed line marks the average of 8 rollouts per question.AdaRoll-GRPO's total rollout count matches vanilla GRPO with G = 8.Source: §D.2 · pp. 18–19, Figure 7 · p. 20
  • Sum-SumSqA concentrates rollouts on a few moderately difficult questions, including one receiving 32, whereas Min-SumSqA avoids overallocating to any single question.The question receiving 32 rollouts had a mean reward near 0.33, yielding both correct and incorrect trajectories.Source: §D.2 · p. 19
  • Min-V overallocates to questions with mean reward one and zero variance, while Sum-V shows the same tendency less aggressively and allocates more uniformly.Questions with intermediate mean rewards and nonzero variance consequently receive fewer rollouts under Min-V.Source: §D.2 · p. 19
  • Both MaxR objectives target questions without a successful rollout; Sum-MaxR favors higher estimated success probabilities, while Min-MaxR does not explicitly do so.A solved question cannot further improve either objective's contribution for that question.Source: §D.2 · p. 19
  • On MATH, Sum-SumSqA again concentrates budget on a small subset, while Min-SumSqA spreads it more evenly; Min-V also favors questions with mean reward one more strongly than Sum-V.Several questions receive the maximum 32 rollouts under Sum-SumSqA, while none receive more than 15 under Min-SumSqA.Source: §D.2 · p. 19
  • Sum-MaxR leaves 4 questions unsolved versus 7 for Min-MaxR, while prioritizing unsolved questions with higher estimated success probability to maximize total solved questions.Min-MaxR instead distributes remaining rollouts more evenly to improve the minimum reward across questions.Source: §D.2 · pp. 19–20

E.1 Rollout Requirement Analysis

The appendix details the rollout-requirement analysis, training settings, multi-GPU implementation, entropy measure, and hardware used.

E.1 Rollout Requirement Analysis

  • For rollout-requirement analysis, the Llama-3.2-3B-Instruct model is evaluated on MATH and GSM8K training data, tracking questions first solved at progressively larger rollout counts.Only unsolved questions receive additional rollouts at each stage, through k = 128; bar heights report the share first solved within each rollout interval.Source: §E.1 · p. 20

E.2 Training Hyperparameters

  • Training uses a validation-selected learning rate, AdamW with a cosine scheduler, batch size 2 per device, and one epoch of full-parameter fine-tuning.The learning rate is selected from five candidate values; num_generations is set to 8, the average rollout count per question for AdaRoll-GRPO.Source: §E.2 · p. 21

E.3 Multi-GPU Training

  • For multi-GPU training, each GPU constructs a tabular value function for the global prompt batch and rollout budget, then replicates it across GPUs.The batch size is B = np and rollout budget N = np × G, where G is the average rollout count per prompt.Source: §E.3 · pp. 21–22

E.4 Entropy Computation for RQ 3

  • Dataset entropy is computed from the pretraining distribution of rollout requirements, using the procedure illustrated in Figure 1.Figure 4 relates this entropy to normalized final-accuracy improvement over vanilla GRPO.Source: §E.4 · p. 22

F Hardware Specification

  • Training experiments use NVIDIA H100 SXM5 GPUs with 80 GB memory; auxiliary analyses use an Apple M1 computer with 16 GB unified memory.The 8B model uses two GPUs, while other models use one.Source: F Hardware Specification · p. 22

Related papers

Loading 2610.37429v1…