Technical Summary: Abstraction of Robot Manipulation Skills via Mixture-of-Experts Diffusion Policies

Problem Statement

Diffusion-based policies have demonstrated remarkable success in single-task robot manipulation by modeling action generation as a denoising process. However, extending these policies to multi-task scenarios presents significant challenges. Current strategies often rely on scaling up model size and demonstration datasets to achieve generalization. This approach incurs prohibitive costs: oversized models suffer from slow inference speeds unsuitable for real-time control, and the required demonstration data grows exponentially with task diversity. Furthermore, existing Mixture-of-Experts (MoE) applications in robotics often fail to explicitly disentangle reusable manipulation skills, leading to entangled behaviors and poor interpretability. The core challenge is enabling effective multi-task generalization under moderate model sizes and low sampling latency while ensuring skills are reusable and transferable.

Methodology: Skill Mixture-of-Experts Policy (SMP)

The authors propose Skill Mixture-of-Experts Policy (SMP), a diffusion-based framework that learns a compact, orthogonal skill basis and composes actions from a sparse subset of experts.

1. State-Adaptive Orthogonal Skill Basis

Instead of using a fixed action basis, SMP learns a state-dependent orthonormal basis $B(s) \in \mathbb{R}^{d \times K}$ (where $K \ll d$).

• Construction: A lightweight neural network generates an unconstrained matrix $W(s)$, which is projected onto the Stiefel manifold using a differentiable thin-QR retraction with sign stabilization. This ensures the basis vectors remain orthonormal and evolve continuously with the robot state (e.g., arm pose, contact), preserving local disentanglement.
• Action Decomposition: Actions are decoded as $a_t = B(s_t)(g_t \odot z_t)$, where $g_t$ are gating weights and $z_t$ are skill coefficients. This orthogonal decomposition ensures that each expert contributes to a distinct, non-overlapping direction in the action space.

2. Sticky Routing and Global Usage

To address the issue of rapid, unstable switching between experts (chatter), SMP employs sticky Dirichlet-Markov dynamics for the gating mechanism.

• Temporal Consistency: The gate $g_t$ is modeled as a Markov process where the current state depends on the previous gate $g_{t-1}$ and a global usage vector $\vartheta$. This encourages piecewise-constant skill activations, aligning with the quasi-stationary phases of manipulation tasks (e.g., grasp, move, release).
• Global Usage: A global vector $\vartheta$ captures overall skill prevalence, preventing the policy from collapsing into a single expert while maintaining balanced utilization across tasks.

3. Variational Training Objective

SMP is trained via variational inference to maximize the evidence lower bound (ELBO), which decomposes into three components:

• Reconstruction Loss ($L_{recon}$): Ensures the synthesized action in the whitened basis matches the ground truth. Crucially, the training uses two coefficient targets: a stop-gradient target for the diffusion surrogate and a gradient-carrying target for reconstructing the action, allowing the skill basis $B(s)$ to be updated without destabilizing the diffusion process.
• Gate Regularization ($L_{gate}$): Aligns the amortized gate posterior with the sticky Dirichlet prior, enforcing temporal stickiness and global usage constraints.
• Coefficient Regularization ($L_{coeff}$): Aligns the skill coefficients with their diffusion priors.

4. Adaptive Expert Activation at Inference

To reduce computational cost, SMP does not evaluate all experts at every step.

• State-Only Router: A lightweight router predicts gate weights based solely on the state $s_t$.
• Sparse Selection: Experts are ranked by their "mass" ($m_i = \bar{g}_{t,i}^2$). The system activates only the top-$k$ experts or the smallest subset that captures a threshold (e.g., 95%) of the total mass. This yields sparse, state-dependent activation, significantly reducing active parameters and latency.

Key Contributions

1. SMP Framework: A diffusion-based MoE policy that explicitly abstracts reusable manipulation skills via a state-dependent orthonormal action basis with sticky routing.
2. Adaptive Activation Strategy: A dynamic selection mechanism that activates a compact subset of experts at inference, reducing computational cost while maintaining action sampling accuracy.
3. Empirical Validation: Comprehensive evaluation in simulation (RoboTwin-2, RLBench-2) and on a real dual-arm platform, demonstrating superior multi-task success rates and lower inference costs compared to strong diffusion baselines.

Experimental Results

The authors evaluated SMP against vanilla diffusion policies (DP, DP3), transformer-based policies (ACT), large foundation models (RDT), and other skill-abstraction or MoE baselines.

• Multi-Task Learning: SMP achieved consistently higher success rates across bimanual tasks (e.g., lifting trays, handovers, sweeping) compared to baselines. While scaling RDT (a 1.2B parameter model) improved success by only 19% over a smaller baseline, SMP achieved higher performance with a significantly smaller footprint.
• Efficiency: SMP reduced inference time and active parameters. It activated approximately 30% of its own parameters (roughly 7% of RDT's total parameters) while maintaining high success. In contrast, dense baselines and FFN-MoE variants suffered from higher latency due to the need to synchronize or evaluate all experts.
• Transfer Learning:
  • Few-Shot Adaptation: When fine-tuned on new tasks with only 10 demonstrations, SMP outperformed baselines by effectively reusing relevant experts without needing to retrain the entire backbone.
  • Skill Composition: In tasks requiring the recombination of learned skills (e.g., putting fries in a skillet), SMP's router-only fine-tuning successfully recomposed skills, whereas baselines exhibited blurred roles and inconsistent phases.
• Real-Robot Performance: On a PiPER dual-arm platform, SMP achieved the highest progress scores on four complex bimanual tasks (e.g., removing pen caps, pouring beans) with lower latency than diffusion baselines.

Significance and Claims

The paper claims that SMP provides a practical path toward scalable, transferable multi-task manipulation. By learning reusable skills once and activating only what is needed, SMP addresses the trade-off between model capacity and inference efficiency. The authors emphasize that the learned experts align with intuitive manipulation roles (left/right arm, pick/move/place phases) and that the adaptive activation mechanism offers a simple knob to trade off accuracy against latency. The work suggests that explicit skill disentanglement and phase-consistent routing are critical for robust bimanual control, offering a more efficient alternative to brute-force model scaling.

The authors acknowledge limitations, noting that the current study focuses on bimanual manipulation with relatively small backbones. Future work is proposed to scale SMP to larger models, broaden evaluation to single-arm and mobile manipulation, and conduct more extensive real-robot studies to assess robustness under sensing noise and domain shift.