Predictive Routing Replay for MoE-based LLM Reinforcement Learning

Mixture of Experts (MoE) Large Language Models (LLMs) achieve strong performance at scale. However, reinforcement learning (RL) on MoE-based LLMs often suffers from training instability. A core issue is off-policy optimization, where current training policy updates rely on trajectories generated by stale behavior policies. In MoE models, \textit{router drift} acts as the routing form of this gap because expert activations can change across model updates, causing unstable importance sampling weights in PPO-style algorithms. Routing Replay (R2) mitigates this issue by freezing the replay route within each reasoning trajectory, but it ignores how the router evolves under off-policy updates and thus causes \textit{router staleness}. To address this limitation, we propose \textbf{Predictive Routing Replay (PR2)}, which augments each router with a lightweight evolution predictor that learns to anticipate short-horizon router evolution. During the rollout phase, we use the predictive routing distribution to apply top-$k$ routing, enabling gradients to reach experts that are likely to become active after updates. During the training phase, we replay the resulting predicted route to retain consistency for stable importance estimation. Theoretical analysis and experiments show that PR2 reduces routing mismatch, improves RL stability, and improves benchmark accuracy.