VPWEM: Non-Markovian Visuomotor Policy with Working and Episodic Memory

This paper introduces VPWEM, a non-Markovian visuomotor policy that combines a sliding window of recent observations with a Transformer-based episodic memory compressor to efficiently retain long-term context for robotic control, achieving significant performance improvements over state-of-the-art baselines on memory-intensive manipulation tasks while maintaining constant computational costs.

The Gist

Imagine you are teaching a robot to play a game of "Shell Game." You hide a ball under one of three cups, shuffle them around, and then ask the robot to find the ball.

If the robot only looks at the cups right now, it has no idea where the ball is. It's like trying to guess the ending of a movie by only looking at the final frame. This is the problem most current robot "brains" face: they have very short memories. They can only remember the last few seconds of what they saw. If a task requires remembering something that happened a minute ago (like which cup the ball started under), these robots fail miserably.

This paper introduces VPWEM, a new way to teach robots that gives them a "super memory" without making their brains too heavy or slow.

Here is how it works, using simple analogies:

1. The Problem: The "Short Attention Span" Robot

Most robots today are like students who can only focus on the teacher's last sentence. If the teacher says, "Pick up the red block," the robot does it. But if the teacher says, "Pick up the red block, but only if you saw me hide it under the blue box three minutes ago," the robot gets confused. It forgets the blue box.

