Mastering Massive Multi-Task Reinforcement Learning via Mixture-of-Expert Decision Transformer
Authors: Yilun Kong
Guozheng Ma
Qi Zhao
Haoyu Wang
Li Shen
Xueqian Wang
Dacheng Tao
Venue: ICML 2025 poster (2025)
Keywords: multi-task RL
offline RL
Decision Transformer
Mixture-of-Experts
Primary area: reinforcement_learning->batchoffline

Abstract: 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.

Mastering Massive Multi- Task Reinforcement Learning via Mixture-of-Expert Decision Transformer
Yilun Kong'
Guozheng Ma Qi Zhao
Haoyu Wang'
Li Shen
Xueqian Wang'
Dacheng Tao
Abstract
90
Despite recent advancements in offline multi-task reinforcement learning (MTRL) have harnessed the powerful capabilities of the Transformer ar- chitecture, most approaches focus on limited number of tasks, with scaling to extremely mas- sive tasks remaining formidable challenge. In this paper, we first revisit the key impact of task numbers on current MTRL method, and further re- veal that naively expanding the parameters proves insufficient counteract the performance degra- dation as the number of tasks escalates. Building upon these insights, mixture-on mixture-of-expere 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 the agent, where we strengthen the Decision Trans- former (DT) backhone 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 scal- ability, successfully extending 160 tasks with superior performance.
Nomaalzed Score
05
MSDT
25
MTDT
PromotDT
HarmoDT
70
20
60
100
120
140
160
Task
Figure 1: Illustration of comparative analysis of normal- ized score across various task numbers. M3DT achieves unparalleled task scalability. An in-depth exploration of these results refers to Section
propose M3DT, novel
These models also prove valuable for multi-task RI. (MTRL) challenges, offering a high-capacity framework capable of accommodating task variances and assimilating knowledge from diverse datasets. Additionally, they pave the way for incorporating innovations from language modeling (Brown et al., 2020) into MTRL methodologies, unlocking new potential for cross-disciplinary advancements.
Drawing inspiration from large language models (LLM), where models harness remarkable generalization capabil- ity to address massive tasks, there's growing interest in the potential of training RL agent to master increasingly diverse tasks. However, the application of these high- se- quential models to massive multi-task RL presents consider- able challenges. Firstly, existing approaches exhibit limited scalability with respect to task numbers. With most studies confined dozens of tasks Atari or Meta- World (Lee et al., 2022; Hu et al., 2024), when scaled to larger number of tasks, their performance degrades significantly; while Gato (Reed et al., 2022) extends the research to over 600 tasks, its performance on control tasks remains suboptimal. Secondly, current research either overlooks the impact of parameter scaling, with most studies confined to very small models (Xu al., 2022). or overly relies on the inherent parameter scalability of the Transformer architecture. Al- though Gato (Reed et al., 2022) and Multi-Game DT (Lee et al., 2022) have experimentally demonstrated the perfor- mance enhancements with expanded model size, a compari- son between the increase in model size and the correspond- ing performance gains suggests that these approaches cannot be considered an efficient method parameter scaling.
1. Introduction
Recent developments, such as Decision Transformer (Chen et al., 2021) and Trajectory Transformer (Janner et al., 2021), have reframed offline reinforcement learning (RL) as se- quence modeling problem, showcasing their ability trans- form large-scale datasets into potent decision-making agents.
'Tsinghua University, China "Nnyang Technological University, Singapore 'Shenzhen Campus of Sun Yat-sen University. Corre- spondence to: Xueqian Wang <wang.xq www.tsinghua.cdu.c.. Li Shen mathshenniie @ gomi..com>.
Proceedings of the 1nnd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Mastering Massive Multi-Task Reinforcement Learning via Mixture-of-Expert Decision Transformer
To address the above challenges, we extend the research see- nario to encompass 160 simulation control tasks. We begin by investigating the pivotal role of task quantity model performance and gradient conflicts, followed by an explo- ration of the impact of model expansion on MTRL. With the evident phenomenon that model performance and gradient conflicts deteriorate as the task scales, our findings indicate that these declines are most pronounced when the task num- ber remains relatively low; once the number of tasks reaches sufficiently large, the performance degradation tends to be- come gradual and steady. Therefore, one of our key insights emerged when viewed from a reverse perspective: reduc- ing the learning task numbers to a sufficiently small scale can significuntly enhance the performance. In our study on model scaling, the results surprisingly reveal that simply increasing the model size rapidly hits the performance ceil- ing. Thus, naively expanding shared parameters to offset the performance degradation from an increasing task numbers proves ineffective. In contrast, expanding parameters while reducing the task numbers results in the most pronounced performance gains. So how can we really minimize the num- ber of tasks to be learned while efficiently scaling model parameters, thereby maximizing performance?
MTRL from the perspective of task numbers and model size, analyze the performance degradation and gradient conflicts with increasing task numbers, and identify the limited effectiveness of naively parameter scaling in handling massive multi-task scenarios. (Section 3)
2. Based on the above insights, we propose M3DT,, a novel framework that enhances the DT architecture with MoE, explicitly assign the grouped task subsets each expert through task grouping, and introduce a three-stage training mechanism for training each mod- ule without interference. By increasing the number of experts, we unlock the parameter scalability for mas- tering massive tasks. (Section 4)
3. We demonstrate the superior performance of M3DT through rigorous testing on a broad spectrum of task scales, analyze its functionality through extensive abla- tion studies, and verify its task scalability and parame- ter scalability. (Section 5)
2. Preliminary
Offline Reinforcement Learning. The goal of RL is to learn a policy no maximizing the expected return E 00 7 RR(s, a.)] in a Markov Decision Process (MDP) (S, A. P. R. y, da), with state space S, action space A, envi- ronment dynamics P(s'|s, a) S S - [0, 1]. reward function R : S x A - R. discount factor y [0. 1), and initial state distribution do (Sutton & Barto, 2018). In the offline setting (Levine et al., 2020), a static dataset D = {(s, a, r.), collected by a behavior policy na, is provided. Offline RL algorithms learn a policy entirely from this static offline dataset, without any online interactions with the environment.
Building upon these insights, we propose M3DT, a MoE- based DT for handling Massive Multi-task RL. The overall framework is shown in Figure 4. We make improvements in both the architecture and optimization of the RL agent. Specifically, we introduce mixxere-of-expert (MoE) in the DT backbone to achieve parameter separation, allowing the backbone to learn shared knowledge across all tasks, while each expert specializes in learning task-specific knowledge from distinct small task subset, which greatly simplifies the training of parameters in experts. This also unlocks ef- fective parameter expansion, as the number of experts is the most efficient way to scale models (Fedus et al., 2022). By increasing the number of experts, we can not only introduc- ing a large number of parameters, but reduce the task load on each parameter subset, which mutually reinforce the model performance. Furthermore, M3DT introduces three-stage training mechanism, sequentially optimize the DT backbone, each expert, and the router, respectively, which allows each module to explicitly learn specialized knowledge without interference, mitigating the severe gradient conflicts encoun- tered when scaling to massive tasks, meanwhile reducing the difficulty of MoE training. Experimental results show that, by increasing the number of experts, on one hand, M3DT can consistently enhance its performance as model expan- sion on the fixed task scale; on the other hand, it exhibits remarkable task scalability, successfully extending to 160 tasks with superior performance.
Multi-Tase RL. In multi-task RL, different tasks can have different reward functions, state spaces, and transition func- tions. Given a specific task T p(T), lusk-specified MDP can be defined as "TTTTT RT,, ; d. In- stead of solving a single MDP, the goal is to find an op- timal policy that maximizes expected return over all tasks: The static
" arg max, TT~p(T) Bain (E 1E==0777 dataset D is correspondingly partitioned into per-task sub- sets as D = Dr, where N is the number of tasks.
Prompt Decision Transformer. The integration of Trans- former Vaswani, 2017) in offline RL for sequence modeling has gained prominence in recent years. Prompt-DT (Xu et al., 2022) extends DT by using task-specific prompts to enhance multi-task learning and few-shot generaliza- tion. Unlike text-based prompt in NLP (Liu et al., 2023; Kong et al., 2025), Prompt-DT employs short trajectories as prompts, which consist of state, action, and return-to-go tu- ples (s","", .aa,f")) f*). providing directed guidance to RL agents with seo-shot demonstrations. Each element marked with the superscript * is relevant to the trajectory prompt. These
In summary, our research makes three significant contribu- tions to the field of MTRL:
1. We rethink the challenges of sequence modeling in
2

---

Mastering Massive Multi-Task Reinforcement Learning via Mixture-of-Expert Decision Transformer
trajectory prompts are much shorter than the task's horizon, encompassing essential information to only identify task, yet inadequate for task imitation. During training with of- fline data, Prompt-DT utilizes iiuput ( as input for each task Ti combining a K -tee trajectory prompt
85
Normalized Score
0.28
Gradient Similarity
80
0.26
T.C
Gradient Similarity
Normalized Score
T; with a normal K-step trajectory Ti,t- The trajectory is formulated as:
75
0.24
0.22
input = (fT,1 5,,1 oi,, KiK- Ki
70
0.20
Fi,-- 1+, Si,t- K+,, Miit-K-ll Pitt, Si,t, ds.t)-
65
The action a is predicted through a prediction head linked to the state token. The training objective aims to minimize the mean-squared loss:
0.18
0.16
10
20
40
80
120
160
Number of Tasks
LDT iuput DD:
K
1
Σ
(as.m - -.(Ts,))mmm T),m))?
(2)
Figure 2: With the number of tasks increases from 10 to 160, both model performance and gradient similarity experience a noticeable decline.
K++
Prompt-DT has become a widely used backbone in recent research on MTRL.
generally aligns with that of model performance, while it drops more rapidly when the task number is low, and, coun- terintuitively, levels off and shows minimal decline after exceeding 40 tasks. Performance Variance: For each run, we test different task sets with the same number of tasks, thus the standard deviation of normalized scores reflects the model's robustness to variations in task combi- nations. For a massive number of tasks (120 to 160), the standard deviation is small, as the variations in task sub- sets from the total 160 tasks are minimal. For a moderate number of tasks (20 to 80), the standard deviation is large, indicating that different task combinations significantly im- pact performance, as similar tasks are easier for learning. For few tasks (10), despite greater variability in sampled task subsets, the standard deviation is much smaller than that in moderate tasks, which underscores that when task number is sufficiently low, the learning process simple enough be affected by the inter-task relationships.
3. Rethinking DT with MTRL
In this section, we delineate two primary challenges of em- ploying DT in MTRL: the inability to scale with task number and model size, laying the groundwork for the motivation behind our method.
3.1. Limited Scalability of Task Numbers
The complexity of MTRL is significantly amplified with increasing task numbers, largely attributed to escalating gradient conflicts. These conflicts stem from optimizing a shared set of parameters across tasks with differing objec- tives, leading to compromises in task-specific optimization. To better understand the impact of task numbers, we use Prompt-DT and construct an offline dataset containing 160 tasks based on Meta-Was World (Yu et al., 2020b), DM Con- trol (Tassa et al., 2018), and Mujoco Locomotion (Todorov et al., 2012), investigating the performance trends as the task number scales from 10 to 160. We measure the model's ss per- formance across three benchmarks using normalized scores. The gradient conftict is assessed by the average cosine simi- larity between the aggregate gradient and the gradients of each task, where lower similarity indicates higher de- gree of gradient conflict. The implementation details are described in Appendix A.
Thus, reverse perspective reveals the key insight: reducing the learning task number, particularly to sufficiently small scale, can significantly enhance model performance. When the task number is sufficiently small, challenges such as gradient conflicts and inter-task relationships become effort- lessly manageable. However, given the predefined number of tasks must be learned, how to reduce the actual learning task burden?
During the expansion of the training tasks, we observe di- vergent trend in the performance across different task scales. The results presented in Figure 2 highlight three distinct phenomena: Normalized Score: In the clear trend where the performance degrades with increasing task numbers, the decline is pronounced when task number is relatively low (below 40 tasks), while it becomes much more gradual once the tasks reach a sufficiently large number (above 80 tasks). Gradient Conflicts: The decrease of gradient similarity
3.2. Limited Scalability of Model Size in MTRL
Since existing models struggle with large-scale tasks in MTRL, straightforward approach is to scale up the model size to enhance its capacity. The supervised learning com- munity has convincingly demonstrated that larger networks lead to improved performance, in particular for language models (Kaplan et al., 2020). To investigate whether this trend holds for the DT model in MTRL,, we conduct exper-
3

---

