MemBodied: Recurrent Associative Memory for Vision-Language-Action Models

AI summary

Overview

Research area: Robotics / embodied AI — specifically episodic memory for Vision-Language-Action (VLA) robot policies.

Technical level: Advanced. The paper assumes familiarity with VLA backbones (π₀, π₀.₅, OpenVLA), flow-matching action experts, and fast-weight / associative memory architectures such as DeltaNet and Gated DeltaNet. The core ideas, however, are describable without the linear-algebra notation.

Scope: The paper proposes MemBodied, a fixed-size two-pathway episodic memory (a recurrent associative state plus a fixed initial-scene anchor) for VLA control, and evaluates it on five memory-dependent RMBench tasks, the four LIBERO suites, and three real-robot tasks.

Note: the supplied paper content is truncated partway through Appendix B.1, so appendix-level details beyond that point are not reflected here.

What This Paper Is About

Most VLA policies predict an action chunk from only the current observation, robot state, and language instruction, discarding what happened earlier in the episode. This fails on history-dependent manipulation: if the robot moves a block and later must return it, the current image may no longer reveal the correct destination — a temporal state aliasing problem where similar observations demand different actions. Keeping past observations in context recovers this information but makes context length, memory consumption, and inference cost grow throughout the episode. The goal is a memory that is updated online, stays fixed in size, and retrieves information according to the policy's current state.

Key Contributions

1. A fixed-capacity episodic memory for VLA control. MemBodied combines (a) a recurrent associative state, stored as one matrix per action-network layer, updated with a gated delta rule after each executed action, and (b) a fixed episode anchor that preserves a compact representation of the initial scene via rank-r_A cross-attention. Both have storage independent of episode length.

2. A memory readout interface comparison. The paper introduces memory-token injection (the retrieved vector is projected and added to a dedicated contextual token before self-attention) and compares it against attention steering (additive corrections to attention query/output) and hierarchical memory (an LSTM-like cell updated alongside the associative matrices).

3. A causal delayed-write schedule with sequence-level training. At each step the policy reads memory and acts; the write happens only after the resulting observation arrives, so the stored value includes the observed consequence of an action without leaking that future observation to the action that caused it. The memory parameters are learned jointly through the policy's native action objective, with no separate memory-prediction loss.

4. An evaluation spanning memory-dependent and fully observable control. Five RMBench tasks, four LIBERO suites, three real-robot tasks, two backbones (π₀ and π₀.₅), plus parameter, latency, and GPU-memory comparisons against NativeMEM.

Main Findings

• RMBench mean success: MemBodied reaches 50.0% mean success across the five evaluated RMBench tasks, exceeding π₀-Stateless (6.4%) by 43.6 percentage points and π₀-FrameStack (14.8%) by 35.2 points. Its advantage over the stateless policy is positive on all five tasks.

• Relative gains reported in the abstract and introduction: 7.81× the mean success rate of a stateless π₀ policy, 2.98× that of vanilla recurrent memory (π₀-Vanilla Recurrent Memory at 16.8%, i.e. 33.2 points below MemBodied), and 1.30× the strongest memory-augmented baseline (NativeMEM at 38.4%) with 10× fewer added parameters.

• Comparison with compressed history: Standalone NativeMEM reaches 38.4% mean success versus 50.0% for standalone MemBodied — an 11.6 percentage point gap on average. Adding MemBodied to NativeMEM improves NativeMEM to 45.2%, still 4.8% below standalone MemBodied. No single method dominates at task level: NativeMEM is stronger on Put Back Block and Swap Blocks; MemBodied is stronger on Rearrange Blocks, Battery Try, and Block Ranking.

• Transfer to a second backbone: With π₀.₅, MemBodied reaches 48.0% mean success versus 12.4% for the π₀.₅ baseline, a 3.87× performance increase over the stateless policy, improving on all five tasks. The largest gain is Rearrange Blocks (13.0% → 94.0%); Put Back Block improves from 11.0% to 38.0%.

• Real-robot results: Across three real-robot tasks, MemBodied increases mean success from 3.33% to 26.67% — an 8.0× improvement — with improvements on every task.

