Scaling Short-Term Memory of Visuomotor Policies for Long-Horizon Tasks

2026-06-14

AI Summary

Existing vision-based motor policies lack effective short-term memory mechanisms, hindering their performance on long-horizon tasks that involve temporarily occluded objects or require delayed responses. To address this limitation, this work proposes PRISM, an architecture that leverages a gated attention mechanism to filter irrelevant historical information and employs a hierarchical Transformer to efficiently compress and model long-range temporal dependencies, enabling the utilization of short-term memory spanning up to two minutes. Additionally, the authors introduce ReMemBench, the first general-purpose memory evaluation benchmark for embodied intelligence. Experimental results demonstrate that PRISM outperforms the strongest baseline by 5%–12% on ReMemBench and surpasses memory-free models as well as state-of-the-art methods such as GR00T-N1-3B and OpenVLA by 11%–15% on RoboCasa and LIBERO benchmarks.

Abstract

Many robotic tasks require short-term memory, whether it's retrieving an object that's no longer visible or turning off an appliance after a set period. Yet, most visuomotor policies trained via imitation learning rely only on immediate sensory input without using past experiences to guide decisions. We present PRISM, a transformer-based architecture for visuomotor policies to effectively use short-term memory via two key components: (i) gated attention, which filters retrieved information to suppress irrelevant details, improving performance by reducing the spurious correlations between the history and current action prediction, (ii) a hierarchical architecture that first compresses local information into compact tokens and then integrates them to capture temporally extended dependencies, improving its compute and memory footprint. Together, these mechanisms enable us to scale short-term memory in visuomotor policies for up to two minutes. To systematically evaluate memory in visuomotor control, we introduce ReMemBench -- a benchmark of eight diverse household manipulation tasks spanning four categories of short-term memory -- designed to foster general memory mechanisms rather than siloed, task-specific solutions. PRISM consistently outperforms prior works, including recurrent architectures, transformers, and their variants -- achieving an absolute improvement of 5%--12% over the strongest baseline. On the RoboCasa and LIBERO benchmarks, it achieves absolute improvements of 11%--15% over its no-memory variant and fine-tuned Vision-Language-Action baselines such as GR00T-N1-3B and OpenVLA, despite not leveraging any large-scale pretraining. Together, PRISM and ReMemBench establish a foundation for developing and evaluating short-term memory-augmented visuomotor policies that scale to long-horizon tasks. Additional materials are available at https://shahrutav.github.io/short-term-memory

Problem

• short-term memory
• visuomotor policies
• long-horizon tasks
• imitation learning
• memory-augmented control

Innovation

• gated attention
• hierarchical architecture
• short-term memory
• visuomotor policies
• transformer-based policy