Mastering Massive Multi-Tass Task Reinforcement Learning via Mixture-of-Expert Decision Transformer
4. Methodology: M3DT
160
61.46
67.19
71.61
71.68
71.65
885
From the previous observations and discussions, our key insight is that through task grouping and parameter sepa- ration, we can assign much fewer tasks each parameter subset, leading to exceptional performance. And by expand- ing the number of parameter subsets while preserving their individual sizes, and simultaneously increasing the number of task groups, we can both reduce the task load on each parameter subset and enhance the model's capacity with more parameters, which enables a scalable solution to han- dle an increasing number of tasks by expanding model size without compromising performance.
120
63.09
70.33
72.49
72.88
73.14
Task Numbers
80
80
65.11
72.22
73.94
75.66
75.70
70.22
74.74
78.17
79.19
78.51
-70
20
75.65
81.12
83.57
84.63
83.95
65
10
82.41
86.01
87.12
88.46
88.17
Based on this insight, we introduce M3DT, which includes two key enhancements in both the architecture and optimiza- tion of the agent to improve the task and parameter scal- ability in MTRL. The first enhancement enables efficient parameter separation and expansion to reduce task load, achieved by incorporating an MoE architecture alongside the FFN in Prompt-DT, as detailed in Section 4.1. While relying solely on the advantages of the MoE structure proves ineffective to solve massive MTRL, the second enhancement includes a task grouping process and three distinct training stages that optimize the agent's learning process, especially for ensuring each parameter subset can effectively handle a specific task subset without interference, as outlined in Sec- tion 4.2. The overall framework of our method is illustrated in Figure 4.
1.47M
5.29M
20.02M
77.78M
173.30M
Model Parameters
Figure 3: As the model size increases, model performance rapidly reaches its ceiling and ceases to improve thereafter. A more detailed domain-wir analysis can be seen in Ap- pendix B.I.
iments with Prompt-DTT, a language model-based method trained in supervised learning for MTRL, of varying sizes across different task scales. Detailed implementation can be found in Appendix A.5. Results in Figure 3 highlight some of the surprising phenomena, which contradict the behaviors typically observed in previous research (Lee et al., 2022; Reed et al., 2022):
In MTRL, increasing the model size of DT swiftly hits the performance ceiling, preventing sustained improve- ments. The horizontal comparison in Figure shows that expanding the model parameters is effective within lim- ited range across all task scales; once the model exceeds 20M parameters, further scaling yields no meaningful im- provements. This is clearly inefficient, as scalability is a key advantage of the Transformer architecture, and the per- formance of approaches relying solely on small models becomes increasingly constrained as the number of tasks scales. This leads to another limitation: Scaling up the model size can not effectively mitigate the performance degradation caused by the increase in task numbers. As shown by the clear decline from bottom-left to top-right in Figure 3, larger task numbers prevent the model from achiev- ing the strong performance seen with fewer tasks, even with increased parameters. more detailed rethinking of who blocks DT from scaling is presented in Appendix B.2.
4.1. Architecture: MoE for reducing Task Load
We propose use an MoE architecture with scalable experts for massive tasks. The MoE structure is characterized by its composition of N modular experts and router, OMAB [O,,.., ON. 0, }. where each expert 0, is responsible for a distinct task subset and the router O,, dynamically assigns weights for these experts without the need for exact task IDs. Accordingly, the parameters of a given expert are updated only by gradients from its corresponding small task subset, fully leveraging the remarkable learning capability for fewer tasks. Mcanwhile, by increasing the number of experts, we can further decrease the task load assigned to each expert, thereby effectively alleviating the severe gradient conflicts caused by the overwhelming task numbers and enhancing the performance.
The complete architecture of M3DT is illustrated in the right of Figure 4. We employ the complete Prompt-DTT (Xu et al., 2022) the backbone. As Appendix B.2 have investigated the crucial role of the feed-forward network (FFN) in trans- former module in multi-task settings, we incorporate the MoE architecture by augmenting each Transformer block with additional experts and the router alongside the pre- served FFN, which can maintain the learned shared knowl- edge from all tasks in the backbone model. We use networks
Despite these discouraging phenomena, the heatmap reveals the most pronounced performance increase from the top- left to the bothom-right. Building on the insights from the previous section, a natural question arises:
How can we effectively minimize the number of tasks to be learned, while efficiently scaling model parameters, thereby maximizing performance?
Key Insight:
4

---

Mastering Massive Multi-Tass Reinforcement Learning via Mixture-of-Expert Decision Transformer
Tasks
Small Subset
Small Subset 2
Small Subset
Tasks
Grouping
Predict Head
Predict Head
Prodict Head
Transformer Bock x
PFN
Expert
Espert 2
Expert
FIN
Expert
Expert
Layer
Multi Hezd Attention
Lager
Layer Norm
Mult-Head Attention
Layer Norm
Embelding
Laver Norm
Embedding
Embedcing
Stage 1.
Stage Additionally introdure Expert module Task Grouping & Expert Training
Stage 3. Additionllyy introduce Router module Router Training
Backbone Training
We train the PromptDT on all tasks as the backbone. Stage 2: Here, we employ task grouping obtain various small task subsets, and introduce expert module into every transformer block of the backbone, with each expert handling specific task subset. In this stage, only the experts optimized. Stage 3: We introduce the router to dynamically assign weights to all experts for training across all tasks. Only the router is optimized in this phase.
Figure Overview of M3DT. Stage 1:
identical to the FFN module experts and employ an MLP as the router. Formally, the forward propagation process of the FFN in M3DT can be simplified as follows:
80
Normalized Score
0.6
Gradient Similarity
Normalized Score
76
0.5
Gradient Similarity
f(x) = x fFEN(1) + fMOE(I)),
(3)
0.4
0.3
[MOE(1) Σ Softmar([e, (x).
(4)
0.2
where fEFN denotes the retained FFN in the backbone, and fo, and fe, represent the i-th expert and the router, respec- tively. By this, the MoE and the backbone work in tandem, enabling the agent to dynamically adjust the utilization of different experts for handling diverse types of tasks.
D.1
PromptDT
PromptDT-MoE
M3DT-R
M3DT-G
Figure 5: The inherent structural benefits of MoE alone are insufficient, while task grouping and three-stage train- ing paradigm can substantially enhance model performance while alleviating gradient conflicts. The number of experts is configured 40.
However, the results Figure suggest that relying solely on the advantage of the MoE is insufficient handle mas- sive MTRL. Upon end-to-end training of the MoB-enhanced PromptDT, gradient conflicts within the expert modules are even more pronounced than those observed in the MLP layers of PromptDT, and the overall performance improve- ments remain marginal. These findings underscore the ne- cessity of refining the training paradigm.
we propose a three-stage training mechanism, with distinct training processes designed for the backhone, experts, and router, each trained separately. The overall training process is illustrated in Figure 4.
Backbone training with minimal gradient conflicts. We first train Prompt-DT on all tasks to capture shared knowl- edge, enabling it to embed information and predict actions across tasks, which serves as the backbone architecture for M3DT. This training process results intense gradient con- flicts, as illustrated in Figure 2. To gain deeper insight into the progression of gradient conflicts during training, we
4.2. Optimization: Three-Stage Training Mechanism
Considering our unique training purpose, namely the ex- plicit task assignment and independent training for each expert, well the inherent challenges of training MoF,
5

---

Mastering Massive Multi-Task Reinforcement Learning via Mixture-of-Expert Decision Transformer
Table I: Comparison of M3DT with baselines of varying sizes across three task scales. M3DT consistently outperforms
other baselines and achieves remarkable task scalability, which is attributed to its parameter scalability, as shown in Figure 6.
Task Scale
10 Tasks
80 Tasks
160 Tasks
Method
Score
Parameters
Score
Parameters
Score
Parameters
MTDT-Small
82.75 = 2.29
1.47M
66.65 = 0.43
1.47M
59.19 + 1.77
1.47M
MTDT-Large
88.92 = 0.98
173.30M
74.38 1 1.33
173.30M
70.65 1 1.67
173.30M
PromptDT-Small
82.41 = 1.47
1.47M
65.11 = 2.65
1.47M
61.46 + 1.78
1.47M
PromptDT-Large
88.17 = 1.41
173.30M
75.70 = 1.62
173.30M
71.65 + 1.13
173.30M
HarmoDT-SSmal
79.75 = 1.32
1.47M
60.71 + 3.91
1.47M
57.27 + 0.84
1.47M
HarmoDT-Large
86.63 = 1.28
173.30M
75.56 = 1.69
173.30M
72.80 + 2.89
173.30M
M3DT-Random (Ours)
89.09 = 1.20
47.87M
78.92 + 1.21
98.37M
76.74 + 0.94
174.12M
M3DT-Gradient (Ours)
89.23 = 1.12
47.87M
80.66 0.97
98.37M
78.21 = 0.47
174.12M
conduct experiments that illustrate how performance and gradient confficts develop across training iterations in Fig- ure (left). Initially, the model's performance improves rapidly with gradient conflicts escalate from low initial value, enabling fast knowledge acquisition. However, as training progresses, gradient conflcct rapidly reach the peak, leading to diminishing performance gain. Based on this, we restrict training of the backbone to the early stage before gradient conflicts reach their peak. This strategy allows the model to efficiently learn shared knowledge, ensuring the shared parameters align with the solution space while pre- venting excessive updates on dominant tasks, thereby mini- mizing negative interference with conflicting tasks. More discussion is detailed in Section 5.2. This approach provides a solid foundation for the subsequent expert training.
and task groups to mitigate the severe performance drop. Figure 5 illustrates that grouping tasks and independently training distinct experts can significantly alleviate gradient conflicts within the expert modules, resulting in notable improvements in the overall performance.
Router training. Finally, we train the router on all tasks, enabling it to dynamically assign weights to different well- trained experts for various tasks based on the hidden states from the backbone, integrating all sub-policies into uni- fied strategy without requiring task IDs. In this stage, we freeze the parameters of both the backbone and all experts, allowing only the router to be optimized. This strategy significantly simplifies the training of the MoE while pre- serving the knowledge already acquired by the model.
Task grouping and experts individually training. We employ task grouping and then train the experts on smaller task subsets, which not only reduces the task load to achieve better performance, but also mitigates the severe gradient conflicts that arise in shared parameters after a certain stage of training. We begin by proposing two task grouping meth- ods: (1) random grouping: based on the satisfactory results on few tasks as illustrated in Figure 2, we randomly divide all tasks into equally sized subsets for naively redueing the task number to lower value; (2) gradient-based grouping: we first calculate the agreement vectors (Hu et al., 2024) (as detailed in Appendix C) for each task in the current backbone as a measure of task similarity, and then apply K-means to group these vectors and the corresponding tasks. After obtaining the task subsets, we introduce a dedicated expert for each subset. The expert module is trained on its specific task subset with the backbone parameters frozen, preserving the shared knowledge while allowing the expert to focus on learning task sask-specific information. By limiting the size of each task subset, we can ensure effective learn- ing on these parameter subsets. Additionally, training each expert independently helps mitigate the issue of imbalanced updates among experts in MoE. Consequently, when dealing with large- tasks, we expand the number of experts
5. Experiments and Analysis
In this section, we conduct extensive experiments answer the following questions: (1) How does M3DT compare to other baselines in the massive multi-task regime? (2) Does M3DT exhibit task scalability and parameter scalability? (3) What makes M3DT effective? The detailed implementa- tions', including the baselines and environments, are illus- trated in Appendix A.
5.1. M3DT helps Task and Parameter Scalability
In this study, we benchmark M3DT and its variants against baselines on different task scales. We compare the perfor- mance of the baselines at their default size and a expanded size. The variants of M3DT considered in this evaluation include M3nT-Random, which employ random grouping to obtain task subsets and train experts; M3DT-Gradient, which utilizes gradient information for grouping.
As shown in Table 1, M3DT-Random surpasses all other methods at all task scales, achieving a 0.1%, 4.3% and 5.4%
OOur code is available at: hattps://githss KongYilun
6

---

Mastering Massive Multi-Tass Reinforcement Learning via Mixture-of-Expert Decision Transformer
Parometer scalability on 80 tasks
Parametar scalability on 160 tasks
Table 2: Ablation study on different training process of M3DT, which illustrate the effectiveness of our dedicated three-stage training mechanism.
Normaltzed Score
Method
Normalized Score
Promeiii 10 tidee
grudion
M3DT-R
76.67 = 0.29
1000 runduv
M3DT-G
77.89 = 0.47
on
1751 0.1)
66.57M (44)
533.9M 6C)
474m5 (8)
72.1234 (1.) (20)
M3DT w/o 3-stage training
71.90 = 0.70
Model Size (Export Num)
Model Size (Expert Num)
M3DT w/o grouping
67.34 = 0.56
Figure 6: By increasing the number of experts, M3DT ef-
M3DT-R w/o expert freezing
71.89 = 0.63
fectively unlocks the parameter scalability, further helping to tackle task scalability.
M3DT-G w/o expert freezing
71.88 + 0.62
improvement in 10, 80 and 160 tasks, respectively, com- pared to the best baseline. The advantages of our method become more pronounced as the number of tasks increases. By employing random grouping, M3DT-RRando already effectively competes with the current state-of-the-ar tech- niques, highlighting the significant effectiveness of introduc- ing parameter subsets and reducing corresponding task load. Furthermore, M3DT-Gradient enhances the performance by identifying better task groups to mitigate learning complex- ity, resulting in substantial gains of 6.6% and 7.5% in 80 tasks and 160 tasks, respectively. M3DT also effectively alleviates the performance degradation with increasing task numbers. Compared to the severe performance drop of around 20% observed in other baselines when scaling tasks numbers from 10 to 160, our approach only experienced 12.3% decline. And M3DT achieves higher score on 160 tasks than other baselines do on 80 tasks, demonstrating remarkable scalability across task numbers.
grouping and three-stage training mechanism, or the inher- ent advantages of the MoE architecture itself, we conduct a detailed ablation study, as illustrated in Table 2. Despite utilizing the MoE structure, the absence of our meticulously designed training strategy leads a substantial performance degradation. Training PromptDT with MoE end-to-end from scratch (i.e. M3DT w/o 3-stage training) only yields results similar to those of a standalone PromptDT with the same parameter scale, while freezing the trained backbone and jointly training all experts and the router on all tasks (i.e. M3DT w/o grouping) results in a worse performance, with score of 67.34. The comparison of gradient conflicts is shown in Figure 5. In addition, after training the experts in groups, simultaneously fine-tuning them when training the router (i.e. both M3DT-R/G w/o expert freezing) also leads to suboptimal results. This further underscores the validity of our entire framework.
The scalability of M3DT in handling massive tasks is at- tributed to its parameter scalability, which can meanwhile reduce the task load on each parameter subset by increasing the number of experts, thereby enhancing overall perfor- mance. We conduct experiments on the number of experts across 80 tasks and 160 tasks, as shown in Figure 6. For 10 tasks, 8 experts already yield excellent results, achieving scores of 89.09 and 89.23 with 47.87M parameters. In 80 and 160 tasks, expanding the parameter size within a certain range through increasing the number of experts can signifi- cantly improve model performance, achieving a 11.2% and 11.7% improvement, separately. Since both M3DT-Random and M3DT-Gras Gradient show diminishing performance gains after reaching 40 experts on 160 tasks, we use 40 experts in subsequent experiments unless stated otherwise. Addition- ally, we analyze the reasons why increasing the number of experts cannot continually improve performance and even- tually reaches a performance ceiling in later experiments.
Training curves of the PromptDr
Ablation on Btackbons trkining duration
Grodiest Conticts
Nantalied Store
2
3 2
22
Pelfo
02
0.
Reraticn Sacibbon Trunap vallaan))
Janaatone BlacMboon Traiding On millooo
Figure 7: (left) Training curves of the PromptDT, where the early stage exhibits mild gradient confficts with swift per- formance gain; (right) The overall performance of M3DT varies with the number of training steps applied to the back- bone, with optimal performance occurring just as gradient conflicts reach the peak.
Is early stopping of backbone training truly effective? As illustrated in Figure 7 (left), the gradient conflicts in the PromptDT backbone progressively escalate with train- ing duration, eventually reaching peak, after which the performance improvements become notably sluggish. To assess the efficacy of early stopping, which is employed to mitigate these severe gradient conflicts at the expense
5.2. Further Analysis
Is grouped training important? To investigate whether the success of M3DT is primarily due to the explicit task
7

