기술 요약: WAM4D

문제 정의

World Action Models (WAMs)는 미래의 시각적 관측값과 실행 가능한 로봇 동작을 결합하여 모델링함으로써 체화된 지능(embodied intelligence)을 위한 유망한 접근 방식으로 부상했습니다. 그러나 기존의 WAM은 주로 2D 비디오 또는 잠재 공간(latent spaces)에서 작동합니다. 이러한 모델들은 시각적으로 그럴듯한 미래의 롤아웃(rollout)을 생성할 수는 있지만, 정밀한 로봇 조작에 필수적인 3D 공간 제약, 가려진 접촉 기하학(occluded contact geometry) 및 객체 역학을 포착하는 데는 종종 실패합니다.

4D 기하학(예: 깊이 지도, 포인트 클라우드)을 WAM에 통합하려는 최근의 시도들은 밀집된 기하학적 표현을 명시적인 추론 대상으로 취급합니다. 이러한 방식은 동작 추론 시 비용이 많이 드는 기하학적 디코딩으로 인해 상당한 계산 오버헤드를 발생시켜 인과적 동작 생성을 늦춥니다. 또한, 밀집된 4D 기하학을 명시적으로 예측하는 것은 모델의 목적을 비디오 예측과 동작 생성 사이의 인과적 결합을 강화하는 것이 아니라 기하학적 재구성 쪽으로 전이시킬 수 있습니다. 핵심 과제는 배포 시의 지연 시간과 복잡성을 초래하지 않으면서, 강력한 기하학적 사전 지식(priors)을 활용하여 공간 추론과 동작 예측을 개선하는 방법입니다.

방법론: WAM4D

저자들은 학습 시에만 작동하는 경량화된 메커니즘을 통해 기하학적 파운데이션 사전 지식을 인과적 비디오-액션 트랜스포머에 통합하는 빠른 4D 월드 액션 모델인 WAM4D를 제안합니다.

1. 공간 레지스터 토큰 (Spatial Register Tokens)

추론 중에 밀집된 기하학을 출력하는 대신, WAM4D는 2D 비디오 잠재 공간과 기하학적 사전 지식 사이의 압축된 가교 역할을 하는 공간 레지스터 토큰을 도입합니다.

• 메커니즘: 학습 중에 학습 가능한 레지스터 토큰은 인과적 히스토리 비디오 특징으로부터 미래의 기하학적 정보를 추출하는 쿼리(query) 역할을 합니다. 이 레지스터들은 멀티뷰 RGB 모자이크의 공간적 배치와 정렬됩니다.
• 증류(Distillation): 레지스터는 깊이 추출 블록(히스토리 비디오 토큰과의 교차 주의 집중/cross-attention)을 통해 업데이트되며, 사전 학습된 기하학 헤드(Depth Anything V3에서 초기화됨)로 투영됩니다. 모델은 이 레지스터들로부터 미래의 깊이 시퀀스를 예측합니다.
• 학습 vs. 추론: 깊이 손실(SmoothL1)은 공유된 비디오-액션 백본으로 역전파되어, 동작 예측에 사용되는 히스토리 특징에 공간적 사전 지식을 증류합니다. 결정적으로, 배포 시에는 레지스터 브랜치, 깊이 블록 및 기하학 헤드가 제거되어, 표준적인 인과적 WAM과 동일한 경량 2D 관측-동작 인터페이스를 남깁니다.

2. 인과적 혼합 주의 집중 (Causal Mixture Attention, MoT)

모델이 동작을 예측하기 위해 미래의 기하학을 훔쳐보는 비인과적 지름길(non-causal shortcuts)을 사용하는 것을 방지하기 위해, 저자들은 Mixture-of-Transformers (MoT) 백본을 위한 인과적 혼합 주의 집중 메커니즘을 설계했습니다.

• 가시성 규칙: 주의 집중 마스크(attention mask)는 비디오, 액션, 기하학 토큰 간의 가시성을 엄격하게 정의합니다.
  • 미래의 액션 토큰은 히스토리 비디오, 히스토리 액션, 그리고 자신의 노이즈가 섞인 미래 상태에 주목(attend)할 수 있습니다.
  • 미래의 액션 토큰은 미래의 비디오 토큰 및 공간 레지스터로부터 **마스킹(masked)**됩니다.
  • 공간 레지스터는 오직 자기 자신과 유효한 히스토리 비디오 토큰에만 주목할 수 있습니다.
• 결과: 이를 통해 기하학적 감독(supervision)이 학습 중에 인과적 비디오 특징을 형성하도록 보장하면서도, 추론 시에는 보조적인 기하학 토큰이 액션 예측기에 노출되지 않도록 합니다.

3. 아키텍처 및 학습

