Latent-IDM A Multistream ckpt-55K 24-ep Eval

TL;DR

EEF 104.99 mm — 전 트랙 최고: v3 multi-layer concat 110.9 mm 대비 −5.9 mm (−5.3%), v1 single 114.2 mm 대비 −9.2 mm (−8.1%)
PnP 태스크 대폭 개선: PnPSinkToCounter −107 mm, PnPMicrowaveToCounter −84 mm — layer-wise semantic 보존 가설 부분 입증
Double-door articulated 악화: CloseDoubleDoor +87 mm — spatial trajectory 태스크에서는 dim-level fusion이 더 적합
gripper acc 83.76%: v3 86.19% 대비 −2.4%p — token 수 3× 증가에 따른 cross-attn dilution 시사
버그 fix: get_action() device/dtype 분기에서 cosmos_pooled_cache_multistream 누락 → KeyError: 'z0'

104.99

EEF mm (Best)

83.76%

Gripper Acc

401M

Params

55K

Ckpt Steps

1 배경 / 목적

Latent-IDM 연구는 Cosmos-Predict2 14B frozen backbone의 hidden representation을 IDM action head에 활용하는 방식을 개선하는 것이 목표다. v3 multi-layer concat (dim-level fusion)에서 EEF 110.9 mm를 얻었지만 Linear projection이 layer별 정보를 충분히 fuse하지 못한다는 가설이 남아있었다.

기존 한계 (v3): 3개 layer hidden을 token 마지막 dim으로 concat (5120×3=15360 → 1024 단일 Linear). 모델이 어느 dim이 어느 layer인지 알 수 없어 layer별 semantic 정보가 소실될 수 있다. 24-ep EEF 110.9 mm로 v1 single 114.2 mm 대비 marginal win (+3%)에 그쳤다.

Multistream 가설: 각 layer에 별도 Linear + SelfAttention(2L) + layer position embedding을 적용해 layer-wise representation을 충분히 학습한 뒤 token 차원으로 concat하여 fusion하면, dim-level fusion보다 풍부한 representation을 만들 수 있다.

2 작업 내용

아키텍처

cache (B, 64, 15360) ← v3 multi cache 재사용 (재캐싱 없음) ↓ split last dim → 3 chunks of (B, 64, 5120) Stream 18: Linear(5120→1024) + layer_emb[0] + SelfAttn(2L) → ctx_18 (B, 64, 1024) Stream 24: Linear(5120→1024) + layer_emb[1] + SelfAttn(2L) → ctx_24 (B, 64, 1024) Stream 30: Linear(5120→1024) + layer_emb[2] + SelfAttn(2L) → ctx_30 (B, 64, 1024) ↓ concat tokens (dim=1) fused (B, 192, 1024) → vl_self_attention(4L) → DiT cross-attn(8L) → action (16, 32) Total params: 401.18M (v3 315M 대비 +86M, +27%) DiT cross-attn KV tokens: 192 vs v3 64 (3×, attention cost ~9×)

신규 파일

gr00t/model/action_head/multi_stream_projector.py — 새 모듈 (n_layers=3 stream)
gr00t/model/action_head/flow_matching_action_head_idm.py — cosmos_pooled_cache_multistream 분기 추가 (forward / encode_visual / get_action 3곳)
IDM_dump/base_latent_A_cached_multistream.yaml — feature_source만 다른 yaml

학습 및 Eval

8 GPU × 60K steps 목표. Cluster 만석으로 PREEMPTED 4회 (Jobs 5635→5879→5989→6055). lineage 끊겨 ckpt-55K로 임시 평가. 60K resume은 Job 6216으로 별도 제출.

버그 fix

KeyError: 'z0': ckpt-55K eval 첫 제출(Job 6189) 1분 만에 실패. flow_matching_action_head_idm.py:745 get_action() device/dtype 분기에 cosmos_pooled_cache_multistream 누락 → cosmos_latent else-branch로 fallback하여 data["z0"] 접근. cosmos_pooled_cache 분기에 multistream을 합쳐 해결. 재제출(Job 6194) 28분 정상 종료.

3 결과

전체 트랙 비교 (24-ep heldout)

트랙	EEF mm	Gripper Acc	Params
Latent-A multistream ckpt-55K	104.99 ⭐	83.76%	401M
Latent-A v3 multi concat {18,24,30}	110.88	86.19%	315M
Latent-A v1 single L27	114.23	84.27%	315M

Per-task 변동 — multistream vs v3 (큰 변동 위주)

Task	v3 mm	multistream mm	Δ
PnPSinkToCounter	229.0	121.9	−107
PnPMicrowaveToCounter	358.3	274.7	−84
PnPCounterToMicrowave	160.7	91.7	−69
TurnOnSinkFaucet	94.8	56.3	−39
PnPStoveToCounter	116.9	99.9	−17
TurnOffSinkFaucet	95.5	81.5	−14
OpenSingleDoor	81.3	71.0	−10
⬇ 악화 구간
CloseDoubleDoor	242.3	328.8	+87
OpenDoubleDoor	150.6	188.5	+38
PnPCounterToStove	61.3	88.7	+27
PnPCounterToSink	93.7	116.2	+23
PnPCabToCounter	90.4	111.6	+21

PnP (sink↔counter, microwave↔counter): multistream이 큰 개선. 미세한 object grasp/place 단계에서 layer-wise semantic 분리가 효과적.

Double-door articulated: multistream이 악화. door angle 같은 spatial trajectory reasoning은 dim-level fusion(v3)이 더 적합.

4 Takeaway

Layer-wise semantic 보존 가설 — 부분 입증

PnP 계열 태스크(object affordance + grasp 정밀도 중요)에서 multistream이 큰 이득. 단일 fusion보다 per-layer attention 후 fusion이 더 풍부한 representation을 제공한다는 가설이 해당 태스크에서 검증됐다.

태스크별 표현 선호가 다르다

Door manipulation처럼 spatial trajectory가 중요한 태스크는 v3 dim-level fusion이 더 적합. 두 방식이 상보적 — 단일 모델 안에서 태스크별 routing이 가능한지 다음 실험 대상.

gripper accuracy 손해 원인: KV token 수 64 → 192 (3×)로 인해 gripper-specific signal이 cross-attn에서 dilution됐을 가능성. gripper head 별도 weighting이 필요할 수 있다.

5K saturation: v1/v3 기준 55K vs 60K 차이 미미 (114.0→114.2, 110.5→110.9). multistream 60K 결과도 ~104-105 mm 예상.

5 Next Steps

60K 완성 (Job 6216 진행 중)

정확한 수치 확정. 55K 대비 ~0.5-1 mm 변동 예상. 완료 후 24-ep eval 재제출.

Hybrid track 후보

Door 태스크는 v3 fusion, PnP 태스크는 multistream → routing layer 학습 또는 mixture-of-experts 형태로 결합 가능한지 설계 검토.

Gripper head weighting

token 수 3× 증가에 따른 cross-attn dilution이 gripper acc 하락 원인인지 분석. gripper signal에 별도 loss weight 부여 또는 gripper-specific attention mask 실험.

Latent-C Layer-routed (후속)

ActionDiT block i가 Cosmos hidden layer f(i)만 KV로 받는 layer-routed 설계 (π0.5 KI 영감). multistream과 달리 block-level 1:1 routing. dense N=12 Phase 1 이후 sparse Phase 2 계획 수립됨.

Latent-IDM A Multistream — Per-Layer SelfAttn + Token Concat ckpt-55K 24-ep Eval