時間距離JEPA：面向潛在世界模型預測控制的規劃感知表示學習

摘要

聯合嵌入預測架構（JEPAs）透過在表徵空間中進行預測而非重建像素來學習世界模型，使其成為從離線示範日誌進行潛在模型預測控制（latent model predictive control）的天然基礎架構。JEPA風格的訓練最佳化短期潛在預測，而規劃則需要根據目標進度對想像中的未來進行多步驟排序。先前的JEPA規劃器通常從嵌入幾何（embedding geometry）繼承該排名，通常是潛在歐幾里得距離（latent Euclidean distance），而這只是表徵學習的副產物，並非從日誌中挖掘出的進度成本。我們提出時序距離JEPA（TD-JEPA），其保留LeWM編碼器-預測器骨幹，並從無獎勵軌跡中挖掘有向時序成本：同軌跡步驟順序提供正目標，跨軌跡配對作為啟發式負樣本，以及一個展開一致性項（rollout-consistency term）以匹配規劃器視野。所挖掘的監督訊號扮演兩個角色：當進度是拓撲性質時，作為部署的規劃成本；以及當接觸幾何主導時，作為改善歐幾里得規劃的表徵訊號。在鎖定評估（locked evaluation）下，部署挖掘的成本將Two-Room的成功率提升至100.0%，相較之下LeWM為97.4%；而在相同時序訓練的檢查點上共享歐幾里得規劃，則使OGB-Cube比LeWM提升14.2個百分點，並改善了Push-T。在鎖定評估下，與LeWM及同期RC-aux基線相比，TD-JEPA在每個環境上都達到或超越了這兩種方法。消融實驗顯示，有向頭（directed head）、跨軌跡負樣本及展開一致性各有貢獻。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.