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