Composition of Memory Experts for Diffusion World Models
For reinforcement learning agents that need to plan and make decisions based on past observations, this work provides a scalable and accurate world model that preserves both local detail and long-term memory.
The paper introduces a diffusion-based world model that decouples future-past consistency into three specialized memory experts (short-term, long-term, and spatial long-term), overcoming the memory trade-off between local detail and long-range fidelity. The method improves temporal consistency, recall, and navigation performance across simulated and real-world benchmarks without quadratic cost.
World models aim to predict plausible futures consistent with past observations, a capability central to planning and decision-making in reinforcement learning. Yet, existing architectures face a fundamental memory trade-off: transformers preserve local detail but are bottlenecked by quadratic attention, while recurrent and state-space models scale more efficiently but compress history at the cost of fidelity. To overcome this trade-off, we suggest decoupling future-past consistency from any single architecture and instead leveraging a set of specialized experts. We introduce a diffusion-based framework that integrates heterogeneous memory models through a contrastive product-of-experts formulation. Our approach instantiates three complementary roles: a short-term memory expert that captures fine local dynamics, a long-term memory expert that stores episodic history in external diffusion weights via lightweight test-time finetuning, and a spatial long-term memory expert that enforces geometric and spatial coherence. This compositional design avoids mode collapse and scales to long contexts without incurring a quadratic cost. Across simulated and real-world benchmarks, our method improves temporal consistency, recall of past observations, and navigation performance, establishing a novel paradigm for building and operating memory-augmented diffusion world models.