Temporal-Distance JEPA: 面向潜在世界模型预测控制的规划感知表示学习

TL;DR TD-JEPA 从离线轨迹中挖掘有序时间距离作为规划成本，与 JEPA 世界模型协同训练，弥合了短程预测与多步规划之间的鸿沟，显著提升潜在 MPC 在多种环境下的成功率。

问题

问题背景

模型预测控制(MPC)与基于联合嵌入预测架构(JEPA)的潜在世界模型正成为从离线演示日志中学习规划策略的有力范式。这类方法在表示空间内预测未来状态,避免像素重构,天然适合高效的多步推理。

现有方法局限

JEPA 训练仅优化短时域潜在预测,而规划需要多步想象并依据目标进展代价对候选轨迹排序。现有 JEPA 规划器通常直接继承嵌入空间的几何度量——例如潜在欧氏距离——作为代价函数。这种代价来自表示学习的副产品,而非从轨迹中显式挖掘的进展衡量,导致两个关键缺陷:

1. 训练-规划不对齐:表示学习的目标(预测下一时刻的潜在状态)与规划所需的最优代价(跨时域的进度判断)存在鸿沟;
2. 忽略时间结构与跨轨迹信息:欧氏距离无法捕捉有向时间顺序,且无法利用不同轨迹间的相对进度关系,在需要拓扑推理的任务中表现受限。

为什么这个问题难/重要

从无奖励的离线轨迹中挖掘有向的时间进展代价极具挑战:必须无监督地发现时间顺序,同时兼顾代价函数的可部署性与规划一致性。这一问题在工业界关注度日益升高——许多实际领域(如机器人操作、自动驾驶仿真)难以手工定义稠密奖励,而从纯观测日志中学习隐式进度度量可显著降低任务设计成本。此外,代价提取本身若能反向约束表示学习,可进一步缩小训练与规划的 gap,提升整体规划成功率。

行业类比

类似从自动驾驶车队的历史行驶数据中挖掘“到目标路口的隐含距离”作为规划代价,替代需要人工标注的稠密奖励或启发式度量,使路径生成更贴合真实驾驶逻辑。

核心洞察

• - **训练-规划成本的对齐挖掘**：现有 JEPA 规划器将欧氏距离作为默认的 rollout 成本，源于表示学习而非规划需求。TD-JEPA 直接从无奖励轨迹中自监督地挖掘有向的时间距离，使成本函数与多步规划目标一致，在拓扑任务上实现了近乎完美的成功率。这种表示学习与成本函数的共设计思想，解决了离线数据中长期规划与短期预测之间的目标错位。
• - **成本设计的双角色机制**：TD-JEPA 揭示了基于环境主导几何（接触 vs. 拓扑）的灵活应用策略：挖掘的时间成本既可直接用作规划成本（Two-Room 场景下成功率从 97.4% 提升至 100%），也可作为表示正则项强化欧氏距离规划（Push-T 中提升接触精度）。这为世界模型的规划成本选择提供了上下文感知的灵活性，避免了单一度量假设的局限。

方法

TD-JEPA 在 LeWM 的 encoder–predictor 主干上引入一个可学习的有向时间距离成本，让规划器直接捕捉轨迹中的进展结构，而不再依赖嵌入空间中不可控的欧氏距离。

输入与架构

• 输入：无奖励的离线演示轨迹，每个轨迹是状态序列。
• 编码器：将原始观测映射到潜在表示 (z)。
• 预测器：在潜在空间执行多步前向预测，生成未来状态表示 (\hat{z})，保留标准 JEPA 短期预测损失。
• 定向成本头：额外的前馈网络，输入两个潜在表示，输出它们之间的标量时间距离。

关键监督设计

1. 同轨迹正样本：从同一轨迹采样先后两个状态，强制成本头输出单调递减的时间距离，正向驱动目标进展。
2. 跨轨迹负样本：随机采样不同轨迹的状态对作为启发式负例，遏制成本头对无关状态的过度泛化。
3. Rollout 一致性：要求预测器在连续多步展开后，成本头仍保持与真实时间距离一致的排序，使学习到的进度度量与规划器的展开长度对齐。

三部分损失与预测损失联合优化，使成本头从数据中挖掘出反映可行进展的定向拓扑成本。

规划部署与双角色

• 直接部署成本：在拓扑主导的环境（如 Two-Room）中，用学习到的有向成本代替欧氏距离进行 CEM/iCEM 规划，成功率达到 100%。
• 表示学习正则：在接触几何主导的任务（如 Push-T）中，将时间距离作为辅助信号强化表示质量，使同一 checkpoint 下欧氏规划的性能提升 14.2 点（OGB-Cube）并改善 Push-T。

与同类工作的差异：TD-JEPA 首次显式挖掘轨迹内在的时间进展作为定向规划成本，并设计 rollout 一致性耦合训练和规划阶段，解决了以往 JEPA 规划器被动依赖欧氏距离、训练与规划脱节的问题。

实验

实验设计

实验聚焦于 离线演示日志（reward-free） 下的潜世界模型预测控制（latent MPC）。在 Two-Room（拓扑导航）、OGB-Cube（物体操作）和 Push-T（接触密集型推动）三个环境上，对 TD-JEPA 进行锁定评估（locked evaluation），即冻结学到的表示与规划代价，仅测试规划器。

