H-JEPA: End-to-End Learning of Hierarchical World Models for Visual Planning

Wancong Zhang; Basile Terver; Michael Rabbat; Yann LeCun; Randall Balestriero

cs.LGOct 5, 2026

H-JEPA: End-to-End Learning of Hierarchical World Models for Visual Planning

Authors: Wancong Zhang, Basile Terver, Michael Rabbat, Yann LeCun, Randall Balestriero

Organizations: NYU · Advanced Machine Intelligence · INRIA Paris · Brown University

Abstract

Long-horizon planning with latent world models requires reasoning across timescales and levels of abstraction. Existing task-agnostic JEPA world models predict and plan at a single timescale or with multiple horizons in one shared latent space. We introduce H-JEPA, an end-to-end recipe for training a hierarchy of action-conditioned JEPAs in which each level predicts farther ahead in its own learned latent space. Planning proceeds top-down: the top level optimizes progress toward the goal, and each level's predictions become subgoals for the planner below it. When factors in the data evolve at separated timescales, higher levels discard fast, unpredictable detail and retain slower task-relevant state. Across four simulated navigation and manipulation environments, hierarchical planning improves over a flat JEPA; on Visual AntMaze, a three-level hierarchy raises success from 18% to 73% using less planner compute. Ablations attribute these gains to both temporal decomposition and higher-level goal representations. With inverse-dynamics supervision, the approach extends to diverse real-robot videos from DROID, where hierarchy improves offline planning fidelity at lower planner compute.

Figures & tables

Figure 1: Hierarchical abstractions support planning at multiple timescales. Left: visualization of hierarchical planning in Visual AntMaze. Level 3 plans toward the final goal in the most abstract latent space. Its first predicted state becomes a subgoal for the level-2 planner; this decomposition recurses to the level-1 planner, which outputs raw actions. The first and last columns decode each level’s start and goal latents, respectively; the middle columns decode its planned predictions. Higher-level latents retain the ant’s position and maze layout while losing leg-pose detail. See § 2.2 for the formal planning procedure. Right: planning success versus planner compute on AntMaze. Curves show Pareto fronts for flat LeWM [ 53 ] and two-/three-level H-JEPA; error bars are standard error (SE) over three paired model and planner seeds (§ 4.1 ).

Panel of Figure 2 — Figure 2: H-JEPA training across temporal scales. (a) With $w_{2}{=}1$ and $s_{2}{=}2$ , level 2 encodes every other level-1 state pointwise and aggregates the intervening action embeddings. Faded columns are retained at level 1 but skipped at level 2. (b) Each level predicts the next latent state conditioned on preceding latent states and actions, with SIGReg regularizing the states against collapse. All levels are trained end-to-end, with gradients from upper-level objectives propagating through lower-level encoders.

Panel of Figure 3 — Figure 3: Planning runs top-down through the hierarchy, and the lower level is steered by subgoals the upper level predicted. Two-level instance ( $L{=}2$ ) of § 2.2 . Current and goal observations $o_{0}$ and $o_{g}$ are encoded bottom-up into initial states $z^{(1)}_{0},z^{(2)}_{0}$ and goal states $g^{(1)},g^{(2)}$ . Level 2 optimizes macro-actions by unrolling $F^{(2)}$ and matching its final prediction to $g^{(2)}$ . Its predicted states $\hat{z}^{(2)}_{i}$ become subgoals for level 1, which optimizes primitive actions and lifts its predictions through $E^{(2)}$ for comparison in level-2 latent space (Eq. ( 6 )).

Figure 4: Selective abstraction across hierarchy levels. Left : MLP-probe NMSE (mean $\pm$ SE over three training seeds) at successive levels of four-level H-JEPAs on Visual AntMaze and FourRoomDistractors. Body state or distractor position becomes less recoverable at higher levels while agent position remains recoverable. Right : one observation decoded from levels $1$ – $3$ ; decoders serve only for visualization. AntMaze’s leg pose becomes less distinct; FourRoom’s blue distractor fades while the red agent remains visible.

