Research map / Label-efficient and robust learning / Recognition and long-tail
Tra-MoE: Learning Trajectory Prediction Model from Multiple Domains for Adaptive Policy Conditioning
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes Tra-MoE, a sparsely-gated top-1 mixture-of-experts trajectory prediction model trained on out-of-domain action-free videos plus in-domain demos, and an adaptive learnable 2D trajectory mask for policy conditioning, evaluated on LIBERO and a real dual-arm robot.
Problem
Trajectory prediction models for robot policy guidance do not effectively benefit from large-scale out-of-domain video data; naively adding such data degrades in-domain performance, and dense scaling does not help.
Approach
Replace several transformer blocks of the ATM track transformer with top-1 sparsely-gated MoE blocks so out-of-domain data can be jointly trained while keeping FLOPs per token constant, and condition the policy on a learnable-embedding 2D trajectory mask concatenated with the image channels.
Key results
- Tra-MoE with out-of-domain data reaches 61.4% average success on LIBERO vs 52.0% for dense baseline with the same extra data (57.6% in-domain only).
- Adaptive learnable mask raises ATM average success from 61.4% to 66.3% on LIBERO.
- On real-world tasks, average success goes from 38% (baseline) to 46% with MoE and 56% with adaptive trajectory mask.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add out-of-domain data to the dense trajectory model | Tra-baseline · LIBERO | average success rate: 57.6 → 52.0 | Hurt |
| Add out-of-domain data to the MoE trajectory model | Tra-MoE · LIBERO | average success rate: 51.8 → 61.4 | Helped |
| Use Tra-MoE instead of ATM with RLBench video data | trajectory model trained with RLBench video data · LIBERO | average success rate: 41.9 → 54.5 | Helped |
| Add router z-loss with λz = 1·10−4 (ablation) | Tra-MoE · LIBERO | average success rate: 56.9 → 61.4 | Helped |
| Add load-balancing loss with λlo−ba = 1·10−7 (ablation) | Tra-MoE · LIBERO | average success rate: 56.9 → 57.3 | Helped |
| Add noise to the gating logits (ablation) | Tra-MoE · LIBERO | average success rate: 56.9 → 55.8 | Hurt |
| Increase the expert count from 1 to 2 (ablation) | Tra-MoE · LIBERO | average success rate: 52.0 → 56.0 | Helped |
| Scale the dense baseline by increasing width (ablation) | Tra-baseline · LIBERO | average success rate: 61.4 → 48.4 | Hurt |
| Scale the dense baseline by increasing depth (ablation) | Tra-baseline · LIBERO | average success rate: 61.4 → 52.5 | Hurt |
| Use an adaptive mask with learnable trajectory-point embeddings (ablation) | ATM policy · LIBERO | average success rate: 61.4 → 66.3 | Helped |
| Replace the adaptive mask with a hand-drawn mask (ablation) | ATM policy · LIBERO | average success rate: 61.4 → 61.5 | No clear effect |
| Add human video data (ablation) | real-world trajectory model · Real-world tasks | average success rate | No clear effect |
Surprising findings
- Naively adding out-of-domain data to the dense baseline lowers in-domain success (57.6 to 52.0), while MoE with the same data raises it.
- Load-balancing loss degraded performance in the ablations, contrary to common MoE practice.
Limitations
- Sparse MoE with small in-domain data overfits and underperforms the dense model (51.8 vs 57.6).
- Adaptive mask improvements are uneven across suites; overlapping trajectory parts cause mapping errors on LIBERO-Goal.
- Real-world evaluation uses only 5 tasks with 50 demos each.
Experimental setup
Datasets: LIBERO LIBERO-90 RLBench real-world dual-arm robot demonstrations human videos
Benchmarks: LIBERO-Spatial LIBERO-Object LIBERO-Goal LIBERO-Long real-world dual-arm tasks
Metrics: average success rate
Baselines: ATM Tra-baseline (dense) RT-Trajectory hand-drawn mask
Builds on: ATM CoTracker ST-MoE Switch Transformer
Closest papers by meaning
- Efficient Robotic Policy Learning via Latent Space Backward Planning (ICML 2025 · similarity 0.87)
- STRAP: Robot Sub-Trajectory Retrieval for Augmented Policy Learning (ICLR 2025 · similarity 0.87)
- Predictive Inverse Dynamics Models are Scalable Learners for Robotic Manipulation (ICLR 2025 · similarity 0.86)
- RT-Trajectory: Robotic Task Generalization via Hindsight Trajectory Sketches (ICLR 2024 · similarity 0.86)
- DiffTORI: Differentiable Trajectory Optimization for Deep Reinforcement and Imitation Learning (NeurIPS 2024 · similarity 0.86)
- Foundation Policies with Hilbert Representations (ICML 2024 · similarity 0.85)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.