Learning Diverse Skills for Behavior Models with Mixture of Experts

2026-01-18

AI Summary

This work addresses the performance degradation in multitask imitation learning caused by interference among tasks. The authors propose a framework that jointly trains a mixture-of-experts (MoE) architecture with an energy-based model, enabling each expert to specialize in distinct subregions of the observation space and thereby learn diverse, disentangled skills. The energy model guides experts to focus on specific observation distributions, promoting specialization without requiring architectural overhauls. This approach is designed as a plug-and-play module compatible with existing imitation learning systems. Evaluated on multiple real-world robotic manipulation tasks, the method significantly outperforms current state-of-the-art techniques and demonstrates superior data efficiency and knowledge transfer capability when fine-tuning to novel tasks.

Technology Category

• Machine Learning: Mixture of Experts (MoE)
• Intelligent Robots: Multimodal Perception & Sensor Fusion
• Multiagent Systems: Multiagent Learning

Application Category

• Responsible Web: Machine-in-the-loop, human agency and autonomy
• Economics, Online Markets and Human Computation: Architectures and workflows that use LLMs for crowd work
• User Modeling, Personalization and Recommendation: Explainable and interpretable methods for personalization

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.

Problem

• imitation learning
• multi-task learning
• task interference
• behavior models
• robotic manipulation

Innovation

• Mixture of Experts
• Imitation Learning
• Energy-based Models
• Multi-task Learning
• Behavior Modeling