TD-JEPA: Latent-predictive Representations for Zero-Shot Reinforcement Learning

Abstract: Latent prediction--where agents learn by predicting their own latents--has emerged as a powerful paradigm for training general representations in machine learning. In reinforcement learning (RL), this approach has been explored to define auxiliary losses for a variety of settings, including reward-based and unsupervised RL, behavior cloning, and world modeling. While existing methods are typically limited to single-task learning, one-step prediction, or on-policy trajectory data, we show that temporal difference (TD) learning enables learning representations predictive of long-term latent dynamics across multiple policies from offline, reward-free transitions. Building on this, we introduce TD-JEPA, which leverages TD-based latent-predictive representations into unsupervised RL. TD-JEPA trains explicit state and task encoders, a policy-conditioned multi-step predictor, and a set of parameterized policies directly in latent space. This enables zero-shot optimization of any reward function at test time. Theoretically, we show that an idealized variant of TD-JEPA avoids collapse with proper initialization, and learns encoders that capture a low-rank factorization of long-term policy dynamics, while the predictor recovers their successor features in latent space. Empirically, TD-JEPA matches or outperforms state-of-the-art baselines on locomotion, navigation, and manipulation tasks across 13 datasets in ExoRL and OGBench, especially in the challenging setting of zero-shot RL from pixels.

• The paper introduces TD-JEPA, a zero-shot RL framework that leverages latent-predictive representations to learn multi-step, policy-conditioned predictions.
• It employs an off-policy temporal-difference loss to pre-train state and task encoders, enhancing value estimation and policy generalization.
• Empirical evaluations across 65 tasks demonstrate that TD-JEPA consistently matches or exceeds state-of-the-art baselines, particularly in pixel-based environments.

TD-JEPA: Latent-Predictive Representations for Zero-Shot Reinforcement Learning

Introduction

The paper "TD-JEPA: Latent-predictive Representations for Zero-Shot Reinforcement Learning" (2510.00739) introduces a novel approach for zero-shot reinforcement learning (RL) by leveraging latent-predictive representations. A key challenge in RL is learning state representations that can effectively capture the dynamics of the environment. Such representations should facilitate efficient value estimation and policy optimization across various tasks. Previous work has utilized latent-predictive representation learning, where algorithms learn a state encoder $\phi(s)$ and a predictor $P$ that estimates the representation of a future state $s'$. This paradigm, known as joint-embedding predictive architecture (JEPA), operates entirely in latent space without rewards or state reconstructions.

Several RL methods have incorporated latent prediction as an auxiliary loss to enhance sample efficiency and generalization. In unsupervised settings, these losses provide a framework for learning latent world models that can solve goal-reaching tasks via test-time planning. The proposed method extends this approach by introducing policy-conditioned, multi-step formulations based on a novel off-policy temporal-difference loss. This objective focuses on representations that are predictive of long-term features critical for value estimation across multiple policies.

Methodology

TD-JEPA aims to instantiate temporal difference latent-predictive representation learning into a zero-shot unsupervised RL algorithm. The approach pre-trains four components: a state encoder, a policy-conditioned multi-step predictor, a task encoder, and a set of parameterized policies, learning end-to-end from offline, reward-free transitions. Unlike previous methods, latent prediction is the core objective, allowing TD-JEPA to learn the necessary components for distilling zero-shot policies.

The predictor approximates successor features, which enables extracting policies mapping encoded observations to optimal actions for any reward functions represented within the learned features' span. Empirically, TD-JEPA's architecture comprises state and task encoders that feed information into policy and action-value function predictors. The training leverages off-policy data, using a temporal-difference loss that updates state encoders and task representations separately.

Figure 1: TD-JEPA trains policies $\pi_z$ parameterized by latents $z$. The predictor, conditioned on $z$, predicts the representations of future states visited by $\pi_z$.

Theoretical Analysis

The paper provides theoretical guarantees for TD-JEPA's representation learning. In an idealized linear setting, the representations avoid collapse with appropriate initialization, recover a low-rank factorization of successor measures, and approximate successor features in latent space. These properties minimize policy evaluation error for any reward, rendering zero-shot optimization feasible.

The theoretical framework extends existing analyses by integrating a novel "gradient matching" argument, linking TD-JEPA with other unsupervised RL methodologies such as forward-backward learning and intention-conditioned value functions.

Empirical Evaluation

TD-JEPA was evaluated on 65 tasks across 13 datasets from ExoRL and OGBench, encompassing domains like locomotion, navigation, and manipulation with proprioceptive and pixel-based observations. Results indicate that TD-JEPA matches or surpasses state-of-the-art zero-shot baselines, particularly in challenging settings like pixel-based learning. The algorithm's architecture allows flexibility in learning distinct state and task representations, enhancing multi-step, policy-dependent dynamics prediction.

Figure 2: Probabilities of improvement: how likely is method X to outperform method Y on a random domain? We report symmetrized 95\% simple bootstrap confidence intervals.

Conclusion

The paper presents TD-JEPA as a robust zero-shot reinforcement learning framework that operates entirely in latent space, producing predictors that align with successor features. Empirical evidence and theoretical analysis show promising results in diverse settings, suggesting that TD-JEPA's latent-predictive methodology effectively integrates with unsupervised RL. Future work may explore asymmetric successor measures and large-scale robotic datasets to further expand the utility of latent-predictive objectives.