Investigating Priors for Informed Latent Planning in World Models
Abstract
World models leveraging the Joint Embedding Predictive Architecture (JEPA) are trained to reach goals zero-shot through latent planning. Most work starts from an uninformed Gaussian prior for planning, while a few recent works show the advantage of learned priors. However, the effectiveness of different priors is not thoroughly studied. In this work, we investigate three priors learned by different policy heads, i.e., a deterministic head (Behavior Cloning), a Gaussian one, and a generative one (Flow Matching). We evaluate these priors on success rate together with planning latency across four world model backbones and six robot environments. Our results show that learned priors can improve the success rate by a large margin, e.g., by percentage points in the Cube environment. Moreover, learned priors exhibit computational efficiency, reducing planning time per action sequence by a factor of to across all environments, highlighting their practical utility. We show that while no single prior dominates in the short-horizon setting, Flow Matching learns a more robust prior most capable of challenging long-horizon planning tasks. Based on our findings, we distill guidance on choosing between priors, and other algorithm settings such as replanning cadence, to inform the design of future latent planning systems for real-world applications.
est. 32% chance this paper gets accepted at ICLR 2027.
What do you think this paper will get?
All positions stay anonymous.