Title: Learning and Planning Multi-Agent Tasks via an MoE-based World Model
Venue: NeurIPS 2025
Accepted: yes
Keywords: Multi-task learning

Multi-agent reinforcement learning

Mixture of Experts

World model
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}.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS AND WEAKNESSES:\n* 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.\n\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-29"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\n# 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.\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\n- 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.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\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.\n\nQUESTIONS:\n1/ 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?\n\nLIMITATIONS:\nYes. 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.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\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.\n\nQUESTIONS:\n1.\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?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-06"}