Learning and Planning Multi-Agent Tasks via an MoE-based World Model
Authors: Zijie Zhao
Zhongyue Zhao
Kaixuan Xu
Yuqian Fu
Jiajun Chai
Yuanheng Zhu
Dongbin Zhao
Venue: NeurIPS 2025 poster ()
Keywords: Multi-task learning
Multi-agent reinforcement learning
Mixture of Experts
World model
Primary area: reinforcement_learning
TL;DR: We propose a novel framework  employs a MoE-based multi-agent world model for learning and planning in multi-task scenarios.

Abstract: Multi-task multi-agent reinforcement learning (MT-MARL) aims to develop a single model capable of solving a diverse set of tasks. However, existing methods often fall short due to the substantial variation in optimal policies across tasks, making it challenging for a single policy model to generalize effectively. In contrast, we find that many tasks exhibit **bounded similarity** in their underlying dynamics—highly similar within certain groups (e.g., door-open/close) diverge significantly between unrelated tasks (e.g., door-open \& object-catch). To leverage this property, we reconsider the role of modularity in multi-task learning, and propose **M3W**, a novel approach that applies mixture-of-experts (MoE) to world model instead of policy, enabling both learning and planning. For learning, it uses a SoftMoE-based dynamics model alongside a SparseMoE-based predictor to facilitate knowledge reuse across similar tasks while avoiding gradient conflicts across dissimilar tasks. For planning, it evaluates and optimizes actions using the predicted rollouts from the world model, without relying directly on a explicit policy model, thereby overcoming the limitations of policy-centric methods. As the first MoE-based multi-task world model, M3W demonstrates superior performance, sample efficiency, and multi-task adaptability, as validated on Bi-DexHands with 14 tasks and MA-Mujoco with 24 tasks. The demos and anonymous code are available at \url{https://github.com/zhaozijie2022/m3w-marl}.

Learning and Planning Multi-Agent Tasks via an MoE-based World Model
Zijie Zhao' 1.2, Zhongyue Zhao", Kaixuan Xu², Yuqian Fu Jiajun Chai Yuanheng Zhu Dongbin Zhao 1 School of Artificial Intelligence, University of Chinese Academy of Sciences Institute of Automation, Chinese Academy of Sciences, Meituun jihaozijie2022, zhaozhongyue2024, xukaixuan2025 fuyuqian2022))0ia. jhaijiaaunSmeituun com, {yuanheng. zhu, dongbin zia. *ooreepponiing authors"
Abstract
Multi-task multi-agent reinforcement learning (MT-MARL) aims develop single model capable of solving diverse set tasks. However, existing methods often fall short due the substantial variation in optimal policies across tasks, making it challenging for a single policy model to generalize effectively. In contrast, we find that many tasks exhibit bounded similarity in their underlying dynamics-highly similar within certain groups (e.g., door-open/close) diverge significantly between unrelated tasks (e.g., door-open & object-catch). To leverage this property. we reconsider the role of modularity in multi-task learning, and propose M3W, novel approach that applies mixture-of-experts (MoE) world model instead of policy, enabling both learning and planning. For learning, it uses a SoftMoE-basee dynamics model alongside SparseMoE-basser predictor facilitate knowledge reuse across similar tasks while avoiding gradient conflicts across dissimilar tasks. For planning, it evaluates and optimizes actions using the predicted rollouts from the world model, without relying directly on explicit policy model, thereby overcoming the limitations of policy-centrie methods. the first MoE-based multi-task world model, M3W demonstrates superior performance, sample efficiency, and multi-task adaptability, as validated on Bi-DexHands with tasks and MA-Mujoco with 24 tasks. The code are available at https: ://i thub. htmm/hhazziiiee2222 marl.
Introduction
The goal of multi-task MARL is develop single model capable of solving multiple tasks, improving both learning efficiency and general- ization performance [3]. The core challenge lies in reusing knowledge across similar tasks while avoiding gradient conflicts between dissimilar tasks [15]. Existing approaches attempted to introduce modularity into policy learning to address this issue, employing techniques such as subtask reorganization [36], skill discov- ery (42, 22], and network decomposition [49].
Policies
Dynamics
Dioe--llnne-Cleeen
Dase-Qpen-twwed
O Eatch-Abrent
0 Eatch-Urderam
Figure 1: T-SNE visualization comparing the distribu- tions of dynamics and policies.
'This work was supported in part by National Natural Science Foundation China under Grants 621 160060 62293541, in part by Beijing Nova Program under Grant 20240484514, 4,, and in part by Beijing Natural Science Foundation under Grant No 4232056, I.223020.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

However, we argue that existing policy-centric methods may fail to fully capture the nuanced relationship among tasks. As visualized in Figure I, the policy distributions of the selected four tasks are mutually distinct. In contrast, we find that dynamies similarity exists within certain task groups (e.g., door-open and door-closs catch-abreast and satch-underarm) but not across different groups (e.g., between door and catch tasks). We refer to this phenomenon as bounded similarity, where tasks exhibit strong local similarity in dynamics within specific groups, but such similarity does not generalize globally across the entire task set.
This finding motivates two core insights. First, in terms of "similarity". modeling shared environ- ment dynamics -rather rather than learning separate task task-specific policies- can better exploit similarity and thus improve learning efficiency. Second, in terms of "bounded", modular structure such as Mixture-of ofree-oo-BBee of--xxerts (MoE) [30] is needed to isolate dissimilar dynamics and prevent negative interference across tasks. Based on these insights, we reconsider the role of modularity multi-task learning, and propose to apply it in the world model instead of the policy.
can
Multi- aguet Tasks
MoE-based World
Planner
Multi- Task Performance
catch
dynamics
93
54:1
Intenn
reward
horizon
12.5 24.5
26.1
50
MARL
13.7
open
dosee
critic
actor
X
actionns
MSW
MACO
HMADD
MAT
MAPTO
HATRPO
Figure 2: Overview. (Left) The framework of M3W. (Right) Performance comparisons.
A world model captures the feature of the environment, enabling the generation of imagined trajecto- ries that can be used for model training to improve sample efficiency [10-12], and for planning to enhance performance [31, 1, 14, 13]. As illustrated in Figure 2, we incorporate a MoE structure into the world model. Specifically, we construct multi-agent dynamics model based on SoftMoE [28] and a reward predictor based on SparseMoE [30]. During execution, the framework bypasses an explicit policy model and directly evaluates and optimizes actions using predicted data generated by the world model. We name this approach as M3W, which employs a MoE-based World model for learning and planning Multi-agent Multi-task scenarios. To the best of our knowledge, this the first work to apply MoE in multi-task world models. The contributions of this paper are as follows:
1) We identify the phenomenon of bounded similarity in multi-task dynamics, and reconsider the role of modularity by applying it to world models rather than policies.
2) We propose novel framework that learns dynamics and rewards via MoE-based world model and performs planning directly on predicted rollouts, without relying on an explicit policy.
3) We evaluate M3W on Bi-DexHands and MA-Mujoco, demonstrating its superior performance, sample efficiency, and interpretability through extensive comparisons, ablation studies, and visualizations.
2 Related Work
Multi-task MARL. The key challenge in MT-MARL how to reuse knowledge across similar tasks while avoiding gradient conflicts among dissimilar tasks [3]. Most existing methods address this issue from the perspective of policy learning [46, 18], exploring solutions such as subtask reorga- nization, skill discovery, and network decomposition. For example, UNMAS [2], TETQmix [49] design network architectures to accommodate the compositional variations of multi-agent system. ODIS [42]. HMASD [40] learn coordination skills, enhancing generalization capabilities through skill combinations: DT2GS [36] employs a hierarchical structure to learn the decomposition and composition of subtasks in the multi-task set. As mentioned in Introduction, these policy-centric methods do not fully exploit the bounded similarity among tasks, and a model-based paradigm offers more promising direction.
Model-based MARL. In multi-agent settings, world models often adopt communication-barn paradigms address unique challenges such partial observability and locally non-stationary

---

dynamics [47]. For instance, AORPO[44] assumes agents can communicate to access opponents' actions, using this information to improve model learning. Building on this idea, MAMBA[7] incorporates Recurrett State-Space Models (RSSM) from DreamerV2[11 1] to construct multi- world model. MACD[1] further introduces an explicit Transformer- based communication module, significantly enhancing model prediction accuracy. MA-TDMPC [14] and MAZero [21] extend TD-MPC [14, 13] and MuZero [31] to the multi-agent context.
MoE in MTRL. As model capacity increases and generalization becomes crucial, Mixture of Experts (MoE) has become a key approach in multi-task reinforcement learning (MTRL) for balancing shared and task-specific knowledge. Early works such as CARE [32] introduced context- based representations differentiate tasks within shared policy, while Paco [35] composed shared and task-specific parameters to enhance transfer. AMESAC [5] further incorporated attention based expert selection to address task heterogeneity, and MOORE [16] imposed orthogonal constraints to encourage expert decoupling. Building on these, recent methods such as MCP [27], D2R [15], and M3 [24] employ MoE to enable hierarchical skill reuse, dynamic routing, and transferable policy learning. Unlike prior works that focus MoE-based policy design, our approach integrates MoE into the world model, achieving modular learning of both dynamics and reward functions.
3 Problem Formulation
cooperative MARL task can be modeled as a decentralized partially observable Markov decision process (dec-POMDP), defined by the tuple M = A, P. R),, where I = (I....,7 n) the set of agents, S is the global state space, O I1=,, O' the joint observation space, the joint action space, Ω S x I - O the observation function, P x A - S is the transition function, R: S x A + R is the reward function, and y is the discount factor [34]. In single-task MARL, the goal is to learn a joint policy π II. n² that maximizes the expected discounted return Σr= r,, In contrast, MT-MARL aims to train a single model that performs well across a finite set of tasks T = {M&}, where each task M₁ is itself a dec-POMDP with task- specifce dynamics and reward functions.
A :==1e - A
Multi-task Settings. Following the general setting of multi-task MARL [43], we: (1) use one-hot encoding for each task as task label, (2) align dimensions of S₂ and O₂ across all tasks by padding zeros at end. (3) apply action space masks for A and (4) normalize or bound reward scales across tasks to facilitate stable learning.
Methodologies
As shown in Figure M3W consists of two parts: MoE-based world model and multi-agent planner. During training, the world model is trained using real environment data to improve the quality of predicted rollouts. During execution, the planner evaluates and optimizes actions based on predictions from the world model.
4.1 Learning a World Model for MT-MARL
Learning Multi-Agent Dynamics as Sequence Prediction. At each timestep, the observation action pairs (of: of) from multiple agents can be arranged into sequence [38], allowing the multi-agent dynamies to be modeled as sequential prediction problem [1, 45). The dynamics takes as input a sequence of length [[o] af), (of aa),................... (of a.)]. and outputs sequence of the same length, 1 1. This formulation treats multi-agent dynamics prediction as a structured sequence modeling problem.
Multi-agent World Model. Establishing dynamics and reward prediction within shared latent space has been shown to effectively enhance the predictive capacity of world models [14, 10]. Following this principle, our world model operates entirely in the latent space and does not include decoder. comprises an encoder, dynamics model, reward predictor, and an actor-critic module. Consistent with existing model-based MARL approaches, we assume that agents can access the observations and actions of others, which can be achieved through single round of communication in practice. To support multi-task generalization, we introduce a learnable task embedding e for each task, which is used as conditioning signal throughout the model. Let 2f denote the latent state of agent at time
3

---

