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

Cosmos3-Edge-Policy-DROID — FastWAM-K (부분 비전 유지), step 12000

nvidia/Cosmos3-Edge-Policy-DROID(full WAM, RoboLab 36.5%)에서 warm start한 부분 비전 유지(leading-K) 실험의 중간 체크포인트입니다. 계획된 12000 iter 중 100% 지점이며, 평가용 스냅샷이지 최종 산출물이 아닙니다.

• step별 하위 폴더(step12000/)에 누적. 루트 README는 최근 step을 가리킵니다.
• EMA 가중치 기준 export, training state 없음 — pretrained model 파일만.

⚠️ 서빙 시 필수 인자

이 체크포인트는 conditioning 프레임 + 선행 생성 2 latent 프레임을 유지하도록 학습됐습니다. 기본 FastWAM 플래그로 서빙하면 학습과 다른 시퀀스가 배포됩니다.

python -m cosmos_framework.scripts.action_policy_server_robolab \
    --checkpoint-path "$ckpt" --port 8000 \
    --format-prompt-as-json True --no-guardrails \
    --keep-generated-vision-frames 2        # ← --drop-generated-vision 이 아님

교차 호스트 합의 불변식: num_kept = (실제 conditioning prefix latent 프레임 수) + K, 실제 총 latent 프레임 수로 클램프. 유지된 K 프레임은 conditioned가 아니고 denoise 대상입니다(pack_vision_tokens가 plan의 condition_frame_indexes_vision으로 마스크를 만들어 프레임 1..K가 mse_loss_indexes에 들어갑니다) — 추론 시 미래 ground truth가 없으므로 그래야 배포가 성립합니다.

왜 K인가

RoboLab 24 arm / 2,304 에피소드 결과, FastWAM(K=0)의 full-WAM 대비 격차는 균일한 열화가 아니라 국소적입니다.

태스크	FastWAM (K=0)	full WAM (K=8)
RubiksCubesInBin	3.1%	62.5%
StackYellowOnRed	6.2%	50.0%
UnstackRubiksCube	100%	87.5%
ToyInBin	29.7%	25.0%

앞의 두 태스크가 13.6%p 격차의 8%p 이상을 설명하고, 뒤의 둘은 baseline을 넘습니다. 잃은 두 태스크의 공통점은 동일 형태 물체 여러 개 중 특정 개체를 지목·조작하는 것입니다. 그 능력은 full-WAM base에 있고 K=0 체크포인트에서는 이미 파괴됐으므로, base를 nvidia로 잡아 되살리는 대신 유지합니다. 부수 효과로 K=8이 학습 없는 앵커가 됩니다(base 자체).

leading-K를 씁니다 — conditioning prefix + 시간적으로 가장 가까운 생성 K프레임. 등간격 스트라이드·마지막 K개·풀링 요약 토큰은 다른 곡선을 주며, 이 저장소의 수치는 전부 leading-K에 대한 것입니다. K=2는 4:1 시간 압축에서 8 픽셀 프레임 = 32스텝/2.13초 action chunk의 앞 ~25%로, RoboLab 클라이언트가 chunk를 소진하기 전에 재요청하는 구간입니다.

토큰 수

트랜스포머에 실제로 들어가는 vision 토큰은 프레임당 340입니다.

K	vision 프레임	vision 토큰	시퀀스 전체(실측)
0 (FastWAM)	1	340	531
2 (이 저장소)	3	1,020	1,211
8 (full WAM)	9	3,060	3,251

전부 RoboLab 서버에서 패킹된 시퀀스를 직접 계측한 값입니다(2026-09-08). 340 × (K+1)이 모든 K에서 정확히 맞고, 시퀀스 전체는 여기에 text와 action 33이 더해진 값입니다 (그 실행의 text는 158; 프롬프트 길이에 따라 변합니다 — 제 학습 배치에서는 164였습니다). mse_loss_indexes가 정확히 340 × K이고 condition_mask가 모든 K에서 1이라는 것도 함께 확인됐습니다 — 유지한 K 프레임 전부가 denoise 대상이고 조건 프레임은 하나입니다.

391이 아닌 이유 (파이프라인 단계 구분)

학습 dataloader의 video 텐서를 실측하면 (3, 33, 544, 736)이고, 그 latent는 34×46 → patch_spatial=2 → 17×23 = 391입니다. 이 값 자체는 맞지만 시퀀스 토큰 수가 아닙니다.

540×640 (원본)
  → transform 반사 패딩        544×736   ← 버킷 VIDEO_RES_SIZE_INFO['480']['4,3']
  → VAE (/16)                  latent 34×46   (= patch 391)
  → _remove_padding_from_latent  latent 34×40   ← image_size=[544,736,540,640] 의 콘텐츠 영역
  → pack_vision_tokens         17×20 = 340/프레임   ← 트랜스포머 입력

패딩은 VAE에 통과시키기 위한 임시 캔버스이고 패킹 전에 제거됩니다 (omni_mot_model.py의 get_data_and_condition() 안, data_batch["image_size"]가 있을 때). 그 함수는 학습·추론이 공유하므로 양쪽 다 340으로 패킹됩니다 — 불일치가 없습니다. 학습 배치에 image_size = [544., 736., 540., 640.]가 실제로 들어 있는 것을 확인했습니다.

따라서 원 모델 카드의 "프레임당 340 / 9프레임 3,060"은 옳습니다. 이 저장소의 이전 버전이 391로 적었던 것은 파이프라인 단계를 명시하지 않아 캔버스 기준 값을 시퀀스 토큰 수로 쓴 것이었습니다. 카드의 text 128만 실측(158~164, 프롬프트 의존)과 다릅니다.

학습 설정

항목	값
experiment	action_policy_droid_edge_ft_fastwam_v1lr + optimizer.lr=1e-5 + K override
K	two_pass_keep_generated_vision_frames=2
base	nvidia/Cosmos3-Edge-Policy-DROID (DCP 변환), 가중치만
two-pass	pass B action FM (action_loss_weight=10), pass A vision FM (loss_scale=1)
데이터	DROID success + failure, keep_ranges_1_0_1 (sim 없음)
optimizer	FusedAdamW, lr 1e-5, 액션 브릿지 3개에 5×
scheduler	LambdaLinear, warm-up 1000, cycle 12000 — 전 구간 학습 (lr 1e-5 → 8.36e-6@3k → 5.91e-6@6k → 3.45e-6@9k → 1.0e-6@12k)
하드웨어	B300 ×2, FSDP shard 2, grad_accum 2 → 글로벌 배치 64

평가 주의

• loss 곡선으로 판정하지 마십시오 — 매칭형 목적함수라 곡선이 평탄합니다.
• open-loop MAE는 이 프로젝트에서 세 번 방향을 반대로 가리켰습니다. 판정 기준은 RoboLab입니다.
• 2태스크(RubiksCubesInBin·StackYellowOnRed) × 8 에피소드 스크리닝은 판별용이며 벤치마크 점수가 아닙니다. 양성이면 12태스크 전체 확인이 필요합니다.
• n=96에서 ±9pp는 노이즈입니다. 체크포인트 비교는 적응 샘플링 (--num-episodes-adaptive 200 --ci-pp-width 0.14)으로 하십시오.