Cross-Epoch Adaptive Rollout Optimization for RL Post-Training

Abstract: 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.

• The paper introduces CERO, a framework that redefines rollout allocation across epochs by leveraging Bayesian prompt modeling to maximize cumulative utility under a fixed budget.
• The methodology utilizes a Fenchel-dual reformulation and a primal-dual online algorithm to dynamically schedule rollouts based on estimated marginal utility.
• Empirical results on diverse models and datasets demonstrate consistent accuracy improvements up to +6.04 absolute points over conventional RL post-training methods.

Cross-Epoch Adaptive Rollout Optimization for RL Post-Training: An Expert Perspective

Motivation and Problem Formulation

Reinforcement learning post-training remains central for enhancing the reasoning abilities of LLMs, particularly in mathematical and algorithmic tasks. Standard RL post-training algorithms, such as GRPO, typically allocate a fixed, uniform number of rollouts per prompt per epoch, neglecting the heterogeneity of prompt informativeness and their time-varying contribution to learning. This uniform allocation paradigm results in substantial inefficiency as some prompts are either too easy (providing little marginal information) or too hard (yielding noisy or under-informative gradients), while only a modest subset drive useful updates.

The paper introduces the Cross-Epoch Rollout Optimization (CERO) framework, which reframes rollout budgeting as a cross-epoch, prompt-level resource allocation problem with fixed global budget constraints and diminishing-returns utility. Formally, each prompt receives rollouts across multiple epochs, with the allocation jointly optimized to maximize cumulative utility subject to a fixed overall rollout budget.

Bayesian Utility Design and Fenchel-Dual Reformulation

CERO leverages Bayesian prompt modeling, maintaining a Beta posterior over the latent pass probability for each prompt. This enables online estimation of:

• Posterior mean: The current pass rate
• Posterior variance: Residual uncertainty about the prompt
• Posterior Bernoulli variance expectation: $$q_{k,i} = \mathbb{E}[p_i(1-p_i)\mid \mathcal{H}_i^k]$$, serving as a proxy for a prompt’s expected informativeness as a function of policy progress

The utility of allocating $n$ rollouts to prompt $i$ is thus $U(q_{k,i}, n) = 1-\exp(-\eta q_{k,i} n)$—a concave, saturating function, ensuring strong diminishing returns and prioritization of under-explored, moderately difficult prompts.

Temporal non-separability (utilities depend on cumulative allocation over epochs) is addressed through a Fenchel-dual formulation, yielding a primal-dual online optimization architecture with prompt-level dual variables and a global budget Lagrange multiplier. This enables online rollout scheduling by thresholding the estimated marginal value against a dynamic global price.

Figure 1: The CERO pipeline: prompt-level dual variables orchestrate within-epoch rollout allocation while a global multiplier enforces budget discipline.

Primal–Dual Online Optimization Algorithm

The core CERO algorithm operates atop an existing RL policy optimization protocol (e.g., GRPO), acting solely on the data collection side:

1. Prompt Evaluation: For each prompt, estimate marginal utility of rollouts via the Beta posterior and compute the dual variable.
2. Rollout Scheduling: Assign rollouts only to prompts whose estimated value exceeds the current budget price; all others receive none.
3. Dynamic Budget Control: Update the global dual variable by comparing actual rollout consumption against the per-epoch budget target, throttling allocation accordingly.
4. Prompt-Level Dual Updates: Use projected online gradient updates for the prompt-level duals, adapted by feedback from the new posterior.
5. Policy Updates: As rollouts accrue, the policy is periodically (e.g., batch-wise) updated using the underlying RL objective (unchanged).

This framework inherently couples prompt selection and rollout budgeting across epochs, redistributing the allowed budget toward epochs/prompts anticipated to supply the highest marginal gradient signal.

Theoretical Guarantees

The foundational theoretical result is a non-asymptotic $O(\sqrt{K})$ regret bound relative to an offline optimal benchmark for the cross-epoch allocation problem. The analysis applies standard online convex optimization (OCO) techniques to both prompt- and budget-level dual spaces and covers the case where prompt utilities are stationary. This represents a significant strengthening over batch-only/post-generation approaches, since cross-epoch coupling is explicitly handled.

Empirical Evaluation and Numerical Results

CERO was implemented atop GRPO and evaluated on a suite of open-weight LLMs (DeepSeek-R1-Distill-Qwen-1.5B, Qwen3-4B-Base, Qwen3-4B-Instruct, Qwen2.5-Math-7B) and four mathematical reasoning datasets (AIME24, AIME25, AIME26, AMC23). Across all settings, CERO consistently outperforms GRPO across accuracy benchmarks, delivering improvements ranging from $+1.53$ to $+6.04$ absolute points, depending on the model and dataset.

Ablation analyses highlighted a key causal advantage: CERO maintains a higher fraction of “effective” prompts (prompt groups exhibiting both pass and fail), and thereby prevents the effective prompt ratio from collapsing during late training. This maintains a stronger learning signal and explains the observed downstream gains.

Figure 2: Accuracy improvement and increased effective prompt ratio for DeepSeek-R1-Distill-Qwen-1.5B, showcasing CERO’s efficacy in mathematical reasoning tasks.

Practical Implications and Integration

CERO is explicitly designed as a generic data collection scheduler—requiring no modification of the RL policy update or reward objective—making it readily integrable into current LLM post-training pipelines. Its negligible computational overhead renders it practical for large-scale deployment. Further, the Bayesian informativeness model can easily generalize to alternative definitions (contingent on reward type), and the core dual optimization approach is extensible to more sophisticated resource constraints.

Limitations and Future Work

The theoretical regret analysis relies on the assumption of prompt-level static utilities, whereas actual implementation involves prompt informativeness scores that are plugged-in and refreshed online, introducing a mild nonstationarity. Addressing full theoretical treatment under policy evolution and extending CERO to settings with soft/hard constraints, non-binary reward models, or on-the-fly batch construction remains open.

Generalization to diverse RL post-training domains—beyond mathematical reasoning or supervised instruction data—will further solidify its versatility.

Conclusion

CERO advances RL-based LLM post-training by principled, adaptive, cross-epoch allocation of rollout budget using Bayesian prompt modeling and primal-dual online optimization. The framework improves sample efficiency and policy update quality through judicious exploitation of prompt-level informativeness, yielding both strong empirical gains and provable regret bounds under fixed-budget constraints. Future directions include robustification to nonstationary prompt utility, exploration of richer resource allocation objectives, and domain transfer.