Skip to content

Learning and Planning Multi-Agent Tasks via an MoE-based World Model

🕒 Published (v1): 2025-01-01 · Source: NeurIPS · Venue: NeurIPS 2025 · link

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

M3W applies Mixture-of-Experts (MoE) to a world model rather than to a policy network for multi-task multi-agent reinforcement learning (MT-MARL). It exploits "bounded similarity"—tasks sharing dynamics within groups but differing across groups—via SoftMoE for dynamics and SparseMoE for reward prediction, then plans directly on latent rollouts using MPPI without an explicit policy. This yields superior performance and sample efficiency across 14 Bi-DexHands and 24 MA-MuJoCo tasks.

Problem

Existing MT-MARL methods apply modularity (skill discovery, network decomposition, MoE) to the policy network. However, policy distributions across tasks are highly divergent even when underlying environment dynamics are similar within task groups. This misalignment prevents effective knowledge reuse and causes gradient conflicts, limiting generalization. No prior work applied MoE to multi-task world models in the multi-agent setting.

Method

M3W learns a latent-space world model conditioned on a learnable per-task embedding, comprising: - SoftMoE dynamics model: treats agent observation-action pairs as tokens; a soft router computes a weight matrix W ∈ ℝⁿˣᴺ that dispatches multi-agent tokens to N MLP experts and aggregates outputs back—enabling soft, differentiable routing that shares experts for similar tasks and isolates them for dissimilar ones. - SparseMoE reward predictor: a sparse router selects top-K self-attention experts (with load-balance loss) to predict scalar shared reward from the agent sequence. - Policy-free planning: at execution time, multi-agent MPPI evaluates action sequences by rolling out the world model in latent space; an auxiliary actor (trained with HASAC) generates candidate actions but is not the deployed policy. - The world model is trained with a self-supervised loss combining dynamics prediction error (MSE in latent space), soft cross-entropy reward loss, and Q-value loss; λ-weighted over rollout horizon H.

Key Contributions

  • Identifies "bounded similarity" in multi-task dynamics: strong intra-group similarity, weak inter-group similarity—distinct from policy distributions.
  • Proposes the first MoE-based world model for MT-MARL (M3W), applying SoftMoE to sequence-to-sequence dynamics and SparseMoE to sequence-to-scalar reward prediction.
  • Achieves planning without an explicit policy by using MPPI directly on world-model rollouts.
  • Demonstrates expert specialization and bounded-similarity exploitation via cosine-similarity of router outputs and t-SNE latent trajectory visualization.

Results

  • Bi-DexHands (14 tasks): M3W average normalized reward 84.1 ± 0.8 vs. MACD 72.0 ± 1.2, HMASD 72.3 ± 1.4, MAT 58.8 ± 1.6, MAPPO 63.6 ± 1.4, HATRPO 59.4 ± 1.6 (all after 50M steps).
  • MA-MuJoCo (24 tasks): M3W achieves highest aggregate performance (exact numbers in appendix); outperforms all baselines on cheetah-run variants (sample efficiency) and hopper tasks (stability post-convergence).
  • World model accuracy: M3W achieves lowest dynamics and reward prediction error among MLP, Transformer, Switch-Transformer, and M3W on six held-out Mujoco tasks.
  • Adding the same MPPI planner to MACD improves MACD performance but a large gap remains, indirectly validating M3W's superior world model quality.

Limitations

  • MPPI-based planner restricts M3W to continuous action spaces; discrete-action tasks are unsupported.
  • Assumes agents can access others' observations/actions (single communication round); not applicable to fully decentralized, no-communication settings.
  • Task identity is provided as a one-hot label; zero-shot generalization to unseen tasks is not addressed.
  • Evaluation is confined to two benchmarks; scalability to very large task sets or heterogeneous agent numbers is explored only partially in appendices.

Relevance to Agentic AI / LLM Agents

M3W demonstrates that MoE-based world models—rather than MoE-based policies—are a more natural locus for modularity in multi-task cooperative agents, with direct implications for designing scalable agentic systems that must handle diverse task distributions without catastrophic interference. The policy-free MPPI planning paradigm resonates with emerging agentic frameworks where an LLM or world model generates rollouts for deliberative planning rather than executing a fixed policy. The "bounded similarity" insight also has architectural relevance for multi-agent LLM systems where tool-use or sub-task dynamics cluster into groups (e.g., code execution vs. web retrieval), suggesting expert routing over shared world representations as a design principle. Finally, the interpretability results (expert specialization, router cosine similarity blocks) provide a blueprint for understanding routing decisions in MoE-based agent architectures.