Figure 5: Larger temporal-frequency gaps correlate with stronger selective abstraction. The x-axis measures temporal separation in each training dataset: the ratio of the characteristic frequencies (spectral centroids) of the fastest- and slowest-varying state components, shown on a log scale. The y-axis shows the change in per-entity probing NMSE from level 1 to level 2; positive values indicate poorer recoverability at level 2. Across environments, larger frequency gaps coincide with greater increases in fast-entity probing error (filled markers), while slow-entity error remains near its level-1 value (open markers). Environments with gaps below

2\times

(shaded region) show little selective abstraction.

Figure 6: Gradient-based planning across compute budgets and hierarchy depths , comparing LeWM [ 53 ] , H-JEPA, and the hierarchical world model baseline HWM [ 92 ] . Error bars in both panels show one SE over three seeds. Top: success versus planner FLOPs per episode, sweeping action-sample counts with other planner settings fixed. FLOPs include forward and backward passes. Bold curves are Pareto fronts. Bottom: success versus hierarchy depth on a separate set of

50

tasks per environment. Planner settings are selected under a

100

-TFLOP planning budget. See § E for protocols and Fig. 9 for the DROID results.

Figure 8: Temporal decomposition of planning costs. We analyze a two-level H-JEPA with

s_{2}=3

. Top: Goal cost in level-1 and level-2 space. For a held-out expert trajectory we encode each frame at both levels, and compute the distance between each timestep’s latent

z^{(\ell)}_{t}

and the goal (final-timestep) latent,

\|z^{(\ell)}_{t}-g^{(\ell)}\|_{2}^{2}

. At level 2 we subsample at the training stride. Bottom: Subgoal tracking cost. The level-2 subgoals are the level-2 latents at the subsampled timesteps. Within each subgoal segment we re-anchor the level-1 rollout from the encoded latent at the segment start, unroll it to the next subgoal, project it into level-2 space, and measure its distance to that subgoal. Within each segment, the cost decreases monotonically toward the next subgoal.

Panel of Figure 9 — Figure 9: Scene diversity makes the inverse-dynamics loss necessary, and a second level adds performance at a lower planner budget. (a) First frames of three Cube and three DROID clips, showing scene variation across episodes. (b) DROID at $5$ fps, Fréchet fidelity (%), showing the slow-feature-collapsed LeWM, LeWM+IDM, then two hierarchies: HWM, whose level 2 is an identity encoder, and H-JEPA, whose level 2 is learned. (c) The one-level and two-level models of (b) swept over planner budgets, Fréchet fidelity against measured planner TFLOPs per episode on a log axis, the bold line being the Pareto front. Two frozen-encoder world models are shown for reference: they are public checkpoints, evaluated at their native $1$ fps on the $15$ clips whose window fits at that rate, so their intervals span three planner seeds. All other intervals are SE over three train seeds.

Appendix figures & tables22 assets

Supplementary material from the paper’s appendix.

Appendix

Symbol	Description
Levels and indices
$\ell$ , $L$	Hierarchy level index, $\ell=1,\ldots,L$ ; level $1$ is the finest
$t$	Discrete time index, at the level’s own rate
$a{:}b$	Half-open slice: the $b{-}a$ indices $a,\ldots,b{-}1$
Observations, latents and actions
$o_{t}$	Image observation at time $t$

Appendix

Table 2: Notation reference. Symbols are introduced in § 2 unless the row names another section. Level superscripts are always parenthesized,

(\ell)

; a bare subscript is a time index and

a{:}b

is a half-open range. Appendix-local symbols are defined where they are used.

	LeWM	Two-level H-JEPA	Three-level H-JEPA	Four-level H-JEPA
Level-1 frames per clip	$4$	$7$	$13$	$25$
Minimum clip length (env. steps)	$16$	$31$	$61$	$121$
Training clips per epoch
FourRoom	$952{,}708$	$915{,}298$ ( $96\%$ )	$840{,}478$ ( $88\%$ )	$690{,}838$ ( $73\%$ )
AntMaze	$9{,}575{,}000$	$9{,}200{,}000$ ( $96\%$ )	$8{,}450{,}000$ ( $88\%$ )	$6{,}950{,}000$ ( $73\%$ )
Push-T	$1{,}981{,}721$	$1{,}701{,}446$ ( $86\%$ )	$1{,}143{,}623$ ( $58\%$ )	$285{,}931$ ( $14\%$ )

Appendix

