Latent World Models Should Plan with Distances Along the Latent Manifold

Latent world models such as LeWM plan by searching for the action sequence whose predicted terminal latent is closest to the goal latent under a distance measured in the latent space, such as the Euclidean distance. Planning with such distances assumes that proximity in the latent space measures how far the state is from the goal. For LeWM, we find that this assumption holds only in a small neighborhood near the goal, while the cost forms spurious low-cost basins that trap the planner. Yet neighborhood-based embeddings show that the latents retain the physical state. We hypothesize that the latents lie near a curved manifold along which proximity differs from proximity in the ambient space. If so, the terminal cost should measure the distance along this latent manifold. We propose the Neighborhood Graph Cost (NGC) to estimate this distance by the shortest path on a mutual nearest-neighbor graph over the latents of training observations. NGC changes only the terminal cost and requires no retraining. Experimental results show that NGC improves the success rate of LeWM on all four tasks, with the gain growing as the goal moves farther from the start. The results support the manifold hypothesis and our argument that latent world models should plan with distances along the latent manifold.