geonmin-kim/Cosmos3-Edge-Policy-DROID-FastWAM-v1lr-step6000

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

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

• 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이 배포하는 함수 그 자체를 최적화합니다.

	baseline (4-step)	이 체크포인트 (1-step)
타워 포워드	8회	1~2회
E2E latency	233 ms	측정 예정
open-loop MAE	0.1108	측정 예정
RoboLab success rate	측정 예정	측정 예정

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

학습 설정 (요약)

항목	값
experiment	action_policy_droid_edge_drift
warm start	nvidia/Cosmos3-Edge-Policy-DROID (릴리스 policy 체크포인트, DCP 변환)
데이터	nvidia/Cosmos3-DROID success 서브셋, train 3,332 에피소드 (held-out 128 제외)
drift	use_drifting_loss=true, G(drifting_gen_per_label)=4, temperatures (0.02, 0.05, 0.2), weight 10.0
optimizer	FusedAdamW, lr 1e-5, wd 0.05, betas (0.9, 0.99)
학습 대상	gen expert + 브릿지(moe_gen, *2llm, llm2*, time_embedder) — 백본 VLM·VAE 동결
scheduler	LambdaLinear, warmup 250, f_max 1.0 → f_min 0.1
병렬화	FSDP shard 4 (4×B200), bf16, batch 4/rank (G=4배 확장 반영)
action	chunk 32, fps 15, joint_pos 8-D, concat_view

전체 설정은 작업 레포의 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-step6000 \
  --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의 상대 경로로 실행 시 자동 다운로드됩니다. 머신 로컬 절대 경로는 포함돼 있지 않습니다.