MolmoAct2-SO101-DriftKeyStone-VTP-Init

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

This is the initialization state of a MolmoAct2 SO-101 fine-tuning run with Drift and Grid Sampler stacked. It is not a trained policy — it is the exact artifact that lerobot-train holds in memory before its first optimizer step, published so the retraining run is reproducible from step 0.

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

Single-method siblings: Drift only · Grid Sampler only

Why the two compose

They cut different halves of the forward pass:

Drift alone brought that total to 56.64 ms — and left Vision + LLM untouched at 6.67 + 18.77 ms, which is precisely the remainder Grid Sampler attacks. So the two are expected to compound rather than overlap: Drift removes the dominant cost, Grid Sampler shrinks what is left.

• Grid Sampler replaces each camera's pooled patch grid with K learned sample points: a scout MLP predicts K 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 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. No teacher, no distillation stage, 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.

Honest expectation

Each method held parity on its own on LIBERO (800 episodes/model):

The combination has not been measured on any embodiment. Two simultaneous approximations can interact in ways neither shows alone, and Grid Sampler is the more fragile of the pair — dropping K from 16 to 8 cost 5.1 points. Nothing in the per-method results licenses assuming the stack is also at parity.

So: treat this variant as the one most in need of its own evaluation, and compare it against the single-method runs, not only against the baseline. If it loses accuracy, the single-method numbers are what tell you which half to back off.

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 :       -> 16
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. They are not free parameters here: apply_norm_tag_metadata() pins both from the so100_so101_molmoact2 norm-tag metadata in the base checkpoint, and it runs after CLI overrides, so passing a different value has no effect. The SO-101 baseline runs this checkpoint is meant to be compared against were trained the same way.

The token budget lives in the processor config as well as the model config — the preprocessor must emit the matching number of image placeholder tokens. Both are set consistently here; loading this checkpoint with a different grid_token_sampler_num_tokens will not work.

Normalizer statistics come from the SO-101 training corpus below, not from the base checkpoint's norm tag.

Training data

Merged SO-101 corpus, 177,661 frames from 11 dataset repos across four categories:

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 instead pins a fixed 6 : 2 : 3 : 1 mix in every batch is available but opt-in, via CATEGORY_META=<sidecar>; it needs batch_size >= 12 for every category to get a slot, and warns if a category is rounded out. Normalizer statistics are computed over the whole corpus and are unaffected by the sampling mode, so this checkpoint is valid for either.

Reproducing / continuing

# repo: nota-github/xpu-lerobot, branch feat/molmoact2-so101-compression
#   (= exp/so101-train, which carries the SO-101 joint-frame support, with the
#    Drift+KeyStone commits from feat/molmoact2-drift-keystone cherry-picked on top)

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

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

# sweep the token budget
GRID_NUM_TOKENS=32 HF_TOKEN=... 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