Value-guided action planning with JEPA world models
This addresses a specific bottleneck in JEPA models for better reasoning in environments, though it appears incremental as it builds on an existing framework.
The paper tackles the problem of limited action planning capability in Joint-Embedded Predictive Architectures (JEPA) world models by shaping their representation space to approximate a goal-conditioned value function with distances between state embeddings, resulting in significantly improved planning performance on simple control tasks.
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.