Table 3: Temporal layout of one training sample and the clips it leaves per epoch. The minimum clip length is the shortest episode segment that yields one training sample, in environment steps, including both endpoint observations. All H-JEPA models use stride two and window one at every level above the first. A clip must lie within one episode (the loader reserves

5\,n_{1}

environment steps per clip), so the number of clips an epoch yields falls as the minimum clip length grows; parentheses give each count as a percentage of the LeWM count in the same environment. Each clip supplies three teacher-forced transitions per level.

Figure 10: Per-level training losses of the four-level H-JEPA models and of the DROID models. Top three rows: unweighted prediction loss and SIGReg statistic of each level of the four-level models over training, and their total training objective, the sum over levels of the prediction loss plus the weighted SIGReg term. Every level of a model uses the same SIGReg weight (

0.18

on FourRoom,

0.09

elsewhere), so the weighted curves differ from the unweighted ones by a constant factor per environment. The seed band of the total objective is narrower than its line. Bottom row: unweighted prediction, SIGReg and inverse-dynamics losses of the DROID flat LeWM+IDM (orange, dashed) and of both levels of the two-level H-JEPA (solid) over their

100

training epochs, and the total training objective, the sum of all terms weighted as in Table 5 (level 1: SIGReg

0.08

, IDM

100

; level 2: SIGReg

0.16

, IDM

50

), summed over both levels for H-JEPA. All curves are the mean over three training seeds with a min/max band, on a log scale.

	FourRoom	AntMaze	Push-T	Cube
Frame size / patch	$224$ / $14$	$64$ / $4$	$224$ / $14$	$224$ / $14$
Command / proprio dim.	$2$ / n/a	$8$ / $27$	$2$ / n/a	$5$ / n/a
Latent dim. $d$	$192$	$384$	$192$	$192$
Upper predictor layers / heads	$4$ / $8$	$5$ / $12$	$6$ / $16$	$6$ / $16$
Upper predictor width / FF width	$192$ / $1024$	$224$ / $1536$	$192$ / $2048$	$192$ / $2048$
SIGReg $\lambda$ , LeWM / H-JEPA	$0.72$ / $0.18$	$0.09$ / $0.09$	$0.09$ / $0.09$	$0.09$ / $0.09$

Appendix

Table 4: Per-environment architecture and training settings. Upper-level predictor dimensions and H-JEPA epoch counts apply to every depth. H-JEPA models are initialized independently of the LeWM baselines.

	DROID
Global
Stop-gradient on latent targets	None
Level-1 to level-2 stride	$3$
Level 1
Encoder	ViT-S/16, CLS
Predictor	$384$ / $12$

Appendix

Table 5: Training hyperparameters on DROID. Predictors are causal transformers, given as width/depth. The grouping row gives the shape of the tensor SIGReg tests:

(T,N,D)

tests normality within each timestep, over the

N

sequences of the batch. The architecture and loss settings do not depend on the training schedule, whether end-to-end or stagewise (§ C.2 ); the optimization block does, and reports the end-to-end headline runs. The peak learning rate is the same at both levels, AdamW runs at

\beta=(0.9,0.999)

, and the dataloader length caps the iterations per epoch, which gives

292

at this batch size.

Environment	transitions (train / eval)	Probed entities
FourRoomDistractors	$160$ k / $40$ k	slow: ego $xy$ (2) fast: distractor $xy$ (2)
Visual AntMaze	$500$ k / $50$ k	slow: global $xy$ (2) fast: body state (27)
Push-T	$250$ k / $2.5$ k	agent $xy$ (2), block $xy$ (2), block orientation (2), agent velocity (2)
OGBench Cube	$250$ k / $40$ k	cube position (3), cube orientation (6), effector position (3), effector yaw (2), gripper (2), arm joint positions (6), arm joint velocities (6)

Appendix

Table 6: Probing data and probed entities per environment. Transition counts are for fitting (train) and scoring (eval). Visual AntMaze is scored on one evaluation set with equal numbers of explore and stitch episodes. On FourRoomDistractors and Visual AntMaze, whose state separates into slow and fast entities (Fig. 12 ), the entities are labeled accordingly. Entity dimensions in parentheses.

