Enhancing Latent World Models with Associative Memory

Yuwei Sun; Siyu Zhu

cs.LGOct 14, 2026

Enhancing Latent World Models with Associative Memory

Authors: Yuwei Sun, Siyu Zhu

Organizations: Shanghai Academy of AI for Science · Fudan University

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.

Figures & tables

Figure 1: The scheme of the memory-enhanced latent world models.

Model	Push-T (%)	Two-Room (%)	Reacher (%)	OGBCube (%)	Average (%)	Params (M)
PLDM [ 13 , 14 ]	78.0 $\pm$ 5.0	97.0 $\pm$ 0.94	78.0 $\pm$ 5.85	65.0 $\pm$ 2.49	79.5	20
LeWM [ 5 ]	96.0 $\pm$ 2.83	87.0 $\pm$ 1.79	86.0 $\pm$ 5.25	74.0 $\pm$ 2.49	85.8	18
DINO-WM [ 6 ]	92.0 $\pm$ 1.63	100	79.0 $\pm$ 5.85	86.0 $\pm$ 2.49	89.3	75
MemWM (Ours)	96.7 $\pm$ 0.94	90.7 $\pm$ 1.89	90.7 $\pm$ 0.94	81.3 $\pm$ 1.89	89.9	19

Table 1: Performance comparison of the conventional latent world model methods on the various planning tasks (Success Rate

\uparrow

Figure 2: Trajectory visualization of MemWM latent planning in the OGBCube dataset.

Figure 3: Visualization of memory slot usage. We observe distinct activation patterns across planning stages. In the PushT environment, the first slot exhibits peak activation when the agent starts to approach the T-shaped block, whereas the fifth slot dominates when the agent initiates pushing the block toward the target area. This phenomenon generalizes across various samples, exhibiting similar transitions during planning.

Model	Push-T (%)	Two-Room (%)	Reacher (%)	OGBCube (%)	Average (%)
MemWM	96.7 $\pm$ 0.94	90.7 $\pm$ 1.89	90.7 $\pm$ 0.94	81.3 $\pm$ 1.89	89.9
W/O temporal diversity loss	94.7 $\pm$ 2.49	90.0 $\pm$ 1.63	88.7 $\pm$ 1.89	73.3 $\pm$ 2.49	86.7
W/O memory loss	94.7 $\pm$ 0.94	89.3 $\pm$ 1.89	84.0 $\pm$ 3.27	66.0 $\pm$ 3.27	83.5
W/O shared params	90.7 $\pm$ 2.31	87.3 $\pm$ 3.06	84.7 $\pm$ 1.15	68.7 $\pm$ 2.31	82.9

Table 2: Performance comparison on the various ablations (Success Rate

\uparrow

), demonstrating the effectiveness of the proposed memory mechanism and memory loss objectives.

Hyperparameter		Push-T (%)	Two-Room (%)	Reacher (%)	OGBCube (%)	Average (%)
Number of Memory Slots	5	93.3 $\pm$ 0.94	89.3 $\pm$ 0.94	86.7 $\pm$ 1.89	76.7 $\pm$ 3.40	86.5
	10	96.7 $\pm$ 0.94	90.7 $\pm$ 1.89	90.7 $\pm$ 0.94	81.3 $\pm$ 1.89	89.9
	15	91.3 $\pm$ 0.94	89.3 $\pm$ 0.94	86.0 $\pm$ 1.63	77.3 $\pm$ 0.94	86.0
	20	94.0 $\pm$ 1.63	88.7 $\pm$ 0.94	85.3 $\pm$ 2.49	74.7 $\pm$ 0.94	85.7
Memory Retrieval Temperature	0.2	93.3 $\pm$ 0.94	90.0 $\pm$ 1.63	88.1 $\pm$ 5.85	74.7 $\pm$ 2.49	86.5
	0.5	96.7 $\pm$ 0.94	90.7 $\pm$ 1.89	90.7 $\pm$ 0.94	81.3 $\pm$ 1.89	89.9

Table 3: Model performance with varying memory hyperparameters (Success Rate

\uparrow

). We employ a memory slot number of 10 and a memory retrieval temperature of 0.5 by default, which result in the optimal prediction performance.

Method	Push-T (%)	Two-Room (%)	Reacher (%)	OGBCube (%)	Average (%)
Sequential (default)	96.7 $\pm$ 0.94	90.7 $\pm$ 1.89	90.7 $\pm$ 0.94	81.3 $\pm$ 1.89	89.9
Reverse	90.0 $\pm$ 1.63	87.3 $\pm$ 0.94	86.7 $\pm$ 5.25	69.3 $\pm$ 1.89	83.3
Separate	90.7 $\pm$ 2.31	87.3 $\pm$ 3.06	84.7 $\pm$ 1.15	68.7 $\pm$ 2.31	82.9

Table 4: Ablation on conditioning methods. We compare our default sequential design against the reverse and the separate conditioning alternatives.

Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

Training Configuration
Hyperparameter	Value
Batch Size	128
Max Epochs	100
Train Split	0.9
Gradient Clipping	1.0
Image Size	224

Appendix

Table 5: The detailed settings for the proposed method including the associative memory component, the latent world model component, and the training objectives.

Explore similar work

Beyond the Next Step: Variable-Length Latent World Models for Long-Horizon Planning
Jun 19, 2026Tianqi Du, Qi Zhang, Yifei Wang +1Latent Dynamics Modeling World Models
State Key Lab of General Artificial Intelligence, School of Intelligence Science and Technology, Peking University · Amazon AGI SF Lab · Institute for Artificial Intelligence, Peking University
Beyond a single latent space: a dual-latent world model for long-horizon planning
Sep 29, 2026Delin Zhao, Zhengrong Yue, Shaobin Zhuang +6World Model-Based Planning Action-Conditioned World Models
Nanjing University · Shenzhen University of Advanced Technology · Shanghai Jiao Tong University +1
Hierarchical Planning with Latent World Models
Apr 3, 2026Wancong Zhang, Basile Terver, Artem Zholus +7Hierarchical Planning Long-Horizon Robotic Manipulation
FAIR at Meta · New York University · Mila - Québec AI Institute +1

Jun 19, 2026J/K move · Enter open · S save

Beyond the Next Step: Variable-Length Latent World Models for Long-Horizon Planning

Tianqi Du, Qi Zhang, Yifei Wang, Yisen Wang

State Key Lab of General Artificial Intelligence, School of Intelligence Science and Technology, Peking University · Amazon AGI SF Lab · Institute for Artificial Intelligence, Peking University

Latent Dynamics Modeling World Models World Model-Based Planning Action-Conditioned World Models Long-Horizon Planning Model-Based Planning+1

Open paper AskPDF arXiv

Recently, world models have emerged as a promising paradigm for building intelligent agents by learning predictive models that estimate future environment states conditioned on observations and actions. In particular, JEPA-style latent world models provide an efficient alternative to pixel space prediction by learning action-conditioned dynamics in compact representation spaces. However, existing latent world models typically rely on one-step prediction and must be recursively rolled out for long-horizon planning, which leads to compounding errors and a mismatch between training objectives and downstream planning tasks. To address this limitation, we propose Variable-length Latent World Models (VLWMs), a framework that learns to predict future latent states conditioned on action sequences of variable lengths. Instead of training only on one-step transitions, VLWMs directly model temporally extended dynamics, allowing the same predictor to evaluate action plans over different horizons. We further introduce a curriculum training strategy that progressively expands the action horizon, stabilizing optimization from short-range dynamics to long-range prediction. At test time, we design planning methods tailored to VLWMs to better exploit their variable-length predictive capabilities. Experiments on long-horizon control tasks show that VLWMs significantly improve latent space world models, achieving 13% average improvement over the state-of-the-art LeWM across different datasets, with especially large gains on tasks requiring extended planning. These results suggest that VLWM provides a simple yet effective paradigm for improving long-horizon prediction and planning in latent world models.

Abstract

Explore similar work

Beyond the Next Step: Variable-Length Latent World Models for Long-Horizon Planning

Beyond a single latent space: a dual-latent world model for long-horizon planning

Hierarchical Planning with Latent World Models