Skip to content

Revisiting Multi-Agent World Modeling from a Diffusion-Inspired Perspective

🕒 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

DIMA reformulates multi-agent world modeling as a conditional denoising process by mapping sequential agent action revelation to the reverse diffusion process, reducing modeling complexity from exponential to linear in the number of agents. The resulting centralized world model achieves state-of-the-art sample efficiency and final returns on MAMuJoCo and Bi-DexHands benchmarks, outperforming both model-free and model-based baselines.

Problem

Existing multi-agent world models face a fundamental tradeoff: centralized approaches scale exponentially with the joint action space (|A|^n), while decentralized approaches model local observation dynamics that misalign with the true global Dec-POMDP transition P(s_{t+1}|s_t, a^{1:n}_t), requiring unsupervised communication/aggregation modules that introduce additional error. Neither adequately captures inter-agent dependencies without prohibitive cost or structural mismatch.

Method

DIMA reframes joint state-transition prediction as a conditional denoising diffusion process. The key insight is that progressively conditioning on each agent's action (one agent per denoising step, in some ordering) mirrors the reverse diffusion process: each revealed action reduces uncertainty about the next global state, analogous to iterative denoising from noise to clean sample.

Formally, DIMA defines a diffusion-inspired decomposition (Assumption 1) where P(s_{t+1}|s_t, a^{1:n}t) factorizes as a product of single-agent conditional denoising steps p_θ(s^{(k-1)}, a^k_t, s_t). The resulting ELBO yields a per-agent denoising score matching objective. Two practical properties are enforced: (i) } | s^{(k)}_{t+1permutation invariance via averaging over random agent orderings during training; (ii) condition-independent noising, allowing continuous-time noise schedules. The denoiser F_θ follows the EDM preconditioner framework (Karras et al.), conditioned on a sliding window of past global states and joint actions. A Transformer-based reward/termination model and a VQ-VAE observation decoder complete the world model. Policy learning uses MAPPO with actor-critic (CTDE), trained entirely on imagined rollouts.

Key Contributions

  • Novel theoretical framing: derivation of an ELBO (Theorem 2) that decomposes multi-agent dynamics prediction into per-agent conditional denoising steps, grounding the diffusion analogy formally.
  • DIMA architecture: centralized world model with O(n · |S|) complexity per denoising step vs. O(|A|^n · |S|) for joint conditioning baselines.
  • Permutation invariance achieved through expectation over random conditioning orders during training (Eq. 9), with empirical validation across ascending/descending/random orderings.
  • State-of-the-art on MAMuJoCo (7 configurations) and Bi-DexHands (4 tasks) in low-data regimes (1M and 300k real steps respectively).
  • Demonstrated superior long-horizon OOD generalization (H=25 rollouts, trained at H=15) with lower compounding error than MARIE and MAMBA.

Results

  • MAMuJoCo: DIMA consistently outperforms MAMBA, MARIE, MAPPO, HAPPO, and HASAC across all 7 agent-partitioning settings (HalfCheetah-2x3/3x2/6x1, Walker-2x3/3x2, Ant-2x4/4x2); qualitative trajectory reconstructions show DIMA maintains coherence at H=15 where MARIE and MAMBA degrade noticeably.
  • Bi-DexHands: DIMA outperforms MAMBA, MAPPO, HAPPO, HASAC on all 4 tasks (ShadowHandPen, DoorOpenInward, DoorOpenOutward, BottleCap); MARIE excluded due to OOM.
  • Accumulated L1 errors on Ant-2x4 (100 trajectory segments):
  • Obs @ H=15: DIMA 4.32±1.44 vs. MARIE 6.62±2.68, MAMBA 6.85±2.51
  • Obs @ H=25: DIMA 15.01±9.88 vs. MARIE 21.64±17.51, MAMBA 38.90±24.48
  • Rew @ H=15: DIMA 2.42±0.93 vs. MARIE 3.54±1.60, MAMBA 3.98±1.54
  • Ablation (ShadowHandBottleCap, 8 runs): Sequential (DIMA) at 100K steps: 251.8±17.3 vs. Joint 234.1±20.6; advantage narrows at 300K steps (249.2±10.7 vs. 255.2±7.0).
  • Ablation (DoorOpenOutward/Inward @ 300K): Sequential 352.4±40.5 / 290.3±30.4 vs. Joint 302.5±76.9 / 235.1±68.1 — sequential advantage persists on harder tasks.

Limitations

  • Scalability to large-scale systems (hundreds of agents) is not demonstrated; the authors acknowledge potential challenges and propose grouping techniques as future work.
  • The permutation invariance property holds empirically up to approximately H=10; visible drift between conditioning orders appears at t=12 in visualizations.
  • Sequential denoising requires n forward passes per state prediction step, which may incur higher inference latency than single-pass joint models despite lower theoretical complexity.
  • Evaluated only in fully cooperative, continuous-control settings; generalization to competitive or discrete-action MARL is not established.
  • MARIE was excluded from Bi-DexHands comparisons due to OOM, limiting direct comparison on that benchmark.

Relevance to Agentic AI / LLM Agents

DIMA establishes a principled method for model-based planning in multi-agent settings—central to agentic systems where multiple AI agents must coordinate in a shared environment. The sequential, agent-wise conditioning framework is directly relevant to research on multi-agent coordination, as it provides a scalable substrate for imagined rollouts without sacrificing global consistency. The diffusion-model backbone parallels trends in using generative models as world models for downstream planning, a paradigm increasingly explored in embodied and tool-using LLM agents. For researchers building agentic systems that require learned environment models (e.g., for lookahead search, synthetic data generation, or sim-to-real transfer), DIMA offers both a theoretical blueprint and an open-source implementation.