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

📄 arXiv: 2505.24378v1 📥 PDF

作者: Yilun Kong, Guozheng Ma, Qi Zhao, Haoyu Wang, Li Shen, Xueqian Wang, Dacheng Tao

分类: cs.LG, cs.AI

发布日期: 2025-05-30

备注: ICML 2025

💡 一句话要点

提出M3DT：通过混合专家决策Transformer掌握大规模多任务强化学习

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 多任务强化学习 混合专家模型 决策Transformer 离线强化学习 任务扩展性

📋 核心要点

  1. 现有MTRL方法在任务数量扩展性方面存在瓶颈，简单增加模型参数无法有效解决性能下降问题。
  2. M3DT通过引入混合专家(MoE)机制，并结合三阶段训练策略，提升模型参数利用率和训练效率。
  3. 实验表明，M3DT在扩展到160个任务时仍能保持卓越性能，验证了其优异的任务可扩展性。

📝 摘要（中文）

尽管离线多任务强化学习(MTRL)领域最近取得了进展，并充分利用了Transformer架构的强大功能，但大多数方法都侧重于有限数量的任务，扩展到极其庞大的任务仍然是一个巨大的挑战。本文首先重新审视了任务数量对当前MTRL方法的影响，并进一步揭示了简单地扩展参数不足以抵消随着任务数量增加而导致的性能下降。基于这些见解，我们提出了M3DT，一种新颖的混合专家(MoE)框架，通过进一步释放模型的参数可扩展性来解决任务可扩展性问题。具体来说，我们增强了agent的架构和优化，使用MoE加强了Decision Transformer (DT)骨干网络，以减少参数子集上的任务负载，并引入了一个三阶段训练机制，以促进高效训练并获得最佳性能。实验结果表明，通过增加专家数量，M3DT不仅在固定任务数量下随着模型扩展而持续提高其性能，而且还表现出卓越的任务可扩展性，成功扩展到160个任务并具有卓越的性能。

🔬 方法详解

问题定义：现有的多任务强化学习方法在扩展到大规模任务时面临性能瓶颈。简单地增加模型参数并不能有效解决随着任务数量增加而导致的性能下降问题。这表明模型容量的利用率不高，或者训练方法无法有效地利用大规模参数。

核心思路：论文的核心思路是通过引入混合专家(MoE)机制，将不同的任务分配给不同的专家子网络处理，从而降低每个专家需要处理的任务负载。同时，设计专门的训练策略，以确保MoE模型能够有效地学习和泛化到新的任务。

技术框架：M3DT基于Decision Transformer (DT)架构，并对其进行了扩展，加入了MoE层。整体框架包含以下几个关键模块：1) Decision Transformer Backbone：负责序列建模和决策；2) Mixture-of-Experts (MoE)层：包含多个专家网络，每个专家负责处理一部分任务；3) Routing Mechanism：负责将不同的任务分配给不同的专家；4) 三阶段训练机制：包括预训练、专家选择和微调三个阶段。

关键创新：M3DT的关键创新在于将MoE机制引入到Decision Transformer中，并设计了专门的三阶段训练策略。与传统的MTRL方法相比，M3DT能够更好地利用模型参数，提高任务可扩展性。此外，三阶段训练机制能够有效地避免MoE模型训练中的一些常见问题，例如专家之间的负载不均衡。

关键设计：M3DT的关键设计包括：1) MoE层的具体实现方式，例如专家网络的数量、专家网络的结构等；2) Routing Mechanism的设计，例如使用Gating Network进行任务分配；3) 三阶段训练机制的具体步骤和参数设置，例如预训练的epochs、专家选择的策略、微调的学习率等。此外，损失函数的设计也至关重要，需要考虑如何平衡不同专家之间的负载，以及如何提高模型的泛化能力。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，M3DT在多任务强化学习任务上显著优于现有方法。通过增加专家数量，M3DT不仅在固定任务数量下随着模型扩展而持续提高其性能，而且还表现出卓越的任务可扩展性，成功扩展到160个任务并具有卓越的性能。具体性能提升数据需要在论文中查找。

🎯 应用场景

M3DT在机器人控制、游戏AI、自动驾驶等领域具有广泛的应用前景。它可以用于训练能够同时完成大量不同任务的智能体，从而提高智能体的通用性和适应性。例如，在机器人控制领域，可以使用M3DT训练一个能够同时完成抓取、导航、操作等多种任务的机器人。

📄 摘要（原文）

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.