기술 요약: AdaJEPA: 적응형 잠재 세계 모델 (An Adaptive Latent World Model)

문제 정의

잠재 세계 모델(Latent world models), 특히 결합 임베딩 예측 아키텍처(Joint-Embedding Predictive Architectures, JEPAs)는 압축된 잠재 공간에서 미래 상태를 예측함으로써 효율적인 계획(planning)을 가능하게 합니다. 이러한 모델들은 일반적으로 보상이 없는 오프라인 궤적(reward-free offline trajectories)으로 학습되며, 목표 조건부 작업(goal-conditioned tasks)을 위해 모델 예측 제어(Model Predictive Control, MPC)와 함께 사용됩니다. 그러나 결정적인 한계가 존재합니다: 세계 모델은 대개 **학습 후 고정(frozen)**된다는 점입니다.

배포 시, 테스트 타임의 분포 변화(예: 노이즈나 조명과 같은 시각적 변화, 또는 마찰 및 질량과 같은 동역학적 변화)로 인해 모델의 예측이 부정확해지면, MPC 플래너는 잘못된 목적 함수를 최적화하게 됩니다. "잠재적 상상(latent imagination)" 속에서는 효과적으로 보이는 행동들이 실제 환경에서는 실패할 수 있습니다. 단 한 단계의 오류가 계획 수평선(planning horizon) 동안 누적되면 상당한 성능 저하를 초래합니다. 기존 방식들은 이를 해결하기 위해 추가적인 전문가 시연, 타겟 도메인 미세 조정 데이터, 또는 복잡한 외부 루프를 요구하는 경우가 많아 실제 환경 적용을 저해합니다.

방법론: AdaJEPA

저자들은 MPC의 폐루프(closed loop) 내에서 **테스트 타임 적응(test-time adaptation)**을 수행하는 적응형 잠재 세계 모델인 AdaJEPA를 제안합니다. 모델을 정적인 상태로 두는 대신, AdaJEPA는 에이전트 자신의 상호작용으로부터 유도된 자기 지도 신호를 사용하여 세계 모델을 지속적으로 재보정합니다.

계획-실행-적응-재계획 루프 (The Plan–Execute–Adapt–Replan Loop)

핵심 메커니즘은 각 MPC 단계에서 다음과 같이 작동합니다:

1. 계획(Plan): 에이전트는 현재의 (잠재적으로 구식인) 세계 모델을 사용하여 잠재 목표 도달 비용을 최소화하는 일련의 행동 시퀀스를 계획합니다.
2. 실행(Execute): 에이전트는 환경에서 첫 번째 액션 청크($a_t$)를 실행합니다.
3. 관측(Observe): 에이전트는 다음 관측값($o_{t+1}$)을 받습니다.
4. 적응(Adapt): 전이 $\{o_t, a_t, o_{t+1}\}$가 리플레이 버퍼에 추가됩니다. 모델 파라미터는 이 새로운 전이에 대한 잠재 예측 오차를 최소화하기 위해 적은 수의 그래디언트 단계(종종 단 한 단계)를 통해 업데이트됩니다.
5. 재계획(Replan): 업데이트된 모델은 즉시 다음 계획 반복에 사용됩니다.

기술적 구성 요소

• 아키텍처: 기본 세계 모델은 센서리 인코더($\mathcal{E}_s$), 액션 인코더($\mathcal{E}_a$), 그리고 예측기($f_\theta$)로 구성됩니다. 이는 픽셀을 재구성하는 대신 미래의 잠재 표현($z_{t+1}$)을 예측하도록 오프라인에서 학습됩니다.
• 적응 손실(Adaptation Loss): 적응은 사전 학습과 동일한 자기 지도 잠재 예측 목적 함수를 사용합니다:
  $$\mathcal{L}_{ada}(\mathcal{B}) = \frac{1}{|\mathcal{B}|} \sum_{(o_i, a_i, o_{i+1}) \in \mathcal{B}} \ell(f_\theta(z_i, \mathcal{E}_a(a_i)), \text{sg}(z_{i+1}))$$
  여기서 $\mathcal{B}$는 최근 전이들의 버퍼이며, $\ell$은 잠재 예측 손실(예: MSE)이고, $\text{sg}(\cdot)$는 붕괴(collapse)를 방지하기 위한 stop-gradient 연산자입니다.
• 파라미터 업데이트: 효율성을 보장하기 위해 적응은 경량화되어 있습니다. 저자들은 전체 네트워크 대신 모델의 작은 부분 집합(예: 인코더와 예측기의 마지막 레이어들)만을 업데이트합니다.
• 버퍼 전략: 리플레이 버퍼 $\mathcal{B}$는 "recent-N"(가장 최근의 $N$개 전이를 유지) 또는 "hard-N"(예측 오차가 가장 큰 전이를 유지)을 통해 관리될 수 있습니다.

