MolmoAct2-SO101-VTP_locaware_topp60_Init

MolmoAct2-SO101-VTP_locaware_topp60_Init

Step-0 checkpoint: compression applied, training not started. (Exp 14)

Grid Sampler vision-token pruning with the location-aware selection strategy (vtp_type="conv2d_instead_of_global_pool"), plus the one-step Drift objective on the action expert. Pruning: top 60%, wrist 미적용.

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

Why the selection strategy changed

The original sampler predicted its K sampling coordinates from a single globally average-pooled feature vector. That input has no spatial structure, so nothing in it can say which region of the image matters. Measured on VTP20/60 and VTPtoponly20/60 rollouts:

This variant scores every location with a small conv head (1×1 to mix channels, 3×3 to fold in neighbouring context) and takes the top-K cells:

• content-dependent — the score comes from each cell's own features, so selection follows the scene
• zero duplicates, structurally — topk returns distinct indices, with no auxiliary loss needed
• interface unchanged — (B, C, H, W) → (B, K, C), so the backbone reassembly and the processor's placeholder count are untouched

A unit check on two different feature maps: the original strategy changed 0 of 78 selected cells, this one changed 92.

Gradient reaches the score head through a straight-through factor w_sel / w_sel.detach(). The forward value is exactly the gathered feature, and the backward path is d(log w_sel)/d(logits), bounded by [−1, 1] because d(log softmax_i)/d(logit_j) = δ_ij − w_j. There is no temperature: topk is invariant to monotone rescaling, so τ would only scale the gradient — a degree of freedom the final conv's weights already have.

This variant

grid_token_sampler_num_tokens = [78,0] — a list is per camera in image_keys order (cam0 = top, cam1 = wrist), where 0 means that camera is not pruned and bypasses the sampler entirely. An int applies the same budget to both.

Note the sampler parameter count is independent of K: score_head is Conv2d(2560→256, 1×1) + Conv2d(256→1, 3×3), whose size does not depend on how many cells are kept. The old scout_mlp emitted K×2 coordinates and so shrank as pruning got more aggressive — a counterintuitive property this design removes.

Full sweep

Exp 13 and 15 spend the same token budget (352), as do 14 and 16 (274) — so the sweep also contrasts concentrating a budget on one camera against splitting it across both, at two budget levels.

Drift and KeyStone

Drift replaces the action expert's multi-step flow-matching objective: G = 8 candidate chunks are drawn per observation from one shared backbone KV context, each pulled toward the demonstration and pushed away from its siblings, with the drift direction computed per action dimension in fp32. Deployment integrates the velocity field once (num_inference_steps=1). No teacher, no distillation, no new parameters.

KeyStone is inference-only and adds no parameters, so it is not a training flag. This config keeps test_time_samples=1; enable it at eval:

--policy.test_time_samples=8 --policy.test_time_clusters=4

num_flow_timesteps stays at 8. It is a flow-matching training setting and is never read on the drift path — the drift loss expands the batch by drifting_gen_per_label instead.

Compatibility

This checkpoint is not weight-compatible with the vtp_type="original" checkpoints: this one has score_head.*, those have scout_mlp.*. Loading either into the other raises on unexpected/missing keys. vtp_type is recorded in the config so both generations coexist; configs written before the field default to "original".

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: vtp_type recorded in both the model and processor config, budgets [78,0] resolved to 1 sampler module(s) on the right camera(s), score_head.* present and scout_mlp.* absent, drift flags set (use_drifting_loss=true, G=8, per-dim, num_inference_steps=1), and the parameter count matching the analytic prediction (5,443,512,817).

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

Reproducing

# repo: nota-github/xpu-lerobot, branch feat/so101-train (>= the vtp_type commit)
EXPORT_INIT=1 INIT_EXPORT_DIR=/tmp/exp14-init \
  GRID_NUM_TOKENS="[78,0]" VTP_TYPE=conv2d_instead_of_global_pool HF_TOKEN=... \
  scripts/train_molmoact2_so101_drift_grid_sampler.sh

# training run it seeds (set OMP_NUM_THREADS=1: preprocessing is single-threaded in the
# main process and thread contention cost 22x in measurement)
OMP_NUM_THREADS=1 GRID_NUM_TOKENS="[78,0]" VTP_TYPE=conv2d_instead_of_global_pool \
  HF_TOKEN=... WANDB_API_KEY=... scripts/train_molmoact2_so101_drift_grid_sampler.sh

References

• Deng et al. Generative modeling via drifting. arXiv:2602.04770
• Zuo. Drift-VLA: Fast Vision–Language–Action Policies with One-Step Drifting. 2026
• Dai et al. Geometry guided self-consistency for physical AI (KeyStone). arXiv:2605.08638
• Grid Sampler (GridS) active token sampling, ICML 2026
• Fang et al. MolmoAct2: Action Reasoning Models for Real-world Deployment. arXiv:2605.02881