Search papers, labs, and topics across Lattice.
This paper introduces SCALE, a method that enhances the latent embeddings of LeWorldModel (LeWM) by aligning their geometric properties with those of DINO-WM, thereby improving planning performance across various tasks. By correlating pairwise latent distances with a standardized state space, SCALE retains the advantages of LeWM's learned encoder while addressing the high-variance issues that affect planning efficiency. The results demonstrate that SCALE consistently outperforms LeWM across five tasks and three planning solvers, highlighting the importance of geometric representation in effective planning.
Planning performance can be dramatically improved by aligning latent embedding geometry with task-relevant state representations, as shown by SCALE's consistent outperformance of LeWM.
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.
