Learning Commute-Time-Preserving World Models for Planning

World models allow agents to plan in latent space by choosing a sequence of actions that most reduces the distance to a given goal state. Thus, planning can benefit from latent representations whose distances mirror commute-times in the environment. The spectral embedding space of the graph Laplacian provides such a representation, if it obeys a specific eigenvalue-dependent scaling. Unfortunately, instantiating the graph Laplacian is intractable in large, continuous environments. Self-supervised learning offers a natural route to such commute-time-preserving embeddings at scale. However, here we show that existing methods, which commonly encourage isotropic representations to prevent representational collapse, tend to degrade the “correct” eigenvalue-dependent scaling, leading to an inaccurate representation of commute times. To address this problem, we introduce a Commute-Time-Preserving World Model (CTWM), combining a latent displacement predictor and a log-determinant regularizer that prevents collapse, which provably recovers the correctly scaled Laplacian representation under deterministic dynamics and at the predictor’s fixed point. In numerical simulations, CTWM matches or outperforms LeWM, a strong task-agnostic baseline, on several complex, continuous goal-reaching benchmarks, while using only half the parameters.