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).