Figure 11: Successful hierarchical planner executions. Selected H-JEPA executions on one evaluation task per environment, using three-level models for FourRoom and AntMaze and two-level models for Cube and Push-T. These checkpoints also appear in Fig. 6 (bottom); the recordings use their original evaluation settings rather than the compute-selected sample counts. Each row shows eight observations sampled uniformly in environment steps, followed by the goal observation supplied to the planner. Rows end at first recorded success, except Cube, which shows the full $30$ -step execution to include the completed lift. Labels give elapsed environment steps.

	FourRoom	AntMaze	Push-T	Cube
Episode budget	$210$	$350$	$90$	$30$
$\eta$ , LeWM	$0.1$	$0.3$	$0.05$	$0.05$
$\eta$ , H-JEPA (L1, upper)	$(0.08,0.1)$	$(0.06,0.1)$	$(0.1,0.075)$	$(0.1,0.075)$
$H$ , LeWM	$23$	$23$	$15$	$4$
$H$ , two-level H-JEPA	$(2,12)$	$(4,9)$	$(2,8)$	$(2,2)$
$H$ , three-level H-JEPA	$(2,2,6)$	$(4,2,5)$	$(2,2,4)$	$(2,2,1)$

Appendix

Table 7: Planning settings for the simulated depth comparison.

S

: candidate sequences;

\eta

: AdamW learning rate;

H

: configured horizon in predicted steps of that level, used by the schedule described above. Tuples list levels from lowest to highest, except the H-JEPA learning-rate row, which lists level one and the rate shared by all upper levels. Episode budgets are in environment steps. Cube’s horizon-one upper solvers are skipped, so their listed candidate counts are unused.

	DROID, $5$ fps
Episode budget	$36$
$\eta$ , LeWM+IDM	$0.01$
$\eta$ , HWM (L1, L2)	$(0.01,0.3)$
$\eta$ , H-JEPA (L1, L2)	$(0.01,0.3)$
$H$ , LeWM+IDM	$36$
$H$ , HWM	$(36,12)$

Appendix

Table 8: Planning settings on DROID.

S

: candidate sequences;

\eta

: AdamW learning rate;

H

: configured horizon in predicted steps of that level. Tuples list levels from lowest to highest. Every level runs AdamW for at most

90

iterations with early stopping when the best cost stalls; the cross-level action cost is

0

and the subgoal weight

\beta

of Eq. ( 6 ) is

0.5

, for both two-level arms. Planning is open-loop: the whole level-1 plan is executed, so the commit width equals the level-1 horizon. Each model keeps one

\eta

and these settings at every planner budget: the compute-Pareto of Fig. 9 varies only the number of candidates, chosen independently per level for the two-level arms. The listed

S

is the operating point of each bar in Fig. 9 .

Panel of Figure 12 — Figure 12: Temporal-frequency separation and selective abstraction. (a) One AntMaze episode: normalized global position (red) and $13$ body-state dimensions (blue) over $81$ level-1 frames (left), and cumulative absolute normalized increments averaged across dimensions (right). Inverse spectral centroids are about $35$ frames for position and $4$ for body state ( $8.3\times$ gap); body state accumulates roughly ten times more variation. Including the omitted velocity dimensions leaves the frequency gap unchanged. (b) Per-entity spectral centroids in cycles per level-1 frame (top, log scale) and paired level-1 and level-2 probe macro NMSE (bottom). The left group shows FourRoom, AntMaze and Humanoid; the right shows Push-T, OGBench Cube and DROID. Probe heads are matched across levels within each environment and differ across environments. Error bars are one SE over three training seeds. All entities in an environment share the same level-2 configuration.

Panel of Figure 13 — Figure 13: Planning with different training datasets on Visual AntMaze. Planning success rate (GD online, mean over three seeds), using explore , stitch , or their union for training. Left: flat level-1 planning, with one model trained on each of the three datasets. Middle: two-level hierarchical planning after stagewise training, covering all nine combinations of level-1 (rows) and level-2 (columns) datasets. Right: two-level hierarchical planning after end-to-end training, with both levels trained on the same dataset in each run, giving only diagonal entries. Grey cells were not run. The end-to-end runs are FLOP-matched to the stagewise ones.

		Level-1 planning: vary cost space				Hierarchical planner
