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

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

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

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

⚠️ 서빙 시 필수 인자

이 체크포인트는 conditioning 프레임 + 선행 생성 1 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 1        # ← --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=1인가 (2026-09-09)

K=2 step9000이 12태스크 프로토콜에서 **task-mean 34.38%, 95% CI [26.17, 42.58]**를 냈습니다 — K=0 plateau(밴드 상한 24.97)를 확실히 넘었고 baseline 36.46%와는 구분되지 않습니다(Fisher p=0.880). 지연시간은 실측 730.7 ms, 무플래그 baseline 1711.2 ms 대비 2.34×입니다.

남은 배포 질문은 더 줄여도 되는가입니다. K=1의 zero-shot 스윕 실측은 627.6 ms이고, K=2와 비슷하게 회복한다면 K=2를 대체합니다. 이 저장소가 그 검정입니다.

K=1의 배포 배수는 아직 확정하지 않습니다. 비율의 분자·분모는 비교 대상인 두 정확도가 측정된 것과 같은 구성에서 와야 합니다. K=2의 2.34×는 그 조건을 만족합니다 — 분모 1711.2 ms는 36.5%를 낸 무플래그 서빙, 분자 730.7 ms는 34.4%를 낸 step9000입니다. K=1은 정확도 측정이 아직 없으므로 짝이 될 분자가 없고, 627.6 ms를 --keep-generated-vision-frames 8 측정치(1771.9 ms)로 나눈 2.82×는 정확도를 한 번도 측정하지 않은 구성에 대한 속도비입니다.

이 구분이 중요한 이유: 같은 구성의 실행 간 변동은 0.59%(1711.2 대 1721.3)인데 구성을 바꾸면 3.55%(1711.2 대 1771.9)로 6배입니다. "분자·분모를 같은 실행에서"라는 일반 규칙보다 구성 일치가 우선하는 자리입니다. 그리고 짝을 잘못 고르면 배수가 커지는 방향이라 눈에 덜 띕니다.

주의: 학습 없는 K=1은 회복하지 않습니다 — zero-shot K=1은 두 태스크에서 0/16이고 "baseline의 절반 회복" 가설이 p=0.005로 배제됐습니다. K=2도 3,000 iteration에서는 노이즈였고 6,000에서 전환됐습니다. 따라서 이 arm의 3k 지점을 근거로 K=1을 판정하면 안 됩니다.

왜 leading-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=1은 4:1 시간 압축에서 4 픽셀 프레임 = 32스텝/2.13초 action chunk의 앞 ~12%입니다.

토큰 수

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

K	vision 프레임	vision 토큰	시퀀스 전체(실측)
0 (FastWAM)	1	340	531
2 (직전 arm)	3	1,020	1,211
1 (이 저장소)	2	680	871
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=1
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 (GPU 3,7), 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)으로 하십시오.