기술 요약: 센서리모터 월드 모델 (Sensorimotor World Models, SMWM)

문제 정의

본 논문은 보상 신호나 태스크 라벨 없이 고차원 관측값(예: 픽셀)으로부터 **잠재 월드 모델(latent world models)**을 학습하는 과제에 대해 다룹다. **결합 임베딩 예측 아키텍처(Joint Embedding Predictive Architectures, JEPAs)**는 픽셀 재구성을 피하고 임베딩 공간에서 직접 미래 상태를 예측함으로써 유망한 프레임워크를 제공하지만, 표현 붕괴(representation collapse) 문제를 겪는다. 인코더와 순방향 역학 모델을 함께 훈련하여 예측 오차를 최소화할 때, 인코더가 모든 관측값을 동일한 상수 임베딩으로 매핑하는 방식으로 학습될 수 있다. 이는 순방향 손실을 최소화하는 사소한 해결책이지만, 다운스트림 제어나 계획 수립에는 무용지물로 만든다.

붕괴를 방지하기 위한 기존의 솔루션들은 다음과 같은 복잡한 메커니즘에 의존한다:

• 사전 훈련된 인코더 동결 (예: DINO-WM).
• 분포 정규화 추가 (예: SIGReg를 통해 임베딩을 등방성 가우시안에 맞춤).
• 스톱 그래디언트(stop-gradients)를 포함한 지수 이동 평균(EMA) 타겟 네트워크 사용.
• 분산-공분산 제약을 포함한 다중 항 손실 함수 사용.

저자들은 이러한 접근 방식들이 유용한 월드 모델 표현이 무엇을 유지해야 하는지에 대한 핵심 질문을 완전히 해결하지 못할 수 있다고 주장한다. **"행동을 위한 지각(perception for action)"**과 **생성적 지각(enactive perception)**의 개념에 기초하여, 본 논문은 표현이 조절 가능한 역학(controllable dynamics)에 중요한 차이는 보존하되, 조절 불가능한 방해 요소(uncontrollable distractors)는 버리도록 형성되어야 한다고 주장한다.

방법론: 센서리모터 월드 모델 (SMWM)

저자들은 오프라인, 보상 없는 궤적 $(o_t, a_t, o_{t+1})$로부터 엔드 투 엔드로 훈련되는 SMWM을 제안한다. 핵심 혁신은 **역 역학 정규화(inverse dynamics regularization)**를 표현 붕괴를 방지하기 위한 유일한 메커니즘으로 사용하는 것이다.

아키텍처

모델은 공동으로 훈련되는 세 가지 구성 요소로 이루어진다:

1. 인코더 ($f_\theta$): 관측값 $o_t$를 압축된 잠재 임베딩 $z_t$로 매핑한다.
2. 순방향 역학 모델 ($g_\phi$): 현재 상태 $z_t$와 행동 $a_t$가 주어졌을 때 다음 잠재 상태 $\hat{z}_{t+1}$을 예측한다.
3. 역 역학 모델 ($h_\psi$): 연속된 두 임베딩 $z_t$와 $z_{t+1}$ 사이에서 취해진 행동 $\hat{a}_t$를 예측한다.

훈련 목적 함수

모델은 다음과 같은 결합 손실 함수를 최소화하도록 최적화된다:
$$\mathcal{L} = \mathcal{L}_{fwd} + \lambda \mathcal{L}_{inv}$$

• 순방향 손실 ($\mathcal{L}_{fwd}$): 예측된 다음 임베딩 $\hat{z}_{t+1} = g_\phi(z_t, a_t)$와 실제 인코딩된 다음 상태 $z_{t+1} = f_\theta(o_{t+1})$ 사이의 평균 제곱 오차(MSE)이다.
• 역 손실 ($\mathcal{L}_{inv}$): 예측된 행동 $\hat{a}_t = h_\psi(z_t, z_{t+1})$와 실제 행동 $a_t$ 사이의 MSE이다.
• 정규화 가중치 ($\lambda$): 역 역학 신호의 강도를 제어하는 단일 하이퍼파라미터이다.

작동 메커니즘

역 손실은 안티-붕괴(anti-collapse) 메커니즘 역할을 한다. 만약 인코더가 상수로 붕괴되면, 역 모델은 모든 전이에 대해 동일한 입력을 받게 되며 오직 상수 행동(행동 분포의 평균)만을 예측할 수 있게 된다. 이 상수 예측 위험보다 역 손실을 더 낮추기 위해서, 인코더는 반드시 전이 $(o_t, o_{t+1})$에 포함된 행동 $a_t$에 대한 정보를 보존해야만 한다. 이는 표현이 행동 관련 정보(조절 가능한 자유도)를 유지하도록 강제하는 동시에, 행동 복구에 영향을 주지 않는 조절 불가능한 방해 요소(외생적 노이즈)를 자연스럽게 필터링하도록 만든다.

