Phase-Aware Mixture of Experts for Agentic Reinforcement Learning

Reinforcement learning (RL) has equipped LLM agents with a strong ability to solve complex tasks. However, existing RL methods normally use a single policy network, causing simplicity bias where simple tasks occupy most parameters and dominate gradient updates, leaving insufficient capacity for complex tasks. A plausible remedy could be employing the Mixture-of-Experts (MoE) architecture in the policy network, as MoE allows different parameters (experts) to specialize in different tasks, preventing simple tasks from dominating all parameters. However, a key limitation of traditional MoE is its token-level routing, where the router assigns each token to specialized experts, which fragments phase-consistent patterns into scattered expert assignments and thus undermines expert specialization. In this paper, we propose Phase-Aware Mixture of Experts (PA-MoE). It first features a lightweight phase router that learns latent phase boundaries directly from the RL objective without pre-defining phase categories. Then, the phase router allocates temporally consistent assignments to the same expert, allowing experts to preserve phase-specific expertise. Experimental results demonstrate the effectiveness of our proposed PA-MoE. Code is available at https://anonymous.4open.science/r/PA-MoE-576C/.

AI agents powered by large language models can solve multi-step tasks, but they often learn simple and frequent actions more easily than harder actions that require careful planning. This paper proposes Phase-Aware Mixture of Experts (PA-MoE), a method that lets an agent use different specialized components for different stages of a task. For example, one component may handle exploration, while another handles object manipulation or recovery from mistakes. Instead of switching specialists for every word or token, PA-MoE keeps the same specialist active across a coherent task phase. This helps the agent learn more stable behaviors. Experiments on household and web-shopping tasks show that PA-MoE improves performance, especially on complex tasks requiring several coordinated steps.