MemWM: Associative Memory-Enhanced Latent World Models for Long-Horizon Planning

Abstract

World models trained via latent dynamics offer a scalable path for embodied planning, yet existing methods suffer from representational collapse, computational overhead, and limited long-horizon capability. We introduce Memory World Model (MemWM), an associative memory-enhanced latent world model that augments JEPA-style prediction with a content-addressable memory bank. MemWM learns high-level historical representations directly from observations and leverages them to condition predictions via adaptive layer normalization, complementing low-level action-conditioned dynamics within a shared latent space. To prevent memory slots from collapsing into identical representations, we propose two complementary losses that encourage structural diversity across slots and balanced utilization across timesteps. On PushT, TwoRoom, Reacher, and OGBCube, MemWM attains an average success rate of 89.9%, matching DINO-WM (89.3%) while using roughly 4$\times$ fewer trainable parameters and no pretrained encoder, and improving over LeWM by 4.1%. Additionally, we extend the horizon to 50 timesteps, double the 25 timesteps of the original evaluation, demonstrating MemWM's robustness to long-horizon generalization.