Source-linked AI summary
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
TL;DR
Long-horizon robotic manipulation requires memory and dynamics modeling, but existing WAMs trade efficient bounded context against costly full-history storage. MemoryWAM addresses this with hybrid persistent memory and outperforms strong VLA and WAM baselines in simulation and the real world while maintaining practical efficiency.
Problem
Efficient WAMs use bounded recent-observation windows that struggle with non-Markovian tasks, while full-history WAMs incur latency and memory costs that grow with sequence length.
Method
MemoryWAM combines sliding-window context, event-boundary anchor frames, compact gist tokens, and tailored attention for efficient persistent memory.
Results
MemoryWAM consistently outperforms strong VLA and WAM baselines on long-horizon, memory-dependent manipulation tasks in simulation and the real world.
Takeaways & Limitations
Hybrid persistent memory retains historical context while maintaining practical inference efficiency for memory-dependent manipulation.
Takeaways & Limitations
MemoryWAM inherits video diffusion models’ limited capacity for semantic understanding and reasoning.
Abstract
from arXiv · showhide
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.
1 Introduction
Existing approaches trade memory capacity against efficiency, motivating MemoryWAM’s hybrid persistent memory for long-horizon manipulation. MemoryWAM combines recent context, event-boundary anchors, and gist tokens, outperforming strong baselines while reducing computational costs.
- Motivation: Fixed-window WAMs are efficient but lack information beyond recent observations, whereas full-history WAMs preserve context at growing latency and memory costs.This trade-off is especially problematic for non-Markovian tasks.
- Method: MemoryWAM combines a sliding observation window, event-boundary anchor frames, and compact gist tokens to retain both detailed and long-range history.The design is inspired by complementary short-term, long-term, and event-boundary memory forms.
- Results: MemoryWAM achieves approximately 70 percentage points higher average success than current-observation or short-term-memory methods and outperforms LingBot-VA.The comparison covers RMBench long-horizon, memory-dependent manipulation tasks, with similar advantages in real-world experiments.
- Contributions: The proposed system integrates sliding-window context, gist tokens, and anchor frames while reducing GPU memory consumption and inference latency.The paper also studies trade-offs among latency, GPU cost, and policy performance.
- Results: MemoryWAM consistently outperforms strong VLA and WAM baselines on long-horizon, memory-dependent manipulation tasks in simulation and the real world.
2 Related Work
Related work spans policy-centric VLAs, dynamics-centric WAMs, and memory mechanisms for sequential modeling. These approaches differ in how they represent temporal structure, world evolution, and long-range context.
- Vision-Language-Action Models: VLAs transfer semantic priors from pretrained vision-language models to robotic manipulation and learn broad visuomotor skills from heterogeneous demonstrations.Most remain policy-centric by mapping observations directly to actions, with temporal structure and physical dynamics only implicitly represented.
- World Action Models: WAMs model world evolution together with robot actions, progressing from future visual-goal prediction toward unified video-action modeling.
- Memory Mechanisms for Sequential Modeling: Sequential-modeling memory mechanisms trade off memory capacity, update efficiency, and retrieval fidelity across recurrent, attention-based, and test-time-training designs.Related long-horizon visual systems face similar memory-efficiency trade-offs.
3 Method
MemoryWAM combines video-action modeling with hybrid persistent memory, retaining detailed recent and anchor-frame context while compressing long-range history into gist tokens. This design supports memory-dependent action generation while reducing sequence-length-dependent storage and attention costs.
- Architecture: The model couples a video DiT for dynamics-aware representations with an action DiT that denoises action tokens using cached video context.Video prediction supplies dense training supervision but is not required during inference; the video-side cache is updated from the current observation.
- Hybrid Memory: MemoryWAM uses a hybrid cache that combines recent observations, event-boundary anchor frames, and compressed long-term gist tokens for memory-dependent action generation.Recent frames preserve local interaction cues, anchor frames preserve task-onset information, and gist tokens summarize older frames.
- Short-term Memory: Short-term memory uses a fixed sliding window to preserve high-fidelity cues such as object motion, contact state, and hand-object configuration while bounding attention cost.The window targets immediate closed-loop control rather than complete historical coverage.
- Event-boundary Memory: Anchor frames retain full visual tokens from task onset because initial scene information may later be occluded or leave the recent observation window.These event-boundary frames complement recent context by preserving task-grounding information.
- Gist Memory: Gist tokens replace direct access to older frame tokens, reducing sequence-length-dependent storage and attention cost from O(N) to O(N/d).For L = 120 visual tokens and M = 8 gist tokens per frame, d = 15 and the long-term KV cache is reduced 15× versus full-history attention.
- Hybrid Memory: A unified attention interface lets action generation combine high-fidelity local context, preserved task-boundary information, and compact long-term history.The attention mask routes later tokens to recent and anchor frames directly while using gist tokens for non-recent, non-anchor frames.
4 Experiments
Experiments evaluate MemoryWAM’s efficiency, policy performance, and hybrid-memory design across simulation and real-world manipulation. The results show that hybrid memory preserves strong task performance while reducing computational costs and that its components provide complementary benefits.
- Comparison of Memory Mechanisms: At 1,600 frames, hybrid memory remains more efficient than both RNN- and TTT-based alternatives.
- Comparison of Memory Mechanisms: MemoryWAM achieves the same 87% Press Button success rate as full attention while being substantially more efficient.The comparison measures single-pass latency and GPU memory for one layer as sequence length grows.
- Simulation Experiments: On RMBench, bounded-window policies π0.5 and FastWAM achieve only 10.4% and 5.9% success, respectively, while MemoryWAM improves the average rate by 4.8 points over LingBot-VA.RMBench evaluates long-horizon, memory-dependent manipulation across nine dual-arm tasks, and Table 1 reports success over 100 rollouts.
- Real-World Experiments: MemoryWAM achieves the best performance on both real-world tasks with substantially lower latency and GPU memory cost than LingBot-VA.LingBot-VA’s high inference latency causes failure on cup swaps in the Shell Game.
- Effectiveness of Design Choices: Removing gist tokens causes the largest performance drop, while removing anchor frames or the sliding window also degrades performance.The ablations indicate complementary roles for long-term history, event-boundary information, and high-fidelity recent observations.
5 Conclusion
MemoryWAM provides efficient persistent memory for long-horizon robotic manipulation, outperforming competitive baselines across simulation and real-world tasks while retaining practical inference efficiency. Its main limitation is inherited from video diffusion models: limited semantic understanding and reasoning.
- Across memory-dependent manipulation tasks in both simulation and the real world, MemoryWAM outperforms competitive VLA and WAM baselines while achieving practical inference efficiency.
- MemoryWAM integrates a sliding observation window, preserved anchor frames, and compact gist tokens to retain historical context without prohibitive computational cost.
- MemoryWAM inherits video diffusion models’ limited capacity for semantic understanding and reasoning.
A.1 Implementation Details
Training uses bfloat16 mixed precision with distributed training, checkpointing, and gradient clipping, while inference maintains a hybrid memory KV cache with recent and persistent context.
- Training uses bfloat16 mixed precision, FSDP, activation checkpointing on every DiT block, and gradient clipping at 1.0.
- For fair comparison with Lingbot-VA, action history is incorporated into the action expert for Swap T, with RoboTwin pretraining used for Observe and Pick Up.
- Inference maintains per-transformer-block hybrid KV caches containing initial sink frames, four recent clean frames, and eight persistent context tokens per past frame.
A.2 Real-World Experiment Details
The real-world experiments use a dual-arm platform with RealSense D455 visual sensing and evaluate memory-dependent tasks designed to test tracking of occluded objects over time. Inference includes approximately 0.3 seconds of inter-chunk control latency.
- The experimental platform uses an ARX dual-arm robot with parallel grippers and a RealSense D455 camera for workspace RGB capture.
- Shell Game requires identifying and picking the cup covering a target cube after randomized cup swaps, testing tracking of occluded objects over time.
- The real-world control protocol includes approximately 0.3 seconds of latency between executed action chunks.