PAMAE: Phase-Aware-MoE Action Experts Towards Reliable Flow-Matching Vision-Language-Action Policies

2026-06-25

AI Summary

Existing vision-language-action (VLA) models struggle to reliably generate actions in multi-stage robotic manipulation tasks due to their reliance on a single shared action expert, which fails to capture stage-specific control patterns. To address this limitation, this work proposes PAMAE—a plug-and-play, phase-aware mixture-of-experts action module—that retains the pretrained VLA backbone while introducing a lightweight phase prediction head and a phase-aware routing mechanism. By integrating a sparse mixture-of-experts architecture with a routing alignment objective, PAMAE enables experts to specialize according to execution phases. A two-stage training strategy is employed to stabilize expert specialization, yielding up to a 9.2% improvement in task success rate on multi-stage simulated tasks. Ablation studies confirm the critical roles of phase-supervised routing and stage-wise optimization.

Abstract

Reliable action generation for multi-stage robotic manipulation remains challenging for Vision-Language-Action (VLA) models. While existing flow-matching VLA policies offer strong multimodal grounding and generalization, they typically employ a single shared action expert, limiting their ability to capture phase-specific control patterns across distinct execution stages. We propose a plug-and-play Phase-Aware Mixture-of-Experts Action Module (PAMAE), as a step towards more reliable phase-consistent action generation. PAMAE replaces the original flow-matching action expert with a sparse expert mixture while preserving the pretrained VLA backbone. PAMAE introduces a phase-aware router that leverages execution-phase cues to allocate action generation across experts, supported by a lightweight phase prediction head and a routing alignment objective. To stabilize specialization, we adopt a two-stage training scheme that first warms up the expert module under the standard flow-matching loss and then optimizes phase-consistent routing under auxiliary supervision. On multi-stage manipulation simulation tasks, PAMAE improves task success by up to \textbf{9.2\%} over strong VLA baselines. Further ablations show that both phase-supervised routing and staged optimization are essential for the observed gains. Our results highlight phase-consistent expert allocation as an effective mechanism for improving the reliability and action quality of flow-matching VLA policies.

Problem

• Vision-Language-Action
• multi-stage manipulation
• action generation
• phase-specific control
• flow-matching

Innovation

• Phase-Aware Routing
• Mixture-of-Experts
• Flow-Matching
• Vision-Language-Action
• Robotic Manipulation