Environment	Depth	Native L1 $\uparrow$	L2 proj. $\uparrow$	L3 proj. $\uparrow$	L4 proj. $\uparrow$	H-JEPA $\uparrow$
FourRoom Distractors	LeWM (flat)	$40.7\pm 6.8$
	2 levels	$9.3\pm 0.7$	$38.7\pm 6.6$			$81.3\pm 6.7$
	3 levels	$43.3\pm 8.7$	$65.3\pm 5.5$	$71.3\pm 4.1$		$96.0\pm 1.1$
	4 levels	$68.7\pm 4.7$	$72.0\pm 3.1$	$66.0\pm 3.1$	$70.7\pm 2.7$	$80.0\pm 4.0$
OGBench Cube	LeWM (flat)	$32.0\pm 6.1$

Appendix

Table 9: Planning-cost ablations across environments. Success (%, mean

\pm

SE, three seeds), using the same protocol as Table 1 . Within each H-JEPA row, only the cost space varies across the four level-1 planning columns; LeWM and hierarchical H-JEPA are references. For FourRoom, bold marks the highest mean among cost spaces. FourRoom gains from projected costs, while Cube and Push-T show no clear benefit.

		HWM		H-JEPA
Environment	Level-1 $\uparrow$	Level-2 $\uparrow$	subgoal $\uparrow$	Level-2 $\uparrow$	subgoal $\uparrow$
Visual AntMaze	0.48	0.52	0.94	0.68	0.94
Push-T	0.75	0.77	0.85	0.77	0.89
OGBench Cube	0.65	0.74	0.79	1.00	0.99
FourRoom Distractors	0.56	0.62	0.92	0.85	1.00

Appendix

Table 10: Prediction-cost monotonicity (

-\mathrm{Spearman}

\|z^{(\ell)}-g^{(\ell)}\|_{2}^{2}

against time; higher is better; mean over three training seeds, held-out episodes). L2 columns score each model’s level-2 encoding against the final goal; subgoal scores a level-1 rollout, re-anchored to the true state at each subgoal (bounded horizon), against the level-2 subgoals (per-segment).

Figure 14: On OGBench Cube, where the outcome is observable, Fréchet fidelity separates the episodes that can succeed from those that cannot, but not the ones that do. One frozen checkpoint,

9

configuration groups

\times

3

planner seeds

\times

50

episodes. (a) One point per planner configuration; in the shaded band the flat and hierarchical planners have equal offline fidelity and opposite grasp outcomes. (b) Every episode: none succeeds below a fidelity threshold, and most above it still fail.

Figure 15: Against a continuous outcome the same offline metric ranks degree of progress. RoboCasa CloseDrawer, one frozen checkpoint,

99

planner configurations

\times

32

held-out episodes, budget

200

steps; the configurations are varied to spread the outcomes, not tuned for task performance.

x

is Fréchet fidelity, computed exactly as on DROID;

y

is graded drawer-closure success, a continuous outcome DROID cannot report. (a) One point per planner configuration, median fidelity against mean graded success. (b) Every episode, with the mean of each fidelity decile joined. The claim rests on the within-configuration rank correlation, which holds episode difficulty and planner settings fixed; the pooled and between-configuration values are inflated by a between-configuration mean shift.

Figure 16: What the DROID planner imagines, on every clip of the evaluation set. Two-level stagewise model, held-out, training seed 1; all $16$ clips ordered best to worst by Fréchet fidelity, down the left column then the right, $+75\%$ down to $+0\%$ at the zero-action floor. Each clip is two strips: above, the logged expert clip; below, the decoded plan, which is the level-1 rollout under the planner’s committed commands with the level-2 macro-plan interleaved into it at every third column (temporal stride $3$ ), ringed in blue. Columns are timesteps. The leftmost is the context frame the planner is given (green); the rest are imagined (red), and blur along the rollout, which is the decoder’s fidelity limit and not a planning failure.

Figure 17: SIGReg restores representation rank; the inverse-dynamics term is what turns rank into planning fidelity. Full grid around the level-1 DROID configuration of table 5 : SIGReg

\in\{0,0.08\}

\times

IDM weight

\in\{0,50,100,200\}

\times

prediction-target detachment,

16

cells, one training seed each, read at epoch

100

. Five SIGReg-off cells (hatched) were stopped once collapsed and are read at their last epoch. Left: level-1 effective rank [ 25 ] of the CLS embedding, max

384

. Right: Fréchet fidelity on the 16-clip set; the grey band holds sub-floor cells, drawn as fixed-depth stubs annotated with their true value. Each planning bar is the mean over planner seeds

