기술 요약: Parsimonious Latents 로의 회귀: 시각적 기반에서 작업 중심의 세계 모델 학습

문제 정의

세계 모델 (World models) 은 에이전트가 행동을 조건으로 미래 역학을 예측할 수 있게 하여, 계획 및 제어의 핵심이 되는 능력을 제공합니다. 그러나 기존 접근법들은 잠재 표현 (latent representation) 에서 근본적인 트레이드오프에 직면해 있습니다:

1. 픽셀 기반 모델 (예: Dreamer, MuZero) 은 픽셀에서 직접 역학을 학습하거나 보상 신호에 의해 형성된 컴팩트한 잠재 변수를 사용합니다. 효과적이지만, 종종 높은 계산 비용을 초래하거나 잠재 공간을 형성하기 위해 온라인 상호작용/보상 감독이 필요합니다.
2. 기반 임베딩 모델 (예: DINO-WM, V-JEPA) 은 고정된 사전 훈련된 시각 인코더를 활용하여 고차원 의미 공간에서 직접 예측합니다. 풍부한 의미 사전 지식과 제로샷 일반화를 제공하지만, 이러한 임베딩에는 제어에 영향을 미치지 않는 과도한 작업 무관 세부 사항 (텍스처, 조명, 배경) 이 포함되어 있습니다.

보상 감독 없이 고정된 궤적에서 학습해야 하는 보상 없는 오프라인 설정에서 이러한 문제들은 더욱 악화됩니다. 원시 기반 임베딩을 상태 공간으로 사용하면 모델이 무관한 변이들에 자원을 낭비하게 되어 모델링 효율성이 떨어지고, 제어 해석 가능성이 감소하며, 특히 고차원 행동 공간 (예: 7 자유도 로봇 팔) 에서 표현 붕괴가 발생할 수 있습니다. 핵심 과제는 온라인 상호작용 없이 이러한 풍부하지만 노이즈가 많은 기반 임베딩에서 컴팩트하고 작업에 충분한 잠재 표현을 추출하는 것입니다.

방법론: TC-WM

저자들은 사전 훈련된 기반 임베딩을 최종 상태 공간이 아닌, 작업 관련 역학을 증류하는 "의미적 발판 (semantic scaffold)"으로 간주하는 **TC-WM(작업 중심 세계 모델)**을 제안합니다.

핵심 아키텍처

TC-WM 은 시각 임베딩 ($x_t$) 과 고유감각 신호 ($s^p_t$) 의 파이프라인에서 작동합니다:

1. 결합 임베딩: 고정된 시각 인코더 (예: DINOv2) 가 이미지를 시각 특징으로 매핑하고, 학습 가능한 선형 계층이 고유감각을 인코딩합니다. 이들이 연결되어 결합 임베딩 $x_t$를 형성합니다.
2. 잠재 인코더 (투사): 선형 인코더가 고차원 결합 임베딩을 컴팩트한 잠재 상태 $z_t$로 투사합니다. 핵심적으로, $z_t$는 두 개의 부분 공간으로 분해됩니다:
   • $z^s_t$ (상태 관련): 물리적 상태와 정렬됨.
   • $z^c_t$ (보조적): 재구성 제약을 충족시키기 위해 잔여 의미 구조 (외관, 배경) 를 보유함.
3. 잠재 역학: 트랜스포머 기반 예측기 ($f_\theta$) 가 컴팩트한 잠재 공간에서 시간적 전이를 모델링합니다 ($\hat{z}_{t+1} = f_\theta(z_{t-H:t}, a_{t-H:t})$).
4. 작업 중심 정렬: 모델은 InfoNCE 와 같은 대비적 목적 함수를 사용하여 $z^s_t$ 부분 공간을 고유감각 상태 $s^p_t$와 정렬합니다. 이는 인코더가 작업 관련 물리적 요인을 노이즈 변이와 능동적으로 분리하도록 강제합니다.
5. 임베딩 디코더: 선형 디코더가 $z_t$에서 원래 시각 임베딩을 재구성합니다. 이는 정규화제 역할을 하여 컴팩트한 잠재 변수가 기반 다양체 (manifold) 에 고정되도록 하고 퇴화된 해를 방지합니다.

학습 목적 함수

총 손실은 네 가지 구성 요소를 결합합니다:

• 잠재 역학 손실 ($L^z_{dyn}$): 잠재 공간에서의 예측 오차를 최소화합니다.
• TC 역학 손실 ($L^s_{dyn}$): 과거 데이터에서 미래 고유감각을 직접 예측하여 작업 중심 역학을 강제합니다.
• 정렬 손실 ($L_{align}$): $z^s_t$와 $s^p_t$를 정렬하는 대비적 손실입니다.
• 재구성 손실 ($L_{rec}$): 시각 구조를 보존하기 위해 $z_t$에서 시각 임베딩을 재구성합니다.

테스트 시간 계획

