SCALE: State-Calibrated Latent Embeddings for JEPA Planning in the Right Geometry

📄 arXiv: 2608.16287v1 📥 PDF

作者: Jiaming Hu, Yan Zheng, Tian Wang

分类: cs.LG

发布日期: 2026-08-17

备注: 15 pages, 2 figures

💡 一句话要点

提出SCALE以提升JEPA规划中的状态校准潜在嵌入

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 状态校准 潜在嵌入 联合嵌入 规划模型 几何特性 反塌缩正则化 机器人规划 智能控制

📋 核心要点

  1. 现有的DINO-WM和LeWM方法在获取非塌缩表示时各有优缺点，导致状态信息的保留和利用不均衡。
  2. SCALE通过将潜在距离与任务相关状态空间的距离进行关联，增强了LeWM的几何特性，而不替换其学习的编码器。
  3. 在五个任务的实验中，SCALE在每个任务的求解器平均表现上均优于LeWM，展示了其有效性。

📝 摘要（中文）

联合嵌入预测世界模型通过对比预测的终端嵌入与目标嵌入来进行规划，现有方法在获取非塌缩表示方面存在不足。本文提出SCALE（状态校准潜在嵌入），旨在将LeWM的端到端表示与DINO-WM的几何特性相结合。SCALE通过将采样的成对潜在距离与标准化的任务相关状态空间中的距离相关联，来实现这一目标。实验结果表明，SCALE在五个任务、三个规划求解器和五个计算预算下均优于LeWM，且在规划时间没有额外开销的情况下，提升了任务求解的平均表现。

🔬 方法详解

问题定义：本文旨在解决现有联合嵌入预测模型在状态信息保留和几何特性方面的不足，特别是DINO-WM与LeWM之间的差异。

核心思路：SCALE通过将潜在嵌入的成对距离与标准化的任务相关状态空间中的距离进行关联，来增强LeWM的几何特性，从而改善规划性能。

技术框架：SCALE的整体架构包括一个轻量级的训练时间正则化器，保持LeWM的编码器不变，同时在训练过程中引入潜在距离的校准。

关键创新：SCALE的主要创新在于通过状态校准潜在嵌入来改善几何特性，这一方法与现有的通过预训练特征空间或反塌缩正则化的方式本质上不同。

关键设计：在设计中，SCALE引入了一个正则化损失函数，以确保潜在距离与状态空间的距离一致，同时保持了LeWM的原有结构和性能。该方法在计算开销上也保持了轻量级。

🖼️ 关键图片

fig_0
fig_1

📊 实验亮点

实验结果显示，SCALE在五个不同任务中均优于LeWM，提升幅度显著，尤其是在高方差方向的规划成本上表现更佳。具体而言，SCALE在每个任务的求解器平均表现上均有提升，展示了其在实际应用中的有效性。

🎯 应用场景

SCALE的研究成果在机器人规划、自动驾驶和智能控制等领域具有广泛的应用潜力。通过提升模型的状态信息利用效率，SCALE能够帮助实现更高效的决策制定和路径规划，进而推动智能系统的自主性和智能化水平。

📄 摘要（原文）

Joint-embedding predictive world models plan by scoring predicted terminal embeddings against a goal embedding using a cost defined on the representation itself. Two prominent strategies for obtaining non-collapsed representations are to inherit a pretrained feature space, as in DINO-WM, and to learn an embedding end to end with anti-collapse regularization, as in LeWorldModel (LeWM) with SIGReg. These strategies show complementary strengths across tasks. Although task-relevant state is decodable from the full embeddings of both models, DINO-WM's leading principal components usually retain substantially more state information than LeWM's. Because Euclidean planning costs are dominated by high-variance directions, this difference affects how strongly state can influence candidate selection. We propose SCALE (State-CAlibrated Latent Embeddings) to give the end-to-end LeWM representation the favorable geometric property observed in DINO-WM. SCALE induces this property by correlating sampled pairwise latent distances with distances in a standardized task-relevant state space, without replacing LeWM's learned encoder. Across five tasks, three planning solvers, and five compute budgets, SCALE improves every task--solver average over LeWM. A latent-to-state regression control matches or exceeds SCALE's full-embedding decodability yet leaves latent--state distance alignment essentially unchanged and yields less consistent planning gains. SCALE adds a single lightweight training-time regularizer and no planning-time overhead. These results show that planning depends not only on whether task-relevant information is present, but also on whether it shapes the geometry consumed by the planner.