PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Royson Lee, Fady Rezk, Titouan Parcollet, Timothy Hospedales, Cristina Cornelio

arXiv:2609.31868v2cs.AIcs.RO

TL;DR

JEPA-based goal-reaching planners are effective over short horizons, and hierarchical planners extend them by predicting intermediate latent sub-goals. METRO-WM retrieves sub-goals from recorded experience instead, and achieves up to 37.33 pp higher long-horizon success than the next-best hierarchical approach while running up to 10.9× faster.

  • Problem

    Flat JEPA-based latent planners are reliable only over short horizons, while hierarchical extensions predict intermediate latent sub-goals to address longer tasks.

    Source: §1 · pp. 1–2

  • Method

    METRO-WM replaces generated sub-goals with retrieved states from recorded episodes, using a graph of observed frames to find routes to goals.

    Source: §1 · p. 2

  • Results

    METRO-WM achieves up to 37.33 pp higher long-horizon success than the next-best hierarchical approach, runs up to 10.9× faster, and requires 13–56× less offline compute.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Retrieving sub-goals from recorded states guarantees their physical plausibility, and graph planning allows routes to be recomputed from the current state after execution errors.

    Source: §1 · p. 2, §7 · p. 9

    Because METRO-WM relies on bridge edges between visually similar states across episodes, the authors anticipate it may struggle in environments with highly diverse layouts.

    Source: §7 · p. 9

Abstract

from arXiv · show

Model-predictive control with Joint-Embedding Predictive Architectures (JEPAs) provides a strong zero-shot goal-reaching planner, but it is only effective over short planning horizons. Hierarchical extensions attempt to bridge this gap by learning a macro planner to predict intermediate latent sub-goals to guide the micro planner. In this work, we demonstrate that unconstrained latent sub-goal prediction is fundamentally flawed. A rigorous evaluation reveals that a leading state-of-the-art macro planner routinely emits physically unrealisable sub-goals. To resolve this, we introduce Metro-WM, a hierarchical framework that issues sub-goals by retrieving genuine states from prior experience rather than generating ungrounded latent vectors. Specifically, Metro-WM constructs a graph whose vertices are observed frames from offline expert demonstrations or random-action trajectories, allowing frames from different episodes to be connected and stitched into routes to the goal. Planning over the full graph also makes the system highly robust to execution errors: if the micro planner drifts off course, Metro-WM instantly finds a new optimal path from the current state. Our experiments show that Metro-WM achieves superior long-horizon success rates of up to 37.33 percentage points over the next best hierarchical approach while being up to 10.9 times faster, requiring both 13-56 times less offline compute and fewer tuned hyperparameters. Additional analysis reveals that Metro-WM finds 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.

1 Introduction

JEPA-based latent planners support zero-shot goal-reaching but are reliable only over short horizons, while hierarchical planners extend their reach by generating intermediate latent sub-goals. METRO-WM instead retrieves realisable sub-goals from recorded episodes and connects observed frames into routes to the goal.

1 Introduction

  • Flat JEPA planners are reliable only over short horizons, as prediction errors accumulate and latent distances become less informative over longer action sequences.Many tasks, including obstacle navigation and multi-step object manipulation, require longer sequences of actions.Source: §1 · pp. 1–2
  • Hierarchical JEPA planners generate intermediate latent sub-goals, but these may not correspond to physically realisable states.METRO-WM replaces generation with retrieval from recorded episodes, guaranteeing that each sub-goal encodes a realisable state.Source: §1 · p. 2
  • METRO-WM achieves up to 37.33 pp higher long-horizon success than the next-best hierarchical approach, runs up to 10.9× faster, and uses 13–56× less offline compute.It builds a graph from expert or random-action episodes and may stitch together multiple trajectories.Source: §1 · p. 2

2 Related work

Prior JEPA-based hierarchical planners generally generate latent sub-goals, while data retrieval has mainly been used in classical reinforcement learning and visual navigation. METRO-WM brings retrieval-based sub-goals to JEPA planning with a non-parametric, training-free macro planner.