t, then the world model are defined as:
Encoder:
2f PE (0
Dynamics:
2
QD Zt, Oz,
Predictor:
Ft
QR Zt, dt E)
Critic:
Qe = Q Zt, at, e)
Actor:
&z = Actor [2),)).
where 2t ........
27] at [a...... ... J. Although Eq.(1) includes an actor, does not serve as an explicit decision-making policy. Instead, it is used solely to generate candidate actions for the planner.
Learning Objectives. Let θg, lp. On th and denote the parameters of the encoder, dynamics model, reward predictor, critic, and actor, respectively. The world model are optimized under a self-supervised framework [10, 13, 1]. and the loss function is defined as:
((0,0)
ΣΣ
Σxx
| + SoftCE (fr,n)+ SoftCE (2r G.)
Dynamic Loss
where = (8g, pD. 8R is the prediction horizon, E (0, 1] is a constant coefficient that balances the contribution of each rollout step, G, the TD target, #1+1 = sg (pE(o(+1) is the encoded next latent state.
Following Eq. (2), the encoder learns transformation from real observations latent representations. The dynamics model and predictor are trained to minimize the prediction loss in latent space, ensuring aligmment between predicted trajectories and real environment transitions.
4.2 Incorporate MoE in Multi-Agent World Model
According to the modeling in Section 4, the dynamics learning formulated as sequence-to- sequence prediction problem, while reward prediction is treated sequence-to-scalle regression problem. To align with these formulations, we implement the dynamics model using SoftMoE [28] and the reward predictor using SparseMoE [30].
Soft Router
Weights
Sparse
Top
Router
dispatching
aggregating
activate
MLP Expert
feel
Salf Expert
MLP Expert
GZ
X2
Self-Artn Expert
multi-agent tokens
Expert
prediction tokens
Input tokens
output tokens
Expert
Figure 3: MoE in world model. The router generates of weights based on the task hand, which is used activating and allocating experts. (Left) SoftMoE-basee multi-agent dynamics model. (Right) SparseMoE- based reward predictor.
SoftMoE for Multi-Agent Dynamics Learning. In sequence [[2]] a!)...., (x?...,,,,,")))))) E (dx(d2+dd)) each pair (zf, a;) can treated as a token. The dynamics model, shown on the left side of Figure 3, is implemented as variant of SoftMoE [28], comprising soft router and set of experts. The soft router computes routing weight matrix as:
W SoftRouter(z:, ae, e) B
which plays two roles: (1) dispatching multi-agent tokens (ei, at) to N input tokens hū for the experts; (2) aggregating the N experts' output tokens hfp into n predicted latent tokens 2441 Here, i l,...,n and & I,
For the first role, at the input stage, we apply sof tmax normalization along the rows of W and compute N input tokens as weighted combinations of the n multi-agent tokens:
exp(WW#) h E, exp( (2) @ at) Rddz+d, k = I,..,,N,
n
(4)

---

where @ denotes the concatenation operation. Each h serves as the input to the k-th expert, implemented as an independent MLP to avoid gradient interference. The experts produce N output tokens h = Expert" (h,) E _d=
For the second role, at the output stage, we apply sof tmax normalization along the columns of W to combine N outputs of the experts and yield the final prediction [ 2141
N
(x)
h E Rd i i= 1,..., E, exp(W
(5)
Compared existing works that directly use MLPs or Transformers for multi-agent dynamics learning [7, 1, 45], our proposed dynamics model offers better adaptability to multi- mulli--ask context. The set of experts can be viewed
knowledge library. while the router can invoke and allocate experts based on the task at hand, effectively leveraging the property of bounded similarity. For similarity aspect, the router assigns similar experts to similar tasks, enhancing knowledge sharing and learning efficiency; for bounded aspect, it assigns different experts to dissimilar tasks, isolating gradients and avoiding conflicts.
MparseMoE for Reward Prediction. Unlike multi-agent dynamics, the reward is shared among all agents, making reward prediction sequence-to-scalar regression problem. As shown on the right side of Figure 3, the reward predictor is implemented as a SparseMoE, consisting of sparse router and set of self-attention experts.
The sparse router generates a length-N vector, representing the relevance scores of N experts for the current task. During training, this vector is transformed into probability distribution via sof tmax, and K experts are selected via roulette sampling. During execution, the top- K experts with the highest scores are directly activated. This allows the sparse router invoke the most suitable experts based on the task at hand. Each expert in the predictor consists of a multi-head attention layer followed by an average pooling layer. The attention mechanism captures inter-agent interactions, while the pooling layer produces a unified reward prediction shared by all agents.
Load-Baland Loss. To mitigate the load imbalance issue in SparseMoE, some experts are overloaded with tasks while others remain inactive, we incorporate a load- balancinn loss in [8].
N
=r(On)= = E,,,, J3 K j=1
Σ MET
1(M selects Expert j), PJ = T
Σ MET
sj(M), (6)
where fú denotes the selection frequency of expert j. P, represents its average task-relevance score, and 8y (M) is the score of expert j for task M.
4.3 Learning and Planning Multi-Agent Tasks
We then integrate the MoE-based world model with an actor-critic module to enable learning and planning multi-agent tasks. key advantage of M3W's planning is that it leverages the predictive capability of the world model directly, without relying explicitly on decision-making policy, thereby avoiding the limitations of policy-centric methods. Action planning in M3W is realized using a multi-agent extension of the model predictive path integral (MPPI) algorithm [39]. At timestep t, given an action sequence at.++H. H. the world model performs rollouts in the latent space generate predicted rewards ft-++ and latent terminal state 2 и By combining these predictions with the critic's estimate of the Q-value at the terminal state, the expected return of each action sequence can be evaluated as:
H-1
Vaneen = HQ (2042(H at, e)
Σ
r PR (Zeth, or+h,e),
where the first term represents the long- terg-temm return, estimated by the critic, and the second term represents the short-term cumulative rewards predicted by the world model. Then the policy of planner can be expressed as Λ Planner arg max Vauta
As shown in Eq. (1), the world model of M3W includes an actor module, The actor does not serve as the explicit decision Decisinn-making making policy, but rather acts as candidate action generator for the planner.
5

---

We adopt the Heterogeneous Multi-Agent Soft Actor-Critic (HASAC) algorithm(20) to optimize the actor, with the loss function defined as:
72 [()) - = X (a7t ( Actor (-1=f,e)) - Q (zz, Ce; e))
(8)
where a the entropy coefficient, H() denotes the entropy function. For the second term, Ze
[=.....,,"" =f'l and a = pæ (o), e) is the encoded initial latent state, zf = qD (Et, at at,e) e)" is the [a)...., a?] and aj Actor e))
predicted states from dynamics model,
Details of the multi-agent planner, pseudocode for learning and planning, and implementation specifics are provided in Appendix A, B, and C, respectively. Further discussion and ablation studies on the role the actor module appears in Appendix E.4.
Experiments
In this section, evaluate the following properties of our method: (a) performance in multi-task learning. particularly in comparison with state-of-the-ar policy-centrie methods; (b) effectiveness of the MoE-based world model, assessing whether the MoE improves prediction accuracy and planning quality; (c) interpretability, including M3W's exploitation of bounded similarity, expert specialization, and the diversity of multi-task behaviors.
Environments. We evaluate M3W and the baselines on on two challenging benchmarks: Bimanual Dexterous Hands (Bi- (Bi-DexHands)44 with 14 tasks, and the multi-agent Mujoco (MA-Mujoco)]26 with 24 tasks. Both benchmarks include a mix similar and diverse tasks, providing a solid testhed for multi-task performance. In Bi-DexHands, two agents control pair of dexterous hands to perform complex manipulations, with high high-dimensional observation and action spaces (up to O E IR 229 .A IR 26). In MA-MuJoCo, agents coordinate by controlling different joints of a MuJoCo robot to accomplish locomotion and control tasks. More detailed descriptions, visualizations and comparisons with the other environments be found in Appendix D.
5.1 Comparisons
Baselines. For the baselines, MACD [1] introduces a transformer- based world model, serving as a baseline for model-based methods. HMASD [40] applies modularity to hierarchical skill discovery and serves as baseline for multi-task methods. MAT [38], MAPPO [41], and HATRPO [19] are among the most popular general MARL algorithms currently available. Additionally, trained separate MAPPO model for each task as a strong baseline.
MSW touns
MACD
HMASD
MAT
MAPPO
HATRPO ... Single MAPPO
100
Catch Abreast
100
Catch Underamm
100
Door Open Outward
Door Closs Inward
Kowrage Bakum
Pecroge Retum
Avcrage Retum
60
50
8
9
20
8
%
10M
20M
S0M
40M
S0M
13M4
20M
30M
40M
50M
10M
20M
S0M
40M
S0M
10M
20M
30M
40M 50M
Training Stops
Trahiing Stops
Training Steps
Trakting Stops
Figure 4: Performance comparisons for selected tasks in Bi- Bi-DexHands. The black dashed line represents the strong baseline under single-task MAPPO, while the remaining bold curves indicate the average returns, with the shaded areas representing the 95% confidence intervals.
Results Bi-DexHands. Figure 4 shows the performance curves of M3W and baselines on the selected tasks in Bi- Di-DexHands, while Table reports results across all tasks. Overall, M3W achieves superior performance, better sample efficiency, and greater multi- task adaptability. After 50M steps, M3W consistently outperforms the baselines on average and closely matches the strong baseline on most individual tasks.
For sample efficiency, consider several complex yet similar tasks, such as catch- aatchaabreast, and catch-underarm. These tasks all involve object throwing and catching between two hands, sharing

---

Table I: Performance comparisons for various tasks in Bi-DexHands, The values following the + indicate the 95% error margins. The maximum values highlighted in bold with blue shading. The rewards normalized the range [0. 100]. the original rewards and normalization methods can found in Appendix E.I.
Task
(our)))))
MACD [I]
HMASD [40]
MAT [38]
MAPPO [41]
HATRPO [19]
over
70.0 1.3
64.4 0,9
62.3 1 0.8
$4.6 1
63.1 1:
42.3 1: 4.3
bottle-cap
89.8 0.6
85.7 1.1
90.1 + 0.6
89.6 0.4 0.4
90.9 +
88.2 0.7
catch-abreast
64.1 4.4
44.5 =.1
19.1
10.0
13.0 3.8
cutch-over2underam
84.8 1.4
79.9 2.0
7
42.9
69.7 + 2.1
70.3 4.6
catch- sich-underarm
87.3 1.3
1.4
78.8 2.2
1.5
52.1 1 1.3
29.2 1.21:44
block-stack
88.1 1:
72.5 = 2.4
60.5 1 7.3
62.5 1 1.3
59.5 1 3.3
62.5
door-close-inwart
91.3 0.7
54.0 6.0
73.5+
56.4
59.3 0.9
58.4 0.8
door-close-outtee outward
95.0 1.2
1.0
85.6
8.0
74.8 + 3.1
86.5 0.5
door-open-inward
91.5 0.4
92.7 = 0.8
+.3
83.7
8
80,6 1.7
doot-open-outward
89.0 2.1
77.5 = 2.5
68.9 1 3.1
6.66
65.8 1 1.6
69.4 2.2
kettle
90.4 0,4
0.9
89.9 + 1.1
8788
84.1 1.0
86.2 1.3
lift-underarm
77.0 2.1
60.8 = 3.3
+ $7.7
2.7
40.5
46.6 + 3.2
pen
83.3 = 3.3
69.6 3.5
68.6 2.7
29.7
38.5
29.9 + 2.9
scissors
76.4 1 2.3
76.7 = 4.3
86.0 + 6.8
+1777477
80.1 + 3.3
75.2 + 2.1
average
84.1 0.8
72.0 1.2 72.3 1: 1.4 58.8 1.8616
63.6 1.4
59.4 1: 1.6
the same observation/action spaces well as similar dynamics. As shown in Figure 4, M3W achieves significantly faster convergence and even surpasses the strong single-task baseline the latter task. This demonstrates that the MoE-basee world model effectively exploits task similarity, enabling rapid acquisition of similar dynamics and improved sample efficiency.
For adaptability in multi-task learning, taking door- outward and poor-clos--inware as examples, they share the same observation/action spaces and dynamics, yet have completely disparate initial states and reward functions. As shown in the last two plots of Figure 4, most baselines struggle to perform well on both tasks concurrentlye..g. HMASD performs slightly better
-owa--ieearrddd but still learning slowly door-open-outward. In contrast, M3W successfully learns both opposing tasks simultaneously. This flexibility attributed to the modular structure of the world model, which effectively mitigates semantic conflicts between tasks.
Results on MA-Mujoco. The performance comparison of the selected tasks in MA-Mujoot can be found in Figure 5, and the complete results in MA-Mujoco can be found in Appendix E.1. The tasks in MA-Mujoco involves multiple robotic suites, each corresponding to multiple tasks. Within the same suite, the tasks share dynamics while differing in their reward functions. While across different suites, the dynamies are also distinct, which is challenge for the model effectively knowledge reuse.
MAPPO
HATRPO
Single MAPPO
CwwcChhhhthhhhch Cha.hhhatt-nBaackkkkkkk RunBack
2ag DaggCChstah-Rarae RunFront
Hopper-Hop
Hoppar-Stand
(betm B. '88
Res0go Ritu
Rdum 8
besoge PanooeePa 88
20
10M
Stepe
Figure S: Performance comparisons for selected tasks in MA-Mujoco.
Figure 5 (left) presents results on two tasks from the cheetah suite, including run run-front and run-back, which encourage the robot tasks share the same dynamics, and M3W achieves higher sample efficiency, indicating that the world model effectively enables knowledge sharing across similar tasks. Figure 5 (right) shows the tasks hopper-hop and hopper-stand, which involve robots that are entirely different from the cheetah. In these tasks, M3W also demonstrates superior performance. Notably, for simple tasks like hopper- stand, only M3W maintains consistent performance post-convergence, whereas other baselines suffer degradation due to interference from other tasks. This further confirms the advantages of M3W addressing multi-task knowledge conflicts.
using only its front-leg back-leg, respectively. These
For statistical analysis of the comparison results, please refer to Appendix E.2.
7

---

5.2 Effectiveness of MoE-based World Model
We then conduct experiments to evaluate the accuracy of the MoE-based world model and its contribution to action planning, providing both direct and indirect evidence of its effectiveness.
Accuracy of the World Model Prediction. We selected six tasks misks-(Cheetail, run, run-bwd, run run-front; Walker2D: run. run-bwd, walk) iinkiiinvolving two Mujoco robots whose dynamics exhibit a clear pattern of bounded similarity. We replace MoE with other alternative architectures, such as Transformer, Switch- TranhhTTranffommer and MLP, then retrain the models, and evaluate their prediction accuracy. Results are reported in Figure Compared other structures, M3W achieves the lowest prediction error and highest stability, along with strong long-range prediction capability. More detailed experimental settings and models' parameter counts can be found Appendix E.3.
For dynamics prediction, MLP performs the worst, because the local dynamics is non-stationary and difficult fit with shared model. Transformer [1, 45] uses attention mechanism to model interactions between agents, learning dynamics from global perspective. While it shows some improvement over MLP, single set of shared parameters still limits adaptability across multiple tasks. Corich-Transformen [8] replaces the ferwarr network (FFN) with MoE, but due to the lack of proper dispatching and aggregation mechanisms, each expert operates only on local observations and actions. As result, it also struggles to capture diverse dynamics. For reward prediction, MLP again exhibits poor performance. Self-Artentinn also suffers from limited generalization due to parameter sharing. Only our method M3W achieves the highest accuracy and stability across tasks.
Dynamics Prediction
Reward Prediction
0.8
M3Wlours)
Nithours saliantion
Swetch-ianstemmr
Transfomer
0.6
0.8
0.5
0.6
0.1
0.0
0.0
ours
Altri
MLP
Horizon
Modells
Figure 6: Accuracy of the world model prediction. (Left) Prediction dynamics. where the height of each indicates the average error, and the horizontal lines represent the 95% confidence intervals. (Right) Prediction errors on rewards.
Contribution to Model-basee Planner. The performance of the planner serves an indirect indicator of the effec- tiveness of the world model. As discussed in Section 4.3, the planner does not involve any learnable parameters its performance relies entirely on the quality of predicted tra- jectories. The more accurate the predictions, the greater the performance boost from planning. To further exam- ine this, equip MACD with the same planner, and the results are presented in Figure 7. We observe that adding the planner does lead to performance improvements for MACD, but significant gap remains between it and M3W. This further supports, from an indirect perspective, the ac- curacy and effectiveness of our MoE-based world model in generating reliable multi-task predictions.
Contribution to Planner
30
Average Reaum
125
MEW
MACO
10
MACO
Traiinn Steps
Figure 7: Performance impact of adding a planner to MACD.
5.3 Interpretability and Visualization
In the following, conduct a series of visualization experiments to illustrate the interpretability of M3W.
Expert Specialization. We visualize the attention distributions of individual experts during the Walker-run task in Figure The results reveal that the highest-scoring expert consistently focuses on the front leg in contact with the ground. It is a reasonable pattern, as these joints are actively

---

t= 2 t=3 ==5
t 7
1
111
Figure 8: Visualization of the attention outputs of experts in the reward predictor on robot joints, where deeper colors indicate higher weights. (Top) Highest-scorins expert: (Bottom) Lowest-scoring expert.
exerting force and play a critical role in determining the robot's movement. In contrast, the lowest- scoring expert exhibits an implausible attention pattern, maintaining focus on the same rear leg across multiple time steps (e.g., at / 5, 6, 7, 8), without adapting to the movement, This contrast highlights that the experts M3W have achieved functional specialization, and that the router has successfully learned to assign the appropriate expert to task at hand.
no tr-1oe
Alount
Opes OutarE
Dae
Clar Octnn
Anbo bed
5dien
vppe--ozel coed
Buttir Cat
O
Figure 9: (Left) Visualization of cosine similarity of router outputs. (Right) t-SNE trajectories of the joint latent states.
Exploiting Bounded Similarity. As discussed in Section 4.2, the experts set can viewed as a knowledge library, with the router invoking and allocating them based on the task at hand. To better understand this mechanism, we visualized the pairwise cosine similarity of router outputs across tasks, shown the left of Figure 9. The resulting pattern exhibits clear diagonal blocks along the main diagonal, which correspond to tasks with similar dynamics. This indicates that the SoftMoE- based dynamics model successfully leverages the bounded similarity among tasks, achieving expert reuse and effective generalization across related tasks.
Multi- Task Behavior Diversity. The right side of Figure 9 visualizes multi-task trajectories in the shared latent space. The results show clearly separated regions for different tasks, indicating that the model captures diverse behaviors across tasks. Furthermore, tasks with similar dynamics tend to cluster together- instance, catch- astchaabreast, inteaii//w///dwwrrdddddd positioned closely, the close/ caeae//ppe-iiee iovennniward///eo outward tasks. This observation aligns with M3W's design motivation and supports its effectiveness multi-task learning.
over2underarm. and catch- -unhhhu
5.4 Other Experimental Analysis
To further examine the characteristics of M3W, we present the following additional analyses:
1) The impact of the actor prior in Appendix E.4;
2) The robustness of M3W to variations in reward scale in Appendix E.5;
3) The inference time cost analysis in Appendix E.6;
4) The scalability evaluation with respect the number of agents Appendix E.7;

---

