Research map / Architectures and efficiency / Mixture of experts
Mastering Massive Multi-Task Reinforcement Learning via Mixture-of-Expert Decision Transformer
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 M3DT, a mixture-of-experts Prompt Decision Transformer trained in three stages (backbone, task-group experts, router) to scale offline multi-task RL to 160 control tasks.
Problem
Offline multi-task RL with Decision Transformers degrades as the number of tasks grows (to 160), and naively scaling shared parameters hits a performance ceiling due to gradient conflicts.
Approach
Augment each Prompt-DT transformer block with a mixture of task-group experts plus a router, and train in three stages (frozen-free backbone up to early stopping, experts on task subsets with backbone frozen, router with everything else frozen) so each expert sees few tasks and gradient conflicts are reduced.
Key results
- M3DT-Gradient reaches 78.21 normalized score on 160 tasks versus 71.65 for PromptDT-Large (173M) and 72.80 for HarmoDT-Large.
- On 80 tasks M3DT-Gradient scores 80.66 vs 75.70 for PromptDT-Large, with 98.37M parameters.
- Ablation: removing the three-stage training drops M3DT to 71.90 and removing grouping to 67.34.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use M3DT with random task grouping instead of the strongest reported baseline | M3DT · 10 tasks | normalized score: 88.17 ± 1.41 -> 89.09± 1.20 | Helped |
| Use M3DT with gradient-based task grouping instead of the strongest reported baseline | M3DT · 80 tasks | normalized score: 75.70± 1.62 -> 80.66± 0.97 | Helped |
| Use M3DT with random task grouping instead of the strongest reported baseline | M3DT · 160 tasks | normalized score: 72.80± 2.89 -> 76.74± 0.94 | Helped |
| Add the three-stage training mechanism (ablation) | M3DT · 80 tasks | normalized score: 71.90± 0.70 -> 76.67± 0.29 | Helped |
| Add task grouping and train experts on their assigned task subsets (ablation) | M3DT · 80 tasks | normalized score: 67.34± 0.56 -> 76.67± 0.29 | Helped |
| Freeze experts during router training (ablation) | M3DT-R · 80 tasks | normalized score: 71.89± 0.63 -> 76.67± 0.29 | Helped |
| Freeze experts during router training (ablation) | M3DT-G · 80 tasks | normalized score: 71.88± 0.62 -> 77.89± 0.47 | Helped |
| Use standard FFN-sized experts rather than full Transformer experts (ablation) | M3DT-Gradient · 80 tasks | normalized score: 76.53 -> 77.89 | Helped |
| Increase the number of experts (ablation) | M3DT · 80 tasks and 160 tasks | normalized score: 11.2% and 11.7% improvement | Helped |
| Use Top-4 sparse routing instead of weighted combination of all experts (ablation) | M3DT · Different expert counts | normalized score | Hurt |
Surprising findings
- Simply increasing DT model size hits a performance ceiling beyond about 20M parameters in multi-task RL.
- Top-4 routing fails to scale with expert count, degrading performance as experts increase.
Limitations
- Fine-grained expert and router architecture not explored; experts are plain FFN copies.
- Held-out task generalization and continual learning are not evaluated.
- Scaling experts increases inference cost; Top-k sparse routing degrades performance.
Experimental setup
Datasets: Meta-World DMControl Ant-dir Cheetah-vel
Benchmarks: Meta-World DMControl Ant-dir Cheetah-vel
Metrics: normalized score success rate gradient similarity
Baselines: MTDT PromptDT HarmoDT
Builds on: Decision Transformer Prompt-DT HarmoDT Top-K routing
Closest papers by meaning
- Mixture-of-Experts Meets In-Context Reinforcement Learning (NeurIPS 2025 · similarity 0.90)
- Multi-Task Reinforcement Learning with Mixture of Orthogonal Experts (ICLR 2024 · similarity 0.90)
- Representation Learning Enables Scalable Multitask Deep Reinforcement Learning (NeurIPS 2026 · similarity 0.89)
- Mixtures of Experts Unlock Parameter Scaling for Deep RL (ICML 2024 · similarity 0.89)
- AMAGO-2: Breaking the Multi-Task Barrier in Meta-Reinforcement Learning with Transformers (NeurIPS 2024 · similarity 0.88)
- DaCe-DT: Data-Centric Offline Multi-Task Reinforcement Learning via Adaptive Prompts and Trajectory Correction for Heterogeneous Tasks (NeurIPS 2026 · similarity 0.88)
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.