Compositional Planning with Jumpy World Models

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 Thakoor et al. (2022), 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 (Farebrother et al., 2025), 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.

Artificial Intelligence (AI) often struggles with complex tasks that require long-term planning. When an AI tries to plan hundreds of tiny, individual actions into the future, small prediction errors quickly pile up and cause the entire plan to fail. Instead of planning one small action at a time, we developed "jumpy world models" that allow the AI to think in terms of entire behaviors. By predicting the final outcome of a complete behavior, essentially "jumping" over the intermediate steps, our system can intelligently mix and match previously learned behaviors to form a complete plan. We also added a consistency mechanism to ensure these long-term "jumpy" predictions remain accurate across different timeframes. This approach allows AI agents to solve entirely new, complex problems by recombining what they already know. In our tests on challenging robotic navigation and object manipulation tasks, planning with these jumpy models improved performance by 200% compared to traditional step-by-step planning.