Learning Goal-Conditioned Subgoals through Frozen Hierarchical World Models
Organizations: Department of Industrial and Systems Engineering, Hosei University
Abstract
Hierarchical world models plan over long horizons by passing predicted states from a coarse model to a lower-level controller. Choosing these states can be expensive: the upper-level planner repeatedly differentiates through candidate macro-action trajectories. We replace this search with a lightweight goal-conditioned network trained through the frozen high-level predictor. The network supplies a conditioning vector, while the predictor generates the next latent waypoint. Training uses hindsight goals and state-prediction losses on offline trajectories, without search-solution targets or macro-action labels. Inference retains an autoregressive high-level rollout, but requires no high-level candidate optimization. We evaluate this approach on the shared-latent HWM checkpoints distributed with H-JEPA, using three backbone seeds and three head initializations per backbone. With a common evaluation protocol, mean success is 70.4% on Cube, 48.7% on Push-T, 100.0% on FourRoom, and 71.1% on Visual AntMaze. The method improves over the corresponding HWM search baseline on three environments and has the same mean on Push-T. Dedicated measurements on an NVIDIA L40S give 1.28-4.39x end-to-end speedups for batched evaluation. Ablations identify state-adaptive conditioning as a consistent contributor, while predictor-mediated subgoals and unrestricted conditions have environment-dependent benefits.
Figures & tables
| Environment | ||||
|---|---|---|---|---|
| Cube | 100 | 90 | 2 | 9,000 |
| Push-T | 400 | 90 | 8 | 36,000 |
| FourRoom | 800 | 90 | 12 | 72,000 |
| AntMaze | 200 | 90 | 9 | 18,000 |
| Environment | HWM search | H-JEPA search | Flat | Linear subgoals | Learned conditions |
|---|---|---|---|---|---|
| Cube | |||||
| Push-T | |||||
| FourRoom | |||||
| AntMaze |
| Variant | Cube | Push-T | FourRoom | AntMaze |
|---|---|---|---|---|
| Full objective | ||||
| Direct subgoals | ||||
| No head state input | ||||
| 4D conditions via | – |
Appendix figures & tables9 assets
Supplementary material from the paper’s appendix.
Appendix
| Environment | Backbone | Head 0 | Head 1 | Head 2 | Mean |
|---|---|---|---|---|---|
| Cube | 42 | 64 | 72 | 70 | 68.7 |
| Cube | 43 | 64 | 70 | 68 | 67.3 |
| Cube | 44 | 76 | 76 | 74 | 75.3 |
| Push-T | 42 | 42 | 44 | 46 | 44.0 |
| Push-T | 43 | 50 | 52 | 60 | 54.0 |
| Push-T | 44 | 42 | 52 | 50 | 48.0 |
| Environment | Planner | 42 | 43 | 44 | Mean SE |
|---|---|---|---|---|---|
| Cube | HWM search | 48 | 44 | 62 | |
| Cube | H-JEPA search | 60 | 40 | 40 | |
| Cube | Flat | 24 | 26 | 28 | |
| Cube | Linear | 4 | 2 | 8 | |
| Push-T | HWM search | 50 | 44 | 52 | |
| Push-T | H-JEPA search | 32 | 52 | 38 |
| Environment | States | params (M) | ||
|---|---|---|---|---|
| Cube | 192 | 2 | 3 | 11.584128 |
| Push-T | 192 | 8 | 9 | 11.584128 |
| FourRoom | 192 | 12 | 13 | 4.837504 |
| AntMaze | 384 | 9 | 10 | 10.247552 |
| Environment | Heads | Median (min) | Range (min) |
|---|---|---|---|
| Cube | 9 | 6.9 | 6.9–8.3 |
| Push-T | 9 | 32.2 | 28.6–43.3 |
| FourRoom | 9 | 24.0 | 23.6–30.4 |
| AntMaze | 9 | 26.0 | 25.7–31.4 |
| Environment | Method | High (s) | Low (s) | Total (s) | Speedup |
|---|---|---|---|---|---|
| Cube | HWM search | 26.189 | 100.80 | 148.44 | – |
| Cube | Learned conditions | 0.068 | 93.65 | 115.66 | |
| Push-T | HWM search | 505.278 | 237.39 | 757.51 | – |
| Push-T | Learned conditions | 0.201 | 207.31 | 221.56 | |
| FourRoom | HWM search | 1963.936 | 768.98 | 2784.24 | – |
| FourRoom | Learned conditions | 0.661 | 586.35 | 634.73 |
| Variant | Head 0 | Head 1 | Head 2 | Mean SD |
|---|---|---|---|---|
| Full objective | 64 | 72 | 70 | |
| Without teacher forcing | 68 | 66 | 68 | |
| Without rollout loss | 74 | 68 | 72 | |
| Without endpoint loss | 70 | 72 | 68 | |
| Teacher forcing only | 68 | 70 | 74 | |
| 76 | 64 | 68 |
| Variant | Head 0 | Head 1 | Head 2 | Mean SD |
|---|---|---|---|---|
| Full objective | 42 | 44 | 46 | |
| Without teacher forcing | 38 | 42 | 54 | |
| Without rollout loss | 36 | 36 | 28 | |
| Without endpoint loss | 38 | 40 | 44 | |
| Teacher forcing only | 24 | 30 | 40 | |
| 42 | 40 | 46 |
| Variant | Head 0 | Head 1 | Head 2 | Mean SD |
|---|---|---|---|---|
| Full objective | 100 | 100 | 100 | |
| Short training segments | 30 | 34 | 36 | |
| Long training segments | 100 | 100 | 100 | |
| Direct subgoals | 98 | 98 | 100 | |
| No head state input | 10 | 4 | 6 |
| Variant | Head 0 | Head 1 | Head 2 | Mean SD |
|---|---|---|---|---|
| Full objective | 72 | 76 | 64 | |
| 4D conditions via | 76 | 66 | 70 | |
| Direct subgoals | 58 | 64 | 56 | |
| No head state input | 2 | 0 | 4 |
World Model-Based Planning
World model-based planning uses models of how environments change to compare imagined actions and choose goal-directed behavior.
See all 245 papers