Value-guided action planning with JEPA world models

TL;DR

This paper enhances JEPA world models for better action planning by shaping their representation space to approximate goal-conditioned value functions, leading to improved planning performance in control tasks.

Contribution

It introduces a novel method to enforce value function approximation through representation shaping during training of JEPA models.

Findings

Significantly improved planning performance on control tasks.

Effective representation shaping aligns value functions with state embedding distances.

Method outperforms standard JEPA models in experiments.

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.