시간적 거리 JEPA: 잠재 세계 모델 예측 제어를 위한 계획 인식 표현 학습

초록

공동 임베딩 예측 아키텍처(JEPA)는 픽셀을 재구성하는 대신 표현 공간에서 예측을 통해 세계 모델을 학습하므로, 오프라인 시연 로그로부터의 잠재 모델 예측 제어를 위한 자연스러운 백본이 된다. JEPA 스타일 훈련은 단기 잠재 예측을 최적화하는 반면, 계획은 목표 진행에 따른 상상된 미래의 다단계 순위 매기기를 필요로 한다. 기존의 JEPA 기반 계획기는 일반적으로 이러한 순위 매기기를 임베딩 기하학, 전형적으로는 표현 학습의 부산물로서 발생하는 잠재 유클리드 거리로부터 상속받으며, 이는 로그에서 추출된 진행 비용이 아니다. 우리는 시간 거리 JEPA(TD-JEPA)를 제안한다. 이는 LeWM 인코더-예측기 백본을 유지하고, 보상 없는 궤적으로부터 방향성 시간 비용을 추출한다. 동일 궤적 단계 순서는 양성 목표를 제공하고, 교차 궤적 쌍은 휴리스틱 음성으로 작용하며, 롤아웃 일관성 항은 계획 수평과 일치시킨다. 추출된 학습 신호는 두 가지 역할을 수행한다. 진행이 위상적인 경우 배포된 계획 비용으로, 그리고 접촉 기하학이 지배적인 경우 유클리드 계획을 개선하는 표현 신호로 기능한다. 고정 평가에서 추출된 비용을 배포하면 Two-Room 성공률이 LeWM의 97.4% 대비 100.0%로 상승하며, 동일한 시간 훈련 체크포인트에서 공유 유클리드 계획을 사용하면 OGB-Cube에서 LeWM 대비 14.2점 향상되고 Push-T도 개선된다. 고정 평가에서 LeWM 및 동시대의 RC-aux 기준과 비교했을 때, TD-JEPA는 모든 환경에서 두 방법과 동등하거나 능가한다. 소거 실험은 방향성 헤드, 교차 궤적 음성, 롤아웃 일관성이 각각 기여함을 보여준다. TD-JEPA는 오프라인 로그에서 시간적 진행 구조를 발견하고 계획 시간 배포와 비용 형태를 공동 설계함으로써 JEPA 세계 모델 계획기의 훈련-계획 격차를 좁힌다. 코드는 https://github.com/HKBU-KnowComp/TD-JEPA에서 확인할 수 있다.

English

Joint-Embedding Predictive Architectures (JEPAs) learn world models by predicting in representation space rather than reconstructing pixels, making them a natural backbone for latent model predictive control from offline demonstration logs. JEPA-style training optimizes short-horizon latent prediction, whereas planning requires a multi-step ranking of imagined futures by goal progress. Prior JEPA planners often inherit that ranking from embedding geometry, typically latent Euclidean distance, which arises as a byproduct of representation learning rather than as a progress cost mined from the logs. We propose temporal-distance JEPA (TD-JEPA), which retains the LeWM encoder--predictor backbone and mines a directed temporal cost from reward-free trajectories: same-trajectory step order supplies positive targets, cross-trajectory pairs act as heuristic negatives, and a rollout-consistency term matches the planner horizon. The mined supervision serves two roles: as the deployed planning cost when progress is topological, and as a representation signal that improves Euclidean planning when contact geometry dominates. Under locked evaluation, deploying the mined cost raises Two-Room success to 100.0% versus LeWM's 97.4%, while shared Euclidean planning on the same temporally trained checkpoint raises OGB-Cube by 14.2 points over LeWM and improves Push-T. Against LeWM and the concurrent RC-aux baseline under locked evaluation, TD-JEPA matches or exceeds both methods on every environment. Ablations show that the directed head, cross-trajectory negatives, and rollout consistency each contribute. TD-JEPA narrows the train--plan gap for JEPA world-model planners by discovering temporal progress structure in offline logs and co-designing cost form with plan-time deployment. Code is available at https://github.com/HKBU-KnowComp/TD-JEPA.