pi05-mem — SO-101 pick-the-remembered-cube, 64 memory tokens, TBPTT K=2
π₀.₅ with a recurrent memory mechanism, fine-tuned from
lerobot/pi05_base on a real-robot
memory task. A bank of 64 learnable memory tokens is carried across timesteps inside an
episode and updated end-to-end with truncated backpropagation through time, so the policy
can condition on an observation that is no longer visible in the current frame.
This is the port of the memory mechanism from
μVLA (a fork of OpenVLA-OFT) onto π₀.₅ as
implemented in LeRobot. Code: avanturist322/muVLA-pi0.5.
Sibling checkpoint: the same recipe with a longer truncation window is at
avanturist/pi05-mem-so101-pick-cube-color-m64-k8-tbptt.
Read the Comparability section before treating the pair as an ablation over K.
Task
Dataset: tttonyalpha/so101_pick_cube_300_task3
(LeRobot v3).
Remember the color of the cube and then pick the matching one
| Robot | SO-101 follower arm |
| Episodes / frames | 300 / 116 016 |
| Control rate | 30 fps |
| Cameras | observation.images.top, observation.images.wrist |
| State / action dim | 6 / 6 |
| Tasks | 1 |
The cue is visible only early in the episode and is gone by the time the grasp has to be committed, which is what makes a recurrent state useful here rather than decorative.
Training
Fine-tuned from lerobot/pi05_base on 8 GPUs (80 GB each) with DDP.
| Memory | |
|---|---|
use_memory | true |
num_mem_tokens | 64 |
memory_update | tbptt |
tbptt_length (K) | 2 |
attention_mask_mode | custom |
memory_write_scale | 1.0 |
memory_init_std | 0.02 |
| Optimization | |
|---|---|
| Batch size | 64 per rank, 512 global |
| Micro-steps | 30 000 |
| Optimizer steps | 15 000 (one per TBPTT window) |
| Frames seen | ~15.36 M (30 000 × 512) |
Action chunk (action_horizon) | 8 |
| Learning rate | 5e-5, cosine, 750 warmup steps, min ratio 0.1 |
| Optimizer | AdamW, β = (0.9, 0.95), ε = 1e-8, weight decay 0.01 |
| Gradient clipping | 1.0 |
| Precision | bfloat16, gradient checkpointing on |
| Vision encoder | frozen |
| Seed | 42 |
| Wall-clock | 33.7 h |
Final training loss 0.0044 (mean of the last 10 logged values: 0.0066); the run started
at 1.918. train_config.json and training_summary.json in this repository are the
verbatim artifacts the run wrote, including the DDP sharding check.
Files
| File | What it is |
|---|---|
model.safetensors | 814 tensors, the full policy including model.memory_module.initial_memory (64 × 2048) |
memory_module.pt | The memory parameters on their own, loaded with strict=True at eval |
memory_meta.json | Memory hyperparameters, auto-detected when the checkpoint is loaded |
config.json | π₀.₅ policy config |
dataset_stats.json | Copy of the training dataset's meta/stats.json — see the warning below |
train_config.json | Every training flag as passed |
training_summary.json | Step counts, losses, wall-clock, DDP check |
memory_module.pt is deliberately redundant with model.safetensors: the safetensors path
goes through π₀.₅'s key remapper with strict=False, so a silent rename would leave the
memory at its random initialization — which looks like "memory does not help" rather than
like a bug. Loading the standalone file with strict=True makes that failure loud.
Normalization statistics are not in the weights
This checkpoint does not carry normalize_inputs / unnormalize_outputs buffers. The
loader takes normalization from the dataset directory, because π₀.₅ discretizes the
normalized state into 256 bins inside the text prompt — evaluating with the wrong quantiles
feeds the model a corrupted prompt and produces plausible-looking but wrong actions, with no
error raised.
So you need the dataset alongside the weights. dataset_stats.json here is a copy of
meta/stats.json for reference and reproducibility, but the loader reads the dataset
directory, not this file.
Usage
git clone https://github.com/avanturist322/muVLA-pi0.5.git && cd muVLA-pi0.5
# environment setup: see the repository README
from huggingface_hub import snapshot_download
from pi05_mem.eval.loader import load_eval_policy
ckpt = snapshot_download("avanturist/pi05-mem-so101-pick-cube-color-m64-k2-tbptt")
data = snapshot_download(
"tttonyalpha/so101_pick_cube_300_task3", repo_type="dataset"
) # supplies the normalization statistics
bundle = load_eval_policy(ckpt, data, device="cuda", dtype="bfloat16")
load_eval_policy reads memory_meta.json, so memory does not have to be re-declared on
the command line: the checkpoint states whether memory is on and how it is configured.
Inference regime defaults to receding horizon for memory checkpoints (n_action_steps=1,
requery every environment step, memory advances once per step exactly as in training). Pass
receding_horizon=False for the open-loop chunked regime.
Comparability with the K=8 sibling
The two checkpoints are not batch-matched, so the difference between them is not a clean ablation over the truncation length:
| K=2 (this) | K=8 (sibling) | |
|---|---|---|
| Global batch | 512 | 256 |
| Optimizer steps | 15 000 | 3 750 |
| Frames seen | ~15.36 M | ~7.68 M |
| Final loss | 0.0044 | 0.0312 |
This run saw twice the frames and took four times as many optimizer steps. Attributing its lower loss to the shorter TBPTT window would be wrong.
Earlier partial upload
An in-progress checkpoint from this same run (step 7 500 of 30 000) was published earlier at
avanturist/pi05-mem-lerobot-config-c-mem-k2-v3.
This repository holds the final step-30 000 weights and supersedes it.
What is not here
No in-environment success rates. Both checkpoints were trained but not yet evaluated on the physical SO-101 or in simulation, so this card reports training loss only. Training loss on a 300-episode imitation dataset is a weak proxy for task success and should not be read as one.
Citation
The memory mechanism this ports:
@article{cherepanov2026muvla,
title={{$\mu$}VLA: On Recurrent Memory for Partially Observable Manipulation in VLA Models},
author={Cherepanov, Egor and Kachaev, Nikita and Zelezetsky, Daniil and Bulatov, Aydar and Pshenitsyn, Artem and Kuratov, Yuri and Skrynnik, Alexey and Panov, Aleksandr I and Kovalev, Alexey K},
journal={arXiv preprint arXiv:2606.12497},
year={2026}
}
The base model is π₀.₅ (arXiv:2504.16054) as implemented
in LeRobot. Weights inherit the
Gemma Terms of Use from lerobot/pi05_base.