MolmoAct2-LIBERO-Drift-Keystone — 1-step action expert
allenai/MolmoAct2-LIBERO의 flow-matching action expert를 Drift(one-step drifting loss,
arXiv:2602.04770)로 재학습하여 추론 시 Euler 적분을 10-step → 1-step으로 줄인 VLA.
KeyStone(기하 자기일관성 best-of-K 선택, arXiv:2605.08638)을 추론 옵션으로 지원.
VLM 백본·expert 구조·체크포인트 포맷은 원본과 동일(추가 파라미터 0). LeRobot 포맷.
성능 (LIBERO 4-suite × task당 20 ep = 800 ep, seed 1000)
| libero_10 | goal | object | spatial | 전체 | |
|---|---|---|---|---|---|
| 원본 baseline (10-step FM) | 95.0 | 95.5 | 99.5 | 98.5 | 97.1 |
| 본 모델 (1-step) | 94.5 | 97.5 | 99.5 | 99.0 | 97.6 |
| 본 모델 (1-step + KeyStone K=8/C=4) | 95.5 | 96.5 | 100.0 | 99.5 | 97.9 |
Latency (B200, batch=1, eager): action head 228.3 → 26.9 ms (8.5×↓), 전체 decision 260.6 → 56.6 ms (4.6×↓), KeyStone +3.7 ms.
학습
- 코드: nota-github/xpu-lerobot @ feat/molmoact2-drift-keystone
- 레시피: drifting loss (G=8, action-dimension grouping, temps {0.02, 0.05, 0.2}, flow time t=0 pin),
allenai/MolmoAct2-LIBERO-Dataset, 10k steps, eff batch 32, bf16, seed 1000 - 스크립트:
scripts/train_molmoact2_libero_drift_ddp2.sh
사용 (lerobot, 위 브랜치 필요 — config의 drifting_/test_time_ 키 파싱)
# 1-step 평가
POLICY_PATH=xpuenabler/MolmoAct2-LIBERO-Drift-Keystone \
scripts/eval_molmoact2_drift_libero.sh full
# KeyStone best-of-8 (거의 공짜: +3.7ms)
POLICY_PATH=xpuenabler/MolmoAct2-LIBERO-Drift-Keystone KEYSTONE_K=8 KEYSTONE_C=4 \
scripts/eval_molmoact2_drift_libero.sh full
핵심 플래그: --policy.num_inference_steps=1 (필수 — 학습된 생성기와 일치),
--policy.test_time_samples=8 --policy.test_time_clusters=4 (KeyStone, 선택).
평가 상세: Confluence "Flow-matching steps 감소를 위한 Drift 적용 및 성능 평가" (page 2426404865).