기술 요약: OA-WAM: 견고한 로봇 조작을 위한 객체 주소 지정 가능 세계 행동 모델

1. 문제 제기

기존의 비전 - 언어 - 행동 (VLA) 및 세계 행동 모델 (WAM) 은 백본과 데이터를 확장함으로써 표준 조작 벤치마크에서 높은 성공률을 달성했습니다. 그러나 이러한 모델은 장면의 경미한 교란 (예: 객체 배치, 카메라 시점, 조명, 또는 로봇 초기 상태의 변화) 하에서는 종종 실패합니다. 저자들은 이러한 취약성을 '객체 주소 지정 가능성 (object addressability)'의 부재로 진단합니다.

전체적인 WAM 에서 예측된 세계 상태는 전체 프레임 이미지, 비디오 토큰 스트림, 또는 글로벌 잠재 공간으로 표현됩니다. 이러한 표현들은 장면의 진화를 포착하지만, 토큰 스트림 내에서 대상 객체의 정체성을 주변 문맥과 얽히게 만듭니다. 장면 기하학이 변할 때, 행동 디코더는 지시문에 명시된 특정 객체를 분리해 내는 데 어려움을 겪으며, 이로 인해 정책이 잘못된 객체에 작용하거나 일반화에 실패하는 '드리프트 (drift)'가 발생합니다. 현재 객체 중심 접근법들은 종종 가림과 혼란에 민감한 동적 슬롯 발견 (예: Slot Attention) 에 의존하거나, 정체성, 외관, 자세를 단일 벡터로 인코딩하여 "어떤 객체에 작용할지"와 "해당 객체가 현재 무엇인지"를 분리하지 못합니다.

2. 방법론: OA-WAM

저자들은 세계 예측을 객체별 슬롯 상태 예측으로 재형상화한 OA-WAM(Object-Addressable World Action Model) 을 제안합니다. 이 아키텍처는 텐서 수준에서 객체 정체성과 시간 가변적 콘텐츠를 엄격하게 분리합니다.

2.1 객체 주소 지정 가능 토큰화

각 프레임은 $N+1$개의 슬롯 (1 개의 로봇 슬롯, $N$개의 객체 슬롯) 으로 분해됩니다. 각 슬롯 벡터 $s^t_k$는 다음을 연결하여 구성됩니다:

• 정체성 주소 ($addr_k$): 언어 레이블과 초기 DINOv3 특징으로부터 에피소드당 한 번 계산된 고정된 32 차원 벡터입니다. 이는 에피소드 동안 절대 업데이트되지 않습니다.
• 시간 가변적 콘텐츠 ($cnt^t_k$): SAM 3 마스크와 DINOv3 특징으로부터 매 단계 갱신되는 256 차원 벡터로, 객체의 현재 상태 (외관, 자세) 를 나타냅니다.
• 보조 임베딩: 프레임 인덱스 ($\pi^t$) 및 역할 레이블 ($\rho_k$).

이러한 슬롯들은 학습 가능한 슬롯 어댑터를 통해 4096 차원 은닉 공간으로 투영되며, 텍스트, 이미지-VQ, 고유 감각, 이전 행동 토큰과 블록-인과적 시퀀스로 융합되어 7B Chameleon 스타일 트랜스포머 트렁크에서 처리됩니다.

2.2 객체 주소 지정 가능 어텐션 (OA 제약)

핵심 아키텍처 혁신은 각 트랜스포머 레이어에서 적용되는 파라미터 없는 제약으로, 행동 디코더가 현재 외관이나 문맥이 아닌 오직 정체성 기반으로 올바른 객체에 어텐션하도록 보장합니다. 이는 두 가지 연산을 통해 달성됩니다:

1. 주소 전용 키 투영: 슬롯 유형 위치의 경우, 키 투영 ($K$) 은 잔여 스트림의 첫 32 차원 (즉, $addr$서브벡터) 만 읽습니다. 쿼리 ($Q$) 와 값 ($V$) 투영은 전체 은닉 상태를 사용합니다. 이는 슬롯 간 어텐션을 오직 정체성 기반으로 라우팅되도록 강제합니다.
2. 레이어별 주소 재설정: 포워드 훅이 각 트랜스포머 블록 후 슬롯 위치에서 잔여 스트림의 첫 32 차원을 캐시된 $addr_k$로 덮어씁니다. 이는 잔여 스트림의 콘텐츠 업데이트로 인해 정체성 주소가 드리프트하거나 손상되는 것을 방지합니다.

2.3 예측 헤드

모델은 단일 순전파를 수행하여 다음을 생성합니다:

• 세계 헤드: 각 슬롯에 대해 다음 프레임 상태를 회귀하며, 콘텐츠 ($\hat{c}^{t+1}$) 와 자세 ($\hat{p}^{t+1}$) 를 모두 예측합니다. 이는 시간적 일관성을 위한 보조 감독을 제공합니다.
• 행동 헤드: 학습 가능한 행동 쿼리 토큰 ($[ACT\_Q]$) 에서 16 단계 연속 행동 청크를 디코딩하는 플로우 매칭 MLP 입니다.
• 보조 이미지-VQ 헤드: 새로운 파라미터 없이 사전 학습을 지원하기 위해 트렁크의 언어 헤드를 재사용하여 다음 프레임 VQ 토큰을 예측합니다.

