ATLAS가 월드모델 계획 잠재표현의 관계 기하를 보존한다

무엇인가

잠재 월드모델은 후보 행동을 굴려 예측한 종단 잠재를 목표 잠재와 비교해 행동을 고른다. 그래서 표현 공간의 기하 구조가 계획 문제의 일부가 된다. 행동적으로 비슷한 상태는 가깝게, 중요한 차이가 있는 상태는 구분되게 남아야 한다. 문제는 기존의 반붕괴(anti-collapse) 정규화가 분산·공분산·전체 분포 같은 통계적 성질만 제약할 뿐, 어떤 상태가 가깝고 먼지를 정해주지 않는다는 점이다. 두 잠재 공간이 똑같이 잘 behaved된 전역 분포를 가져도 상태 간 관계는 전혀 다를 수 있다. 저자들은 LeWM에서 이 간극을 실측한다. 학습 데이터로부터의 거리(kNN novelty)가 계획 실패를 얼마나 잘 예측하는지 단계별로 재보니, PushT에서 실패 예측 AUROC가 인코더의 평균 풀링 패치 특징에서는 0.76인데 계획에 실제로 쓰이는 최종 잠재에서는 0.44로 떨어진다. 상태가 낯설다는 정보가 인코더 안에는 있는데 계획용 표현으로 넘어오면서 크게 약해진 것이다.

어떻게 동작하나

ATLAS는 이 문제를 학습 목적함수에서 직접 겨냥한다. 전체 손실은 예측 손실, WEMReg, OOD-recovery 세 항의 합이다. 구조는 LeWM의 잠재 예측 아키텍처를 따른다. 학습 가능한 ViT 인코더가 관측을 d차원 계획 잠재로 매핑하고, 예측기가 이전 H개의 잠재-행동 쌍을 입력받아 다음 잠재를 예측한다. 첫 번째 구성 요소인 WEMReg(Wasserstein embedding matching)는 잠재 분포의 전역 스케일과 모양을 잡는다. 단위 방향 u를 무작위로 뽑아 스칼라 사영 u^T z의 분포와 표준 정규분포 사이의 Wasserstein-2 거리 제곱을 구하고, 이를 방향에 대해 평균낸다. 실제 구현에서는 미니배치마다 L개의 무작위 방향을 쓰고, 각 방향에서 사영값을 정렬해 표준 정규분포의 분위수에 맞춘다. 정렬된 사영 샘플과 가우시안 목표의 불일치가 곧 1차원 Wasserstein 비용이므로 닫힌 형태의 배치 목적함수가 나온다. LeWM이 쓰는 SIGReg는 유한한 주파수 집합에서 특성함수를 맞추는 방식인데, 논문의 정리 1은 유한 주파수에서는 구별되지 않는 비가우시안 분포가 존재하지만 WEMReg 값은 0보다 크다는 것을 보인다. 즉 유한 주파수 매칭은 놓치는 비가우시안 구조를 사영 Wasserstein 매칭은 잡아낸다.

무엇과 다른가

두 번째 구성 요소인 OOD-recovery는 상태 간 관계를 옮긴다. 각 프레임의 마지막 층 패치 토큰 K개의 평균을 관계적 앵커로 삼는다. 이 표현에서 측정한 novelty가 계획 실패를 훨씬 잘 예측하기 때문이다. 미니배치에서 계획 잠재와 패치 표현 각각에 대해 좌표별로 배치 표준편차로 표준화하고 쌍별 유클리드 거리를 계산한 뒤, 모든 쌍별 거리를 배치 평균으로 나눠 전체 스케일 차이를 제거한다. 그리고 두 정규화 거리 행렬의 제곱오차 평균을 손실로 쓴다. 패치 표현에는 stop-gradient가 걸려 고정된 목표 역할을 한다. 이 손실은 패치 특징 자체를 재현하라고 요구하는 게 아니라 상대적 기하를 보존하라고 요구한다. 명제 2는 관계 보존과 주변 분포 보정이 서로를 결정하지 않는 비중복 제약임을 보인다. 정리 3은 이 목적함수들을 계획 안정성에 연결한다. 계획 잠재와 앵커 표현에서 계산한 평균 kNN novelty 점수의 차이는 관계 왜곡 ε_rel로 상한이 잡히고, 유한 후보 집합과 유계 잠재 스케일 불일치 아래에서 예측 잠재 거리로 고른 후보의 앵커 기하 기준 후회(regret)는 2[e/ℓ + (1+η)ε_rel + ηM]로 묶인다. 예측 손실은 종단 롤아웃 오차 e를, 관계 목적은 ε_rel을, WEMReg는 스케일 불일치 η를 줄인다는 대응이 여기서 나온다.

어떻게 쓰나

