Latent-IDM A v3 Multi-layer — 3-layer Concat 24-ep Eval

TL;DR

avg EEF 110.9 mm: v1 single (block 27, 114.2 mm) 대비 −3.3 mm (+3% 개선)
gripper acc 86.2%: v1 대비 +1.9%p 개선
비용 증가: cache 3× (28 GB → 84 GB), 학습 throughput −26% (365 → 269 samples/s)
task 분산 큼: CloseDoubleDoor −98 mm, PnPMicrowaveToCounter +145 mm — 평균 개선이 안정적이지 않음
권장 다음 방향: spatial resolution ↑ (4×4 → 8×8) — multi-layer보다 spatial bottleneck이 더 critical

110.9

avg EEF mm

86.2%

Gripper Acc

0.0114

Train Loss

60K

Steps

84 GB

Cache 크기

1 배경 / 목적

2026-04-30 분석(analysis-pixel_vs_latent_parity)에서 Latent-A v1 (single layer, block 27)이 Pixel v2_8gpu와 사실상 동등(131.8 vs 134.8 mm, MG-30 기준)임을 확인했다. 이전에 latent의 "2.5× 개선"으로 보였던 것은 batch 64 vs 1024 비교의 오해석이었다.

단일 layer feature의 한계를 극복하기 위한 첫 번째 시도로 multi-layer feature concat을 검증. Cosmos DiT 36-layer 중 mid-to-late 영역 3개(blocks 18, 24, 30)의 hidden state를 per-token으로 concat하여 action-relevant signal의 풍부도를 높이는 것이 목표.

가설: 단일 semantic layer(block 27)보다 mid-late 다층 정보 결합이 다양한 task에 걸쳐 더 robust한 action 예측을 가능하게 할 것.

선행 작업: v1 (block 27 single, 114.2 mm, 2026-04-28), v2 (action+state 동시 예측, 127.9 mm, 2026-04-30). 이번 v3은 v1 아키텍처에서 feature 추출 방식만 변경.

2 작업 내용

Layer 선택: blocks {18, 24, 30}

36-layer Cosmos DiT에서 전체의 50%, 67%, 83% 위치. AC-2 cosine similarity가 검증된 block 27(75%) 주변의 mid-to-late semantic 영역으로 선정. 전체 layer 사용 시 cache 비용이 너무 커서 3개로 한정.

처리 구조 (per-token concat)

Cosmos DiT (1번 forward) ├ block 18 hook → hidden_18 (K, T, H, W, 5120) ├ block 24 hook → hidden_24 (K, T, H, W, 5120) └ block 30 hook → hidden_30 (K, T, H, W, 5120) ↓ spatial pool 4×4 each ↓ concat last dim pooled (K, T, 4, 4, 15360) ← cache .npz ↓ temporal slice (a_off-aware) per-token (64 tokens × 15360) ↓ Linear(15360 → 1024) ← latent_projector ↓ vl_self_attention(4) + DiT(8) predicted action (16, 32)

설계 선택 비교

옵션	설명	채택	이유
A. concat per token	dim만 3× (15360)	✅ 채택	가성비 최고, v1 코드 무수정
B. stack separate tokens	N_tok 3× = 192, layer pos embed 가능	미채택	attention KV 비용 3× 증가
C. multi-stream	layer별 parallel head	미채택	model architecture 큰 변경
D. learnable weighted sum	3 scalar 학습	미채택	5120 dim 그대로, model capacity 작음

신규 파일 (v1 코드 무수정)

파일	역할
`scripts/cache_cosmos_multi_layer_pooled_mg30.py`	N layer hook → spatial pool → concat
`scripts/run_cache_multi_layer_pooled_mg30.sh`	MG-30 train cache (4 GPU sbm)
`scripts/run_cache_multi_layer_pooled_heldout.sh`	heldout cache
`IDM_dump/base_latent_A_cached_multi.yaml`	yaml (`latent_in_dim=15360`만 변경)
`scripts/run_idm_train_robocasa_latent_A_multi_smoke.sh`	30-step smoke
`scripts/run_idm_train_robocasa_latent_A_multi_60k_8gpu.sh`	60K 본 학습
`scripts/run_eval_latent_A_multi_24ep.sh`	24-ep eval

# yaml 변경점 (v1 대비 latent_in_dim 하나만 다름) latent_in_dim: 15360 # v1: 5120 (single layer) # architecture 변경 없음 — dataset이 D 크기를 자동 감지

Dataset은 v1 cached 그대로 재사용 가능. RobocasaLatentIDMDataset이 npz 로드 시 D 크기를 자동 감지(5120 → 15360)하므로 코드 변경 불필요.

3 결과

전체 트랙 비교 (24-ep heldout eval, MG-30 기준)

트랙	데이터	Steps	Train Loss	avg EEF mm	Gripper Acc	Leakage
Pixel B0	MG-30 (batch 64)	100K	0.094	332	—	없음
Pixel v2_8gpu	MG-30 (batch 1024)	60K	—	134.8	—	없음
HF 공식	MG-3000	—	—	113	—	미확인
Latent-A v2 (action+state)	MG-30	60K	0.0096	127.9	83.4%	없음
Latent-A v1 (single L27)	MG-30	60K	0.0119	114.2	84.3%	없음
Latent-A v3 multi (L18,24,30)	MG-30	60K	0.0114	110.9	86.2%	없음

