Robocasa Wan T2V 2-view+PT 파이프라인 구현

TL;DR

신규 파이프라인: Wan2.1-T2V-1.3B 기반 robocasa 2-view RGB + PT row 공동 생성, 기존 i2v width-concat 코드 무변경 유지
First-frame conditioning: width 축 conditioning 제거 → temporal frame 0 clean clamp 트릭, 학습/추론 모두 동일 경로
코드량: 신규 ~1,070 lines (sft_trainer 380L, dataset 290L, pt_extractor 180L), 수정 ~122 lines, 전체 syntax 검사 통과
데이터셋 범위: plan(CoffeePressButton only) → 구현 후 전체 single_stage 24 HDF5 (23 tasks) 로 확장, 70,267 train demos

1,070

New Lines

Files Changed

70K

Train Demos

Tasks

1 배경 / 목적

기존 EgoX2/exp3는 Wan2.2-I2V-A14B (14B MoE, dual-expert) + DROID/EgoExo4D mp4 기반 width-concat 파이프라인으로, exo 절반은 clean conditioning, ego 절반만 generation target이었다. 이 구조는 ablation iteration이 느리고 (14B MoE 학습 비용) 실제 로봇 manipulation 도메인 신호가 없었다.

기존 한계: Wan2.2-I2V-A14B 14B 모델 → 단일 GPU에서 step당 ~4s 이상, DROID/EgoExo4D 데이터는 egocentric 카메라 페어링이지만 prompt diversity 부족

목표: Wan2.1-T2V-1.3B (single-expert) 로 다운사이즈하고, robocasa HDF5 (task-conditioned 로봇 demo, ep_meta의 lang 텍스트 포함) 를 도입하여 multi-view + multi-modal joint generation 실험 cycle을 단축.

설계 선택 근거: 기존 i2v width-concat 헬퍼를 그대로 놔두고 wan_t2v/ 서브모듈로 분리 → DROID/EgoExo4D 학습 경로 회귀 위험 없음

2 작업 내용

핵심 아키텍처 결정

2×2 HCPT grid: height-concat × width-concat으로 두 카메라 뷰(agentview_left + eye_in_hand)와 RGB row + PT row를 하나의 latent tensor로 처리.

INPUT shape (pixel): [B, 3, F=81, H=432, W=768] → view1 RGB 216×384 | view2 RGB 216×384 (top row) → view1 PT 216×384 | view2 PT 216×384 (bottom row) Latent shape (8× spatial, 4× temporal): [B, 16, F=21, H=54, W=96] First-frame clean conditioning (temporal): noisy[:, :, 0:1, :, :] = latent[:, :, 0:1, :, :] # frame 0 고정 loss = F.mse_loss(v_pred[:, :, 1:], v_target[:, :, 1:]) # frame 0 제외

신규 파일 6개

파일	Lines	역할
`core/finetune/models/wan_t2v/sft_trainer.py`	380	WanT2VSftTrainer + WanFirstFrameClampedT2VPipeline. compute_loss에서 frame 0 clamp + frame 0 loss 제외. validation_step은 prompt + first_frame_latent만 소비.
`core/finetune/datasets/robocasa_hdf5.py`	290	RobocasaHDF5Dataset(_BaseWidthConcatDataset). h5py로 HDF5 열어 random temporal crop (81 frames), ep_meta["lang"] prompt, PT .pt 로드 → 2×2 grid latent 빌드 + safetensors 캐시.
`scripts/extract_robocasa_pt.py`	180	depth → camera-frame XYZ 변환. fovy→(fx,fy,cx,cy) 계산, 모든 T_demo 프레임 추출, norm.json 저장.
`core/finetune/models/wan_t2v/lora_trainer.py`	8	WanT2VLoraTrainer(WanT2VSftTrainer) + register("wan-t2v","lora",…)
`configs/robocasa_t2v.yaml`	60	model_name: wan-t2v, dataset_type: robocasa_hdf5, train_resolution: 81×432×768, rank: 128
`scripts/run_robocasa_t2v.sh`	25	sbm 호출용 launcher. #SBATCH 헤더 없음. GPU 수 자동 감지해 torchrun/python 분기.

수정 파일 5개

