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

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

1. Problem Statement

Robotic tasks often require short-term memory (spanning seconds to minutes) to handle partial observability, such as retrieving objects that are no longer visible or timing actions like turning off an appliance. However, most current visuomotor policies trained via imitation learning are myopic, relying solely on immediate sensory inputs without leveraging past experiences.

Extending standard transformer-based policies to include longer context windows to address this presents two critical challenges:

1. Spurious Correlations: Naively attending to long histories introduces irrelevant information (e.g., past distractor locations), leading models to learn false correlations that degrade performance when the context changes.
2. Computational Scalability: The computational and memory footprint of standard attention mechanisms scales linearly or quadratically with context length. Processing high-dimensional visual inputs over long temporal sequences (e.g., minutes) becomes prohibitively expensive.

2. Methodology: PRISM

The authors propose PRISM (Policy Architecture with Short-Term Memory), a transformer-based architecture designed to effectively utilize short-term memory through two key innovations:

• Gated Attention (Information Filtering): To mitigate spurious correlations, PRISM introduces a post-attention gating mechanism. This component selectively filters retrieved information by modulating the output of each attention block using a learnable gate conditioned on current input features. This allows the policy to adaptively suppress irrelevant details from the history, reducing noise in the action prediction.
• Hierarchical Architecture (Efficiency): To address scalability, PRISM employs a hierarchical design.
  • Local Encoding: A local memory encoder processes sensory inputs and compresses them into compact tokens using a Perceiver-style resampler.
  • Global Attention: A global memory encoder computes attention over these compressed tokens across time.
  • Fusion: The resulting global information is broadcast back and fused with the original local tokens via residual connections.
  • Benefit: This approach reduces the computational complexity from $O((n \cdot k_{tot})^2)$ to approximately $O(n \cdot k_{tot} \cdot m + (n \cdot k_{tot}/m)^2)$, where $m$ is the compression factor. It also significantly reduces the memory footprint at test time by caching only the compressed summary tokens.

PRISM is trained using standard behavior cloning (imitation learning) without auxiliary losses or external memory modules, mapping sequences of past sensory inputs and actions to future actions in a closed-loop setting.

3. Key Contributions

• PRISM Architecture: A novel policy architecture that combines gated attention and hierarchical summarization to scale short-term memory up to two minutes (at 2–3 Hz) while maintaining computational efficiency and robustness to noise.
• REMEMBENCH: A new benchmark comprising eight diverse household manipulation tasks spanning four categories of short-term memory derived from cognitive science:
  • Spatial: Recalling object locations (e.g., retrieving hidden fruit).
  • Prospective: Retaining intentions over a delay (e.g., cooking meat for a specific duration).
  • Object-Associative: Recalling object-location associations (e.g., returning a washed item to its specific original saucer).
  • Object-Set: Maintaining and updating sets of multiple objects (e.g., counting breadsticks moved to a microwave).
  • Unlike prior benchmarks that focus on single aspects or indirect memory tests, REMEMBENCH explicitly requires memory to succeed, as immediate sensory inputs are insufficient.

4. Experimental Results

The paper evaluates PRISM against recurrent architectures (LSTM, Mamba), segment-level recurrent transformers (Transformer-XL), and other transformer variants (Linear Attention, Scene Memory Transformer) on REMEMBENCH and standard benchmarks.

• REMEMBENCH Performance: PRISM achieves an average success rate of 41%, outperforming the strongest baselines (SMT at 36%, PTP at 28%, Mamba at 25%) by an absolute margin of 5%–12%.
  • Scaling: Performance scales with memory size; increasing the window from 1 to 512 timesteps improves success rates from 16% to 42% without saturation.
  • Ablation: Removing the gating mechanism causes a 16-point drop in performance, confirming its necessity for filtering irrelevant history.
• Standard Benchmarks (ROBOCASA & LIBERO):
  • On ROBOCASA, PRISM with memory (n=256) achieves 43% success, an 11% absolute improvement over its no-memory variant and a 11% improvement over the fine-tuned GR00T-N1-2B baseline.
  • On LIBERO, PRISM achieves 89% average success, a 15% improvement over the fine-tuned OpenVLA baseline and a 12% improvement over its no-memory variant.
  • The authors note that adding memory helps resolve action ambiguity in visually identical states (e.g., distinguishing "just grasped" from "about to place"), reducing multimodality even in tasks not explicitly designed to test memory.
• Real-World Evaluation: In a real-world adaptation of the "Wash and Return to Container" task, PRISM achieved 30% success, compared to 15% for a transformer baseline without memory and 0% for a policy with no memory.
• Efficiency: At a context length of 512 steps, PRISM reduces peak GPU memory usage by 80% and runtime by 40% compared to full global attention and hybrid sliding-window approaches.

5. Significance and Claims

The paper claims that PRISM and REMEMBENCH establish a foundation for developing and evaluating short-term memory–augmented visuomotor policies capable of handling long-horizon tasks.

• Generalization: By covering four distinct memory categories, REMEMBENCH encourages the development of general memory mechanisms rather than siloed, task-specific solutions.
• Scalability: PRISM demonstrates that short-term memory can be scaled to minutes of interaction without the prohibitive costs associated with naive context window extensions.
• Robustness: The gated attention mechanism proves essential for preventing the model from learning spurious correlations in noisy, long-horizon data.

The authors conclude that while current success rates on REMEMBENCH (around 30–40%) indicate that long-horizon tasks remain challenging, the consistent gains from memory augmentation across diverse architectures and benchmarks validate the approach. They identify future directions such as integrating persistent long-term memory and leveraging internet-scale pretraining for memory modules, but emphasize that their current work focuses strictly on the mechanics of short-term memory within the imitation learning framework.