# Mastering Massive Multi-Task Reinforcement Learning via Mixture-of-Expert Decision Transformer

arXiv: https://arxiv.org/abs/2505.24378 · HTML: https://mlsanity.com/paper/2505.24378/

## Key facts

- Cited by 15 papers in the map (data through 14 Oct 2026).
- Placed in MoE in reinforcement learning and robotics, one of 12 directions in Mixture-of-experts.
- Builds on 37 of its 41 references that are in the map.
- First posted on arXiv 30 May 2025.

## Placed in

- Mixture-of-experts / [MoE in reinforcement learning and robotics](https://mlsanity.com/moe/moe-in-reinforcement-learning-and-robotics/)
- Offline RL and imitation learning from data / [Decision transformers and sequence modeling](https://mlsanity.com/offlinerl/decision-transformers-and-sequence-modeling/)

## Builds on (most cited first, 30 of 37)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networks](https://mlsanity.com/paper/1703.03400/) (2017)
- [Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor](https://mlsanity.com/paper/1801.01290/) (2018)
- [Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer](https://mlsanity.com/paper/1701.06538/) (2017)
- [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://mlsanity.com/paper/2101.03961/) (2021)
- [Pre-train, Prompt, and Predict: A Systematic Survey of Prompting Methods in Natural Language Processing](https://mlsanity.com/paper/2107.13586/) (2021)
- [Offline Reinforcement Learning: Tutorial, Review, and Perspectives on Open Problems](https://mlsanity.com/paper/2005.01643/) (2020)
- [GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding](https://mlsanity.com/paper/2006.16668/) (2020)
- [Decision Transformer: Reinforcement Learning via Sequence Modeling](https://mlsanity.com/paper/2106.01345/) (2021)
- [Gradient Surgery for Multi-Task Learning](https://mlsanity.com/paper/2001.06782/) (2020)
- [Meta-World: A Benchmark and Evaluation for Multi-Task and Meta Reinforcement Learning](https://mlsanity.com/paper/1910.10897/) (2019)
- [DeepMind Control Suite](https://mlsanity.com/paper/1801.00690/) (2018)
- [A Generalist Agent](https://mlsanity.com/paper/2205.06175/) (2022)
- [Scaling Vision with Sparse Mixture of Experts](https://mlsanity.com/paper/2106.05974/) (2021)
- [Offline Reinforcement Learning as One Big Sequence Modeling Problem](https://mlsanity.com/paper/2106.02039/) (2021)
- [TD-MPC2: Scalable, Robust World Models for Continuous Control](https://mlsanity.com/paper/2310.16828/) (2023)
- [Conflict-Averse Gradient Descent for Multi-task Learning](https://mlsanity.com/paper/2110.14048/) (2021)
- [Multi-Game Decision Transformers](https://mlsanity.com/paper/2205.15241/) (2022)
- [Multimodal Contrastive Learning with LIMoE: the Language-Image Mixture of Experts](https://mlsanity.com/paper/2206.02770/) (2022)
- [Just Pick a Sign: Optimizing Deep Multitask Models with Gradient Sign Dropout](https://mlsanity.com/paper/2010.06808/) (2020)
- [Prompting Decision Transformer for Few-Shot Policy Generalization](https://mlsanity.com/paper/2206.13499/) (2022)
- [Multi-Task Reinforcement Learning with Soft Modularization](https://mlsanity.com/paper/2003.13661/) (2020)
- [Diffusion Model is an Effective Planner and Data Synthesizer for Multi-Task Reinforcement Learning](https://mlsanity.com/paper/2305.18459/) (2023)
- [M3ViT: Mixture-of-Experts Vision Transformer for Efficient Multi-task Learning with Model-Accelerator Co-design](https://mlsanity.com/paper/2210.14793/) (2022)
- [ProMP: Proximal Meta-Policy Search](https://mlsanity.com/paper/1810.06784/) (2018)
- [Merging Multi-Task Models via Weight-Ensembling Mixture of Experts](https://mlsanity.com/paper/2402.00433/) (2024)
- [Mixtures of Experts Unlock Parameter Scaling for Deep RL](https://mlsanity.com/paper/2402.08609/) (2024)
- [Multi-Task Reinforcement Learning with Mixture of Orthogonal Experts](https://mlsanity.com/paper/2311.11385/) (2023)

## Built on by (most cited first, 15 of 15)

- [Mixture-of-Experts Meets In-Context Reinforcement Learning](https://mlsanity.com/paper/2506.05426/) (2025)
- [ACE-Brain-0.5: A Unified Embodied Foundational Model for Physical Agentic AI](https://mlsanity.com/paper/2607.04426/) (2026)
- [Guided by the Experts: Provable Feature Learning Dynamic of Soft-Routed Mixture-of-Experts](https://mlsanity.com/paper/2510.07205/) (2025)
- [DyGRO-VLA: Cross-Task Scaling of Vision-Language-Action Models via Dynamic Grouped Residual Optimization](https://mlsanity.com/paper/2605.17486/) (2026)
- [ExToken: Structured Exploration for Efficient Vision-Language-Action Reinforcement Fine-tuning](https://mlsanity.com/paper/2607.12931/) (2026)
- [Mixture-of-World Models: Scaling Multi-Task Reinforcement Learning with Modular Latent Dynamics](https://mlsanity.com/paper/2602.01270/) (2026)
- [Generalizable Multi-Task Learning for Wireless Networks Using Prompt Decision Transformers](https://mlsanity.com/paper/2606.04328/) (2026)
- [Conditional Sequence Modeling for Safe Reinforcement Learning](https://mlsanity.com/paper/2602.08584/) (2026)
- [CORE-MTL: Rethinking Gradient Balancing via Causal Orthogonal Representations](https://mlsanity.com/paper/2606.02221/) (2026)
- [Representation Learning Enables Scalable Multitask Deep Reinforcement Learning](https://mlsanity.com/paper/2606.05555/) (2026)
- [CTS-MoE: Implicit Terrain Adaptation via Mixture-of-Experts for Perceptive Locomotion](https://mlsanity.com/paper/2606.19633/) (2026)
- [Beyond Isolation: Unlocking Reinforcement Learning Component Synergy for Sample-Efficient Continuous Control](https://mlsanity.com/paper/2608.07086/) (2026)
- [WiSDoM: Wireless Sparse Decision Transformer with Mixture-of-Experts for Multi-Task Mobile Network Optimization](https://mlsanity.com/paper/2609.00284/) (2026)
- [Hierarchical Wasserstein Merging for Multi-Domain Multi-Task Learning: From Specialists to a Generalist](https://mlsanity.com/paper/2609.06406/) (2026)
- [DaCe-DT: Data-Centric Offline Multi-Task Reinforcement Learning via Adaptive Prompts and Trajectory Correction for Heterogeneous Tasks](https://mlsanity.com/paper/2610.11085/) (2026)

## Abstract (arXiv)

Despite recent advancements in offline multi-task reinforcement learning (MTRL) have harnessed the powerful capabilities of the Transformer architecture, most approaches focus on a limited number of tasks, with scaling to extremely massive tasks remaining a formidable challenge. In this paper, we first revisit the key impact of task numbers on current MTRL method, and further reveal that naively expanding the parameters proves insufficient to counteract the performance degradation as the number of tasks escalates. Building upon these insights, we propose M3DT, a novel mixture-of-experts (MoE) framework that tackles task scalability by further unlocking the model's parameter scalability. Specifically, we enhance both the architecture and the optimization of the agent, where we strengthen the Decision Transformer (DT) backbone with MoE to reduce task load on parameter subsets, and introduce a three-stage training mechanism to facilitate efficient training with optimal performance. Experimental results show that, by increasing the number of experts, M3DT not only consistently enhances its performance as model expansion on the fixed task numbers, but also exhibits remarkable task scalability, successfully extending to 160 tasks with superior performance.
