AdaJEPA: 적응형 잠재 공간 모델

AdaJEPA: An Adaptive Latent World Model

잠재 공간 모델은 고차원 관찰 데이터를 기반으로 미래 상태를 예측하여 계획 수립을 가능하게 합니다. 그러나 이러한 모델은 일반적으로 테스트 단계에서 고정되어 사용됩니다. 이로 인해 예측 정확도가 떨어지면 계획이 실패할 수 있으며, 특히 테스트 시간 동안 데이터 분포가 변화하는 경우 더욱 그렇습니다. 이러한 문제를 해결하기 위해, 본 논문에서는 모델 예측 제어(MPC)의 폐쇄 루프 내에서 테스트 단계 적응을 수행하는 적응형 잠재 공간 모델인 AdaJEPA를 제안합니다. AdaJEPA는 학습 후 첫 번째 행동 단계를 계획하고 실행한 다음, 관찰된 다음 상태 변화를 자기 지도적 적응 신호로 사용하고 업데이트된 모델을 사용하여 재계획합니다. 이러한 폐쇄 루프 업데이트는 추가적인 전문가 데모 없이도 지속적으로 세계 모델을 재조정합니다. 다양한 목표 달성 작업에서 AdaJEPA는 MPC 재계획 단계당 단 한 번의 경사 하강법만으로도 계획 성공률을 크게 향상시킵니다.

Latent world models enable planning from high-dimensional observations by predicting future states in a compact latent space. However, these models are typically kept frozen at test time: when their predictions become inaccurate, planning can fail, especially under test-time distribution shift. To address this, we propose AdaJEPA, an adaptive latent world model that performs test-time adaptation within the closed loop of model predictive control (MPC). After training, AdaJEPA plans and executes the first action chunk, uses the observed next-state transition as a self-supervised adaptation signal, and replans with the updated model. This closed-loop update continuously recalibrates the world model without additional expert demonstrations. Across a range of goal-reaching tasks, AdaJEPA substantially improves planning success with as few as one gradient step per MPC replanning step.