---

Mastering Massive Multi-Tass Reinforcement Learning via Mixture-of-Expere Decision Transformer
85.0
Alrrage Tingla Partoemace waineay 2 parmetaes
82.5
76
Nomaltzd Score
74
Normalzzed Soore
80.0
77.5
70
75.0
72.5
M
Overall Performance Oracle Expert Selection
64
MBDTC MSDD
70.0
M
16
26
32
40
&.mal
Export Num
67.5
MODTG G Small
Figure 9: Top-K routing selects the router's 's top-k outputs, applies softmax to obtain probabilities, and computes the weighted sum of the corresponding expert outputs. M3DT with Top-4 router fails to scale with the number of experts.
25
50
75
100
125
150
175
200
Model Size (M)
Figure 8: Performance curves of M3DT with increasing model size (i.e., the number of experts) across different base sizes, where each point denotes the addition of 8 experts.
the dashed and solid lines reflects the performance loss at- tributed to the router. When the number of experts is small, the router can easily allocate weights across the few ex- perts, resulting in minimal performance loss or even better results. As the number of experts increases, the difficulty of assigning weights to the experts grows, resulting in a pro- gressively larger performance gap, which peaks when the expert performance continues to improve while the overall performance plateaus or even drops. Thus, it is reasonable and tolerable that the framework's 's overall performance may be lower than the ideal case of perfect expert switch, as increasing the number of experts inherently raises the diffi- culty of optimal weight allocation. (3) Backbone + Expert: Expert performance also tends to plateau when the number of experts becomes sufficiently large. This is primarily due to the shared knowledge learned by the backbone across all tasks is limited, which restricts further performance im- provements, regardless of the expert's 's capabilities. Addi- tionally, when the number of experts is large enough and each expert already faces a sufficiently small task subset, further increasing the number of experts yields diminishing returns in reducing the task load.
of some performance gains, we conduct a comparison of M3DT using PromptDT trained for varying steps as back- bone. The results are depicted in Figure 7 (right). M3DT consistently achieves strong performance across various training durations, with early stopping at 400k steps, just as gradient conflicts peak, yielding the most optimal results. This demonstrates the robustness of M3DT in the initial training phase. Training for only 200k steps prevents the backbone from fully learning the shared knowledge, causing the parameters to deviate from the optimal solution space, resulting in suboptimal performance. On the other hand, continuing training beyond 400k steps, when gradient con- flicts are extremely severe, also leads to a decline in M3DT's performance. This is likely due to the parameters become overfitted to tasks whose gradients dominate, causing them to deviate further from the solution space of other conflicting tasks. Introducing expert modules at this time for grouped task learning also falls victim this issue, as they unable to learn the specific knowledge on these conflicting tasks, thus failing to alleviate the performance drop.
In this study, we in- troduce: Oracle Expert Selection: manually selecting each test task's corresponding expert based on Task ID, individu- ally evaluating their performance on each task, calculate the averaged results as the upper-bound performance of M3DT, as task IDs are unavailable in overall method evaluation and real-world applications: Small: a scaled- -own version of M3DT, where the width of backbone, experts, and router are all proportionally reduces. As shown in Figure 8, the main factors influencing the parameter scalability M3DT are threefold: (1) Model width: With the number of experts increasing, M3DT-Small achieves significant performance improvements with modest growth in parameters. How- ever, due to the constrains of the small model width on the capacity of each module, its scalability to the number of experts is poor, reaching the performance ceiling at only 24 experts. As a result, its parameter scalability is inherently constrained. (2) Router: The performance gap between
What affects parameter expansion?
Router design. We compare the Top-K routing (Shazeer et al., 2017) strategy to investigate whether reducing the number of activated experts can enhance the scalability of expert number. In this experiment, we employed Top-4 routing, and the results are presented in Figure 9. Although our proposed three-stage training mechanism significantly simplifies the training of MoE, Top-4 router fails to scale the number of experts, and performance deteriorates as expert number increases. This is attributed to improper routing load balancing, where certain routes are excessively optimized, combined with the instability in router training induced by sparsity. This finding is consistent with prior work demon- strating the difficulty in scaling up deep RL networks with Top-K router (Obando-Cero) et al., 2024). While it is pos- sible some losses (Riquelme et al., 2021; Mustafa et al., 2022) may result in better Top-K performance, this find- ing suggests that M3DT benefits from having a weighted combination of all experts from all task subsets.
8

---

