geonmin-kim/Cosmos3-Edge-Policy-DROID-FastWAM-v2

Cosmos3-Edge-Policy-DROID — FastWAM two-pass + IsaacSim (v2), step 3000

nvidia/Cosmos3-Edge-Policy-DROID의 two-pass FastWAM fine-tune (v2) 중간 체크포인트입니다. 계획된 100000 iter 중 3% 지점이며 학습은 계속 진행 중일 수 있습니다. 평가용 스냅샷이지 최종 산출물이 아닙니다.

• 이 저장소는 step별 하위 폴더(step3000/)에 체크포인트를 누적합니다. 루트 README는 가장 최근에 업로드된 step을 가리킵니다.
• EMA 가중치 기준 export (use_ema_weights: true)
• training state 없음 — optimizer/scheduler/trainer 상태는 업로드하지 않고 pretrained model 파일만 포함합니다. 따라서 다른 호스트에서의 "재개"는 정확한 resume이 아니라 가중치 warm start입니다 (optimizer 모멘트 초기화, lr 스케줄 재시작).

무엇을 학습하는가

--drop-generated-vision(FastWAM식 추론)은 WAM 시퀀스에서 미래 영상 토큰 약 3,060개를 제거해 latency를 233 ms → 185 ms(−21%)로 줄이지만, 학습 없이 적용하면 정확도가 붕괴합니다 (open-loop MAE 0.111 → 0.197, state-hold 수준). v2는 그 배포 시퀀스에 모델을 적응시키되, 영상 예측 능력을 버리지 않도록 매 스텝 두 개의 pass를 함께 돕니다.

pass	시퀀스	손실	가중치
B	[text | cond-frame vision | action] (배포와 동일하게 pruned)	action flow-matching	action_loss_weight = 10.0
A	33-frame 전체 latent clip	vision flow-matching	loss_scale = 1.0

pruning이 데이터셋이 아니라 모델 내부(inference_drop_generated_vision)로 옮겨졌기 때문에 데이터셋은 전체 클립을 그대로 내보내고 (action_only_vision=false, encode_exact_durations=[33]), 두 shape가 매 스텝 번갈아 나타납니다.

데이터

풀	비율	내용
DROID	3	nvidia/Cosmos3-DROID success + failure 전체, 공식 keep_ranges_1_0_1 윈도 필터 적용. failure 에피소드는 video 목적함수에만 기여합니다 (is_failure 태그 → action flow-matching 항에서 가중치 0, 단 action token은 시퀀스에 유지되어 pass A의 video 분기가 action-conditioned로 남음).
sim	1	nvidia/PhysicalAI-Robotics-Manipulation-SingleArm (Franka Panda, IsaacSim) 4,211 에피소드 / 300,665 프레임을 DROID 스키마로 변환. keep-ranges 필터 없음(해당 파일은 DROID의 gs:// URI로 키가 잡혀 있음), success split만 존재.

sim 비율을 3:1로 보수적으로 잡은 이유: sim 로봇은 프레임당 이동량이 DROID 텔레오퍼레이션의 약 6배(0.0386 vs 0.0061 rad)이고 DROID 정규화 범위를 벗어난 관절 포즈를 방문합니다 (joint[4]가 프레임의 14.8%에서 min-max 밖). sim 비중을 더 키우면 정책이 큰 스텝 쪽으로 끌립니다. 변환기는 소스의 관절 델타를 절대 목표값으로 적분합니다 — 데이터셋 카드의 "joint motion"을 액면 그대로 받았다면 액션 레이블이 조용히 오염됐을 항목입니다.

RankPartitionedDataLoader는 rank 단위로 분할하므로 4 GPU에서 3:1은 "세 rank가 DROID, 한 rank가 sim"을 의미합니다 — 비율의 해상도가 world size입니다.

학습 설정

항목	값
experiment	action_policy_droid_edge_ft_fastwam_v2
warm start	geonmin-kim/Cosmos3-Edge-Policy-DROID-FastWAM-v2-step3000 (가중치만, optimizer 상태 없음)
max_iter	100000 (체크포인트 3,000 iter마다)
optimizer	FusedAdamW, lr 2e-4, action2llm/llm2action/action_modality_embed에 5x multiplier
scheduler	LambdaLinear, warmup 250, f_max 1.0 → f_min 0.1 (cycle 25,000). 25k 이후 lr 2e-5로 유지되는 장기 저-lr 구간
학습 대상	gen expert + 브릿지 — 백본 VLM·VAE 동결
정밀도 / 병렬화	bfloat16, FSDP data_parallel_shard_degree=4, activation checkpointing full
하드웨어	NVIDIA B300 ×4

사용법

이 저장소는 step별 하위 폴더 구조라 서버가 repo id만으로 자동 다운로드하지 못합니다. snapshot_download로 원하는 step만 받아 로컬 경로를 넘기십시오.

from huggingface_hub import snapshot_download

path = snapshot_download(
    "geonmin-kim/Cosmos3-Edge-Policy-DROID-FastWAM-v2",
    allow_patterns="step3000/*",
)
ckpt = f"{path}/step3000"

# RoboLab 정책 서버 (RTX GPU 필수 — Isaac Sim 렌더러가 RT 코어를 요구합니다)
python -m cosmos_framework.scripts.action_policy_server_robolab \
    --checkpoint-path "$ckpt" --port 8000 \
    --format-prompt-as-json True --no-guardrails --drop-generated-vision

평가에 대한 경고

• loss 곡선으로 판정하지 마십시오. 매칭형 목적함수라 성공/실패 모두 곡선이 평탄합니다.
• open-loop MAE는 신뢰하지 마십시오. 이 프로젝트에서 세 번 방향을 반대로 가리켰습니다. 스모크 테스트 용도로만 쓰고, 판정 기준은 RoboLab입니다.