Compositional Planning with Jumpy World Models

Jesse Farebrother; Matteo Pirotta; Andrea Tirinzoni; Marc G. Bellemare; Alessandro Lazaric; Ahmed Touati

cs.LGFeb 23, 2026

Compositional Planning with Jumpy World Models

Authors: Jesse Farebrother, Matteo Pirotta, Andrea Tirinzoni, Marc G. Bellemare, Alessandro Lazaric, Ahmed Touati

Organizations: Mila – Québec AI Institute · McGill University · FAIR at Meta

Abstract

The ability to plan with temporal abstractions is central to intelligent decision-making. Rather than reasoning over primitive actions, we study agents that compose pre-trained policies as temporally extended actions, enabling solutions to complex tasks that no constituent alone can solve. Such compositional planning remains elusive as compounding errors in long-horizon predictions make it challenging to estimate the visitation distribution induced by sequencing policies. Motivated by the geometric policy composition framework introduced in arXiv:2206.08736, we address these challenges by learning predictive models of multi-step dynamics -- so-called jumpy world models -- that capture state occupancies induced by pre-trained policies across multiple timescales in an off-policy manner. Building on Temporal Difference Flows (arXiv:2503.09817), we enhance these models with a novel consistency objective that aligns predictions across timescales, improving long-horizon predictive accuracy. We further demonstrate how to combine these generative predictions to estimate the value of executing arbitrary sequences of policies over varying timescales. Empirically, we find that compositional planning with jumpy world models significantly improves zero-shot performance across a wide range of base policies on challenging manipulation and navigation tasks, yielding, on average, a 200% relative improvement over planning with primitive actions on long-horizon tasks.

Figures & tables

Figure 1: Percentage point change (

\uparrow

) over zero-shot policies. We compare: action-level planning ( ActionPlan ) with a world model; generalized policy improvement (GPI) and compositional planning ( CompPlan ; ours) with GHMs.

Appendix figures & tables11 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 2: Two proposal distributions for goal-conditioned compositional planning. Left: GoalCondProposal samples subgoal sequences by chaining GHM predictions conditioned on the goal

g\in\mathsf{S}

, guiding the agent toward the target. Right: Unconditional proposal samples from the behavior policy’s GHM that can be trained alongside

m^{\pi_{z}}_{\gamma}

by periodically setting

z=\varnothing,a=\varnothing

Figure 3: Success rate (

\uparrow

) of base policies

\pi_{g}

(Zero Shot) and compositional planning as in Equation ( 6 ) with GHMs ( CompPlan , ours) averaged over tasks.

Figure 4: Success rate (

\uparrow

) of planning vs. zero shot. We consider action-level planning ( ActionPlan ) with a world model; generalized policy improvement (GPI) and compositional planning ( CompPlan ; ours) with GHMs.

Figure 5: Percent point change (

\uparrow

) of planning over zero shot. We consider action-level planning ( ActionPlan ) with a world model; generalized policy improvement (GPI) and compositional planning ( CompPlan ; ours) with GHMs.

Figure 6: Success rate (

\uparrow

) of CompPlan with different base policies when planning at every step of every 5 steps. We report mean and standard deviation 3 seeds for GHM training.

Figure 7: Success rate (

\uparrow

) of CompPlan with different base policies when maximizing over

(a_{1},z_{1},\ldots z_{n})

(Eq. 15 ) or

(z_{1},\ldots,z_{n})

(Eq. 16 ). We report mean and standard deviation 3 seeds for GHM training.

Figure 8: Success rate (

\uparrow

) of CompPlan with different base policies when using the conditional or unconditional GHMs to propose subgoals. We report mean and standard deviation 3 seeds for GHM training.

Panel of Figure 9 — Figure 9: Qualitative plots of GHM samples on antmaze-giant task 1 at different horizons $(0.99,0.995,0.998)$ from td-flow and td-hc with the last row depicting the ground truth discounted occupancy. As can be seen in the figure, performance is comparable at smaller horizons with td-hc doing a much better job at capturing the true distribution as the horizon increases.

