기술 요약: 프리즘틱 세계 모델 (PRISM-WM)

1. 문제 제기

로봇 분야에서의 모델 기반 강화학습 (MBRL) 과 계획 (planning) 은 물리적 역학의 하이브리드적 성격에 의해 근본적으로 제한받습니다. 특히 보행 및 조작과 같은 실세계 시스템은 접촉, 충격, 정지 대 미끄러짐과 같은 이산적 사건으로 단절된 연속 운동을 포함합니다. 이러한 상호작용은 기본 벡터장에서의 날카로운 불연속성을 수반하는 비부드럽고 매우 비선형적인 역학을 생성합니다.

전통적인 잠재 세계 모델은 일반적으로 글로벌 연속성을 강제하는 단일 구조 (monolithic) 신경망을 사용합니다. 이론적으로 보편적 근사가 가능함에도 불구하고, 이러한 모델들은 실제로는 뚜렷한 역학 모드를 과도하게 평활화 (over-smooths) 하는 평균화된 표현을 학습합니다. 궤적 최적화 계획자에게 있어, 이러한 과도한 평활화는 물리적 경계에서 특히 장기적 예측 (long-horizon lookaheads) 동안 누적 오차를 유발합니다. 그 결과 다음과 같은 문제가 발생합니다:

• 접촉 경계 전반에 걸친 스펙트럼 흐림 (Spectral blurring) 으로 인한 심각한 기울기 불일치.
• 고차원 구성 공간에서의 누적 드리프트 (Compounding drift) 로 인한 장기 계획의 신뢰성 상실.
• 표준 모델이 안정적 및 불안정적 영역을 구분하지 못해 접촉이 풍부한 환경에서 발생하는 정책 붕괴 (Policy collapse).

2. 방법론: 프리즘틱 세계 모델 (PRISM-WM)

이러한 한계를 해결하기 위해 저자들은 복잡한 하이브리드 역학을 구성 가능한 원소로 분해하도록 설계된 구조화된 아키텍처인 PRISM-WM을 소개합니다. 이 모델은 새로운 잠재 직교화 (latent orthogonalization) 목적 함수를 갖춘 맥락 인식형 전문가 혼합 (Mixture-of-Experts, MoE) 프레임워크를 기반으로 구축됩니다.

핵심 아키텍처

• 맥락 인식 게이트 (Context-Aware Gating): 게이트 네트워크 ($G$) 는 상태 $z_t$, 행동 $a_t$, 그리고 선택적으로 작업 임베딩 $e_{task}$로 구성된 컨텍스트 벡터 $c_t$를 수신하여 소프트맥스 정규화된 가중치 $w_t$를 생성합니다. 이 메커니즘은 명시적 지도 없이도 (예: 비행 대 지지) 활성 물리 영역을 암시적으로 식별합니다.
• 전문가 (Specialized Experts): 상태는 $K$개의 전문가 네트워크 $\{E_k\}$로 라우팅됩니다. 각 전문가는 특정 국소 역학 영역 내의 연속 흐름을 학습하는 데 특화되어, 상충되는 모드 간의 간섭을 방지합니다.
• 잔차 예측 (Residual Prediction): 다음 잠재 상태는 현재 상태에 전문가 출력의 가중 합을 더하여 예측됩니다:
  $$\hat{z}_{t+1} = z_t + \sum_{k=1}^{K} w_{t,k} \cdot v_k(x_t)$$
  여기서 $v_k$는 $k$번째 전문가가 예측한 잔차 역학을 나타냅니다.

주요 혁신: 잠재 직교화

모드 붕괴 (mode collapse, 전문가들이 중복된 특징을 학습하는 현상) 를 방지하고 다양성을 보장하기 위해, PRISM-WM 은 미분 가능한 그람 - 슈미트 (Gram-Schmidt) 직교화 계층을 도입합니다.

• 각 전문가가 penultimate-layer 벡터 $u_k$를 출력한 후, 모델은 안정화된 그람 - 슈미트 과정을 적용하여 직교 기저 벡터 $v_k$를 생성합니다.
• 이는 전문가 표현이 슈테ifel 다양체 (Stiefel manifold) 위에 놓이도록 강제하여, 잔차 역학을 위한 최대 랭크의 비중복 부분 공간을 확보합니다.
• 물리적 접촉 모드는 순서가 없으나, 이러한 직교화는 전문가들이 잔차 다양체의 고유한 구성 요소를 포착하도록 보장하며, 게이트 네트워크가 가중 합을 통해 순서가 없는 물리 모드를 재구성할 수 있게 합니다.

계획과의 통합

PRISM-WM 은 TD-MPC 계획 프레임워크에 통합됩니다. 복잡한 역학을 비중복 전문가로 인수분해함으로써, 모델은 한 단계 모델 잔차를 $\epsilon$에서 약 $\epsilon/K\alpha$로 감소시킵니다. 이러한 구조적 분해는 누적 드리프트를 완화하여 고정밀도 장기 롤아웃 (rollouts) 과 정책 학습을 위한 안정적인 기울기 전파를 가능하게 합니다.

