[논문리뷰] Temporal-Distance JEPA: Plan-Aware Representation Learning for Latent World Model Predictive Control

링크: 논문 PDF로 바로 열기

본 논문은 JEPA(Joint-Embedding Predictive Architectures) 기반의 Latent World Model이 가지는 'Train-Plan Gap' 문제를 해결하기 위해, 보상 없이 오프라인 로그에서 직접적인 시간적 진행 구조를 추출하는 TD-JEPA(Temporal-Distance JEPA) 프레임워크를 제안한다.

1. Key Terms & Definitions (핵심 용어 및 정의)

• TD-JEPA: 오프라인 로그에서 directed temporal cost를 마이닝하여 Latent MPC의 계획 능력을 향상시키는 JEPA 확장 모델이다.
• Latent MPC: 학습된 Latent Space 내에서 미래의 임베딩 궤적을 예측하고, 보상 함수나 거리 지표를 최적화하는 동작 시퀀스를 탐색하는 방식이다.
• MRN (Metric-Residual Network): 비대칭적인 목표 도달 거리(asymmetric reachability)를 표현하기 위해 Symmetric Component와 Directed Residual을 결합하여 cost를 계산하는 모듈이다.
• Train-Plan Gap: JEPA 모델이 학습 시에는 단기 예측 성능에 최적화되지만, 실제 계획 시에는 임베딩의 기하학적 구조(Euclidean distance)에 의존하게 되어 장기적인 목표 달성 순위를 정확히 반영하지 못하는 불일치 현상이다.
• Rollout-Consistency: Planner가 사용하는 오픈 루프 호라이즌과 동일한 길이로 학습 시 궤적 예측 일관성을 강제하여, 학습과 계획 간의 아키텍처적 간극을 줄이는 기법이다.

2. Motivation & Problem Statement (연구 배경 및 문제 정의)

기존의 LeWM(LeWorldModel)과 같은 JEPA 기반 모델들은 학습 시에는 효율적인 표현력을 제공하지만, 테스트 시 계획(planning) 단계에서는 단순히 latent space 내의 Euclidean Distance를 목표 점수(cost)로 사용한다 [Figure 1]. 이러한 기하학적 거리는 실제 작업의 진행 단계(steps-to-goal)나 비대칭적인 도달 가능성을 완벽하게 대변하지 못하므로, 결과적으로 MPC가 최적의 미래 시퀀스를 순위 매기는 데 한계를 드러낸다. 저자들은 보상이 없는 로그 데이터 자체에 포함된 시간적 순서(temporal order)를 학습하여 이러한 순위 매기기 문제를 개선하고자 한다.

Figure 1 — Push-T에서의 temporal-gap 정렬 비교

3. Method & Key Results (제안 방법론 및 핵심 결과)

TD-JEPA는 기존 LeWM의 백본을 유지하면서, Directed Temporal Cost를 마이닝하기 위한 세 가지 핵심 요소(같은 궤적 내 순서쌍을 활용한 positive 타겟, 교차 궤적을 이용한 negative heuristic, horizon-aligned rollout consistency)를 도입한다 [Figure 2]. 제안된 MRN 헤드는 학습된 비대칭 cost 함수 **$d_{\psi}$**를 출력하며, 이는 위상적 작업(topology-dominated)에서는 직접적인 계획 비용으로, 접촉 위주의 작업(contact-rich)에서는 Euclidean 계획을 보완하는 표현 학습 신호로 사용된다. 실험 결과, TD-JEPA는 Two-Room 환경에서 100.0%의 성공률을 기록하며 LeWM(97.4%)을 상회하였다. 또한, OGB-Cube 작업에서는 동일한 체크포인트에서 Latent $\ell_{2}$ 계획을 수행하여 LeWM 대비 14.2점의 성능 향상을 달성하였다. 모든 환경에서 RC-aux 베이스라인과 비교했을 때 동등하거나 더 우수한 성능을 보여주며 제안 방식의 범용성을 입증하였다.

Figure 2 — TD-JEPA 학습 및 MPC 파이프라인

4. Conclusion & Impact (결론 및 시사점)

본 논문은 오프라인 데이터를 통해 계획 중심의 시간적 진행 구조를 명시적으로 학습하는 것이 Latent World Model의 성능에 필수적임을 입증하였다. 특히, 작업의 특성(토폴로지 vs. 접촉 기반)에 따라 학습된 cost의 적용 방식을 다르게 설계해야 한다는 점은 향후 Autonomous Machine Intelligence 연구에 중요한 디자인 지침을 제공한다. 이 연구는 모델이 단순히 데이터를 재구성하거나 짧은 미래를 예측하는 것을 넘어, 계획 단계에서의 실질적인 목표 달성 능력을 실시간으로 제어할 수 있는 경로를 제시하였다.

⚠️ 알림: 이 리뷰는 AI로 작성되었습니다.