Latent Geometry Beyond Search: Amortizing Planning in World Models

TL;DR

This paper introduces a method to amortize planning in world models by learning a latent inverse-dynamics model, significantly reducing computation while maintaining or improving performance across various environments.

Contribution

It demonstrates that structured latent spaces enable replacing iterative planning with a learned inverse-dynamics model, simplifying control in vision-based world models.

Findings

The proposed GC-IDM matches or exceeds CEM performance in most benchmarks.

Per-decision cost is reduced by 100-130x with the new method.

The approach is robust across different test-time planners.

Abstract

Modern vision-based world models can represent observations as compact yet expressive latent manifolds, but fast goal-oriented planning in these spaces remains challenging. This raises a central question: when does a learned representation simplify control, rather than merely enabling prediction? We study this question in a pretrained LeWorldModel, whose latent geometry is regularized for smoothness and uniformity. Our key insight is that, under such geometry, planning can be amortized into a latent inverse-dynamics mapping instead of requiring online search. We therefore replace iterative planning with a lightweight Goal-Conditioned Inverse Dynamics Model (GC-IDM) that maps the current latent state, goal latent state, and remaining horizon directly to the next action. Empirically, across four benchmark environments spanning navigation, contact-rich manipulation, and continuous control, our controller matches or exceeds CEM in seven of eight environment-protocol settings while reducing per-decision cost by 100-130x. A broader sweep over test-time planners (CEM, MPPI, iCEM, and gradient-based methods) shows that this result is not specific to a particular optimizer. These findings suggest that much of the structure recovered by test-time planning is already locally encoded in the latent representation. More broadly, our results indicate that sufficiently structured latent spaces can shift part of the planning burden from online optimization to learned inference.