JEPA for RL: Investigating Joint-Embedding Predictive Architectures for Reinforcement Learning
This work applies a self-supervised learning method to reinforcement learning, which is incremental as it extends JEPA to a new domain.
The paper adapts Joint-Embedding Predictive Architectures (JEPA) to reinforcement learning from images, addressing model collapse and demonstrating results on the Cart Pole task.
Joint-Embedding Predictive Architectures (JEPA) have recently become popular as promising architectures for self-supervised learning. Vision transformers have been trained using JEPA to produce embeddings from images and videos, which have been shown to be highly suitable for downstream tasks like classification and segmentation. In this paper, we show how to adapt the JEPA architecture to reinforcement learning from images. We discuss model collapse, show how to prevent it, and provide exemplary data on the classical Cart Pole task.