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

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

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∈Sg\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γπzm^{\pi_{z}}_{\gamma} by periodically setting z=∅,a=∅z=\varnothing,a=\varnothing .
Figure 3: Success rate ( ↑\uparrow ) of base policies πg\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 (a1,z1,…zn)(a_{1},z_{1},\ldots z_{n}) (Eq. 15 ) or (z1,…,zn)(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.
Figure 9: Qualitative plots of GHM samples on antmaze-giant task 1 at different horizons (0.99,0.995,0.998)(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.
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.
MethodHyperparameterAntmazeCube
CRL GC-TD3 GC-1SFQL Distillation coefficient{0.1,0.15,0.2,0.3,0.4}\{0.1,0.15,0.2,0.3,0.4\}{0.7,0.8,0.9,1,3}\{0.7,0.8,0.9,1,3\}
Discount factor{0.995,0.997}\{0.995,0.997\} for giant 0.990.99 otherwise{0.99,0.995}\{0.99,0.995\} for cube-4 0.990.99 otherwise
Gradient steps{1M,3M}\{1\text{M},3\text{M}\}500k500\text{k} for cube- {1,2}\{1,2\} 1M1\text{M} for cube-3 3M3\text{M} for cube-4
GC-BCDiscount factor{0.99,0.995}\{0.99,0.995\} for giant {0.98,0.99}\{0.98,0.99\} for medium {0.98,0.99}\{0.98,0.99\} for large{0.95,0.96}\{0.95,0.96\} for cube-1 {0.96,0.97}\{0.96,0.97\} for cube-2 {0.96,0.97,0.98}\{0.96,0.97,0.98\} for cube-3 {0.96,0.98,0.99}\{0.96,0.98,0.99\} for cube-4
Gradient steps125k125\text{k}500k500\text{k} for cube-1 1M1\text{M} for cube-2 2M2\text{M} for cube-3 3M3\text{M} for cube-4
HFBC SHARSALookahead[25,50][25,50][25,50][25,50]
Appendix
Table 10: Base policy hyperparameters. Parameters in {  }\{\;\} denote sweeps performed over the values inside brackets.
 HyperparameterValue
Flow Matching ( Lipman et al., 2023 )Probability PathConditional OT ( σ=0\sigma{=}0 )
Time SamplerU([0,1])\mathcal{U}([0,1])
ODE SolverEuler
ODE dt\mathrm{d}t (train) / steps0.10.1 / 1010
ODE dt\mathrm{d}t (eval) / steps0.050.05 / 2020
Network (U-Net) ( Ronneberger et al., 2015 )tt -Positional Embedding Dim.256256
Appendix
Table 11: Hyperparameters for Geometric Horizon Model pre-training.
More in this topic

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

Explore similar work

CardsList
Oct 5, 2026

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

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.

Wancong Zhang, Basile Terver, Michael Rabbat +2

NYU · Advanced Machine Intelligence · INRIA Paris +1

World Models for RoboticsWorld ModelsWorld Model-Based PlanningJoint-Embedding Predictive ArchitectureLong-Horizon PlanningHierarchical World Models+1
AskPDF
Jul 30, 2026

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

Building generalizable agents for diverse applications remains a fundamental challenge. While imitation learning-based policies succeed in specific training environments, they often fail to generalize to novel scenes and tasks. In this work, we propose World Action Planner, a robot planning system that leverages the reasoning capabilities of Vision-Language Models (VLMs) and the physical grounding of a multi-task pose-image conditioned world model. Our system enables an agent to propose initial action plans and iteratively refine them via optimization and search, reasoning over imagined world model rollouts. We demonstrate that our approach achieves superior performance across compositional tasks, new layouts, and zero-shot generalization scenarios, significantly outperforming state-of-the-art end-to-end policy models such as VLAs and WAMs. Project website at worldactionplanner.github.io

Xiangcheng Zhang, Yilun Du

Harvard University

World Models for RoboticsWorld Model-Based PlanningAction-Conditioned World ModelsRobot Policy GeneralizationWorld Models for VLA Models
AskPDF
Jun 8, 2026

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

Hierarchical PlanningLong-Horizon Robotic ManipulationWorld Model-Based PlanningJoint-Embedding Predictive ArchitectureLong-Horizon PlanningModel-Based Planning+1
AskPDF