MemoryWAM: Hybrid Memory in Robotics

• MemoryWAM is a hybrid robotic action model that integrates short-term frames, event-boundary anchors, and gist tokens to maintain long-range contextual information.
• It employs a Mixture-of-Transformers framework combining video and action diffusion transformers to achieve robust manipulation policies with sharply reduced latency and GPU memory.
• Experimental results in simulation and real-world tasks demonstrate that MemoryWAM outperforms traditional fixed-window and full-history models while supporting real-time 10 Hz control.

MemoryWAM is a world action model (WAM) for robotic manipulation that implements efficient persistent memory via a hybrid architecture inspired by human cognitive memory. It addresses the core challenge in robotic policy learning of maintaining long-horizon history and dynamics modeling without compromising real-time inference efficiency. By blending recent frames, event-boundary anchor frames, and compact “gist” tokens summarizing long-range context, MemoryWAM enables manipulation policies to condition on both detailed short-term context and compressed long-term information, achieving strong empirical performance in both simulation and real-world environments while sharply reducing inference latency and GPU memory consumption (Yang et al., 18 Jun 2026).

1. Motivation and Problem Setting

MemoryWAM is motivated by the limitations of previous WAMs in handling non-Markovian tasks in robotic manipulation. Many real-world problems—such as tracking occluded objects or executing actions dependent on information presented early in the episode—require policies to access information that may have occurred hundreds or thousands of steps prior. Fixed-window WAMs (e.g., FastWAM, Cosmos-Policy) provide constant inference cost by conditioning only on the latest $N$ observations, but fail when relevant cues disappear from the short-term window. Autoregressive full-history WAMs (e.g., LingBot-VA, DreamZero) maintain the entire observational history, resulting in $\mathcal{O}(N)$ scaling in both latency and GPU RAM as the sequence length $N$ grows, which becomes impractical for long episodes.

Cognitive studies suggest a memory system comprising short-term buffers, abstract long-term “gist” traces, and special encoding at event boundaries. MemoryWAM adopts this tri-level structure to break the trade-off: it provides efficient, scalable memory while preserving task-relevant context over arbitrarily long episodes (Yang et al., 18 Jun 2026).

2. Architecture and Hybrid Memory System

MemoryWAM is structured as a Mixture-of-Transformers (MoT) framework composed of a video diffusion transformer ($\Phi_v$) and an action diffusion transformer ($\Phi_a$):

• Video DiT ($\Phi_v$): Consumes VAE-encoded frame latents $z_t$, updating a key-value (KV) memory cache $\mathcal{C}_t^v$.
• Action DiT ($\Phi_a$): Predicts action sequences given noise-injected tokens $x_\tau^a$ and context from $\mathcal{O}(N)$0.

The hybrid memory cache consists of three disjoint segments:

1. Short-term window: Stores all tokens from the most recent $\mathcal{O}(N)$1 frames.
2. Event-boundary anchors: Retains the initial $\mathcal{O}(N)$2 frames (often corresponding to task onset or manually specified boundaries).
3. Gist tokens: For every frame that is neither recent nor an anchor, only $\mathcal{O}(N)$3 “gist” tokens are retained. These gist tokens are produced via a series of learnable queries at each DiT layer, attending to frame tokens and the prior cache—thus compressing both the frame and preceding history into $\mathcal{O}(N)$4 tokens per frame. With $\mathcal{O}(N)$5, $\mathcal{O}(N)$6, the effective compression ratio is $\mathcal{O}(N)$7.

Formally, the video-side cache at step $\mathcal{O}(N)$8 is

$\mathcal{O}(N)$9

with gist token updates given by

$N$0

where $N$1 are the $N$2 tokens for frame $N$3.

3. Attention Mechanism and Masking

MemoryWAM’s transformers utilize standard scaled dot-product attention but with a custom boolean mask enforcing the hybrid memory constraints:

• During video-side updates, new tokens are appended to the cache, with gist tokens attending to their associated frame’s tokens and all earlier memory, recursively absorbing long-range features.
• In the action prediction phase, action tokens can attend to:
  • All tokens in the most recent window,
  • All anchor frames (full tokens),
  • All gist tokens (summary for non-recent, non-anchor frames),
  • but not to full tokens of non-anchor, non-recent frames.