• Inference efficiency: MemBodied achieves 91.9% lower inference latency and uses 9.5% less peak GPU memory than NativeMEM. It adds 40M parameters (1.26% of π₀), roughly one-tenth of NativeMEM's 415M (12.81% of π₀).

• Fully observable manipulation (LIBERO): MemBodied reaches 95.1% four-suite mean success versus π₀'s 94.2%. Spatial is 96.2 (−0.6), Object 97.8 (−1.0), Goal 95.8 (unchanged), and Long 90.6%, exceeding π₀'s 85.2% by 5.4 percentage points. The improvement is skewed toward the long-horizon suite.

• Readout interface ablation: Among anchor-free variants, memory-token injection reaches 37.6% mean success, versus 23.2% for hierarchical memory and 20.8% for attention steering.

• Episode anchor ablation: The full two-pathway memory reaches 50.0%, versus 37.6% without the anchor. The anchor improves four of five tasks — the largest gain is Swap Blocks (16.0% → 56.0%) — but reduces Block Ranking (32.0% → 22.0%), a task that depends more on tracking evolving progress than recalling the initial scene.

• Write-value composition: Combining the subsequent visual observation with the action summary achieves 46.7% mean success versus 42.7% for vision-only and 34.7% for action-only memory. On Battery Try the combined value reaches 30.0%, versus 24.0% for vision alone and 22.0% for action alone.

• Qualitative analysis: Manual annotation of 50 seed-matched rollout pairs per task shows MemBodied and π₀-Stateless behave similarly when the current observation suffices, but diverge when past information is required. On Put Back Block, MemBodied reduces wrong-pad returns from 37/50 to 13/50, while both policies complete the preceding block movement and button press in every rollout.

Methodology in Plain English

The researchers treat manipulation as a history-dependent sequential decision problem and maintain an episode memory ℰₜ = (ℳₜ, A) that persists across policy calls.

