---
title: 'Rollout-GDRO: Dynamic Compute Allocation'
url: https://www.emergentmind.com/topics/rollout-gdro
type: topic
---

# Rollout-GDRO: Dynamic Compute Allocation

Rollout-GDRO is a compute-neutral allocation strategy introduced within the Group Distributionally Robust Optimization (GDRO) framework for reinforcement learning post-training of large language models (LLMs). It dynamically reallocates the number of policy rollouts per prompt across difficulty-defined groups, maximizing gradient variance reduction for hard reasoning tasks under a strict mean rollout budget. Rollout-GDRO contrasts with standard uniform rollout strategies by adapting computation to the evolving difficulty distribution of prompts, yielding significant robustness gains on heterogeneous, heavy-tailed reasoning data [2601.19280].

## 1. Motivation for Dynamic Rollout Allocation

Traditional Group Relative Policy Optimization (GRPO) allocates a fixed number of rollouts $n$ per sampled prompt, maintaining uniform exploration across the dataset. In LLM reasoning, the data distribution is heavy-tailed: most prompts become “solved” (i.e., exhibit low variance) early, whereas a minority of hard prompts continue to present high estimation uncertainty. Persisting with uniform rollout allocation leads to wasted compute on simple, well-understood prompts and inadequate coverage of the hard tail, thereby hindering policy improvement precisely where it is most needed. Rollout-GDRO addresses this inefficiency by making the rollout count a function of the prompt's difficulty group, while strictly preserving the total mean rollout budget $\bar n$.

## 2. Game-Theoretic Formulation of Rollout-GDRO

At each training step, an online pass@k classifier partitions prompts into $B$ dynamic “difficulty bins.” Let $\hat q_t(b)$ denote the empirical proportion of prompts in bin $b$ for the current batch. The goal is to allocate, for each bin $b$, an integer number of rollouts $n_b \in \{ n_{\min}, \dots, n_{\max} \}$ so as to maximize empirical utility under a compute constraint:
\[
\max_{\{n_b\}} \sum_{b=1}^B \hat q_t(b)\, \hat J_b(\theta; n_b) \quad \text{s.t.} \quad \sum_{b=1}^B \hat q_t(b)\, n_b = \bar n,
\]
where $\hat J_b(\theta; n_b)$ is the empirical bin-utility (negative loss) defined as
\[
\hat J_b(\theta; n_b) = -\frac{1}{|\mathcal{B}_{b,t}|} \sum_{x_i \in \mathcal{B}_{b,t}} \left[ \frac{1}{n_b} \sum_{j=1}^{n_b} \ell_{i,j}(\theta) \right],
\]
with $\ell_{i,j}(\theta)$ the GRPO loss for rollout $j$ and prompt $x_i$. The corresponding Lagrangian introduces a shadow price $\mu$ (compute dual variable):
\[
\mathcal{L}(\{n_b\}, \mu) = -\sum_{b=1}^B \hat q_t(b)\, \hat J_b(\theta; n_b) + \mu \left(\sum_{b=1}^B \hat q_t(b)\, n_b - \bar n \right).
\]
Each bin $b$ faces a penalized bandit loss:
\[
L_b(n) = -\hat J_b(\theta; n) + \mu n.
\]

## 3. Variance Proxy and the Square-Root Allocation Law

The effect of $n_b$ on policy improvement manifests through variance reduction in Monte Carlo gradient estimation. Under mild bounded-difference conditions, the per-prompt gradient variance for bin $b$ is bounded as:
\[
\mathrm{Var}[\hat g(x; \theta, n_b)] \le \frac{v_b(\theta)}{n_b},
\]
where $v_b(\theta)$ is the intrinsic gradient variance for bin $b$. The overall batch-level variance proxy is:
\[
\mathrm{VarProxy}(\{n_b\}; \theta) = \sum_{b=1}^B \hat q_t(b) \frac{v_b(\theta)}{n_b}.
\]
Relaxing $n_b$ to $\mathbb{R}_{>0}$ leads to a convex optimization problem whose Karush-Kuhn-Tucker (KKT) solution yields the "square-root law":
\[
n_b^* = \bar n \cdot \frac{\sqrt{v_b(\theta)}}{\sum_{j=1}^B \hat q_t(j)\, \sqrt{v_j(\theta)}},
\]
implying that bins with larger variance $v_b$ are assigned disproportionately more rollouts, scaling as $\sqrt{v_b}$ [2601.19280].

