Can hierarchical JEPA world models improve long-horizon visual planning?

One shared latent space forces a tradeoff

Existing task-agnostic JEPA world models predict and plan within a single latent space, often at a single timescale. The paper identifies two consequences. First, at a single timescale a long-horizon task must be rolled out as many fine-grained steps, so prediction error compounds and the action search space grows. Second, one shared latent space has to serve both low-level dynamics and goal matching: a latent that carries every fast-varying detail may model low-level dynamics well but gives a poor goal-matching cost when the goal is more abstract than the state, such as a location to reach rather than a pose to match.

H-JEPA's premise is that a hierarchy divides this work. Each higher level predicts over a longer horizon and keeps in its latent space only what remains predictable at that horizon, while lower levels handle fast features over short strides for short-range prediction and low-level control. The paper motivates this with cortical hierarchies of intrinsic timescales, where slower-varying representations sit higher up, but the tested claim is narrower: when factors in the data evolve at separated timescales, higher levels discard fast, unpredictable detail and retain slower task-relevant state.

Each level predicts in its own latent space

H-JEPA is an end-to-end recipe for training a hierarchy of action-conditioned JEPAs, with no reward and no reconstruction objective. Each level predicts the future within its own learned latent space, at a coarser temporal stride and a higher degree of abstraction than the level below. Planning proceeds top-down: the top level optimizes progress toward the goal, and each level's predictions become subgoals for the planner below it. Concretely, the first predicted latent is passed down as a subgoal, and the lower level's rollout is encoded back into the upper latent space for comparison.

Every active level optimizes candidate action sequences by differentiating through its latent rollout, using AdamW for at most 90 iterations with early stopping when the best cost stalls, and returns the lowest-cost candidate. Level-1 candidates are initialized from a standard Gaussian in normalized action space. Upper levels have no native action bounds, so their search is anchored to the macro-actions seen in training: each upper-level action encoder keeps its last 4000 training macro-actions, candidates are initialized from a per-dimension Gaussian fitted to them, and after every optimizer step candidates are clipped to that sample's per-dimension 2nd to 98th percentiles. The latent-distance cost scores the last two predicted states at level 1 and the last state at upper levels.

Execution follows a commit schedule: the agent replans at environment times t0 = 0 < t1 < t2 < ..., encodes the observation at each replanning time, runs the top-down solve, executes the first actions of the level-1 plan, and discards the rest. All models execute two level-1 action blocks, 10 environment steps, before replanning from scratch. Level 1 advances 5 environment steps per predicted step, and each upper level doubles that interval. Committing one command per phase is fully closed-loop control; committing an entire plan in one phase is open-loop single-shot execution.

Visual AntMaze success rises from 18% to 73%

The headline comparison is on Visual AntMaze, where a three-level hierarchy raised planning success from 18% to 73% relative to a flat JEPA, and did so using less planner compute. The paper reports this as a success-rate improvement across four simulated navigation and manipulation environments, in which hierarchical planning improved over flat planning. The compute comparison matters as much as the success number: the gain is not bought by spending more search at test time.

The reported operating points are specific. Episode budgets are 210 environment steps on FourRoom, 350 on AntMaze, 90 on Push-T and 30 on Cube. Configured horizons in predicted steps differ by depth: for example, the three-level H-JEPA uses horizons (2, 2, 6) on FourRoom, (4, 2, 5) on AntMaze, (2, 2, 4) on Push-T and (2, 2, 1) on Cube, listed from lowest to highest level. Candidate counts per level also vary, such as (50, 200, 200) for the three-level model on FourRoom. On the compute-sweep tasks, LeWM and H-JEPA select the candidate allocation with the highest observed mean success rate under 100 mean planner TFLOPs per episode, breaking ties by lower compute, and those settings are then evaluated without further tuning on a separate 50-task set.

Temporal decomposition and abstract goal scoring

Ablations attribute the gains to two complementary mechanisms rather than to hierarchy as such. The first is temporal decomposition: splitting a long-horizon task into easier subproblems at coarser strides. The second is higher-level goal representations: hierarchical latent spaces let higher-level planners score progress at the goal's own level of abstraction, instead of forcing an abstract goal through a latent that also carries fast-varying detail.