2 Related work

  • Existing JEPA-based hierarchical approaches predominantly generate sub-goals, while retrieval from episodic memory or topological maps has mainly appeared in classical reinforcement learning and visual navigation.Examples cited include [Zhang et al., 2026, Caselli et al., 2026, Masip et al., 2026] for latent planning and [Eysenbach et al., 2019, Emmons et al., 2020] for reinforcement learning retrieval.Source: §2 · p. 3
  • METRO-WM adapts retrieval-based sub-goals to JEPA latent planning to guide a zero-shot latent micro planner while preventing sub-goal hallucination.Its frame-based edge weights are calibrated from empirical dataset quantiles, and its macro planner is non-parametric and training-free.Source: §2 · p. 3

3 Preliminaries

A flat latent planner repeatedly predicts action-conditioned latent outcomes and selects actions by their distance to the encoded goal. Hierarchical planners add a macro-level sub-goal, but unconstrained latent optimization can place it outside the set of encodings of physical states.

3 Preliminaries

  • A flat latent planner uses a frozen encoder and dynamics model to plan toward an encoded goal, with actions chunked into blocks of five environment frames.The planner replans at each step using a sampling optimizer.Source: §3 · p. 3
  • The flat micro planner is reliable only at short horizons because prediction errors accumulate and latent distances saturate, preventing reliable ranking of distant candidates.The paper identifies saturation as a bottleneck even with a perfect predictor.Source: §3 · p. 3
  • A hierarchical macro planner emits a sub-goal that replaces the final goal for the micro planner, but optimizing over the full latent space can produce a physically non-existent state.Realisable latents form a thin subset of the full latent space because the underlying physical state space is often low-dimensional.Source: §3 · p. 3

4 Are macro latent sub-goals realisable?

The paper evaluates sub-goal realisability without relying on a learned decoder, using the environment renderer and frozen encoder instead. Its evaluation finds that Hi-LeWM-C sub-goals on CUBE and PUSHT remain far from any nearby physical state under the reported search.

4 Are macro latent sub-goals realisable?

  • The decoder-free evaluation protocol uses the environment renderer and frozen encoder, avoiding the confound of decoder reconstruction ability and remaining independent of the provided demonstrations.The measure is the realisability residual.Source: §4 · p. 3
  • Figure 2 contrasts renderer-drawn nearest realisable states for Hi-LeWM-C sub-goals with METRO-WM sub-goals retrieved as real training-data states.The caption states that constrained micro actuation will not progress toward the final goal with Hi-LeWM-C sub-goals, but will with METRO-WM sub-goals.Source: Figure 2 · p. 4
  • Because CEM only upper-bounds the realisability residual, the authors tune its hyperparameters toward lower residuals and check that a larger search budget leaves the median unchanged.These checks are reported as safeguards against under-searching.Source: §4 · p. 4
  • Hi-LeWM-C sub-goals have median realisability residuals of 0.9 on CUBE and 2.3 on PUSHT, with residuals staying above zero as search budget grows.The authors state that the residual scale is not interpretable in isolation; a residual remaining above zero indicates no realisable state exists.Source: §4 · p. 4

5 METRO-WM: sub-goal selection by shortest paths

METRO-WM replaces predicted latent sub-goals with a shortest-path search over recorded experience, selecting observed states as sub-goals. Its graph connects feasible within-episode transitions and cross-episode bridges, with costs calibrated in frames.

5 METRO-WM: sub-goal selection by shortest paths

  • METRO-WM replaces latent sub-goal prediction with shortest-path search over prior experience to select physically plausible states that are useful and on a route to the goal.The search selects recorded frames as sub-goals for the micro planner.Source: §5 · p. 4