3. 주요 기여

1. 아키텍처 분해: 단일 구조 글로벌 근사에서 벗어나, 맥락 인식형 MoE 아키텍처와 잠재 직교화를 결합하여 하이브리드 역학을 분해하는 PRISM-WM 의 도입.
2. 직교화 목적 함수: 미분 가능한 그람 - 슈미트 직교화를 사용한 새로운 학습 목적 함수로, 전문가 다양성을 강제하고 모드 붕괴를 방지하며, 잠재 공간을 이산적 물리 전이와 구조적으로 정렬시킴.
3. 이론적 및 실증적 검증: 이러한 분해가 고차원 공간에서 누적 오차를 감소시키고, 계획 (TD-MPC) 과 직접 정책 학습 (PWM) 모두에 대한 손실 지형 (loss landscape) 의 기하학을 개선함을 입증.
4. Zero-Shot 전이 능력: 특정 역학 모드 (예: 마찰 대 수직 지지) 를 격리함으로써, 직교 기저가 서로 다른 접촉 물리 (예: 소프트 접촉 IsaacGym 에서 강성 접촉 MuJoCo 로) 를 가진 시뮬레이터 간에 모델이 일반화될 수 있음을 보여주는 증거.

4. 실험 결과

저자들은 고차원 휴머노이드, 다중 작업 설정, Zero-shot 시뮬레이션 간 전이를 포함한 연속 제어 벤치마크에서 PRISM-WM 을 평가했습니다.

• 고차원 보행 (DiffRL 및 Humanoid-Bench): PRISM-WM 은 단일 구조 베이스라인 (TD-MPC2, DreamerV3, PWM, SAC) 에 비해 더 높은 샘플 효율성과 점근적 성능을 달성했습니다. 특히, 비부드러운 접촉 사건의 평균화로 인해 이족 보행 작업 (Humanoid, SNU-Humanoid) 에서 정책 붕괴를 경험한 단일 구조 모델들과 달리, PRISM-WM 은 안정적인 최적화를 유지했습니다.
• 다중 작업 일반화: 30 개 작업 DeepMind Control Suite 벤치마크에서 PRISM-WM 은 평균 정규화 점수 0.531을 달성하여, 단일 구조 TD-MPC2 베이스라인 (0.430) 대비 23.5% 향상을 보였습니다. 이 모델은 30 개 작업 중 26 개에서 베이스라인을 능가하여 이질적 역학 간의 부정적 전이를 효과적으로 완화했습니다.
• Zero-Shot 시뮬레이션 간 전이: IsaacGym 에서 MuJoCo 로 전이할 때, 단일 구조 MLP 베이스라인은 12% 성공률로 47.8%의 성능 저하를 겪었습니다. PRISM-WM 은 이 저하를 11.8% 로 줄이고 88%의 성공률을 달성했습니다. 직교화는 모델이 글로벌 역학 다양체를 훼손하지 않고 특정 전문가 (예: 수평 마찰) 에 대한 가중치를 조정할 수 있게 했습니다.
• 안정성 및 효율성: 제거 실험 (Ablation studies) 은 직교화가 제약 없는 MoE 모델에서 관찰되는 보상 추정 저하를 방지함을 확인했습니다. NVIDIA RTX 4090 에서의 추론 효율성 테스트 결과, PRISM-WM (K=4) 은 최소한의 메모리 오버헤드로 단일 구조 베이스라인 (7,499 FPS) 을 약간 능가하는 7,712 FPS를 달성했습니다.

5. 중요성 및 주장

본 논문은 PRISM-WM 이 모델 기반 에이전트를 위한 기초 모델로서 궤적 최적화 알고리즘을 위한 고정밀도 기질 (high-fidelity substrate) 을 제공한다고 주장합니다. 복잡한 물리를 직교 원소로 분해함으로써, 이 모델은 다음과 같은 이점을 제공합니다:

• 접촉이 풍부한 환경에서 신뢰할 수 있는 장기 추론을 가능하게 하는 계획 시간 범위 내의 장기 드리프트 완화.
• 정책 학습을 위한 기울기 전파 안정화로, 파손된 손실 지형을 기능적으로 부드러운 볼록 혼합으로 변환.
• 과도한 평활화로 인해 표준 세계 모델이 실패하는 하이브리드 시스템에서 강건한 의사결정 가능.

저자들은 복잡한 물리를 직교 원소로 분해하는 것이 "접촉이 풍부한 제어의 숙달에 근본적으로 필수적"이라고 결론지었습니다. 현재 프레임워크는 전문가 수 $K$의 수동 조정이 필요하지만, 이 작업은 미래의 구체적 에이전트 (embodied agents) 를 위한 범용적이고 물리 기반의 세계 모델을 구축하기 위한 구조적 기초를 확립합니다.