MolmoAct2-SO101-DriftKeyStone-Init

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

This is the initialization state of a MolmoAct2 SO-101 fine-tuning run whose action expert has been switched from multi-step flow matching to the one-step Drift objective. 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.

To use a trained policy, take a checkpoint from the training run this one seeds.

What the compression does

	Baseline	This variant
Action-expert training objective	flow-matching MSE over num_flow_timesteps	Drift loss (attraction to demo + repulsion between sibling samples)
Inference	integrate the velocity field 8–10×	1 Euler step (num_inference_steps=1)
New parameters	—	none
Backbone / vision encoder / KV-cache bridge / action interface	unchanged	unchanged

Drift trains a one-step generator directly, so there is no teacher model and no distillation stage. G = 8 candidate chunks are drawn per observation from one shared backbone KV context; each is pulled toward the demonstration and pushed away from its siblings, with the drift direction computed per action dimension in fp32. The discrete (autoregressive) supervision path is untouched — Drift applies only to the continuous action expert.

KeyStone (test_time_samples / test_time_clusters) is a geometry-guided test-time selector: draw K one-step candidates, k-means them, execute the medoid of the dominant cluster, with a unimodality guard. It is inference-only, adds no parameters, and is therefore not enabled in this checkpoint's config — turn it on at eval time.

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          # deployed as a single Euler step
action_mode              : continuous -> both
setup_type               : "single so100/so101 robotic arm in molmoact2"
                           -> "single SO-101 5-DOF robotic arm with gripper"
control_mode             : "absolute joint pose" (unchanged)

num_flow_timesteps is deliberately left 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. The knob that makes deployment one-step is num_inference_steps=1.

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.

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

Reference results (LIBERO, same compression, different embodiment)

From docs/results_drift_keystone_libero.md in the training repo — 4 suites × 10 tasks × 20 episodes = 800 episodes per model:

Model	Inference	Overall success
MolmoAct2-LIBERO baseline	10-step FM	97.1%
no-prune control, same recipe	10-step FM	97.9%
Drift	1-step	97.6%
Drift + KeyStone (K=8, C=4)	1-step ×8	97.9%

Module latency on one B200 (batch 1, eager, CUDA graph off):

Config	Vision	LLM	Action head	Total
10-step FM	6.89 ms	19.18 ms	228.27 ms	260.62 ms
Drift 1-step	6.67 ms	18.77 ms	26.88 ms	56.64 ms
Drift + KeyStone K=8	6.71 ms	19.14 ms	28.11 ms	60.32 ms

Action head 8.5× faster, end-to-end 4.6× faster. These are LIBERO numbers; SO-101 parity has not yet been measured — that is what the run seeded by this checkpoint is for.

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

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-init HF_TOKEN=... \
  scripts/train_molmoact2_so101_drift_ddp2.sh

# the training run it seeds (2 GPUs, 18k steps, per-device batch 16)
HF_TOKEN=... WANDB_API_KEY=... scripts/train_molmoact2_so101_drift_ddp2.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
• Fang et al. MolmoAct2: Action Reasoning Models for Real-world Deployment. arXiv:2605.02881