Skip to content

Back to the catalog

Bayes-Adaptive MDPs for Rollout Allocation to Enhance GRPO Training
arxiv.orgpaper

Bayes-Adaptive MDPs for Rollout Allocation to Enhance GRPO Training

AdaRoll-GRPO is a method that allocates GRPO rollout budgets per question by difficulty instead of uniformly. It frames allocation as a sequential decision problem and learns an allocation policy with a Bayes-adaptive MDP, computed offline so training adds no online optimization cost. Experiments on several reasoning datasets and LLM architectures show gains over standard GRPO at the same total rollout budget.

TL;DR Proposes AdaRoll-GRPO, which uses a Bayes-adaptive MDP to assign rollouts by question difficulty and beats uniform-budget GRPO at equal cost.
Open original ↗
Document

What it claims

AdaRoll-GRPO, from the University of Toronto, Vector Institute and KAIST (arXiv, October 2026), reallocates GRPO rollouts inside each batch. The total stays at N = B·G, the same budget vanilla GRPO uses. The motivation is measured, not assumed. With Llama3.2-3B-Instruct, 41.48% of MATH training questions are solved on the first rollout. Another 28.5% need eight or more rollouts before any correct answer appears, and 4.55% need more than 128. On GSM8K, 74.51% are solved at the first attempt and about 4.3% need eight or more. Under uniform G = 8, those hard questions return all-zero rewards and contribute nothing to the update.

How it works

Rewards are binary: 1 if the final answer is correct, 0 otherwise. Allocation is first posed as an integer program that maximizes expected training utility subject to the allocations summing to N. That problem is online, because rewards are only seen after rollouts are drawn. So it is recast as a finite-horizon MDP. At each step the action picks one question to receive the next rollout, and the observed success or failure updates the state.

Each question's chance of producing a correct answer is unknown. The paper therefore uses a Bayes-adaptive MDP: belief distributions over those success probabilities are folded into an augmented belief state. Value iteration with a closed-form Bellman update solves it offline. The resulting policy is stored as a lookup and adds no optimization during training.

flowchart TD
  A[Belief state for batch] --> B[Policy picks question j]
  B --> C[Generate one rollout]
  C -->|correct| D[Update belief: success]
  C -->|incorrect| E[Update belief: failure]
  D --> F{Budget N used?}
  E --> F
  F -->|no| B
  F -->|yes| G[GRPO update]

The policy depends on the chosen utility. The page defines three:

  • U1: the minimum per-question reward variance.
  • U2: the minimum per-question sum of squared advantages.
  • U3: the number of questions with at least one correct rollout.

Appendix A lists three more: the minimum of maximum rewards, the sum of variances, and the sum of advantage contributions. In the paper's worked GSM8K example, uniform G = 8 is replaced by G = 4 for a question solved on its 2nd attempt and G = 12 for one first solved on its 10th.

The authors present the method as orthogonal to objective changes such as DAPO, Dr. GRPO, GSPO and TIC-GRPO. They say the same about advantage-estimation changes such as MC-GRPO, KRPO and MS-GRPO, and state that it can be combined with them.

Limits

The fetched text stops at the start of Section 5. It contains no experimental results: no benchmark scores, no baselines, no model list beyond Llama3.2-3B-Instruct for the difficulty analysis, and no data on which utility function works best. The claim of consistent gains comes only from the abstract.

The fetched text also omits several practical details:

  • the belief prior;
  • the batch size and budget the offline policy was solved for;
  • whether the policy must be recomputed when B or N changes.

The formulation assumes binary correctness rewards, and graded or learned rewards are not covered. The page gives no code link or licence.

Organisations named here

Technologies named here

Techniques named here

Related in training