Learning Diverse Skills for Behavior Models with Mixture of Experts

TL;DR

This work tackles multi-task imitation learning in robotics by introducing Di-BM, a Mixture-of-Experts policy where each expert has an energy-based observation distribution and a shared gating network. The approach enables automatic specialization of experts to distinct regions of the observation space, with a diffusion-policy action model that handles multi-modal behaviors. Empirical results in simulation and real hardware show Di-BM surpasses strong baselines and offers data-efficient transfer to new tasks through post-training. The method is designed to be plug-and-play with existing imitation-learning frameworks, and visualizations confirm that experts allocate to different task phases, illustrating learned primitive skill decomposition. Overall, Di-BM advances scalable, reusable skill learning for complex robotic manipulation.

Abstract

Imitation learning has demonstrated strong performance in robotic manipulation by learning from large-scale human demonstrations. While existing models excel at single-task learning, it is observed in practical applications that their performance degrades in the multi-task setting, where interference across tasks leads to an averaging effect. To address this issue, we propose to learn diverse skills for behavior models with Mixture of Experts, referred to as Di-BM. Di-BM associates each expert with a distinct observation distribution, enabling experts to specialize in sub-regions of the observation space. Specifically, we employ energy-based models to represent expert-specific observation distributions and jointly train them alongside the corresponding action models. Our approach is plug-and-play and can be seamlessly integrated into standard imitation learning methods. Extensive experiments on multiple real-world robotic manipulation tasks demonstrate that Di-BM significantly outperforms state-of-the-art baselines. Moreover, fine-tuning the pretrained Di-BM on novel tasks exhibits superior data efficiency and the reusable of expert-learned knowledge. Code is available at https://github.com/robotnav-bot/Di-BM.

Figures (9)

• Figure 1: Complex robotic tasks are hypothesized to be decomposed into a set of primitive skills, which are mastered by different experts.
• Figure 2: The overall policy consists of an encoder, a gating network, and an action model. During training, the gating network assigns appropriate data to each expert for specialization, whereas during inference it selects the most suitable expert to handle the current observation.
• Figure 3: Visualization of $\pi(e|o)$ across different tasks, where the shaded regions indicate the currently dominant expert.
• Figure 4: $\pi(o|e)$ under different $\beta$ settings. Here, $\pi(o|e)=\exp(g_{\phi}(o,e))/Z_e$, where $Z_e=\sum_{i=0}^{N-1}\exp(g_{\phi}(o,i))$ is estimated from 100 random samples in the dataset. A smaller $\beta$ leads all experts to "slack off".
• Figure 5: Performance on novel tasks with different training data ratios. "from scratch" is trained only on new data, "org pretrain" is fine-tuned from the pretrained original DiffusionPolicy, "vanilla MoE pretrain" is fine-tuned from the pretrained vanilla MoE variant, and "Di-BM pretrain" is fine-tuned from our pretrained model.