JEPA 世界模型规划中的表示几何问题：SCALE 方法

Authors:Jiaming Hu, Yan Zheng, Shi Bo, Tian Wang, Florian Dubost, Alejandro Mottini, Junze Liu, Arvind Srinivasan, Kai Zhong, Kun Qian, Sharon Gao, Qingjun Cui

View PDF HTML (experimental)

Abstract:Joint-embedding predictive world models support planning through latent predictions, but unconstrained joint training can collapse distinct observations to identical embeddings. Two prominent strategies for avoiding collapse are to inherit pretrained features, as in DINO-WM, or to learn representations end-to-end with anti-collapse regularization, as in LeWorldModel (LeWM). Yet avoiding collapse does not ensure that latent distances distinguish outcomes in ways that matter for the task. In object manipulation, for example, success depends on the object's position and orientation relative to the goal. Such task-relevant state information can remain accurately decodable while barely influencing latent distance. The resulting planning cost may fail to reflect how close a predicted outcome is to the task goal. In this paper, we propose SCALE (State-CAlibrated Latent Embeddings), a method that correlates sampled pairwise latent distances with distances in task-relevant state space. Added to LeWM's existing objective, SCALE preserves its architecture, requires privileged state only during training, and adds no planning-time computation. We show that SCALE improves planning success over LeWM across manipulation and navigation tasks with multiple solvers and provide a comprehensive analysis of how SCALE reshapes representation geometry to support planning.

Comments:	19 pages, 3 figures
Subjects:	Machine Learning (cs.LG)
Cite as:	arXiv:2608.16287 [cs.LG]
	(or arXiv:2608.16287v3 [cs.LG] for this version)
	https://doi.org/10.48550/arXiv.2608.16287 arXiv-issued DOI via DataCite

Submission history

From: Jiaming Hu [view email]
[v1] Mon, 17 Aug 2026 08:58:55 UTC (470 KB)
[v2] Wed, 26 Aug 2026 17:10:59 UTC (470 KB)
[v3] Tue, 6 Oct 2026 05:01:03 UTC (1,657 KB)