2.4 학습

모델은 세 단계로 학습됩니다:

1. 단계 0: OA 제약이 강제된 혼합 코퍼스 (웹 이미지 - 텍스트 + 로봇 데이터) 에서 7B 트렁크의 사전 학습.
2. 단계 I: 고정된 트렁크 가중치 위에서 슬롯 어댑터와 세계 헤드를 정렬.
3. 단계 II: 트렁크에 LoRA 를 사용한 전체 시스템 파인튜닝 및 행동 헤드 학습. 손실 함수는 행동 플로우 매칭, 세계 예측 MSE, 그리고 방해 객체 순열 불변성 및 역할 정렬을 위한 보조 정규화 항을 결합합니다.

3. 주요 기여

1. 객체 주소 지정 가능성의 공식화: 본 논문은 기존 WAM 들이 기하학적 교란 하에서 견고성 실패를 겪는 근본 원인이 안정적이고 쿼리 가능한 객체 인터페이스의 부재임을 규명합니다.
2. 아키텍처적 분리: OA-WAM 은 추가 토큰을 도입하거나 동적 슬롯 발견에 의존하지 않고, 텐서 수준의 메커니즘 (고정된 주소 + 주소 전용 라우팅) 을 통해 "어떤 객체에 작용할지"와 "해당 객체가 무엇인지"를 구조적으로 분리합니다.
3. 최첨단 견고성: 이 모델은 레이아웃, 카메라, 로봇 초기 상태 변화를 특히 테스트하는 LIBERO-Plus 벤치마크의 기하학적 축에서 새로운 SOTA 결과를 달성합니다.

4. 실험 결과

저자들은 OA-WAM 을 LIBERO, SimplerEnv, LIBERO-Plus 세 가지 벤치마크에서 평가합니다.

• 표준 벤치마크: OA-WAM 은 분포 내 (in-distribution) 작업에서 SOTA VLA 및 WAM 베이스라인과 맞먹거나 초과하는 성능을 보입니다.
  • LIBERO: 평균 성공률 97.8% (VLA-JEPA 의 97.2% 대비).
  • SimplerEnv (WidowX): 평균 성공률 79.3% (CoWVLA 의 76.0% 대비).
• 견고성 (LIBERO-Plus): OA-WAM 은 가설과 가장 부합하는 기하학적 축 (카메라, 로봇 초기, 레이아웃) 에서 새로운 SOTA 를 설정합니다.
  • Geo-Avg (카메라/로봇/레이아웃): 84.3%로, $\pi0.5$ 대비 +4.8% 향상.
  • 특정 축: 카메라에서 80.5% (Cosmos-Policy 대비 +4.7%), 로봇 초기에서 89.6%.
  • 센서 노이즈: 모델은 센서 노이즈에서 뒤처집니다 (-17.1%, 이전 최선 대비). 저자들은 이를 정책 자체의 실패가 아닌 상류 지각 스택 (슬롯 추출) 의 실패로 귀결합니다.
• 메커니즘 검증:
  • 결손 분석: OA 제약 (마스크 또는 재설정 훅) 을 제거하면 기하학적 견고성이 급격히 떨어집니다 (예: 카메라에서 -13.3%). 반면 분포 내 성능은 크게 변하지 않아 편향이 OOD(분포 외) 특정적임을 확인합니다.
  • 인과적 개입: 추론 시 대상 슬롯의 주소를 방해 객체의 주소와 교환하면, OA-WAM 의 궤적은 새로운 대상으로 이동하며 코사인 정렬도가 0.87입니다. 반면 모든 전체적 베이스라인은 정렬도 $\le 0.09$를 보여, OA-WAM 의 행동이 명시적 주소 부분 공간에 기반함을 입증합니다.

5. 중요성 및 주장

본 논문은 세계 모델을 주소 지정 가능한 객체 상태로 명시적으로 구성하는 것이 장면 교란 하에서 견고한 로봇 조작을 위한 효과적인 경로임을 주장합니다. 타겟 선택을 고정된 언어 기반 정체성 벡터에 고정함으로써, OA-WAM 은 견고성을 학습 데이터의 다양성 결과가 아닌 아키텍처의 구조적 속성으로 만듭니다.

저자들은 이 접근법이 정책이 장면이 어떻게 재배치되든 상관없이 작업 관련 객체를 신뢰성 있게 쿼리할 수 있는 "개입 가능한 객체 수준 인터페이스"를 제공한다고 강조합니다. 그들은 한계를 인정하며, 현재 이 접근법은 시뮬레이터 전용이며 고정된 지각 스택이 실패할 때 (예: 과도한 센서 노이즈 하 또는 투명/가려진 객체의 경우) 성능이 저하된다고 지적하지만, 아키텍처적 기여가 라우팅 메커니즘을 지각 및 정책 단계로부터 성공적으로 분리했다고 주장합니다.