Tra-MoE: Learning Trajectory Prediction Model from Multiple Domains for Adaptive Policy Conditioning

TL;DR

This paper introduces Tra-MoE, a sparse Mixture of Experts model for trajectory prediction across multiple domains, enhancing generalization and control in robotic systems through adaptive policy conditioning.

Contribution

We propose Tra-MoE, a sparse MoE architecture with adaptive policy conditioning, improving trajectory prediction from diverse out-of-domain data for robotic applications.

Findings

Tra-MoE outperforms dense models in accuracy and generalization.

Sparse MoE maintains constant FLOPs while scaling to large parameters.

Adaptive policy conditioning improves action prediction flexibility.

Abstract

Learning from multiple domains is a primary factor that influences the generalization of a single unified robot system. In this paper, we aim to learn the trajectory prediction model by using broad out-of-domain data to improve its performance and generalization ability. Trajectory model is designed to predict any-point trajectories in the current frame given an instruction and can provide detailed control guidance for robotic policy learning. To handle the diverse out-of-domain data distribution, we propose a sparsely-gated MoE (\textbf{Top-1} gating strategy) architecture for trajectory model, coined as \textbf{Tra-MoE}. The sparse activation design enables good balance between parameter cooperation and specialization, effectively benefiting from large-scale out-of-domain data while maintaining constant FLOPs per token. In addition, we further introduce an adaptive policy conditioning technique by learning 2D mask representations for predicted trajectories, which is explicitly aligned with image observations to guide action prediction more flexibly. We perform extensive experiments on both simulation and real-world scenarios to verify the effectiveness of Tra-MoE and adaptive policy conditioning technique. We also conduct a comprehensive empirical study to train Tra-MoE, demonstrating that our Tra-MoE consistently exhibits superior performance compared to the dense baseline model, even when the latter is scaled to match Tra-MoE's parameter count.