mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-27 11:46:04 +00:00
feat(smolvla): add MEM visual memory
This commit is contained in:
+15
-1
@@ -111,7 +111,21 @@ dataset = LeRobotDataset("lerobot/robomme")
|
||||
|
||||
The env wrapper exposes `pixels/image` and `pixels/wrist_image` as observation keys. The `features_map` in `RoboMMEEnv` maps these to `observation.images.image` and `observation.images.wrist_image` for the policy. State is exposed as `agent_pos` and maps to `observation.state`.
|
||||
|
||||
The dataset's `image` and `wrist_image` columns already align with the policy input keys, so no renaming is needed when fine-tuning.
|
||||
The published dataset uses the raw keys shown above. Policies that expect canonical LeRobot keys should map them at training time. For example, the pretrained SmolVLA checkpoint expects three cameras, while RoboMME provides two:
|
||||
|
||||
```bash
|
||||
uv run lerobot-train \
|
||||
--policy.path=lerobot/smolvla_base \
|
||||
--policy.empty_cameras=1 \
|
||||
--dataset.repo_id=lerobot/robomme \
|
||||
'--rename_map={"image":"observation.images.camera1","wrist_image":"observation.images.camera2","state":"observation.state","actions":"action"}'
|
||||
```
|
||||
|
||||
At evaluation time the environment wrapper has already converted observations to canonical keys. Only the two camera suffixes need to be aligned with the checkpoint:
|
||||
|
||||
```bash
|
||||
'--rename_map={"observation.images.image":"observation.images.camera1","observation.images.wrist_image":"observation.images.camera2"}'
|
||||
```
|
||||
|
||||
## Action Spaces
|
||||
|
||||
|
||||
@@ -76,6 +76,37 @@ Fine-tuning is an art. For a complete overview of the options for finetuning, ru
|
||||
lerobot-train --help
|
||||
```
|
||||
|
||||
### Experimental MEM visual memory
|
||||
|
||||
SmolVLA can optionally fuse a short history of camera frames using the
|
||||
space-time separable vision encoder from [MEM](https://arxiv.org/abs/2603.03596).
|
||||
Every fourth SigLIP layer adds causal attention across time for matching image
|
||||
patches. Historical tokens are discarded inside the vision tower, so the VLM
|
||||
receives the same number of image tokens as the baseline policy.
|
||||
|
||||
The option is disabled by default. The following example uses six observations
|
||||
spaced one second apart for a 10 fps dataset:
|
||||
|
||||
```bash
|
||||
lerobot-train \
|
||||
--policy.path=lerobot/smolvla_base \
|
||||
--policy.use_visual_memory=true \
|
||||
--policy.visual_memory_frames=6 \
|
||||
--policy.visual_memory_stride=10 \
|
||||
--policy.visual_memory_temporal_attention_every=4 \
|
||||
--policy.freeze_vision_encoder=false \
|
||||
--policy.train_expert_only=false \
|
||||
--dataset.repo_id=${HF_USER}/mydataset \
|
||||
--output_dir=outputs/train/my_smolvla_mem
|
||||
```
|
||||
|
||||
`visual_memory_stride` is measured in dataset or environment steps. Keep all
|
||||
other settings and the random seed fixed when comparing against a baseline.
|
||||
Because the public SmolVLA checkpoint was not pretrained with this temporal
|
||||
attention pattern, this is a post-training-only MEM ablation; the MEM paper
|
||||
reports that memory-aware pretraining performs better than introducing visual
|
||||
memory only during task-specific fine-tuning.
|
||||
|
||||
<p align="center">
|
||||
<img
|
||||
src="https://cdn-uploads.huggingface.co/production/uploads/640e21ef3c82bd463ee5a76d/S-3vvVCulChREwHDkquoc.gif"
|
||||
|
||||
Reference in New Issue
Block a user