Abstracting Robot Manipulation Skills via Mixture-of-Experts Diffusion Policies

TL;DR

The paper tackles the difficulty of scaling diffusion-based robot manipulation to multi-task settings. It introduces SMP, a diffusion-based mixture-of-experts policy that learns a compact state-adaptive orthogonal skill basis and uses sticky routing with adaptive expert activation to compose actions from a small, task-relevant subset of experts, enabling fast inference. The authors formulate a variational objective that couples reconstruction in the whitened skill space, gate regularization, and coefficient priors, and demonstrate that adaptive activation preserves policy quality while reducing computation. Empirically, SMP achieves higher multi-task and transfer success with markedly lower inference cost than large diffusion baselines, on both simulated bimanual tasks and real dual-arm experiments, and exhibits stable, phase-consistent skill reuse across tasks. These results suggest a practical path toward scalable, transferable multi-task manipulation by learning reusable skills once and activating only what is needed.

Abstract

Diffusion-based policies have recently shown strong results in robot manipulation, but their extension to multi-task scenarios is hindered by the high cost of scaling model size and demonstrations. We introduce Skill Mixture-of-Experts Policy (SMP), a diffusion-based mixture-of-experts policy that learns a compact orthogonal skill basis and uses sticky routing to compose actions from a small, task-relevant subset of experts at each step. A variational training objective supports this design, and adaptive expert activation at inference yields fast sampling without oversized backbones. We validate SMP in simulation and on a real dual-arm platform with multi-task learning and transfer learning tasks, where SMP achieves higher success rates and markedly lower inference cost than large diffusion baselines. These results indicate a practical path toward scalable, transferable multi-task manipulation: learn reusable skills once, activate only what is needed, and adapt quickly when tasks change.

Figures (10)

• Figure 1: Overview of SMP.Top: Bimanual rollout of “put card in drawer” with key steps (1)–(5). Middle: Skill decomposition by arm and phase: the state-adaptive orthonormal skill basis and sticky routing yields spatial specialization (left/right), and organizes behavior into pick, adjust, reach, and release with corresponding end-effector (EE) actions. Bottom: Gate values over time show sparse, phase-consistent activation—only a few experts are active per step for efficient sampling.
• Figure 2: Skill Mixture-of-Experts Policy (SMP) Training Framework.Left (a): During training, raw observations are encoded into state features, which generate an unconstrained matrix $W(s)$. A QR retraction produces a state-adaptive orthogonal basis $B(s)$. Actions are reconstructed via $B(s)(g\odot z)$, where $g$ are sticky-gated weights and $z$ are diffusion-based coefficients. The model is trained with reconstruction, diffusion, gate regularization, and alignment losses. Right (b): Illustration of the state-adaptive basis across timesteps: as the robot moves, the basis vectors adjust with the state, while sticky gates preserve consistent expert roles (e.g., translation and rotation).
• Figure 3: Multi-task learning in RoboTwin-2 and RLBench-2. SMP partitions bimanual control into an orthonormal skill basis and routes with sticky gates. Across tasks, the same experts are reused for left- and right-arm primitives and for pick–move–place phases, with few switches and long segments. Gate traces reveal sparse, phase-consistent activation, and cross-task skill reuse, indicating that actions are composed from a small, task-relevant subset of experts.
• Figure 3: Success Rate in Few-shot Transfer
• Figure 5: Real-robot experiments with four manipulation policies. Left: SMP executes 4 bimanual manipulation tasks. Right: Progress score $\uparrow$ of each task averaged in 10 trials.