임베딩 공간의 기하학적 구조를 규정하는(예: 가우시안으로 강제하는) 분포 정규화와 달리, 역 역학은 특정 기하학적 사전 확률을 부과하지 않고 표현을 태스크에 근거한 양(행동)에 고정시킨다.

주요 기여

1. SMWM 아키텍처: 역 역학을 유일한 안티-붕계 메커니즘으로 사용하여, 재구성이나 보상 신호, 동결된 인코더, 또는 복잡한 분포 사전 확률 없이도 엔드 투 엔드로 훈련 가능한 단순한 센서리모터 월드 모델을 제공한다.
2. 표현 특성: 학습된 표현이 단순히 붕괴되지 않은 것뿐만 아니라, 구체적으로 조절 가능한 자유도를 추적하고, 공간 기하학(데카르트 및 주기적 구조 모두)을 보존하며, 조절 불가능한 방해 요소를 필터링한다는 실증적 증거를 제시한다.
3. 다운스트림 성능: 2D 및 3D 제어 태스크에서 경쟁력 있는 계획 성능을 입증하여, 표현의 구조적 이점이 효과적인 목표 조건부 계획(goal-conditioned planning)으로 이어진다는 것을 보여준다.

실험 결과

통제된 테스트베드 (Dot World & Sprite World)

• 내재 차원 회복: "Dot World" 환경에서 모델은 64차원 잠재 공간에 임베딩되어 있음에도 불구하고, 픽셀로부터 상태 공간의 실제 2D 내재 차원을 성공적으로 회복했다.
• 방해 요소 필터링: 조절 불가능한 "방해 객체"(무작위로 움직임)가 도입되었을 때, 인코더는 이를 무시했다. 유효 잠재 차원은 전체 관측 차원이 아닌 조절 가능한 자유도와 일치했다.
• 제어 의존적 지각: "Sprite World" 실험에서 모델은 포즈 변수(위치, 방향)가 조절 가능할 때만 이를 보존하도록 학습했다. 방향이 조절 불가능해지자, 모델은 방향을 평균화하여 방향에 대한 흐릿한 재구성을 나타냈다.

계획 벤치마크

저자들은 네 가지 오프라인 제어 태스크(TwoRoom(2D 내비게이션), Reacher(연속 제어), Push-T(2D 조작), OGBench-Cube(3D 조작))에 대해 SMWM을 평가했다.

• 베이스라인: SIGReg(가우시안 정규화를 사용하는 강력한 JEPA 베이스라인), Forward-only(역 손실 없음), Random 행동 샘플링과 비교하였다.
• 성능:
  • SMWM은 세 가지 2D 태스크에서 SIGReg과 대등한 성능을 보였다.
  • SMSM은 3D OGBench-Cube 태스크에서 SIGReg을 능가했다 (성공률 84% vs 59%).
  • 두 정규화 방법 모두 Forward-only 절제 실험보다 성능이 현저히 높았으며, 이는 안티-붕계 메커니즘의 필요성을 확인시켜 준다.
• 잠재 기하학: 프로빙 실험 결과, SMWM 임베딩은 물리적 상태 기하학(예: Reacher의 주기적 관절 각도에 대한 원통형 구조, 내비게이션에 대한 데카르트 맵)을 반영하는 저차원 부분 공간을 형성함을 밝혔다.

의의 및 주장

본 논문은 SMWM이 잠재 월드 모델 학습에서 **"행동을 위한 지각"**의 원칙적이고 단순한 구현임을 주장한다. 에이전트의 행동을 복구하기에 충분한 정보를 잠재 상태가 보존하도록 강제함으로써, 모델은 자연스럽게 환경의 조절 가능한 자유도를 향해 표현을 편향시킨다.

핵-테이크(Key takeaways)는 다음과 같다:

• 단순성: 역 역학만으로도 재구성, 보상 신호, 동결된 인코더 또는 복잡한 분포 사전 확률 없이 JEPA의 엔드 투 엔드 훈련을 안정화하기에 충분하다.
• 구조적 편향: 이 방법은 표현이 단순한 예측 상관관계가 아니라, 환경의 인과 구조(행동 및 그 효과)와 정렬된 해석 가능한 형태를 갖도록 유도한다.
• 실용적 유용성: 학습된 표현은 특히 다른 정규화 베이스라인들이 어려움을 겪는 복잡한 3D 환경에서 효과적인 잠재 계획 수립을 지원한다.

저자들은 연속적인 관측으로부터 행동을 복구할 수 있다는 가정과, 관측 이력 없이는 단일 프레임에서 식별할 수 없는 양(예: 속도)을 포착하지 못하는 현재의 한계를 인정한다. 그러나 SM-WM을 오프라인 데이터로부터 실행 가능한(actionable) 월드 모델을 학습하기 위한 견고하고 보상이 필요 없는 접근법으로 제시한다.