μVLA: Memory Tokens Double Success Rate from 42% to 84%

Picture trying to assemble a puzzle where each glance only reveals a tiny, shifting corner of the image. That is exactly the challenge robots face with partial observability — the camera doesn't see the whole scene, objects get occluded behind the arm, and targets slide out of frame. Traditional VLA models have no working memory: each timestep they get a fresh view and nothing else, no record of what they saw a moment ago.

Researchers at AIRI (Artificial Intelligence Research Institute, Moscow) tackled this with μVLA (arXiv 2606.12497) — a surprisingly elegant fix: append a small set of learnable recurrent memory tokens to the transformer backbone and carry them across timesteps. The result: success rate on MIKASA-Robo jumps from 42% to 84% with no auxiliary losses, no new datasets, no architectural redesign.

The Problem: Memoryless VLA Policies

Current VLA models like OpenVLA-OFT treat each observation independently — a Markovian policy. This works when the camera sees everything, but fails when:

• Objects become occluded: the target box is now hidden behind the gripper, but the robot still needs its location
• Targets leave the camera's field of view: the target moves, the camera doesn't track it
• Tasks require remembering earlier state: "pick the cube whose color you saw at the start" — by pickup time, the cue is gone from the frame

MIKASA-Robo was designed specifically to stress-test this weakness: 32 manipulation tasks across 12 groups, with deliberate partial observability. Baseline OpenVLA-OFT without memory achieves only 0.42 average success rate on 5 training tasks — and just 0.07 on held-out tasks.

The Core Idea: Memory Tokens Inside Self-Attention

The solution is deceptively simple. Instead of redesigning the architecture, the authors do one small thing: insert m learnable tokens into the attention layer's input sequence and carry those tokens across timesteps within the same episode.

The input sequence with memory looks like:

[BOS] [VIS] [PROPRIO] [M_t] [TEXT] [ACT] [STOP]

