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.
Figures & tables
Figure 1: One cycle of Metro -WM. Rather than generating sub-goals, Metro -WM retrieves them from previously seen episodes. The macro planner searches the graph whose vertices represent observed frames encoded as latent states, with e1…e5 in the figure indicating the specific episodes containing these frames. Directed temporal edges (solid grey) connect consecutive frames within each episode, while symmetric bridge edges (dashed grey) link visually similar frames across episodes (example endpoints in Appendix Figure 7 ). All edge costs are expressed in frame units: Temporal edges have unit cost, while bridge-edge costs are converted into frame-equivalent costs derived from empirical dataset quantiles. Given a goal zgoal , we use Dijkstra’s algorithm to search backward from zgoal to find the lowest-cost route from every graph vertex. At each planning cycle, Metro -WM add the current observation zobst to the graph, and identifies the shortest route (blue) to zgoal (which may stitch together multiple episodes). The furthest frame along this route that is within the micro planning horizon is then selected as the next sub-goal zsubt . After execution (solid orange), planning resumes from the reached state zobst+1 (see Appendix Figure 4 ).
Figure 2: Decoder-free subgoal analysis. Top row: closest realizable states to the non-realisable sub-goals emitted by Hi-LeWM-C macro planner (Eq. 3 ), drawn by the environment’s own renderer. Bottom row: Our Metro -WM sub-goals, which are real states retrieved from the training dataset. Micro actuation - constrained to realisable states - will not progress toward the final goal when pursuing sub-goals issued by Hi-LeWM-C, but will progress when using Metro -WM sub-goals.
Figure 3: Performance and cost comparison between Metro -WM and baselines across varying budgets. Four benchmarks × three goal distances (planning horizons) d .
Metro -WM preprocessing (1 GPU)
Hi-LeWM-C training (4 GPUs)
benchmark
encode
index
total
grad. steps
wall
GPU-hours
PushT
17.1
6.0
23.1 min
076,215
4.01 h
16.0
Cube
34.2
4.7
38.9 min
118,125
9.09 h
36.4
TwoRoom
07.6
1.4
09.0 min
008,370
0.50 h
02.0
Reacher
24.4
4.7
29.1 min
118,125
6.58 h
26.3
Table 1: One-time offline cost per benchmark. Metro -WM’s “encode” is one forward pass of the frozen encoder over every recorded frame; “index” is the exact nearest-neighbour table plus the edge-cost computation. Hi-LeWM-C’s training figures are wall clock at the original authors’ recorded recipe; GPU-hours are wall clock × devices.
Metro -LeWM
Metro -PLDM
benchmark
re-entry
w/o re-entry
re-entry
w/o re-entry
PushT
62.67±4.62
49.33±3.06
44.00±5.29
39.33±6.43
Cube
92.67±3.06
84.00±0.00
89.33±5.03
88.67±2.31
TwoRoom
100.00±0.00
76.67±4.16
100.00±0.00
61.33±5.77
Reacher
100.00±0.00
99.33±1.15
100.00±0.00
98.00±0.00
Table 2: Success (%) with and without graph re-entry, d=75 .
Success (%)
Arrival Ratio
Approach
Benchmark
Metro -WM
oracle
Metro -WM
oracle
Metro -LeWM
PushT
62.67±4.62
64.00±2.00
2.71
2.89
Cube
92.67±3.06
72.00±4.00
0.44
1.32
TwoRoom
100.00±0.00
96.67±2.31
0.93
1.65
Reacher
100.00±0.00
93.33±3.06
0.65
2.38
Metro -PLDM
PushT
44.00±5.29
72.00±4.00
3.08
3.26
Table 3: Comparison with a privileged oracle , where the sub-goals come from the query’s episode own trajectory at inference. Arrival ratio is the median, over the successful episodes, of frames-to-success divided by the expert demonstration’s own first-hitting time; <1 means arriving sooner than the demonstration .
Approach
Corpus Size
Cube
PushT
Reacher
TwoRoom
Metro -LeWM
100%
92.67±3.06
62.67±4.62
100.00±0.00
100.00±0.00
10%
90.00±5.29
46.00±2.00
98.67±1.15
100.00±0.00
3%
79.33±1.15
38.00±2.00
100.00±0.00
100.00±0.00
1%
72.00±2.00
31.33±4.16
100.00±0.00
100.00±0.00
Metro -PLDM
100%
89.33±5.03
44.00±5.29
100.00±0.00
100.00±0.00
10%
90.67±1.15
29.33±5.03
100.00±0.00
100.00±0.00
Table 4: Success (%) against corpus size, d=75 . 1% is 187 episodes on PushT and 100 elsewhere.
Figure 4: Metro -WM’s graph re-entry. This follows Fig. 1 : after the micro planner stops at zobst+1 , the next macro step re-enters the graph there following the optimal shortest route from zobst+1 to zgoal .
Appendix figures & tables6 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 5: Top row: closest realisable Hi-LeWM-C’s sub-goals from its macro planner found by Eq. 3 and drawn by the environment’s own renderer. Bottom row: our method’s, METRO-WM, sub-goals which are retrieved real states of other demonstrations in the dataset.
Figure 6: Latent distance stops resolving time within a few dozen frames. Median ∥zt−zt−d∥ between two frames of the same demonstration d frames apart (band: interquartile range), divided by the median distance between frames of two different , randomly drawn demonstrations. A curve at 1 cannot tell a state d frames away from a random frame of the corpus.
Figure 7: A bridge, in pixels, on each benchmark. Each row is part of one route computed by Metro -LeWM for one query of the evaluation set. The ends are the query’s zobs and zgoal and every panel between them is a recorded frame the route traverses, and ⋯ marks the stretches not shown. The route runs along one demonstration (blue), crosses a bridge, and continues along a different one (red). The two heavily outlined frames are that bridge’s endpoints, which comes from different demonstrations in the dataset.
Approach
Success (%)
Δ
LeWM (ft)
15.33±3.06
+1.33
Metro -LeWM
72.00±3.58
+9.33
Oracle
81.33±3.06
+17.33
Appendix
Table 5: We fine-tune PushT LeWM, dubbed LeWM (ft) to get a stronger d=25 model and then evaluate it at d=75 . Δ is the gain over performance of the released LeWM weights, on the same 50 episodes.
benchmark
no. of episodes/demonstrations
no. of frames
PushT
18,685
2,261,996
Cube
10,000
1,970,000
TwoRoom
10,000
880,809
Reacher
10,000
1,970,000
Appendix
Table 6: Total number of demonstrations and frames for each benchmark in our experiments
d
macro: samples × iterations, top- k
micro: samples × iterations, top- k
25
900×20 , 10
300×30 , 150
50
1500×40 , 10
900×30 , 150
75
1200×60 , 10
1200×30 , 150
Appendix
Table 7: CEM hyperparameters of the hierarchical baselines, variance scale 1.0 at both levels. They come from the authors’ released matrix, which covers PushT and Cube ; we apply the same rows on TwoRoom and Reacher .