The associative state (ℳₜ). Each action-network layer keeps its own matrix. After the robot executes an action chunk and the environment returns the next observation, the model builds a "write value" that fuses two things: a pooled encoding of the observed visual consequence (each camera's patch tokens are queried by the projected robot state and averaged across cameras) and a summary of the action chunk that produced it. This value is written into the matrices using a gated delta rule — a retention gate decides how much of the old state survives, and a write gate controls a correction that compares the new value against what is already associated with the write key. Writes update existing matrices rather than appending observations, so state size never grows. Memory rank was set to r = 128.

The anchor (A). Because repeated delta updates can overwrite fine-grained early details, a separate pathway keeps a compact reference to the first observation: frozen vision tokens are average-pooled from a 16×16 patch grid down to a 4×4 grid per camera and concatenated across cameras. At each call, the current observation queries this anchor through rank-r_A cross-attention (r_A = 64), and the result is projected, repeated across the action horizon, and used to condition the action representation. The anchor does not lengthen the policy prefix or change token positions or attention masks.

Reading memory. The reading query comes from the state-token representation entering each layer. The retrieved vectors are projected to the action-network width, scaled by α_mem/r, gated by a scalar derived from the current state, and added to a dedicated memory token that participates in self-attention. This memory token is re-inserted fresh at every policy call; the recurrent information lives in the matrices, not in a carried token.

Causal scheduling and training. Read → act → execute → write. At inference the policy caches preceding layer states and the sampled action chunk from the final denoising step and performs the write when the next observation arrives, so training and inference share the same causal interpretation. Because the write is delayed, training is sequence-level: each sample holds N observation-action pairs separated by one action horizon of H environment steps. Image encoding and policy-input construction are parallelised across batch and sequence, while memory state is propagated sequentially, and gradients flow through the full sequence so later action losses optimise earlier memory operations.

Evaluation setup. RMBench was used with 50 rollouts per task across three M(1) tasks (put_back_block, rearrange_blocks, swap_blocks) and two M(n) tasks (battery_try, block_ranking_try). LIBERO used four suites of ten tasks each, 50 rollouts per task — 500 per suite and 2,000 rollouts per checkpoint. Baselines include π₀-Stateless, π₀-FrameStack, π₀-Hint, π₀-Vanilla Recurrent Memory, π₀-μ-VLA, published Diffusion Policy / ACT / X-VLA / π₀.₅ results, and NativeMEM integrated into the same setup. Locally trained policies share the same π₀ backbone, LoRA adaptation settings, and optimizer steps.

Why This Matters

Impact on research. The paper reframes memory for robot control as an interface-design problem: whether the associative state is exposed as attention correction, an auxiliary recurrent cell, or contextual token content materially changes performance (20.8% / 23.2% / 37.6% anchor-free). It also shows that a fixed-size associative state can outperform an explicit compressed-history representation while adding roughly one-tenth the parameters, and that gains on memory-dependent tasks need not come at the cost of general manipulation performance — the LIBERO four-suite mean stays comparable at 95.1% versus 94.2% while LIBERO-Long improves by 5.4 points.

Real-world applications.

• Household and service robots that must return objects to their original locations or restore an arrangement after an interruption.
• Warehouse and logistics picking, where the correct destination depends on where an item was taken from rather than what is currently visible.
• Multi-step assembly or inspection tasks where a brief early event (a button press, a placement) determines a later decision many steps away.
• Long-horizon tabletop manipulation on physical platforms, where the reported real-robot mean success improved from 3.33% to 26.67%.

Industry relevance. The efficiency profile is the practical hook: 91.9% lower inference latency, 9.5% less peak GPU memory, and 40M added parameters versus NativeMEM's 415M. For deployed robots running policies on constrained onboard hardware, a memory that does not grow with episode length and does not require an extra video-history encoding step is directly relevant to throughput and cost.

Future Directions

• Better combination of explicit video history with associative memory. The paper reports that adding MemBodied to NativeMEM reaches 45.2%, still below standalone MemBodied's 50.0%, and attributes the gap partly to history tokens drawing attention away from retrieved associative content. The authors name this as a direction for future work.
• Making the episode anchor adaptive rather than fixed. The anchor reduces Block Ranking from 32.0% to 22.0% because that task depends on tracking evolving progress rather than recalling the initial scene. A mechanism that decides when the initial-scene reference is useful remains open.
• Resolving the task-level split between memory types. NativeMEM is stronger on Put Back Block and Swap Blocks while MemBodied is stronger on Rearrange Blocks, Battery Try, and Block Ranking, so a principled way to select or route between compressed history and associative memory is unresolved.
• Extending the evaluation. The paper notes nine RMBench tasks in the benchmark while evaluating five, and its largest LIBERO improvement is concentrated in the long-horizon suite; whether the approach scales to more memory-complexity categories and longer fully observable horizons is not reported.

Target Audience

Robotics and embodied-AI researchers working on VLA policies and memory-augmented control; engineers deploying manipulation policies under latency, memory, or parameter budgets; and graduate students familiar with transformer-based policy architectures who want a concrete case study in transferring an associative fast-weight memory from language modelling to continuous robot control. Readers looking for a purely conceptual introduction to VLA models should start elsewhere, since the paper presumes knowledge of flow-matching action experts and delta-rule memory.

Authors’ abstract

Vision-Language-Action models provide a strong foundation for general-purpose robot control, yet a vast majority of policies do not preserve and leverage episode-level information beyond the current observation. This limitation is consequential in history-dependent manipulation tasks that depend on information available only in past observations. Retaining past observations in context can aid in recovering this information, but at the significant cost of ever-growing, bloated context and inference latency. We thus introduce MemBodied, a fixed-size episodic memory with two complementary components: an associative state that records interactions across policy calls and an episode anchor that preserves a compact representation of the initial scene as a reference. At each policy call, the model conditions action generation on the current input and the memory components, rather than directly using past observations. Across five evaluated RMBench tasks requiring memory, MemBodied achieves $7.81\times$ the mean success rate of a stateless policy and $2.98\times$ of vanilla recurrent memory, while outperforming the strongest memory-augmented baseline by $1.3\times$ with $10\times$ fewer added parameters. On the fully observable LIBERO-Long suite, it reached 90.6%, a 5.4% improvement over the stateless $π_0$ policy. These findings support MemBodied as a practical alternative to expanding the policy context for history-dependent manipulation.