Stabilizing MoE Reinforcement Learning by Aligning Training and Inference Routers

Reinforcement learning (RL) has emerged as a crucial approach for enhancing the capabilities of large language models. However, in Mixture-of-Experts (MoE) models, RL training tends to be unstable and can even collapse entirely. We analyze the training-inference consistency of MoE models and identify a notable discrepancy in routing behaviors between the two phases. To address this issue, we propose Rollout Routing Replay (R3), a novel and effective method that records routing distributions from the inference engine and replays them during training. R3 significantly reduces training-inference policy KL divergence and mitigates extreme discrepancies without compromising training speed. Extensive experiments on various settings confirm that R3 succeeds in stabilizing RL training, preventing collapse and outperforming strong baselines. R3 is orthogonal to most policy optimization algorithm improvements, allowing it to be used in conjunction with them. We believe this work can offer a new solution for stabilizing RL in MoE model.

Today's most powerful AI models are often built like a team of specialists: instead of one giant brain doing everything, a "router" sends each piece of a problem to whichever experts are best suited to handle it. This design, used in many top models, makes them much more efficient. But it creates a tricky problem — when we try to improve these models through trial-and-error training (reinforcement learning), training becomes very unstable and often breaks down. We found the cause. A model writes its answers on one software system but learns from them on another, and tiny number differences cause the router to pick different experts in each step. As a result, the experts that actually did the work aren't always the ones that get the feedback. Our method, Rollout Routing Replay, fixes this by recording which experts the router used while writing an answer and reusing exactly those choices during training. It follows a simple rule — whoever does the work gets the feedback — and prevents breakdowns with almost no added cost.