mu-VLA (OpenVLA-OFT + recurrent memory), LIBERO 4 suites, 64 memory tokens, TBPTT K=8
mu-VLA extends OpenVLA-OFT with a recurrent memory module: a set of learnable memory tokens injected into the multimodal prefix of the transformer sequence and carried across environment steps. Gradients flow through a truncated backpropagation-through-time (TBPTT) window; the memory state is reset at episode boundaries.
This checkpoint corresponds to experiment 7 in the mu-VLA experiment matrix: a single policy trained jointly on all four LIBERO suites through the episodic dataloader, with 64 memory tokens and a TBPTT window of K=8.
LIBERO is fully observable, so memory is not required to solve it. This checkpoint is here to show that adding recurrent memory does not cost anything on standard manipulation benchmarks, which is the control for the POMDP results on MIKASA-Robo.
Sibling checkpoints:
- mu-vla-openvla-oft-mikasa-robo-5-tasks-m64-k8-tbptt
- mu-vla-openvla-oft-mikasa-robo-5-tasks-m64-k2-tbptt
Memory configuration
| Setting | Value |
|---|---|
num_mem_tokens | 64 |
memory_update | tbptt |
tbptt_length | 8 |
attention_mask_mode | custom |
ema_alpha | 0.1 (unused when memory_update=tbptt) |
The same values are shipped in memory_meta.json so that the evaluation scripts in the code
repository pick them up automatically.
Training setup
Base model: openvla/openvla-7b, fine-tuned with the OpenVLA-OFT recipe (L1 regression action
head, proprioception projector, two input camera views, no FiLM, no diffusion).
| Setting | Value |
|---|---|
| Checkpoint step | 150000 |
| Dataset | LIBEROVLAEpisodicDataset, libero_combined |
| Suites | libero_spatial_no_noops, libero_object_no_noops, libero_goal_no_noops, libero_10_no_noops |
| GPUs | 8 x NVIDIA A100 80GB |
| Batch size | 4 per GPU |
| LoRA rank | 32 |
| Learning rate | 5e-4, cosine schedule, 2000 warmup steps, minimum ratio 0.1 |
| Image augmentation | on |
| Gradient checkpointing | on |
One policy covers all four suites. Action normalisation statistics are computed jointly and stored
in dataset_statistics.json.
Evaluation
Protocol: 500 episodes per suite (all tasks in the suite), seed 7, receding horizon inference (the model is queried every step and only the first action of the predicted chunk is executed). Uncertainty is the standard error over the 500 trials.
| Suite | Success rate |
|---|---|
libero_spatial | 0.930 ± 0.011 (465/500) |
libero_object | 0.994 ± 0.004 (497/500) |
libero_goal | 0.966 ± 0.008 (483/500) |
libero_10 | 0.958 ± 0.009 (479/500) |
| average | 0.962 (1924/2000) |
Files
Loading this checkpoint requires more than the base transformer weights. The extra heads are stored as separate PyTorch state dicts, matching the layout the training and evaluation scripts expect.
| File | Purpose |
|---|---|
model-0000*-of-00004.safetensors, model.safetensors.index.json | merged vision-language backbone |
lora_adapter/ | LoRA adapter as saved by PEFT, before merging |
action_head--150000_checkpoint.pt | L1 regression action head |
proprio_projector--150000_checkpoint.pt | proprioception projector |
memory_module--150000_checkpoint.pt | memory module, including initial_memory |
memory_meta.json | memory hyperparameters for evaluation auto-detection |
dataset_statistics.json | action normalisation statistics |
optimizer--150000_checkpoint.pt, scheduler--150000_checkpoint.pt | optimizer and scheduler state, for resuming training |
configuration_prismatic.py, modeling_prismatic.py, processing_prismatic.py | custom model code, memory token injection and the custom attention mask live here |
Usage
The checkpoint is not loadable with upstream transformers. It needs the mu-VLA code, which
depends on a forked transformers that supports the custom attention mask, and on
LIBERO for the environments.
huggingface-cli download mu-vla/mu-vla-openvla-oft-libero-4-tasks-m64-k8-tbptt \
--local-dir ./mu-vla-libero-m64-k8
Then point the evaluation script at the downloaded directory; memory hyperparameters are read from
memory_meta.json. See the code repository, CognitiveAISystems/mu-vla, for setup and the exact
command line.
Attribution
mu-VLA is a fork of OpenVLA-OFT by Moo Jin Kim, Chelsea Finn and Percy Liang (arXiv:2502.19645), which in turn builds on OpenVLA. The recurrent memory module and the episodic dataloaders are the contribution of this work. Released under the MIT license, following upstream.