1

2

and

3

of one trained model, and its error bar is the standard error (SE) over those three planner seeds.

Figure 18: The SIGReg and inverse-dynamics coefficients are chosen jointly, once per level. DROID at $5$ fps; Fréchet fidelity (%, higher is better), mean and standard error over three training seeds, each averaged over three planner seeds. Left , the level-1 grid on single-level models; right , the level-2 grid on stagewise models with level 1 frozen. Both use the gradient-based planner of our headline results, and share one color scale. The heavy outline marks the coefficients we keep; at level 1 they lie within one standard error of the best cell.

		$\mathcal{T}_{\mathrm{EP}}$
Shards	$\\|\nabla\\|$ (ours)	ours	single-GPU reference
$1\times 256$	$0.349154178012$	$1.372619178543$	$1.372619$
$2\times 128$	$0.349154178012$	$1.372619178543$	$1.180770$
$4\times 64$	$0.349154178012$	$1.372619178543$	$1.186233$

Appendix

Table 11: Our distributed SIGReg is exactly invariant to how a fixed total batch is sharded; a single-GPU implementation is not. One fixed Gaussian matrix (

n=256

D=32

M=64

directions) sharded contiguously across

W

ranks in float64 . The reference is the SIGReg released with LeWM [ 53 ] , formula-identical to ours but with its prefactor set by the local batch and no collective; its code documents itself as single-GPU. The data follow the target, so its

14\%

drift is launch topology, not signal.

Figure 19: Scene diversity sends latent variance between episodes. Within-episode (the slow-feature degree

\chi

of Definition 1 ) and between-episode shares of the latent variance. Bars are

\pm

1 standard error (SE) over three training seeds.

$\lambda$	$\lambda_{\mathrm{sim}}$	Onset $\uparrow$	Sens. $<0.05$ $\uparrow$	Best fidelity $\uparrow$	Final fidelity $\uparrow$	Final sens.	$\chi$
$9.77\times 10^{-6}$	$0$	$7$	$13$	$-0.64$ (2.74)	–	–	–
$1.95\times 10^{-5}$	$0$	$21$	$20$	$3.35$ (1.40)	–	–	–
$3.91\times 10^{-5}$	$0$	$27$	$26$	$0.63$ (5.49)	–	–	–
$7.81\times 10^{-5}$	$0$	$34$	$39$	$5.86$ (5.10)	–	–	–
$1.56\times 10^{-4}$	$0$	$48$	$57$	$8.38$ (2.29)	$-81.03$	$0.014$	$0.011$
$1.56\times 10^{-4}$	$-0.05$	–	none	$6.46$	$\mathbf{2.44}$ (5.80)	$0.301$ (0.001)	$0.17$

Appendix

Table 12: Training without the inverse-dynamics term: the SIGReg ladder (top) and time repulsion at its largest coefficient (middle). Level-1 models at

5

fps,

\gamma=0

, three train seeds,

100

epochs. These runs use the SIGReg grouping that places time in the sample axis, so they carry the

\Delta

of Proposition 2 ; the barrier is what the ladder probes. “Onset” is the first epoch at which Fréchet fidelity (Eq. ( 7 )) drops clearly below the zero-action floor, the next column the epoch at which the predictor’s action sensitivity first falls below

0.05

(seed spread

\pm 2

epochs). Best fidelity is the highest before turning over; final fidelity and sensitivity are read over epochs

98

–

100

;

\chi

