Responses

Q1: The training cost and inference time of the world model compare to policy-based methods. (Weakness.1, Question.2)

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:

Note 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.

Q2: The generalizability of the world model to other scenarios is not thoroughly analyzed. (Weakness.2)

A2: The factors that hinder the generalization of our method to other scenarios include the full observation assumption and the continuous action space limitation.

• 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:

  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.

  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.

  • 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.

    	full	mask	decentralized
    Dynamics Error ↓	0.025±0.003	0.033±0.005	0.044±0.005
    Reward Error ↓	0.004±0.001	0.012±0.001	0.013±0.001
    Return ↑	32.1±0.4	31.3±0.6	31.1±0.3

• Continuous Action Space: There are two modifications to extend our method to discrete action spaces:

  1. Use one-hot actions or bitactions [6] to transform a discrete action into a continuous one.
  2. Replace the MPPI control with Cross Entropy Method (CEM) or Monte Carlo Tree Search (MCTS).

  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.

  CEM Procedure: Assume the action dimension is $n$,

  • S1: Initialize a discrete action distribution, $p^{(0)}=1/n$
  • S2: Sample $M$ action sequences, and compute their values according to (Eq.7, L185, main text)
  • S3: Select the top $m$ elite action sequences $a^∗$
  • S4: For each action dimension $j$, update the distribution $p^{(k+1)}(a[j])=\frac{1}{m} \sum \mathbb{I}(a[j]=a^*)$
  • Repeat steps S2–S4 until convergence.

Q3: Why the actor is not used as an explicit policy? What are the advantages and disadvantages? (Question.1)

A3:

1. 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.

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

   	actor+plan	plan-only	actor-only
   Return ↑	84.1 ± 0.8	81.4 ± 3.5	76.2 ± 2.1

3. 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.

These 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.

Summary

Thank you for your valuable review and constructive feedback. We appreciate your recognition of the motivation, technical soundness and strong performance of our work.

We also acknowledge the affirmations from other reviewers R1(DXnW), R2 (gkLA), and R3 (mmFi) on:

• 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).

• 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).

• 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).

• Interpretability: (R1: enhancing interpretability through expert specialization, visualizations strengthen the interpretability of M3W), (R3: interpretability analysis shows the mechanism of sharing experts).

We 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.

References

[1] TD-MPC2: Scalable, Robust World Models for Continuous Control, ICLR, 2024.

[2] Sparse Imagination for Efficient Visual World Model Planning, arXiv, 2025.

[3] Aligning Credit for Multi-Agent Cooperation via Model-based Counterfactual Imagination, AAMAS, 2024.

[4] Decentralized Transformers with Centralized Aggregation are Sample-Efficient Multi-Agent World Models, TMLR, 2025.

[5] Masked Autoencoders Are Scalable Vision Learners, CVPR, 2022.

[6] MA2E: Addressing Partial Observability in Multi-Agent Reinforcement Learning with Masked Auto-Encoder, ICLR, 2025.

[7] INS: Interaction-aware Synthesis to Enhance Offline Multi-agent Reinforcement Learning, ICLR, 2025.

Dear Reviewer k53v,

Thank you very much for your kind follow-up and for acknowledging our rebuttal. We truly appreciate your recognition of our efforts and your willingness to raise the rating. We would also like to reaffirm that all the clarifications and improvements in the rebuttal will be carefully incorporated into the camera-ready version.

As a quick note, we noticed that the review interface currently does not reflect an updated score ( it is still shown as 4 ). If convenient, could you kindly double-check whether the rating has been successfully submitted on your end (via the "Edit" button in the Official Review section)?

We sincerely appreciate your time and support throughout the review process.

Best regards,

Authors of Submission#6412

I appreciate the authors' effort in the rebuttal and they addressed my concerns and confusion. I believe if contents of rebuttal could be incorporated into the the paper, it could really help paint a better picture of the contribution and findings. Consequently, I will raise my rating.

Responses

Q1: Whether the improvements are due to the capacity of the model rather than the design choices? (Weakness.1, Limitation)

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}$.

Following 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.

Q2: Missing an analysis of the number of experts. (Weakness.3, Question)

A2: Following your constructive suggestion, we conducted an additional ablation study on the number of experts, and the results show:

• 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.

• 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.

Detailed numerical results are provided in the following table.

Q3: The demonstrated concept and tools are known and not original. (Weakness.2)

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:

• 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’).

• 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).

Q4: Insufficient discussion of MTRL and MoE-related literature, plus minor typos. (Weakness.4, Weakness.5)

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.

We 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.

The following are the relevant paragraphs that are expected to be added：

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.

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.

Summary

Thank you for your valuable review and constructive feedback. We appreciate your recognition of the motivation, technical soundness and strong performance of our work.

We also acknowledge the affirmations from other reviewers R1(DXnW), R2 (gkLA), and R4 (k53v) on:

• 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).

• 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).

• 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).

• Interpretability: (R1: enhancing interpretability through expert specialization, visualizations strengthen the interpretability of M3W).

We 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.

References

[1] Mixtures of Experts Unlock Parameter Scaling for Deep RL, ICML, 2024.

[2] Multi-task reinforcement learning with context-based representations, ICML, 2021.

[3] Paco: Parameter-compositional multi-task reinforcement learning, NeurIPS, 2022.

[4] Multi-task reinforcement learning with attention-based mixture of experts, IEEE RAL, 2023.

[5] Multi-Task Reinforcement Learning with Mixture of Orthogonal Experts, ICLR. 2024.

Thank you for your thoughtful follow-up and your valuable suggestions.

