Pi0.5 P3 Mix NEW IDM — Droid IDM Ablation 3 Variants 완료

TL;DR

목적: IDM 품질(= training/inference visual 분포 정합)이 downstream Pi0.5 policy SR에 미치는 영향을 단일 변수 ablation으로 정량화
설계: 코드 / config / main dataset (mg30_real_h50 720ep) / seed 완전 동일 — aux dataset만 OLD _v2 vs NEW _v2_droid 교체
완료: init0 (Job 2077, 5h02m) / mid (Job 2078, 4h59m) / random (Job 2079, 4h54m) 모두 COMPLETED. 24-task PoC eval 필요.

3/3

Jobs COMPLETED

20K

Steps/variant

4 GPU

Per job

~5h

Per variant

1 배경 / 목적

5/12 Droid IDM re-dump 완료 (Keh0t0/dreamgen_robocasa_p3_{init0,mid,random}_v2_droid)로 새 synthetic dataset 3 stage가 HF에 준비됨. 기존 OLD H50 IDM 기반 mix 학습 (Jobs 1198/1199/1200, 5/11 완료)의 정확한 ablation 짝으로 같은 mix 패턴을 새 droid IDM dataset으로 재현하여 IDM 품질이 downstream Pi0.5 policy 성능에 미치는 영향을 비교한다.

비교 설계

항목	OLD (Jobs 1198-1200)	NEW (이번, Jobs 2077-2079)
aux dataset	`dreamgen_robocasa_p3_<V>_v2`	`dreamgen_robocasa_p3_<V>_v2_droid`
추출 IDM	H50 IDM (128×128 center-crop)	Droid IDM (siglip-only, 384×216 native)
학습 코드	동일 (DAVIAN H50 ref)
main dataset	동일 (`mg30_real_h50` 720ep)
mix 비율	main:aux = 5:5 (batch weight)
step / seed	동일 (20K steps, seed=1000)

변수 1개: aux dataset의 IDM 품질 (H50-128-crop OOD vs droid-siglip native). 다른 모든 조건 고정.

2 작업 내용

1. HF 데이터셋 사전 검증

3개 repo 모두 존재 확인 (last_modified 5/12 14:13~14:21). meta/info.json 검증: fps=16, codebase=v3.0, 1440 ep, 535680 fr, action[12]/state[16]/3 cameras. OLD _v2와 schema 완전히 동일 → Pi0.5 input compatible. v3.0 tag 3개 생성.

2. Multi-rank race 방지 prefetch

4 rank 동시 download 시 meta/info.json race 방지를 위해 1-process prefetch 선행. 3개 모두 ~6-9초로 완료, 각 ~600 MB.

3. Launch script 구성

scripts/launch_pi05_b2_p3_mix_new_idm.sh ← 기존 launch_pi05_b2_p3_mix.sh 미러 주요 차이 3개: --aux_dataset_repo_id=Keh0t0/dreamgen_robocasa_p3_${P3_VARIANT}_v2_droid (_droid 추가) output_dir = outputs/pi05_b2_p3_${P3_VARIANT}_new_idm_4gpu wandb run name = mg30_real_p3_${P3_VARIANT}_mix_new_idm_4gpu Env 3종 동일: LD_LIBRARY_PATH (torchcodec), SSL_CERT_FILE (certifi), HF_TOKEN

4. 3 variants 제출 (sbmr 자동 requeue × 5)

Job 2077 --qos=own --gres=gpu:4 -c 32 --mem=800GB P3_VARIANT=init0 → COMPLETED (5h 02m) Job 2078 --qos=own --gres=gpu:4 -c 32 --mem=800GB P3_VARIANT=mid → COMPLETED (4h 59m) Job 2079 --qos=extra --gres=gpu:4 -c 32 --mem=800GB P3_VARIANT=random → COMPLETED (4h 54m) core-own 8 GPU quota = init0(4) + mid(4) 정확히 채움. random은 extra (preempt 가능, sbmr requeue 5 enabled).

3 결과

Job 완료 현황

Job	QOS	Variant	Elapsed	State
2077	core-own	init0	5h 02m	COMPLETED
2078	core-own	mid	4h 59m	COMPLETED
2079	core-extra	random	4h 54m	COMPLETED

3/3 variants 모두 20K steps 정상 완료. random은 extra QOS임에도 preempt 없이 완주. 예상 mix 로그 (main 206570 fr + aux 535680 fr, w=0.50 각) 정상 진행 확인됨.

학습 설정 요약

Main Dataset

mg30_real_h50 (720ep)

Aux Dataset

_v2_droid (1440ep×3)

Mix Ratio

5:5 batch weight

Steps

20K

GPU / Job

4 × B200

Checkpoints

5K / 10K / 15K / 20K

SR 수치 pending: 24-task PoC eval 미실시. ckpt-20K 기준 eval 필요. OLD H50-IDM mix (Jobs 1198-1200)와 동일 eval 스크립트로 직접 비교 예정.

4 Takeaway

IDM 품질 → Pi0.5 downstream SR: 첫 번째 단일 변수 ablation

DreamGen 파이프라인에서 IDM action label quality가 policy 성능에 어떤 영향을 주는지 코드/config/seed를 모두 고정한 첫 번째 청정 실험이다. OLD H50 IDM의 action collapse (close task에서 max 0.18~0.58 saturate)는 sim eval에서 motion 붕괴로 확인됐고, NEW droid IDM은 10/13 task full ±1 정상화됐다.

이 ablation의 SR 차이가 크다면 IDM training distribution alignment가 DreamGen quality의 핵심 bottleneck으로 확정된다. 차이가 작다면 Pi0.5 multi-task 학습에서 action scale이 어느 정도 자기 보정됨을 시사.

5 Next Steps

24-task PoC eval @ ckpt 20K (최우선)

run_pi05_eval_24tasks.sh 에 b2_new_idm_{init0,mid,random} 분기 추가. OLD mix (Jobs 1198-1200) ckpt-20K와 동일 eval 조건으로 SR 비교. 기대: NEW droid IDM이 OLD 대비 SR 향상 (action range 정상화 효과).

만약 NEW = OLD SR이면

IDM action quality가 아닌 다른 bottleneck 존재 — real 분포와 sim init의 mismatch, video cam mismatch (synthetic droid view vs real agentview), 또는 fps 16 vs 20 scale issue. EnvUtils + reset_to PoC 결과와 교차 검증 후 다음 ablation 설계.

wandb 6-run overlay 비교

<V>_mix_4gpu vs <V>_mix_new_idm_4gpu 6 runs step-axis overlay로 in-loop SR 곡선 시각 비교. 학습 초기 (step 1000~5000) 에서 이미 갈라지는지 확인.