## 4. No-Regret Primal–Dual Controller and Algorithmic Implementation

The allocation is realized as a discrete zero-sum game between a primal player, distributing rollout arms via an EXP3P/entropic mirror-descent update, and a dual player, updating the shadow price $\mu$ by projected gradient ascent on the budget violation. At each step:
- Each bin $b$ samples rollout arm $n_{t,b}$ from a mixture of current distribution $p_{t,b}$ and uniform exploration.
- The exact mean budget constraint is enforced (via dynamic programming if necessary).
- Bin-utility $\hat J_b(\theta; n_{t,b})$ is estimated.
- Primal updates use bandit loss $-\hat J_b(\theta; n) + \mu_t n$.
- The dual variable $\mu_{t+1}$ is updated as $\mu_{t} + \alpha_\mu (\bar n_t - \bar n)$, projected to a prescribed interval.
The algorithm guarantees (Theorem B.11 in the paper) that the saddle-point gap of the averaged iterates is of order $O\left((\log K + \mu_{\max}^2) / \sqrt{T}\right)$ after $T$ training steps, ensuring no-regret and near-optimality for the original budgeted variance minimization.

## 5. Pseudocode Summary

A compact version of Rollout-GDRO is as follows:

```python
# bins 1…B, arms N=[n_min,...,n_max], mean budget n_bar, dual lr alpha_mu, primal lr eta_p, exploration gamma_p
initialize p_{1,b} = Uniform(N) for all b; mu_1 = 0
for t = 1 to T:
    # compute empirical bin fractions
    observe batch, compute q_t(b) for each bin
    # primal: sample arms
    for each bin b:
        n_{t,b} ~ (1-gamma_p)*p_{t,b} + gamma_p*Uniform(N)
    # enforce budget constraint
    adjust {n_{t,b}} using dynamic programming so sum_b q_t(b) n_{t,b} = n_bar
    # collect rollouts and compute GRPO losses
    for each bin b and prompt:
        run n_{t,b} rollouts and compute ell_{i,j}
    # estimate bin-loss
    compute J_b(theta; n_{t,b}) for each b
    # primal update
    for each b, n in N:
        hat_L_{t,b}(n) = -J_b(theta; n) + mu_t*n
        p_{t+1,b}(n) ∝ p_{t,b}(n) * exp(-eta_p * hat_L_{t,b}(n))
    # dual update
    n_bar_t = sum_b q_t(b) n_{t,b}
    mu_{t+1} = mu_t + alpha_mu * (n_bar_t - n_bar)  # project if needed
return theta_T, p_{T,b}
```

## 6. Empirical Evaluation and Performance

Rollout-GDRO was evaluated on the DAPO-14.1k math reasoning dataset, using pass@8 as the online difficulty classifier and Qwen3-Base models at 1.7B, 4B, and 8B parameter scales. Under a mean budget of $\bar n = 4$ (with allowed arms $n \in [2,12]$), Rollout-GDRO achieved substantial improvements in pass@8 accuracy at no additional sampling compute:

| Model Size | GRPO   | Rollout-GDRO | Relative Gain |
|------------|--------|--------------|---------------|
| 1.7B       | 50.74% | 56.14%       | +10.64%       |
| 4B         | 56.31% | 62.27%       | +10.59%       |
| 8B         | 62.04% | 67.75%       |  +9.20%       |

## 7. Emergent Allocation Patterns and Qualitative Analysis

Rollout-GDRO exhibits several qualitatively distinct behaviors:

- **Budget Frontier**: The allocation decouples compute from data frequency; rare, hard bins can receive 3×–10× the rollouts assigned by uniform baselines.
- **Variance Reduction**: The weighted standard-error proxy (WSE) is reduced by 37.1%, 22.6%, and 33.4% for 1.7B, 4B, and 8B models, respectively, relative to uniform rollout allocation.
- **Staircase Allocation**: The shadow price $\mu$ induces abrupt transitions between discrete rollout arm choices for bins (“staircase” patterns in allocation snapshots).
- **Multiplier Effect**: At mid-training, bins containing fewer than 20% of prompts may command over 80% of the total rollout budget.

These findings collectively support the conclusion that dynamically adapting the rollout budget using a GDRO-style shadow-price controller induces an emergent curriculum, shifting computational resources toward evolving hard-task frontiers and yielding large robustness improvements for LLM reasoning, all without increasing total sampling compute [2601.19280].

Source: https://www.emergentmind.com/topics/rollout-gdro