파일	+Lines	변경 내용
`core/finetune/schemas/args.py`	+30	dataset_type Literal에 "robocasa_hdf5" 추가 + 7개 robocasa 필드
`core/finetune/trainer.py`	+75	prepare_dataset robocasa 분기 + _build_robocasa_first_frame_latent() + _maybe_run_validation robocasa 경로
`core/finetune/models/_base/width_concat.py`	+50	기존 i2v 헬퍼 보존 + build_2x2_grid_t2v_target_shape(), apply_first_frame_clean_noise() 추가
`core/finetune/datasets/__init__.py`	+2	RobocasaHDF5Dataset export
`infer.py`	+15	--task wan-t2v 분기 + _main_wan_t2v()

재사용 헬퍼 (변경 없음)

_BaseWidthConcatDataset의 _encode_pixel_video, _apply_4dnex_norm, _cache_dirs 그대로 상속. load_pointmap, build_pt_init_video 그대로 사용. transformer_2 None 자동 감지 로직 그대로 통과.

3 결과

Syntax 검사 통과: 신규 6개 파일 모두 python -m ast parsing 통과. 실제 import + smoke test는 사용자 환경에서 별도 진행.

코드 규모 요약

신규 코드

~1,070 lines

수정 코드

~172 lines

plan 항목 완료

11 / 11

Syntax 검사

6 / 6 통과

구현 vs. Plan 차이점

항목	Plan (260427)	실제 구현
데이터셋 범위	CoffeePressButton만 (1 task)	single_stage 전체 23 tasks (24 HDF5)
split 방식	mask_subset_key (1000_demos)	build_robocasa_split.py로 사전 생성된 split JSON (70,267 train / 240 val)
temporal crop	random_temporal_crop: True	False로 변경 (결정론적 첫 81 frames)
PT norm	per-demo per-view min/max	동일 (norm.json)

Smoke Test (2026-04-30)

checkpoint-4, logs/, validation/ 생성 확인 latent shape [1, 16, 21, 54, 96] ✓ noisy[:,:,0] == latent[:,:,0] (frame 0 clamp) ✓ step 4 validation mp4 생성 ✓

4 Takeaway

핵심 설계 결정

기존 i2v 코드(EgoX2_exp1/exp2, rgb branch)를 전혀 건드리지 않고 wan_t2v/ 서브모듈로 완전 분리. 이 결정으로 exp1/exp2 학습 경로의 회귀 위험이 0이 됐고, t2v 파이프라인 독립적으로 ablation 가능해짐.

Width 축 conditioning 제거 + temporal first-frame clean trick: Wan2.1-T2V-1.3B의 표준 16-channel input을 건드리지 않고 i2v 효과를 얻는 최소 침습 방법. 추후 Wan2.2-T2V-A14B (dual-expert)로 scale-up 시 동일 코드 경로 사용 가능.

robocasa 도입으로 EgoX v2 프레임워크가 처음으로 로봇 manipulation 도메인의 task-conditioned 비디오 생성 실험을 지원하게 됨. ep_meta의 lang 필드 덕분에 별도 캡셔닝 없이 의미 있는 text-video 쌍 확보.

5 Next Steps

4DNeX scalar post-norm 일관성 검토: _build_robocasa_first_frame_latent의 PT셀에 _apply_4dnex_norm 후처리가 누락 — 학습 latent와 first_frame_latent 간 PT 분포 불일치 가능성. validation 출력 PT셀 확인 후 수정 여부 결정.
Phase 2: 3rd view 추가 (agentview_right) — 1×3 width-concat으로 stereo agentview + wrist 조합 → 3D consistency 향상 여부 ablation.
여러 task HDF5 합치기: robocasa_hdf5_paths가 list라 자연 확장. KitchenStove, Drawer 등 task 다양화 → prompt diversity 확보. (이미 전체 single_stage로 확장됨 ✓)
Wan2.2-T2V-A14B scale-up: 코드는 transformer_2 자동 감지 분기 존재 → ckpt 교체 + boundary_ratio 튜닝만 하면 됨.
pose/action conditioning: HDF5의 actions (T,12), obs/robot0_eef_pos 를 추가 conditioning으로 활용.

Robocasa Wan T2V 2-view + PT 파이프라인 구현