ICML2026

Gradient-Aware Scheduling: Coupling Curriculum and Staleness for Async Reinforcement Learning

Xinyu Zhang

摘要

Asynchronous reinforcement learning (RL) accelerates LLM training through parallel data collection, but introduces policy lag: experiences are collected under stale policy weights, which destabilizes learning—especially on hard tasks. We identify why: gradient variance scales exponentially with task difficulty under staleness, because hard tasks have narrow solution spaces corresponding to sharp loss-landscape curvature (high Hessian eigenvalues). We formalize this as a staleness budget optimization problem and prove that the optimal allocation follows an exponential decay: η(d)=ηbaseeλd\eta^*(d) = \eta_{\text{base}} \cdot e^{-\lambda d} where λ=α/2\lambda = \alpha/2 is half the Hessian growth rate. Building on this principle, we propose GAS (Gradient-Aware Scheduling), a drop-in recipe for fast and stable asynchronous RL with three components: (i) a bandit-based curriculum that selects task difficulty by learning signal (ACB); (ii) execution-aware staleness budgets that keep slow tasks from bottlenecking rollout (EAAS); and (iii) curriculum–staleness coupling that tightens the staleness budget as difficulty rises (CSC). Our mechanistic analysis validates the theoretical predictions: the "safe zone" of gradient coherence follows the derived exponential boundary. On code generation benchmarks, GAS improves Pass@1 from 39.7% to 60.1% while training at 2.3×2.3\times the throughput of synchronous GRPO, showing that matching each task's staleness budget to its difficulty (and thus to its loss-landscape curvature) makes asynchronous RL both fast and stable.