LAR-MoE: Latent-Aligned Routing for Mixture of Experts in Robotic Imitation Learning
Authors: Ariel Rodriguez, Chenpan Li, Lorenzo Mazza, Rayan Younis, Ortrun Hellig, Sebastian Bodenstedt, Martin Wagner, Stefanie Speidel
Organizations: Department of Translational Surgical Oncology, NCT/UCC Dresden, a partnership between DKFZ, MF UKD, TUD, HZDR Dresden, Germany. · Cluster of Excellence-CeTI, TUD, Germany. · Department of Visceral, Thoracic and Vascular Surgery, Faculty of Medicine and University Hospital Carl Gustav Carus, TUD, Germany.
Imitation learning enables robots to acquire manipulation skills from demonstrations, yet deploying a policy across tasks with heterogeneous dynamics remains challenging, as models tend to average over distinct behavioral modes present in the demonstrations. Mixture-of-Experts (MoE) architectures address this by activating specialized subnetworks, but requires meaningful skill decompositions for expert routing. We introduce Latent-Aligned Routing for Mixture of Experts (LAR-MoE), a two-stage framework that decouples unsupervised skill discovery from policy learning. In pre-training, we learn a joint latent representation between observations and future actions through student-teacher co-training. In a post-training stage, the expert routing is regularized to follow the structure of the learned latent space, preventing expert collapse while maintaining parameter efficiency. We evaluate LAR-MoE in simulation and on hardware. On the LIBERO benchmark, our method achieves a 95.2% average success rate with 150M parameters. On a surgical bowel grasping and retraction task, LAR-MoE matches a supervised MoE baseline without requiring any phase annotations, and transfers zero-shot to ex vivo porcine tissue. Our findings suggest that latent-aligned routing provides a principled alternative to supervised skill decomposition, enabling structured expert specialization from unlabeled demonstrations.
Figures & tables
Fig. 2 : Overview of LAR-MoE. In pre-training, our method learns a joint latent representation of observations and future actions via student–teacher co-training. In post-training, we anchors expert routing to the learned latent structure by freezing the student model and using a regularization strategy. The action experts are implemented using a simple transformer decoder architecture. The language promts are encoded using MiniLM-L6.
Fig. 3 : Expert activation over time during an ex-vivo bowel retraction rollout. The heatmap displays the activation weights of all experts at each timestep ( ≈33ms temporal resolution). The color bar labeled Po denotes the expert with the highest activation at each timestep chosen by the router network, illustrating the temporal specialization and switching behavior of the mixture-of-experts policy. The color bar labeled Hu represents a human-annotated task phase segmentation of the same rollout.
Fig. 4 : Ablation studies on the LIBERO benchmark [ liu2023libero ] . (a) Effect of freezing the student encoder (+F) and applying latent-alignment regularization (+R) on LAR-MoE16 success rate. (b) Success rate as a function of expert count.
Fig. 5 : Illustration of the surgical bowel grasping and retraction task. The first frame shows the robot instrument, awaiting the start of the task. The remaining figures correspond to individual phases, numbered in the upper-left corner. The robot-controlled instrument is labeled in white ( Robot ) and the surgeon-operated instrument in yellow ( Surgeon ). Phase 1: the surgeon indicates the target grasping region. Phase 2: the robot grasps the indicated bowel segment. Phase 3: the robot holds its position while the surgeon grasps the opposite end. Phase 4 is not shown, as it represents the transitional stretching motion between Phase 3 and Phase 5. Phase 5: the robot maintains tension throughout the remainder of the procedure.
Policy
Phase Success
E2E
Reach
Grasp
Retract
ACT
16/20
13/20
12/20
10/20
ACT + Supervised MoE
20/20
20/20
19/20
17/20
LAR-MoE16 (Ours)
20/20
18/20
18/20
17/20
TABLE II : On-hardware success rates for the bowel grasping and retraction task on a bowel phantom.
Fig. 6 : Spatial expert assignment during 10 rollouts on (a) phantom and (b) ex vivo porcine tissue. Each color corresponds to a different expert and indicates the most active expert at the respective spatial position in millimeters . Expert colors are shared across both visualizations to enable direct comparison.
We consider the problem of learning compositional robot policies end-to-end from expert demonstrations, without any pre-specified notion of task decomposition or hierarchy. We ask whether a VLA trained with a simplified Mixture-of-Experts (MoE) action head can emergently learn to decompose tasks into reusable, interpretable primitives. We find that learned experts are heavily reused across tasks and consistently correspond to qualitatively distinct low-level behaviors, suggesting that the router implicitly learns to perform high-level sequencing while experts serve as compositional primitives. Our MoE matches the task performance of a monolithic baseline while demonstrating meaningful expert specialization, a step toward modular, interpretable robot policies that emerge from data alone.
Department of Computer Science, Princeton University, Princeton, NJ, USA · Department of Electrical and Computer Engineering, Princeton University, Princeton, NJ, USA
Diffusion-based policies have established a new standard for precise robotic manipulation but face a critical scalability bottleneck: high-performance models are computationally expensive, while lightweight alternatives often fail to generalize across diverse multi-task environments. Mixture-of-Experts (MoE) architectures offer a promising path to efficiency by activating only a subset of parameters. However, existing MoE routing mechanisms typically rely on low-level noise or latent statistics, ignoring the compositional nature of manipulation tasks. This can fragment reusable behaviors across experts, limiting interpretability and transferability. We introduce Semantically Structured Mixture-of-Experts Diffusion Policy (SMoDP) for compositional robotic manipulation, a framework that grounds expert specialization in semantic task structure. SMoDP leverages a lightweight, inference-time skill predictor, supervised by offline annotations from Vision-Language Models (VLMs), to route action chunks to experts specialized for specific behavioral phases. To ensure robust assignment, we propose a dual contrastive alignment strategy that grounds multi-modal observations in language-defined skill semantics (Inter-modal) while enforcing routing consistency across visually distinct but functionally related behaviors (Intra-modal). Our approach outperforms representative diffusion and MoE-based baselines on multi-task benchmarks with significantly improved parameter efficiency and demonstrates effective compositional transfer to novel tasks through parameter-efficient fine-tuning. Project website: https://deng-cy20.github.io/SMoDP/
Chengyu Deng, Guanqi Chen, Yizhou Chen +4
The University of Hong Kong · Southern University of Science and Technology
While MoE augments VLA via expert specialization, router suffers from ineffective expert routing owing to the kinematic heterogeneity of actions across manipulation tasks and, even worse, the unavailability of the kinematic signals at inference time. In this work, we first observe that most semantically distinct manipulation tasks reduce to multiple kinematic archetypes. Motivated by this finding, we propose Kinematics-supervised explicit routing (KinRT), a new paradigm that shifts from implicit, observation-driven expert routing to explicit, kinematics-guided expert dispatching. Specifically, we perform kinematic clustering on action trajectories into multiple kinematically coherent groups, whose IDs serve as ground truth to supervise the training of the router; at inference time, the router dispatches experts only using visual-language observations, without any reliance on action kinematics. KinRT actually introduces an asymmetric bridging mechanism that distills the task kinematics from the action space in training into the observation space at inference. In addition, to assess KinRT's cross-platform generalization, we build an economical, Do-It-Yourself robot (DIYRobot) platform from scratch using 3D-print technology (< 2,000USD). Extensive experiments demonstrate KinRT's superiority over both dense and MoE-featured VLAs by more than 23.26% on RoboTwin benchmark and 20.27% on our introduced DIYRobot platform. Our code and DIYRobot platform will be open-sourced.
Tianhang Yang, Yanze Zheng, Junjie Wang +3
Tsinghua Shenzhen International Graduate School, Tsinghua University, Shenzhen, China · Pengcheng Lab, Shenzhen, China