기술 요약: DWM: 잠재 세계 모델에서 세계 효과와 행동의 분리

문제 정의

잠재 세계 모델(Latent world models)은 현대적 모델 기반 제어의 기초로서, 픽셀 관측으로부터 직접적인 계획 수립과 정책 학습을 가능하게 합니다. 그러나 현재의 행동 조건부 정식화(action-conditioned formulations)는 **감독 수준의 한계(supervision-level limitation)**를 겪고 있습니다. 즉, 단일하고 차별화되지 않은 타겟을 사용하여 다음 잠재 상태를 예측한다는 점입니다. 이로 인해 모델은 에이전트의 행동에 의해 유도된 변화와 환경의 고유한 역학으로부터 발생하는 변화 모두를 하나의 단일한 학습 신호로 흡수해야만 합니다.

실제 물리 시스템에서 전이는 종종 두 가지 이질적인 원인으로부터 발생합니다:

1. 행동 유도 효과(Action-driven effects): 에이전트의 제어에 의해 유도되는 변화.
2. 세계 효과(World effects): 널(null) 행동 하에서도 지속되는 행동 불변적 변화 (예: 중력에 의한 미끄러짐, 관성, 접촉 반동 또는 지속적인 드리프트).

이들을 하나의 타겟으로 융합하면 잠재 전이 내에서 두 구성 요소가 서로 얽히게 됩니다. 결과적으로, 모델은 관찰된 변화를 그 근본 원인에 명확하게 귀속시킬 수 없습니다. 이러한 한계는 정적인 벤치마크(예: 평평한 테이블 위의 밀기)에서는 숨겨져 있을 수 있으나, 지속적인 자율 역학이 존재하는 환경에서는 치명적이 되어 부정확한 다단계 롤아웃(multi-step rollouts)과 최적화되지 않은 계획 수립으로 이어집니다.

방법론: 분해된 세계 모델 (DWM)

저자들은 기본 아키텍처나 추론 파이프라인을 변경하지 않고 세계와 행동 효과의 분리를 구현하는 훈련 프레임워크인 **DWM (Decomposed World Model)**을 제안합니다.

핵심 아키텍처

DWM은 표준적인 행동 조건부 잠재 세계 모델의 구성 요소(인코더 $\phi$, 행동 조건부 예측기 $g$, 그리고 프로젝션 헤드)를 그대로 유지합니다:

• 예측 헤드 ($h_p$): 베이스라인으로부터 유지됩니다. 전체 다음 잠재 상태 $\hat{z}$를 예측하며, 추론 시 사용되는 유일한 헤드입니다.
• 세계 헤드 ($h_w$): 예측 상태 $r_t$에 병렬로 부착된 보조적이고 경량화된 MLP입니다. 이는 행동 불변적 전이 성분($\hat{z}_w \approx z_t + \Delta z^{world}_t$)을 포착하도록 훈련됩니다.
• 행동 유도 성분(Action-Driven Component): 잔차(residual) $\hat{z}_a := \hat{z} - \hat{z}_w$로 정의됩니다.

모델은 $\hat{z} = \hat{z}_w + \hat{z}_a$라는 항등식을 강제하여, 예측된 전이를 행동 불변적 세계 성분과 상보적인 행동 유도 성분으로 명시적으로 분해합니다.

훈련 목적 함수

전체 손실 함수는 표준 세계 모델 손실과 두 개의 보조 항을 결합합니다:
$$L = L_{WM} + \lambda_{wc}L_{wc} + \lambda_{orth}L_{orth}$$

1. 세계-대조 손실 (World-Contrastive Loss, $L_{wc}$): 세계 헤드에 적용되는 InfoNCE 목적 함수입니다. 동일한 상태 이력에 서로 다른 현재 행동(행동 섭동)을 입력하여 양의 쌍(positive pairs)을 구성함으로써, $\hat{z}_w$가 행동에는 불변하면서도 서로 다른 상태들에 대해서는 변별력을 갖도록 강제합니다.
2. 직교성 제약 (Orthogonality Constraint, $L_{orth}$): $|\cos(\hat{z}_w, \hat{z}_a)|$로 정의되며, 이 항은 행동 불변 성분과 행동 유도 성분이 잠재 공간에서 독립적인 방향을 점유하도록 유도하여, 행동 효과가 세계 표현 속으로 흡수되는 것을 방지합니다.

