Cosmos Policy latent 슬롯 구조와 사전학습 분포 괴리

TL;DR

"t와 t+K만 예측하고 중간은 안 한다"는 주장은 사실: 데이터로더가 디스크에서 읽는 프레임은 {t, t+K} 딱 2장. 중간 K-1장은 JPEG 디코딩조차 안 함.
latent 시간축이 "시간"이 아니라 "슬롯 ID": state_t=9가 [blank, proprio, wrist, primary, action, fut.proprio, fut.wrist, fut.primary, value]. 같은 이미지를 4번 복제해 한 슬롯을 채워 video tokenizer를 image tokenizer로 무력화.
가장 심한 OOD는 따로 있음: 액션/proprio/value 슬롯은 VAE를 아예 통과하지 않고 raw 숫자를 16×28×28=12,544 원소에 타일링해 써넣음. 사전학습 DiT가 본 적 없는 latent 값.
결론: "world model을 policy로 쓴다"가 아니라 "video pretraining을 initialization으로 쓴다"가 정확한 표현. dynamics prior는 실제로 사용되지 않음.

사용 프레임 / 샘플

3 / 9

VAE 미통과 슬롯

12,544

latent 원소 / 슬롯

16/32/50

K (LIBERO/RoboCasa/ALOHA)

future denoise step

1 배경 / 목적

Cosmos Policy는 Cosmos-Predict2 2B video2world 모델을 백본으로 쓴다고 알려져 있다. "t ~ t+K 프레임을 예측하는 비디오 모델을 쓰면서, 실제로는 t와 t+K만 예측하고 중간은 예측하지 않는다"는 주장이 있어 코드로 검증하고, 사실이라면 그것이 원래 video model의 학습 분포와 얼마나 어긋나는지를 정량적으로 파악할 필요가 있었다.

왜 중요한가: 이 구조를 정확히 알아야 (a) 레포를 어디까지 재사용할 수 있는지, (b) 중간 프레임을 쓰는 변형이 의미 있는 실험인지 판단할 수 있다.

2 작업 내용

레포 코드를 직접 읽어 latent 시퀀스 구성, 조건/생성 대상 분리, 액션·proprio·value의 latent 주입 방식을 추적했다.

파일	확인 내용
`config/experiment/cosmos_policy_experiment_configs.py:116,236,321`	`state_t`, `chunk_duration`, conditional frame 수
`datasets/robocasa_dataset.py:827-966`	로드하는 프레임 인덱스, latent 슬롯 배치 순서
`utils/utils.py:41`	`duplicate_array` = `np.stack([arr]*4)`
`models/policy_text2world_model.py:52-160`	`replace_latent_with_action` / `_proprio`
`models/policy_video2world_model.py:150-373`	conditional mask 설정, gt_frames 주입
`experiments/robot/cosmos_utils.py:557-620`	latent → 액션/value 역추출
`README.md:20-95`	inference denoising step 설정

3 결과

3.1 latent 시퀀스 레이아웃 (LIBERO, `state_t=9`)

[0] blank [1] proprio [2] wrist [3] primary | [4] action [5] fut.proprio [6] fut.wrist [7] fut.primary [8] value <------------ 조건 (clean, frame_replace) -----------> <------------------- denoise 대상 -------------------> state_t=9 tokenizer.chunk_duration=33 # 1 blank + 32 = 8슬롯 x 4장 num_duplicates_per_image=4 # WAN2.1 tokenizer: 4 images per latent frame min_num_conditional_frames=4 # 1 blank + 3 조건(proprio, wrist, primary) max_num_conditional_frames=4 conditioning_strategy="frame_replace" sigma_conditional=0.0 # 조건 latent에는 노이즈 없음

WAN2.1 tokenizer가 temporal 4:1 압축이므로 픽셀 4장 = latent 1프레임. 여기에 같은 이미지를 4번 복제해 한 슬롯을 채운다. 결과적으로 latent 프레임 1개 = 시간 스텝 1개가 아니라 의미 항목 1개가 된다.

3.2 디스크에서 읽는 프레임은 정확히 2장

# robocasa_dataset.py:827-837 future_frame_idx = relative_step_idx + self.chunk_size max_possible_idx = episode_data["num_steps"] - 1 if future_frame_idx > max_possible_idx: future_frame_idx = max_possible_idx frames_needed = {relative_step_idx, future_frame_idx} # <-- 원소 2개

확인: t+1 … t+K-1은 JPEG 디코딩조차 하지 않는다. 중간 프레임 dynamics는 학습 신호에 전혀 들어가지 않는다.

3.3 (핵심) 액션 / proprio / value는 VAE를 통과하지 않는다

