MolmoAct2-SO101-VTPprun20-DriftKeystone-Init

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

Grid Sampler vision-token pruning at a conservative 20.4% ratio (156 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 → 156 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 156 learned sample points: a scout MLP predicts 156 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: 156 tokens/camera, 20.4% 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:

1. the same pruning ratio without Drift, and
2. 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 :       -> 156
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 156 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,129,976, 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 156 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

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

# regenerate this exact init checkpoint
EXPORT_INIT=1 INIT_EXPORT_DIR=/tmp/vtpprun20-drift-init GRID_NUM_TOKENS=156 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=156 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