Hierarchical models speed up goal planning with learned subgoals

Learning Goal-Conditioned Subgoals through Frozen Hierarchical World Models

Computer Vision and Pattern Recognition

Summary

Planning steps in complex tasks can be slow because computers try lots of possibilities to decide what to do next. The authors show a way to speed this up by training a special network to pick subgoals quickly, without searching through many options. Their approach learns from past experiences and works well on several simulated navigation and manipulation tasks. It runs faster than previous methods while keeping good success rates.

What this means in practice

  • For roboticists: Accelerate long-horizon planning for robots by using learned subgoals without costly search, improving efficiency in real-time control tasks.
  • For video game ai developers: Enable faster decision-making for AI agents in complex virtual environments by integrating learned hierarchical world models that reduce planning overhead.

Authors

Xiang Kaibo, Liu Qingfeng

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.