--- license: apache-2.0 library_name: lerobot pipeline_tag: robotics tags: - robotics - lerobot - molmoact2 - vla - so101 - drift - keystone - one-step-flow-matching base_model: lerobot/MolmoAct2-SO100_101-LeRobot --- # 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=` (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 ```bash # 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=/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