Mastering Massive Multi-Task Reinforcement Learning via Mixture-of-Expert Decision Transformer
6. Related Work
8. Conclusion
Multi-Task RL Multi-task RL aims to learn a shared policy for a diverse set of tasks. One of the most straightforward approaches to MTRL is formulate the multi-task model as task-conditional sequence modeling (Xu et al., 2022; Reed et al., 2022; Lee et al., 2022; He et al., 2023; Hu et al., 2024). Some methods also focus on handling gradient conflicts among different tasks (Yu et al., 2020a; Chen et al., 2020; Liu et al., 2021). On the other hand, some methods employed a dedicated shared structure to leverage the shared knowledge (D' Eramo et al., 2024; Yang al., 2020; Sun et al., 2022). While most of these methods address tens or even less tasks, focus scaling to significantly larger number of tasks, achieving task scalability.
In this study, we first delve into the challenges of sequence modeling in MTRL, analyze the performance degradation and gradient conflicts with increasing task numbers, and identify the limited effectiveness of naively scaling parame- ter in MTRL. Based on these insights, we introduce M3DT, novel approach designed unlock the parameter scalabil- ity for handling massive tasks. By employing task grouping and MoE, M3DT significantly reduces the task load as- signed to each parameter subset while enhancing the overall performance through the expansion of experts. We pro- pose a three-stage training mechanism that allows explicit task assignment to each expert, enabling sequential training of different modules without interference. Our empirical evaluations across diverse task scales underscore M3DT's superior performance compared existing baselines, estab- lishing its state-of-ther -frth-rr- effectiveness in MTRL scenarios.
ofxtture-ofExpertt MoEs have recently helped scaling language models up to trillions of parameters thanks to their modular nature (Lepikhin et al., 2020; Fedus et al., 2022). MoBs also help performance in multi-task settings (Fan et al., 2022; Ye & Xu, 2023; Dou al., 2024). There have been few works exploring MoEs in RL for single (Akrour et al., 2021; Obando-Ceron et al., 2024) and multi-task learning (Hendawy et al., 2023; Huang et al., 2024). While M3DT explicitly assign tasks to specific experts and train each expert independently to handle multi-task problems.
Impact Statement
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
A Acknowledgement
7. Limitation
This work supported by STI 2030-Major Projects (No. 20212D0201405) Shenzhen Basic Research Project (Natu- ral Science Foundation) Basic Research Key Project (NO, (CY1202412021244300AA National Research Foundation, Singapore, under its NRF Professorship Award No. NRF. P2024-001.
We primarily focus on the challenges in MTRL, analyze the scalability of task numbers and parameter size, and propose a paradigm that increases model parameters to achieve task scalability. However:
(1) we did not focus on fine -naiiin network architecture design, such as the expert and router modules. Fur- ther optimization of these components could potentially improve performance while reducing model size.
References
Akrour, R., Tateo, D., and Peters, J. Continuous action reinforcement learning from a mixture of interpretable experts. IEEE Transactions on Pattern Analysis and Ma- chine Intelligence, 44 44(10):6795-6806, 2021.
(2) While we leave held-out task generalization and con- tinual learning unexplored, our method's 's modular pa- rameterization and group-wise learning naturally sup- ports solutions like held-out task adaptation via expert learning, dynamic router-based skill composition for held-out tasks, and forgetting mitigation mechanisms for the router for past tasks. Since generalization and continual learning are highly critical in RI., resolving them presents promising future works.
Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al. Language models are few-shot learners. Advances in neural information processing systems, 33: 1877-1901, 2020.
(3) Scaling experts with task count raises inference costs. While we prioritize algorithm performance above all else, we experimented with activating Top-k experts with fixed k to control inference overhead, but it de- grades the performance. A more tailored Top-k gating mechanism, better aligned with our three-stage training paradigm, could enhance the practical applicability of our method.
Chen, L., Lu, K., Rajeswaran, A., Lee, K., Grover, A., Laskin, M., Abbeel, P., Srinivas, A., and Mordatch, I. De- cision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems, 34:15084-15099 2021.
Chen, Z., Ngiam, J., Huang, Y., Luong, T., Kretzschmar, H., Chai, Y., and Anguelov, D. Just pick sign: Opti-

---

Mastering Massive Multi-Task Reinforcement Learning via Mixture-of-Expert Decision Transformer
mizing deep multitask models with gradient sign dropout. Advances in Neural Information Processing Systems, 33: 2039-2050, 2020.
Huang, S., Zhang, Z., Liang, T., Xu, Y., Kou, Z., Lu, C., Xu, G., Xue, Z., and Xu, H. Mentor: Mixture-of-expert network with task-oriented perturbation for visual rein- forcement learning. arXiv preprint arXiv:2410.14977 2024.
D' D'Eramo, C., Tateo, D., Bonarini, A., Restelli, M., and Peters, J. Sharing knowledge in multi task deep rein- forcement learning. arXiv preprint N1i:240010956// 2024.
Janner, M., Li. Q., and Levine. S. Offline reinforcement learning as one big sequence modeling problem. Ad- vances in neural information processing systems, 34: 1273-1286, 2021.
Dou, S., Zhou, E., Liu, Y., Gao, S., Shen, W., Xiong, L., Zhou, Y.. Wang, X., Xi, Z., Fan, X., et al. Loramoe: Alleviating world knowledge forgetting in large language models via moe-style plugin. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume Long Papers), pp. 1932-1945, 2024.
Kaplan, J., McCandlish, S., Henighan, T., Brown, T. B., Chess, B., Child, R., Gray, S., Radford, A., Wu, J., and Amodei, D. Scaling laws for neural language models. arXiv preprint SxXi::0001008366 2020.
Kong, Y., Mao, H., Qi, Z., Zhang, B., Ruan, J., Shen, L., Chang, Y., Wang, X., Zhao, R., and Tao, D. QPO: Query- dependent prompt optimization via multi-loop offline re- inforcement learning. Transactions on Machine Learn- ing Research, 2025. ISSN 2835-8856. URL https: / /poenrevieww net forum?iruiiooe corrmmPi.--bgggoTTTT.T... M.MMM..TT..............
Fan, Z., Sarkar, R., Jiang, Z., Chen, T., Zou, K., Cheng, Y., Hao, C., Wang, Z., et al. vtvit: Mixture-of-expert vision transformer for efficient multi-task learning with model- accelerator co-design. Advances Neural Information Processing Systems, 35:28441 -2841,, 2022.
Lee, K.-H., Nachum, O., Yang, M. S., Lee, L., Free- man, D., Guadarrama, S., Fischer, I., Xu, W., Jang, E., Michalewski, H., et al. Multi-game decision transformers. Advances in Neural Information Processing Systems, 35: 27921-27936 2022.
Fedus, W., Zoph, B., and Shazeer, N. Switch transformers: Scaling trillion parameter models with simple and ef- ficient sparsity. Journal of Machine Learning Research, 23( 23(120):1-39) 2022.
Finn, C., Abbeel, P., and Levine, S. Model-agnostic meta- learning for fast adaptation of deep networks. In Interna- tional conference on machine learning, pp. 1126-1135. PMLR, 2017.
Lepikhin, D., Lee, H., Xu, Y., Chen, D., Firat, O., Huang, Y., Krikun, M., Shazeer, N., and Chen, Z. Gshard: Scaling giant models with conditional computation and automatic sharding. arXiv preprint arXiv:2006. 16668, 2020.
Haarnoja, T., Zhou, A., Abbeel, P., and Levine, S. Soft actor-critic: Off-policy maximum entropy deep reinforce- ment learning with stochastic actor. In International conference on machine learning, pp. 1861-1870. PMLR, 2018.
Levine, S., Kumar, A., Tucker, G., and Fu, J. Offline rein- forcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint arXiv:2005.00664 2020.
Liu, B., Liu, X., Jin, X., Stone, P., and Liu, Q. Conflict- averse gradient descent for multi-task learning. Advances in Neural Information Processing Systems, 34:18878 18890, 2021.
Hansen, N., Su, H., and Wang, X. Td-mpc2: Scalable, ro- bust world models for continuous control. arXiv preprint arXiv:2310. T.t::230..18222 16828, 2023.
Liu, P., Yuan, W., Fu, J., Jiang, Z., Hayashi, H., and Neubig, G. Pre-train, prompt, and predict: A systematic survey of prompting methods in natural language processing. ACM Computing Surveys, 55(9):1-35, 2023.
He, H., Bai, C., Xu, K., Yang. Z., Zhang, W., Wang, D., Zhao, B., and Li, X. Diffusion model is an effective planner and data synthesizer for multi-task reinforcement learning. Advances in neural information processing systems, 36:64896-6499 64917, 2023.
Mustafa, B., Riquelme, C., Puigcerver, J., Jenatton, R., and Houlsby, N. Multimodal contrastive learning with limoe: the language-image mixture of experts. Advances in Neural Information Processing Systems, 35:9564-9574 2022.
Hendawy, A., Peters, J., and D' E'Eammo, C. Multi-task rein- forcement learning with mixture of orthogonal experts. arXiv preprint arXiv:2311.11388 2023.
Hu, S., Fan, Z. Shen, L., Zhang, Y., Wang, Y., and Tao, D. Harmodt: Harmony multi-task decision trans- former for offline reinforcement learning. arXiv preprint arXiv:2405. 18080, 2024.
Obando-Ceron, J., Sokar, G., Willi, T., Lyle, C., Farebrother, J., Foerster, J., Dziugaite, G. K., Precup, D., and Castro, P. S. Mixtures of experts unlock parameter scaling for deep rl. arXiv preprint arXiv:2402 200609940,.08609 2024.
10

---

Mastering Massive Multi-Task Reinforcement Learning via Mixture-of-Expert Decision Transformer
Reed, S., Zolna, K., Parisotto, E., Colmenarejo, S. G., Novikov, A., Barth-Maron, G., Gimenez, M., Sulsky, Y., Kay, J., Springenberg, J. T., et al. A generalist agent. arXiv preprint arXiv:2205.061175. 2022.
Advances in Neural Information Processing Systems, 33: 2020a.
Yu, T., Quillen, D.. He, Z.. Julian, R., Hausman, K., Finn, C., and Levine, Meta- deta-world: benchmark and evalua- tion for multi-task and meta reinforcement learning. In Conference robot learning, pp. 1094-1100, PMLR, 2020b.
Riquelme, C., Puigcerver, J., Mustafa, B., Neumann, M., Jenatton, R., Susano Pinto, A., Keysers, D., and Houlsby, N. Scaling vision with sparse mixture of experts. Ad- vances in Neural Information Processing Systems, 34: 8583-8595, 2021.
Rothfuss, J., Lee, D., Clavera, I,, Asfour, T., and Abbeel, P. Promp: Proximal meta-policy search. arXiv preprint arXiv:1810.06774. 2018.
Shazeer, N., Mirhoseini, A., Maziarz, K., Davis, A., Le, Q. Hinton, G., and Dean, Outrageously large neural networks: The sparsely-gated mixture-of-expprts layer. arXiv preprint arXiv: 06538, 2017.
Sun, L., Zhang, H., Xu, W. and Tomizuka, M. Paco: Parameter-composttonna multi-task reinforcement learn- ing. Advances in Neural Information Processing Systems, 35:21495-21500 2022.
Sutton, R. S. and Barto, A. G. Reinforcement learning: An introduction. MIT press, 2018.
Tassa, Y., Doron, Y., Muldal, A., Erez, T., Li, Y., Casas, D. d. L., Budden, D., Abdolmaleki, A., Merel, J., Lefrancq, A. al. Deepmind control suite. arXiv preprint arXiv:1801.00660 2018.
Todorov, E., Erez, T., and Tassa, Y. Mujoco: A physics engine for model-based based control. 2012 IEEE/RSJ Inter- national Conference on Intelligent Robots and Systems, pp. 5-5550225050 2012. doi: 10.1 200//00..2002..6550...
Vaswani, A. Attention is all you need. Advances in Neural Information Processing Systems, 2017.
Xu, M., Shen, Y., Zhang, S., Lu, Y., Zhao, D., Tenenbaum, J., and Gan, C. Prompting decision transformer for few- shot policy generalization. In international conference on machine learning, pp. 24631-24644 PMLR, 2022.
Yang. R., Xu, H., Wu, Y., and Wang, X. Multi-task rein- forcement learning with soft modularization. Advances in Neurul Information Processing Systems, 33:4767-47777, 2020.
Ye, and Xu, D. Taskexpert: Dynamically assembling multi-task representations with memorial mixture-of- experts. In Proceedings of the IEEE/CVF International Conference Computer Vision, pp. 21828-21837, 2023.
Yu, T., Kumar, S., Gupta, A., Levine, S., Hausman, K., and Finn, C. Gradient surgery for multi-tas learning.

---

Mastering Massive Multi-Task Reinforcement Learning via Mixture-of-Expert Decision Transformer
A. Experimental Details
A.I. Detailed Environments
We consider a total of 160 continuous control tasks from 3 task domains: Meta-World (Yu et al., 2020b), DMControl (Tassa et al., 2018), Mujoco Locomotion (Todorov al., 2012). This section provides an exhaustive introduction of the tasks considered, including their observation and action dimensions, and the calculation of normalized score. Our goal is not propose a new benchmark with 160 tasks, but rather to use sufficiently large number of tasks explore the impact of task quantity.
1. META-WORLD
The Meta-World benchmark encompasses a diverse array of 50 distinct manipulation tasks, unified by shared dynamics. These tasks involve a Sawyer robot engaging with a variety of objects, each distinguished by unique shapes, joints, and connective properties. The complexity of this benchmark lies in the heterogeneity of the state spaces and reward functions across tasks, as the robot is required to manipulate different objects towards varying objectives. The robot operates with a 4- 4-dimensional fine- fine-grained action imput at each timestep, which controls the 3D positional movements of its end effector and modulates the gripper's 's openness. The state space is unified into 39 dimension. In its original configuration, the NorldWorld environment is set with fixed goals, format that somewhat limits the scope and realism of robotic learning applications. To address this and align with recent advancements in the field, as noted in works by (Yang et al., 2020; Sun et al., 2022: He et al., 2023; Hu et al., 2024), we have modified all tasks to incorporate a random-goal setting. For the offline dataset, we follow the works (He et al., 2023; Hu et al., 2024) and utilize their dataset with the near-optimal trajectories, which consists of the experience from random to expert (convergence) in RepaRy (Haarnoja al., 2018). The primary metric for evaluating performance in this benchmark is the average success rate across all tasks, providing comprehensive measure of the robotic system's is adaptability and proficiency in varied task environments. We directly use the success rate each task as its normalized score.
A.1.2. DMCONTROL
The tasks in DMControl involve significantly more diverse embodiments, state spaces, action spaces, and reward functions, which greatly increases their complexity. We consider a total of 30 continuous control tasks in the DMControl domain, including 19 original DMControl tasks and 11 new (custom) tasks created specifically for M3DT benchmarking, following the work (Hansen et al., 2023). We directly use the dataset collected by (Hansen et al., 2023), and we only use the first 2,000 trajectories for each task to ensure consistency in dataset size with other tasks. We list all used DMControl tasks in Table 3. For evaluation, we linearly scale the original reward range of [0,1000] to [0,100], using it as our normalized score.
A.1.3. MUJOCO LOCOMOTION
In this paper, we also employ a diverse array of meta-RL control tasks to construct a dataset with a sufficient number of tasks for exploring the challenges of MTRL when confronted with a large number of tasks. We directly utilize the datasets proposed by Xu et al. (2022). The tasks are detailed as follows:
Cheetah-vel (Finn et al., 2017): It defines 40 unique tasks, each associated with specific goal velocity, uniformly distributed between 0 and 3 m/s. The agent's performance is assessed based the 12 error relative to the target velocity, with a penalty for deviations. These 40 tasks share a unified state space of 20 and an action space of 6. Based on the reward ranges of these environments, we linearly map the return values within the interval [-100, -30] to the normalized range of [0, 100] as our normalized scores, while returns outside this range are directly capped at 0 or 100, respectively.
Ant-dir (Rothfuss et al., 2018): We also use 40 tasks in this domain, each with goal direction uniformly sampled in two-dimensional plane. The agent, an 8-jointed ant, is incentivized to attain high velocity in the designated direction. The state space for these tasks has dimensionality of 27, and the action space consists of 8 dimensions. We linearly map the return values within the interval [0, 500] to the normalized range of [0, 100], while returns outside this range are directly capped at 0 or 100, respectively, to calculate our normalized scores.
By using normalized scores, we can align tasks with initially inconsistent evaluation metrics, enabling us to assess the model's ' ability to simultancously tackle multiple tasks. For the calculation of average normalized score, we assign equal weight to each of the 160 tasks and compute the average of normalized scores.
12

---

Mastering Massive Multi-Task Reinforcement Learning via Mixture-of-Expert Decision Transformer
Table 3: DMControl tasks used in this paper.
Task
Observation dim
Action dim
New?
Acrobot Swingup
6
N
Cartpole Balance
5
1
N
Cartpole Balance Sparse
5
1
N
Cartpole Swingup
5
1
N
Cartpole Swingup Sparse
5
1
N
Cheetah Jump
17
6
Cheetah Run
17
6
Cheetah Run Back
17
6
Y
Cheetah Run Backwards
17
6
Y
Cheetah Run Front
17
6
Y
Cup Catch
8
2
N
Cup Spin
8
2
Y
Finger Spin
9
2
N
N
Finger Turn Easy
12
2
N
Finger Turn Hard
12
2
Fish Swim
24
5
N
Hopper Hop
15
4
N
Hopper Hop Backwards
15
4
Y
Hopper Stand
15
4
N
Pendulum Spin
3
1
Y
Pendulum Swingup
3
1
N
Reacher Easy
6
2
N
Reacher Hard
6
2
N
Reacher Three Easy
8
3
Y
Reacher Three Hard
8
3
Y
Walker Run
24
6
N
Y
Walker Run Backwards
24
6
Walker Stand
24
6
N
Walker Walk
24
6
Walker Walk Backwards
24
6 Y
To address the inconsistency in state and action spaces across tasks, we zero-pad all states and actions their largest respective dimensions (i.e. and respectively). and mask out invalid action dimensions in predictions made by the policy during both training and inference.
A.2. Baselines
We compare our proposed M3DT with the following DT-based baselines.
We extend the DT architecture (Chen al., 2021) learn from multitask data. Specifically, MTDT concatenates embedding and state as the input tokens, where is the encoding of task ID. In evaluation, the reward-to-go and task ID are fed into the Transformer to provide task-specific information. Leveraging the scalability of the Transformer architecture, we compare the performance this method at both its default size (1.47M) and expanded size (173.30M).
MTDT:
PromptDT (Xu et al., 2022): PromptDT built on DT aims to learn from multi-task data and generalize the policy to unseen tasks. It leverages short task trajectories as prompts to guide the model in identifying the current task. PromptDT generates actions based on the trajectory prompts and reward-to-go. compare the performance of this method both its default size (1.47M) and expanded size (173.30M).
HarmoDT (Hu et al., 2024): HarmoDT built on PromptDT aims to lean a specific mask for each task, effectively shielding the model's 's parameters that conflict most with the task, thereby mitigating the severe gradient conflicts in
13

---

Mastering Massive Multi-Task Reinforcement Learning via Mixture-of-Expert Decision Transformer
MTRL. During evaluation, HarmoDT also generates actions based the trajectory prompts and reward-to-go. We compare the performance of this method at both its default size (1.47M) and expanded size (173.30M).
A.3. Task Selection
To mitigate the impact of task selection and combination on performance, use different task combinations for each seed in scenarios involving 10, 20, 40, 80, and 120 tasks. To fairly analyze the effects solely caused by task quantity, maintain consistent average task difficulty across various task scales and seeds. Specifically, we train 160 separate PromptDT models for each of the 160 tasks and compute the score for each task, which serves as the task difficulty for that task. When selecting tasks for different task scales, we ensure that the average task difficulty of the chosen task set aligns with that of the complete set of 160 tasks. We use the 1.47M PromptDT model to compute the task difficulty and derive the task sets for all task scales. Subsequently, all experiments are conducted using the same task set to eliminate potential bias.
In summary, achieve methodological rigor in task selection, follow three key criteria:
Performance Balance: Based on single-task normalized score, the selected task subsets maintain average scores comparable to the full set of 160 tasks, to eliminate the influence of task difficulty;
Domain Ratio Preservation: We make the subset composition different task scales approximately maintains the in order to mitigate domain-specific biases;
original domain ratio
Subset Diversity: We employ distinct task subsets for each run with the same task scale to maximize coverage of the full task pool, eliminating spurious correlations between specific tasks.
A.4. Gradient Simalarity
We compute the gradients for each task and calculate the mean gradient across all tasks. The gradient similarity is defined as the average cosine similarity between this mean gradient and tasks' gradients. During training, we record this metric every le4 steps to obtain similarity curve. After smoothing to the stabilized phase of training, we use the plateau value as our final gradient similarity. Gradient conffict calculated (1 gradient similarity).
A.5. Implementation of Model Expansion
In this experiment, we do not invest significant effort into exploring the optimal structure for each baseline model at different parameter scales. Instead, adopt their default number of layers and attention heads, only increasing the width of the models, i.e the dimension of their hidden states. Although previous research (Kaplan al., 2020) have shown that the model structure has only a marginal impact on performance when the parameter size remains constant, we conduct experiments with PromptDT to verify whether our expanding approach is reasonable in the context of MTRL. The results are shown in Figure 10, where dmodet denotes the width of the model, Wayer and Whead denote the number of layers and attention heads, respectively. In MTRL, the model performance is also weakly depends on the model architecture. We summarize the specific structures of each scaled model Table
Table 4: DMControl tasks used this paper.
Parameters
Layers
Attention Heads
Model Width
1.47M
128
5.29M
6
8
256
20.02M
6
8
512
77.78M
8
1024
173.30M
1536
A.6. Experimental Setups, Hyper-parameters and Resources
In this section, we introduce the implementation for M3DT. We employ the PromptDT with 5.29M parameters as the backbone of M3DT, with the structure outlined in Table 4. The structures our introduced experts are identical to that of the FFN in the backbone. We employ 5-MLP as our router. The padded input and output dimensions are 39 and
14

---

Mastering Massive Multi-Tass Task Reinforcement Learning via Mixture-of-Expert Decision Transformer
80
Normalized Score
75
173M
77M
10¹
10
10a
10a
Aspect Ratio maye
Attention Head Dim dmode//Nhead
Figure 10: Performance depends very mildly on model shape when the total number of parameters is held fixed.
respectively, as illustrated A.1. .. All experiment in this paper are run with 3 seeds. The specific model parameters and hypee-paramettrs utilized our training process are outlined in Table 5. We use NVIDIA GeForce RTX 4090 to train and evaluate each model except HargeDDT-Lare, while it is trained and evaluated on NVIDIA A100 40G due to its substantial resource requirements.
Table Hyper-parameters of M3DT in our experiments.
Parameter
Value
Number of layers
Number of attention heads
Hidden dimension
256
Number of experts
[8,16,24.32,40048]
Nonlinearity function
ReLU
Batch size
16
Prompt length K
20
Dropout
0.1
Learning rate
1.0c-4
Optimizer
Adam
Total rounds
le6
-Backbone training rounds
4e5
-Expert training rounds
2c5
-Router training rounds
4e5
B. More Experiments and Further Analysis
B.1. Domain-wise Analysis of the Scalability of Task Numbers and Model Size
We separately report the scores in respective domains to eliminate the potential conclusion bias caused by varying difficulty levels across domains. The model is still trained on mixed tasks from all four domains (MW + DMC + Ant + Cheetah). The Total number on the y-axis denotes the overall trained tasks from the four domains, while the Domain number in parentheses indicates the number tasks belonging to specific domains the training set, and the score represents the average performance across these tasks. Both task scalability and parameter scalability exhibit the same trends as demonstrated in the original paper.
For the parameter scalability, the performances on MW+DMC exhibit the same performance patterns as observed in Fig. 3. For Ant+Cheetah, PromptDT achieves near-optimal performance at original size when trained on limited tasks, while also shows the same pattern of rapidly reaching performance saturation with increased parameters under large-scale tasks. These collectively confirm that including Ant and Cheetah as 80 tasks does not introduce bias to our parameter scalability conclusion.
15

---

Mastering Massive Multi-Task Reinforcement Learning via Mixture-of-Expert Decision Transformer
Normalized Score on MW+DMC
Normalized Score on Ant+ Cheetah
Total (Domain) Task Numbers
160/80)-
35.56
46.95
53.63
53.51
53.17
Total (Domain) Task Numbers
160(80)
87.37
87.43
90.06
90.11
120/60)
53.41
58.57
120(60)
86.52
89.11
90.07
70
92
42.09
59.71
59.40
88.15
89.69
90.66
600
53.01
40(20)
89
20/10)
70.60
75.35
20(10)
91.64
92.52
70.81
77.76
82.69
83.61
82.84
-40
10(5)
94.07
94.41
93.38
93.13
93.61
1.47M
5.29M
20.02M
77.78M
173.30M
1.47M
5.29M
20.02M
77.78M
173.30M
Model Parameters
Model Parameters
Figure 11: The task scalability and parameter scalability exhibit the same trends demonstrated in Section 3.
From the perspective of task scalability, with the criteria of Domain Ratio Consistency ensuring proportional task sampling across different scales, both MW+DMC and Ant+Cheetah scores exhibit identical scaling patterns: sharp drop from 10 to 80 tasks and plateau from 80 to 160 tasks, which further confirms our conclusion.
B.2. Who blocks DT from scaling in MTRI.
To further systematically investigate the limiting factors in DT model scaling, we conduct module-wise dimensional expansion to analyze each component's ' impact on MTRL performance. The results are illustrated in Figure 12.
Through layer-wise analysis gradient similarity across Transformer blocks, we identify significantly stronger gradient conflicts MLP compared to Attention, as shown in Figure 12a. And Figure 12b and Figure 12c indicate that separately expanding the dimensions of either MI.P or attention layers demonstrates measurable efficacy in mitigating gradient conflicts within their respective modules, albeit to varying degrees. Additionally, seen in Figure 12d, independently enhancing the dimensionality of either the MLP attention layers increase the overall gradient similarity of the model.
Nonetheless, the results in Figure 1 12d suggest that despite separate parameter expansion, the model's 's overall gradient conflict still converges to a fairly high value. In terms of model performance, while an initial expansion of the MLP yields some improvement, continued enlargement results in performance degradation. On the other hand, expanding the Attention module directly causes performance drop, suggesting that the default-sized Attention already possesses sufficient representation modeling capability. Given that the MLP a key component for enhancing performance and mitigating gradient conflicts, but merely expanding its parameters yields limited benefits, what approach should we take?
B.3. Domain-wise Main Results
We also separately report the normalized scores of our method and baselines in respective domains in Table and Table 7, M3DT shows more superior performance on complex tasks (MW+DMC).
B.4. Training Time Overhead and Computational Cost
Our experiments are conducted on RTX 4090 GPU. the first stage, we train PromptDT-5M for 4e5 steps, which takes approximately 5.2 hours. the second stage, since each expert trains independently on its dedicated task group, all experts train parallel. We record the training time of the overall second stage training single expert for approximately 1.8 hours. In the third stage, we train the router for 4e5 steps, which takes 17.2 hours. The primary reason for the long training time is our code implementation' 's lack of parallel computation for experts. Code optimization may reduce this training time. Thus, the total training time of our method amounts to approximately 24,2 hours. In comparison, training both PromptDT-1731 173M and MTDT-I 173M requires similar durations of around 21.3 hours, whereas training HarmoDT-173M demands 95.6 hours.
16

