Cosmos3-Edge-Policy-DROID-FastWAM-v1-step25000

Cosmos3-Edge-Policy-DROID — Drift 1-step distillation, step 25000 (중간 체크포인트)

nvidia/Cosmos3-Edge-Policy-DROID를 1-step 생성기로 증류하는 학습(Step 4)의 iter 25000 중간 체크포인트입니다. 계획된 25000 iter 중 100% 지점이며, 학습이 계속 진행 중일 수 있습니다. 평가용 스냅샷이지 최종 산출물이 아닙니다.

• EMA 가중치 기준 export (use_ema_weights: true)
• training state 없음 — optimizer/scheduler/trainer 상태 제외, pretrained model 파일만
• 시퀀스는 full WAM 유지 (video token 제거 없음). video token 제거와 결합한 변형은 Step 5에서 별도로 학습합니다.

무엇을 바꾸는가

baseline 추론은 4-step UniPC × CFG 2분기 = 타워 포워드 8회로 액션을 만듭니다 (B200에서 denoise 110ms / E2E 233ms). Drift objective는 σ=1 고정에서 1-step 생성기를 직접 학습 — 즉 num_inference_steps=1이 배포하는 함수 그 자체를 최적화합니다.

SO-101에서 같은 방식으로 증류했을 때 open-loop MAE가 오히려 개선됐습니다 (26.96 → 5.45): 증류가 샘플링 스텝 수뿐 아니라 CFG로 인한 분포 이탈 문제까지 해결했기 때문입니다.

학습 설정 (요약)

전체 설정은 작업 레포의 docs/droid_training_config.md 참조.

사용법

추론 코드는 cosmos-framework(commit 5e67049) + 패치가 필요합니다. 패치는 geonmin-kim/Cosmos3-Edge-Policy-DROID-FastWAM의 inference/ 폴더에 있습니다.

CUDA_VISIBLE_DEVICES=0 python -m cosmos_framework.scripts.action_policy_server_robolab \
  --checkpoint-path geonmin-kim/Cosmos3-Edge-Policy-DROID-Drift-step25000 \
  --port 8000 --format-prompt-as-json True --no-guardrails

1-step 증류 모델이므로 추론 시 num_inference_steps=1 / guidance_scale=1.0으로 맞춰야 합니다 — 학습된 함수와 다른 스텝 수/CFG로 돌리면 증류 이득이 사라집니다. --drop-generated-vision은 이 체크포인트에 쓰지 마세요 (full WAM 시퀀스로 학습됨).

RoboLab 클라이언트(Isaac Sim, RTX GPU 필수):

./.venv/bin/python policies/cosmos3/run.py \
  --task <TASK> --num-envs 10 --headless --host <server> --port 8000

파일 구성

consolidated safetensors(2 shard) + index, config.json, checkpoint.json(policy 메타데이터: action_chunk_size 32 / conditioning_fps 15.0 / domain_name droid_lerobot), tokenizer·processor 5종 + chat_template.jinja, vision_encoder/. Wan2.2 VAE는 config의 상대 경로로 실행 시 자동 다운로드됩니다. 머신 로컬 절대 경로는 포함돼 있지 않습니다.