• 백본: 비디오 VAE(Wan2.2)와 비디오 및 액션 스트림 모두에 대한 플로우 매칭(flow-matching) 목적 함수를 사용하는 LingBot-VA의 인과적 비디오-액션 백본을 기반으로 구축되었습니다.
• 목적 함수: 전체 손실은 레지스터 브랜치에 적용되는 깊이 감독 손실($\mathcal{L}_{depth}$)과 비디오($\mathcal{L}_{video}$) 및 액션($\mathcal{L}_{action}$)에 대한 조건부 플로우 매칭 손실의 결합으로 구성됩니다.
• 구현: 3-뷰 모자이크(헤드 + 두 개의 손목 카메라)를 사용하며 절대적인 엔드 이펙터 포즈(end-effector poses)를 예측합니다.

주요 기여

1. WAM4D 모델: 압축된 공간-레지스터 인터페이스를 통해 기하학적 파운데이션 사전 지식을 인과적 비디오-액션 표현으로 전달하며, 추론 시 밀집된 기하학 디코딩을 피하는 빠른 4D 월드 액션 모델입니다.
2. 인과적 혼합 주의 집중: 모달리티별 가시성을 정의하는 MoT 백본을 위한 새로운 주의 집중 메커전즘으로, 이를 통해 학습 중에는 기하학적 감독을 가능하게 하면서 배포 시에는 경량화된 인과적 정책 경로를 유지합니다.
3. 경험적 검증: RoboTwin 2.0 시뮬레이션 스위트와 실제 환경의 AstriBot S1 태스크에 대한 포괄적인 평가를 통해, 공간적 일관성과 성공률을 높이면서도 추론 효율성을 저해하지 않음을 입증했습니다.

실험 결과

본 논문은 $\pi_0, \pi_{0.5}$와 같은 VLA 정책 및 최근의 비디오-액션 모델(Motus, LingBot-VA, Fast-WAM)을 포함한 강력한 베이스라인들과 WAM4D를 비교 평가합니다.

• RoboTwin 2.0 (시뮬레이션): WAM4D는 깨끗한 설정(clean settings)에서 93.8%, 무작위 설정(randomized settings)에서 **89.9%**의 성공률을 달성하여 최신 성능의 베이스라인들(예: LingBot-VA: clean 92.9%, rand 91.6%)과 대등하거나 능가하는 성능을 보였습니다. 또한 Motus(1516ms)나 LingBot-VA(843ms)와 비교하여 경쟁력 있는 추론 지연 시간(~525ms)과 더 낮은 VRAM 사용량(9.71 GiB)을 유지했습니다.
• 실제 환경 (AstriBot S1): 네 가지 조작 태스크(접시 들기, 병 놓기, 펜 캡 제거, LEGO 분류)에서 WAM4D는 평균 서브 액션 성공률 **90%**를 기록하며, $\pi_{0.5}$(74%) 및 Fast-WAM(80%)을 능가하는 가장 높은 성과를 거두었습니다.
• 절제 연구 (Ablation Studies):
  • 레지스터 배치: 중간 레이어 레지스터(layers 12, 14, 16, 18)가 기하학적 증류와 시각적 특징 보존 사이의 최적의 균형을 제공했습니다.
  • 기하학 헤드: (Depth Anything V3에서 초기화된) 학습 가능한 사전 학습된 기하학 헤드를 사용하는 것이 무작위 초기화나 고정된 사전 학습 헤드를 사용하는 것보다 훨씬 더 나은 결과를 얻었으며, 이는 강력한 기하학적 사전 지식을 적응시키는 것의 필요성을 확인시켜 줍니다.
  • 인과적 제약: 인과적 마스킹을 제거(양방향 레지스터)했을 때 성공률은 약간 향상되었으나 기하학적 지표가 저하되고 복잡성이 증가하여, 단방향 설계의 타당성을 입증했습니다.

의의 및 주장

본 논문은 WAM4D가 공간적 일관성과 추론 효율성 사이의 트레이드오프를 성공적으로 해결했다고 주장합니다. 기하학을 추론 시의 출력이 아닌 학습 시의 읽기(readout)로 처리함으로써, WAM4D는 다음과 같은 효과를 거둡니다:

• 객체의 범위, 가려진 표면, 접촉 기하학에 대한 모델의 추론 능력을 향상시킵니다.
• 2D 비디오-액션 모델의 경량 배포 특성을 유지하여 실시간 로봇 제어에 적합하게 만듭니다.
• 실행 중 명시적인 3D 재구성을 요구하지 않고도 기하학적 파운데이션 사전 지식이 인과적 비디오 특징으로 효과적으로 증류될 수 있음을 보여줍니다.

저자들은 모델에 명시적인 장기 객체 메모리가 부족하여, 객체가 가려질 때 긴 자기 회귀적(autoregressive) 롤아웃에서 정체성 불일치가 발생할 수 있다는 한계를 인정했습니다. 그러나 이 모델은 환경으로부터 신선한 관측값을 지속적으로 수신하므로, 이것이 폐쇄 루프 제어(closed-loop control)에서의 정책 성공을 저해하지는 않는다고 강조합니다.