---

Mastering Massive Multi-Tas Task Reinforcement Learning via Mixture-of-Expert Decision Transformer
Attention
Attention-256, MLP-256
MLP
Attention-256, MLP-2048
MLP MLP Layer Gradient Similarity
Attention-256, MLP-4096
Attention-2585 MLP-8192
Gradient Similanty
0.12
0.10
Layer Index (a)
Layer Index (b)
72
Attention-25e MLP-256
0.22
MLP Attention
Attention Layer Gradient Simllarity
Attention- MLP-256
Overall Model Gradient Similarity
Attention-4096. MLP-256
0.21
8192, MLP-256
Normalized zedalizzed Score
0.20
66
0.19
0.18
64
0.12
0.17
0.10
0.16
2
256
2048
4096
8192
Layer Index
Dimensian (d)
(c)
Figure 12: (a) MLP layers exhibit more severe gradient conflicts compared to Attention layers within each Transformer layer. (b) Expanding the dimension MLP can mitigate gradient conflicts in the MLP layers. (c) Expanding the dimension of Attention also helps alleviate gradient conflicts in the Attention layers. (d) Despite dimensional expansion across different modules, the model's overall gradient conflicts converge to a severe level, with limited practical performance gains.
B.5. Expert design.
Our proposed MoE architecture replaces the FFN with MoE in each transformer block. This is based what is common practice when adding MoEs to transformer architectures, but is by means the only way to utilize MoEs. Here we investigate a variant: Big: Each expert is a full transformer architecture, where the embedding layer and prediction layer are shared in the backhone. However, B3DDT-Gradient-Big only results normalized score of 76.53, compared to M3DT-Random scored 77.89, which confirms our intuition that employing normal MoE performs better.
C. A Agreement Vector
This section elucidates the utilization of the Agreement Vector, proposed Hu al. (2024), metric for task grouping. For each task T. the agreement score vector is defined as follows: A(T.) @ & Σ fi. where 9i denotes the gradient of the parameters calculated from task T; and N denotes the total number of tasks. This vector reflects the gradient similarity between the task-speeific and the average gradients, and further indicate the task similarity.
17

---

