Source-linked AI summary
PAC: Progress-Augmented Advantage Curriculum for Multi-Task Reinforcement Learning of LLMs
Yuanqiang Yu, Yanzhao Zheng, Zhentao Zhang, Tianze Xu, Chao Ma, Jihuai Zhu, Jiashun Liu, Xinle Deng, Baohua Dong, Hangcheng Zhu, Ruohui Huang
TL;DR
Multi-task RL post-training must allocate rollouts across tasks whose usefulness changes, while advantage-only signals may miss whether updates improve rewards. PAC combines advantage-derived learnability with reward-derived progress in a Bayesian Thompson Sampling controller, and it improves sample efficiency and final performance across multi-level and multi-domain reasoning settings. The method is evaluated under verifiable outcome-level rewards and task-level allocation granularity.
Problem
Existing multi-task RL post-training pipelines use fixed or manually designed mixtures, while online curricula may prioritize update magnitude without measuring resulting reward gains.
Method
PAC combines advantage-derived learnability and recent reward-derived progress into task utilities used by a Bayesian Thompson Sampling controller to allocate GRPO rollouts.
Results
PAC improves sample efficiency and final performance across multi-level and multi-domain reasoning, reaching baseline-level averages earlier and achieving higher curriculum averages than comparison curricula.
Takeaways & Limitations
Jointly tracking update potential and actual reward progress yields an effective online curriculum for LLM RL post-training within the evaluated settings.
Takeaways & Limitations
PAC is validated with outcome-level verifiable rewards and task-level arms; sparse or noisy rewards and substantial within-arm diversity remain scope boundaries.
Abstract
from arXiv · showhide
Reinforcement learning (RL) is used to improve the reasoning abilities of LLMs, while training data span heterogeneous tasks. However, most RL post-training pipelines rely on fixed or manually designed task mixtures, even though task usefulness changes as training progresses. Online curriculum methods often define learnability by update magnitude, ignoring whether the update translates into reward gains, which can misallocate rollout budget toward tasks with large but ineffective updates. We propose PAC, a Progress-Augmented Advantage Curriculum for multi-task RL of LLMs that combines two task-level signals: advantage-derived learnability, which measures the magnitude of the policy update a task can induce, and recent reward gains, which show whether those updates have improved task performance. A Bayesian Thompson Sampling controller uses these signals to allocate rollouts across tasks during GRPO training. We evaluate PAC under two settings: a multi-level reasoning setting and a multi-domain reasoning setting. PAC improves sample efficiency and final performance: it reaches comparable validation scores with fewer rollout steps and achieves higher final averages than random sampling and advantage-based curriculum baselines in both settings. These results show that jointly tracking advantage signals and actual reward gains yields an effective online curriculum for LLM post-training.
1 Introduction
Multi-task RL post-training must allocate limited rollouts as task usefulness changes with the policy. PAC addresses the gap by combining update potential with actual reward progress in a Bayesian online curriculum.
- Heterogeneous LLM RL mixtures create an online curriculum problem: trainers must allocate limited rollout budgets across tasks as policies change.
- GRPO makes rollout allocation consequential because it samples multiple responses, evaluates verifiable rewards, and derives relative advantages for each update.
- Advantage-only curricula can prioritize high-variance tasks whose large updates do not produce reward gains, while undersampling tasks with moderate advantages and steadily improving rewards.
- PAC combines advantage-derived learnability with reward-derived progress to estimate task utility and adapt rollout allocation through Bayesian Thompson Sampling.
- PAC is validated in multi-level and multi-domain reasoning settings, improving sample efficiency and final performance over random and curriculum baselines.
2 Related Work
Prior curriculum research spans supervised fine-tuning and reinforcement learning, but the cited approaches do not fully address online task allocation for evolving LLM RL policies.
- Curriculum Learning for SFT: SFT curricula typically use static mixtures or precomputed properties such as task balance, data quality, difficulty, and instruction diversity.
- Curriculum Learning for SFT: Because SFT targets are fixed, these methods do not address RL post-training allocation when task usefulness shifts as the policy evolves.
- Curriculum Learning for RL: RL curricula adapt tasks using episodic return, learning progress, or estimated future learning potential, with methods including Prioritized Level Replay, T3S, and Distral.
3 PAC: Progress-Augmented Advantage Curriculum
PAC models task selection as a non-stationary allocation problem and estimates each task’s short-horizon utility by combining GRPO update potential with recent reward progress. A Bayesian controller uses these estimates to adapt rollout allocation under uncertainty.
- Problem Formulation: PAC treats each task dataset as a bandit arm whose utility can become learnable, saturate, or remain too difficult as the policy changes.
- Online Task Utility Estimation: The ideal utility is the expected short-horizon reward gain from assigning a training sample to a task, approximated online because cross-task marginal gains are not directly observable.
- Online Task Utility Estimation: PAC factorizes task utility into update magnitude and reward improvement per unit update, motivating separate learnability and conversion signals.
- Advantage-Derived Learnability: Advantage-derived learnability uses mean absolute group-normalized advantages over a history window to estimate the update a task can still provide.
- Reward-Derived Progress: Reward-derived progress fits a least-squares linear trend to each task’s recent mean rewards and normalizes the fitted slopes across tasks.
- Reward-Derived Progress: PAC uses advantage-derived learnability as the base utility and modulates it with normalized progress, keeping the modifier neutral when progress is near zero.
- Bayesian Curriculum Controller: The controller maintains Gaussian beliefs over task utilities, weighting observations by sample-supported precision while preserving uncertainty for sparsely sampled arms.
- Bayesian Curriculum Controller: PAC inflates posterior variance after updates so stale evidence does not suppress exploration as task utilities change.
4 Experiments
PAC is evaluated in multi-level and multi-domain reasoning settings against single-task, random, and curriculum baselines. Across these experiments, PAC improves rollout sample efficiency and final average validation performance, with ablations showing contributions from progress, advantage, and Thompson Sampling.
- Experimental settings: PAC is evaluated on multi-level reasoning with Countdown, Zebra, and ARC, and on multi-domain reasoning spanning mathematics, code generation, and symbolic logic.The multi-level setting uses Qwen2.5-3B and Qwen2.5-7B; the multi-domain setting uses Qwen2.5-7B and Qwen2.5-32B.
- Baselines and evaluation: The baselines include single-task training, uniform random sampling, and SEC, which selects task categories using average absolute advantage as the bandit reward.Final results are averaged over three runs with different random seeds.
- Sample efficiency: PAC reaches curriculum-baseline final averages with fewer rollout steps and maintains higher final averages in the multi-domain setting.In the multi-level setting, PAC matches SEC after about 170 and 143 steps for Qwen2.5-3B and Qwen2.5-7B, versus SEC’s 225 and 227 steps, yielding roughly 1.3× and 1.6× gains; in the multi-domain setting, PAC gains approximately 1.2× and 1.3× over DUMP.
- Multi-level results: PAC achieves the best multi-level average, improving over SEC by 13.1% for Qwen2.5-3B and 8.5% for Qwen2.5-7B.These results are reported for transfer to harder validation splits, while SEC and DUMP remain competitive on selected individual splits.
- Multi-domain results: PAC achieves the best curriculum average in the multi-domain setting, improving over DUMP by 6.2% for Qwen2.5-7B and 1.9% for Qwen2.5-32B.PAC maintains stronger average performance across math, code, and logic while improving on harder math benchmarks as well as code and logic benchmarks.
- Ablation study: Ablations show that progress is dominant: removing it lowers the multi-domain average by 9.5%, while removing advantage and Thompson Sampling lowers it by 3.0% and 4.2%.Without progress, allocation becomes advantage-only and continues investing in arms whose rewards have saturated; without Thompson Sampling, allocation concentrates rapidly after some arms obtain higher utility estimates.
5 Conclusion
PAC adapts rollout allocation in multi-task LLM RL by combining update potential with reward progress. Across multi-level and multi-domain reasoning, it improves sample efficiency and final validation performance.
- PAC combines advantage-derived learnability with reward-derived progress and uses Bayesian Thompson Sampling to adapt task allocation under uncertainty.
- PAC reaches baseline-level validation averages with fewer rollout steps than random and advantage-based baselines across both reasoning settings.
- PAC achieves higher final validation averages across model sizes and experimental settings.
- The findings support tracking both update potential and reward conversion when designing adaptive RL post-training curricula.
Limitations
PAC is evaluated under verifiable outcome-level rewards and at task-level granularity, with empirical coverage limited to selected reasoning, code, and logic settings and modest arm counts.
- Reward signal assumptions: PAC is validated with outcome-level verifiable rewards, including rule-based checkers, symbolic verifiers, and unit-test execution.These rewards provide relatively dense observations for computing window-level mean rewards.
- Reward signal assumptions: Under sparser, noisier, or longer-horizon rewards, the linear trend estimator may need replacement by a more robust progress estimator.The stated concern is sensitivity to reward noise and non-monotonic learning dynamics.
- Task-level granularity: PAC allocates at task-level granularity, so internal variation within a single arm cannot be exploited by the curriculum.Prompt-level or instance-level extension remains open and would require posterior structures supporting many more arms.
- Empirical scope: Experiments cover backbones up to 32B parameters, reasoning, code, and logic mixtures, and at most nine task arms.Instruction following, dialogue, multilingual, safety, larger arm sets, and alternative advantage normalizations were not directly studied.
A.2 Proof and Online Estimator
The proof links mean absolute advantage to the scale of a task-conditioned GRPO update under bounded, smoothly varying score-gradient norms. The online estimator therefore uses advantage magnitude as a learnability proxy.
- Proof: The policy gradient for a task aggregates prompt-conditioned gradients into a task-conditioned gradient whose scale depends on absolute advantage and score-gradient norms.
- Proof: Under the theorem’s smoothness assumption, score-gradient norms vary more slowly across tasks than reward separation captured by normalized advantages.
- Proof: Mean absolute advantage is therefore the main task-dependent term controlling the GRPO update scale.
- Online estimator: The advantage-derived learnability score estimates the advantage-driven update signal task i can still provide under the current policy.
B.2 Training and Sampling
PAC samples from predefined task arms, records rewards and advantages during a uniform cold start, then uses sequential Thompson Sampling with posterior updates and variance inflation.
- Training configuration: Each training step constructs a batch of B = 256, and each selected prompt receives G = 8 sampled responses scored by a verifiable reward function.
- Task arms and sampling: Task arms follow difficulty splits in the multi-level setting and domain-specific splits in the multi-domain setting.
- PAC state updates: PAC initializes each arm with Gaussian posterior mean µ_i = 0 and uses a uniform T_cold = 50-step cold-start period.During cold start, rewards and advantages are recorded so each arm receives comparable initial evidence.
- PAC state updates: After cold start, sequential Thompson Sampling draws utility values from arm posteriors to form batches and updates beliefs after sampling.
- Utility estimation: PAC computes learnability from step-level mean absolute group-normalized advantages and progress from step-level mean rewards over a history window.Insufficient recent observations set the fitted progress slope to zero before cross-task normalization.
B.4 Validation Performance at Individual Difficulty Levels
PAC preserves performance on easier difficulty levels while delivering its largest gains on harder held-out levels in the multi-level setting.
- The main multi-level evaluation emphasizes the extremely hard split to measure transfer beyond the training difficulty range.
- Validation scores are averaged across Countdown, Zebra, and ARC at four difficulty levels, including the extremely hard Ext. level.The evaluation uses Qwen2.5-3B after 500 training steps, with held-out validation examples disjoint from RL training data.
- 12.8% and 13.1% improvements on Hard and Ext., respectively, relative to the strongest baseline in each column.PAC is within 1.1% of the strongest baseline on Simple and Medium.
B.5 History-Window Sensitivity
History-window length controls a trade-off between smoothing noisy rollout signals and responding quickly to changing task utility, with W = 16 providing the best balance.
- 21.2% lower average accuracy at W = 4 than the default, falling from 0.302 to 0.238.The short window can overreact to transient fluctuations because it uses only four recent steps and reward observations.
- 10.3%, 19.2%, and 20.5% reductions at W = 64, 128, and 256, respectively, relative to W = 16.Long windows retain observations from earlier policies and can delay reallocation as task utility changes.
- W = 32 reaches 0.296, only 2.0% below the default 0.302, indicating stability near the default.W = 16 provides the best balance between smoothing rollout noise and tracking changing task utility.
- The utility ablations change either the observation driving the controller or the allocation rule while keeping the remaining training configuration fixed.
- Ours w/o adv uses reward-derived progress alone, whereas Ours w/o prog recovers an advantage-only utility matching the signal type used by SEC and DUMP.
- Ours w/o ts replaces the posterior-based Thompson Sampling controller with Boltzmann allocation over current utility scores, without uncertainty modeling or variance inflation.
C.2 Per-Task Results
Per-task ablations show that reward progress, learnability, and uncertainty-aware allocation each contribute under different multi-domain conditions.
- Ours w/o prog undersamples Math because the saturating K&K arm continues to dominate the advantage-only utility, losing most ground on MATH500 and AMC22–23.
- Ours w/o adv stays close to the full model on easier benchmarks but loses accuracy on AIME24 and BigCodeBench, where reward gains are slower and progress estimates noisier.The learnability term helps distinguish slow reward gains from noisy slopes in this regime.
- Ours w/o ts records the lowest K&K accuracy because Boltzmann Sampling collapses that arm’s share and lacks variance inflation to revisit it.
D Analysis of Curriculum Dynamics in Multi-Domain Training
PAC’s curriculum reallocates rollout budget as task utility changes, combining advantage-derived learnability with reward-derived progress to avoid over-investing in saturated tasks.
- Curriculum dynamics: Figure 4 tracks sampling probability, advantage-derived learnability, and reward-derived progress for Math, Code, and K&K Logic in 50-step windows.The vertical marker identifies the end of the uniform cold-start stage at step 50, after which Thompson Sampling allocates rollouts.
- Curriculum dynamics: After Thompson Sampling activates, PAC first concentrates budget on K&K Logic, then rotates budget toward Math and Code after step 200, leaving K&K near one-quarter of the batch.
- Reward-derived progress: K&K Logic’s learnability remains comparable to other arms despite declining by roughly half, so advantage-only scheduling would continue over-allocating there.
- Reward-derived progress: K&K Logic’s progress decays rapidly near its reward ceiling, lowering fused utility and redirecting rollouts toward tasks where updates still produce reward gains.
- Reallocation mechanism: Math’s learnability and progress decay slowly, while Code’s progress strengthens once rollouts reliably pass unit-test verification, increasing their fused utility.
- Validation gain: PAC’s mid-training reallocation coincides with separation from DUMP and accounts for a 6.2% cross-domain average improvement over advantage-only baselines for Qwen2.5-7B.