Conclusion
In this paper, we propose M3W, which leverages a MoE-based world model for learning and planning multi-agent tasks. Motivated by the finding of bounded similarity, we reconsider the role of modularity multi-task learning and propose to integrate it into the world model rather than the policy network. Specifically, a SoftMoE-based dynamics model and SparseMoE- reward predictor are used to isolate dissimilar task dynamics while promoting knowledge reuse among similar tasks. By directly performing action planning synthetic trajectories generated by the world model, M3W overcomes the limitations of policy-centric methods. Experimental results on Bi-DexHands and MA-Mujoco demonstrate that M3W significantly improves both performance and sample efficiency. Comprehensive ablation studies and visualizations further confirm the critical role of the MoE-based architecture in capturing bounded similarity and enabling effective multi-task planning.
Currently, M3W is only applicable tasks with continuous action spaces, due to the limitations of the MPPI-based planner. Future work could integrate M3W with Monte Carlo Tree Search (MCTS) or Cross Cross-Entroppy Method (CEM)- based planner accommodate tasks in discrete action spaces.
10

---

References
[1] Jiajun Chai, Yuqian Fu, Dongbin Zhao, and Yuanheng Zhu. Aligning credit for multi-agent cooperation via model-based counterfactual imagination. In AAMAS, pages 281-289, 2024.
[2] Jiajun Chai, Weifan Li, Yuanheng Zhu, Dongbin Zhao, Zhe Ma, Kewu Sun, and Jishiyu Ding. Unmas: Multiagent reinforcement learning for unshaped cooperative scenarios. IEEE Transactions Neural Networks and Learning Systems, 34(4):2093-2104, 2023.
[3] Jiajun Chai, Zijie Zhao, Yuanheng Zhu, and Dongbin Zhao. A survey of cooperative multi aulti--geett reinforcement learning for multi-task scenarios. Artificial Intelligence Science and Engineering. 1(2)98-121, 2025.
[4] Yuanpei Chen, Tianhao Wu, Shengjie Wang, Xidong Feng, Jiechuan Jiang, Zongqing Lu, Stephen McAleer, Hao Dong, Song-Chun Zhu, and Yaodong Yang. Towards human-level bi- manual dexterous manipulation with reinforcement learning. In Advances in Neural Information Processing Systems, volume 35, pages 5150-51633. Curran Associates, Inc., 2022.
[5] Guangran Cheng, Lu Dong, Wenzhe Cai, and Changyin Sun. Multi-task reinforcement learning with attention-based mixture experts. IEEE Robotics and Automation Letters, 8(6):8812 3819, 2023.
[6] Junha Chun, Youngjoon Jeong, and Taesup Kim. Sparse imagination for efficient visual world model planning. arXiv preprint arXiv:2506.0199, 2025.
[7] Vladimir Egorov and Alexei Shpilman. Scalable multi-agent model-based reinforcement learning. In Proceedings of the 21st International Conference on Autonomous Agents and Multiagent Systems, AAMAS '22, page 381-390, Richland, SC, 2022. International Foundation for Autonomous Agents and Multiagent Systems.
[8] William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research, 23(120):1-39, 2022.
[9] Arthur Gretton, Karsten M Borgwardt, Malte J Rasch, Bernhard Schölkopf, and Alexander Smola. A kernel two-sample test. The journal of machine learning research, 13(1):723-773 2012.
[10] Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Leamiing behaviors by latent imagination. In International Conference on Learning Represenia- tions, 2020.
[11] Danijar Hafner, Timothy P Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. In International Conference on Learning Representations, 2021.
[12] Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse control tasks through world models. Nature, 640(8059):647-653 2025.
[13] Nicklas Hansen, Hao Su, and Xiaolong Wang. TD-MPC2: Scalable, robust world models for continuous control. In The Twelfth International Conference on Learning Representations, 2024.
[14] Nicklas A Hansen, Hao Su, and Xiaolong Wang. Temporal difference learning for model predictive control. In Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 8387-8406. PMLR, 17-23 Jul 2022.
[15] Jinmin He, Kai Li. Yifan Zang, Haobo Fu, Qiang Fu, Junliang Xing, and Jian Cheng. Not all tasks are equally difficult: Multi-task deep reinforcement learning with dynamic depth routing. Proceedings of the AAAI Conference on Artificial Intelligence, 38(1 2812376-12388 Mar. 2024.
[16] Ahmed Hendawy, Jan Peters, and Carlo D'Eramo. Multi-task reinforcement learning with mixture of orthogonal experts. In The Twelfth International Conference on Learning Represen- tations, 2024.

---

[17] Hanchi Huang. Deheng Ye, Li Shen, and Wei Liu. Curriculum based asymmetric multi-task reinforcement learning. IEEE transactions on pattern analysis and machine intelligence. 45(6):7258-772 2022.
[18] Shangjing Huang, Zijie Zhao, Yuanheng Zhu, and Dongbin Zhao. Adaptive multi aggett coordination among different team attribute tasks via contextual meta--riinforemment learning. In 2024 IEEE 13th Data Driven Control and Learning Systems Conference (DDCLS), pages 1780-1787, 2024.
[19] Jakub Grudzien Kuba, Ruiqing Chen, Muning Wen, Ying Wen, Fanglei Sun, Jun Wang, and Yaodong Yang. Trust region policy optimisation in multi-agent reinforcement learning. In International Conference on Learning Representations, 2022.
[20] Jiarong Liu, Yifan Zhong, Siyi Hu, Haobo Fu, QIANG FU, Xiaojun Chang, and Yaodong Yang. Maximum entropy heterogeneous--ggen agent reinforcement learning. In The Twelfth International Conference on Learning Representations, 2024.
[21] Qihan Liu, Jianing Ye, Xiaoteng Ma, Jun Yang, Bin Liang, and Chongjie Zhang. Efficient multi-agent reinforcement learning by planning. In The Twelfth International Conference on Learning Representations, 2024.
[22] Xin Liu, Yaran Chen, Guixing Chen, Haoran Li, and Dongbin Zhao. Balancing state explo- ration and skill diversity in unsupervised skill discovery. IEEE Transactions on Cybernetics, 55(5):2234-222) 2025.
[23] Ryan Lowe, YI WU, Aviv Tamar, Jean Harb, OpenAl Pieter Abbeel, and Igor Mordatch. Multi-agent actor-critic for mixed cooperative-competiie environments. 30, 2017.
[24] Linghui Meng, Jingqing Ruan, Xuantang Xiong, Xiyun Li, Xi Zhang, Dengpeng Xing, and Bo Xu. M3: Modularization for multi-task and multi-agent offline pre-training. In AAMAS, pages (624-1633 2023.
[25] Georgios Papoudakis, Filippos Christianos, Lukas Schäfer, and Stefano V Albrecht. Bench- marking multi-agent deep reinforcement learning algorithms in cooperative tasks. In Thirty-fifth Conference on Neural Information Processing Systems Datasets and Benchmarks Track (Round 1), 2021.
[26] Bei Peng. Tabish Rashid, Christian Schroeder de Witt, Pierre-Alexandee Kamienny, Philip Torr, Wendelin Boehmer, and Shimon Whiteson. Facmac: Factored multi-agent centralised policy gradients. In Advances in Neural Information Processing Systems, volume 34, pages Curran Associates, Inc., 2021.
[27] Xue Bin Peng, Michael Chang, Grace Zhang, Pieter Abbeel, and Sergey Levine. Mcp: Learning composable hierarchical control with multiplicative compositional policies. 32, 2019.
[28] Joan Puigcerver, Carlos Riquelme Ruiz, Basil Mustafa, and Neil Houlsby. From sparse to soft mixtures of experts. In The Twelfth International Conference Learning Representations, 2024.
[29] Mikayel Samvelyan, Tabish Rashid, Christian Schroeder de Witt, Gregory Farquhar, Nantas Nardelli, Tim G. J. Rudner, Chia-Man Hung. Philip H. S. Torr, Jakob Foerster, and Shimon Whiteson. The starcraft multi-agent challenge. page 2186-2188, 2019.
[30] Noam Shazeer, * AAaalia Mirhoseini, *Krzysztof Maziarz, Andy Davis, Quoe Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gan -ated mixture-of- experts layer. In International Conference on Learning Representations, 2017.
[31] David Silver, Julian Schrittwieser, Karen Simonyan, Ioannis Antonoglou, Aja Huang, Arthur Guez, Thomas Hubert, Lucas Baker, Matthew Lai, Adrian Bolton, Yutian Chen, Timothy Lillicrap, Fan Hui, Laurent Sifre, George van den Driessche, Thore Graepel, and Demis Hassabis. Mastering the game of go without human knowledge. Nature, 55046666))666)33433999 2017.
12

---

[32] Shagun Sodhani, Amy Zhang, and Joelle Pineau. Multi-tase reinforcement learning with context-based representations. In Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 9767-9779. PMLR, 18-24 Jul 2021.
[33] Ryan Sullivan, Akarsh Kumar, Shengyi Huang, John Dickerson, and Joseph Suarez. Reward scale robustness for proximal policy optimization via dreamerv3 tricks. 36: 36:1335-1362, 2023.
[34] Jingbo Sun, Songjun Tu, Haoran Li, Xin Liu, Yaran Chen, Ke Chen, Dongbin Zhao, al. Unsupervised zero-shot reinforcement learning via dual-value forward-backward representation. In The Thirteenth International Conferencee on Learning Representations, 2025.
[35] Lingfeng Sun, Haichao Zhang, Wei Xu, and Masayoshi TOMIZUKA. Paco: Parameter-
compositional multi-task reinforcement learning. 35:21495-215577, 2022.
[36] Zikang Tian, Ruizhi Chen, Xing Hu, Ling Li, Rui Zhang, Fan Wu, Shaohui Peng. Jiaming Guo, Zidong Du, Qi Guo, and Yunji Chen. Decompose a task into generalizable subtasks in multi-agent reinforcement learning. In Advances in Neural Information Processing Systems, volume 36, pages 78514-78532. Curran Associates, Inc., 2023.
[37] Saran Tunyasuvunakool, Alistair Muldal, Yotam Doron, Siqi Liu, Steven Bohez, Josh Merel, Tom Erez, Timothy Lillicrap, Nicolas Heess, and Yuval Tassa. dmcontrol: Software and tasks for continuous control. Software Impacts, 6: 100022, 2020.
[38] Muning Wen, Jakub Kuba, Runji Lin, Weinan Zhang, Ying Wen, Jun Wang, and Yaodong Yang. Multi-agent reinforcement learning is sequence modeling problem. In Advances in Neural Information Processing Systems, volume 35, pages 16509-16521, Curran Associates, Inc., 2022.
[39] Grady Williams, Andrew Aldrich, and Evangelos Theodorou. Model predictive path integral control using covariance variable importance sampling. arXiv preprint arXiv: 149, 2015.
[40] Mingyu Yang. Yaodong Yang, Zhenbo Lu, Wengang Zhou, and Houqiang Li. Hicrarchical multi-agent skill discovery. 36:61759-61776, 2023.
[41] Chao Yu, Akash Velu, Eugene Vinitsky, Jiaxuan Gao, Yu Wang, Alexandre Bayen, and YI WU, The surprising effectiveness of ppo in cooperative multi-agent games. Advances in Neural Information Processing Systems, volume 35, pages 24611 14611-244 1-6224. Curran Associates, Inc., 2022.
[42] Fuxiang Zhang, Chengxing Jia, Yi-Chen Li, Lei Yuan, Yang Yu, and Zongzhang Zhang. Discovering generalizable multi-agent coordination skills from multi-task offline data. In The Eleventh International Conference on Learning Representations, 2023.
[43] Junkai Zhang, Jinmin Hc, Yifan Zhang, Yifan Zang, Ning Xu, and Jian Cheng. Bi-level knowledge transfer for multi-task multi-agent reinforcement learning. In The Thirty-ninth Annual Conference on Neural Information Processing Systems, 2025.
[44] Weinan Zhang, Xihuai Wang, Jian Shen, and Ming Zhou. Model-based multi-agent policy optimization with adaptive opponent-wise rollouts. In Proceedings of the Thirtieth International Joint Conference Artificial Intelligence, IJCAI-21, 2021.
[45] Yang Zhang, Chenjia Bai, Bin Zhao, Junchi Yan, Xiu Li, and Xuelong Li. Decentralized transformers with centralized aggregation are sample-efficient multi-agent world models. CoRR, abs/2406. 15836, 2024.
[46] Zijie Zhao, Yuqian Fu, Jiajun Chai, Yuanheng Zhu, and Donghin Zhao. Meta learning task representation in multiagent reinforcement learning: From global inference to local inference. IEEE Transactions on Neural Networks and Learning Systems, 4908):1400811421, 2025.
[47] Zijie Zhao, Honglei Guo, Shengqian Chen, Kaixuan Xu, Bo Jiang, Yuanheng Zhu, and Dongbin Zhao. Empowering multi-robot cooperation via sequential world models. arXiv preprint arXiv:2509.1309: 2025.
13

---

[48] Lianmin Zheng, Jiacheng Yang, Han Cai, Ming Zhou, Weinan Zhang, Jun Wang. and Yong Yu. Magent: A many-agent reinforcement learning platform for artificial collective intellignnce In Proceedings of the AAAI conference on artificial intelligence, 2018.
[49] Yuanheng Zhu, Shangjing Huang, Binbin Zuo, Dongbin Zhao, and Changyin Sun. Multi-task multi-agent reinforcement learning with task-entity transformers and value decomposition training. IEEE Transactions on Automation Science and Engineering, 22:9164-999 9177, Nov. 2024.
14

---

Appendices
Multi-agent Planner
16
B Pseudocode
17
Implement Details
17
C.1 Discrete Regression of Reward and Value
17
C.2 Symlog Predictions
18
Percentile Sealing
18
C,4 SimNorm Normalization
18
C,5 Hyperparameters
18
Additional Description of Environments
18
D.1 Description of Bi-DexHands tasks
18
D.2 Description of MA-Mujoco tasks
19
D.3 Visualization of Tasks
19
D.4 Why Choose Bi-DexHands and MA-Mujoco?
19
D.5
Bounded Similarity in Other Environments
21
Additional Experimental Results
22
E.1
Additional Comparison Results
22
E.2
Statistical Analysis of Comparison Results
................................
22
E.3 Details of Ablation Settings...
................................
E,4
Additional Ablation Studies on Actor Prior
24
E.5
Robustness to Reward Scale
25
E.6 Time Cost Analysis with an Early-Stopping Planner
26
E,7
Scalability with the Number of Agents
27
15

---

