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:

	measured	meaning
coordinate drift between frames	2.2–4.3 px	a patch is 16 px — positions barely moved
success vs failure difference	1.3–2.3 px	frame noise; no content dependence
duplicate tokens	11–29%	K=156 gave only ~111 distinct positions
cell coverage	37–64%	a third to two thirds of the image never sampled

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

	value
pruning	top 60%, wrist 미적용
tokens kept, top camera	78 / 196
tokens kept, wrist camera	196 (미적용, sampler 우회)
image tokens total (2 cams)	274 / 392 (30.1% fewer)
sampler modules	1
new parameters	1,316,609 (1,316,609 per sampler)
total parameters	5,443,512,817

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	pruning	image tokens	samplers	checkpoint
13	top 20%	352	1	locaware_topp20
14	top 60%	274	1	locaware_topp60
15	both 10%	352	2	locaware_top10_wristp10
16	both 30%	274	2	locaware_top30_wristp30

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

bash

# 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

geonmin-kim/MolmoAct2-SO101-VTP_locaware_topp60_Init

MolmoAct2-SO101-VTP_locaware_topp60_Init