Step-0 checkpoint: compression applied, training not started. (Exp 15)
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·wrist 각 10%.
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·wrist 각 10%
tokens kept, top camera
176 / 196
tokens kept, wrist camera
176 / 196
image tokens total (2 cams)
352 / 392 (10.2% fewer)
sampler modules
2
new parameters
2,633,218 (1,316,609 per sampler)
total parameters
5,444,829,426
grid_token_sampler_num_tokens = 176 — 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.
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:
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 176 resolved to 2 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,444,829,426).
Not smoke-tested through training steps — no training was run for this variant.
Reproducing
bash
1# repo: nota-github/xpu-lerobot, branch feat/so101-train (>= the vtp_type commit)2EXPORT_INIT=1INIT_EXPORT_DIR=/tmp/exp15-init \3GRID_NUM_TOKENS="176"VTP_TYPE=conv2d_instead_of_global_pool HF_TOKEN=... \4 scripts/train_molmoact2_so101_drift_grid_sampler.sh
56# training run it seeds (set OMP_NUM_THREADS=1: preprocessing is single-threaded in the7# main process and thread contention cost 22x in measurement)8OMP_NUM_THREADS=1GRID_NUM_TOKENS="176"VTP_TYPE=conv2d_instead_of_global_pool \9HF_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
