MemoryWAM: Efficient World Action Modeling with Persistent Memory

Sizhe Yang, Juncheng Mu, Tianming Wei, Chenhao Lu, Xiaofan Li, Linning Xu, Zhengrong Xue, Zhecheng Yuan, Dahua Lin, Jiangmiao Pang, Huazhe Xu · · 2026

Summary

Robust robotic manipulation in the real world requires not only an understanding of the current observation, but also memory and dynamics modeling.

Abstract Summary

Robust robotic manipulation in the real world requires not only an understanding of the current observation, but also memory and dynamics modeling. World action models (WAMs) possess these capabilities by jointly modeling visual foresight and actions conditioned on both current and historical observations, making them a promising paradigm for robotic manipulation. However, existing WAMs face a fundamental trade-off: methods with efficient inference typically condition only on a bounded window of recent observations and therefore struggle in non-Markovian environments, whereas methods that preserve long histories incur time and space costs that grow substantially with sequence length. To address this challenge, we introduce MemoryWAM, a world action model with efficient persistent memory. MemoryWAM uses a hybrid memory design that combines recent frames, event-boundary anchor frames, and compact gist tokens that summarize long-range history. A tailored attention mechanism enables retrieval of both detailed short-term context and compressed long-term context, supporting memory-dependent decision-making with reduced inference latency and GPU memory usage. Across long-horizon, memory-dependent manipulation tasks in both simulation and the real world, MemoryWAM outperforms strong vision-language-action (VLA) and WAM baselines while maintaining favorable computational efficiency.

Key Points

Proposes MemoryWAM for robotics.
Evaluated on real or simulated robotic tasks.
Primary contribution in ar-vr.

Method

Prior approaches in vision-language-action models (VLAs) [ 7 , 6 ] and world action models (WAMs) [ 39 , 62 ] typically map recent observations to actions, relying on a bounded temporal window $o_{t-N:t}$ and task instruction $l$: While effective for short-horizon tasks, they struggle in non-Markovian environments where decisions depend on long-range history.

A straightforward approach to retain long-term history is to preserve the full KV cache of all past observations within an autoregressive Transformer [ 33 , 60 , 50 ] .

Method — MemoryWAM: Efficient World Action Modeling with Persistent Memory — Figure 2: MemoryWAM adopts an MoT architecture with a video DiT and an action DiT. Video prediction provides dense supervision of dynamics modeling during training and is not required during inference. For persistent memory, MemoryWAM preserves tokens from initial anchor frames and recent frames, and compresses long-range history into a small set of gist tokens. This hybrid memory enables non-Markovian decision-making while maintaining low inference latency and GPU memory cost.

Experiments

MemoryWAM aims to equip world action models with efficient persistent memory, enabling end-to-end execution of long-horizon manipulation tasks with low inference cost. In this section, we evaluate MemoryWAM from three complementary perspectives: efficiency, policy performance, and design effectiveness.

We first describe the implementation details of MemoryWAM, including model architecture, training setup, and inference protocol in Sec. 4.1 . We then conduct a systematic study of memory mechanisms in Sec.

Experiments — MemoryWAM: Efficient World Action Modeling with Persistent Memory — Figure 4: Comparison of memory mechanisms. We compare full attention, test-time training (TTT), recurrent neural networks (RNNs), and our hybrid memory mechanism with respect to (a) single-pass inference latency and (b) GPU memory usage as functions of sequence length, and (c) success rates on the Press Button task. Latency and GPU memory usage are measured for a single layer.

Sources and Demonstrations

Figures are reproduced from Yang et al. See the full paper for experimental details and the project page, when the paper links one, for demonstrations and videos.