The paper's probe analysis gives the mechanism a measurable signature. For each state entity, the authors compute the spectral centroid of its amplitude-normalized trajectory stream at the level-1 sampling rate, pooled over episodes, and define the entity-frequency gap as the ratio of the largest to smallest entity centroid. Humanoid, AntMaze and FourRoom have gaps of 3.1 to 12.7 times, and there body-state or distractor decoding worsens at level 2 while agent position remains near its level-1 floor. Push-T, OGBench Cube and DROID have gaps of only 1.2 to 1.6 times and show little selective abstraction, with comparable entity recoverability at both levels. The relevant comparison is the within-entity change between levels, since absolute NMSE differs across environments, particularly for DROID's natural video.

A single AntMaze episode illustrates the separation: inverse spectral centroids are about 35 frames for global position and 4 for body state, an 8.3-times gap, and body state accumulates roughly ten times more variation. The paper is careful that frequency separation alone does not establish a difference in predictability, since fast periodic motion can remain predictable; the per-entity probes are what relate the separation empirically to content retained at level 2.

A second level helps only with structured data

The paper tests whether training each hierarchy level on a different type of data improves planning, using stagewise training that freezes each lower level before training the next. With a two-level hierarchy on Visual AntMaze's explore and stitch dataset types and their union, matched data shows a second level helps only when the data is temporally structured. Along the diagonal, adding an abstract level improves long-horizon planning on stitch and union: end-to-end hierarchical planning success rises to 21 from the single-level 11 on stitch, and to 54 from 19 on union. On explore, neither two-level model beats the single-level baseline of 18, with stagewise and end-to-end both at 11.

The paper reads this as the level-2 model being unable to form useful abstractions when the training data lacks the temporal structure the hierarchy is meant to exploit. This is a scope condition on the headline result: the Visual AntMaze gain from 18% to 73% is reported for the paper's evaluation setting, while the dataset comparison shows that the benefit of an added level depends on whether the data itself separates into slow and fast factors.

Inverse dynamics extends the recipe to DROID

With an inverse-dynamics term, the method extends to DROID, a real-robot manipulation corpus whose scene, lighting and objects change every episode. There, hierarchy improves offline planning fidelity at lower planner compute. The level-1 model used for this is LeWM+IDM, following the use of inverse dynamics as a way to retain information about the moving agent, with SIGReg providing marginal anti-collapse regularization.

The DROID setting is materially different from the simulated environments and should not be read as closed-loop robot control. DROID is offline and open-loop: one plan covers the whole window, so the commit width equals the horizon and nothing is replanned. It is scored on 16 clips with three planner seeds and three training seeds. The episode budget is 36, the level-1 horizon is 36 and the level-2 horizon is 12, with 16 candidate sequences at each level for both HWM and H-JEPA. The subgoal weight is 0.5 for both two-level arms, and the cross-level action cost is 0. The paper's own stated next step is to test whether these gains translate to closed-loop control on a physical robot.

Where the hierarchical claim holds and stops

The paper positions H-JEPA against HWM, a task-agnostic JEPA world model that plans hierarchically but whose predictors over different horizons share one latent space. The paper's framing is that HWM decomposes the horizon but cannot match objectives in more abstract spaces, and that H-JEPA differs on exactly the axis of a hierarchy of distinct learned latent spaces. The four-axis comparison in the paper places H-JEPA as the only listed method that is pixel-based, task-agnostic, reconstruction-free and built on a hierarchy of learned latent spaces.

The practical reading is a priority order rather than a blanket endorsement. The strongest evidence is the Visual AntMaze comparison, where a three-level hierarchy raised success from 18% to 73% at lower planner compute, and the ablations tie that gain to temporal decomposition plus higher-level goal representations. The dataset comparison narrows when an extra level pays off: on stitch and union it does, on explore it does not. The DROID result extends the recipe to real-robot video but only for offline, open-loop planning fidelity. A reader deciding whether to adopt a hierarchy should therefore first ask whether their data separates into slow and fast factors, and whether their evaluation is closed-loop control or offline plan scoring, because those two conditions are what the paper's own results turn on.

Sources used in this answer