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 MemWM, a memory-enhanced latent world model that augments JEPA-style prediction with a content-addressable memory bank. MemWM learns high-level representations directly from observations and uses 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 time steps. On PushT, TwoRoom, Reacher, and OGBCube, MemWM attains an average success rate of 89.9%, matching DINO-WM (89.3%) while using 4$\times$ fewer trainable parameters and no pretrained encoder, and improving over LeWM by 4.1%.