Driving · 2026-10-20

Learning Goal-Conditioned Subgoals through Frozen Hierarchical World Models

Xiang Kaibo, Liu Qingfeng

AI-assisted summary based on the paper. Consult the linked source for methods and limitations. Benchmark numbers appear after editor verification.

TL;DR

The authors replace expensive high-level macro-action search in hierarchical world models (HWM) with a lightweight goal-conditioned head trained through a frozen high-level predictor. The head outputs predictor conditioning from (latent state, goal, remaining-step encoding). Training uses hindsight goals and state-prediction losses on offline trajectories (no search-solution targets or macro-action labels). Inference keeps autoregressive rollout but removes high-level candidate optimization.

Why it matters

High-level plan optimization can dominate evaluation time in hierarchical latent world models. This work evaluates whether learning predictor conditions through a frozen hierarchical model can match or improve success while eliminating high-level candidate optimization, and reports batched evaluation speedups on NVIDIA L40S along with ablations on condition design choices.

Method

Learn a trainable goal-conditioned head gθ that produces frozen coarse predictor inputs; the coarse predictor F2 generates the next latent waypoint, while the low level MPC planner is unchanged.
Train only the head by differentiating through the frozen predictor using hindsight goals on offline trajectories, with teacher-forced transition loss, autoregressive rollout loss, and a separate endpoint loss; no macro-action labels or search-solution targets.
At inference, perform the autoregressive high-level rollout using the learned head without high-level candidate optimization; use the first predicted waypoint as the low-level target and execute primitive actions with the released low-level optimizer.

Limitation

The authors report that their comparisons only replace the fixed HWM upper-level search under a common protocol and “do not establish dominance over all possible search settings.”

Abstract (from arXiv)

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.

Publication record

Bibliographic evidence describes where and when a record appeared. It does not validate scientific claims.

First public release: 2026-10-19 · Indexed by Pub-AI: 2026-10-20 UTC

preprint