MolmoAct2-SO101-VTP_toponly_p60_Init

Step-0 checkpoint: compression applied, training not started.

Grid Sampler vision-token pruning on the overhead (top) camera only, at 60.2% — 78 of 196 tokens kept. The wrist camera is left untouched: it keeps its full 196-token pooled grid and does not pass through a sampler at all.

The Grid Sampler weights are randomly initialized. This is a training seed, not something to evaluate.

Why top-only, and why "untouched" is not the same as K=196

The two views carry different kinds of information — the overhead view gives global layout (where the cube and box are), the wrist view gives fine detail near the gripper. Pruning one and not the other isolates which view tolerates compression, which is directly actionable for on-device deployment.

A camera marked "not pruned" here bypasses the sampler entirely. Asking for K=196 instead would still run the ActiveTokenSampler: it predicts 196 coordinates, bilinearly resamples the feature map at them, and adds a coordinate embedding. Same token count, but not an identity — and it adds randomly initialized parameters to a path that needed none. This checkpoint takes the bypass, verified by asserting the wrist tokens come out torch.equal to the pooled grid that went in.

Baseline This variant
top camera tokens 196 78 (learned sample points)
wrist camera tokens 196 196 (unchanged, sampler bypassed)
image tokens total 392 274 (30.1% fewer)
Grid Sampler modules — 1 (top only)
New parameters — 2,049,948 (randomly initialized)
Action expert / flow matching unchanged unchanged

Where this sits

Variant top wrist image tokens Grid Sampler modules
no pruning 196 196 392 0
toponly p20 156 196 352 1
toponly p40 117 196 313 1
toponly p60 78 196 274 1
VTP (both cams, K=16) 16 16 32 1 (shared)

This variant: top 78 tokens (60.2% pruned), wrist untouched.

Honest expectation

The premise "the wrist view must stay intact" is not supported by the evidence that exists. On LIBERO the Grid Sampler was applied uniformly to both cameras (agentview and wrist) and K=16 — 91.8% pruning on both — held parity with the unpruned control (97.1 vs 97.9), collapsing only at K=8 (92.0). So the wrist view tolerated aggressive pruning there.

LIBERO is simulation with a different camera rig and task, so SO-101 may differ, and that is exactly what this sweep measures. But do not expect top-only pruning to be safer than uniform pruning on the strength of intuition alone.

Note also the ceiling: leaving one camera intact caps the token reduction at half of what uniform pruning reaches. With num_state_tokens=256 plus the prompt in the sequence, this variant's 30.1% image-token cut is a smaller fraction of the total sequence than it looks.

To attribute a result, compare against uniform pruning at a similar total token count, not only against the unpruned baseline.

Config

use_grid_token_sampler        : false -> true
grid_token_sampler_num_tokens :       -> [78, 0]     # [top, wrist]; 0 = not pruned
action_mode                   : continuous -> both
setup_type                    : "single SO-101 5-DOF robotic arm with gripper"

The per-camera list is honoured in three places that must agree, or the image tokens scatter into the language sequence misaligned — silently, with no error: the model builds one sampler per pruned camera, the processor emits that many image placeholder tokens per camera in the same order, and the sequence-length budget sums the per-camera counts. All three were asserted equal (274) before this checkpoint was exported.

chunk_size / n_action_steps stay at 30: apply_norm_tag_metadata() pins both from the so100_so101_molmoact2 norm tag after CLI overrides, so passing anything else has no effect.

Training data

Merged SO-101 corpus, 177,661 frames from 11 dataset repos (5 cm teleop 21.6%, 5 cm DAgger 11.2%, 2 cm teleop 62.1%, 2 cm DAgger 5.0%). Batches are drawn by uniform shuffling by default; a balanced sampler pinning a 6:2:3:1 per-batch mix is opt-in via CATEGORY_META. Normalizer statistics come from the whole corpus and are unaffected by the sampling mode.

Status of verification

Verified at construction: config and processor agree on [78, 0], exactly one sampler exists (for the top camera), scout_mlp.2.weight is (156, 512), the parameter count matches the analytic prediction, and an unpruned camera passes its grid through unchanged.

Not smoke-tested through training steps — no training was run for this variant.

Reproducing

# repo: nota-github/xpu-lerobot, branch feat/so101-train
EXPORT_INIT=1 INIT_EXPORT_DIR=/tmp/toponly-p60 GRID_NUM_TOKENS="[78,0]" HF_TOKEN=... \
  scripts/train_molmoact2_so101_grid_sampler.sh

# training run it seeds
GRID_NUM_TOKENS="[78,0]" HF_TOKEN=... WANDB_API_KEY=... \
  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
Downloads last month
7
Safetensors
Model size
5B params
Tensor type
F32
·
BF16
·
Video Preview
loading

Model tree for geonmin-kim/MolmoAct2-SO101-VTP_toponly_p60_Init

Finetuned
(14)
this model

Paper for geonmin-kim/MolmoAct2-SO101-VTP_toponly_p60_Init