To fix this, engineers used to just tell the robot to "remember everything." But this is like trying to carry a library in your backpack. It gets too heavy (too much computer power needed), too slow, and the robot gets overwhelmed by irrelevant details (like the color of the wall, which doesn't matter).

2. The Solution: The "Two-Brain" System

VPWEM solves this by giving the robot two types of memory, inspired by how human brains work:

• Working Memory (The Sticky Note): This is the robot's short-term memory. It keeps the last few seconds of video and sensor data right in front of its eyes. It's like a sticky note on your computer screen with the immediate instructions.
• Episodic Memory (The Diary): This is the long-term memory. Instead of keeping a video of the entire past hour (which is huge), the robot uses a special "Memory Compressor."

3. The Magic Ingredient: The "Memory Compressor"

Think of the Contextual Memory Compressor as a brilliant editor or a librarian.

• How it works: As the robot moves through the world, it sees thousands of images. Every time an image moves out of its "Sticky Note" (Working Memory), the Compressor takes that image and asks: "Is this important?"
• The Compression: If the robot saw a red ball being hidden, the Compressor doesn't save the whole video of the ball moving. It writes a tiny, perfect summary sentence in a "Diary" (Episodic Memory) that says: "Red ball hidden under left cup at 10:00 AM."
• The Result: The robot doesn't need to carry a 1-hour video file. It just carries a few pages of "Diary entries" that summarize the whole history. This keeps the robot's brain light and fast, but it still knows the whole story.

4. How the Robot Uses It

When the robot needs to make a move (like reaching for a cup), it looks at two things:

1. The Sticky Note: What is happening right now?
2. The Diary: What happened earlier that matters?

It combines these two to make a smart decision. It's like a detective solving a crime: they look at the crime scene (Working Memory) and check their case file for clues found hours ago (Episodic Memory).

Why This Matters

The authors tested this on robots doing tricky tasks, like:

• The Shell Game: Remembering where a ball was hidden after it was covered.
• Mobile Manipulation: A robot arm on a moving cart that has to clean a table, remembering where the dirty dishes were placed earlier.

The Results:

• On tasks requiring long-term memory, this new method beat the best existing robots by 20%.
• It didn't make the robot slower or require a supercomputer. It actually ran efficiently because it didn't waste space on useless memories.

The Bottom Line

VPWEM is like giving a robot a photographic memory that knows how to summarize. Instead of trying to remember every single second of its life (which is impossible), it learns to remember the story of what happened. This allows robots to finally tackle complex, real-world tasks that require patience, planning, and remembering the past.

Technical Summary

Here is a detailed technical summary of the paper "VPWEM: Non-Markovian Visuomotor Policy with Working and Episodic Memory."

1. Problem Statement

Current imitation learning policies for robotics typically rely on Markovian assumptions or short-context histories (e.g., the last 2–10 observations). While effective for simple tasks, these approaches fail in non-Markovian environments where:

• Long-term dependencies are required (e.g., remembering an object's initial position after it has been occluded).
• Sensor limitations or environmental stochasticity make the current observation insufficient for decision-making.
• Simply increasing the context window leads to computational intractability (quadratic $O(L^2)$ complexity in Transformers) and overfitting to spurious correlations (the "copycat" problem), causing catastrophic failures under distribution shifts.

The core challenge is how to equip robot policies with long-term memory that is computationally efficient, bounded in size, and capable of filtering irrelevant historical noise, mimicking human cognitive processes (working memory vs. long-term episodic memory).

2. Methodology: VPWEM Framework

The authors propose VPWEM (Visuomotor Policy with Working and Episodic Memory), a framework that augments diffusion-based policies with a dual-memory system inspired by the human hippocampus.

A. Dual-Memory Architecture

1. Working Memory (Short-term):

   • Maintains a sliding window of recent observation tokens ($L$ steps).
   • Acts as a First-In-First-Out (FIFO) buffer.
   • Provides immediate context for the policy.

2. Episodic Memory (Long-term):

   • Designed to store compressed information from observations that have "fallen out" of the working memory window.
   • Contextual Memory Compressor: A Transformer-based module that recursively processes out-of-window observation tokens.
     • Mechanism: It uses self-attention over a cache of past summary tokens and cross-attention over a cache of historical observations.
     • Output: It converts a potentially infinite history of observations into a fixed number ($M$) of episodic memory tokens.
     • Training: The compressor is trained jointly with the policy to learn to filter out nuisance variables and retain only task-relevant information.

B. Action Generation (Diffusion Policy)

The framework is instantiated on Diffusion Policies:

• Conditioning: The action generation process is conditioned on both the Working Memory ($w_t$) and the Episodic Memory ($e_\tau$).
• Process:
  1. Encoding: New frames are encoded; out-of-window frames are pushed to the observation cache.
  2. Compression: The compressor updates the episodic memory tokens based on the new history.
  3. Denoising: A noise predictor (Transformer decoder) iteratively denoises a random action sequence, guided by the combined memory tokens, to produce the final action chunk.

C. Training and Inference Efficiency

• Training: Uses a behavior cloning loss. To handle varying trajectory lengths, samples are grouped by length, and the compressor input is constructed via random subsampling of the history to ensure robustness. Gradients are detached from the cache management to prevent backpropagation through time, reducing memory costs.
• Inference: Operates with nearly constant memory and computation per step, regardless of the total episode length, because the episodic memory size is fixed.

3. Key Contributions

1. Novel Framework: Proposes a Transformer-based Contextual Memory Compressor that recursively condenses infinite historical tokens into a fixed-size episodic memory, serving as a dynamic summary of the full trajectory.
2. Instantiation on Diffusion Policies: Successfully integrates this memory mechanism into state-of-the-art diffusion policies (specifically DP and MaIL), redesigning training pipelines to condition action generation on both short-term (working) and long-term (episodic) contexts.
3. Empirical Validation: Demonstrates that explicit memory modules outperform simply scaling context windows, achieving superior performance in memory-intensive tasks while maintaining efficiency.

4. Experimental Results

The authors evaluated VPWEM on three benchmarks: MIKASA (memory-intensive manipulation), MoMaRT (mobile manipulation), and Robomimic (mostly Markovian).

• MIKASA (Memory-Intensive Tasks):
  • VPWEM outperformed state-of-the-art baselines (including Diffusion Policies, VLA models like OpenVLA, and MemoryVLA) by >20% on average.
  • Tasks like ShellGameTouch (hiding an object) and RememberColor required long-term recall, where standard short-context policies failed.
• MoMaRT (Mobile Manipulation):
  • Achieved an average 5% improvement over baselines (DP and MaIL) on long-horizon mobile tasks.
  • Efficiency: Compared to increasing the context window of standard Diffusion Policies (which drastically increases GPU memory and inference time), VPWEM added only ~2.24M parameters and maintained low inference latency while achieving a higher success rate (58.3% vs. ~49% for long-context baselines).
• Robomimic (Markovian Tasks):
  • Performed on par with baselines, proving that the memory mechanism does not degrade performance in tasks where long-term memory is unnecessary.

5. Significance and Impact

• Solving Non-Markovian Control: VPWEM addresses a critical bottleneck in robotic learning: the inability of current policies to handle tasks requiring long-term temporal reasoning without incurring prohibitive computational costs.
• Biological Inspiration: By mimicking the human brain's separation of working memory and episodic memory (via the hippocampus), the method offers a scalable solution for lifelong learning and complex task execution.
• Efficiency vs. Performance Trade-off: The paper demonstrates that intelligent compression of history is superior to brute-force expansion of context windows. It enables robots to "remember" the entire episode history with constant memory overhead, making real-time deployment of long-horizon policies feasible.
• Generalizability: The approach is orthogonal to existing policy improvements (like better vision encoders) and can be applied to various base architectures (Diffusion, Mamba, etc.).