World models improve planning by preserving commute-time distances
Learning Commute-Time-Preserving World Models for Planning
Machine Learning
Summary
Planning how to reach a goal can be improved if the agent’s internal map reflects how long it takes to move between locations. The authors found that existing methods that try to create this kind of map often lose important distance information. They propose a new approach called Commute-Time-Preserving World Models (CTWMs) that keeps these distances accurate, helping agents plan better in complex environments. Their tests show CTWMs can perform as well or better than existing methods while using fewer parameters.
What this means in practice
- •For robotics teams: Enhance robot navigation by creating internal maps that accurately reflect travel time between locations for improved goal-oriented planning.
- •For video game developers: Develop AI agents capable of better pathfinding in complex virtual worlds by leveraging commute-time-preserving latent models.
Authors
Michael Hauri, Peter Buttaroni, Fabian A. Mikulasch, Friedemann Zenke
Abstract
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 Commute-Time-Preserving World Models (CTWMs), combining a latent displacement predictor and a log-determinant regularizer that prevents collapse, which provably recover the correctly scaled Laplacian representation under reversible deterministic dynamics and at the predictor's fixed point. In numerical simulations, CTWM matches or outperforms LeWM, a task-agnostic baseline, on several complex, continuous goal-reaching benchmarks, while using half the parameters.