Mastering Massive Multi- Task Reinforcement Learning via Mixture-of-Expert Decision Transformer
Table 6: Normalized score MW+DMC
Task Scale
10 Tasks
80 Tasks
160 Tasks
Method
Score
Parameters
Score
Parameters
Score
Parameters
MTDT-Small
70.97
1.47M
45.89
1.47M
35.34
1.47M
MTDT-Large
84.32
173.30M
59.01
173.30M
50.86
173.30M
PromptDT-Small
70.81
1.47M
42.09
1.47M
35.46
1.47M
PromptDT-Large
81.24
173.30M
59.40
173.30M
53.07
173.30M
HarmoDT-Small
65.13
1.47M
34.35
1.47M
28.05
1.47M
HarmoDT-Large
80.55
173.30M
59.91
173.30M
56.81
173.30M
M3DT-Random (Ours)
84.75
47.87M
66.93
98.37M
62.48
174.12M 12M
M3DT-Gradient (Ours)
84.77
47.87M
70.52
98.37M
65.90
174.12M
Table 7: Normalized Ant+Cheetah
Task Scale
10 Tasks
80 Tasks
160 Tasks
Method
Score
Parameters
Score
Parameters
Score
Parameters
MTDT-Small
94.09
1.47M
86.82
1.47M
83.04
1.47M
MTDT-Large
93.89
173.30M
89.30
173.30M
90.43
173.30M
PromptDT-Small
94.47
1.47M
88.15
1.47M
87.37
1.47M
PromptDT--Larg
93.61
173.30M
90.66
173.30M
90.11
173.30M
HarmoDT-Small
93.83
1.47M
86.31
1.47M
86.54
1.47M
HarmoDT-LLare
92.61
173.30M
90.75
173.30M
90.19
173.30M
M3DT-Random (Ours)
93.34
47.87M
90.55
98.37M
90.97
174.12M
M3DT-Gradient (Ours)
93.52
47.87M
90.41
98.37M
90.53
174.12M
18


