LAR-MoE: Latent-Aligned Routing for Mixture of Experts in Robotic Imitation Learning

TL;DR

LAR-MoE introduces a two-stage latent-aligned routing framework for mixture-of-experts in robotic imitation learning, enabling unsupervised skill discovery and improved task generalization across heterogeneous dynamics.

Contribution

The paper presents a novel two-stage framework that decouples skill discovery from policy learning, using latent-aligned routing to improve expert specialization without supervision.

Findings

Achieves 95.2% success rate on LIBERO benchmark with 150M parameters.

Matches supervised MoE performance on surgical tasks without phase annotations.

Enables zero-shot transfer to ex vivo tissue.

Abstract

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.