is Eq. ( 10 ) on the final checkpoint (one seed for the repulsion runs, three for the IDM-on row). Ladder rows carry no terminal value (at most two seeds reach epoch

100

Method	Pixels (In)	Task-Agnostic World Models	Reconstruction Free	Hierarchy of Learned Latent Spaces
Classical hierarchical MPC [ 23 , 51 , 45 ]	✗	✓	✗	✗
Director [ 32 ]	✓	✗	✗	✗
Puppeteer [ 36 ]	✓	✗	✓	✗
SPlaTES [ 28 ]	✗	✗	✓	$\sim$
Schiewer et al. [ 69 ]	✗	✗	✗	✓
THICK [ 27 ]	✓	$\sim$	✗	✗

Appendix

Table 13: Model-based hierarchical control methods along four axes that define H-JEPA. Pixels : operates from high-dimensional visual observations rather than a low-dimensional state. Task-agnostic : reward-free training with no task signal; the task objective (a goal or a planning cost) is defined only at test time. Reconstruction Free : learns representations and dynamics with no reconstruction objective, predicting latent targets rather than decoding observations. Hierarchy of learned latent spaces : a stack of distinct latent spaces, each a coarser abstraction of the one below that discards detail (vs. a single/shared space, parallel spaces, or a hierarchy of decisions only). ✓ yes, ✗ no,

\sim

partial. H-JEPA is the only method satisfying all four; HWM, our baseline, differs on exactly one axis.

Explore similar work

Jun 8, 2026cs.AI

FF-JEPA: Long-Horizon Planning in World Models with Latent Planners

Joint Embedding Predictive Architectures (JEPAs) have shown promising world modeling capabilities, enabling planning in latent space by optimizing action trajectories using methods like the Cross-Entropy Method (CEM). These methods are, however, too computationally expensive and ineffective for long-horizon planning. Furthermore, these methods typically require an explicit image of the goal state, which is not always possible in real-world tasks. In this work, we tackle these limitations by proposing Forward-Forward-JEPA (FF-JEPA), a hierarchical approach leveraging two forward dynamics models. Alongside a standard action-conditioned forward model, we introduce an action-free latent planner that predicts the next subgoal given the current state. This approach removes the need for goal images and enables long-horizon planning by decomposing complex trajectories into a sequence of tractable, short-term optimization problems. Preliminary results on PushT demonstrate that FF-JEPA successfully overcomes flat world models' long-horizon collapse, highlighting this approach as a promising direction for goal-free planning.

Sergi Masip, Jonathan Swinnen, Yutong Hu +2

PSI, ESAT, KU Leuven · RAM, MECH, KU Leuven

Long-Horizon Planning Low-Latency Latent Planning Latent World Models Planning World Models S-Jepa

Ask PDF

Sep 3, 2026cs.RO

Toward Physically Grounded JEPA World Models for Goal-Conditioned Robotic Planning

Action-conditioned JEPA world models enable planning toward visually specified goals without reconstructing future pixels, yet latent prediction alone does not explicitly encourage the learned representations to retain information relevant to robotic control. We introduce an end-to-end JEPA world model that augments latent prediction with inverse dynamics (IDM) and state alignment (SA). While inverse dynamics discourages latent collapse and makes latent transitions informative of the actions that produced them, state alignment grounds consecutive representations in their associated physical configuration and motion. Across four benchmark tasks, our model attains the highest success rates on TwoRoom (100%), PushT (98%), and OGBench-Cube (87%), while performing comparably to LeWorldModel on Reacher. Our ablation further shows that adding state alignment consistently improves planning success over IDM alone across all four tasks. Although LeWorldModel, our primary baseline, attains higher average straightening on OGBench-Cube, transition-subspace analysis shows that its transition energy is concentrated in a substantially lower-dimensional subspace. Our state-aligned model exhibits a higher effective transition dimension than LeWorldModel and improves planning over IDM alone, supporting state alignment as an effective complement to inverse dynamics for robotic planning.

Muyuan Liu, Yue Huang, Zheng Liang +1

GENISOM AI, Beijing, China

Inverse Dynamics Robot Planning World Models

Ask PDF

Jun 30, 2026cs.LG

AdaJEPA: An Adaptive Latent World Model

Latent world models enable planning from high-dimensional observations by predicting future states in a compact latent space. However, these models are typically kept frozen at test time: when their predictions become inaccurate, planning can fail, especially under test-time distribution shift. To address this, we propose AdaJEPA, an adaptive latent world model that performs test-time adaptation within the closed loop of model predictive control (MPC). After training, AdaJEPA plans and executes the first action chunk, uses the observed next-state transition as a self-supervised adaptation signal, and replans with the updated model. This closed-loop update continuously recalibrates the world model without additional expert demonstrations. Across a range of goal-reaching tasks, AdaJEPA substantially improves planning success with as few as one gradient step per MPC replanning step.

Ying Wang, Oumayma Bounou, Yann LeCun +1

New York University · AMI Labs

Latent World Models Model Predictive Control World Models World Test-Time Adaptation

Ask PDF