Metro-WM: Long-Horizon Latent Planning with Realisable Sub-Goals

Model-predictive control with Joint-Embedding Predictive Architectures (JEPAs) is a strong zero-shot goal-reaching planner but only works well over short planning horizons. Hierarchical extensions try to bridge this by learning a macro planner that predicts intermediate latent sub-goals to guide a micro planner. The paper argues this approach is fundamentally flawed: a rigorous evaluation shows a leading state-of-the-art macro planner routinely emits physically unrealisable sub-goals.

Metro-WM instead issues sub-goals by retrieving genuine states from prior experience rather than generating ungrounded latent vectors. It builds a graph whose vertices are observed frames from offline expert demonstrations or random-action trajectories. Frames from different episodes can be connected and stitched into routes to the goal. Because planning is done over the full graph, the system is robust to execution errors: if the micro planner drifts off course, Metro-WM instantly finds a new optimal path from the current state.

Experiments show Metro-WM achieves long-horizon success rates up to 37.33 percentage points higher than the next best hierarchical approach while being up to 10.9x faster and requiring 13–56x less offline compute and fewer tuned hyperparameters. Additional analysis finds it produces shorter paths than the offline demonstrations, outperforms an oracle relying on the query's own demonstration, and maintains robust performance under extremely sparse dataset conditions.