ATG-MoE: Autoregressive trajectory generation with mixture-of-experts for assembly skill learning

TL;DR

ATG-MoE is an end-to-end autoregressive model with mixture-of-experts architecture that learns multiple assembly skills from demonstrations by integrating multi-modal inputs for robust and flexible robotic manipulation.

Contribution

It introduces a unified model that combines multi-modal feature fusion, autoregressive trajectory generation, and mixture-of-experts for efficient multi-skill learning in robotic assembly tasks.

Findings

Achieves 96.3% grasp success rate in simulation

Attains 91.8% overall success rate in assembly tasks

Demonstrates strong generalization and multi-skill integration

Abstract

Flexible manufacturing requires robot systems that can adapt to constantly changing tasks, objects, and environments. However, traditional robot programming is labor-intensive and inflexible, while existing learning-based assembly methods often suffer from weak positional generalization, complex multi-stage designs, and limited multi-skill integration capability. To address these issues, this paper proposes ATG-MoE, an end-to-end autoregressive trajectory generation method with mixture of experts for assembly skill learning from demonstration. The proposed method establishes a closed-loop mapping from multi-modal inputs, including RGB-D observations, natural language instructions, and robot proprioception to manipulation trajectories. It integrates multi-modal feature fusion for scene and task understanding, autoregressive sequence modeling for temporally coherent trajectory generation, and a mixture-of-experts architecture for unified multi-skill learning. In contrast to conventional methods that separate visual perception and control or train different skills independently, ATG-MoE directly incorporates visual information into trajectory generation and supports efficient multi-skill integration within a single model. We train and evaluate the proposed method on eight representative assembly skills from a pressure-reducing valve assembly task. Experimental results show that ATG-MoE achieves strong overall performance in simulation, with an average grasp success rate of 96.3% and an average overall success rate of 91.8%, while also demonstrating strong generalization and effective multi-skill integration. Real-world experiments further verify its practicality for multi-skill industrial assembly. The project page can be found at https://hwh23.github.io/ATG-MoE