결정적으로, 보조 세계 헤드는 훈련 후에 폐기되므로 추론 시의 파이프라인은 베이스라인과 동일하게 유지됩니다.

주요 기여

1. 감독의 한계 식별: 본 논문은 단일 타겟 감독이 세계 모델에서 에이전트 유도 변화와 자율적 환경 역학을 구분하지 못해 얽힌 표현을 생성한다는 점을 밝혀냈습니다.
2. 새로운 벤치마크 스위트 (W-Variants): 저자들은 표준 작업에 지속적인 행동 불변 역학(중력에 의한 미끄러짐, 수직 평면 중력, 지속적인 드리프트)을 주입한 세 가지 표준 벤치마크의 제어된 변형인 PushT-W, Reacher-W, TwoRoom-W를 도입했습니다. 또한 상태 의존적 진동 역학이 있는 Ball-in-Cup에서도 평가를 수행했습니다.
3. DWM 프레임워크: 기존의 잠재 세계 모델에 보조 세계 헤드와 상보적인 분리 목적 함수를 추가하는 훈련 프레임워크로, 백본이나 추론 방식의 변경을 요구하지 않습니다.
4. 실증적 검증: DWM이 세계 효과가 존재하는 설정에서 계획 성공률과 잠재 역학 정확도를 개선하는 동시에, 표준 작업에서는 성능을 유지함을 입증했습니다.

실험 결과

저자들은 CEM(Cross-Entropy Method) 계획법을 사용하여 제안된 벤치마크에서 DWM을 LeWM 베이스라인과 비교 평가했습니다.

• 계획 성공률: W-변형 작업에서 DWM은 CEM 계획 성공률에서 평균 절대 개선치 **13.1%**를 달 기록했습니다:
  • PushT-W: 32.0% $\to$ 44.0% (+12.0 pp)
  • Reacher-W: 60.0% $\to$ 70.7% (+10.7 pp)
  • TwoRoom-W: 60.0% $\to$ 76.7% (+16.7 pp)
  • Ball-in-Cup: 36.8% $\to$ 42.8% (+6.0 pp)
• 베이스라인 성능: 세계 효과가 없는 원래의 "평평한(flat)" 대응 작업에서는 DWM이 단일 헤드인 LeWM과 대등한 성능을 유지하여, 보조 손실이 세계 효과가 미미할 때 성능을 저하시키지 않음을 보여주었습니다.
• 표현 진단: 분석 결과, 세계 헤드는 거의 행동 불변적이 된 반면(행동 섭동 하에서 분산이 수 차례 감소), 예측 헤드는 행동에 대한 민감도를 유지했습니다.
• 예측 품질: DWM은 1단계 및 horizon-20 자기회귀 예측 오차를 크게 줄였습니다 (예: PushT-W에서 -89.8% 오프라인 손실). 이는 분해가 잠재 역학을 더 정확하게 만든다는 것을 입증합니다.
• OOD(Out-of-Distribution) 일반화: 중력 벡터가 변화된 PushT-W 환경에서, DWM은 높은 계획 성공률(44.0%)을 유지한 반면 베이스라인은 성능이 저하되었습니다(30.0%). 또한 학습된 행동 유도 성분은 중력 변화에도 안정적이었습니다 (코사인 유사성 $\approx$ 0.999).

의의 및 주장

본 논문은 훈련 신호 수준에서 행동 불변적 세계 효과와 행동 유도 변화를 분리하는 것이 물리적으로 현실적인 제어 상황의 행동 조건부 잠재 세계 모델에 유용한 귀납적 편향(inductive bias)으로 작용한다고 주장합니다.

저자들은 DWM이 아키텍처의 전면적인 개편 없이도 기존의 감독 전략이 가진 근본적인 한계를 해결한다는 점을 강조합니다. "세계가 스스로 무엇을 하는가"와 "에이전트가 무엇을 변화시키는가"를 명시적으로 분리함으로써, DWM은 더 정확한 잠재 역학을 학습할 수 있게 하며, 이는 더 나은 다단계 계획 수립과 에이전트의 행동과 무관하게 지속되는 환경 역학에 대한 강건함으로 직결됩니다. 결과적으로, 이 분해 방식은 에이전트의 행동과 독립적으로 진화하는 환경이 존재하는 작업에서 매우 중요하다는 것을 시사합니다.