5.1 The graph of recorded frames

  • The graph uses encoded recorded frames as vertices, directed unit-cost temporal edges for demonstrated transitions, and symmetric cross-episode bridges between similar frames.Bridges allow a route to leave one demonstration and continue along another.Source: §5.1 · p. 4

5.2 Bridge edge weights

  • Bridge costs convert latent distance into frame costs using empirical distance quantiles by frame gap, with conservative pricing intended to avoid un-executable shortcuts.The quantile curve is kept below saturation within the waypoint budget, and its running maximum prevents larger latent gaps from being charged fewer frames.Source: §5.2 · p. 5

5.3 Graph search and sub-goal selection

  • For each query, backward Dijkstra search computes cost-to-go and optimal successors across the graph, while a one-step Bellman backup connects the changing current state.The planner selects the farthest vertex along the optimal route within waypoint budget H; if execution drifts, the current state can query an optimal continuation without recalculating a path, and disconnected cases fall back toward the closest reachable vertex to the goal.Source: §5.3 · p. 5

5.4 Summary: end-to-end pipeline

  • The pipeline encodes the corpus and builds the weighted graph offline, computes query-specific graph costs once per episode, then re-enters the graph after each macro step.At each macro step it selects a sub-goal within H and hands it to the micro planner.Source: §5.4 · p. 6

6 Evaluation

Evaluation across four goal-conditioned benchmarks compares frozen LeWM and PLDM micro planners with hierarchical baselines over goal distances and solver budgets. METRO-WM’s long-horizon results include stronger, faster performance and lower offline cost, alongside analyses of re-entry, route length, and dataset density.

6.1 Experimental Setup

  • The setup uses frozen LeWM and PLDM micro planners across four benchmarks, evaluates 50 query episodes per benchmark with three planner seeds, and sweeps budgets and planning horizons.The benchmarks are PUSHT, OGBench-CUBE, TWOROOM, and REACHER.Source: §6.1 · p. 6

6.2 Main Result

  • Both METRO-WM variants beat flat planning with the same world model, and METRO-LeWM’s d=75 gain over LeWM is +28.67–46.67, while gains at d=25 are small except on CUBE.Source: §6.2 · pp. 6–7
  • At d=75, METRO-LeWM beat Hi-LeWM-C by +16.67, +18.00, +37.33, and +4.00 points across PUSHT, CUBE, TWOROOM, and REACHER, ran 2.2–10.9× faster, and reached 100% on TWOROOM and REACHER.Source: §6.2 · pp. 6–7, Figure 3 · p. 7
  • METRO-WM’s graph construction costs 9–39 minutes on one GPU with no gradient steps, while macro-planner training costs 13–56× as much GPU time.The macro-planner training costs serve all three hierarchical baselines per benchmark.Source: §6.2 · p. 7, Table 1 · p. 8

6.3 Analysis

  • Fine-tuning a stronger micro planner and plugging it into METRO-WM leads to a performance boost over using the provided frozen planners.Source: §6.3 · pp. 7–8
  • At d=75, committed routes reduce performance relative to re-entry, except in REACHER, where the difference is negligible and goals are often reachable within two macro plans.Re-entry lets METRO-WM find a new optimal remaining route from the micro planner’s current state.Source: §6.3 · p. 8, Table 2 · p. 8
  • In CUBE, REACHER, and TWOROOM, METRO-WM reaches goals sooner than the query’s own demonstration by stitching multiple demonstrations into shorter paths.In PUSHT, both METRO-WM and the oracle take longer paths than the demonstration.Source: §6.3 · p. 8, Table 3 · p. 9
  • With 1% of recorded episodes, METRO-LeWM remains similar to the next-best hierarchical baseline on PUSHT and CUBE, while 100 episodes suffice to solve TWOROOM and REACHER every time.Success on PUSHT and CUBE drops as the number of demonstrations decreases.Source: §6.3 · p. 8, Table 4 · p. 9

7 Conclusion

