Phase-Aware Mixture of Experts for Agentic Reinforcement Learning

Shengtian Yang; Yu Li; Shuo He; Yewen Li; Qingpeng Cai; Peng Jiang; Lei Feng

cs.AIFeb 19, 2026

Phase-Aware Mixture of Experts for Agentic Reinforcement Learning

Authors: Shengtian Yang, Yu Li, Shuo He, Yewen Li, Qingpeng Cai, Peng Jiang, Lei Feng

Organizations: School of Computer Science and Engineering, Southeast University, Nanjing, China · Kuaishou Technology, Beijing, China · College of Computing and Data Science, Nanyang Technological University, Singapore

Abstract

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.

Figures & tables

Figure 1: Single-Policy Networks vs. Our PA-MoE. (a) Single-policy networks exhibit severe simplicity bias: simple tasks (pick_and_place) occupy 75% of parameters while complex tasks (heat/cool/clean requiring multi-step tool interaction) receive only 5%. Parameter occupancy is measured as the fraction of training batches where each task category contributes

>

50% of the batch loss. (b) PA-MoE achieves balanced parameter allocation (

\sim

30% per expert) and uniformly high performance across all task difficulties (96%-92%-98%).

Figure 2: Routing Granularity Comparison. (a) Expert switches per episode across routing strategies. Token-level MoE causes excessive fragmentation (45 switches), while trajectory-level MoE is overly coarse (3 switches). PA-MoE strikes a balance with 8 switches per episode. (b) Visualization of phase-level expert assignment showing semantic coherence: the same expert (indicated by color) is maintained within each contiguous behavioral phase, with switches occurring only at phase boundaries.

Figure 3: PA-MoE Architecture. Upper panel: Phase-Aware Router (Sec. 3 ) processes observation $o_{t}$ and goal $g$ via cross-attention, and action history $h_{t}$ via LSTM, to select expert $k^{*}=\arg\max p_{i}$ . Balance loss $\mathcal{L}_{\text{bal}}$ ensures uniform expert utilization. Lower panel: Phase-Aware MoE Execution shows agent-environment interaction. The router selects expert $k^{*}$ from $K$ LoRA-based experts sharing a frozen base model (gray). The selected expert generates action via policy $\pi_{\text{exp}}^{k^{*}}(a_{t}|s_{t},h_{t},g)$ . Training optimizes both router and experts jointly via RL loss $\mathcal{L}_{\text{RL}}$ and diversity loss $\mathcal{L}_{\text{div}}$ (Sec. 3 ).

Closed-Source Model ‡
Type	Method	ALFWorld							WebShop
		Pick	Look	Clean	Heat	Cool	Pick2	All	Score	Succ.
Prompting	GPT-4o †	75.3	60.8	31.2	56.7	21.6	49.8	48.0	31.8	23.7
Prompting	Gemini-2.5-Pro †	92.8	63.3	62.1	69.0	25.4	58.7	60.3	42.5	35.9
Qwen2.5-1.5B-Instruct
Prompting	Qwen2.5 †	5.9	5.5	3.3	9.7	4.2	0.0	4.1	23.1	5.2

Table 1: Performance comparison on ALFWorld and WebShop benchmarks. For ALFWorld, we report the average success rate (%) for each subtask as well as the overall result. For WebShop, we report both the average score and the average success rate (%). We report mean and standard deviation over 3 random seeds for RL methods. Prompting results are deterministic single runs. “GiGPO w/o std” refers to variants without advantage standardization, respectively.

$K$	Pick	Pick2	Look	Heat	Cool	Clean	All
0	94.3	87.5	62.5	94.4	89.5	83.3	88.3
2	97.1	91.7	62.5	100.0	73.7	100.0	91.4
4 (ours)	97.1	91.7	75.0	100.0	89.5	95.8	93.8
6	91.4	70.8	62.5	100.0	84.2	91.7	85.9

