Grid Sampler vision-token pruning at a conservative 40.8% ratio (116 of 196 tokens kept per camera) plus the one-step Drift objective on the action expert. This is the initialization state of a MolmoAct2 SO-101 fine-tuning run, published so the run is reproducible from step 0.

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

Why the two compose

They cut different halves of the forward pass, so stacking compounds rather than overlaps. From the LIBERO breakdown (B200, batch 1):

Stage	Baseline	Cut by
Vision	6.89 ms	Grid Sampler (196 → 116 tokens per camera)
LLM	19.18 ms	Grid Sampler (shorter prefix to attend over)
Action head	228.27 ms	Drift (8–10 integration steps → 1)
Total	260.62 ms

Drift alone took that total to 56.64 ms while leaving Vision + LLM at 6.67 + 18.77 ms — precisely the remainder Grid Sampler attacks.

Grid Sampler replaces each camera's fixed pooled patch grid with 116 learned sample points: a scout MLP predicts 116 normalized (x, y) coordinates from the globally pooled feature map, features are bilinearly sampled there, and a coordinate encoder adds a positional embedding.
Drift replaces the action expert's flow-matching MSE with a one-step objective: G = 8 candidate chunks per observation from one shared backbone KV context, each pulled toward the demonstration and pushed away from its siblings, drift direction computed per action dimension in fp32. No teacher, no distillation, no new parameters.
KeyStone is inference-only: draw K one-step candidates, k-means them, execute the medoid of the dominant cluster with a unimodality guard. Not a training flag — enable at eval with --policy.test_time_samples=8 --policy.test_time_clusters=4.

Where this sits in the sweep

Drift + Grid Sampler at four pruning ratios:

Pruning	Tokens/camera	Drift + Grid Sampler checkpoint
20.4%	156	VTPprun20-DriftKeystone
40.8%	116	VTPprun40-DriftKeystone
60.2%	78	VTPprun60-DriftKeystone
91.8%	16	DriftKeyStone-VTP

Single-method reference: Drift only, no pruning.

This variant: 116 tokens/camera, 40.8% pruned, Drift on.

Honest expectation

On LIBERO each method held parity on its own (Drift 97.6, GridSampler K=16 97.1, control 97.9), and GridSampler K=8 collapsed to 92.0 — so the pruning axis is known to have a cliff, and it was never mapped between 196 and 16 on any embodiment.

Neither the combination nor these intermediate ratios have been measured anywhere. Two simultaneous approximations can interact in ways neither shows alone. Compare this checkpoint's trained result against:

the same pruning ratio without Drift, and
Drift without pruning,

so that a regression can be attributed to one axis rather than guessed at.

Config deltas vs. `lerobot/MolmoAct2-SO100_101-LeRobot`

use_drifting_loss             : false -> true
drifting_gen_per_label        :       -> 8
drifting_perdim_loss          :       -> true      # action-dimension grouping
drifting_temperatures         :       -> (0.02, 0.05, 0.2)
num_inference_steps           : null  -> 1         # single Euler step
use_grid_token_sampler        : false -> true
grid_token_sampler_num_tokens :       -> 116
action_mode                   : continuous -> both
setup_type                    : "single so100/so101 robotic arm in molmoact2"
                                -> "single SO-101 5-DOF robotic arm with gripper"

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. num_inference_steps=1 is what makes deployment one-step.

chunk_size and n_action_steps stay at 30: apply_norm_tag_metadata() pins both from the so100_so101_molmoact2 norm-tag metadata and runs after CLI overrides, so passing a different value has no effect. The SO-101 baselines this is compared against were trained the same way.

The token budget lives in the processor config as well as the model config. Both are 116 here; loading this checkpoint with a different grid_token_sampler_num_tokens will not work.

New parameters: the Grid Sampler's scout MLP + coordinate encoder (2,088,936, randomly initialized). Drift contributes none — the parameter count is identical to the Grid-Sampler-only variant at the same K.

Normalizer statistics come from the SO-101 training corpus below.

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 pins a fixed 6 : 2 : 3 : 1 mix per batch is available but opt-in, via CATEGORY_META=<sidecar> (needs 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.

Status of verification

Verified at construction: the policy instantiates with the sampler present, the model config and processor config agree on 116 tokens, and the parameter count matches the analytic prediction for this K.

It has not been smoke-tested through training steps — no training was run for this variant. The K=16 equivalent (DriftKeyStone-VTP-Init) was smoke-tested and trains cleanly, and K only changes a layer width, so the risk is low but not zero.

Reproducing / continuing

bash

# repo: nota-github/xpu-lerobot, branch exp/so101-train

# regenerate this exact init checkpoint
EXPORT_INIT=1 INIT_EXPORT_DIR=/tmp/vtpprun40-drift-init GRID_NUM_TOKENS=116 HF_TOKEN=... \
  scripts/train_molmoact2_so101_drift_grid_sampler.sh

# the training run it seeds (2 GPUs, 18k steps, per-device batch 16)
GRID_NUM_TOKENS=116 HF_TOKEN=... WANDB_API_KEY=... \
  scripts/train_molmoact2_so101_drift_grid_sampler.sh

# eval with KeyStone on
POLICY_PATH=<ckpt>/pretrained_model KEYSTONE_K=8 KEYSTONE_C=4 \
  scripts/eval_molmoact2_drift_libero.sh full

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-VTPprun40-DriftKeystone-Init

MolmoAct2-SO101-VTPprun40-DriftKeystone-Init