Regarding your question about why we reported prediction error instead of environment performance: the primary reason is that evaluating performance (i.e., return) requires running a complete training loop involving planning-based rollouts in the environment, which is significantly more time-consuming than computing prediction errors. Given the time constraints of the rebuttal phase, we chose prediction error as a more efficient proxy for model quality.

Importantly, as our method relies on model-based planning, the accuracy of dynamics and reward prediction directly impacts downstream performance. A more accurate world model enables more reliable rollouts, which in turn leads to better planning outcomes.

That said, we fully agree that including performance metrics would make the analysis more comprehensive. Therefore, we report the corresponding average returns in the table below.

We thank the reviewer for the constructive feedback and promise that, if the paper is accepted, we will include these additional experiments in the camera-ready version.

We would like to thank you again for your constructive feedback.

Following our previous commitment, we have now completed the additional experiments to evaluate the performance metrics (average return) corresponding to the number of experts. The new results, which have been added to our comment, confirm that the trends observed in prediction error indeed correlate with performance.

We hope this more complete analysis, which was prompted by your excellent feedback, addresses your concerns and provides a more thorough understanding of our work, and we welcome any further suggestions and disscussions. With this in mind, we would appreciate it if you would consider raising the score.

Responses

Q1: How can M3W adapt to decentralized or partially observable settings given its reliance on full observability? (Weakness.1 & Question.1)

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:

1. 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.

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.

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 and limited communication with minimal performance degradation.

Here are some formulation details of the communication cache.

At 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}\\}$.

Note 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.

Q2: Discrepancy between Eq.(2)’s MSE loss and the implementation’s two-hot cross-entropy. (Weakness.2 & Question.2)

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:

Responses

Q1: Does the concept of bounded similarity generalize to other environments with with different characteristics? (Weakness.1)

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.

• 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.

• 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).

• 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.

• For clarity, only the upper triangle of the symmetric MMD matrix is shown:

  SMAC	2s3z	3s5z	MMM	MMM2
  2s3z		0.15	0.58	0.69
  3s5z			0.58	0.69
  MMM				0.09
  MMM2

  MPE	adversary	push	reference	spread
  adversary		<0.01	0.40	0.38
  push			0.39	0.38
  reference				0.03
  spread

  RWARE	left-easy	left-hard	right-easy	right-hard
  left-easy		<0.01	0.13	0.12
  left-hard			0.14	0.13
  right-easy				<0.01
  right-hard

Q2: Could the proposed approach be extended to partially observable and discrete-action environments, and how? (Weakness.2 & Question.3)

A2: We start from two perspectives: partially observable and discrete actions.

• 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:

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

  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.

  	full	mask	decentralized
  Dynamics Error ↓	0.025±0.003	0.033±0.005	0.044±0.005
  Reward Error ↓	0.004±0.001	0.012±0.001	0.013±0.001
  Return ↑	32.1±0.4	31.3±0.6	31.1±0.3

• Discrete Action Space: There are two modifications to extend our method to discrete action spaces:

  1. Use one-hot actions or bitactions [6] to transform a discrete action into a continuous one.
  2. Replace the MPPI control with Cross Entropy Method (CEM) or Monte Carlo Tree Search (MCTS).

  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.

  CEM Procedure: Assume the action dimension is $n$,

  • S1: Initialize a discrete action distribution, $p^{(0)}=1/n$
  • S2: Sample $M$ action sequences, and compute their values according to (Eq.7, L185, main text)
  • S3: Select the top $m$ elite action sequences $a^∗$
  • S4: For each action dimension $j$, update the distribution $p^{(k+1)}(a[j])=\frac{1}{m} \sum \mathbb{I}(a[j]=a^*)$
  • Repeat steps S2–S4 until convergence.

Q3: Inference time cost of baselines. (Weakness.3)

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.

Q4: How would the performance and computational cost change if the heterogeneous MoE design is replaced with a uniform one? (Question.1)

A4: The heterogeneous design was chosen for two main reasons:

1. 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.
2. SparseMoE for reward prediction: SoftMoE is an n-to-n architecture, whereas reward prediction is an n-to-1 mapping, so we choose SparseMoE.

We 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.

Q5: Does the performance deteriorates as the number of agents increases? (Question.2)

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.

• 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.
• 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.
• Future Work: We plan to further incorporate permutation-invariance module or autoregressive mechanism to better handle tasks with a varying number of agents.

Summary

Thank you for your valuable review and constructive feedback. We appreciate your recognition of the motivation, novelty, strong performance, and interpretability of our work.

We also acknowledge the affirmations from other reviewers R2 (gkLA), R3 (mmFi), and R4 (k53v) on:

• 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).

• 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).

• 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).

• Interpretability: (R3: interpretability analysis shows the mechanism of sharing experts).

We 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.

References

[1] A kernel two-sample test, JMLR, 2012.

[2] Aligning Credit for Multi-Agent Cooperation via Model-based Counterfactual Imagination, AAMAS, 2024.

[3] Decentralized Transformers with Centralized Aggregation are Sample-Efficient Multi-Agent World Models, TMLR, 2025.

[4] Masked Autoencoders Are Scalable Vision Learners, CVPR, 2022.

[5] MA2E: Addressing Partial Observability in Multi-Agent Reinforcement Learning with Masked Auto-Encoder, ICLR, 2025.

[6] INS: Interaction-aware Synthesis to Enhance Offline Multi-agent Reinforcement Learning, ICLR, 2025.

[7] TD-MPC2: Scalable, Robust World Models for Continuous Control, ICLR, 2024.

[8] Sparse Imagination for Efficient Visual World Model Planning, arXiv, 2025.

[9] From Sparse to Soft Mixtures of Experts, ICLR, 2024.