Table 2: Ablation on number of experts (

K

K{=}0

uses a single shared LoRA adapter without routing. All six ALFWorld task types shown. Bold indicates best per column; underline indicates ties with best. We report the average success rate (%) for each subtask as well as the overall result. See Appendix F for

K{=}3,5

results and standard deviations.

Routing	Pick(%)	Heat(%)	All(%)
Token-level	88.6	86.1	85.7
Trajectory-level	94.3	91.7	88.5
Phase-level (ours)	97.1	100.0	93.8

Table 3: Comparison of routing granularities. Switches denotes average expert switches per episode, measured at the environment step level (i.e.,

\sum_{t}\mathbf{1}[z_{t}\neq z_{t+1}]

). Trajectory-level assigns one expert per episode at

t{=}0

based on initial observation, with minimal within-episode adaptation. See Appendix C for token-level baseline details.

Router Configuration	Success Rate (%)
w/o action history	89.2
w/o goal attention	90.5
w/o both	86.8
Full router (ours)	93.8

Table 4: Ablation on router components.

	E1	E2	E3	E4
Role	Explore	Manipulate	Navigate	Recover
Action (%)	73	82	64	71

Table 5: Emergent expert specialization on ALFWorld (role and activation frequency).

Appendix figures & tables20 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 4: Gradient conflict analysis. (a) Phase-specific gradients projected onto top-2 principal components show pairwise conflicts: Explore and Interact gradients form angles exceeding 90°, and no two phases produce aligned gradients. (b) Gradient conflict score throughout training, defined as the average negative cosine similarity between phase gradients. Single policy maintains high conflict (

>0.4

) while PA-MoE reduces conflict to near-zero by epoch 50. (c) Pareto analysis of simple vs. complex task performance. PA-MoE achieves Pareto optimality while GiGPO baseline suffers degraded complex task performance.

Figure 5: Entropy mismatch analysis. All entropy values reported in bits (log base 2). (a) Policy entropy across phases: optimal phase-specific entropy (blue), single policy entropy (coral), and PA-MoE entropy (green). (b) Action distribution modality for exploration (diffuse,

H{=}3.5

bits), interaction (peaked,

H{=}0.5

bits), and single policy (intermediate,

H{=}2.3

bits). (c) Entropy evolution over episode timesteps. (d) Absolute entropy deviation from optimal by phase.

Configuration	Success(%)	$\Delta$
Full (both losses)	93.8 ±2.4	–
w/o $\mathcal{L}_{\text{div}}$	88.3 ±3.1	$-5.5$
w/o $\mathcal{L}_{\text{bal}}$	87.5 ±2.8	$-6.3$
w/o both	86.5 ±3.5	$-7.3$

Appendix

Table 6: Impact of regularization terms on ALFWorld. Mean

\pm

std over 3 seeds.

Method	Success(%)	$\Delta$
GiGPO (baseline)	86.1 ±4.7	–
+ PCGrad	87.3 ±3.9	+1.2
+ GradNorm	86.8 ±4.2	+0.7
+ CAGrad	87.9 ±3.6	+1.8
GiGPO + PA-MoE	93.8 ±2.4	+7.7

Appendix

Table 7: Comparison with gradient surgery methods. Mean

\pm

std over 3 seeds.

$K$	Pick	Look	Heat	Cool	Clean	Pick2	All
0	94.3 ±1.8	62.5 ±5.1	94.4 ±3.2	89.5 ±4.1	83.3 ±4.8	87.5 ±3.9	88.3 ±2.1
2	97.1 ±1.2	62.5 ±4.8	100 ±0.0	73.7 ±5.6	100 ±0.0	91.7 ±3.4	91.4 ±1.9
3	95.7 ±1.5	83.3 ±5.2	97.9 ±2.1	85.1 ±4.3	94.6 ±3.1	87.5 ±4.2	91.7 ±2.0
4	97.1 ±1.4	75.0 ±3.2	100 ±0.0	89.5 ±4.4	95.8 ±6.4	91.7 ±4.1	93.8 ±2.4
5	94.8 ±1.6	70.8 ±4.1	98.2 ±1.8	87.3 ±4.8	93.4 ±5.2	89.6 ±3.8	91.2 ±2.2
6	91.4 ±2.1	62.5 ±5.4	100 ±0.0	84.2 ±5.1	91.7 ±4.9	70.8 ±5.8	85.9 ±2.8

Appendix

Table 8: Ablation on number of experts (

K

). Mean

\pm

std over 3 seeds. We report the average success rate (%) for each subtask as well as the overall result.

Architecture	Heat(%)	Cool(%)	All(%)
MLP only	83.3	68.4	84.2
GRU + CrossAttn	97.9	87.4	92.8
LSTM + CrossAttn	100	89.5	93.8
Transformer + CrossAttn	98.6	88.1	93.2
Bi-LSTM + CrossAttn	99.3	89.2	93.7

Appendix

Table 9: Router architecture comparison

Expert	Phase	Activation	Entropy (bits)
Expert 1	Exploration	73.2%	3.5
Expert 2	Manipulation	81.7%	0.5
Expert 3	Navigation	64.1%	1.5
Expert 4	Recovery	71.3%	2.0

Appendix

Table 10: Emergent expert specialization. Entropy in bits.

Subtask	Freq.	Baseline	PA-MoE	$\Delta$
Simple search	32%	84.3%	91.2%	+6.9%
Multi-attribute	28%	62.1%	78.6%	+16.5%
Option selection	25%	71.4%	84.9%	+13.5%
Price comparison	15%	49.2%	73.8%	+24.6%
Overall	100%	67.4%	82.3%	+14.9%

Appendix

Table 11: WebShop performance by subtask type

Algorithm	Task	Baseline	PA-MoE	$\Delta$
PPO	Clean	57.1%	50.0%	$-7.1$ %
PPO	Cool	46.4%	36.8%	$-9.6$ %
PPO	Pick2	47.4%	37.5%	$-9.9$ %
RLOO	Pick	88.3%	74.3%	$-14.0$ %
GRPO	Cool	59.7%	50.0%	$-9.7$ %
GRPO	Pick2	53.5%	45.8%	$-7.7$ %

Appendix

Table 12: Task-algorithm combinations with performance variations

$\lambda_{s}$	Val SR(%)	Sw/Ep
0.00	85.2	$\sim$ 36
0.01	90.6	$\sim$ 20
0.05 (ours)	93.8	$\sim$ 8
0.10	91.4	$\sim$ 5
0.20	87.5	$\sim$ 3

Appendix

Table 13: Sensitivity of PA-MoE to the switching penalty coefficient

\lambda_{s}

. Sw/Ep denotes average expert switches per episode.

Statistic	Value
Episodes exhibiting $A{\rightarrow}B{\rightarrow}A$ recurrence	85.4%
Average expert switches per episode	8.4
Average expert revisitations per episode	2.8
Average phase (segment) length	4.8 steps

Appendix

Table 14: Phase recurrence and segmentation statistics over 500 ALFWorld validation episodes (

\lambda_{s}=0.05

K{=}4

). A revisitation occurs when an episode contains a pattern

E_{i}\rightarrow E_{j}\rightarrow E_{i}

with

j\neq i

Configuration	LoRA Params	Val SR(%)
Vanilla GiGPO	0	86.1 ±4.7
Single LoRA $r{=}32$ ( $K{=}0$ )	1 $\times$	88.3 ±2.1
Single LoRA $r{=}128$	4 $\times$	91.4 ±1.6
PA-MoE ( $K{=}4$ , $r{=}32$ )	4 $\times$	93.8 ±0.8

Appendix

Table 15: Capacity-matched comparison on ALFWorld. “Params” is the LoRA parameter budget relative to a single

r{=}32

adapter.

Task	Phases	Single $r{=}128$ (%)	PA-MoE(%)
Pick (single-phase)	1	97.1	97.1
Heat (multi-phase)	3+	94.4	100.0
Cool (multi-phase)	3+	78.9	89.5
Look (multi-phase)	2+	50.0	75.0

Appendix

Table 16: Per-task comparison between capacity-matched single LoRA (

r{=}128

) and PA-MoE (

K{=}4

r{=}32

) on ALFWorld. PA-MoE’s advantage concentrates on multi-phase tasks.

Routing strategy	Sw/Ep	Val SR(%)
Token-level MoE	$\sim$ 45	85.7
Trajectory-level MoE	$\sim$ 3	88.5
Random switching ( $p{=}0.21$ )	$\sim$ 8.4	90.6
Fixed-interval ( $N{=}5$ )	$\sim$ 8	92.2
PA-MoE (learned phases)	$\sim$ 8.4	93.8

Appendix

Table 17: Frequency-matched routing baselines on ALFWorld. “Fixed-interval” switches expert every

N{=}5

steps. “Random (

p

)” switches at each step independently with probability

p

, set to match PA-MoE’s average frequency. Sw/Ep is average expert switches per episode.

Config	Description	Routing	Regularization	Sw/Ep	Val SR(%)
A	Token MoE (baseline)	token	none	$\sim$ 45	85.7
B	Token + switching penalty	token	$\mathcal{L}_{\text{switch}}$ only	$\sim$ 20	87.5
C	Token + all PA-MoE reg.	token	all	$\sim$ 15	88.2
D	Phase routing, no penalty ( $\lambda_{s}{=}0$ )	phase	all except $\mathcal{L}_{\text{switch}}$	$\sim$ 36	85.2
E	PA-MoE (full)	phase	all	$\sim$ 8	93.8

Appendix

Table 18: Controlled decomposition of routing granularity and regularization on ALFWorld. All regularization terms (

\mathcal{L}_{\text{switch}}

\mathcal{L}_{\text{bal}}

\mathcal{L}_{\text{div}}

, temperature annealing) are included or excluded uniformly across configurations.

Gradient Estimator	Val SR(%)
REINFORCE + STE (ours)	93.8
REINFORCE + Gumbel-Softmax ( $\tau{=}0.5$ )	93.6

Appendix

Table 19: Comparison of gradient estimators for the switching indicator. The primary expert-selection gradient is REINFORCE in both cases; only the estimator for the indicator term differs.

Router Input	Val SR(%)
Full (observation + goal + history)	93.8
Observation + goal only	92.2
History only	91.4
Observation only	88.7

Appendix

Table 20: Router input stream ablation on ALFWorld. Each row keeps the indicated input streams active and zero-pads the others.

Component	Per-step cost
Router forward (LSTM + cross-attention + MLP)	1.3 ms
LoRA expert switch (pointer swap)	0.007 ms
LLM generation ( $\sim$ 50 tokens for one action)	438 ms
Total overhead vs. LLM generation	0.30% wall-clock

Appendix

Table 21: Per-step inference overhead breakdown. “LoRA switch” is the cost of swapping the active LoRA adapter via a pointer swap; weights are pre-loaded.

Configuration	Bypass rate	Val SR(%)	Router latency
Always run router	0%	93.8	1.3 ms/step
Bypass at $\theta{=}0.9$	72%	93.8	0.85 ms/step ( $-$ 35%)

Appendix

Table 22: Confidence-thresholded inference bypass on ALFWorld. Bypass rate is the fraction of steps that skip the router forward; latency reduction is measured relative to always running the router.

Measurement	Within-segment	Across-segment	Uniform reference
Top-1 expert switching rate (consecutive tokens)	93%	—	98.3% (59/60)
Per-segment routing entropy (bits)	4.1 (out of 5.9)	5.0 (out of 5.9)	5.9
Fraction of uniform entropy	70%	85%	100%

Appendix

Table 23: Token-level routing fragmentation in Qwen1.5-MoE-A2.7B (24 MoE layers, 60 experts, top-4). “Within-segment” is computed inside a single semantic unit (e.g., one reasoning step or one function body). “Across-segment” is computed across segment boundaries.

Reinforcement Learning

Reinforcement learning studies how agents learn sequential decisions from rewards and interaction, balancing exploration, performance, safety, and adaptation.

See all 1,717 papers

Structuring MoE Expert Selection for Agentic Reinforcement Learning
Oct 5, 2026Bolian Li, Ting-Yao Hu, Cheng-Yu Hsieh +3Agentic RL Expert Routing
Apple · Purdue University
Expert-Space Exploration in MoE Reinforcement Learning
Sep 11, 2026Hongyi He, Zhenghao Lin, Xiao Liu +3Reinforcement Learning Expert Routing
Tsinghua University · Microsoft Research
Beyond Global Routing Aggregation: Phase-Aware Expert Merging for MoE Vision-Language Models
Aug 5, 2026Hongyu Zhang, Cheng Yan, Xiang Xia +1Vision-Language Models Mixture of Experts
University of Science and Technology of China, Hefei, China

Oct 5, 2026J/K move · Enter open · S save

Structuring MoE Expert Selection for Agentic Reinforcement Learning

Bolian Li, Ting-Yao Hu, Cheng-Yu Hsieh, Sanjoy Chowdhury+2

Apple · Purdue University

Agentic RL Expert Routing RL Post-Training Mixture-of-Experts Models

Open paper AskPDF arXiv

Long-horizon LLM agents are frequently implemented using sparse mixture-of-experts (MoE) models, yet the co-design of agentic behavior and MoE structures remains underexplored. In this work, we comprehensively study the connections between agentic post-training and MoE expert selection. In off-the-shelf MoE models, we observe expert selection exhibits a specialized structure that naturally aligns with agentic trajectories. Specifically, expert routing overlaps more between turns where the agent performs semantically similar operations (e.g., READ, UPDATE) than between turns with differing operations. However, standard RL algorithms ignore this specialization, allowing the MoE routing to go uncontrolled during training, which empirically limit task performance and inference efficiency. To address this, we introduce a hierarchical routing control framework for agentic tasks. We explicitly encourage turn-level expert selections to align with agentic operations while regularizing token-level expert selections to maintain local consistency. To resolve stability issues that arise during post-training with the proposed methods, we further introduce an entropy-gated control mechanism. Overall, our routing control framework achieves over 10-point improvements in success rate on all evaluated benchmarks. These results demonstrate that agentic trajectory structure provides an effective signal for optimizing MoE capacity during RL post-training.

Abstract

Reinforcement Learning

Explore similar work

Structuring MoE Expert Selection for Agentic Reinforcement Learning

Expert-Space Exploration in MoE Reinforcement Learning

Beyond Global Routing Aggregation: Phase-Aware Expert Merging for MoE Vision-Language Models