where M_t ∈ ℝ^{m×d} is the memory state at timestep t (m tokens of dimension d, matching the backbone's hidden size). Each attention layer processes the full sequence, and after the forward pass:

• Memory token outputs become M_{t+1} — passed to the next step
• Observation and action tokens are processed normally for action prediction

A single forward pass simultaneously reads memory from t-1 and writes the recurrent state for t. No separate pass, no extra module.

Memory-Action Guard: Preventing Shortcut Learning

A key insight: if memory tokens can attend to previous action tokens, the model may learn to simply copy actions rather than encoding meaningful scene information — what the authors call action leakage.

The fix is a memory-action guard: an attention mask that blocks memory tokens from attending to action tokens. Memory can attend to: observations, proprioception, language, and prior memory state. This forces M_{t+1} to encode task-relevant scene information rather than taking the easy shortcut.

Architecture in Detail

μVLA instantiates on top of OpenVLA-OFT — the speed-optimized fine-tuned version of OpenVLA, built on a LLaMA2 7B backbone and DINO + SigLIP vision encoders. Three parameters define a μVLA variant:

Two update rules:

TBPTT (Truncated Backpropagation Through Time): Gradients flow directly through K steps. Memory receives explicit learning signal — more expressive but uses more VRAM.

EMA (Exponential Moving Average): Memory is updated as M_in[t+1] = α × M_out[t] + (1-α) × M_in[t]. No backprop through time; lighter on VRAM but weaker gradient signal.

Initialization: M_0 (at episode start) is a shared learnable parameter — not zeros, not random noise. The model learns a good "starting state" for memory.

Installation

μVLA requires a custom transformers fork to support the memory-augmented backbone. This is pinned in pyproject.toml.

Hardware requirements:

• Inference: 1 GPU with ~16 GB VRAM
• Training: 1-8 GPUs with 40-80 GB VRAM (bfloat16); single 80GB GPU: --batch_size 4 --use_gradient_checkpointing True

git clone https://github.com/CognitiveAISystems/muVLA.git
cd muVLA
uv sync --python 3.10

After setup, follow the benchmark-specific docs:

• MIKASA-Robo: SETUP.md then MIKASA.md (sections 2, 3, 4)
• LIBERO: SETUP.md then LIBERO.md (sections 1, 2, 3, 4)

Pre-trained Checkpoints

Four checkpoints are available on Hugging Face (mu-vla org), all with merged LoRA weights (~16GB each):

from huggingface_hub import snapshot_download

# Best checkpoint for memory tasks: m=64, K=2, TBPTT
snapshot_download(
    "mu-vla/mu-vla-openvla-oft-mikasa-robo-5-tasks-m64-k2-tbptt",
    local_dir="./checkpoints/mu-vla-k2"
)

Full checkpoint list:

Training μVLA from Scratch

μVLA is trained end-to-end on task-specific data with LoRA rank 32 on top of OpenVLA-OFT. Memory tokens are randomly initialized and learned jointly with fine-tuning — no separate pretraining stage.

# Best configuration: m=64, K=2, TBPTT, cosine LR
uv run python train.py \
  --base_model openvla/openvla-7b \
  --dataset mikasa_robo \
  --memory_tokens 64 \
  --tbptt_length 2 \
  --use_lora true \
  --lr_schedule cosine \
  --batch_size 4

# Memoryless ablation baseline
uv run python train.py \
  --base_model openvla/openvla-7b \
  --dataset mikasa_robo \
  --use_memory false \
  --use_lora true

Key training notes:

• Cosine LR schedule consistently outperforms constant schedule — treat as required
• Round-robin episodic dataloader is essential to maintain temporal order within episodes
• Memory tokens are volatile in early epochs — run at least 3 seeds, report the average

Experimental Results

MIKASA-Robo: Partial Observability Benchmark

MIKASA-Robo's flagship task is RememberColor — the robot sees a colored cube at the start of the episode, then must manipulate based on that color after the cue is no longer visible. Five training task variants are used for primary evaluation; additional tasks assess held-out generalization.

Why K=2 beats K=8 — a non-obvious finding:

Longer TBPTT windows (K=8) let gradients flow farther back in time, but they also create a harder credit assignment problem: the model struggles to identify which past observation caused the current action outcome. K=2 hits a sweet spot — enough temporal lookahead to learn cue-recall, not so much that gradient signal gets diluted.

On RememberColor5 (the hardest variant): μVLA K=2 achieves 0.93 vs baseline 0.35–0.40.

LIBERO: No Regression on Fully Observable Tasks

On the fully observable LIBERO benchmark (4 suites), μVLA m=64, K=8 achieves 96.2% average success — matching or exceeding memoryless OpenVLA-OFT. Memory tokens don't hurt when they're not needed; the model effectively learns to ignore them in Markovian settings.

Causal Analysis: Proving Memory Is Functionally Used

Two experiments verify memory is causally necessary, not decorative:

Noise Injection Test: Inject random noise into M_t at inference time, nothing else changed. RememberColor5 drops from 0.94 → 0.09 under K=2. Corrupt memory → robot forgets the cue → task fails completely.

Freeze-First Test: Lock memory to M_1 (first frame) for the entire episode. Cue-recall tasks (where the cue is front-loaded) remain stable; tasks requiring continuous scene tracking collapse. Memory encodes exactly the type of information each task class needs.

The cosine distance plot shows memory changing sharply at task-phase transitions (when an object disappears, when the task moves to a new stage) and stabilizing when the scene is static — meaningful behavior, not noise.

A Critical Inference Detail

Receding-horizon control is mandatory for cue-recall tasks: query the model at every environment step and execute only the first action from the predicted chunk. Running action chunking (chunk ≥ 2) causes recall performance to collapse.

The reason: memory must be updated from every observed frame. If you execute several actions before querying again, the memory representation goes stale.

# CORRECT: query every step
for step in range(episode_length):
    action_chunk = model.predict(obs, memory_state)
    memory_state = action_chunk.memory_output   # update memory
    env.step(action_chunk[0])                   # execute only first action

# WRONG: execute the full chunk before re-querying
action_chunk = model.predict(obs, memory_state)
for action in action_chunk:
    env.step(action)  # memory is stale for 7 of these 8 steps

Comparison with Other Memory Approaches

μVLA is not the only way to add memory to VLA models:

μVLA sits at the sweet spot of simplicity and effectiveness: no retrieval, no auxiliary loss, no new dataset — just add tokens and fine-tune on existing task data.

When to Use μVLA

Good fit:

• Tasks with clear partial observability — narrow field-of-view cameras, occlusion-prone setups
• Manipulation requiring state memory across multiple steps (RememberColor-style tasks)
• You're already using OpenVLA-OFT and want a quick, principled improvement

Not the right tool:

• Fully observable tasks (LIBERO-style): memory gain is near zero
• Inference with large action chunks (chunk ≥ 2): conflicts with the receding-horizon requirement
• Very tight VRAM budgets: TBPTT K=2 adds ~15-20% memory overhead at training time

Expert Takeaway

What makes μVLA stand out is not the number 84% by itself — it's the scientific rigor behind it. Each variable is ablated cleanly (m, K, update rule), causal mechanism is verified with intervention experiments, and limitations are disclosed honestly (K=2 loses ~75% recall beyond training horizon; novel memory structures transfer poorly). This is the kind of paper you can trust to deploy from, not just cite.

The paper's deepest contribution may actually be MIKASA-Robo + μVLA as a template for how to properly benchmark and ablate memory mechanisms in VLA research. Other papers claiming "our model has memory" now have a rigorous baseline to beat.

Related Posts

• OpenVLA-OFT: Fine-Tuning VLA for Fast Robot Manipulation
• MemoryVLA++: Dual Memory and World Models for Robotic Manipulation
• FM-VLA: Force Memory Token for Contact-Rich Manipulation