MolmoAct2-SO101-VTP-Init
Step-0 checkpoint: compression applied, training not started.
This is the initialization state of a MolmoAct2 SO-101 fine-tuning run with
Grid Sampler vision-token pruning enabled. It is not a trained policy — it
is the exact artifact that lerobot-train holds in memory before its first
optimizer step, published so the retraining run is reproducible from step 0.
The Grid Sampler weights in this checkpoint are randomly initialized. This checkpoint is a training seed, not something to evaluate.
What the compression does
MolmoAct2 feeds the VLM one pooled 14×14 patch grid (196 tokens) per camera. Grid
Sampler replaces that grid with grid_token_sampler_num_tokens tokens whose
locations are predicted rather than fixed: a scout MLP reads the globally pooled
feature map and emits K normalized (x, y) coordinates, features are bilinearly
sampled at those coordinates, and a coordinate encoder adds a positional
embedding. With 2 cameras that is 392 → 32 image tokens at K=16, shortening the
prefix the LLM attends over.
| Baseline | This variant | |
|---|---|---|
| Image tokens per camera | 196 (14×14 pooled grid) | 16 (learned sample points) |
| Image tokens total (2 cams) | 392 | 32 |
| New parameters | — | scout MLP + coordinate encoder (randomly initialized) |
| Action expert / flow matching | unchanged | unchanged |
Config deltas vs. lerobot/MolmoAct2-SO100_101-LeRobot
use_grid_token_sampler : false -> true
grid_token_sampler_num_tokens : -> 16
action_mode : continuous -> both
setup_type : "single so100/so101 robotic arm in molmoact2"
-> "single SO-101 5-DOF robotic arm with gripper"
control_mode : "absolute joint pose" (unchanged)
chunk_size and n_action_steps stay at 30. They are not free parameters
here: apply_norm_tag_metadata() pins both from the so100_so101_molmoact2
norm-tag metadata in the base checkpoint, and it runs after CLI overrides, so
passing a different value has no effect. The SO-101 baseline runs this checkpoint
is meant to be compared against were trained the same way.
The token budget is part of the processor config as well as the model config —
the preprocessor has to emit the matching number of image placeholder tokens.
Both are set consistently in this checkpoint; loading it with a different
grid_token_sampler_num_tokens will not work.
Normalizer statistics come from the SO-101 training corpus (see below), not from
the base checkpoint's so100_so101_molmoact2 norm tag.
Why K=16
From docs/results_drift_keystone_libero.md in the training repo — LIBERO, 4
suites × 10 tasks × 20 episodes = 800 episodes per model, 10-step flow matching:
| Model | Overall success |
|---|---|
| MolmoAct2-LIBERO baseline | 97.1% |
| no-prune control, same recipe | 97.9% |
| GridSampler K=8 | 92.0% |
| GridSampler K=16 | 97.1% |
| GridSampler K=32 | 96.6% |
K=8 falls off a cliff; K=16 holds parity with the unpruned control at 12× fewer
image tokens, and K=32 buys nothing over K=16. Those are LIBERO numbers on a
different embodiment — SO-101 parity has not yet been measured, which is what the
run seeded by this checkpoint is for. Override with GRID_NUM_TOKENS to sweep.
Training data
Merged SO-101 corpus, 177,661 frames from 11 dataset repos across four categories:
| Category | Content | Share of frames | Optional batch ratio |
|---|---|---|---|
| A-1 | 5 cm teleop | 21.6% | 6 |
| A-2 | 5 cm DAgger | 11.2% | 2 |
| A-3 | 2 cm teleop | 62.1% | 3 |
| A-4 | 2 cm DAgger | 5.0% | 1 |
By default batches are drawn by uniform shuffling, so composition follows the
frame shares above — 2 cm teleop dominates at 62%. A balanced batch sampler that
instead pins a fixed 6 : 2 : 3 : 1 mix in every batch is available but
opt-in, via CATEGORY_META=<sidecar> (requires batch_size >= 12 for every
category to get a slot). Normalizer statistics are computed over the whole corpus
and are unaffected by the sampling mode, so this checkpoint is valid for either.
Reproducing / continuing
# repo: nota-github/xpu-lerobot, branch feat/molmoact2-so101-compression
# regenerate this exact init checkpoint
EXPORT_INIT=1 INIT_EXPORT_DIR=/tmp/vtp-init HF_TOKEN=... \
scripts/train_molmoact2_so101_grid_sampler.sh
# the training run it seeds (2 GPUs, 18k steps, per-device batch 16)
HF_TOKEN=... WANDB_API_KEY=... scripts/train_molmoact2_so101_grid_sampler.sh
# sweep the token budget
GRID_NUM_TOKENS=32 HF_TOKEN=... scripts/train_molmoact2_so101_grid_sampler.sh
References
- Grid Sampler (GridS) active token sampling, ICML 2026
- Fang et al. MolmoAct2: Action Reasoning Models for Real-world Deployment. arXiv:2605.02881