A
Multi-agent Planner
We provide detailed implementation of the multi-agent planner mentioned in Section 4.3. The planning process is illustrated in Figure 10.
World Model Rollout
Initial Encode
Evaluation
V Oaae en. Za-0 y Caam ()
enc
Elite Selection
Actor
elite
Prior
Action Distribution Update
Action Sampling
Figure 10: The process model-based planning.
Action Sampling. At each decision timestep t, a diagonal Gaussian distribution is used rep- resent the action distribution for each agent, and Np action sequences are sampled, d ~ N (plstn (o)++H)))) " where P E IR AIx Jhrt BlAAx are the mean and standard deviation the action distribution for agent
We incorporate N₂ < Np action sequences with policy priors into the candidate set, which are obtained directly from the actor. This inclusion aids accelerating the convergence the planner during the early iterations, while ensuring the foundational performance of the model-based planner.
Actor Prior.
World Model Rollout. After action sampling, the world model of M3W performs Hl-step rollouts in the latent space, resulting in No predicted trajectories of length H1, one of which is denoted as T { 27 [o] ......., a.] }
Evaluation. Subsequently, these trajectories evaluated based the value function using H-step return. Specifically, the value of trajectory I'm is given by Eq. (7).
Elite Selection. Based on the estimated value estimated, the sampled trajectories are ranked, and the top M trajectories are selected as elite set, one which denoted I",
Action Distribution Update. The process of updating the action distribution can be formally expressed Eq. (9). in which a'm exp [ (Vr. maXme Vrr:.) denotes the weight of each
clite trajectory, and 7 is a temperature coefficient.
M Σ Em=1 am
Σ am (I (m
(9)
DM O'm
M
M
Om
Iteration.
The above process is iterated K, times to derive the final action distribution.
The detailed hyperparameters used in the model-based planner summarized in Table 2.
16

---

Table 2: The Notations and Values of hyperparameters in the planner.
Hyperparameters
Notations
Value
Hyperparameters
Notations
Value
rollout horizon
sampling actions
N
512
planning iterations
clites
M
64
temperature
0.5
actor prior samples
N₂
24
B
Pseudocode
We present the pseudocode for M3W training and planning, as shown in Algorithm and Algorithm 2, respectively.
Algorithm 1 Model Training
Input: Task T {M&} replay buffer B, task embeddings ex: parameterized networks Φ, t for actor and critic function, 0R, lp, OR for encoder, dynamics model and reward predictor for episode ....3..... do do Get real data ([o:]e=1m [a]]i=im: Tt: (li=mm) by interacting with the environment Add transition into buffer: U (loilimtoss [af]i=ton: Tty [0 10(+
for step / = 1, 2,
end for
for epoch =
do for task M₁ do
Sample trajectories from B
Update 0g. Or,0 s/ by minimizing Eq. (2) & Eq. (6)
Update á by minimizing Eq. (8).
end for end for end for
Algorithm 2 Model Planning
Input: learned parameters 0g,0D, 0p o, vi, hyperparameters H, Kp, T, Np, M, N initial dis- tribution;
for step 1, do
Get environment observation lojli=1:n and encode them to latent space: zt Pe(o,, θg) for iteration 2, 3. ..... K2 do
Sample Np actions a N (Mist n (ohae)))1
Sample N7 actions from actor n² T2/,, i,
Get
predicted
data
by world
model
rollouts,
Tm
[[+, . 27]
[a.,....,")) a²] (r)
Evaluate the trajectories by Eq. (7) and select M clite action sequences 1" = arg max Vr.,
Update action distribution following Eq. (9)
end for
end for
C Implement Details
C.1 Discrete Regression of Reward and Value
We discretize the reward predictions and value estimates, transforming the scalar regression into a discrete regression enhance robustness [12, 13, 33). Taking reward prediction as an example, we partition the range of rewards into B bins, and the scalar reward can be transformed into B-dimensional two-hot vector. The two-hot encoding process is defined as follows: for a given scalar value we locate the adjacent bins by < r bu+1. and then assign weights to the two corresponding
17

---

positions based on linear interpolation. Formally, the encoding is defined as:
1b&+1
-r|/|bx+1 if d = k
bul.
two-hot(r)d+
1bx r1/|bu+1 bul, if k
Σ 16; sr)
(10)
0,
else
During training, we used the soft cross-entropy loss to optimize the discretized reward and value estimation. For Eq. (2), its last two loss terms are replaced as follows:
(an(zr, az, e) rr) exp(yR(#1, at, e)a) x two-hot(ri)a
(11)
explqreZz az, c)a)
C.2 Symlog Predictions
To further enhance numerical stability, we apply the symmetric logarithmic (symlog) and exponential (symexp) transformations. These transformations are particularly effective in scenarios where target values exhibit high variance, such as in reward prediction value estimation. They preserve continuity and differentiability. enabling stable learning while maintaining meaningful relationships between values. These transformations are defined in [12] as:
symlog x))
sign(x) In([æ] 1), symexp(æ) sign(=)(exp((=)) 1).
(12)
C.3 Percentile Scaling
We adopt percentile-based scaling method [13]. This method dynamically adjusts the scaling of Q-values by calculating the difference between the 5th and 95th percentiles of the returns in each batch. The scaling factor is updated smoothly to maintain numerical stability.
C.4 SimNorm Normalization
We adopt SimNorm (Simplicial Normalization) [13] to construct the latent space. projects the latent state onto a simplex of fixed dimensionality using softmax operation. SimNorm is particularly effective in sigh--immnsional spaces, as promotes sparse representations, mitigates information loss degradation, and preserves the structural integrity of the learned representation. The SimNorm transformation is defined as:
exp(xestL/T)
SimNorm(z) = .gi.....991 99 = Σ' exp(#xs+L/T)
(13)
where the simplex dimension, the temperature parameter, N dz/ the number of simplex segments, and 2 is the i-th segment of the input vector.
C.5 Hyperparameters
The hyperparameters for the model-based planner have been provided in Appendix A. Here, present the hyperparameters regarding the the training and network structures in Table and Table
Table 3: The training hyperparameters.
Hyperparameters
Value
Hyperparameters
Value
Hyperparameters
Value
buffer size
le6
batch size
256
train interval
step balance
0.5
5c-4
encoder Ir
1.5c-4
n-step return
10
gamma
0.99
D Additional Description of Environments
D.1 Description of Bi-DexHands tasks
Bimanual Dexterous Hands (Bi-DexHands) [4] widely adapted cooperative multi-agent envi- ronment, in which two agents control two ShadowHands to perform dexterous manipulation tasks.
18

---

Table 4: The network configurations.
Hyperparameters
Value
Hyperparameters
Value
Hyperparameters
Value
task dim
96
latent dim
encoder size
[256]
Som dim
8
num experts
16
experts size
[512, 512]
predictor K
2
actor & critic size
[512, 512]
num bins
101
scale p q coef
0.01
dynamics coef
20
reward coef
0.1
0.1
entropy coef
0.01
We selected 14 representative tasks that comprise the Bi-DexHands task set, which are visualized on the top Figure 11. Bi-DexHands features a high-dimensional observation and action space, up to O E R229 AA E Ra6 respectively. Table 5 presents a list of tasks in Bi- BinDexHands, including the dimensions of the observation and action spaces, along with a brief description of each task.
Table 5: Task descriptions in BinDsxHands.
X do X d,
Task
Description
Two palm-up hands oriented oppositely transfer object; hand holds and throws the object while the other hand catches it (the dogrees of froedom at the hand's base are frozen).
211 x 20
Over
Two hands collaborate open the bottle cap: one band holds the bottle while the other hand removes the cap.
Cap
2. 229 x 26
Stack
Two hands each pick one block and stack them together.
Twa palm up hands oriented in the direction pass object; hand holds the object while the other hand catches it
Catch-Ahreast
Over2Unneraam
palmn up Two palm-up hands onientdd oppositlly transter an object, one hand holds the object while the base are frozen).
vertical
object
223 26
Catch-Underaam
catches it (the degrees freedom at the
Close-Tnwee Inward
Two hands close a door that opens outward by pulling it inward.
Close-Outward
opens inward by pushing
Open-Inward Inward
Iwo hands open a door by pulling it inward.
Open-Outward
hands open a door by pushing outward.
Two hands pouring water collaboratively. one hand lifts the kettle and pours water into bucket
2
Kettle
beld by the other hand.
Lin Underamm
Two hands grasp handle pol and lift to specilied position.
Pen
Two hands cooperate to open the cap from pen.
Scissors
lugeiher open pair scisoos.
D.2 Description of MA-Mujoco tasks
Multi-agent Mujoco (MA-Mujoco) [26] is an extended version of Mujoco designed for multi-agent cooperation. It divides Mujoco suite into multiple agents based on joints and parts, allowing agents to collaboratively control robots to complete tasks. In this context, we selected 24 tasks to form the MA-Mujoco task set, with some of these tasks being adapted and expanded from DMControl [37]. These tasks exhibit significant degree of shared dynamics, making them an ideal testing ground for evaluating knowledge transfer and reuse multi-task learning algorithms. The visualizations of
the tasks are presented on the bottom of Figure 11, and the descriptions of the tasks are provided in Table
D.3
Visualization of Tasks
D.4
Why Choose Bi-DexHands and MA-Mujoco?
From the perspective multi-task learning, we compare Bi-DexHands [4] and MA-Mujoco [26] with other multi-agent environments, including SMAC(StarCraft Multi-Agent Challenge) [29],
MPE(Multi-Agent Particle Environment) [23], and RWARE(Multi Robot Warehouse) [25], which visualized in Figure 12.
compare these environments from multiple aspects, as summarized in Table 7. evident that Bi- Di-DexHands and MA-Mujoco are indeed challenging multi-task multi-agent environments.
19

---

Table 6: Task descriptions in MA-Mujoco.
Suite
do da
Taask
Description
The front and back legs of the half cheetaah are treated as two agenis that collaborate to forward.
Tun
control the rabot
2x7x3
run-bwd
The robot runs backwards
jump
legs the ground jump while maintaining
stationary position.
run-hack
The robot runs forward using only the back logs.
Chostah
run-front
The robot runs forward using only the front legs.
Each joint of the half-cheetah is treated as an agent, with the other settings being the those Zag-cheetah-run.
sun-bwd
see Zag-cheetah-run bwd
6x4x1
jump
see cheetah jump
nun-hack
see Saag-cheetal--runbbee
run-font
sec 2ag-choctall-ruo--
The upper
humanoid
agents collaborate
run
to conirol the robol to run lorward.
Hamanoid
xx
stand
The humanoid stands and maintains balance.
walk
The humamoid walks forward.
standup
The humancid lying
controlled stand up.
The left and right legs of the walker2d are treated as two agents that collaborate to robot
Walker2D
2x7x3
sun-bwd
The walker2D runs backwards.
stand
The walker2D stands and maintains balance.
walk
The walker21y walks forward.
walk-bbad
The walker2D walks backwards.
The two joints of the swimmer treated agents that collaborate control the zubot swinn forward.
Swimmer
2x8x1
swim-bwd
The rohot backward.
The three joints of the hopper are treated three agents that collaborate to sobot forward.
hop
Hopper
3x4x1
stand
The Hopper stands and maintains balance.
The two joints of the reacher are treated as two agents that collaborate to extend the the target.
Reacher
2x3x1
reach
robot's 's arm
Table 7: Comparison of multi-task multi agentt environments.
Bi-DexHands
MA-Mujoco
SMAC
MPE
RWARE
Semantics
X
X
Diverse Dynamics Heterogeneous
X
X
X
X
1. Semantics. In Bi-DexHands, tasks exhibit varying semantic relationships, including similar, opposite, and unrelated tasks. For instance, catch-abreast and aarchhunderarm are semantically similar, close-inward and open-outware are semantically opposite, while tasks like pen and scissors are entirely unrelated to others. A similar situation is observed in MA-Mujoco, where tasks such as walker-run and walker-walk are quite similar, and walker--rubaackwardd is opposite to them, while feacher-reade is entirely unrelated to others. In contrast, SMAC and RWARE, the semantics of different tasks are fundamentally the same. For example, in SMAC, multi-task setting is primarily reflected in variations in agent types or numbers, but with the common semantic goal of "defeating the enemy to achieve victory." In RWARE, tasks differ in the layout of shelf positions, but share the common goal of "transporting goods designated locations."
2. Diverse Dynamics. In Bi- Bi-DexHare Hands, variations in manipulated objects and the freezing of certain degrees of freedom lead to changes in dynamics. For example, in the tasks, the degrees of freedom at the hand's 's base are frozen, resulting dynamics that differ from those of catch tasks. Similarly, in MA-Mujoco, different suites exhibit varying dynamics: walker-run and walker-walk share similar locomotion dynamics, but are unrelated to those in reacher. In contrast, the dynamics in other environments tend be quite simple, where almost all tasks share identical dynamics. This simplicity limits their utility in assessing models' ability generalize across dynamic variations.
3. Heterogeneous. Both Bi-DexHands and MA-Mujoco feature heterogeneous agents. In Bi-DexHands, two agents control the left and right hands separately, each with its own set
20

---

Over
Bottle
Catch Abrost
Catch Underam
Catch OverUnnderamm
Block Stack
Close-Inward
Close-Outward
Open-Imeard
Open-Cuttwar
Kettle
Lift Underamm
Pen
Swissors
2ag- Seag-Cheetah-Ree Run
Zag- Seee-Chheeeheteeeee RunBwd
2ag-Cheetah- Zag-Cheetah- Rhecehettaaa-naack
2ag Cheetah RunFront
6ag Cheotal Run
6ag Chestah RunBad
bag Cheetah Jump
fag Cheetah RunBack
6ag Cheetah RunFront
Reacher Reach
Humanoid RundWalk
Humanoid Stand
Humanoid Standup
Walker2D Run
Walker2D RunBed
Walker2D
WWllkBwd
Swimmer--Swmmee
Hopper-Hop
Hopper Stand
Figure 11: Visualization of tasks. (Top) Bi-DexHands tasks. (Bottom) MA-Mujoco tasks.
5m, 6m
2s3z
(a) SMAC
(b) MPE
(c) RWARE
Figure Rendering tasks in other multi-task multi-agent environments.
of observations and actions. In MA-Mujoco, different agents are responsible for controlling different joints of the robots, (e.g., hips vs. knees), leading to partial observability and coordination requirements. This agent-level heterogeneity poses higher challenge for world modeling and multi-agent generalization.
Considering the above metrics, Bi-DexHands and MA-Mujoco exhibit diverse semantic relationships, varying dynamics, and significant agent heterogeneity, making them ideal for evaluating MT-MARL algorithms.
D.5 Bounded Similarity in Other Environments
We believe that the bounded similarity phenomenon not limited to continuous-comtrol settings but represents general property of multi-task environments. To further illustrate, we conducted experiments similar Figure on other environments. Maximum Mean Discrepancy (MMD) [9] used to measure pairwise similarity between task dynamics. The results are presented in Table 8, which demonstrate that tasks within the same group (e.g., MMM & MMM2) consistently exhibit much lower MMD scores than across-group pairs (e.g.. 2s3z & MMM), indicating clear intra-group
21

---

similarity and inter-group divergence. This confirms that bounded similarity is consistently observed across diverse environments, validating its generality.
Table 8: Pairwise MMD scores among tasks in different environments (lower indicates higher similarity).
SMAC
2s3z
3s5z
MMM
MMM2
2s3z
0.15
0.58
0.69
3s5z
0.58
0.69
MMM
0.09
MMM2
MPE
adversary
push 0.01
reference
spread
adversary
0.40
0.38
push
0.39
0.38
reference
0.03
spread
RWARE
left-easy
left-hard
right-easy
right-hard
left-casy
0.01
0.13
0.12
left-hard
0.14
0.13
right-easy
< 0.01
right-hard
E Additional Experimental Results
E.1
Additional Comparison Results
Due the varying reward ranges across different tasks, the multi-task performance in Table is presented after normalization, which can be formulated as,
Traw
bmin
Fnorm
D
bmin
x 100.
(14)
The boundaries in different tasks vary, which are reported Table 9,
Table 9: The boundaries of tasks in Bi- Bi-DexHands.
Task
buin
Dmax
Task
bmin
dmax
Task
buii
bmax
block-stack
150
325
bottle-cap
280
aitch-abreast
catch-over-underarm
catch-underarm
15
close-inward
close-outward
0
1000
open-inward
0
400
open-outward
0
600
kettle
-200
50
lift-underarm
-60
275
over
0
300
pen
0
125
scissors
-50
400
We present the comparison results of all tasks in MA-Mujoco in Table 10.
E.2
Statistical Analysis of Comparison Results
Across the 38 tasks from Bi-DexHands and MA-Mujoco, M3W achieves the best performance on 71% (27 out of 38) of the tasks. Conventionally, achieving the best result on more than 50% of tasks is considered statistically significant [17]. In addition, we conducted statistical analyses of the comparative results shown in Table and Table 10, with detailed outcomes summarized in Table 11. From the perspective of statistical significance, M3W demonstrates significant improvements (p < 0.05) on 27 out of 38 tasks. We also employed the Probability Superiority (PoS) to assess the relative performance of the algorithms. Comparing M3W to the best-pen best-performing baselines, M3W shows explicit advantage (PoS 64%) on tasks, overall advantage (PoS > 50%) on 27 tasks, and explicit disadvantage (PoS 36%) only 6 tasks.
22

---

Table 10: Comparision results multi-task MA-Muujco.
Suite
()W()
MACD
HMASD [40]
[38]
[41]
HATRPO [19]
run
31.4 = 0.5
25.6 1.2
14.7 3.0
5.9 0.3
4.3 0.3
4,6 = 0.2
run hankhacwwanss
33.0
16.3
0.2
2a-Choctah
jump
40.8 0.9
36.6 2.3
37.1 1.0
1.00
28.7 1.8
1844
run backleg
32.5
17.7 1.6
1
5.8 0.7
722 0.8
1.8 0.1
munfrromtleg
28.5
194 2.0
204
59
7
3.7
run
94
8.6
11.0
7.7
4.6
2.6 - 0.2
run wan-hackwards
12.3 123:004
7.5 0.7
18.1 - 0.7
14.2
9.8
5.6 + 0.3
6a Cheetah
jump
39.1 + 0.3
32.8 2 1.0
37.8 0.7
41.1 0.7
31.9:09
15.4 0.5
run backlac
13
4: 3.5 2
run-froud frontleg
16.6 -
14.2 0.8
18.2 + 0.3
1..5
9.2
run stand walk standup
35.6
d:
2.9 0.2
20.8 0.3
10.0 - 0.5
20.8 + 0.2
62
10.0 0.2
6.9 = 0.2
Humanaid
50.5
10.2 1::00. 0.6
22.7 1: 3.2
14.5 0.8
20.3 0.6
4,8 0.3
21.6
20.6
11.1 #. 0.3
23.4 1: 0.3
12.5 1.8
11.0 1.2
run hackwards
22.5
22.6 1.9
33.5 4.6
5.2
2.8
0.6 + 0.1
2488
14.3 2.2
19.3 4:
7:15::05
4,9 1 0.2
Walker2D
stand
90.3
0.8
86.0 +6.0+:0 1.7
83.8 118:10 1.9
60.1 2.5
66.4 2.9
67.4 3.7
walk-backwarrds
42.1
1,9
32.4 = 144
26.8 2.6
25.8 1.5
21.8 1.4
23.0 + 0.9
46.7
5.8
Swimmer
swim-backwards
46.6 =
$7.9 = 1.0
514
40.4 0.6
48.1
11.0 1.4
Hopper
hop
65.4
37.4
9.1 2.4
stand
87.2 1.1
77.1 2: 0.9
81.0 1.6
78.9 d: 0.8
76.7 : 1.0
78.3 0.8 0.8
Reacher
reach
87.4
68.2
2.6
67.7 2.3
61.8 2.0
Average
40.2
29.5 1.1
33.1 + 1.1
24.3 d: 1.0
24.1 d: 1.0
15.7 d: 1.0
Table 11: Statistical analysis across all tasks.
Task
Baseline
p-value
PoS
Task
Baseline
p-value
Po8
Bi-DexHands
over
MACD
<.0001
75.2%
bottle-cap
MAPPO
0.004
47.3%
cutch-abreast
MACD
<0.001
77.8%
calch-over2underram
MACD
72.2%
catch-underarm
HMASD
<0.001
68.9%
block-stack
MACD
<0.001
74.1%
door-close-inwand
HMASD
<00001
68.1%
door-closc-outtee
MACD
<0.001
72.6%
door-open-inward
MACD
0.069
45.5%
dour-open--otwwar
MACD
<0.001
89.1%
kettle
HMASD
0.075
55.1%
lift-underarm
MACD
< 0.001
76.7%
pen
MACD
<0.001
79.8%
scissors
HMASD
<0.001
49.2%
MA-Mujoco
2-chectab-run
MACD
<0.001
67.5%
2-chectah-runbwwl
MACD
<0.001
73.5%
scheetah-jump
HMASD
<0.001
67.9%
2-cheetab-rrubuat
HMASD
<0.001
2-cheetab-runfee
HMASD
<0.001
79.4%
6-cheetah-run
HMASD
<0.001
36.4%
6-cheetab-runbwd
HMASD
<00.01
3.9%
6-cheetah-jump
MAT
<0.001
27.3%
Schetahh-unbaade
HMASD
<0.001
86%
6-chectah- runfront
HMASD
<0.001
32.9%
humanoid-run
HMASD
0.032
98.9%
bumamoid-stand
HMASD
<0.041
65.2%
humanoid-wer walk
HMASD
<0.001
100.0%
humanoid-stanee -ttandup
MAT
<0.001
10.8%
walker2d-run
HMASD
<0.001
39.3%
walker2d-rubbwd
HMASD
<00,001
84.9%
walker2d-stand
MACD
<0.001
75.2%
walker21-walk
HMASD
<0.001
92.8%
walker2d-walle walkbbwd
MACD
<0.001
100.0%
swimmer-swim
MAT
0.007
31.0%
swimmer-swibbud
MACD
<0.001
hopper-hop
HMASD
<0.001
100.0%
hopper-stand
HMASD
<0.001
73.4%
reacher-reach
MACD
<0.001
90.4%
Probability of Superiority (PoS) non-parametric measure used quantify the likelihood that
randomly selected score from one algorithm (e.g., M3W) is greater than randomly selected score from another (e.g.. baseline). Formally, given two sets of performance scores X and Y, drawn from two algorithms respectively, PoS is defined as:
PoS(X,Y) P( 4) + for a E X
In our analysis, we interpret:PoS: 64% as an explicit advantage of one method over another, PoS > 50% as overall advantage, and PoS < 36% explicit disadvantage. This metric especially
23

---

useful when comparing algorithms across multiple stochastic tasks, as it avoids assumptions about the underlying performance distributions and offers an interpretable measure of practical dominance.
E.3 Details of Ablation Settings
Experimental Settings. We selected six tasks from MA-Mujoco involving two robot suites. In the task set [Cheetah: run, run-bud, run-front; Walker2D: run, run-bwd, walk), although the semantics and objectives differ across tasks, their dynamics exhibit bounded similarity. We collected 25K steps of transitions for each task using random actions, resulting in combined dataset of 150K transitions, which was used consistently across all models for training. Training hyperparameters are the same with Table 3. Models are trained for 2,500 iterations, during which all modules are frozen except the ablated components. Specifically, for dynamics prediction experiments, the unfrozen modules include the encoder, dynamics model, and task embedding, with the loss defined as follows:
H
C(0g,0D.e) Σ
x 90 (Ze. (Ee ar. -
(15)
For reward prediction experiments, the unfrozen modules include the encoder, reward model, and task embedding, with the loss defined as:
C(0g,0n.e) = Σ X (qn (Ee, a) r2t1)2.
(16)
In Eq. (15) and Eq. (16), the loss is backpropagated through the latent variable, where x = pe(o), e).
Parameter Sizes. We report the parameter sizes of all models in Table 12. The results reveal a trade-off between model size and prediction error. M3W increases the parameters count by +1 +12% (dynamics), +76% (reward), but achieves substantial reduction in prediction error by -53% (dynamics), -78% (reward).
Table 12: Model sizes and prediction errors of different models. ST means Switch Transformer, SA means Self-Attention.
M3W
ST
Transformer
SA
MLP
Size
707,072
630,768
192,016
35,648
Dynamies
Error
0.042 1: 0.007
0.067
0.009
0.119 0.011
0.422 4 0.058
Size
267,280
151,616
33,600
Reward
Error
0.14 + 0.03
0.59 0.12
0.81 + 0.26
The training curves of M3W across the six tasks are shown in Figure 13. The results reveal that the dynamics error curves naturally cluster into two distinct groups, corresponding to the Cheetah and Walker2D task suites, respectively. This observation further supports our claim regarding the presence of bounded similarity in multi-task dynamics.
E.4 Additional Ablation Studies on Actor Prior
As shown in Eq.1, M3W incorporates both an actor and critic, though their importance and roles differ significantly. The actor is used solely to generate candidate actions for the planner and does not act as a policy during decision-making return of the planning horizon in Eq. (7).
contrast, the critic is crucial, as estimates the expected
Table presents an ablation study on the actor prior. the planner-only variant, the planner does not rely on the actor prior, and all candidate actions are generated randomly. The results indicate that the actor prior contributes to improving action quality and reducing variance. However, M3W remains effective even without the actor prior, demonstrating the robustness of the planning component.
24

---

Cheetah: Run
RunBwd
RunFront
Walker:
Run
RunBwd
Walk
7
1.4
6
Dynamics Error (x0.000 (0.001)
1.2
Reward Error (xO)
5
4
0.8
3
0.6
0.4
0.2
0.0
0.5K
1K
1.5K
2K
2.5K
0
0.5K
1K
1.5K
2K
2.5K
Update Times
Update Times
Figure 13: Training curves M3W dynamics error and reward error.
Table 13: Ablation studies actor prior.
Task
Default
Planner-Only
Task
Default
Planner-Only
Bi-DexHands
over
70.0 + 1.3
67.3 1.6
bottle-cap
89.8 +
89.2 + 0.5
catch-abreast
64.1 1 4.4
60.3 1: 5.1
catch-ovve2underrm
84.8 1 1.4
83.8 1 2.0
catch-underarm
87.3 1 1.3
83.6 1 13.8
block-stack
88.1 1 3.1
84.1 1 2.9
door-close-nnward
91.3 1 0.7
90.2 1.2
door-close-outward
95.0 -1 1.2
94.8 1 1.5
door-open-inge oowan-
91.5
91.0 1.6
door-open-outwart
89.0
87.1 + 3.2
kettle
90.4 0.4
89.9 0.4
lift lift-underamm
77.0
73.3 + 2.1
83.3 3.3
76.3 - 6.8
scissors
76.4 1
69.9 + 6.6
average
84.1 + 0.8
81.4 3.5
MA-Mujoco
2-chectah-run
31.4 + 0.5
30.9
0.7
2-chcetah-runbwd
33.0 0.9
32.8 + 0.9
2-chectah-jump
40.8 0.9
39.6 + 1.5
2-chectah-runback
32.5
31.3
2-cheetah-runfront
28.5 0.8
0.9
6-cheetah-run
9.4
9.2
6-cheetah-runbwd
12.3 0.4
12.0 0.8
6-cheetah-jump
0.
0.3
-cheetah-runback
1.4
16.3
6-cheetah-runfront
1.66
1.64
humanoid-run
35.6
35.0 1.2
humanoid-stame
2.3
2.34
humanoid walk
50.5
50.1 2.6
seramoddsstnddun
21.6
20,9
walker2d-run
2.60 0.4
walker2d-runbwd
2..
0.42
walker2d-stand
9..3
0,9
walker2d-walt
7..7
7.6
walker2d-walkbwd
42.1 1.9
41.2
swimmer-swim
2.7
45.3
swimmer-swimbwd
46.6 1.1
45.9 - 0.8
hopper-hop
65.4 + 0.6
63.9
+ 1.1
hopper-stand
87.2 1.1
85.9 1.99000 0.7
reacher-reach
87.4 + 1.0
86.0 2.
average
40.2 =: 1.1
39.4 + 1.3
E.5 Robustness to Reward Scale
For policy-centric approaches, significant discrepancies in reward scales across tasks can lead to learning imbalance, where the policy is dominated by tasks with larger reward magnitudes. Since M3W includes reward predictor as an essential component, natural question arises: Is M3W robust variations in reward scale across tasks?
To assess robustness M3W this regard, we constructed a task set consisting of hopper (reward scaled by x100) and hopper-hop (unchanged). The challenge lies in the fact that stand is much simpler behavior but provides a disproportionately high reward, which can easily lead models to fall into local optimum. The learning curves shown in Figure 14.
25

---

M3W (ours)
MACD
HMASD
MAT
MAPPO
HATRPO
Sland
Hop
10000
50
8000
40
Average Retum
Average Retum
6000
4000
20
2000
10
0.2M
0.4M
0.6M
0.8M
0.2M
0.4M
0.6M
0.8M
1M
Training Steps
Training Steps
Figure 14: The results validate M3W's robustness to reward scales.
The results show that while all methods converge quickly on the stand task, only M3W retains the plasticity to continue learning the more complex hop behavior. This demonstrates M3W's robustness to large reward scale discrepancies. Its resilience can be attributed to three key factors:
1. Reward predictor with task embedding: M3W's reward predictor takes task embeddings as input, enabling it to differentiate between tasks. long as each task maintains internally consistent reward scaling, M3W can adjust its predictions accordingly.
2. Symlog-base based numerical transformation: As described Appendix C.2, when target values are large in magnitude, the symlog transformation stabilizes training by reducing resolution in high-value ranges.
3. Planning- based training paradigm: M3W does not explicitly rely policy function, which reduces the risk of becoming trapped in local optima.
E.6 Time Cost Analysis with an Early-Stopping Planner
Inference Time Cost. We report the per-step execution time single RTX A6000 GPU (Figure 15, left), where H denotes the rollout horizon and Kp is the number planner iterations. Under the default settings, M3W achieves execution time of approximately 25.8 ms per step the Bi-DexHands environment, which is well below the 50 ms threshold and satisfies the real-time requirements of most robotie tasks.
As shown in Table 14, while policy-centric methods avoid online planning and thus yield lower inference times at the cost of flexibility, our method remains comparable to, even faster than, existing planner-based approaches [13, 6]. Moreover, the carly-stopping heuristic further reduces the inference time to below 10 ms per step, making real-time applications feasible.
Table 14: Per-step inference time (ms) on Bi- Hi-DexHands.
Method
default
carly-stop
MACD
HMASD
MAT
MAPPO
HATRPO
Time Cost (ms)
25.8
9.8
1.4
3.4
2.2
1.1
1.1
Early-Stopping Planner. Planning trades additional computation for performance improvement. As the number of iterations increases, action quality and stability typically improve, but so does computational cost. We observe a clear diminishing return: the benefit of additional iterations decreases as action quality improvements taper off. To address this, we propose an early-stopping mechanism based on KI. divergence, which monitors the divergence between consecutive action distributions to assess convergence. When the change falls below a predefined threshold, the planner terminates early. shown on the right side of Figure 15, setting the threshold 0.5 achieves 62.1% speedup with only 1.9% performance drop.
26

---

100
H=3
H=5
Time Gast Stop (ms)
50
H=7
palicy
Average Retur
default
60
Time Cosl Step
30
40
20
20-
Average Return
2
8
10
default
0.1 0.25
0.5
1.0
actor
Number Planning Iterations
Thrsshold
Figure 15: Time cost analysis. (Left): Time cost of M3W with different rollout horizoms H/ and planner iterations Kº Bi-DexHands. (Right): Evaluation of average return and time cost for the early-stopping planner.
E.7 Scalability with the Number of Agents
In M3W, we only concatenate zf and of for each agent i, not concatenate all agents' spaces together. So prevents dimensional explosion as the number of agents increases. To validate the scalability of M3W, we conducted evaluations on Many-Agent-Mujjoon and MAgent [48], with results shown in Figure 16.
M3W (ours)
MACD
HMASD
MAT
MAPPO
HATRPO
10-agent-Swimmer
56-agent MAgent
100
100
80
80
Average Return
Average Return
60
40
0.5M
1M
1.5M
2M
25K
50K
75K
100K
Training Steps
Training Steps
Figure Additional Comparisons Paagent-Swimmer 56-agent MAgent, which confirm the scalability M3W with the number of agents.
NeurIPS Paper Checklist
Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The abstract and introduction introduce the concept bounded similarity in multi-task dynamics, prompting a reconsideration of the role of modularity in multi- task learning. Based on this perspective, Sec. 4 outlines M3W, a method that integrates Mixture-of-Expert (MoE) into the world model.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
27

---

The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: We acknowledge that the MPPI-based planner only applicable to continuous action spaces. This limitation is discussed in Sec. 6, where we also suggest potential directions for future improvements.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full of assumptions and a complete (and correct) proof?
Answer:
[NA]
Justification:
This paper does not include theoretical results
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged provide short proof sketch provide intuition.
28

---

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: This paper contains all relevant details of our implementation to reproduce the results.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access to hosted model (e.g., in the of a large language model), releasing of a model checkpoint, or other means that are appropriate the research performed.
While NeurlPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open- open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers have some path reproducing verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification:
The code are available at https://gttauu.ccm//hhoooj/////
m3v-marl
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (kttps://nips.cd cc/
for more details.
29

---

While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https:
/ /nips. wwtpppccc//cccccccccccccccccccct gwwcc/ggg////e/eeeeee /tteggguuddeesoeeuubniiii SoodeSubmissionPP
for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
The relevant details are provided in the paper.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: All of the results are accompanied by error bars, confidence intervals, or statistical significance tests.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results accompanied by enror bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report I-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
30

---

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: See Appendix E.6.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https ptttps://neurree http://nerrisss.coo http:///c//cccccccccccccteeeeeee
Answer: [Yes]
Justification: We have read the NeurlPS Code of Ethics and the paper conforms to the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g.. if there is special consid- cration due to laws or regulations in their jurisdiction..
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: The paper is concerned with foundational research and is theoretical in nature with no direct societal impact.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinfornation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
31

---

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
there are negative societal impacts, the authors could also discuss possible mitigation strategies (e-g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of MI.).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: The paper is theoretical nature and does not involve release of data or code and hence poses no such risks
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: We properly linked and cited the base model and source data corpus used in our experiments.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. wwl.cooo/aaaaaa com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is available online, the authors encouraged to reach out the asset's 's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
32

---

Answer:
[NA]
Justification: No assets are released in this paper.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
There were no such experiments or research
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not involve research with human subjects
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, originality of the research, declaration is not required.
Answer: [No]
33

---

Justification:
LLM is used only for writing, editing, or formatting purposes.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to LLM policy (https: : http://nurriss.c.eeeo http:r/pp....o/ooffeeeee/e////// for what should or should not be described.
34


Reviews:
[{"review_id": "ABb0xUMZc0", "reviewer": "Reviewer_DXnW", "summary": "This paper presents M3W, a mixture-of-experts world model for multi-task, multi-agent reinforcement learning. The authors point out that different tasks often need very different policies yet still share parts of the same physics, a phenomenon they call bounded similarity. To use this overlap, they place modularity in the world model instead of in the policy. A SoftMoE module predicts the next latent state, while a SparseMoE module predicts rewards; a planner then searches in latent space. On the Bi-DexHands and MA-Mujoco suites, M3W improves return and sample efficiency on opposing tasks when compared with recent modular-policy baselines like MACD and HMASD. Visual analyses further show that different experts specialize in intuitive joints and task groups, lending interpretability to the bounded-similarity idea.", "questions": "1.\tHave the authors conducted comparative experiments using a uniform SoftMoE or uniform SparseMoE setup? If the heterogeneous design is removed, how would the performance and computational cost change?\n2.\tTable 9 indicates that M3W achieves lower performance on the 6a-Cheetah task. Does this imply that M3W’s performance deteriorates as the number of agents increases?\n3.\tCould the proposed approach be adapted to partially observable environments? If so, what specific modifications would be necessary?", "limitations": "Yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n\n1. M3W moves modularity from the policy to an MoE-based world model, allowing tasks that share underlying physics but need different behaviors to reuse dynamics while avoiding the gradient clashes that often sink multitask RL.\n\n2. The paper first introduces the MoE-based world model for multi-task settings, combining a SoftMoE dynamics predictor and a SparseMoE reward predictor. Task-conditioned routing selectively re-uses experts, reducing gradient interference and enhancing interpretability through expert specialization.\n\n3. The paper incorporates a load-balance loss to equalize expert utilization, preventing overload of popular experts and enhancing training robustness, thereby avoiding the expert-imbalance pitfalls observed in early MoE works.\n4. Experiments on two challenging benchmarks-Bi-DexHands and MA-Mujoco show that M3W achieves the top return, with improvements significant over model-based and model-free baselines.\n\n5. The paper provides attention distributions of individual experts, cosine-similarity of router outputs revealing task-relatedness, and t-SNE visualizations of latent trajectories—strengthening the  interpretability of M3W.\n\n\nWeakness:\n\n1. The evaluation primarily focuses on two benchmarks which share similar continuous control settings. It remains unclear whether the concept of bounded similarity highlighted by the authors applies broadly to other types of environments, particularly those with different characteristics.\n\n2. Because the planner relies on MPPI control, M3W is only applicable to tasks with continuous action spaces. Consequently, the method cannot yet be deployed in many popular multi-agent benchmarks such as SMAC or Atari-style games. Moreover, M3W further assumes full observability, restricting applicability in partially observable domains.\n\n3. Even after applying the authors’ early-stopping heuristic, each control step still takes a considerable amount of time on an RTX A6000. Moreover, the paper reports no GPU-time for competing baselines, making direct comparison impossible.", "was_revised": true, "final_justification": "I appreciate the authors’ thoughtful response in their rebuttal, which effectively clarified my initial concerns and resolved my confusion. Accordingly, I have decided to increase my rating.", "rebuttal": "`Responses`\n---\n---\n> **Q1: Does the concept of bounded similarity generalize to other environments with with different characteristics? (Weakness.1)**\n\n**A1:** We believe that the bounded similarity phenomenon is not limited to continuous-control settings but represents **a general property of multi-task environments.** To further illustrate, we conducted experiments similar (Fig. 1, L27, main text) on other environments.\n\n+ **Experimental setting:** We conducted additional experiments on three representative environments discussed in (D.4 *Why Choose Bi-DexHands and MA-Mujoco?*, L541, Appendix) — **SMAC** (discrete actions, partial observability, heterogeneous agents), **MPE** (continuous actions, full observability, homogeneous agents), and **RWARE** (mixed difficulty and partial observability)—covering substantially different action spaces, observability assumptions, and agent heterogeneity.\n\n+ **Method:** Task trajectories can be treated as samples of the underlying dynamics. We measured pairwise similarity between task dynamics using Maximum Mean Discrepancy (MMD) [1] between task trajectories (with lower scores indicating higher similarity).\n\n+ **Results:** Across all environments, tasks within the same group (e.g., MMM & MMM2) consistently exhibit much lower MMD scores than across-group pairs (e.g., 2s3z & MMM), indicating clear intra-group similarity and inter-group divergence. This confirms that **bounded similarity is consistently observed across diverse environments, validating its generality.**\n\n+ For clarity, only the upper triangle of the symmetric MMD matrix is shown:\n\n  SMAC|2s3z|3s5z|MMM|MMM2\n  -|-|-|-|-\n  **2s3z**||**0.15**|0.58|0.69\n  **3s5z**|||0.58|0.69\n  **MMM**||||**0.09**\n  **MMM2**||||\n\n  MPE|adversary|push|reference|spread\n  -|-|-|-|-\n  **adversary**||**<0.01**|0.40|0.38\n  **push**|||0.39|0.38\n  **reference**||||**0.03**\n  **spread**||||\n\n  RWARE|left-easy|left-hard|right-easy|right-hard\n  -|-|-|-|-\n  **left-easy**||**<0.01**|0.13|0.12\n  **left-hard**|||0.14|0.13\n  **right-easy**||||**<0.01**\n  **right-hard**||||\n\n&nbsp;\n\n---\n> **Q2: Could the proposed approach be extended to partially observable and discrete-action environments, and how? (Weakness.2 & Question.3)**\n\n**A2:** We start from two perspectives: partially observable and discrete actions.\n\n+ **Partial Observability:** We acknowledge that the full-observability assumption is indeed a common limitation of existing model-based MARL methods [2,3]. However, M3W can be naturally extended to partial observability by two simple modifications:\n  1. **Training-time masking:** Inspired by Masked AutoEncoder [4] (and its applications in MARL [5]), we randomly mask other agents with a certain probability to simulate communication failures, thereby enhancing robustness.\n  2. **Planning-time communication cache:** The planner caches previously predicted rollouts and reuses them in case of communication failure, replacing missing ground-truth inputs to maintain trajectory continuity.\n  \n  We tested the above extensions on *BiDexHands-Over* and confirmed their effectiveness. In *mask* group, other agents are masked with 50% probability, and communication fails with 50% probability during execution. **These results show that M3W can adapt to partial observability with minimal performance degradation.**\n\n  ||full|mask|decentralized|\n  -|-|-|-\n  Dynamics Error ↓|0.025±0.003|0.033±0.005|0.044±0.005\n  Reward Error ↓|0.004±0.001|0.012±0.001|0.013±0.001\n  Return ↑|32.1±0.4|31.3±0.6|31.1±0.3\n\n+ **Discrete Action Space:** There are two modifications to extend our method to discrete action spaces:\n  1. Use one-hot actions or bitactions [6] to transform a discrete action into a continuous one.\n  2. Replace the MPPI control with Cross Entropy Method (CEM) or Monte Carlo Tree Search (MCTS).\n  \n  During training, discrete actions are converted into one-hot representations to facilitate understanding by the world model. During execution, a CEM-based planner is adopted, and its procedure is reported below. Similar to the above, we evaluated these extensions on the classic discrete tasks SMAC ‘3m’ and RWARE ‘left-easy’, achieving a **100% win rate** and a reward **> 7.75/100 steps**.\n\n  **CEM Procedure:** Assume the action dimension is $n$,\n  + S1: Initialize a discrete action distribution, $p^{(0)}=1/n$\n  + S2: Sample $M$ action sequences, and compute their values according to (Eq.7, L185, main text)\n  + S3: Select the top $m$ elite action sequences $a^∗$\n  + S4: For each action dimension $j$, update the distribution $p^{(k+1)}(a[j])=\\frac{1}{m} \\sum \\mathbb{I}(a[j]=a^*)$\n  + Repeat steps S2–S4 until convergence.\n\n&nbsp;\n\n---\n> **Q3: Inference time cost of baselines. (Weakness.3)**\n\n**A3:** The table below reports inference time (ms/step) on Bi-DexHands. While policy-centric methods avoid online planning and thus achieve lower inference times at the cost of flexibility, our method remains comparable to, or even faster than, existing planner-based approaches [7,8]. Moreover, our early-stopping heuristic further reduces the inference time to < 10 ms/step, making real-time applications feasible.\n\n||Ours(default)|Ours(early-stop)|MACD|HMASD|MAT|MAPPO|HATRPO\n|-|-|-|-|-|-|-|-\nTime Cost (ms)|25.8|9.8|1.4|3.4|2.2|1.1|1.1\n\n&nbsp;\n\n---\n> **Q4: How would the performance and computational cost change if the heterogeneous MoE design is replaced with a uniform one? (Question.1)**\n\n**A4:** The heterogeneous design was chosen for two main reasons:\n1. **SoftMoE for dynamics modeling:** As noted in [9], SoftMoE is fully differentiable and effectively captures token-wise correlations, making it suitable for modeling complex inter-agent interactions.\n2. **SparseMoE for reward prediction:** SoftMoE is an n-to-n architecture, whereas reward prediction is an n-to-1 mapping, so we choose SparseMoE.\n\nWe additionally conducted experiments replacing the heterogeneous design with uniform SoftMoE and uniform SparseMoE. The results show that uniform SoftMoE increases reward prediction error and time cost, while uniform SparseMoE worsens dynamics modeling. Thus, the heterogeneous design achieves the best trade-off between accuracy and efficiency.\n\n||Dynamics Error ↓|Reward Error ↓|Time Cost(ms/step) ↓\n-|-|-|-\nHeterogeneous|**0.042±0.007**|**0.147±0.004**|25.8\nUniform-SoftMoE||0.151±0.004|30.1\nUniform-SparseMoE|0.063±0.008||**19.2**\n\n&nbsp;\n\n---\n> **Q5: Does the performance deteriorates as the number of agents increases? (Question.2)**\n\n**A5:** our method achieves second-best performance on 6a-Cheetah. We believe the slight performance drop is not due to scalability limitations, but rather due to limited training samples for 6-agent tasks in the MA-MuJoCo benchmark.\n\n+ **Scalability Analysis:** As shown in (E.7 *Scalability with the Number of Agents, L663*, Appendix), we evaluated M3W on *10-agent Swimmer* and *56-agent MAgent*, where it maintained stable performance and outperformed all baselines. This demonstrates that M3W can scale to significantly larger numbers of agents.\n+ **Why the Drop on 6a-Cheetah?** In (Section 4.1 *Learning Multi-Agent Dynamics as Sequence Prediction*, L108, main text), we model the multi-agent dynamics as a sequence prediction problem, which makes our approach more sensitive to the number of agents compared to policy-centric methods. Most MA-MuJoCo tasks involve 2–3 agents, and 6-agent tasks are underrepresented, leading to slightly worse sample efficiency.\n+ **Future Work:** We plan to further incorporate permutation-invariance module or autoregressive mechanism to better handle tasks with a varying number of agents.\n\n&nbsp;\n\n\n`Summary`\n---\n---\n\nThank you for your valuable review and constructive feedback. We appreciate your recognition of the **motivation**, **novelty**, **strong performance**, and **interpretability** of our work.\n\nWe also acknowledge the affirmations from other reviewers R2 (gkLA), R3 (mmFi), and R4 (k53v) on:\n\n+ **Motivation and novelty:** (**R2:** *identifies the novel insight of ‘bounded similarity’*; *introduce a MoE within world model is a conceptual shift in MT-MARL*), (**R3:** *well-structured and strongly motivated*), (**R4:** *proposes a new perspective on modularity*).\n\n+ **Method design:** (**R2:** *allows better exploitation of inter-task structure;* *modular components selectively share or isolate knowledge*), (**R4:** *MoE world model effectively captures bounded similarity among tasks*).\n\n+ **Strong performance:** (**R2:** *outperforms baselines in both return and sample efficiency*), (**R3:** *really strong results*), (**R4:** *comprehensive ablation studies and visualizations*; *superior performance and sample efficiency*).\n\n+ **Interpretability:** (**R3:** *interpretability analysis shows the mechanism of sharing experts*).\n\nWe hope that our responses have fully addressed your concerns, and we welcome any further suggestions and discussions. **If our responses have resolved your concerns, we would sincerely appreciate your consideration for raising the score.**\n\n&nbsp;\n\n\n`References`\n---\n---\n\n[1] A kernel two-sample test, JMLR, 2012.\n\n[2] Aligning Credit for Multi-Agent Cooperation via Model-based Counterfactual Imagination, AAMAS, 2024.\n\n[3] Decentralized Transformers with Centralized Aggregation are Sample-Efficient Multi-Agent World Models, TMLR, 2025.\n\n[4] Masked Autoencoders Are Scalable Vision Learners, CVPR, 2022.\n\n[5] MA2E: Addressing Partial Observability in Multi-Agent Reinforcement Learning with Masked Auto-Encoder, ICLR, 2025.\n\n[6] INS: Interaction-aware Synthesis to Enhance Offline Multi-agent Reinforcement Learning, ICLR, 2025.\n\n[7] TD-MPC2: Scalable, Robust World Models for Continuous Control, ICLR, 2024.\n\n[8] Sparse Imagination for Efficient Visual World Model Planning, arXiv, 2025.\n\n[9] From Sparse to Soft Mixtures of Experts, ICLR, 2024."}, {"review_id": "0fed01YrWx", "reviewer": "Reviewer_gkLA", "summary": "The paper proposes M3W, a novel framework for multi-task multi-agent reinforcement learning (MT-MARL) that leverages a Mixture-of-Experts architecture applied to the world model instead of the policy. M3W is motivated by the observation of bounded similarity across task dynamics—tasks can share highly similar dynamics within groups but differ greatly across groups. To exploit this, M3W employs a SoftMoE-based dynamics model and a SparseMoE-based reward predictor, allowing for modular learning that promotes knowledge reuse across similar tasks while isolating dissimilar ones to prevent gradient interference. Unlike policy-centric approaches, M3W performs model-based planning by generating imagined trajectories from the world model and directly optimizing actions without an explicit policy.", "questions": "1/ The model assumes that each agent can access the observations and actions of all other agents through communication. However, in real-world decentralized or partially observable environments (e.g., autonomous driving), such global observability is often infeasible due to sensing limitations and communication constraints. How would M3W perform in more realistic settings with limited or no inter-agent communication? Could the authors discuss possible extensions to support decentralized training or partial observability?\n\n2/  There appears to be a mismatch between the loss formulation in Eq. (2), which uses mean squared error for reward and Q-value regression, and the actual implementation described in Appendix C.1, which uses discretized two-hot encoding and soft cross-entropy loss. Could the authors clarify this discrepancy? Would it be possible to revise Eq. (2) to reflect the actual training objective for better transparency and reproducibility?", "limitations": "Yes. The author(s) identified the limitation that their method is only applicable to continuous action spaces.\n\nThere may be additional limitations, as listed in the \"Strengths and Weaknesses\" section, that the author(s) could consider including in the paper.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n\n1/ The paper introduces a MoE architecture within the world model instead of applying modularity to policies (as most prior works do). This is (to my best knowledge) a conceptual shift in multi-task multi-agent reinforcement learning. This enables both dynamics prediction and reward modeling to be modular and task-adaptive, allowing better exploitation of inter-task structure.\n\n2/ The paper identifies the novel insight of 'bounded similarity', where task dynamics are locally similar within certain groups but globally different across others, and effectively uses this to design modular components that selectively share or isolate knowledge.\n\n3/ Regarding experiments: The method achieves state-of-the-art performance on two challenging benchmarks: Bi-DexHands and MA-Mujoco. It outperforms policy-centric baselines like MAPPO, MAT, and HMASD in both average return and sample efficiency.\n\nWeaknesses:\n\nThese are potential areas for future improvement and do not necessarily represent weaknesses:\n\n1/ The model assumes that each agent can access the observations and actions of all other agents through communication (e.g., a shared communication round). However, this assumption may not hold in real-world decentralized or partially observable environments. For example, in autonomous driving scenarios with multiple vehicles (agents), each vehicle typically has limited local sensing and cannot directly observe the internal states or actions of nearby vehicles due to communication delays, privacy constraints, or sensor occlusions. This makes full observability and synchronized communication impractical, limiting the direct applicability of M3W to such domains.\n\n2/ A potential inconsistency arises between the loss formulation Eq. (2) and the actual implementation details described in Appendix C.1. Eq. (2) expresses the reward and Q--alue losses as standard mean squared errors, implying scalar regression, the appendix clarifies that both targets are discretized using a two-hot encoding scheme and trained via soft cross-entropy loss. This discrepancy suggests that the formulation in Eq. (2) does not accurately reflect the implementation.\n\nTo improve transparency, the authors should revise or annotate Eq. (2) to indicate the use of discretized targets and soft classification loss.", "was_revised": true, "final_justification": "The authors have thoroughly addressed my questions.", "rebuttal": "`Responses`\n---\n> **Q1: How can M3W adapt to decentralized or partially observable settings given its reliance on full observability? (Weakness.1 & Question.1)**\n\n**A1:** We acknowledge that the full-observability assumption is indeed a common limitation of existing model-based MARL methods [1,2]. However, M3W can be naturally extended by two simple modifications:\n\n1. **Training-time masking:** Inspired by Masked AutoEncoder [3] (and its applications in MARL [4]), we randomly mask other agents with a certain probability to simulate communication failures, thereby enhancing robustness.\n   \n2. **Planning-time communication cache:** The planner caches previously predicted rollouts and reuses them in case of communication failure, replacing missing ground-truth inputs to maintain trajectory continuity.\n\nWe tested the above extensions on *BiDexHands-Over* and confirmed their effectiveness. In *mask* group, other agents are masked with 50% probability, and communication fails with 50% probability during execution. **These results show that M3W can adapt to partial observability and limited communication with minimal performance degradation.**\n\n||full|mask|decentralized|\n-|-|-|-\nDynamics Error ↓|0.025±0.003|0.033±0.005|0.044±0.005\nReward Error ↓|0.004±0.001|0.012±0.001|0.013±0.001\nReturn ↑|32.1±0.4|31.3±0.6|31.1±0.3\n\nHere are some formulation details of the communication cache. \n\nAt any time step $t$, the planner receives the observation $o^i\\_t$ and $o^{-i}\\_t$, which are encoded into ground-truth latent states $z^i\\_t$ and $z^{-i}\\_t$ ($-i$ denotes other agents). Starting from $z^{i}\\_t$ and $z^{-i}\\_t$, the world model performs a rollout to generate a latent trajectory of length $H$, $\\Gamma\\_1 = \\\\{\\hat{z}^i\\_{t+1:t+H}, \\hat{z}^{-i}\\_{t+1:t+H}\\\\}$. Ideally, at time $t+1$, the planner would again retrieve $o^i\\_{t+1}$, $o^{-i}\\_{t+1}$ from the environment and generate a new rollout, $\\Gamma\\_2 = \\\\{\\hat{z}^i\\_{t+2:t+H+1}, \\hat{z}^{-i}\\_{t+2:t+H+1}\\\\}$.\n\nNote that $\\Gamma\\_1$ and $\\Gamma\\_2$ share overlapping predictions. Therefore, we store $\\Gamma\\_1$ in a cache, and when communication fails at $t+1$, the planner reuses the cached prediction $\\hat{z}^{-i}\\_{t+1}$ to replace the missing ground-truth $z^{-i}\\_t$. Therefore, the model has the ability to cope with communication interruptions.\n\n&nbsp;\n\n---\n> **Q2: Discrepancy between Eq.(2)’s MSE loss and the implementation’s two-hot cross-entropy. (Weakness.2 & Question.2)**\n\n**A2:** As you pointed out, the description in (C.1 *Discrete Regression of Reward and Value*, L493, Appendix) accurately reflects the actual implementation. We will update Eq.(2) in the revised version as follows:\n\n$$\\mathcal{L}(\\theta, \\psi) = \\sum\\_i^n \\sum\\_t^H \\lambda^t \\left(\n\\underbrace{\\left\\| \\hat{z}^i\\_{t+1} - z^i\\_{t+1} \\right\\|\\_2^2}\\_{\\textrm{Dynamic Loss}} + \\underbrace{\\textrm{SoftCE}\\left( \\hat{r}\\_{t+1} - r\\_{t+1} \\right)}\\_{\\textrm{Reward Loss}} + \\underbrace{\\textrm{SoftCE}\\left( \\hat{q}\\_t - G\\_t \\right)}\\_{\\textrm{Q Loss}}\n\\right)$$\n\nwhere,  $\\hat{z}^i\\_{t+1}=q\\_D(\\mathbf{z\\_t}, \\mathbf{a\\_t}, e)^i, \\hat{r}^i\\_{t+1}=q\\_R(\\mathbf{z\\_t}, \\mathbf{a\\_t}, e), \\hat{q}\\_t = Q(\\mathbf{z\\_t}, \\mathbf{a\\_t}, e)$\n\n&nbsp;\n\n\n`Summary`\n---\n\nThank you for your valuable review and constructive feedback. We appreciate your recognition of the **motivation**, **model design** and **strong performance** of our work.\n\nWe also acknowledge the affirmations from other reviewers R1(DXnW), R3 (mmFi), and R4 (k53v) on:\n\n+ **Motivation and novelty:** (**R1:** *moves modularity from policy to world model*, *first introduce MoE-based world model for multi-task setting*), (**R3:** *well-structured and strongly motivated*), (**R4:** *proposes a new perspective on modularity*).\n\n+ **Method design:** (**R1:** *task-conditioned routing reduces gradient interference*, *load-balance loss enhances training robustness*), (**R4:** *MoE world model effectively captures bounded similarity among tasks*).\n\n+ **Strong performance:** (**R1:** *achieves the top return*, *with significant improvements over baselines*), (**R3:** *really strong results*), (**R4:** *comprehensive ablation studies and visualizations*; *superior performance and sample efficiency*).\n\n+ **Interpretability:** (**R1:** *enhancing interpretability through expert specialization*, *visualizations strengthen the interpretability of M3W*), (**R3:** *interpretability analysis shows the mechanism of sharing experts*).\n\nWe hope that our responses have fully addressed your concerns, and we welcome any further suggestions to improve this work. **If our clarifications have resolved your concerns, we would sincerely appreciate your consideration for raising the score.**\n\n&nbsp;\n\n`References`\n---\n---\n\n[1] Aligning Credit for Multi-Agent Cooperation via Model-based Counterfactual Imagination, AAMAS, 2024.\n\n[2] Decentralized Transformers with Centralized Aggregation are Sample-Efficient Multi-Agent World Models, TMLR, 2025.\n\n[3] Masked Autoencoders Are Scalable Vision Learners, CVPR, 2022.\n\n[4] MA2E: Addressing Partial Observability in Multi-Agent Reinforcement Learning with Masked Auto-Encoder, ICLR, 2025."}, {"review_id": "YIZiVomBJO", "reviewer": "Reviewer_mmFi", "summary": "This paper discusses the paradigm of Multi-Task Multi-Agent Reinforcement Learning (MT-MARL) where a group of agents are cooperatively solving a set of tasks while employing a single model. In this paper, the authors focused on modeling a world model instead of a centralized policy by utilizing a mixture of experts architecture that helps in sharing knowledge between similar tasks while disentangling sharing across dissimilar ones. This work is motivated by the phenomenon, referred to as bounded similarity, where a group of similar tasks exhibits similarity in their underlying dynamics, while across groups dissimilarities appear. The proposed approach, named M3W, has been evaluated on Bi-DexHands and MA-Mujoco, showing strong performance and sample efficiency. In addition, interpretability analysis has been conducted to show the mechanism of sharing experts among a certain group of similar tasks.", "questions": "- What is the effect of scaling the model down or up (in terms of the number of experts) with the number of tasks or the number of agents?", "limitations": "- I believe a major limitation of the proposed method, which I believe this work did not discuss well, is the scalability of this approach due to the high parameter count. I understand the improved results by this method, but this raises the question of whether this is due to the capacity of the model and not the design choices.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "# Strengths\n\n- In my opinion, this paper is well-structured and strongly motivated. In addition, I appreciate the flow of the paper and the quality of writing.\n- The proposed method connects recent advances in supervised learning (e.g, using SoftMoE and SparseMoE) and exploits their benefits in an interesting problem in RL.\n- The proposed method sounds and demonstrates really strong results on two different benchmarks.\n\n# Weaknesses\n\n- The proposed method utilizes a complex structure with many parameters. In the appendix, the parameter count is listed for each method, showing the computational burden of the proposed method.\n- The demonstrated concept and tools are known and not original. Yet, the technical aspect of the paper and the presented results dominate that.\n- One important ablation study is missing, which might have a strong impact on the work. An analysis of the number of experts needed is essential to understand the scalability of this method. I believe that not discussing it is affecting the fairness of the evaluation.\n- Another important weakness to highlight is the weak and insufficient discussion of the MTRL literature and the MoE in the MTRL literature. Since the problem tackled in this work is MT-MARL, the literature of each component should have been discussed in more depth. For instance, many works [1, 2, 3, 4] that utilize MoE or multiple models in MTRL were never mentioned.\n- Minor typos and errors, in Line 196, $q_D$ should be predicting the next latent state $z_{t+1}^i$ instead of the current. In Line 205, double \"on\".\n\n[1] Sodhani, Shagun, Amy Zhang, and Joelle Pineau. \"Multi-task reinforcement learning with context-based representations.\" International Conference on Machine Learning. PMLR, 2021.\n\n[2] Sun, Lingfeng, et al. \"Paco: Parameter-compositional multi-task reinforcement learning.\" Advances in Neural Information Processing Systems 35 (2022): 21495-21507.\n\n[3] Cheng, Guangran, et al. \"Multi-task reinforcement learning with attention-based mixture of experts.\" IEEE Robotics and Automation Letters 8.6 (2023): 3812-3819.\n\n[4] Hendawy, Ahmed, Jan Peters, and Carlo D'Eramo. \"Multi-Task Reinforcement Learning with Mixture of Orthogonal Experts.\" The Twelfth International Conference on Learning Representations.", "was_revised": true, "final_justification": "Since the authors addressed my concerns and questions, I will increase my score.", "rebuttal": "`Responses`\n---\n---\n\n> **Q1: Whether the improvements are due to the capacity of the model rather than the design choices? (Weakness.1, Limitation)**\n\n**A1:** We agree that increased parameter count contributes to some extent, but **the main driver is the modular design**. To further disentangle the effect of modularity from parameter scaling, we conducted an experiment similar to Figure 5 in [1]. We keep the number of experts $m$ fixed but reduce the dimension of each expert to $\\frac{1}{m}$. \n\nFollowing the experimental settings in (E.3 *Details of Ablation Settings*, L601, Appendix), the results shows that, SoftMoE achieves significantly lower prediction error compared to a dense MLP model with a similar parameter count. **These results indicate that performance gains primarily stem from the modular design, not merely from increasing parameters.**\n\n|Model|Parameters|Prediction Error\n-|-|-\nSoftMoE|707,072|0.042±0.007\nSoftMoE(÷m)|127,232|0.064±0.016\nDensen|128,290|0.213±0.051\n\n&nbsp;\n\n---\n\n> **Q2: Missing an analysis of the number of experts. (Weakness.3, Question)**\n\n**A2:** Following your constructive suggestion, we conducted an additional ablation study on the number of experts, and the results show:\n\n+ **Dynamics prediction:** The dynamics error decreases rapidly as the number of experts increases but plateaus after ~16 experts across all task scales. This is because SoftMoE’s fully differentiable routing effectively utilizes additional experts while maintaining stable performance.\n\n+ **Reward prediction:** The reward error first decreases and then increases as the number of experts grows. Too few experts underfit the reward function, whereas too many slow down training due to sparse routing.\n\nDetailed numerical results are provided in the following table.\n\n| | 6 tasks | | 12 tasks| | 24 tasks| |\n|-|-|-|-|-|-|-|\n|Experts| Dynamics | Reward | Dynamics | Reward | Dynamics | Reward |\n| 4 | 0.043±0.007 | 0.218±0.070 | 0.061±0.010 | 0.421±0.131 | 0.113±0.019 | 0.644±0.198 |\n| 8 | 0.043±0.007 | **0.114±0.027** | 0.048±0.008 | 0.211±0.084 | 0.098±0.017 | 0.502±0.113 |\n| 16| **0.042±0.007** | 0.147±0.035 | **0.046±0.007** | **0.151±0.039** | 0.067±0.011 | **0.185±0.054** |\n| 24| 0.043±0.007 | 0.176±0.043 | 0.048±0.007 | 0.217±0.086 | 0.062±0.010 | 0.268±0.094 |\n| 32| **0.042±0.007** | 0.209±0.049 | **0.046±0.007** | 0.288±0.093 | **0.060±0.010** | 0.365±0.101 | \n\n&nbsp;\n\n---\n\n> **Q3: The demonstrated concept and tools are known and not original. (Weakness.2)**\n\n**A3:** Thank you for recognizing the technical soundness of our work. We would like to explain that the main novelty of our paper lies in **its conceptual and methodological contributions** rather than introducing new tools or network architectures. Specifically:\n\n+ **Bounded similarity:** We identify the phenomenon of bounded similarity in multi-task settings, offering a new perspective on modeling task relationships — as emphasized by **R1** (*the authors point out … called bounded similarity*) and **R2** (*identifies the novel insight of ‘bounded similarity’*).\n\n+ **MoE-based world model:** We apply MoE to the world model instead of the policy, introducing a new paradigm for modularity in MT-MARL — as emphasized by **R2** (*conceptual shift in MT-MARL*) and **R4** (*a new perspective on modularity in MTRL*).\n\n&nbsp;\n\n---\n\n> **Q4: Insufficient discussion of MTRL and MoE-related literature, plus minor typos. (Weakness.4, Weakness.5)**\n\n**A4:** Thank you for your careful and constructive review. Following your suggestion, we will expand the related work section to include recent MoE-based MTRL approaches and clarify their differences from our work.\n\nWe also appreciate your attention. The typos in Line 196 and Line 205 have been corrected, and we will carefully proofread the entire paper to avoid similar errors.\n\nThe following are the relevant paragraphs that are expected to be added：\n\n> As model capacity grows and generalization becomes more important, Mixture of Experts (MoE) has gained attention in multi-task reinforcement learning (MTRL) for balancing shared and task-specific knowledge.\n> \n> CARE[2] proposed context-based representations to differentiate tasks within a shared policy, laying the foundation for the integration of MoE in MTRL. Building on this, Paco[3] introduced a parameter-compositional MTRL method, which  combines shared and task-specific parameters, enhancing knowledge transfer. More structurally aligned with traditional MoE, AMESAC[4] proposed an attention-based expert selection mechanism, dynamically selecting different experts to handle task heterogeneity. Furthermore, MOORE[5] developed a MoE framework with orthogonal constraints, promoting functional decoupling among experts to improve generalization and sample efficiency.\n\n\n&nbsp;\n\n`Summary`\n---\n---\n\nThank you for your valuable review and constructive feedback. We appreciate your recognition of the **motivation**, **technical soundness** and **strong performance** of our work.\n\nWe also acknowledge the affirmations from other reviewers R1(DXnW), R2 (gkLA), and R4 (k53v) on:\n\n+ **Motivation and novelty:** (**R1:** *moves modularity from policy to world model*, *first introduce MoE-based world model for multi-task setting*), (**R2:** *identifies the novel insight of ‘bounded similarity’*; *introduces an MoE-based world model … a conceptual shift in MT-MARL*), (**R4:** *proposes a new perspective on modularity*).\n\n+ **Method design:** (**R1:** *task-conditioned routing reduces gradient interference*, *load-balance loss enhances training robustness*), (**R2:** *allows better exploitation of inter-task structure;* *modular components selectively share or isolate knowledge*), (**R4:** *MoE world model effectively captures bounded similarity among tasks*).\n\n+ **Strong performance:** (**R1:** *achieves the top return*, *with significant improvements over baselines*), (**R2:** *outperforms baselines in both return and sample efficiency*), (**R4:** *comprehensive ablation studies and visualizations*; *superior performance and sample efficiency*).\n\n+ **Interpretability:** (**R1:** *enhancing interpretability through expert specialization*, *visualizations strengthen the interpretability of M3W*).\n\nWe hope that our responses have fully addressed your concerns, and we welcome any further suggestions to improve this work. **If our clarifications have resolved your concerns, we would sincerely appreciate your consideration for raising the score.**\n\n&nbsp;\n\n`References`\n---\n---\n\n[1] Mixtures of Experts Unlock Parameter Scaling for Deep RL, ICML, 2024.\n\n[2] Multi-task reinforcement learning with context-based representations, ICML, 2021.\n\n[3] Paco: Parameter-compositional multi-task reinforcement learning, NeurIPS, 2022.\n\n[4] Multi-task reinforcement learning with attention-based mixture of experts, IEEE RAL, 2023.\n\n[5] Multi-Task Reinforcement Learning with Mixture of Orthogonal Experts, ICLR. 2024."}, {"review_id": "Jah2MoVCDw", "reviewer": "Reviewer_k53v", "summary": "The paper presents M3W, a novel approach for multi-task multi-agent reinforcement learning that applies a mixture-of-experts (MoE) structure to the world model instead of the policy. It identifies the phenomenon of bounded similarity in multi-task dynamics and leverages this property to improve learning efficiency and planning performance.", "questions": "1. Could you provide more detailed explanation on why the actor in the world model is not used as an explicit decision-making policy? What are the advantages and disadvantages of using it solely for generating candidate actions for the planner compared to using it as a direct policy? Is there any experimental comparison?\n2. How does the training cost and inference time of the world model compare to policy-based methods? Could you provide some analysis or experimental results on this aspect?", "limitations": "The authors have mentioned some limitations of the current work, such as the applicability to continuous action spaces only. However, more discussion on the potential limitations regarding the scalability of the world model would be beneficial.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "* Strengths:\n1. Proposes a new perspective on modularity in multi-task learning by applying MoE to the world model, which effectively captures bounded similarity among tasks.\n2. Demonstrates superior performance and sample efficiency in experiments on Bi-DexHands and MA-Mujoco.\n3. Provides comprehensive ablation studies and visualizations to support the effectiveness of the MoE-based architecture.\n* Weaknesses:\n1. Limited discussion on the training cost and inference time of the world model compared to policy-based methods, which is crucial for understanding the practicality of the approach.\n2. The generalizability of the world model to other scenarios is not thoroughly analyzed.", "was_revised": true, "final_justification": "I have increased my score, as the authors have addressed my concerns.", "rebuttal": "`Responses`\n---\n---\n\n> **Q1: The training cost and inference time of the world model compare to policy-based methods. (Weakness.1, Question.2)**\n\n**A1:** As reported in (E.6 *Time Cost Analysis with an Early-Stopping Planner*, L649, Appendix), we have already provided the inference time of our method. Following your suggestions, we additionally report the training and inference time of the baselines in the table below:\n\n||Ours(default)|Ours(early-stop)|MACD|HMASD|MAT|MAPPO|HATRPO\n|-|-|-|-|-|-|-|-\nInference (ms)|25.8|9.8|1.4|3.4|2.2|1.1|1.1\nTrain (h)|~19||~34|~14|~11|~8|~14\n\nNote that planning-based methods inherently trade additional computation for improved performance, which aligns with our design objectives. Compared to policy-centric baselines, our method incurs higher computational cost but achieves substantially better sample efficiency and performance. Furthermore, its inference time is comparable to, or even faster than, other existing planner-based approaches [1,2]. Moreover, our early-stopping heuristic further reduces the inference time to **< 10 ms/step**, making real-time applications feasible.\n\n&nbsp;\n\n---\n\n> **Q2: The generalizability of the world model to other scenarios is not thoroughly analyzed. (Weakness.2)**\n\n**A2:** The factors that hinder the generalization of our method to other scenarios include the full observation assumption and the continuous action space limitation.\n\n+ **Partial Observability:** We acknowledge that the full-observability assumption is indeed a common limitation of existing model-based MARL methods [3,4]. However, our method can be naturally extended to partial observability by two simple modifications:\n  \n  1. **Training-time masking:** Inspired by Masked AutoEncoder [5] (and its applications in MARL [6]), we randomly mask other agents with a certain probability to simulate communication failures, thereby enhancing robustness.\n   \n  2. **Planning-time communication cache:** The planner caches previously predicted rollouts and reuses them in case of communication failure, replacing missing ground-truth inputs to maintain trajectory continuity.\n   \n  + **Preliminary results:** We tested the above extensions on *BiDexHands-Over* and confirmed their effectiveness. In *mask* group, other agents are masked with 50% probability, and communication fails with 50% probability during execution. **These results show that our method can adapt to partial observability with minimal performance degradation.**\n\n    ||full|mask|decentralized|\n    -|-|-|-\n    Dynamics Error ↓|0.025±0.003|0.033±0.005|0.044±0.005\n    Reward Error ↓|0.004±0.001|0.012±0.001|0.013±0.001\n    Return ↑|32.1±0.4|31.3±0.6|31.1±0.3\n\n+ **Continuous Action Space:** There are two modifications to extend our method to discrete action spaces:\n  1. Use one-hot actions or bitactions [6] to transform a discrete action into a continuous one.\n  2. Replace the MPPI control with Cross Entropy Method (CEM) or Monte Carlo Tree Search (MCTS).\n  \n  During training, discrete actions are converted into one-hot representations to facilitate understanding by the world model. During execution, a CEM-based planner is adopted, and its procedure is reported below. Similar to the above, we evaluated these extensions on the classic discrete tasks SMAC ‘3m’ and RWARE ‘left-easy’, achieving a **100% win rate** and a reward **> 7.75/100 steps**.\n\n  **CEM Procedure:** Assume the action dimension is $n$,\n  + S1: Initialize a discrete action distribution, $p^{(0)}=1/n$\n  + S2: Sample $M$ action sequences, and compute their values according to (Eq.7, L185, main text)\n  + S3: Select the top $m$ elite action sequences $a^∗$\n  + S4: For each action dimension $j$, update the distribution $p^{(k+1)}(a[j])=\\frac{1}{m} \\sum \\mathbb{I}(a[j]=a^*)$\n  + Repeat steps S2–S4 until convergence.\n\n&nbsp;\n\n---\n\n> **Q3: Why the actor is not used as an explicit policy? What are the advantages and disadvantages? (Question.1)**\n\n**A3:** \n\n1. **Reason:** We believe **the dynamics model to be more reliable and transferable across tasks than the actor.** As you noted in your Summary and Strengths.1, the bounded similarity phenomenon implies that different tasks share similar underlying dynamics but have highly diverse policies, making the learned actor less robust for multi-task generalization.\n\n2. **Advantages:** The planner achieves considerably higher performance. (E.4, *Additional Ablation Studies on Actor Prior*, L620, Appendix) compared planner performance when using actor-generated candidate actions versus random candidates. Following your suggestions, we additionally included a direct actor-only policy as a baseline. Results are shown below:\n\n    ||actor+plan|plan-only|actor-only\n    -|-|-|-\n    Return ↑| **84.1 ± 0.8** | 81.4 ± 3.5 | 76.2 ± 2.1\n\n3. **Disadvantages:** Directly using the actor as the decision-making policy offers better time efficiency. To mitigate the planner’s higher time cost, we introduced an early-stopping heuristic in (E.6 *Early-Stopping Planner*, L655, Appendix), which terminates planning when the KL-divergence between consecutive action distributions falls below a threshold. With a threshold of 0.5, **it achieved a 62.1% speedup with only a 1.9% performance drop.**\n\nThese results show that the actor serves best as an action generator for the planner, while the planner remains the primary decision-making module for robust multi-task performance.\n\n&nbsp;\n\n\n`Summary`\n---\n---\n\nThank you for your valuable review and constructive feedback. We appreciate your recognition of the **motivation**, **technical soundness** and **strong performance** of our work.\n\nWe also acknowledge the affirmations from other reviewers R1(DXnW), R2 (gkLA), and R3 (mmFi) on:\n\n+ **Motivation and novelty:** (**R1:** *moves modularity from policy to world model*, *first introduce MoE-based world model for multi-task setting*), (**R2:** *identifies the novel insight of ‘bounded similarity’*; *introduces an MoE-based world model … a conceptual shift in MT-MARL*), (**R3:** *well-structured and strongly motivated*).\n\n+ **Method design:** (**R1:** *task-conditioned routing reduces gradient interference*, *load-balance loss enhances training robustness*), (**R2:** *allows better exploitation of inter-task structure;* *modular components selectively share or isolate knowledge*).\n\n+ **Strong performance:** (**R1:** *achieves the top return*, *with significant improvements over baselines*), (**R2:** *outperforms baselines in both return and sample efficiency*), (**R3:** *really strong results*).\n\n+ **Interpretability:** (**R1:** *enhancing interpretability through expert specialization*, *visualizations strengthen the interpretability of M3W*), (**R3:** *interpretability analysis shows the mechanism of sharing experts*).\n\nWe hope that our responses have fully addressed your concerns, and we welcome any further suggestions to improve this work. **If our clarifications have resolved your concerns, we would sincerely appreciate your consideration for raising the score.**\n\n\n&nbsp;\n\n`References`\n---\n---\n\n[1] TD-MPC2: Scalable, Robust World Models for Continuous Control, ICLR, 2024.\n\n[2] Sparse Imagination for Efficient Visual World Model Planning, arXiv, 2025.\n\n[3] Aligning Credit for Multi-Agent Cooperation via Model-based Counterfactual Imagination, AAMAS, 2024.\n\n[4] Decentralized Transformers with Centralized Aggregation are Sample-Efficient Multi-Agent World Models, TMLR, 2025.\n\n[5] Masked Autoencoders Are Scalable Vision Learners, CVPR, 2022.\n\n[6] MA2E: Addressing Partial Observability in Multi-Agent Reinforcement Learning with Masked Auto-Encoder, ICLR, 2025.\n\n[7] INS: Interaction-aware Synthesis to Enhance Offline Multi-agent Reinforcement Learning, ICLR, 2025."}]

Author rebuttal:

Meta-review: This paper introduces M3W, a mixture-of-experts world model for multi-task multi-agent reinforcement learning. The core conceptual contribution is shifting modularity from policies to the world model, motivated by the idea of bounded similarity across task dynamics. The method combines a SoftMoE dynamics predictor and a SparseMoE reward predictor, with task-conditioned routing to enable selective sharing of knowledge across related tasks. Planning is performed using imagined rollouts from the world model rather than relying on an explicit policy.

The strengths are clear. Reviewers consistently recognized the novelty of applying MoE modularity within the world model, which represents a meaningful conceptual shift in MT-MARL. The empirical evaluation on Bi-DexHands and MA-MuJoCo is strong, with consistent improvements in return and sample efficiency compared to policy-centric baselines. The interpretability analyses (expert specialization, similarity matrices, visualizations) provide additional insight into the bounded similarity phenomenon. The rebuttal further strengthened the work by supplying ablations (heterogeneous vs. uniform MoE, number of experts, actor-only vs. planner-only), extensions to partial observability and discrete action spaces, and additional results on scalability and computational cost. These clarifications satisfied reviewers, three of whom raised their scores after rebuttal.

The main limitations are scope and practicality. Evaluation is limited to continuous-control multi-agent benchmarks, leaving open how well the approach translates to discrete or partially observable domains. While the authors provided preliminary extensions, further validation in diverse environments would strengthen claims of generality. The method also carries higher computational cost compared to policy-based baselines, though early-stopping mitigates this to some extent. Finally, the related work discussion, particularly on prior MoE applications in MTRL, was initially thin but will be expanded in the camera-ready version.

Overall, this is a technically solid and well-motivated contribution. The paper offers a novel perspective on modularity, demonstrates strong empirical results, and provides thoughtful analyses and extensions. The limitations are acknowledged and do not outweigh the contributions.
Decision: Accept (poster)