주요 기여

1. 계획 단계에서의 최초 JEPA 적응: 저자들의 지식에 따르면, 이는 JEPA 세계 모델을 사전 학습이나 오프라인 미세 조정에만 의존하는 대신 계획 단계 동안 적응시키는 첫 번째 연구입니다.
2. 폐루프 자기 지도 학습: 이 방법은 실행된 행동의 즉각적인 결과를 자기 지도 신호로 활용하여, 보상 라벨, 전문가 시연, 또는 별도의 데이터 수집 단계의 필요성을 제거합니다.
3. 샘플 및 계산 효율성: 이 프레임워크는 최소한의 계산 오버헤드로 상당한 성능 향상을 달 achieves 합니다. 많은 실험에서, MPC 재계획 단계당 단 한 번의 그래디언트 단계만으로도 충분합니다.
4. 분포 변화에 대한 강건성: 이 접근 방식은 시각적 변화(노이즈, 블러, 색상 변화), 동역학적 변화(질량, 감쇠), 구조적 변화(보지 못한 미로 레이아웃)에 대해 강건함을 입증합니다.

실험 결과

저자들은 다양한 훈련-테스트 분포 변화 하에서 PushT(블록 밀기) 및 PointMaze(내비게이션) 벤치마크를 통해 AdaJEPA를 평가했습니다.

• In-Distribution 성능: 훈련 데이터와 유사한 작업에서도, AdaJEPA는 성공률을 개선하거나(예: 특정 PushObj 형상에서 >20% 이득) 저하 없이 베이스라인 성능을 유지함으로써, In-distribution 적용이 안전함을 증명했습니다.
• Out-of-Distribution (OOD) 성능:
  • 형상 변화: 보지 못한 블록 형상에 대해, AdaJEPA는 고정된 모델에 비해 계획 성공률을 거의 두 배로 높였습니다.
  • 시각적 변화: 가우시안 블러, 솔트 앤 페퍼 노이즈, 어두운 조명 하에서 명확한 이점을 보여주었습니다.
  • 동역학 및 레이아웃 변화: PointMaze에서 AdaJEPA는 낮은 질량이나 높은 감쇠로 인한 실패로부터 회복하였으며, 보지 못한 미로 레이아웃(예: 특정 레이아웃 변화에서 +25.3% 이득)에 대해 성공률을 크게 개선했습니다.
• 지연 시간(Latency): 적응 단계는 파라미터의 작은 부분 집합만을 업데이트하기 때문에 재계획당 무시할 만한 지연 시간(약 0.01–0.03초)을 도입합니다. 또한, 예측 오차로부터 더 빠르게 회복함으로써 에이전트는 종종 더 적은 총 재계획 횟수로 목표에 도달하여 전체 실행 시간을 단축합니다.
• 데이터 부족에 대한 일반화: 데이터가 부족한 상황에서, 테스트 타임 적응은 제한된 훈련 커버리지를 보완했습니다. 예를 들어, 훈련 데이터가 제한적인 경우에도 AdaJEPA는 16배 더 많은 궤적을 사용하여 학습된 고정 모델보다 우수한 성능을 보였습니다.

의의 및 주장

본 논문은 잠재 세계 모델이 학습 후 고정된 상태로 남아있어서는 안 되며, 배포 중에 마주치는 경험으로부터 계속해서 개선되어야 한다고 주장합니다.

• 생물학적 영감: 이 접근 방식은 시스템(예: 소뇌)이 변화하는 입력과 동역학에 따라 행동을 조정하는 생물학적 센서모터 적응(sensorimotor adaptation)에 근거를 두고 있습니다.
• 데이터 스케일링의 보완: 더 크고 다양한 훈련 데이터가 일반화 성능을 높이지만, AdaJEPA는 배포 중 모델을 정교화하기 위한 샘플 효율적인 경로를 제공하며, 이는 데이터 스케일링을 보완하는 역할을 합니다.
• 향후 방향: 저자들은 이 경량 테스트 타임 적응을 지속적 학습(continual learning) 및 능동 학습(active learning)과 결합하여, 시간이 지남에 따라 세계 모델의 커버리지를 확장하고 변화하는 환경에서 더 탄력적인 인지와 계획을 가능하게 하는 것이 다음 단계라고 제안합니다.

저자들은 자신들의 방법이 특정 세계 모델 구현에 구애받지 않으며, 다양한 아키텍처(ResNet, Transformer, DINO 기반)와 플래너(Gradient Descent, CEM) 전반에 걸쳐 일관된 개선을 보여준다는 점을 강조합니다.