losterberg3.github.io/unimem-vla

Abstract

Mulitmodal Memory and control, one backbone

While Vision-Language-Action (VLA) models have leveraged internet-scale pretraining and task-focused finetuning to achieve strong performance on long-horizon tasks, they often struggle with non-Markovian tasks that require memory. Existing approaches to memory typically involve additional Vision-Language-Models (VLMs) for long-term memory management, introducing a memory bottleneck and a fractured training pipeline. Conditioning on multiple historical frames can provide the VLA with access to more descriptive features of past scenes, but can degrade performance if frames are chosen at arbitrary, fixed intervals. To address these limitations, we present UniMem, a framework that unifies high-level, multimodal memory and low-level control under one backbone. UniMem employs an event classifier for memory updates, a keyframe encoder for dense spatial memory, and a keyframe caching technique to minimize overhead during policy rollouts. We evaluate UniMem across five simulation and four hardware tasks targeting sequential and spatial memory, demonstrating that our unified, single-model system outperforms fixed-interval image sampling baselines (93.4% vs. 68.2%) in simulation and hierarchical baselines (80.0% vs. 43.5%) in hardware, while offering faster inference and a simple training pipeline for easy adoption.

Overview

State-of-the-art VLA policies act entirely on the current camera frame — they have no architectural memory. That's fine for short, fully-observable tasks, but it breaks down the moment a task requires knowing what already happened: state-of-the-art policies still fail at tasks as simple as picking something up and putting it back where it was. Two moments in a task can look nearly identical while demanding completely different actions depending on history — a failure mode we call perceptual aliasing. Prior methods either add a second, slower VLM to plan over history (a hierarchical, decoupled pipeline) or feed the policy a fixed window of arbitrarily sampled recent frames (which invites spurious correlations and doesn't scale). UniMem instead adds a lightweight event classifier directly on top of a VLA's own latent features. When it detects a task-critical milestone, it (1) appends a short phrase to a running textual memory and (2) caches the current multi-view image as a keyframe — both fed back into the same backbone that produces actions. Because keyframe hidden states are cached rather than recomputed, this rich multimodal memory costs almost nothing at inference time.

🧭

Event-Driven Memory

A lightweight classifier head reads the VLA's own latent state to detect task milestones and autoregressively updates a language prompt and keyframe cache — no separate memory module or transformer.

🧩

Multi-Modal Conditioning

Textual memory anchors task progress; keyframe memory supplies fine-grained spatial detail. Ablations show both are necessary and complementary — neither alone is enough.

⚡

Real-Time Keyframe Caching

Keyframe hidden states are cached prior to temporal attention, giving ~90ms single-frame-like latency and a 6× speedup over hierarchical memory baselines.

Method

How UniMem works

Overview of UniMem. An event classifier ($f_\phi$) detects sparse sub-task transition events from the backbone latent space, autoregressively updating textual memory and a cache of precomputed keyframe hidden states. This routes to our keyframe encoder and a tokenizer, providing event-driven, multimodal memory directly to the backbone and action expert while unifying memory and control in one self-sustaining, low-latency VLA.

Event classification & textual memory

UniMem is built on top of π_0.5, an open-source VLA with a PaliGemma backbone and a Gemma action expert. A lightweight MLP head is attached to the backbone's final-layer latent representation and trained jointly with the policy to predict a task-relevant event (e.g. "grabbed box") or a null class at every step. Because this head shares gradients with the backbone, it doubles as auxiliary supervision — pushing the model's interanl features toward whatever is relevant for memory and action. Detected events are appended in natural language to a running textual memory that conditions every subsequent step, giving the policy an explicit and inexpensive signal of its own progress.

Keyframe encoder & caching

Text alone can't disambiguate everything — e.g. spatial questions like which cup or bin. So every time the event classifier fires, UniMem also captures the current wrist and exterior images as a keyframe, maintaining a small rolling bank (capped at 3–4 past milestones) rather than a fixed-stride sequence of arbitrary frames. A modified SigLIP encoder interleaves causal temporal self-attention every 4 layers, and once a keyframe's hidden state has propagated through, it's cached in raw form — so adding a new frame only requires shifting positional embeddings, never recomputing spatial attention over the whole history. This keeps inference near single-frame speed even as visual memory accumulates.

Automated data labeling pipeline diagram — **Automated data labeling.** An LLM agent (Claude Sonnet 5) is prompted to generate a script that labels demonstrations with events and textual memory strings directly from action signatures (e.g. a gripper closing means an object was grasped). Multiple frames in a window $W_i$ are labeled with the corresponding event $e_t$, and $\mathcal{M}_t$ is only updated once this window has passed. At training time, a keyframe from each past $W_i$ along with $I_t$ is used to build $\mathcal{H}_t$. A human spot-checks a subset of labeled rollouts and the script is refined if needed — no manual frame-by-frame annotation required.

Experiments

Nine memory-critical tasks

UniMem is evaluated on five simulation tasks (robosuite, 7-DoF Franka Panda) and four hardware tasks (UFACTORY xArm6), each designed to require sequential and/or spatial memory that a purely reactive policy cannot solve.

Overview of simulation and hardware tasks — **Overview of tasks.** (Top) Manipulation tasks in robosuite. (Bottom) Real-world tasks on the xArm6 setup.

Real-World Results

Hardware tasks (xArm6)

UniMem is benchmarked against MemER — a hierarchical VLM + π_0.5 system — as the primary hardware baseline, plus no-memory, text-only, and keyframe-only ablations. UniMem reaches 80.0% average success versus MemER's 43.5% due to conditioning the low-level policy on memory directly at every control step rather than through an intermediate subtask command.

