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