개선: v1 single (114.2 mm) → v3 multi (110.9 mm) — −3.3 mm (~3%), gripper acc +1.9%p. train loss도 0.0119 → 0.0114로 소폭 개선.

Per-task 상세 비교 (v1 single → v3 multi)

Task	v1 single (mm)	v3 multi (mm)	Δ (mm)	변화
CloseDoubleDoor	340.1	242.3	−97.8	큰 개선
OpenSingleDoor	126.2	81.3	−44.9	개선
CoffeeSetupMug	73.9	56.9	−17.0	개선
CoffeePressButton	38.8	31.4	−7.4	개선
TurnOnStove	28.4	37.7	+9.3	약간 악화
PnPSinkToCounter	98.4	229.0	+130.6	큰 악화
PnPMicrowaveToCounter	213.0	358.3	+145.3	큰 악화

분산 경고: 평균 −3.3 mm 개선의 실체는 Door 계열 outlier 감소 + PnP 계열 대폭 악화의 상쇄. "멀티레이어가 전반적으로 낫다"고 볼 수 없음.

비용 비교

항목	v1 single	v3 multi	비율
Cache 크기 (mg30 + heldout)	28 GB	84 GB	3×
Cache 생성 시간 (4 GPU)	mg30 23min, heldout 45min	mg30 23min, heldout 45min	동일
Latent dim	5,120	15,360	3×
학습 throughput	365 samples/s	269 samples/s	−26%
60K 학습 wall time	~12–13 h	9 h 15 min	빠른 worker 이용

비용 대비 효과: cache 3×, throughput −26% 투자에 대해 avg −3.3 mm 개선. marginal ROI.

4 Takeaway

멀티레이어 concat의 한계 진단

avg 3% 개선은 확인됐지만 비용 대비 marginal하고 task별 분산이 크다. 근본 원인으로 세 가지가 가능하다:

1. Layer 다양성 부족 — blocks 18, 24, 30은 모두 mid-to-late 영역(50~83%)이라 비슷한 abstraction level. {6, 18, 30, 35}처럼 더 다양한 spread가 필요할 수 있음.

2. Linear projection이 layer 정보를 보존 못함 — 15360 → 1024로 압축 시 layer-wise structure가 소멸. Layer-specific positional embedding이나 attention-based fusion이 필요.

3. Spatial resolution bottleneck — 4×4 spatial pool이 task-relevant detail(small object 위치 등)을 이미 심하게 압축. PnP 계열 악화의 원인이 spatial 정보 손실일 가능성 높음. multi-layer를 추가해도 같은 spatial 한계가 지속됨.

v1 코드 재사용성 검증: 새 파일 7개만으로 multi-layer 인프라 구축, v1 dataset의 D 자동 감지 기능 덕분에 코드 0 수정. 이 패턴은 이후 spatial resolution 변경 실험에서도 동일하게 적용 가능.

latent 트랙 전반 재검토 포인트: pixel(134.8 mm) ≈ latent v1(131.8 mm) ≈ v3 multi(110.9 mm). MG-30 기준에서 latent 트랙의 추가 비용(Cosmos forward, 대용량 cache)이 정당화되려면 spatial resolution 개선이나 OOD transfer 시나리오에서 유의미한 차이가 나와야 한다.

5 Next Steps

미해결 한계

PnP 계열(Sink, Microwave) task에서 v3 multi가 v1 대비 +130~145 mm 악화. 원인 미확정 — small object 위치를 요구하는 task에서 4×4 spatial pool이 너무 거칠 가능성. 또는 15360 → 1024 linear projection이 PnP-relevant feature를 v1보다 더 많이 버릴 가능성.

Option 1 (우선): Spatial resolution ↑ (4×4 → 8×8, single layer)

multi-layer보다 spatial bottleneck이 더 critical하다는 가설 검증. 단일 layer(block 27)에서 pool을 4×4 → 8×8로 변경 → token 수 64 → 256, latent dim 5120 그대로. cache 크기 4×, GPU memory 증가 예상. 60K 재학습 필요.

Option 2: Layer 조합 변경 ({6, 18, 30, 35})

현재 {18, 24, 30}에서 더 다양한 abstraction spread로 교체. early layer (block 6, ~17%)를 포함해 low-level texture/shape 정보 추가. cache 크기 동일(4 layer이면 20 GB × 4 = 80 GB), 60K 재학습 필요.

Option 3 (중기): Pi0.5 fine-tuning on DreamGen P2

Phase 4 진입 — IDM 트랙과 병행하여 DreamGen P2 1440-ep synthetic data로 Pi0.5 policy fine-tuning. IDM 개선과 독립적으로 진행 가능. 사용자로부터 Pi0.5 코드 제공 대기 중.

Latent-IDM A v3 — Multi-layer (blocks {18, 24, 30}) Concat 60K + 24-ep Eval