TL;DR
- 동기: DreamGen P2는 Cosmos로 16-frame(~1s) 청크만 생성 → Pi0.5 B1 학습 신호가 전체 task trajectory를 커버하지 못해 SR 4.7%에 그침
- 접근: Chained autoregressive inference — iter N의 마지막 프레임을 iter N+1의 init frame으로 연결, task별 필요 iter 수를 max episode length에서 역산
- 파일럿 결과: Job 10659 COMPLETED (33min), 전환 자연스러움 확인. 사용자 영상 피드백: 2-iter(11.6s) → CloseDoubleDoor 한 쪽 문만 닫힘 → N=5 필요
- 비용 추정: 1440 ep × 평균 4 iter × 10 min/iter (8 GPU) ≈ 160 GPU-hr (~20 hr wall-clock)
1 배경 / 목적
Pi0.5 B1(DreamGen P2 synthetic-only, 1440 ep)을 24-task 평가한 결과 SR 4.7% — B0 Real(17.2%) 대비 3.7× 열세. 단순히 학습 데이터 품질 문제가 아니라 P2 데이터 구조 자체의 한계가 의심됐다.
P2의 구조적 한계
16-frame chunk 문제: Cosmos Predict2는 한 번 inference 시 93 frame을 생성하지만, 사용 설정상 init frame에서 시작해 16 frame만 유효 학습 데이터로 활용. 대부분의 RoboCasa task는 완료까지 200~800+ frame이 필요한데, 16-frame chunk는 task의 초기 일부만 커버.
padding/truncation 학습 신호: chunk 끝부분에서 policy는 trajectory가 끝난 것처럼 학습 → action이 mid-task에서도 "task end" 패턴 출력 가능.
목표: 기존 P2 생성 파이프라인을 수정해 task 완료까지 연속 프레임을 생성하는 "P3 autoregressive" 데이터셋 구축. 이를 IDM으로 action dump해 Pi0.5 재학습.
2 작업 내용
핵심 아이디어: Chained Autoregressive Inference
기존 P2: init_frame → Cosmos → 93 frames (1회, 16 frame 사용)
P3 방식: init_frame → Cosmos → 93 frames → [last frame] → Cosmos → 93 frames → ... (N회 반복, 전체 사용)
새 스크립트: cosmos-predict2/examples/video2world_gr00t_autoregressive.py
핵심 로직:
- for iter_i in range(num_iterations):
frames = cosmos_inference(init_frame=last_frame_of_prev_iter)
all_frames.extend(frames)
last_frame = frames[-1] # 다음 iter의 init
출력: ep당 (num_iterations × 93) frame 연속 비디오
Launch: scripts/run_cosmos_datagen_autoregressive.sh
task별 num_iterations 계산
MG-3000 dataset의 task별 max episode length (20fps 기준) → Cosmos fps=16 변환 → iter당 93 frame으로 나눠 올림.
공식: iters = ceil(max_ep_len_at20fps × (16/20) / 93)
| Task | max@20fps | Iters | Task | max@20fps | Iters |
|---|
| CloseDoubleDoor | 580 | 5 | OpenDoubleDoor | 868 | 8 |
| CloseDrawer | 268 | 3 | OpenDrawer | 288 | 3 |
| CloseSingleDoor | 347 | 3 | OpenSingleDoor | 426 | 4 |
| CoffeePressButton | 163 | 2 | TurnSinkSpout | 216 | 2 |
| CoffeeServeMug | 512 | 5 | CoffeeSetupMug | 426 | 4 |
| PnPCabToCounter | 489 | 5 | PnPCounterToCab | 376 | 4 |
| PnPCounterToMicrowave | 550 | 5 | PnPMicrowaveToCounter | 442 | 4 |
| PnPCounterToSink | 692 | 6 | PnPSinkToCounter | 374 | 4 |
| PnPCounterToStove | 443 | 4 | PnPStoveToCounter | 429 | 4 |
| TurnOffMicrowave | 325 | 3 | TurnOnMicrowave | 286 | 3 |
| TurnOffSinkFaucet | 442 | 4 | TurnOnSinkFaucet | 405 | 4 |
| TurnOffStove | 287 | 3 | TurnOnStove | 361 | 4 |
평균 ~4 iters/task. 최소 2 (CoffeePressButton, TurnSinkSpout), 최대 8 (OpenDoubleDoor).
실행 이력 (SLURM)
| Job ID | 상태 | Elapsed | 비고 |
|---|
| 9973 | CANCELLED | 00:00 | -c 144 초과 요청 오류 |
| 10628 | PREEMPTED | 34min | sub partition, preempt |
| 10659 | COMPLETED ✓ | 33min | sub-extra 4 GPU, 파일럿 성공 |
| 12043 | PREEMPTED | 10min | sub partition, 재시도 중 preempt |
3 결과
파일럿 검증 결과
전환 자연스러움 확인: Chained autoregressive inference에서 iter 간 프레임 전환이 시각적으로 자연스럽게 이어짐 (AC 통과 기준 cos_sim ≥ 0.85 수준 추정). 14B Cosmos 모델이 마지막 프레임을 conditioning으로 받아 연속성 유지.
사용자 영상 피드백: 2-iter (11.6s) CloseDoubleDoor
N=2 불충분: 2-iter 생성 시 CloseDoubleDoor에서 두 문 중 한 쪽 문만 거의 닫힘. Task 완료 기준(두 문 모두 닫힘)에 미달. CloseDoubleDoor의 계산된 필요 iters는 5 (max_ep=580 frame).
생성 속도 비교
| GPU 수 | 시간/iter | 전체 1440 ep × 4 iter |
|---|
| 4 GPU | ~30min | ~720 hr (불가) |
| 8 GPU | ~10min | ~160 hr (~20 hr wall) |
8 GPU가 4 GPU 대비 3× 빠름. 14B Cosmos 모델 특성상 8 GPU 권장.
PREEMPT 빈도 높음: sub/sub-extra partition에서 core 우선 작업에 의해 자주 선점됨. sbmr (자동 재시도) 필수. 총 4회 시도 중 1회만 완료.
4 Takeaway
DreamGen P3의 의미
Pi0.5 B1(synthetic-only) 4.7% SR의 핵심 원인 중 하나로 P2 데이터의 16-frame chunk 구조가 지목됐다. Autoregressive chained generation이 기술적으로 동작함을 파일럿이 증명했다. P3 데이터셋이 완성되면:
- 각 episode가 task 완료까지의 전체 trajectory를 담은 synthetic video를 가짐
- Pi0.5가 "task 끝까지 어떻게 해야 하는가"를 학습할 수 있음
- B1 SR 4.7% → 유의미한 개선 기대 (B0 Real 17.2% 대비)
비용 대비 기대 효과: ~160 GPU-hr(20 hr wall-clock)로 1440 ep P3 데이터셋 생성 가능. P2 생성(Cosmos 6.5h + IDM 3h) 대비 ~20× 많은 compute지만, 학습 quality 개선 경로로 타당.
5 Next Steps
즉시: 5-iter 파일럿 검증
CloseDoubleDoor N=5로 생성 → 두 문이 모두 닫히는지 확인. 성공 시 build_datagen_seeds.py에 24-task iter dict 반영 후 full batch_input.json 생성.
sbmr 10 "bash scripts/run_cosmos_datagen_autoregressive.sh" \
--gres=gpu:8 -c 112 --mem 1600GB --qos=core-extra
# 5-iter CloseDoubleDoor pilot → 영상 검증
P3 full-scale 생성 (검증 통과 후)
1440 ep × task별 2~8 iter. 예상 wall-clock ~20 hr (8 GPU, sbmr 자동 재시도). PREEMPT 빈도 고려해 sbmr retry 10 이상 설정. 완료 후 IDM action dump 필요.
IDM State 통합 (Step 3 병행)
P3 video에서 Latent-IDM A v2로 state field도 함께 dump (action 12-dim + state 16-dim). 단, heldout schema (8-dim joint+gripper) vs v2 schema (base_to_eef 중심) 불일치 해소 필요 → 전용 IDM 학습 또는 schema 변환 레이어. P3 + IDM-state → Pi0.5 B2' 학습.
미해결 한계
- Cosmos가 생성한 long-horizon video가 실제로 task를 완료하는 동작을 보이는지 미검증 (5-iter 파일럿 목적)
- task별 num_iterations은 max episode length 기반 — 실제 필요 iters는 task difficulty에 따라 다를 수 있음 (OpenDoubleDoor 8 iters 추정 → 실제 더 필요할 수도)
- sub partition PREEMPT 빈도 → core partition 잔여 용량 확인 후 core-extra 우선 사용 권장