METRO-WM addresses unrealisable latent sub-goals with shortest-path planning over recorded frames, while dynamic graph re-entry mitigates execution errors. Its entry backup computes an exact shortest-path cost for the graph.

7 Conclusion

  • METRO-WM uses frame-priced shortest paths over a dataset graph to guarantee plausible sub-goals and mitigate micro-planner execution errors through dynamic re-entry.It often finds paths shorter than the offline demonstrations, but may struggle to stitch trajectories across highly diverse layouts.Source: §7 · p. 9
  • The conclusion anticipates difficulty in environments with highly diverse layouts because bridge edges connect visually similar states across episodes.The paper identifies accommodating topological diversity as future work.Source: §7 · p. 9

A Graph entry

  • Adding the current state with outgoing edges to candidate neighbors lets Eq. 7 compute its exact shortest-path cost without changing existing graph costs.The backup takes the minimum over candidate neighbors’ entry-edge costs plus their existing cost-to-go.Source: A Graph entry · p. 12
  • The entry backup introduces no approximation error for any candidate set; the neighbor count k controls how densely the current state connects to the graph.The guarantee applies to the graph on which the backup is run.Source: A Graph entry · p. 12

B Realisability search of macro latent sub-goals

The realisability search tests whether generated latent sub-goals correspond to plausible simulator states, while METRO-WM’s algorithm plans over recorded frames and replans at each macro step.

B Realisability search of macro latent sub-goals

  • Fourfold more compute barely reduces the median realisability residual: CUBE changes from 0.92 to 0.89, and PUSHT from 2.31 to 2.29.The residual plateaus above zero, which the authors interpret as failure to find a plausible solution.Source: B Realisability search of macro latent… · p. 12
  • Figure 5 contrasts Hi-LeWM-C’s closest realisable macro sub-goals with METRO-WM’s sub-goals, which are retrieved real states from other demonstrations.The figure renders the sub-goals using the environment’s renderer.Source: Figure 5 · p. 13

C METRO-WM’s Algorithm

  • The algorithm preprocesses the corpus and builds its graph offline, then runs goal-directed Dijkstra search and selects a route waypoint at each macro replanning step.It re-encodes the state actually reached before replanning and invokes the micro planner with the selected sub-goal.Source: C METRO-WM’s Algorithm · p. 13

D Latent metric saturation

Latent distance distinguishes frames only across relatively small temporal gaps; at larger gaps, it saturates near the distance between frames from different demonstrations. The planner therefore prices graph edges in frames.

D Latent metric saturation

  • For PUSHT, the normalized latent gap rises from 0.29 to 0.52 between 5 and 10 frames, but only from 0.98 to 1.0 between 50 and 75 frames.At a curve value of 1, latent distance cannot distinguish a frame that far away from a random frame in the corpus.Source: D Latent metric saturation · p. 13, Figure 6 · p. 14
  • The saturation point depends on the benchmark and encoder; beyond it, latent cost cannot indicate which frame is closer in time.The authors distinguish this representation property from rollout behavior.Source: D Latent metric saturation · p. 13
  • Because latent distance saturates and prediction errors compound, the authors explain that the micro planner fails over longer horizons and price graph edges in frames.The figure plots median within-demonstration latent distance divided by the distance to a random frame from another demonstration, with interquartile-range bands.Source: D Latent metric saturation · p. 14, Figure 6 · p. 14

E Additional Results

Figure 7 depicts routes that traverse recorded frames, sometimes switching between demonstrations through a bridge between their endpoints.

E.1 What a bridge looks like

  • Each row shows part of one METRO-LeWM route for an evaluation query, with the query’s observed and goal states at the route’s ends.Ellipses mark stretches of the route omitted from the figure.Source: Figure 7 · p. 14
  • The depicted route follows one demonstration, crosses a bridge, and continues along another demonstration.The heavily outlined bridge endpoints come from different demonstrations.Source: Figure 7 · p. 14
  • Compare the outlined bridge endpoints to see where the route connects frames from separate demonstrations.The remaining route panels are recorded frames traversed between the query’s observed and goal states.Source: Figure 7 · p. 14

