论文

SCALE:状态校准的潜在嵌入,用于正确几何形状的 JEPA 规划

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

模型训练预训练

摘要

联合嵌入预测世界模型通过使用表示本身定义的成本根据目标嵌入对预测的终端嵌入进行评分来进行规划。获得非折叠表示的两个重要策略是继承预训练的特征空间(如 DINO-WM),以及学习具有抗折叠正则化的端到端嵌入(如使用 SIGReg 的 LeWorldModel (LeWM))。这些策略显示了跨任务的互补优势。尽管任务相关状态可以从两个模型的完整嵌入中解码,但 DINO-WM 的主要主成分通常比 LeWM 保留更多的状态信息。由于欧几里得规划成本由高方差方向主导,因此这种差异会影响状态对候选选择的影响程度。我们提出 SCALE(状态校准潜在嵌入)来为端到端 LeWM 表示提供在 DINO-WM 中观察到的有利几何特性。 SCALE 通过将采样的成对潜在距离与标准化任务相关状态空间中的距离相关联来引入此属性,而无需替换 LeWM 的学习编码器。在五个任务、三个规划求解器和五个计算预算中,SCALE 比 LeWM 提高了每个任务求解器的平均值。潜在状态回归控制匹配或超过 SCALE 的完全嵌入可解码性,但使潜在状态距离对齐基本不变,并产生不太一致的规划增益。 SCALE 添加了一个轻量级训练时间正则化器,并且没有规划时间开销。这些结果表明,规划不仅取决于任务相关信息是否存在,还取决于它是否塑造了规划器所消耗的几何形状。