Latent world models improved with memory for better planning
Enhancing Latent World Models with Associative Memory
Machine LearningComputer Vision and Pattern Recognition
Summary
Planning how to act in a changing environment can be tricky for AI because it needs to remember details over time and think ahead. The authors introduce a way to help AI models keep better long-term memories by adding a special memory system that matches pieces of information by content. This helps the models make more accurate predictions over longer times and work faster with fewer resources. They tested this approach in different virtual tasks and found it performed better than older methods. The new method also avoids common problems like collapsing representations and inefficient memory use.
What this means in practice
- •For robotics teams: Enable robots to improve planning and decision-making over long sequences of actions using efficient memory-enhanced models.
- •For game developers: Enhance AI characters’ ability to predict and plan across many steps in complex game environments through better memory integration.
Authors
Yuwei Sun, Siyu Zhu
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 a memory-enhanced latent world model that augments JEPA-style prediction with a content-addressable memory bank. Our approach learns high-level representations directly from observations, operating at multiple temporal scales within a shared latent space, and conditions predictions via adaptive layer normalization. To prevent collapse and ensure effective memory usage, we design two complementary losses that encourage slot diversity and balanced temporal utilization. Experiments on PushT, Tworoom, Reacher, and OGBCube demonstrate that our method achieves state-of-the-art performance with a lightweight architecture, outperforming conventional latent world models on long-horizon planning tasks.