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

Temporal-Distance JEPA: Plan-Aware Representation Learning for Latent World Model Predictive Control

합쳐진 임베딩 예측 아키텍처(JEPAs)는 픽셀을 재구성하는 대신 표현 공간에서 예측함으로써 세계 모델을 학습하며, 이는 오프라인 데모 로그로부터의 잠재 모델 예측 제어를 위한 자연스러운 기반이 됩니다. JEPA 스타일의 학습은 단기 잠재 예측을 최적화하는 반면, 계획은 목표 진행에 따른 여러 단계의 가상 미래 시퀀스를 순위화해야 합니다. 기존의 JEPA 플래너는 종종 이러한 순위를 임베딩 기하학, 일반적으로 표현 학습의 부산물로 발생하는 잠재 유클리드 거리에서 가져옵니다. 본 논문에서는 보상 없이 수집된 트래젝토리로부터 방향성 시간 비용을 추출하는 Temporal-Distance-JEPA를 제안합니다. 동일 트래젝토리 내 단계 순서는 긍정적인 목표를 제공하고, 다른 트래젝토리 쌍은 휴리스틱 부정으로 사용되며, 롤아웃 일관성 항은 플래너의 수평과 일치하도록 설계됩니다. 추출된 감독 신호는 두 가지 역할을 합니다. 첫째, 진행이 토폴로지적일 때 배포되는 계획 비용으로 사용되며, 둘째, 표현 학습 신호로서 접촉 기하학이 지배할 때 유클리드 계획을 개선합니다. 고정 평가 환경에서 추출된 비용을 사용하는 것은 LeWM의 97.4%에 비해 Two-Room 문제에서 성공률을 100.0%로 향상시킵니다. 또한, 동일한 시간적으로 학습된 체크포인트에서 공유 유클리드 계획은 LeWM보다 OGB-Cube 문제를 14.2점 개선하고 Push-T 성능을 향상시킵니다. 고정 평가 환경에서 Temporal-Distance-JEPA는 LeWM 및 동시대 RC-aux 기준선과 비교하여 모든 환경에서 동일하거나 뛰어넘는 성능을 보입니다. 실험 결과, 방향성 헤드, 트래젝토리 간 부정 샘플링, 롤아웃 일관성이 각각 기여하는 것으로 나타났습니다. Temporal-Distance-JEPA는 오프라인 로그 내 시간적 진행 구조를 발견하고 계획 시 배포와 함께 비용 형식을 공동 설계함으로써 JEPA 세계 모델 플래너의 학습과 계획 간 격차를 줄입니다. 관련 코드는 https://github.com/HKBU-KnowComp/Temporal-Distance-JEPA 에서 확인할 수 있습니다.

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, 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, Temporal-Distance-JEPA matches or exceeds both methods on every environment. Ablations show that the directed head, cross-trajectory negatives, and rollout consistency each contribute. Temporal-Distance-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/Temporal-Distance-JEPA.