Reviews:
[{"review_id": "ZbU1pg12ed", "reviewer": "Reviewer_is3p", "summary": "The authors study the problem of distilling a large multi-task offline RL dataset into a single policy via Prompt Decision Transformer (DT). The authors first study the scalability of Prompt DT with respect to both tasks and model size. These experiments provide a (very clear) demonstration of the theory that multi-task RL datasets diminish performance on each individual task due to conflicting gradients and show that larger models do not necessarily fix the problem. A common solution is to split parameters between tasks. This is usually done with separate output heads, but the authors take a different approach and use a Mixture of Experts model to split learning. This **M3DT** architecture is trained in a three-stage process (base model, experts, then router). Experiments find M3DT improves multi-task performance relative to several other DT methods and ablations.\n\n\n## Update After Rebuttal\n\nThe authors' explanation of some missing details and additional results breakdowns do partially address my concerns. I think it is a mistake to setup an empirical multi-task experiment like this and make half the task set solvable by meta-RL methods that do not need to address multi-task RL challenges. Results emphasize an under-discussed method for task subset selection and score normalization. \n\n**The main conclusion from the results is that splitting the architecture of a multi-task policy into independent components improves performance, which is a well established pattern. Mixture of Experts is a different and simple way to go about this (Fig 7), but the simplicity is somewhat undercut by a complex three-stage training process. New figures in the rebuttal confirm MoE performs worse than the joint single-task baseline**. I understand and appreciate the desire to avoid task IDs, but disagree that the original task set makes this an unachievable oracle. The lack of task IDs would be a more valid limitation if randomized variants of Meta-World + DMC taks were also being used, or if the domains were more difficult to distinguish from a single observation. The rebuttal task transfer results show limited zero-shot benefit from implicit task identification in MW+DMC, which is expected, but further support the idea that these tasks are not similar enough to demand ground-truth task IDs.\n\nI am very torn between a score of 2&3 but will maintain my original score of 2 partially because my concerns are quite different from other reviews.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate the reviewer's thorough reading of our manuscript and your valuable comments regarding out task setup and method. These comments have helped us realize that we inadvertently omitted several critical details in our original manuscript. We believe that supplementing and clarifying this information will significantly enhance the quality of our paper and we hope these can address your concerns.\n\n**Regarding the use of Ant-dir and Cheetah-vel as 80 tasks**, our intention is not to propose a new benchmark, but to ensure sufficient task quantity to investigate how task scale affects MTRL scalability. To achieve methodological rigor in task selection, we follow three key criteria:\n\n(1) Performance Balance: Based on single-task score, the selected task subsets maintain average scores comparable to the full 160 set, eliminating the influence of task difficulty\n\n(2) Domain Ratio Preservation: The subset composition at different task scales approximately maintains the original domain ratio (MW:DMC:Ant:Cheetah=5:3:4:4), mitigating domain-specific biases\n\n(3) Subset Diversity: We use distinct task subsets for each run with the same task scale to maximize task coverage, eliminating correlations between specific tasks.\n\nThese eliminate potential bias caused by the relative ease of Ant and Cheetah tasks. We have separately reported the scores for MW+DMC and Ant+Cheetah tasks,which are consistent with the average scores across160 tasks in manuscript, validating the rationality of task selection and reliability of our findings.  Let me know if link not work https://anonymous.4open.science/r/ICML_rebuttal_11746/revised_fig2.png\n\n>Q1. Calculation for Average Normalized Score\n\nA1.  We assign equal weight to each of the 160 tasks and compute the average normalized scores. Ant and Cheetah collectively account for 50% of the average score.\n\n>Q2. \"In MTRL, increasing ...\" maybe because PromptDT has nearly perfect performance at original size. \"In the clear ...\" is caused by repeatedly sampling the same task.\n\nA2. In seperated results, both task scalability and parameter scalability exhibit the same trends as demonstrated in the original paper. Detailed analyse please see https://anonymous.4open.science/r/ICML_rebuttal_11746/revised_fig2.png.\n\n>Q3. \"Random ...\" maybe because the conflict between them are not significant enough.\n\nA3. We conduct experiments by separately training PromptDT-1.47M on 80 MW+DMC and 80 Ant+Cheetah tasks to show gradient conflicts and normalized scores,see https://anonymous.4open.science/r/ICML_rebuttal_11746/Q4.png. The results show that training on either standalone domain exhibits significantly weaker gradient conflicts. \n\n>Q4. \"M3DT...\" M3DT appears to learn several single-task experts and a switch function.\n\nA4. M3DT fundamentally differs from training single-task policies and a switch function: (1) our router dynamically assigns weights to experts at token-level, rather than switching to the single best expert at task-level (dashed line in Fig.7), and Fig.7 demonstrates our router outperforms oracle switching when expert count is small. (2) reducing tasks per expert doesn't guarantee better performance, as evidenced by the drop when scaling from 48 to 56 experts (175M to 200M in Fig.7 dashed line), which shows contribution that identifying the optimal task load for given parameter subsets and strike a trade-off between expert performance and gating difficulty.\n\n>Q5. How would the main results change if Cheetah-Vel and Ant-Dir were completely ignored?\n\nA5. M3DT shows more superior performance on complex tasks(MW+DMC). See https://anonymous.4open.science/r/ICML_rebuttal_11746/domain_main_result.png\n\n>Q6. Evaluating on held-out tasks, a key advantage of Prompt DT rather than one-hot identifying task.\n\nA6. We first emphasize that both M3DT and PromptDT operate without task ID and only use trajectory prompt during held-in task evaluation (see Rebuttal A8), whereas MTDT and HarmoDT require task IDs to function. For held-out evaluation, see https://anonymous.4open.science/r/ICML_rebuttal_11746/generalization.png\n\n>Q7. At least a significant reduction in parameters(vs. switching single-task experts)\n\nA7. We think the core challenge in MTRL lies in the limitation that simply increasing parameters fails to address large-scale tasks. **Only after achieving competent performance at large-scale tasks does parameter reduction become meaningful**. But parameter efficiency also remains a valuable future work, particularly through refined architectural designs for experts and router. While single-task agents totally require 235.2M parameters with an average score of 84.35, our method achieves a 26% reduction in parameters.\n\n>Q8. Figure 7 is a confusing result. Why use the router? Compare the average single-task performance?\n\nA8. Please see Reviewer aD2s Q4A4. We add the averaged single-task performance in the revised Fig.7 at https://anonymous.4open.science/r/ICML_rebuttal_11746/revised_fig7.jpg", "claims_and_evidence": "The main claims of the paper are supported by many experiments. However, I have several questions and concerns about the conclusions from the experiments, discussed below.", "methods_and_evaluation_criteria": "To aggregate results across domains with different reward functions, the authors report every result as a normalized score. However, they are missing breakdowns of the normalized score by domain. To reach 160 tasks, the authors use tasks from three domains. 50 tasks are from Meta-World, and 30 are from the DM Control Suite (DMC). These choices make sense because previous work has demonstrated the key gradient conflict issue here (especially in Meta-World). In the main text, the authors state the third domain is “Mujoco Locomotion,” and I think a reasonable guess from that description would be the standard gym tasks that are similar to the DMC. Appendix A.1.3 clarifies: **80 of the 160 total tasks used are minor variations of two locomotion problems**. \n\n**I am very skeptical that “40 tasks” from Cheetah-Vel  and “40 tasks” from Ant-Dir should count as 80 tasks in the context of this paper**. A meta-learning paper (including the Prompt DT paper the dataset is taken from) would call this “40 tasks”, but this is a formality of the meta-RL problem statement. Meta-RL is interested in identifying an unknown MDP at test-time, while many of the baselines here are provided with the task ID upfront. The bar for a unique \"task\" is much lower in meta-RL. In any case, these are very saturated meta-RL benchmarks: **it is easy to solve all 40 of these tasks online with a small model and without task IDs. Gradient conflicts are a non-issue**. When the standard of a new task is that it creates a distinct optimization objective, as in this paper, I think we should treat these \"80 tasks\" as two tasks. This makes the 160 task count a bit misleading. **My main concern is that the normalized score is actually 50% (80 tasks / 160) reliant on the performance in these two problems**. The writing does not make this totally clear, but reporting the score for 120 and 160 tasks as separate metrics suggests this is the case. If so, it significantly changes many of the takeaways upon a re-read. For example:\n\n> **Conclusion**: *“In MTRL, increasing the model size of DT swiftly hits the performance ceiling”*. **Reinterpreted**: yes, maybe because 50% of the eval metric relies on 2 tasks where Prompt-DT has nearly perfect performance at its original model size. This could explain why the margins are thin across the paper.\n\n> **Conclusion**: *“In the clear trend where the performance degrades with increasing task numbers, the decline is pronounced when task number is relatively low (below 40 tasks), while it becomes much more gradual once the tasks reach a sufficiently large number (above 80 tasks).”* **Reinterpreted**: the diminishing rate of decrease is caused by repeatedly sampling the same task. This is complicated by Appendix A.3, where the authors clarify that task sampling is not uniform. Instead, the subsets are designed such that the average performance of single-task DTs is consistent with the average performance of single-task DT on all 160 tasks. If that metric is 50% Cheetah/Ant, and PromptDT performance is nearly 100% successful on those tasks, perhaps the samples are biased by Cheetah/Ant, but too few details are given to analyze this. The subsets used to measure performance vs. task count might be skewed in some way. \n\n\n> **Conclusion**: Random task groupings work almost as well as the more sophisticated gradient-similarity task groupings (Tables 1 & 2). **Reinterpreted**: randomly spreading the Cheetah/Ant tasks among the experts is fine because the conflict between them are not significant enough to stop Transformers of this size from learning with standard online RL, much less supervised DT.\n\n\n> **Conclusion**: *“M3DT… shows diminishing performance gains after reaching 40 experts on 160 tasks.”* **Reinterpreted**: By the time we reach 40 independently trained models on 82 effective tasks, gradient conflicts are no longer an issue. We are approaching the baseline of learning a switch function between single-task expert models. Ideally, MTRL might be better than that method by positive transfer, but offline methods often treat it as an oracle. From that perspective, 48 experts work the best (Fig 5)  because it is closest to 1 expert per task. The authors set out to distill a multi-task dataset into a single policy, but the solution is to set up a MoE architecture and scale it up 40x until it resembles training a separate policy for every task. \n\n\n**It would be very useful to see the results with Ant-Dir and Cheetah-Vel removed.**\n\nMy apologies for the long discussion, but this is a rare case where I felt an appendix detail may substantially change my opinion on a paper's results, and that not all reviews may comment on this.", "theoretical_claims": "N/A", "experimental_designs_or_analyses": "Content that might normally go here was discussed in a section above.", "relation_to_broader_scientific_literature": "This paper relates to important topics in multi-task optimization in RL and the use of large-scale sequence models for policy learning.", "essential_references_not_discussed": "Nothing critical.", "other_strengths_and_weaknesses": "The first few figures in this paper clearly demonstrate the core MTRL optimization challenge and would support many existing works in this area. The method is a good example of applying ideas from large-scale LLMs to RL, an emerging trend and an important area for study.\n\nThis line of work involves training single-task agents from scratch on every task and then distilling their behavior into a single policy. The main benefit of this might be generalization to unseen tasks or at least a significant reduction in parameters to handle the training set (vs. switching between the single-task experts). The paper does not discuss generalization or fine-tuning to new tasks, and the architecture almost directly scales with task count.\n\nFigure 7 is a confusing result. The MoE experts outperform the overall model (including the router), which should be bad (why use the router then?), but the writing acts like this is falling short of an oracle upper bound. The standard upper bound is usually the aggregate results of (smaller) single-task specialist agents (e.g., Multi-Game DT, Gato). Because the experts are mostly redundant architectures to a base model that trains on its own, it is not 100% clear to me whether Figure 7 is reporting something different than this. Regardless, **Appendix A.3 makes it clear that, at some point, the authors trained a single-task Prompt DT model on all 160 tasks and used its score for balancing the dataset. Can we compare the performance of M3DT against those scores?**", "other_comments_or_suggestions": "The paper would benefit from evaluating on held-out tasks, which should be a key advantage of using Prompt DT as the base model rather than one-hot identifying every task.\n\nMinor: The Cheetah-Vel reference is [MAML](https://arxiv.org/abs/1703.03400), and the Ant-Dir reference is [ProMP](https://arxiv.org/abs/1810.06784).", "questions_for_authors": "Please clarify whether the 80 Cheetah-Vel/Ant-Dir tasks receive equal weight in the normalized score as the 80 distinct Meta-World+DMC tasks. How would the main results change if Cheetah-Vel and Ant-Dir were completely ignored?\n\nCould you clarify Figure 7 and the aggregate score of the 160 single-task Prompt DTs? See the last paragraph of “Other Strengths and Weaknesses” for context.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "h49lJbyOIG", "reviewer": "Reviewer_BLnM", "summary": "This paper introduces M3DT, a novel mixture‐of-experts (MoE) extension of the Decision Transformer designed to tackle the scalability challenges in massive multi-task reinforcement learning. The method leverages task grouping, dedicated expert modules, and a three-stage training mechanism to reduce gradient conflicts and enable effective parameter scaling, demonstrating improved performance across up to 160 tasks on challenging benchmarks.\n\n**Update after rebuttal**--I appreciate the authors’ thorough revision and the additional explanation addressing my concerns. I have decided to keep my original score.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are greatly appreciative of your recognition of our work. We offer our responses to address your concerns as follows and will address the formatting issues in Other Comments in the revised manuscript.\n>Q1. how “massive” these benchmarks are.\n\nA1. Current standard MTRL algorithms typically handle a limited number of tasks: offline methods are generally confined to 50 Meta-World tasks, while some online methods are restricted to MT10 tasks. In contrast, our approach integrates 160 tasks from Meta-World, DMC, and MuJoCo Locomotion.\n\n>Q2. Strengthening theoretical analysis would add depth to the work.\n\nA2. Thanks for your suggestions. But our method is proposed based on observed experimental phenomena, making theoretical analysis difficult. Neither DT nor MoE we used has been theoretically analyzed. While theoretical methods are rare in MTRL compared to MTL, our future work may focus on developing a theoretically grounded MTRL algorithm.\n\n>Q3. How online MT methods would perform in massive tasks.\n\nA3. Recent studies have shown offline MTRL methods outperform online methods on Meta-World MT50 benchmark, as seen in HarmoDT and MTDiff. This performance gap may widen significantly when scaling to larger task sets, where online methods often degrades. As task number increases, the required wall-time overhead for environment interaction grows substantially, leading to a significant expansion in training time.\n\n>Q4. Discuss more about MoE-based MTRL approaches.\n\nA4. Thanks for your suggestion. We will supplement and discuss [1,2] in the revised manuscript.\n\n[1] MoE-Loco: Mixture of Experts for Multitask Locomotion\n\n[2] Multi-Task Reinforcement Learning With Attention-Based Mixture of Experts\n\n\n>Q5. How you measured “gradient conflict”? Did you observe the frequent gradient conflicts with more tasks?\n\nA5. We compute the gradients for each task and calculate the mean gradient across all tasks. The gradient similarity is defined as the average cosine similarity between this mean gradient and all tasks' gradients. During training, we record this metric every 1e4 steps to obtain a similarity curve. After smoothing to the stabilized phase of training, we use the plateau value as our final gradient similarity. Gradient conflict is calculated as (1 - gradient similarity). Yes, our results confirm that the gradient conflicts intensifies with more tasks.\n\n>Q6. How the three-stage training specifically prevents overfitting to dominant tasks?\n\nA6. First stage, we terminate full-task training when performance plateaus and gradient conflicts peak (Fig.6(left)) to avoid overfitting to dominant tasks. Second stage, we use task grouping to reduce task numbers each expert learns, mitigating gradient conflicts and preventing overfitting. Third stage, we use well-trained experts to balance router learning on each task.\n\n>Q7. Could the approach be extended to online MT settings, what modifications would be necessary?\n\nA7. M3DT can be readily extended to online MTRL by simply replacing the backbone DT's training objective and optimization with those of Online DT and iterating through all tasks for parameter optimization. However, as task numbers increase, the time cost of interacting with all tasks becomes prohibitive, hence we still recommend offline as the primary paradigm.\n\n>Q8. When to use M3DT-Random or M3DT-Gradient? What’s the criteria?\n\nA8. M3DT-G consistently outperforms M3DT-R. The latter is designed as a baseline to validate that even random task grouping improves performance by reducing the learning task load per parameter subset. M3DT-G requires only two additional steps, computing gradients across all tasks and performing K-means clustering, which account for a little increase in training time. We recommend use M3DT-G as default approach.\n\n>Q9. Regarding the router, what specific design choices (architecture of the MLP) led to its effectiveness?\n\nA9. Thanks to our analysis of task quantity and the proposed algorithm, we achieve strong results using a basic 5-layer MLP with ReLU activation for the router, without any specialized design. We believe optimizing the router structure or routing mechanism could further enhance performance or extend our method to generalization and continual learning scenarios, which we leave as future work.\n\n>Q10. Normalized score is a key evaluation metric across diverse tasks. The normalization process used for each benchmark?\n\nA10. As stated in Appendix A.1, we perform score normalization separately for each domain's tasks. For MW, we use the success metric. For DMC, we map the original range [0, 1000] to [0, 100]. For Cheetah-vel, we map [-100, -30] to [0, 100]. For Ant-dir, we map [0, 500] to [0, 100]. For the latter two domains, scores outside the original range are clipped to 0 or 100. For overall average normalized score across all tasks, we assign equal weight to each task, computing the algorithm's final score as the uniform average over all 160 tasks.", "claims_and_evidence": "The authors claim that naive parameter scaling fails as task numbers increase and that their proposed MoE approach, with its task-specific grouping and staged training, significantly mitigates performance degradation. These claims are supported by extensive experiments and ablation studies comparing M3DT with existing baselines, although additional discussion on the statistical significance of improvements would be beneficial.", "methods_and_evaluation_criteria": "The method is well-motivated and appropriate for addressing multi-task RL scalability. The use of standard benchmarks (Meta-World, DMControl, Mujoco, though mixed!) and normalized performance metrics strengthens the experimental evaluation. But my main concern is how “massive” these benchmarks are.", "theoretical_claims": "While the paper offers empirical insights into gradient conflicts and parameter scaling, it lacks rigorous formal proofs. Strengthening the theoretical analysis, particularly regarding the benefits of the three-stage training mechanism, would add depth to the work.", "experimental_designs_or_analyses": "The experimental design is comprehensive, including comparisons with state-of-the-art DT-based baselines and detailed ablation studies. Also, I’m curious how online multi-task methods would perform in these “massive” tasks.", "relation_to_broader_scientific_literature": "M3DT builds clearly on recent advances in DT and MoE architectures for multi-task RL. It addresses a significant gap in scaling multi-task RL, and the authors relate their work to relevant literature, although a discussion of some recent MoE approaches in reinforcement learning could provide additional context.", "essential_references_not_discussed": "It would be great if the authors could discuss more about the closely related MoE-based multi-task RL approaches (i.e., Hendawy et al. (2023) and Huang et al. (2024)).", "other_strengths_and_weaknesses": "Strengths include a novel approach to reducing gradient conflicts, a well-structured three-stage training process, and robust experimental validations. On the downside, parts of the presentation are dense and could benefit from clearer exposition, and further theoretical justification of the methods would strengthen the paper.", "other_comments_or_suggestions": "- Font size in the figures is too small.\n- Inconsistency: Mixture-of-Expert v.s. Mixture-of-Experts\n- In table 1 (and check others as well), please double check to bold the values that are in a statistically significant range.", "questions_for_authors": "- Could you elaborate more on how you measured “gradient conflict”? Did you observe the frequent occurence of gradient conflicts in cases with more tasks?\n- Can you elaborate on how the three-stage training mechanism specifically prevents overfitting to dominant tasks?\n- Could the approach be extended to online multi-task settings, and if so, what modifications would be necessary?\n- When should we use M3DT-Random? and M3DT-Gradient? What’s the criteria suggested by the authors?\n- Regarding the router component, what specific design choices (e.g., architecture of the MLP) led to its effectiveness in dynamically allocating expert weights?\n- The normalized score is a key evaluation metric across diverse tasks. Could you elaborate on the normalization process used for each benchmark?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "HpLfj9wI9p", "reviewer": "Reviewer_aD2s", "summary": "Transformer-based models have recently shown success in offline reinforcement learning by framing the problem as a sequence modeling problem. Moreover, offline multi-task reinforcement learning (MTRL) has benefited from the high capacity of these models for solving complex and diverse tasks. Nevertheless, as the number of tasks massively increases, the overall performance starts to drop tremendously. The naive scaling of the transformer-based models doesn't counteract the drop in performance caused by scaling the number of tasks. The performance of the model saturates eventually as the model capacity increases. Accordingly, we are not gaining from the scalability of the transformer-based models in MTRL as done in the supervised learning setting. In this work, the authors propose a different way to scale the model by learning different experts for different groups of tasks. With that, scaling the model is in the form of increasing the number of experts, hence increasing the number of tasks to be learned. Moreover, this reduces the conflicts between tasks since the number of tasks per expert can be lower. The proposed method, named M3DT, consists of three stages of learning: a backbone transformer model, experts, and a router model. The approach has been evaluated on 160 tasks, a mix of tasks from metaworld, dm-control, and continuous control problems from Mujoco, against relevant and recent offline MTRL baselines with a transformer-based architecture.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate your recognition of our work. We offer our responses to address your concerns as follows and will supplement these details and correct the typo in the revised manuscript.\n\n>Q1. The limitations are not well-discussed. It is important to understand how to fix some limitations or extend this work to future directions.\n\nA1. We primarily focus on the challenges in MTRL, analyzing the scalability of task numbers and parameter size, and proposing a paradigm that increases model parameters to achieve task scalability. However:\n\n(1) we did not focus on fine-grained network architecture design, such as optimizing the expert and router architecture, which could potentially further improve performance and reduce model size.\n\n(2) While we leave held-out task generalization and continual learning unexplored, our method’s modular parameterization and group-wise learning naturally support solutions like held-out task adaptation via expert learning, dynamic skill composition for held-out tasks, and forgetting mitigation for the router for past tasks. Addressing these areas presents promising future work.\n\n(3) Scaling experts with task count raises inference costs. While we prioritize performance, we experimented with activating Top-k experts with fixed k to control inference overhead, but it degrades performance.  A more tailored Top-k gating mechanism, better aligned with our three-stage training paradigm, could enhance our practical applicability.\n\n>Q2. The training time for the proposed approach compared to the baselines.\n\nA2. Our experiments were conducted on RTX 4090. We train PromptDT-5M for 4e5 steps in first stage, taking around 5.2h. In second stage, all experts are trained independently on their dedicated task group in parallel, taking around 1.8 hours (the time for training a single expert). In third stage, we train the router for 4e5 steps, taking 17.2h. The reason for the long training time in third stage is our code's lack of parallel computation for experts. Code optimization may reduce this training time. **The total training time of M3DT-174M is around 24.2 hours. For comparison, training PromptDT-173M and MTDT-173M each takes around 21.3 hours, while training HarmoDT-173M requires 95.6 hours.**\n\n>Q3. In Fig.7, the performance eventually saturates even when the number of experts increases.\n\nA3. As stated in Section 5.2 (line 402-415), the performance plateaus for two main reasons:\n\n(1) The router's weight allocation becomes increasingly difficult with more experts;\n\n(2) Each expert's task load becomes small enough to achieve optimal performance, preventing further improvement in individual expert performance (dashed line). \n\nThus, the overall performance eventually saturates.\n\n>Q4. The expert performance seems better than the performance with the router; why do we need a router then? \n\nA4. Sorry for your confusion. The \"Expert Performance\" is obtained by manually selecting each test task's corresponding expert based on Task ID and individually evaluating their performance. However, this approach is impractical for large-scale tasks in our overall method evaluation and real-world applications: (1) the specific task IDs are usually unavailable, (2) manual expert selection is prohibitively expensive or impossible. Therefore, a unified policy is essential. The router automatically assigns weights to each expert based on the backbone's hidden states, integrating all sub-policies into a unified policy without needing task ID. This necessitates the router. We consider it reasonable the framework's overall performance may be lower than the ideal case of perfect expert switch, as increasing the number of experts raises the difficulty of optimal weight allocation. We will rename \"Expert erformance\" to \"Oracle Expert Selection Performance\" in the revised manuscript and add the above explanation.\n\n>Q5. Why freeze the predict head while the corresponding input has been changed because of the training of the router?\n\nA5. **Through experiments, we evaluate the performance of M3DT-G with simultaneous training of both router and predict head, achieving a score of 76.86.** This result is slightly lower than our original algorithm's 77.89, but outperforms other ablation variants. The rationale for freezing the predict head is:\n\n(1) The predict head trained in the first stage already contains knowledge for all tasks, and continued training of the backbone parameters across all tasks would lead to overfitting due to severe gradient conflicts (as shown in Fig.6); \n\n(2) In the second-stage expert training, we also freeze the predict head, so all learned experts' outputs remain compatible with the predict head's input. In the third stage, the router performs weighted sum of all expert outputs (with weights normalized by Softmax to sum to 1), which means the entire MoE's final output can be viewed equivalent to a single expert's output, thereby preserving compatibility with the predict head's expected input.", "claims_and_evidence": "- In my opinion, the claims presented in this work are quite clear and well-motivated. \n- The claims are supported by two studies (in Figures 2 and 3), which makes them very convincing.", "methods_and_evaluation_criteria": "- I find the proposed approach, named M3DT, sound for tackling the problem of multi-task reinformcenet learning in the offline setting.\n- The massive number of tasks used for evaluation is interesting and quite motivating for the effectiveness of the approach.", "theoretical_claims": "- There are no theoretical claims presented in this work.", "experimental_designs_or_analyses": "- I highly appreciate the amount and the diversity of the experiments and ablations done in this work. \n- As mentioned before, I like the two studies presented in Figures 2 and 3.", "relation_to_broader_scientific_literature": "- I believe this work is important to close the gap between supervised learning and reinforcement learning in terms of the scaling of the transformer-based models. \n- I agree with the authors that offline MTRL has been limited to dozens of tasks without looking into more realistic scenarios that we expect from offline RL.\n- Hopefully, this will affect the online MTRL setting to handle such massive number.", "essential_references_not_discussed": "- I have no recommendations for papers to be cited.", "other_strengths_and_weaknesses": "- I found some weaknesses in this work,\n  - The limitations of this approach are not well-discussed.\n  - One important limitation, in my opinion, is the training time since M3DT requires 3 stages of training. I would advise the authors to indicate the training time for the proposed approach, compared to the baselines.\n  - I believe the future works were never clearly stated in this work. It is important to understand how to fix some limitations in this work or how to extend this work to future directions.\n  - In Figure 7, it seems that the performance of the proposed approach eventually saturates even when the number of experts increases. I would encourage the authors to comment on that.", "other_comments_or_suggestions": "- In section 2 (Preliminaries), subsection 2 (Multi-Task RL), second line, there is a small typo in funct(u)ions.", "questions_for_authors": "- How long is the training for M3DT compared to the other baselines?\n- The expert performance seems better than the performance with the router; why do we need a router then? Am I missing something?\n- Why do you freeze the *predict head* while the corresponding input has been changed because of the training of the router in the third phase?\n\nI would be happy to raise my score if my concerns are addressed.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "J2BK68Bfvc", "reviewer": "Reviewer_NqBs", "summary": "This paper studies the Massive Multi-Task Reinforcement Learning problem and proposes M3DT. When the number of tasks is large, it shows great performance improvement compared with former methods.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank the reviewer for the thoughtful feedback and questions that will surely turn our paper into a better shape. \n\n>Q1. The logic in insight \"Reducing ...\" is not smooth. Fig.2 does not tell us reducing 160 tasks to a small number of groups can enhance the performance. First, the performance may highly depend on the grouping strategies. Second, there may still exist conflict within task groups.\n\nA1. (1) As stated in Section 3.1 (bolded \"Performance Variance\"), we test on different task sets each run. The STD reflects the model’s robustness to variations in task combinations. When the task count is reduced to 10, the model exhibits minimal STD, demonstrating the good performance is **robust to the grouping strategies**.\n\nNote: The specific criteria for constructing task subsets can be found in Reviewer is3P top\n\n(2) The gradient similarity on 10 tasks is 0.29 in Fig.2, indicating **gradient conflicts still exist within task groups**.  However, compared to massive tasks, the **gradient conflict is significantly reduced** when training on 10 tasks, and model achieves better scores.\n\nTherefore, by interpreting Fig.2 from the **reverse perspective (from right to left)**, we derive our robust insight. Importantly, this insight does not advocate focusing solely on small-scale tasks. Given the predetermined task count the model must solve, how to reduce task numbers to be learned is precisely the problem we raised in Section 3.2 and addressed in our method.\n\n\n>Q2. In Figure 4, the model is drawn bottom-top, but the task sets appear above output layer with downward arrows, which may cause confusion.\n\nA2. Thanks for identifying this. We upload the revised figure at link(https://anonymous.4open.science/r/ICML_rebuttal_11746/revised_fig4.JPG) and will replace in the revised manuscript.\n\n>Q3. How to determine the number of groups?\n\nA3. We align the number of task groups with experts, with each expert dedicated to a distinct task subset. Increasing the number of groups = more experts = more model parameters. Through experimental analysis in Fig.5 and Fig.7, we identify the optimal group (expert) counts for different task scales, which reflects the trade-off between task subset learning difficulty and gating difficulty. We will add these details to the revised manuscript.\n\n>Q4. What is the grouping frequency?\n\nA4. Task grouping is performed **only once after completing the first stage**. The overall algorithm workflow is backbone training in first stage, task grouping, each expert individually training on its corresponding fixed task group in second stage, and router training in third stage. The algorithm concludes after the third stage. **Each stage executes once** in the entire training process, with detailed training iterations per stage specified in Appendix A.5. We will supplement this in pseudocode in revised manuscript.\n\n>Q5. M3DT-G needs to calculate all gradients. Computation cost is high.\n\nA5. Since task grouping is performed only once, M3DT-G requires merely a single additional computation of all tasks' gradients followed by K-means clustering, which is negligible compared to the overall training cost.\n\n>Q6. Is there any special treatment for tasks within the same group?\n\nA6. We do **not** apply any special processing or treatment to any of the tasks.\n\n>Q7. In Figure 5 left, why is performance drop from 98M to 123M?\n\nA7. As analyzed in Section 5.2 (line 402-415) and Fig.7, the overall performance is a trade-off between expert performance and gating difficulty. When the expert count exceeds a threshold, further increasing the expert count no longer improves expert performance, while the complexity of router assigning weights to experts continues to improve. Thus, the performance plateaus or even degrades.\n\n>Q8. Why do not authors simply group tasks and train MoE?\n\nA8. Through experiments, we evaluate the ablation: group tasks and train MoE (simultaneously training router and task-ID-matched expert while freezing other experts), which achieved **71.21**, compared to 77.89 in Table 2. The performance drop primarily stems from the instability caused by alternating optimization of different experts. It also fails to leverage the the second stage, where experts can be trained independently and in parallel, resulting in significantly prolonged training duration.\n\n>Q9. I feel the methods are general and can be applied in supervised MTL. What specific challenges in MTRL do authors find and solve? \n\nA9. In RL, task distinctions are more pronounced and decision-making process is more sensitive compared to regression and classification. While scaling up model parameters has proven effective for handling numerous tasks in other fields, it falls short in RL. This is the core challenge that our paper addresses. Our method can be readily adapted to other supervised MTL domains by simply replacing the backbone with other task-specific models. We hope this work can inspire new research in MTL and advance the field.", "claims_and_evidence": "In section 3.1, the authors propose an insight \"Reducing the learning task number, particularly to a sufficiently small scale can significantly enhance model performance\" according to Figure 2. However, I do not feel the logic is smooth here. \nIn my opinion, Figure 2 only shows that the more tasks there are, the more conflict and worse performance they have. \nThis paper studies massive MTRL where the number of tasks is larger, say 160. Then Figure 2 does not tell us reducing 160 tasks to a small number of groups, say 20, can enhance the performance. First, the performance may highly depend on the grouping strategies. Second, there may still exist conflict within task groups. Thus, in terms of the presentation, this claim is not appropriate.", "methods_and_evaluation_criteria": "Yes, the evaluation and methods make sense to me.", "theoretical_claims": "There is no theoretical analysis in this paper.", "experimental_designs_or_analyses": "This paper studies the massive multi-task RL setting where the number of tasks could be 160, which is sound to me.", "relation_to_broader_scientific_literature": "This paper is highly related to current literature in this area.", "essential_references_not_discussed": "I am not sure whether all essential references are included.", "other_strengths_and_weaknesses": "1. In Figure 4, the model architecture is drawn from bottom to top. However, the small figures for those tasks are drawn above the output layer with a \"down arrow\". This may raise some confusion.\n2. The description of the method/algorithm is not clear, especially for the grouping part. I have the following questions but failed to find out the answer in the paper: \na) How to determine the number of groups? b) What is the grouping frequency, such as grouping every epoch? c) For M3DT-G, it needs to calculate all gradients. Then, the computation cost is pretty high. d) Is there any special treatment for tasks within the same group? For example, do authors treat them equally important and use linear scalarization perhaps? \n3. In Figure 5 on the left, why is there a performance drop from 98M to 123M?\n4. I appreciate the ablation study in section 5.2 to support the method design. However, I do not understand why to separate the expert training and router training because they are trained simultaneously in the original MoE with load-balancing loss. Why do not authors simply group tasks and train MoE layers (router+experts)?\n5. Lastly, this paper studies the MTRL problem. Nevertheless, I feel the methods are general and can be applied in supervised MTL as well. What specific challenges in MTRL do authors find and solve? The novelty in the current version looks like using grouping and MoE in MTL.\n\nI am happy to change my mind if these questions are answered or if I misunderstand the paper.", "other_comments_or_suggestions": "Please check weaknesses.", "questions_for_authors": "Please check weaknesses.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper tackles the challenge of scaling offline multi-task reinforcement learning (MTRL) to a massive number of tasks, proposing M3DT, a Mixture-of-Experts (MoE) framework built upon the Decision Transformer (DT). The authors first demonstrate empirically that naively increasing model size offers diminishing returns as task numbers grow due to issues like gradient conflicts. M3DT addresses this by incorporating MoE layers into the DT architecture and employing a three-stage training process (backbone, experts, router) to manage task complexity and improve scalability, showing results on up to 160 tasks.   

Reviewers acknowledged the importance of the problem and appreciated the empirical studies highlighting the limitations of standard scaling. The MoE approach was seen as a relevant direction. However, significant concerns were raised, particularly regarding the experimental setup and interpretation.   

Reviewer is3p strongly criticized the task composition, noting that 80 out of the 160 tasks were minor variations of two locomotion problems (Ant-Dir, Cheetah-Vel). This reviewer argued that these should effectively count as only two tasks in the context of MTRL (as opposed to meta-RL), and their inclusion significantly skews the results and conclusions about scalability, performance saturation, and the effectiveness of grouping strategies. While the authors provided breakdowns by domain in the rebuttal, Reviewer is3p remained unconvinced, arguing the setup obscures rather than clarifies the challenges of massive MTRL and that the core finding (parameter splitting helps) is already well-established. This reviewer also found the MoE approach somewhat undermined by the complex three-stage training and maintained a weak reject score.   

Reviewer NqBs also initially recommended weak reject, questioning the logical connection between the initial experiments and the proposed solution, finding the method description (especially task grouping and training stages) unclear, and asking about the specific MTRL challenges addressed beyond general multi-task learning. The authors provided clarifications on the logic, grouping mechanism (performed once after stage 1), training rationale, and MTRL challenges. Reviewer NqBs acknowledged the rebuttal but did not provide further feedback or update their score.   

Reviewers BLnM and aD2s were more positive. Reviewer BLnM (Weak Accept) found the method well-motivated but sought clarification on gradient conflict measurement, the 3-stage training benefits, online applicability, router design, and score normalization. Reviewer aD2s praised the motivation and experiments but noted the lack of discussion on limitations and training time, also questioning the performance saturation and the expert vs. router performance discrepancy in Figure 7. The authors addressed these points, providing details on gradient similarity calculation, the 3-stage training rationale, training times, performance saturation explanations, and clarifying Figure 7 by distinguishing oracle expert selection from the router's role. These responses satisfied Reviewer BLnM (who maintained Weak Accept) and Reviewer aD2s (who raised their score to Accept).

While the paper addresses an important problem in scaling MTRL and proposes a relevant MoE-based approach, the significant concerns raised by Reviewer is3p regarding the experimental design's validity cast doubt on the central claims about mastering "massive" MTRL. The heavy reliance on numerous variations of just two locomotion tasks potentially increases the task count and might overstate the method's scalability and effectiveness in diverse, large-scale multi-task settings. Although the authors provided additional results and clarifications that satisfied two reviewers (leading to one score increase), the fundamental critique about the task composition from Reviewer is3p was not fully resolved.
Decision: Accept (poster)