主要对比基线：

• LeWM：传统 JEPA 规划器，采用隐空间欧氏距离作为规划代价
• RC-aux：同期工作，引入辅助表示学习

为解耦代价设计的影响，实验区分两种部署模式：

1. 直接使用挖掘出的 有向时序代价（temporal cost）
2. 在同一个基于时序训练的 checkpoint 上，仍使用 共享欧氏距离 规划

关键发现

• 拓扑主导环境（Two-Room）：部署挖掘的时序代价将成功率从 LeWM 的 97.4% 提升至 100.0%，显示时序结构对远离目标的规划至关重要。
• 接触主导环境（OGB-Cube、Push-T）：即使不直接使用时序代价，仅将时序信号作为表示学习的一部分，也能显著提升欧氏规划性能——OGB-Cube 指标提高 +14.2 点，Push-T 也有改善。
• 消融实验表明，有向头、跨轨迹负样本和 rollout 一致性损失三项设计各自均有增益。

与基线对比的深度解读

TD-JEPA 在全部三个环境上匹配或超越了 LeWM 与 RC-aux。其优势并非来自更强的模型容量，而是 挖掘离线数据中隐式的时序进度结构，并与规划代价协同设计。

当环境中目标距离与拓扑顺序强相关（如 Two-Room），直接使用时序代价可消除欧氏度量的短视偏差；当物理接触约束更为关键（如 Push-T），纯粹的时序代价反而不够，此时它主要充当表示正则化器，增强欧氏空间中的可规划性。这种根据任务特性选择代价部署方式的灵活性，是 TD-JEPA 相比固定代价的基线更鲁棒的根本原因。

行业影响

落地场景

TD-JEPA 为基于离线数据的世界模型规划提供了一种更精准的进度衡量方式，尤其适用于只能通过演示轨迹进行模仿学习、但环境奖励稀疏或无奖励的场景。典型落地包括：

• 仓储物流机器人：从人类操作员的历史拣选轨迹中学习，无需精细奖励设计，利用时间距离成本在潜在空间中规划出接近最优的导航与操作序列。
• 自动驾驶仿真与离线强化学习：利用录制的大量驾驶日志训练世界模型，在潜在空间内进行多步预测与轨迹排序，提升决策的安全性与合规性。
• 内容推荐与用户引导：将用户行为序列视为轨迹，挖掘时间进度特征，优化推荐 slot 排序策略。

商业价值

• 降低数据标注与奖励工程的成本：TD-JEPA 直接从无奖励轨迹中挖掘时间距离成本，省去了为每个任务手工设计密集奖励的繁重工作，尤其适合多任务、动态变化的环境。
• 提升离线学习到在线部署的决策质量：通过缩小训练时的短程预测与规划时的多步排序之间的差距，使模型在未见过的场景中表现得更为稳健，直接带来任务成功率与效率的跃升（在 Two-Room 任务中成功率已达 100%）。
• 加速产品迭代：基于离线日志即可持续改进规划策略，无需频繁上线试错，对风险敏感行业（如自动驾驶）价值显著。

与现有产品/工作流的接口

TD-JEPA 以 JEPA 编码器-预测器骨干 为基础，可无缝接入现有的潜在世界模型训练流程。集成要点：

1. 数据管道：输入为无奖励的轨迹序列（状态-动作对），无需改动日志采集方式。
2. 训练阶段：在原有 JEPA 损失上叠加定向时间距离头、跨轨迹负样本对比及 rollout 一致性正则项，训练开销可控。
3. 部署形式：训练后产出两个可直接使用的成本：拓扑进度占优时部署挖掘的时间成本，接触几何占优时用其作为表征信号增强欧氏规划，可自适应切换或融合。
4. 与 MPC 框架集成：可将该成本函数嵌入 CEM / iCEM 这样的通用模型预测控制求解器，实现即插即用的规划升级。

具体用例：某电商仓库的移动抓取机器人，利用历史操作日志训练 TD-JEPA 世界模型，在无需重新标注奖励的情况下，将货架间导航的成功率提高 2.6 个百分点（类比 Two-Room 提升），并缩短平均路径长度；同时，该模型提供的表征改善还使得 Push-T 这类精密操作任务的成功率与容错性大幅提升。

局限

• - **任务依赖性较强，缺乏统一部署策略**：论文指出纯时间成本在接触密集型任务（如 Push-T）上需要与欧氏几何结合，而拓扑主导任务可直接使用。实际应用中需要人工判断何时部署哪种成本形式，缺少自适应机制。若未来能自动检测任务特性并切换成本，将更具普适性。
• - **评估环境有限且离线数据假设较强**：实验仅在 Three-Room、OGB-Cube 和 Push-T 三个模拟环境上进行，且离线轨迹来自单一策略的演示。论文未探讨数据分布偏移、次优数据或噪声轨迹对时间监督挖掘的影响，实际部署时的鲁棒性和泛化能力尚待验证。
• - **对比基线范围较窄**：与 LeWM 和 RC-aux 的比较体现了在 JEPA 范式内的进步，但缺少与 Dreamer 系列等主流 latent MPC 方法、或基于 contrastive learning 的规划方法的直接对比，无法全面衡量 TD-JEPA 在更广泛模型基类中的竞争力。