학습된 작업 중심 잠재 공간은 다양한 계획 전략을 지원합니다:

• 저차원 제어 (내비게이션): 잠재 공간에서 행동 시퀀스를 최적화하기 위해 교차 엔트로피 방법 (CEM) 을 사용하는 모델 예측 제어 (MPC).
• 고차원 조작: 잠재 궤적에서 실행 가능한 행동을 생성하기 위해 역동역학 모델 (IDM) 과 결합된 잠재 확산 계획 (LDP).

주요 기여

1. 작업 중심 분해: 고차원 기반 임베딩을 물리적 상태 역학과 의미적 노이즈 변수를 명시적으로 분리함으로써 컴팩트하고 작업에 충분한 잠재 공간으로 분해하는 방법을 제시합니다.
2. 능동적 식별 가능성: 메커니즘 다양성을 수동적으로 가정하여 식별 가능성을 확보하는 기존 연구와 달리, TC-WM 은 고유감각과의 대비적 정렬을 사용하여 서로 다른 물리적 상태에 대한 서로 다른 조건부 분포를 능동적으로 유도하며, 이론적으로 아핀 변환까지 작업 중심 요인의 복원을 보장합니다.
3. 보상 없는 오프라인 학습: 보상 신호 없이 시각 관찰과 고유감각 부가 정보에만 의존하여 오프라인 환경에서 고품질 세계 모델링 및 계획을 달성합니다.
4. 이론적 보장: 저자들은 약한 가정 (맥락적 역학, 잠재 - 관측 변동성, 충분한 정렬) 하에 학습된 컴팩트한 잠재 변수가 식별 가능하며 재파라미터화까지 실제 물리적 상태를 복원한다는 정리를 제공합니다.

실험 결과

저자들은 TD-MPC2, DreamerV3, MuZero, DINO-WM 을 포함한 베이스라인과 비교하여 Maze, Wall, Push-T, Lift, Can, Square, Reacher, Cheetah, Hopper 의 9 개 오프라인 벤치마크에서 TC-WM 을 평가했습니다.

• 세계 모델 품질: TC-WM 은 거의 모든 작업에서 가장 낮은 잠재 롤아웃 예측 오차 (MSE) 를 달성했으며, 특히 장기 역학의 정확성이 중요한 조작 환경에서 두드러졌습니다. 이미지 공간 재구성은 경쟁력이 있지만, 주요 이득은 잠재 역학의 안정성과 정확성에 있습니다.
• 계획 성능:
  • 내비게이션 (저차원 행동) 에서 TC-WM 은 TD-MPC2 와 같은 강력한 베이스라인과 동등한 성능을 보입니다.
  • 조작 (고차원, 접촉이 많은 행동) 에서 TC-WM 은 모든 베이스라인을 크게 능가합니다. 예를 들어, Robomimic 의 "Lift" 작업에서 TC-WM 은 54% 의 성공률을 달성한 반면, DINO-WM 은 29%, TD-MPC2 는 31% 였습니다.
• 표현 분석:
  • 선형 프로빙: TC-WM 의 잠재 공간에서 훈련된 선형 회귀 모델은 베이스라인보다 높은 $R^2$ 점수로 고유감각 및 객체 상태를 복원하여, 잠재 변수가 물리적으로 기반을 두고 있음을 확인합니다.
  • 붕괴 방지: TC-WM 은 원시 임베딩을 직접 사용할 때 표현 붕괴를 겪는 DINO-WM 에 비해 잠재 공간의 더 나은 활용 (더 높은 유효 순위) 을 보여줍니다.
  • 강건성: TC-WM 은 시각적 교란 (가우시안 노이즈, 컬러 지터) 하에서 DINO-WM 보다 더 나은 성능을 유지하며 SSIM 감소 폭이 더 작습니다.

중요성 및 주장

이 논문은 기반 임베딩이 직접적인 상태 표현이 아니라 의미적 발판으로 활용되는 것이 가장 효과적이라고 주장합니다. 이러한 임베딩을 물리적 신호와 정렬된 작업 중심 부분 공간으로 투사함으로써 에이전트는 다음과 같은 이점을 얻을 수 있습니다:

• 간결한 잠재 변수: 제어에 필요한 역학을 보존하면서 관련 없는 시각적 세부 사항을 필터링하는 차원 축소.
• 개선된 제어: 노이즈와 관련 없는 변이로 인해 표준 기반 모델이 실패하는 복잡하고 고차원 환경에서 향상된 계획 성능.
• 이론적 건전성: 고정된 시각 인코더에서 작업 관련 요인을 복원하기 위해 선형 투사와 대비적 정렬이 충분하다는 형식적 근거.

저자들은 TC-WM 을 기반 모델의 일반화 능력과 물리적 에이전트에 필요한 제어 가능성 사이의 가교로 위치시키며, 특히 보상 없는 오프라인 학습 시나리오에서의 격차를 해결합니다.