Value-guided action planning with JEPA world models
Matthieu Destrade, Oumayma Bounou, Quentin Le Lidec, Jean Ponce, Yann LeCun
TL;DR
The paper tackles action planning with JEPA world models by shaping the embedding space so that the distance between state embeddings approximates the negative goal-conditioned value for reaching a target, enabling more effective planning with MPC. It formalizes this through an IQL-based objective where $V_\theta(s,g) = -\|\mathcal{E}_\theta(s) - \mathcal{E}_\theta(g)\|_2$, and experiments with Sep and joint training regimes, including a quasi-distance variant. Across offline Wall and Maze tasks, value-guided representations—especially the quasi-distance approach—consistently improve planning success over standard JEPA methods, with dataset characteristics (e.g., WB vs WS) influencing outcomes. The work highlights locality issues and data coverage as key factors, suggesting hierarchical representations and targeted data collection to better capture long-range state relations and improve planning in more complex or stochastic environments.
Abstract
Building deep learning models that can reason about their environment requires capturing its underlying dynamics. Joint-Embedded Predictive Architectures (JEPA) provide a promising framework to model such dynamics by learning representations and predictors through a self-supervised prediction objective. However, their ability to support effective action planning remains limited. We propose an approach to enhance planning with JEPA world models by shaping their representation space so that the negative goal-conditioned value function for a reaching cost in a given environment is approximated by a distance (or quasi-distance) between state embeddings. We introduce a practical method to enforce this constraint during training and show that it leads to significantly improved planning performance compared to standard JEPA models on simple control tasks.