Panel of Figure 10 — Figure 10: Qualitative visualization for an episode of cube-4-task-5 . The robots and cubes shown with transparency indicated GHM samples from the first policy in the GSP.

Method	Hyperparameter	Antmaze	Cube
CRL GC-TD3 GC-1S	FQL Distillation coefficient	$\{0.1,0.15,0.2,0.3,0.4\}$	$\{0.7,0.8,0.9,1,3\}$
	Discount factor	$\{0.995,0.997\}$ for giant $0.99$ otherwise	$\{0.99,0.995\}$ for cube-4 $0.99$ otherwise
	Gradient steps	$\{1\text{M},3\text{M}\}$	$500\text{k}$ for cube- $\{1,2\}$ $1\text{M}$ for cube-3 $3\text{M}$ for cube-4
GC-BC	Discount factor	$\{0.99,0.995\}$ for giant $\{0.98,0.99\}$ for medium $\{0.98,0.99\}$ for large	$\{0.95,0.96\}$ for cube-1 $\{0.96,0.97\}$ for cube-2 $\{0.96,0.97,0.98\}$ for cube-3 $\{0.96,0.98,0.99\}$ for cube-4
GC-BC	Gradient steps	$125\text{k}$	$500\text{k}$ for cube-1 $1\text{M}$ for cube-2 $2\text{M}$ for cube-3 $3\text{M}$ for cube-4
HFBC SHARSA	Lookahead	$[25,50]$	$[25,50]$

Appendix

Table 10: Base policy hyperparameters. Parameters in

\{\;\}

denote sweeps performed over the values inside brackets.

	Hyperparameter	Value
Flow Matching ( Lipman et al., 2023 )	Probability Path	Conditional OT ( $\sigma{=}0$ )
	Time Sampler	$\mathcal{U}([0,1])$
	ODE Solver	Euler
	ODE $\mathrm{d}t$ (train) / steps	$0.1$ / $10$
	ODE $\mathrm{d}t$ (eval) / steps	$0.05$ / $20$
Network (U-Net) ( Ronneberger et al., 2015 )	$t$ -Positional Embedding Dim.	$256$

Appendix

Table 11: Hyperparameters for Geometric Horizon Model pre-training.

Long-Horizon Planning

Long-horizon planning studies how agents coordinate actions over extended sequences to achieve goals amid changing states and uncertainty.

See all 60 papers

H-JEPA: End-to-End Learning of Hierarchical World Models for Visual Planning
Oct 5, 2026Wancong Zhang, Basile Terver, Michael Rabbat +2World Models for Robotics World Models
NYU · Advanced Machine Intelligence · INRIA Paris +1
World Action Planner: Generalizable Decision-Making with Action-Conditioned World Models
Jul 30, 2026Xiangcheng Zhang, Yilun DuWorld Models for Robotics World Model-Based Planning
Harvard University
FF-JEPA: Long-Horizon Planning in World Models with Latent Planners
Jun 8, 2026Sergi Masip, Jonathan Swinnen, Yutong Hu +2Hierarchical Planning Long-Horizon Robotic Manipulation
PSI, ESAT, KU Leuven · RAM, MECH, KU Leuven

Oct 5, 2026J/K move · Enter open · S save

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

Wancong Zhang, Basile Terver, Michael Rabbat, Yann LeCun+1

NYU · Advanced Machine Intelligence · INRIA Paris+1

World Models for Robotics World Models World Model-Based Planning Joint-Embedding Predictive Architecture Long-Horizon Planning Hierarchical World Models+1

Open paper AskPDF arXiv

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.

Abstract

Long-Horizon Planning

Explore similar work

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

World Action Planner: Generalizable Decision-Making with Action-Conditioned World Models

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