# policy_text2world_model.py:112 계열 latent_elements = latent_channels * latent_h * latent_w # 16 * 28 * 28 = 12,544 num_repeats = (latent_elements + num_action_elements - 1) // num_action_elements repeated_action = flat_action.repeat(1, num_repeats)[:, :latent_elements] new_x0[batch_indices, :, action_indices, :, :] = result

데이터로더는 해당 슬롯에 검은 이미지를 넣어 tokenizer로 인코딩하지만, 모델이 그 latent를 버리고 정규화된 raw 숫자를 타일링해 덮어쓴다. 역추출은 반대로 평균 (cosmos_utils.py:598, 619).

슬롯	원본 원소 수	타일링 반복 횟수	추출 방식
action chunk (16×7)	112	112×	chunk별 reshape 후 평균
proprio (9-dim)	9	~1,394×	—
value (scalar)	1	12,544×	전체 원소 평균

부수 효과: 이 중복 인코딩은 diffusion 노이즈에 대한 robustness를 준다. 동시에 검은 이미지 슬롯의 tokenizer 인코딩 연산은 그대로 낭비된다.

3.4 사전학습 분포와의 괴리 — 4겹

층위	사전학습(Predict2) 분포	Cosmos Policy 입력	OOD 정도
tokenizer 입력	16fps 연속 비디오	동일 이미지 4장 복제 (motion=0)	낮음 — 정지샷은 valid video
DiT 시간축 의미	t, t+1, t+2 … 연속 시간	슬롯 ID (proprio/wrist/primary/action/…)	높음 — temporal RoPE 의미 재정의
DiT 시간축 내용	단일 시점 연속	latent마다 카메라 하드컷 + 검은 프레임	높음
latent 값 자체	VAE encoder 출력 (smooth, 채널 상관)	손으로 써넣은 raw 숫자 타일링	최대

3.5 inference의 denoising step 비대칭

num_denoising_steps_action=5 num_denoising_steps_future_state=1 # <-- 미래 관측은 1스텝 num_denoising_steps_value=1

해석: p(o_t+K | o_t, ℓ)는 중간 경로가 여럿이라 multimodal인데 1스텝만 denoise하므로 사실상 conditional mean(블러)을 뽑는다. 미래 이미지 head는 강한 supervision이 아니라 약한 regularizer로 봐야 한다.

4 Takeaway

"world model을 policy로 쓴다"는 부정확하다

실제로 이전되는 것은 dynamics prior가 아니라 (a) 로봇 씬에 대한 시각 표현, (b) 잘 조건화된 diffusion transformer 초기화, (c) 검증된 latent space다. 정직한 표현은 "video pretraining을 initialization으로 쓴다"에 가깝다.

그럼에도 동작하는 이유: zero-shot이 아니라 full finetuning(2B 전체, lr=1e-4, FSDP, 30K step LR decay cycle)이기 때문. 슬롯 의미론을 새로 배울 용량이 충분하다.

공정하게 평가할 부분: 4장 복제 트릭은 영리하다 — video tokenizer를 재학습 없이 image tokenizer로 전환하는 방법이다. 비판할 지점은 이게 아니라 시간축 재정의와 인공 latent다.

5 Next Steps

미해결 — video 사전학습의 기여도 미검증

dynamics를 안 쓰는데도 video init이 image diffusion init(SD3/FLUX 등) 대비 우위가 있는지에 대한 ablation을 레포에서 찾지 못했다. 이게 없으면 "world model 덕분"이라는 주장은 근거가 약하다. 먼저 논문에 해당 ablation이 있는지 확인.

실험 1 — 중간 프레임 활용

t…t+K를 실제 latent 시퀀스로 넣으면 (a) 사전학습 분포와 재정렬되고 (b) dynamics supervision을 실제로 받는다. 가설: K가 큰 ALOHA(K=50, 30fps 기준 ~1.7초 점프)에서 이득이 가장 클 것. 비용: 시퀀스 길이 증가 → 학습/추론 비용 상승. 먼저 LIBERO(K=16)에서 state_t를 늘려 소규모 검증.

실험 2 — 인공 latent 제거

액션/proprio/value를 별도 head나 cross-attention token으로 빼면 latent 분포 괴리 3개 슬롯이 사라진다. 다만 diffusion으로 액션을 뽑는 구조적 장점(multimodal action distribution 표현)을 잃을 수 있어 trade-off 확인 필요.

측정 안 한 것

검은 이미지 슬롯의 tokenizer 인코딩 낭비 — 실측 오버헤드 프로파일링 미실시.

Cosmos Policy는 비디오를 생성하지 않는다 — latent 슬롯 구조와 사전학습 분포와의 괴리