Task	MemER	N.M.	T.O.	K.O.	Ours
HammerMeasure	87	13	53	53	87
BeanScoop	67	0	27	20	93
TableClean	13	0	0	47	80
TapScoopPour	7	7	7	27	60
Hardware Average	43.5	5.0	21.8	36.8	80.0

Success rates (%), N=15 rollouts per task/condition.

MemER — hierarchical VLM + VLA baseline N.M. — No Memory (vanilla π_0.5) T.O. — Text-Only memory K.O. — Keyframe-Only memory Ours — full UniMem

HammerMeasure

"Slowly measure the width of the hammer using the tape measure, controlling the retract to ensure the hook does not slip off the hammer."

🎬 HammerMeasure — UniMem (Ours) Video coming soon

BeanScoop

"Pick up the spoon, put three scoops of beans into the bowl, and then place the spoon back."

🎬 BeanScoop — UniMem (Ours) Video coming soon

TableClean

"Pick and place the bottle off the 60cm × 80cm table, grab the sponge, wipe the bottle's original location, and then place the sponge back."

🎬 TableClean — UniMem (Ours) Video coming soon

TapScoopPour

"Wait for a human to tap one of eight cups, grab the spoon, scoop the beans, pour the beans into the tapped cup, and then place the spoon back."

🎬 TapScoopPour — UniMem (Ours) Video coming soon

Cumulative subtask success rate breakdown — **Cumulative success rates.** Stage-by-stage breakdown for BeanScoop, TableClean, and TapScoopPour, showing where each ablation and baseline drops off within the task.

Simulation Results

Simulation tasks (robosuite)

In simulation, UniMem is compared against π_0.5 with a video-encoder baseline (V.E.) that samples frames at fixed 6-second intervals using the same keyframe encoder architecture — isolating how arbitrary frame sampling performs relative to event-driven keyframes. UniMem reaches 93.4% average success versus 68.2% for the fixed-interval baseline.

Task	π_0.5+V.E.	N.M.	T.O.	K.O.	Ours
UpDown	84	52	96	92	100
UpDown3Times†	16	22	93	23	96
OccludedTap	96	60	88	100	96
UpDownSpatial‡	49	6	30	52	79
PlateRecall	96	8	20	96	96
Sim Average	68.2	29.6	65.4	72.6	93.4

Success rates (%), N=25 rollouts per task/condition. (†) reports mean subtask success due to long-horizon complexity. (‡) reports continuous spatial accuracy.

V.E. — Video Encoder (fixed-interval frames) N.M. — No Memory (vanilla π_0.5) T.O. — Text-Only memory K.O. — Keyframe-Only memory Ours — full UniMem

UpDown

"Pick up the box and then put it back down once."

🎬 UpDown — UniMem (Ours) Video coming soon

UpDown3Times

"Pick up the box and then put it back down three times."

🎬 UpDown3Times — UniMem (Ours) Video coming soon

OccludedTap

"Pick up the box, place it in a bin, retract so you can't see inside the bin, and then decide which bin contains the box."

🎬 OccludedTap — UniMem (Ours) Video coming soon

UpDownSpatial

"Pick up the box at various locations on a table, and then place it back in its original location."

🎬 UpDownSpatial — UniMem (Ours) Video coming soon

PlateRecall

"Pick up the box from one of 4 plates, place it to the side, and then tap which plate originally had the box."

🎬 PlateRecall — UniMem (Ours) Video coming soon

Efficiency

Real-time inference via keyframe caching

By caching pre-computed keyframe hidden states prior to temporal self-attention, UniMem avoids re-encoding the entire visual history at every control step. Benchmarked on an RTX 4090 against a no-caching ablation, single-frame π_0.5, and MemER's dual-system architecture, maintaining a 16-keyframe memory context across four camera streams adds only ~25ms of latency beyond the single-frame, 2-camera base policy — keeping single-frame-like (~90ms) inference and a 6× speedup over hierarchical memory baselines.

Inference latency vs. context window size on an RTX 4090, for 2 camera streams (left) and a simulated bimanual setup with 4 camera streams (right).

Discussion

Conclusion & limitations

UniMem shows that long-horizon textual and visual memory can be unified within a single VLA backbone — matching or beating hierarchical, decoupled memory systems while training with a single-stage pipeline and running at near single-frame inference speed. Across nine tasks in simulation and hardware, this translates to higher success rates, simpler training, and low-latency real-world rollouts.

Horizon length. UniMem has been validated over multi-minute rollouts, not yet over tasks spanning tens of minutes or hours; memory editing mechanisms like pruning and consolidation are a promising next step.
Offline labeling. The keyframe/event extraction pipeline currently relies on automated offline labeling scripts rather than a policy that autonomously decides what's worth remembering.
Single memory timescale. UniMem maintains one temporal context spanning several minutes; distinguishing short- from long-term memory — including memory that lives in the policy's weights, not just its input stream — remains open.

Citation

BibTeX

@misc{osterberg2026unimemunifyingmultimodalmemory,
  title        = {UniMem: Unifying Multimodal Memory and Control for Vision-Language-Action Models}, 
  author       = {Lars Osterberg and Maggie Wang and Mac Schwager},
  year         = {2026},
  eprint       = {2608.22869},
  institution  = {Stanford University},
  archivePrefix= {arXiv},
  primaryClass = {cs.RO},
  url          = {https://arxiv.org/abs/2608.22869}
}