This masking ensures detailed access to local context, exact access at critical event boundaries, and compacted summaries for deep history. The mask is symmetric between training and inference, so no distributional shift arises.

4. Computational Efficiency

MemoryWAM’s hybrid memory sharply reduces both space and time complexity relative to prior full-history approaches:

Method	# Tokens	GPU RAM Growth	Latency (1,600 frames)
Full-history WAM	$N$4	$N$5 MB	$N$6 ms
Sliding-window WAM	$N$7 (const.)	High fixed/weak accuracy	Variable
MemoryWAM	$N$8	$N$9 MB	$\Phi_v$0 ms

Here $\Phi_v$1 is the sequence length, $\Phi_v$2 the number of per-frame tokens, $\Phi_v$3 the compression ratio, and $\Phi_v$4 the number of gist tokens per frame. Both GPU usage and per-layer attention cost scale as $\Phi_v$5, not $\Phi_v$6, enabling episodes with thousands of steps (Yang et al., 18 Jun 2026).

5. Experimental Results

MemoryWAM’s efficacy is demonstrated in both simulation and real-robot scenarios:

RMBench (Simulation)

• 9 dual-arm manipulation tasks, each with long-horizon memory dependencies.
• Trained on 50 demonstrations per task; evaluated over 100 test episodes.
• Success rates (mean over 9 tasks):

| Model | Success (%) | |---------------|-------------| | $\Phi_v$7 (no memory, VLA) | 10.4 | | FastWAM (short window) | 5.9 | | LingBot-VA (full history) | 78.2 | | MemoryWAM | 83.0 |

Press Button: MemoryWAM achieves 87% versus LingBot-VA’s 84%, with threefold lower inference latency.

Real Robot Experiments

• Shell Game (occluded object tracking): MemoryWAM 18/20, LingBot-VA 13/20, $\Phi_v$8 5/20.
• Look & Press (delayed counting): MemoryWAM 15/20, LingBot-VA 14/20, $\Phi_v$9 0/20.

MemoryWAM operates at a real-time 10 Hz control frequency ($\Phi_a$0 s per chunk), supporting closed-loop robot deployment.

6. Training and Inference Workflow

Training loop:

1. For each episode $\Phi_a$1:
   • Encode frames: $\Phi_a$2.
   • Construct an interleaved sequence of noisy frame latents and action chunks.
   • Apply the same hybrid memory mask in both $\Phi_a$3 and $\Phi_a$4 as will be used at inference.
   • Minimize the total flow-matching loss $\Phi_a$5.

Inference loop:

1. Initialize the cache with empty memory, populate anchor frames.
2. For each timestep $\Phi_a$6:
   • Observe $\Phi_a$7, encode $\Phi_a$8.
   • Update $\Phi_a$9 via $\Phi_v$0.
   • Maintain the short-term window; manage anchor set and gist cache.
   • Aggregate keys/values from anchors, window, and gist tokens.
   • Predict $\Phi_v$1, execute in the robot.
   • Collect new frame observations for the next loop.

7. Benefits, Limitations, and Future Directions

Benefits

• Long-range context retention: Gist tokens reliably preserve early task-relevant cues (e.g., initial object positions, instructions, transiently visible numbers).
• Computational efficiency: Both latency and GPU RAM usage scale as $\Phi_v$2, supporting arbitrarily long real-world episodes.
• Complementarity: Each memory component (anchors, window, gist) contributes unique value; ablations confirm performance degrades when any component is removed, most notably the gist.
• Real-time performance: Matches full-history models’ success rates (e.g., 87% vs. 84%) while using less than 20% of GPU RAM for long horizons.

Limitations

• Diffusion transformer backbone imparts inherited weaknesses in semantic reasoning and counting.
• Anchors are statically defined at task onset; dynamic event-boundary detection could enhance memory selectivity.
• Potential for future work in dual-system (symbolic + diffusion) reasoning, dynamic compression schedules, and large-scale joint video+action pretraining.

A plausible implication is that MemoryWAM’s hybrid memory paradigm may generalize to other domains that balance bounded computational resources with the need for persistent, context-rich memory in sequential decision-making (Yang et al., 18 Jun 2026).