실험은 TwoRoom(2D 내비게이션), PushT(접촉이 많은 2D 조작), OGBench-Cube(3D 로봇 조작) 세 가지 오프라인 제어 과제에서 이뤄졌다. 비교 대상은 LeWM, DINO-WM, FAST-WM, PLDM이다. 계획은 동일한 목표 조건 프로토콜과 receding-horizon CEM(후보 300개, 정제 30회, 엘리트 30개)으로 수행하고, 목표는 시연 궤적에서 50스텝 앞을 택한다. 평가 시드 5개(42~46)마다 200 에피소드를 돌려 평균과 표준편차를 보고하며, ATLAS는 모든 과제에서 λ=3.0(WEMReg), μ=0.1(OOD-recovery)을 쓴다. 각 과제·시드마다 200 에피소드를 학습 상태 지지집합으로부터의 거리 중앙값으로 나눠 저 novelty(ID)와 고 novelty(OOD) 절반으로 분리했다. 결과적으로 ATLAS는 세 과제 모두에서 ID와 OOD 양쪽 평균 성공률을 LeWM보다 높였다. 개선폭은 PushT와 OGBench-Cube에서 완만하고 TwoRoom에서 크다. TwoRoom에서 ID는 74.8%에서 87.0%로, OOD는 41.2%에서 61.0%로 올라 가장 큰 향상이 고 novelty TwoRoom 에피소드에서 나왔다. 개선이 OOD에만 국한되지 않고 모든 과제의 저 novelty 부분집합에서도 나타난다는 점을 저자들은 강조한다. 다만 전체 베이스라인 대비 성능은 과제 의존적이다. PushT에서는 두 novelty 부분집합 모두 ATLAS가 최고 평균 성공률을 냈고 OGBench-Cube에서는 여러 방법이 비슷하다. TwoRoom에서는 DINO-WM이 크게 앞서지만, 이는 약 1억 4200만 장의 외부 이미지로 사전학습된 동결 DINOv2-S/14 인코더를 쓰기 때문이며 ATLAS는 외부 시각 사전학습 없이 과제 데이터만으로 학습된다. 저자들은 DINO-WM을 목적함수 수준의 통제된 비교가 아니라 강한 사전학습 표현의 이점을 보여주는 참조점으로 취급한다.

전제와 한계

절제 실험은 SIGReg/WEMReg와 OOD-recovery 켜기/끄기의 2×2 조합으로, 아키텍처·데이터·학습 예산·플래너를 모두 고정하고 목적함수 항만 바꿨다. LeWM은 SIGReg에 OOD-recovery를 끈 조합, ATLAS는 WEMReg에 OOD-recovery를 켠 조합이다. 과제별로 두 목적의 효과가 다르게 나타난다. TwoRoom에서는 WEMReg나 OOD-recovery 중 하나만으로도 LeWM 대비 큰 향상이 나오고 전체 ATLAS는 더 강한 단일 구성 요소와 비슷한 수준이다. PushT에서는 OOD-recovery의 개별 기여가 더 크고 전체 모델이 가장 좋다. OGBench-Cube에서는 OOD-recovery는 LeWM보다 낫지만 WEMReg 단독은 그렇지 않고, 전체 목적함수가 가장 강한 평균 성능을 회복한다. ATLAS는 세 과제 모두에서 최고 평균 성공률을 달성하지만 더 강한 단일 구성 요소 대비 이득은 작다. 이는 두 목적이 서로 다른 표현 성질을 제약하되 하류 이득이 반드시 가산적이지는 않고 표현 병목에서 겹칠 수 있음을 시사한다. 진단 실험도 정리 3의 세 요인과 맞물린다. PushT에서 ATLAS의 계획 잠재는 실패 예측 신호를 0.75까지 유지해 패치 표현의 0.77에 근접했고, WEMReg를 쓰는 변형은 공분산 오차와 사영 분포 불일치가 특히 TwoRoom에서 줄었으며, ATLAS는 측정한 모든 지평에서 LeWM보다 낮은 다단계 잠재 예측 오차를 보였다.

개발자 관점에서 이 논문이 주는 실무적 신호는 명확하다. JEPA 스타일 월드모델이나 잠재 공간 계획기를 만들 때, 반붕괴 정규화가 잘 걸려 있다는 사실만으로 계획 표현이 건강하다고 가정하면 안 된다. 인코더 중간 표현에서 계산한 novelty나 kNN 거리 신호가 최종 계획 잠재에서도 살아 있는지 직접 진단해보는 것이 첫 단계다. 신호가 죽어 있다면 ATLAS처럼 중간 표현의 정규화 쌍별 거리를 계획 잠재에 증류하는 항을 추가하는 것이 비교적 저렴한 처방이 된다. 여기에 분포 보정이 필요하면 유한 주파수 특성함수 매칭 대신 사영 Wasserstein 매칭을 고려할 수 있고, 배치 안에서 정렬과 분위수 매칭만으로 닫힌 형태 목적함수를 구현할 수 있다는 점도 실무적으로 유용하다. 다만 두 항의 이득이 항상 더해지지는 않으므로, 자신의 과제에서 어느 쪽이 실제로 기여하는지 2×2 절제로 확인해야 한다.

한계와 전제도 분명하다. 저자들은 표현·롤아웃 진단을 인과적 설명이 아니라 메커니즘 지향 진단으로 규정한다. PCA 시각화는 2차원 사영일 뿐이라 정량적 kNN 진단을 보조하는 정성적 근거로만 쓴다. 절제에서 전체 ATLAS가 더 강한 단일 구성 요소를 크게 앞서지 못한다는 점, 즉 두 목적의 이득이 겹칠 수 있다는 점도 인정한다. DINO-WM 비교는 동결된 대규모 사전학습 인코더를 쓰므로 통제된 목적함수 비교가 아니라고 명시한다. 또한 이론의 후회 상한은 고정된 유한 후보 집합과 유계 잠재 스케일 불일치를 전제로 하며, 정확한 가정과 정의·증명은 부록 C.1에 있다.