E.2 A stronger micro planner

Fine-tuning improves short-horizon LeWM, but gains at long horizons emerge when paired with macro sub-goal selection. The related work distinguishes METRO-WM’s retrieval of recorded states from generated sub-goals, and the experimental materials document reproducibility resources.

E.2 A stronger micro planner

  • Plugging fine-tuned LeWM into METRO-WM and ORACLE yields substantial gains, while flat planning at d = 75 shows negligible performance differences despite better d = 25 success.LeWM (ft) improves short-horizon success from 93.33 ± 5.03 to 97.33 ± 1.15 at d = 25.Source: §E.2 · pp. 14–15, Table 5 · p. 15

F Extended related work

  • Graph retrieval with recorded-state nodes guarantees realizable sub-goals; METRO-WM brings this approach from navigation and goal-conditioned RL into latent world-model planning.Prior graph-based retrieval includes visual navigation [Savinov et al., 2018, Shah et al., 2021] and goal-conditioned RL [Eysenbach et al., 2019, Emmons et al., 2020].Source: F Extended related work · p. 15
  • Latent macro planners can emit sub-goals that encode no real state; restricting search to demonstration macro-actions still does not guarantee realizability, unlike METRO-WM’s recorded-frame retrieval.Masip et al. [2026] report decoded sub-goals that omit the actor, while Caselli et al. [2026] find that restricting search does not restrict its output.Source: F Extended related work · p. 15
  • Pixel-space planners generate visual sub-goals, but visual realism does not guarantee physical plausibility; METRO-WM instead retrieves recorded frames.Examples include time-agnostic prediction [Jayaraman et al., 2019], HVF [Nair & Finn, 2020], and diffusion-based SuSIE [Black et al., 2024].Source: F Extended related work · p. 15

G Detailed experimental setup

  • The authors provide experiment details and open-source code, and use benchmarks and frozen micro planners downloaded from the LeWM repository.The code is available at https://github.com/SamsungLabs/metro-wm; benchmarks and frozen planners come from https://github.com/lucas-maes/le-wm.Source: G Detailed experimental setup · p. 16

G.1 Benchmarks

The experimental setup covers four benchmarks, two frozen micro planners, three hierarchical baselines, and METRO-WM’s parameter sweeps. It also describes the evaluation resources and training procedures used for the comparisons.

G.1 Benchmarks

  • The four benchmarks are PUSHT, CUBE, TWOROOM, and REACHER, spanning pushing, robotic manipulation, doorway navigation, and joint-space reaching.REACHER’s released corpus comes from a random policy, so no recorded episode reaches the goal on its own.Source: §G.1 · p. 16, Table 6 · p. 16

G.2 Micro planners

  • LeWM and PLDM serve as frozen micro planners on all benchmarks, using released weights and shared architecture and CEM settings.The encoder produces a 192-dimensional frame state; CEM uses 300 samples and top-k 30, with iteration counts differing for PUSHT and the other benchmarks.Source: §G.2 · p. 16

G.3 Hierarchical baselines

  • The hierarchical baselines share a frozen LeWM micro planner and trained macro predictor, but differ in macro-action proposal and replanning; the experiments use released or reconstructed checkpoints.The baselines include Hi-LeWM and Hi-LeWM-C variants; training details and baseline CEM settings are specified, with the same released matrix rows applied to TWOROOM and REACHER.Source: §G.3 · pp. 16–17

G.4 METRO-WM’s hyperparameters

  • METRO-WM tunes only waypoint budget H and micro plan length h, fixes graph degree k = 4, and evaluates a sweep of (h, H) settings across benchmarks.The sweep took under a day on 8 NVIDIA A100 80GB GPUs for each micro planner; frontier configurations are provided for reproducibility.Source: §G.4 · pp. 17–18

Related papers

Loading 2609.31868v2…