---

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

Bagatella M, Pirotta M, Touati A, et al. TD-JEPA: Latent-predictive Representations for Zero-Shot Reinforcement Learning. arXiv preprint arXiv:2510.00739, 2025.

https://openreview.net/forum?id=SzXDuBN8M1

https://github.com/facebookresearch/td_jepa

RLZero-Shot RL40+ICLR 2026 OralCCF-AMeta FAIR苏黎世联邦理工学院马克斯·普朗克智能系统研究所

---

TD-JEPA：面向零样本强化学习的潜在预测表示

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.

潜在预测是一种让智能体通过预测自身潜变量进行学习的方法，已经成为机器学习中训练通用表示的一种强大范式。 在强化学习（RL）中，这种方法已被用于为多种设置定义辅助损失，包括基于奖励的强化学习和无监督强化学习、行为克隆以及世界建模。 现有方法通常局限于单任务学习、单步预测或在策略轨迹数据，而我们表明，时序差分（TD）学习能够从离线、无奖励的转移数据中学习可预测多个策略长期潜在动态的表示。 在此基础上，我们提出 TD-JEPA，将基于 TD 的潜在预测表示用于无监督强化学习。 TD-JEPA 直接在潜在空间中训练显式状态编码器和任务编码器、策略条件多步预测器以及一组参数化策略。 这使其能够在测试时对任意奖励函数进行零样本优化。 在理论上，我们表明，TD-JEPA 的一种理想化变体在适当初始化下能够避免坍缩，并学习捕获长期策略动态低秩分解的编码器，而预测器则在潜在空间中恢复其后继特征。 在经验上，TD-JEPA 在 ExoRL 和 OGBench 的 13 个数据集上，针对运动、导航和操作任务达到或超过最先进基线，尤其是在具有挑战性的基于像素的零样本强化学习设置中。

1. Introduction

Learning effective state representations is a core challenge in reinforcement learning (RL). Useful representations should capture the dynamics of the environment in a way that supports efficient value estimation and policy optimization across tasks. A promising line of work is latent-predictive (a.k.a. self-predictive) representation learning, an instance of the joint-embedding predictive architecture (JEPA) paradigm. These algorithms jointly learn a state encoder $\varphi (s)$ and a predictor $P$, i.e., a latent dynamics model estimating the representation of a future state $s^{\prime }$: $P(\varphi (s))\simeq \varphi (s^{\prime })$. Latent-predictive methods thus perform self-supervised learning entirely in latent space without any reward or reconstruction of (possibly high-dimensional) states.

学习有效的状态表示是强化学习（RL）中的核心挑战。 有用的表示应当捕获环境动态，从而支持跨任务的高效价值估计与策略优化。 一种很有前景的研究方向是潜在预测（亦称自预测）表示学习，它是联合嵌入预测架构（JEPA）范式的一种实例。 这些算法联合学习一个 状态编码器 $\varphi (s)$ 和一个 预测器 $P$；后者是估计未来状态 $s^{\prime }$ 表示的潜在动态模型：$P(\varphi (s))\simeq \varphi (s^{\prime })$。 因此，潜在预测方法完全在潜在空间中执行 自监督 学习，不需要任何奖励，也不需要重构可能具有高维度的状态。

Several RL methods leverage latent prediction as an auxiliary loss to improve sample efficiency and generalization in reward-based learning, behavior cloning, and curiosity-driven exploration. As latent-predictive losses do not require any reward, they have been recently used for unsupervised RL: Assran et al., Zhou et al. and Sobal et al. learn latent world models that can solve goal-reaching tasks via test-time planning, whereas Jajoo et al. learn a state encoder from trajectory data to define the space of tasks used to optimize zero-shot unsupervised policies.

若干强化学习方法将潜在预测用作辅助损失，以提高基于奖励的学习、行为克隆和好奇心驱动探索中的样本效率与泛化能力。 由于潜在预测损失不需要任何奖励，它们最近被用于无监督强化学习：Assran 等人、Zhou 等人和 Sobal 等人学习潜在世界模型，并通过测试时规划解决目标到达任务；Jajoo 等人则从轨迹数据中学习状态编码器，用它定义优化零样本无监督策略的任务空间。

图1：TD-JEPA 训练由潜变量 z 参数化的策略 π_z。预测器以 z 为条件，预测 π_z 访问的未来状态表示（左）。通过 TD 训练时，预测器（右侧箭头）近似每个策略的后继特征，即访问状态表示（圆）的加权重心（星形）。

This paper proposes a novel way to instantiate latent-predictive representations for unsupervised RL. While previous methods have largely focused on either one-step dynamics, single-task/single-policy training, or relied on on-policy data, we introduce a policy-conditioned, multi-step formulation based on a novel off-policy temporal-difference loss. This objective encourages representations that are predictive not only of immediate transitions, but also of long-term features relevant for value estimation across multiple policies. This property makes such representations and the associated predictors particularly well-suited for integration with off-policy, successor-feature based approaches to zero-shot unsupervised RL.

本文提出一种为无监督强化学习实例化潜在预测表示的新方法。 以往方法大多聚焦于单步动态、单任务或单策略训练，或者依赖在策略数据；我们则提出一种基于新型离策略时序差分损失的策略条件多步形式。 该目标鼓励表示不仅预测即时转移，还预测与多个策略的价值估计相关的长期特征。 这一性质使此类表示及其预测器尤其适合与基于后继特征的离策略零样本无监督强化学习方法结合。

We thus instantiate temporal difference latent-predictive representation learning into TD-JEPA, a zero-shot unsupervised RL algorithm which pre-trains four components: a state encoder, a policy-conditioned multi-step predictor, a task encoder, and a set of parameterized policies, all of which are learned end-to-end from offline, reward-free transitions. Departing from previous approaches, latent prediction is not merely an auxiliary loss, but rather the core objective that enables TD-JEPA to learn all the components needed to distill zero-shot policies. In fact, the predictor may be leveraged as an approximation of successor features (see Figure 1) to extract policies mapping encoded observations to optimal actions for all reward functions in the span of the learned features. This enables TD-JEPA to perform zero-shot policy optimization for any downstream reward, entirely in latent space.

因此，我们将时序差分潜在预测表示学习实例化为 TD-JEPA。TD-JEPA 是一种零样本无监督强化学习算法，它预训练四个组件：状态编码器、策略条件多步预测器、任务编码器和一组参数化策略；所有组件都从离线、无奖励的转移数据中进行端到端学习。 与以往方法不同，潜在预测不只是辅助损失，而是使 TD-JEPA 能够学习提取零样本策略所需全部组件的核心目标。 事实上，可以将预测器用作后继特征的近似（见 图1），从而提取把编码观测映射到最优动作的策略，并覆盖所学特征张成空间中的所有奖励函数。 这使 TD-JEPA 能够针对任意下游奖励执行零样本策略优化，而且 完全在潜在空间中完成。

Theoretically, for an idealized version of TD-JEPA with linear predictors, we show that 1) the representations do not collapse with a suitable initialization; 2) they recover a low-rank factorization of the successor measures of the trained policies, while the predictor approximates successor features in latent space; 3) they minimize an upper bound on the policy evaluation error for any reward, thus making zero-shot optimization possible. These results build on a novel "gradient matching" argument that extends and generalizes existing theoretical analyses of latent-predictive representations, and connect TD-JEPA with other unsupervised RL methods such as forward-backward and intention-conditioned value functions.

在理论上，对于采用线性预测器的 TD-JEPA 理想化版本，我们表明：1) 在适当初始化下，表示不会坍缩；2) 表示能够恢复所训练策略后继测度的低秩分解，而预测器在潜在空间中近似后继特征；3) 表示最小化任意奖励下策略评估误差的上界，从而使零样本优化成为可能。 这些结果基于一种新颖的“梯度匹配”论证，它扩展并推广了现有潜在预测表示理论分析，同时将 TD-JEPA 与前向-后向方法和意图条件价值函数等其他无监督强化学习方法联系起来。

Empirically, we evaluate TD-JEPA on 65 tasks across 13 datasets from ExoRL and OGBench, covering locomotion, navigation, and manipulation with both proprioceptive and pixel-based observations. TD-JEPA matches or outperforms state-of-the-art zero-shot baselines across these settings, in particular when learning from pixels, which has proven to be one of the most challenging settings for unsupervised RL so far. Moreover, we ablate several dimensions of the algorithm, demonstrating the importance of learning representations that are predictive of multi-step policy-dependent dynamics, and the advantage of training distinct state and task encoders. Finally, we show that learned representations can be easily reused for offline or online RL, improving over zero-shot policies and learning from scratch.

在经验上，我们在 ExoRL 和 OGBench 的 13 个数据集、65 个任务上评估 TD-JEPA，涵盖采用本体感知观测和像素观测的运动、导航与操作任务。 TD-JEPA 在这些设置中达到或超过最先进的零样本基线，尤其是在从像素学习时；这已被证明是迄今无监督强化学习中最具挑战性的设置之一。 此外，我们对算法的多个维度进行消融，证明学习能够预测多步策略依赖动态的表示非常重要，并展示分别训练状态编码器和任务编码器的优势。 最后，我们表明，学到的表示可以轻松复用于离线或在线强化学习，相比零样本策略和从头学习均有所改进。

2. Preliminaries

We consider a reward-free Markov Decision Process $\mathcal{M}=(\mathcal{S},\mathcal{A},P,\gamma )$, where $\mathcal{S}$ and $\mathcal{A}$ are state and action spaces, $P$ is the probability measure over next states when taking action $a$ in state $s$ as $P(\mathrm{d}{s}^{\prime }\mid s,a)$, and $\gamma \in [0,1)$ is a discount factor. Executing a Markov policy $\pi :\mathcal{S}\to \mathrm{Prob}(\mathcal{A})$ induces an unnormalized distribution over visited states, which is referred to as the successor measure:

我们考虑一个无奖励马尔可夫决策过程 $\mathcal{M}=(\mathcal{S},\mathcal{A},P,\gamma )$，其中 $\mathcal{S}$ 和 $\mathcal{A}$ 分别是状态空间和动作空间，$P$ 是在状态 $s$ 下采取动作 $a$ 时下一状态的概率测度，记为 $P(\mathrm{d}{s}^{\prime }\mid s,a)$，而 $\gamma \in [0,1)$ 是折扣因子。 执行马尔可夫策略 $\pi :\mathcal{S}\to \mathrm{Prob}(\mathcal{A})$ 会在访问状态上诱导一个非归一化分布，称为 后继测度：

$$M^{\pi }(\mathcal{X}\mid s,a)=\sum _{t=0}^{\mathrm{\infty }}{\gamma }^{t}Pr(s_{t+1}\in \mathcal{X}\mid s,a,\pi ),\mathrm{\forall }\mathcal{X}\subseteq \mathcal{S}.$$

Given a reward function $r:\mathcal{S}\to \mathbb{R}$ and a policy $\pi$, the action-value function $Q_r^{\pi }(s,a)$ measures the cumulative discounted reward obtained by the policy over an infinite horizon, i.e., $Q_r^{\pi }(s,a)=\mathbb{E}[\sum _{t=0}^{\mathrm{\infty }}{\gamma }^{t}r(s_{t+1})\mid s,a,\pi ]$. Action-value functions are connected to successor measures via

给定奖励函数 $r:\mathcal{S}\to \mathbb{R}$ 和策略 $\pi$，动作价值函数 $Q_r^{\pi }(s,a)$ 衡量该策略在无限时域中获得的累计折扣奖励，即 $Q_r^{\pi }(s,a)=\mathbb{E}[\sum _{t=0}^{\mathrm{\infty }}{\gamma }^{t}r(s_{t+1})\mid s,a,\pi ]$。 动作价值函数通过下式与后继测度相联系：

$$\begin{array}{rl}{Q}_r^{\pi }(s,a)& ={\int }_{s^{+}\in \mathcal{S}}{M}^{\pi }(\mathrm{d}{s}^{+}\mid s,a)r(s^{+})\\ & ={\mathbb{E}}_{s^{+}\sim M^{\pi }(\cdot \mid s,a)}[r(s^{+})].\end{array}$$

This shows a convenient linear decomposition of $Q_r^{\pi }$ into the reward function and the dynamics induced by $\pi$. Standard RL agents aim at finding reward-maximizing policies ${\pi }_r^{\star }(s)\in \arg \underset{a\in \mathcal{A}}{max}{Q}_r^{\star }(s,a)$, where $Q_r^{\star }(s,a):=\underset{\pi }{max}{Q}_r^{\pi }(s,a)$.

这给出了 $Q_r^{\pi }$ 关于奖励函数和策略 $\pi$ 所诱导动态的一种便利线性分解。 标准强化学习智能体旨在寻找奖励最大化策略 ${\pi }_r^{\star }(s)\in \arg \underset{a\in \mathcal{A}}{max}{Q}_r^{\star }(s,a)$，其中 $Q_r^{\star }(s,a):=\underset{\pi }{max}{Q}_r^{\pi }(s,a)$。

Latent-predictive representations. In high-dimensional settings, state encoders $\varphi :\mathcal{S}\to {\mathbb{R}}^{d_{\varphi }}$ may be learned to ease the estimation of action-value functions. For instance, if an encoder $\varphi$ is such that $Q_r^{\pi }(s,a)=\varphi (s{)}^{\mathrm{\top }}{w}_{a,r}^{\pi }$ for some vector $w_{a,r}^{\pi }\in {\mathbb{R}}^{d_{\varphi }}$, then the RL process reduces to learning vectors in ${\mathbb{R}}^{d_{\varphi }}$ rather than high-dimensional functions $Q_r^{\pi }(s,a)$. Latent-predictive learning has been shown to be an effective approach for this problem. In the simplest formulation, latent-predictive representations capture the one-step latent dynamics of a policy $\pi$ by minimizing the loss

潜在预测表示。 在高维设置中，可以学习 状态编码器 $\varphi :\mathcal{S}\to {\mathbb{R}}^{d_{\varphi }}$ 来简化动作价值函数的估计。 例如，若编码器 $\varphi$ 满足 $Q_r^{\pi }(s,a)=\varphi (s{)}^{\mathrm{\top }}{w}_{a,r}^{\pi }$，其中某个向量 $w_{a,r}^{\pi }\in {\mathbb{R}}^{d_{\varphi }}$，那么强化学习过程就从学习高维函数 $Q_r^{\pi }(s,a)$ 化简为学习 ${\mathbb{R}}^{d_{\varphi }}$ 中的向量。 潜在预测学习已被证明是解决这一问题的有效方法。 在最简单的形式中，潜在预测表示通过最小化下列损失来捕获策略 $\pi$ 的单步潜在动态：

$${\mathcal{L}}_{\text{one-step}}(\varphi ,T)={\mathbb{E}}_{\begin{array}{c}s\sim \rho ,a\sim \pi (\cdot \mid s)\\ s^{\prime }\sim P(\cdot \mid s,a)\end{array}}\left[{\Vert T(\varphi (s))-\mathrm{sg}(\varphi (s^{\prime }))\Vert }^2\right].$$

Here, $T:{\mathbb{R}}^{d_{\varphi }}\to {\mathbb{R}}^{d_{\varphi }}$ is a (possibly non-linear) predictor of the latent one-step dynamics induced by $\varphi$ and policy $\pi$, and $\mathrm{sg}(\varphi )$ denotes stop-gradient. Notably, optimizing for this loss does not require any decoding or reconstruction, and it only relies on an unsupervised dataset $\mathcal{D}=\{(s,a,s^{\prime })\}$. Different instantiations of this approach have been shown both empirically and theoretically to produce representations that accurately approximate action-value functions or policies.

其中，$T:{\mathbb{R}}^{d_{\varphi }}\to {\mathbb{R}}^{d_{\varphi }}$ 是由 $\varphi$ 和策略 $\pi$ 诱导的潜在单步动态预测器，它可以是非线性的；$\mathrm{sg}(\varphi )$ 表示停止梯度。 值得注意的是，优化该损失不需要任何解码或重构，并且只依赖一个 无监督 数据集 $\mathcal{D}=\{(s,a,s^{\prime })\}$。 这一方法的不同实例在经验和理论上均被证明可以产生能够准确近似动作价值函数或策略的表示。

Successor-features and zero-shot unsupervised RL. Considering a state encoder $\psi :\mathcal{S}\to {\mathbb{R}}^{d_{\psi }}$ and the associated space of linear rewards ${\mathcal{R}}_{\psi }=\{r(s)=\psi (s{)}^{\mathrm{\top }}z\mid z\in {\mathbb{R}}^{d_{\psi }}\}$, Q-values for any reward function $r(s)=\psi (s{)}^{\mathrm{\top }}{z}_r\in {\mathcal{R}}_{\psi }$ can be written as

后继特征与零样本无监督强化学习。 考虑状态编码器 $\psi :\mathcal{S}\to {\mathbb{R}}^{d_{\psi }}$ 及其对应的线性奖励空间 ${\mathcal{R}}_{\psi }=\{r(s)=\psi (s{)}^{\mathrm{\top }}z\mid z\in {\mathbb{R}}^{d_{\psi }}\}$，任意奖励函数 $r(s)=\psi (s{)}^{\mathrm{\top }}{z}_r\in {\mathcal{R}}_{\psi }$ 的 Q 值可以写为

$$\begin{array}{rl}{Q}_r^{\pi }(s,a)& ={\int }_{s^{+}\in \mathcal{S}}{M}^{\pi }(\mathrm{d}{s}^{+}\mid s,a)\psi (s^{+}{)}^{\mathrm{\top }}{z}_r\\ & ={\mathbb{E}}_{s^{+}\sim M^{\pi }(\cdot \mid s,a)}[\psi (s^{+}){]}^{\mathrm{\top }}{z}_r\\ & :=F_{\psi }^{\pi }(s,a{)}^{\mathrm{\top }}{z}_r.\end{array}$$

Here, $F_{\psi }^{\pi }(s,a)\in {\mathbb{R}}^{d_{\psi }}$ captures the successor features of $\pi$. The majority of unsupervised zero-shot RL methods learn successor features $F(s,a;z)\approx F_{\psi }^{{\pi }_z}(s,a)$ for a set of parameterized policies $\{{\pi }_z(s){\}}_{z\in \mathcal{Z}}$, with $\mathcal{Z}\subseteq {\mathbb{R}}^d$, that are trained to be optimal for all rewards in ${\mathcal{R}}_{\psi }$, i.e., ${\pi }_z(s)\approx \arg \underset{a}{max}F(s,a;z{)}^{\mathrm{\top }}z$, where $F(s,a;z{)}^{\mathrm{\top }}z$ is an approximation of $Q_r^{\star }(s,a)$ for $r(s)=\psi (s{)}^{\mathrm{\top }}z$. At test time, given a reward function $r$, a vector $z_r\in {\mathbb{R}}^{d_{\psi }}$ is first obtained by projecting $r$ onto ${\mathcal{R}}_{\psi }$, and the associated policy ${\pi }_{z_r}$ is then returned.

其中，$F_{\psi }^{\pi }(s,a)\in {\mathbb{R}}^{d_{\psi }}$ 捕获策略 $\pi$ 的 后继特征。 大多数无监督零样本强化学习方法为一组参数化策略 $\{{\pi }_z(s){\}}_{z\in \mathcal{Z}}$ 学习后继特征 $F(s,a;z)\approx F_{\psi }^{{\pi }_z}(s,a)$，其中 $\mathcal{Z}\subseteq {\mathbb{R}}^d$；这些策略被训练为对 ${\mathcal{R}}_{\psi }$ 中所有奖励均最优，即 ${\pi }_z(s)\approx \arg \underset{a}{max}F(s,a;z{)}^{\mathrm{\top }}z$，而当 $r(s)=\psi (s{)}^{\mathrm{\top }}z$ 时，$F(s,a;z{)}^{\mathrm{\top }}z$ 是 $Q_r^{\star }(s,a)$ 的近似。 在测试时，给定奖励函数 $r$，首先通过将 $r$ 投影到 ${\mathcal{R}}_{\psi }$ 上得到向量 $z_r\in {\mathbb{R}}^{d_{\psi }}$，随后返回对应策略 ${\pi }_{z_r}$。

Given the role played by $\psi$ in defining the space of tasks of interest, with an abuse of terminology, we will refer to $\psi$ as a task encoder. On the other hand, we shall call state encoder a map $\varphi :\mathcal{S}\to {\mathbb{R}}^{d_{\varphi }}$ that is used to embed states before feeding them into different networks (e.g., we will train successor features $F_{\psi }^{\pi }(\varphi (s),a)$ and policies $\pi (\varphi (s))$ in the latent space given by $\varphi$). While the zero-shot methods cited so far train the task encoder $\psi$ in different ways, and do not train any explicit state encoder $\varphi$, the next section will show how multi-step policy-dependent latent-predictive learning can be used to train both simultaneously.

鉴于 $\psi$ 在定义关注的任务空间中所起的作用，我们略微滥用术语，将 $\psi$ 称为 任务编码器。 另一方面，我们把映射 $\varphi :\mathcal{S}\to {\mathbb{R}}^{d_{\varphi }}$ 称为 状态编码器；它在状态输入不同网络之前对其进行嵌入，例如，我们将在 $\varphi$ 给出的潜在空间中训练后继特征 $F_{\psi }^{\pi }(\varphi (s),a)$ 和策略 $\pi (\varphi (s))$。 目前提到的零样本方法以不同方式训练任务编码器 $\psi$，但不训练任何显式状态编码器 $\varphi$；下一节将展示如何使用多步、策略依赖的潜在预测学习同时训练二者。

3. Latent-Predictive Temporal-Difference Representations

We begin by showing how the latent-predictive loss of Equation (3) can model multi-step and policy-dependent dynamics, and how temporal difference (TD) learning allows learning from offline transition data. We will then expand this idea to learn separate state and task embeddings, and finally show how it can be instantiated as a zero-shot unsupervised RL method.

我们首先展示公式（3）的潜在预测损失如何建模多步、策略依赖的动态，以及时序差分（TD）学习如何允许从离线转移数据中学习。 随后，我们扩展这一思想以学习分离的状态嵌入和任务嵌入，最后展示如何将其实现为零样本无监督强化学习方法。

3.1 Multi-step policy-conditioned latent prediction

Let $\{{\pi }_z{\}}_{z\in \mathcal{Z}}$ be a family of policies parameterized by $z\in \mathcal{Z}$, and $\mathcal{D}=\{(s,a,s^{\prime })\}$ be a dataset of transitions. We train a state encoder $\varphi :\mathcal{S}\to {\mathbb{R}}^{d_{\varphi }}$ and a policy-dependent predictor $T_{\varphi }:{\mathbb{R}}^{d_{\varphi }}\times \mathcal{A}\times \mathcal{Z}\to {\mathbb{R}}^{d_{\varphi }}$ to be latent-predictive of the long-term dynamics of the policies $\{{\pi }_z\}$, i.e.,

令 $\{{\pi }_z{\}}_{z\in \mathcal{Z}}$ 表示一族由 $z\in \mathcal{Z}$ 参数化的策略，并令 $\mathcal{D}=\{(s,a,s^{\prime })\}$ 表示转移数据集。 我们训练状态编码器 $\varphi :\mathcal{S}\to {\mathbb{R}}^{d_{\varphi }}$ 和 策略依赖 预测器 $T_{\varphi }:{\mathbb{R}}^{d_{\varphi }}\times \mathcal{A}\times \mathcal{Z}\to {\mathbb{R}}^{d_{\varphi }}$，使其对策略 $\{{\pi }_z\}$ 的 长期 动态进行潜在预测，即

$$\begin{array}{rl}{\mathcal{L}}_{\text{MC-JEPA}}(\varphi ,T_{\varphi })& ={\mathbb{E}}_{\begin{array}{c}(s,a)\sim \mathcal{D},z\sim \mathcal{Z}\\ s^{+}\sim M^{{\pi }_z}(\cdot \mid s,a)\end{array}}\\ & \left[{\Vert T_{\varphi }(\varphi (s),a,z)-\mathrm{sg}(\varphi (s^{+}))\Vert }^2\right].\end{array}$$

Here, MC-JEPA stands for Monte-Carlo (MC) JEPA loss, as on-policy samples $s^{+}\sim M^{{\pi }_z}(\cdot \mid s,a)$ are needed for all policies of interest. Intuitively, $T_{\varphi }(\varphi (s),a,z)$ tries to predict future latent states visited by the policy ${\pi }_z$. More formally, predictors trained via minimization of ${\mathcal{L}}_{\text{MC-JEPA}}(\varphi ,T_{\varphi })$ approximate the successor features of $\varphi$ in the latent space induced by $\varphi$ itself.

这里，MC-JEPA 表示蒙特卡洛（MC）JEPA 损失，因为所有关注策略都需要在策略样本 $s^{+}\sim M^{{\pi }_z}(\cdot \mid s,a)$。 直观而言，$T_{\varphi }(\varphi (s),a,z)$ 试图预测策略 ${\pi }_z$ 访问的未来潜在状态。 更形式化地说，通过最小化 ${\mathcal{L}}_{\text{MC-JEPA}}(\varphi ,T_{\varphi })$ 训练的预测器，在 $\varphi$ 自身诱导的潜在空间中近似 $\varphi$ 的后继特征。

Proposition 1. For any $\varphi$ and $T_{\varphi }$, we have the following equivalence.

命题1。 对任意 $\varphi$ 和 $T_{\varphi }$，都有下列等价关系。

$$\begin{array}{rl}{\mathcal{L}}_{\text{MC-JEPA}}(\varphi ,T_{\varphi })& ={\mathbb{E}}_{(s,a)\sim \mathcal{D},z\sim \mathcal{Z}}\\ & \left[{\Vert T_{\varphi }(\varphi (s),a,z)-\mathrm{sg}(F_{\varphi }^{{\pi }_z}(s,a))\Vert }^2\right]+\mathrm{const}.\end{array}$$

Given the connection between Q-functions and successor features (Equation (4)), this result crucially relates multi-step latent prediction with value estimation across multiple policies. More precisely, it implies that the predictor enables policy evaluation and optimization of rewards in the span of $\varphi$, as we detail at the end of this section. Since $F_{\varphi }^{{\pi }_z}$ is the successor features of $\varphi$, with the terminology introduced in Section 2, $\varphi$ is used both as a state encoder, i.e., to embed states passed to the predictor, and as a task encoder, i.e., defining a space of reward functions.

由于 Q 函数与后继特征之间存在联系（公式（4）），该结果将多步潜在预测与多个策略的价值估计关键地联系起来。 更准确地说，它意味着预测器能够在 $\varphi$ 的张成空间中进行策略评估和奖励优化，我们将在本节末尾详述。 由于 $F_{\varphi }^{{\pi }_z}$ 是 $\varphi$ 的后继特征，按照第2节引入的术语，$\varphi$ 同时用作 状态编码器（即嵌入传给预测器的状态）和 任务编码器（即定义奖励函数空间）。

Unfortunately, this loss cannot be estimated on off-policy data since it requires sampling from the successor measures of the given policies. We can however leverage the previous result and the fact that successor features admit a Bellman equation $F_{\varphi }^{{\pi }_z}(s,a)={\mathbb{E}}_{s^{\prime }\sim P(\cdot \mid s,a),a^{\prime }\sim {\pi }_z(s^{\prime })}[\varphi (s^{\prime })+\gamma F_{\varphi }^{{\pi }_z}(s^{\prime },a^{\prime })]$ to define a temporal-difference version of the previous loss:

遗憾的是，该损失无法在离策略数据上估计，因为它需要从给定策略的后继测度中采样。 不过，我们可以利用前述结果以及后继特征满足贝尔曼方程 $F_{\varphi }^{{\pi }_z}(s,a)={\mathbb{E}}_{s^{\prime }\sim P(\cdot \mid s,a),a^{\prime }\sim {\pi }_z(s^{\prime })}[\varphi (s^{\prime })+\gamma F_{\varphi }^{{\pi }_z}(s^{\prime },a^{\prime })]$ 这一事实，定义前述损失的时序差分版本：

$$\begin{array}{rl}{\mathcal{L}}_{\text{TD-JEPA}}(\varphi ,T_{\varphi })& ={\mathbb{E}}_{\begin{array}{c}(s,a,s^{\prime })\sim \mathcal{D},z\sim \mathcal{Z}\\ a^{\prime }\sim {\pi }_z(\cdot \mid s^{\prime })\end{array}}\\ & \left[{\Vert \begin{array}{rl}{T}_{\varphi }(\varphi (s),a,z)& -\mathrm{sg}(\varphi (s^{\prime }))\\ & -\gamma \mathrm{sg}(T_{\varphi }(\varphi (s^{\prime }),a^{\prime },z))\end{array}\Vert }^2\right].\end{array}$$

Unlike the Monte Carlo loss of Equation (5), ${\mathcal{L}}_{\text{TD-JEPA}}$ only requires sampling one-step transitions and actions from the given policies, and it can thus be estimated from off-policy, offline datasets.

不同于公式（5）的蒙特卡洛损失，${\mathcal{L}}_{\text{TD-JEPA}}$ 只需要从给定策略采样单步转移和动作，因此可以从离策略离线数据集中估计。

3.2 Training separate state and task representations

While in Equations (5) and (7) the same encoder $\varphi$ is used for both state and task representations, these need not be the same in practice. Consider, for instance, a robot navigating a building: useful state representations may capture low-level dynamical information critical for control (e.g., joint positions and velocities), while task representations could abstract higher-level contextual features, such as the building's topology. In this case, a single representation might be either too complex, or too abstract: having flexibility over the dimensionality and content of each representation would be desirable.

尽管公式（5）和（7）对状态表示与任务表示使用相同编码器 $\varphi$，但实践中二者不必相同。 例如，考虑一个在建筑物中导航的机器人：有用的状态表示可能捕获对控制至关重要的低层动态信息（如关节位置和速度），而任务表示可以抽象更高层的上下文特征，例如建筑物的拓扑结构。 在这种情况下，单一表示可能过于复杂，也可能过于抽象；能够灵活控制每种表示的维度和内容会更为理想。

We thus now introduce an asymmetric variant that trains a distinct encoder $\psi :\mathcal{S}\to {\mathbb{R}}^{d_{\psi }}$ to define the set of reward functions of interest (i.e., as a task encoder). We first redefine the predictor as $T_{\varphi }:{\mathbb{R}}^{d_{\varphi }}\times \mathcal{A}\times \mathcal{Z}\to {\mathbb{R}}^{d_{\psi }}$ and the latent-predictive Monte-Carlo loss to train $\varphi$ and $T_{\varphi }$ as

因此，我们引入一种非对称变体，训练一个独立编码器 $\psi :\mathcal{S}\to {\mathbb{R}}^{d_{\psi }}$ 来定义所关注的奖励函数集合，即把它用作 任务 编码器。 我们首先把预测器重新定义为 $T_{\varphi }:{\mathbb{R}}^{d_{\varphi }}\times \mathcal{A}\times \mathcal{Z}\to {\mathbb{R}}^{d_{\psi }}$，并将用于训练 $\varphi$ 和 $T_{\varphi }$ 的潜在预测蒙特卡洛损失重新定义为

$${\mathcal{L}}_{\text{MC-JEPA}}(\varphi ,T_{\varphi },\psi )={\mathbb{E}}_{\begin{array}{c}(s,a)\sim \mathcal{D},z\sim \mathcal{Z}\\ s^{+}\sim M^{{\pi }_z}(\cdot \mid s,a)\end{array}}\left[{\Vert T_{\varphi }(\varphi (s),a,z)-\mathrm{sg}(\psi (s^{+}))\Vert }^2\right].$$

This makes $T_{\varphi }$ map states encoded through $\varphi$ to the long-term dynamics of a policy ${\pi }_z$ in the latent space induced, this time, by $\psi$. Similar to Proposition 1, $T_{\varphi }$ approximates the successor features $F_{\psi }^{{\pi }_z}(s,a)$ of $\psi$ in the latent space induced by $\varphi$. Symmetrically, we train $\psi$ together with an additional predictor $T_{\psi }:{\mathbb{R}}^{d_{\psi }}\times \mathcal{A}\times \mathcal{Z}\to {\mathbb{R}}^{d_{\varphi }}$. To do so, we follow existing literature--according to which joint representations should be predictive of each other--and train $\psi$ and $T_{\psi }$ through the same latent-predictive loss with the roles of $\varphi$ and $\psi$ inverted, i.e., ${\mathcal{L}}_{\text{MC-JEPA}}(\psi ,T_{\psi },\varphi )$. While some existing works use forward-in-time sampling to train one representation and backward-in-time for the other, we use two forward-in-time losses. We further discuss this difference in the appendix.

这使 $T_{\varphi }$ 将经 $\varphi$ 编码的状态映射到策略 ${\pi }_z$ 的长期动态，而这一次对应的潜在空间由 $\psi$ 诱导。 与命题1类似，$T_{\varphi }$ 在 $\varphi$ 诱导的潜在空间中近似 $\psi$ 的后继特征 $F_{\psi }^{{\pi }_z}(s,a)$。 对称地，我们将 $\psi$ 与附加预测器 $T_{\psi }:{\mathbb{R}}^{d_{\psi }}\times \mathcal{A}\times \mathcal{Z}\to {\mathbb{R}}^{d_{\varphi }}$ 一起训练。 为此，我们遵循联合表示应当彼此可预测的既有研究，并通过交换 $\varphi$ 与 $\psi$ 角色后的相同潜在预测损失训练 $\psi$ 和 $T_{\psi }$，即 ${\mathcal{L}}_{\text{MC-JEPA}}(\psi ,T_{\psi },\varphi )$。 一些现有工作使用时间正向采样训练一个表示、时间反向采样训练另一个表示，而我们使用两个时间正向损失。 我们在附录中进一步讨论这一差异。

As before, we can then design an off-policy TD variant of this loss,

与之前一样，我们可以进一步设计该损失的离策略 TD 变体：

$$\begin{array}{rl}{\mathcal{L}}_{\text{TD-JEPA}}(\varphi ,T_{\varphi },\psi )& ={\mathbb{E}}_{\begin{array}{c}(s,a,s^{\prime })\sim \mathcal{D}\\ z\sim \mathcal{Z},a^{\prime }\sim {\pi }_z(\cdot \mid s^{\prime })\end{array}}\\ & \left[{\Vert \begin{array}{rl}{T}_{\varphi }(\varphi (s),a,z)& -\mathrm{sg}(\psi (s^{\prime }))\\ & -\gamma \mathrm{sg}(T_{\varphi }(\varphi (s^{\prime }),a^{\prime },z))\end{array}\Vert }^2\right].\end{array}$$

Thus, $\varphi$ and $T_{\varphi }$ are optimized via ${\mathcal{L}}_{\text{TD-JEPA}}(\varphi ,T_{\varphi },\psi )$, while $\psi$ and $T_{\psi }$ are optimized via ${\mathcal{L}}_{\text{TD-JEPA}}(\psi ,T_{\psi },\varphi )$.

因此，$\varphi$ 和 $T_{\varphi }$ 通过 ${\mathcal{L}}_{\text{TD-JEPA}}(\varphi ,T_{\varphi },\psi )$ 优化，而 $\psi$ 和 $T_{\psi }$ 通过 ${\mathcal{L}}_{\text{TD-JEPA}}(\psi ,T_{\psi },\varphi )$ 优化。

3.3 TD-JEPA representations for zero-shot RL

The relationship between the learned predictors and successor features suggests a seamless instantiation of TD-JEPA as a zero-shot unsupervised RL algorithm. Redefining the policy parameter space $\mathcal{Z}$ as the task embedding space (i.e., $\mathcal{Z}\subseteq {\mathbb{R}}^{d_{\psi }}$), we train latent policies such that

所学预测器与后继特征之间的关系表明，可以自然地将 TD-JEPA 实例化为零样本无监督强化学习算法。 我们将策略参数空间 $\mathcal{Z}$ 重新定义为任务嵌入空间（即 $\mathcal{Z}\subseteq {\mathbb{R}}^{d_{\psi }}$），并训练满足下式的潜在策略：

$${\pi }_z(\varphi (s))=\arg \underset{a}{max}{T}_{\varphi }(\varphi (s),z,a{)}^{\mathrm{\top }}z,\mathrm{\forall }z\in \mathcal{Z}.$$

This decision additionally grounds $\psi$ as task encoder, and breaks the symmetry that could arise from the two encoders $\varphi$ and $\psi$ being trained through similar latent-predictive objectives. Since $T_{\varphi }(\varphi (s),z,a)\simeq F_{\psi }^{{\pi }_z}(s,a)$ (Proposition 1), this produces optimal policies for all rewards in the span of $\psi$, learned directly from state representations $\varphi (\cdot )$. At test time, given an inference dataset of rewarded samples ${\mathcal{D}}_{\mathrm{rwd}}=\{(s,r)\}$, the optimal policy ${\pi }_{z_r}$ can be retrieved by computing $z_r$ through linear regression, e.g. through the closed-form solution

这一选择还进一步确立了 $\psi$ 作为任务编码器的作用，并打破了两个编码器 $\varphi$ 和 $\psi$ 通过相似潜在预测目标训练时可能产生的对称性。 由于 $T_{\varphi }(\varphi (s),z,a)\simeq F_{\psi }^{{\pi }_z}(s,a)$（命题1），它可以直接从状态表示 $\varphi (\cdot )$ 学到对 $\psi$ 张成空间中所有奖励均最优的策略。 测试时，给定带奖励样本的推理数据集 ${\mathcal{D}}_{\mathrm{rwd}}=\{(s,r)\}$，可以通过线性回归计算 $z_r$，进而检索最优策略 ${\pi }_{z_r}$；例如使用下列闭式解：

$$\begin{array}{rl}{z}_r& =\arg \underset{z}{min}{\mathbb{E}}_{(s,r)\sim {\mathcal{D}}_{\mathrm{rwd}}}[(r-\psi (s{)}^{\mathrm{\top }}z{)}^2]\\ & ={\mathbb{E}}_{s\sim {\mathcal{D}}_{\mathrm{rwd}}}[\psi (s)\psi (s{)}^{\mathrm{\top }}{]}^{-1}{\mathbb{E}}_{(s,r)\sim {\mathcal{D}}_{\mathrm{rwd}}}[\psi (s)r(s)].\end{array}$$

Algorithm Algorithm 1 describes TD-JEPA, which combines ${\mathcal{L}}_{\text{TD-JEPA}}$ with stabilization strategies, e.g. target networks and covariance regularization. We remark that latent prediction is not auxiliary: it is the core objective that trains encoders and predictors, from which zero-shot policies can be directly distilled.

算法1 描述了 TD-JEPA，它将 ${\mathcal{L}}_{\text{TD-JEPA}}$ 与目标网络、协方差正则化等稳定化策略结合。 我们强调，潜在预测并非辅助目标：它是训练编码器和预测器的核心目标，零样本策略可以从这些组件中直接提取。

Algorithm 1: TD-JEPA for zero-shot RL

Inputs: Dataset $\mathcal{D}$, batch size $B$, regularization coefficient $\lambda$, networks $\pi$, $T_{\varphi }$, $\varphi$, $T_{\psi }$, $\psi$.

1. Initialize target networks: $T_{\varphi }^{-}\leftarrow T_{\varphi }$, ${\varphi }^{-}\leftarrow \varphi$, $T_{\psi }^{-}\leftarrow T_{\psi }$, ${\psi }^{-}\leftarrow \psi$.
2. while not converged do
3.   // Sample training batch
4.   $\{(s_i,a_i,s_i^{\prime }){\}}_{i=1}^B\sim \mathcal{D}$, $\{z_i{\}}_{i=1}^B\sim \mathcal{Z}$, $\{a_i^{\prime }{\}}_{i=1}^B\sim \{\mathrm{sg}(\pi ({\varphi }^{-}(s_i^{\prime }),z_i)){\}}_{i=1}^B$.
5.   // Compute latent-predictive losses
6.   ${\hat{\mathcal{L}}}_{\text{TD-JEPA}}(\varphi ,T_{\varphi },\psi )=\frac{1}{2B}\sum _i\parallel T_{\varphi }(\varphi (s_i),a_i,z_i)-\mathrm{sg}({\psi }^{-}(s_i^{\prime }))-\gamma \mathrm{sg}(T_{\varphi }^{-}({\varphi }^{-}(s_i^{\prime }),a_i^{\prime },z_i)){\parallel }^2$.
7.   ${\hat{\mathcal{L}}}_{\text{TD-JEPA}}(\psi ,T_{\psi },\varphi )=\frac{1}{2B}\sum _i\parallel T_{\psi }(\psi (s_i),a_i,z_i)-\mathrm{sg}({\varphi }^{-}(s_i^{\prime }))-\gamma \mathrm{sg}(T_{\psi }^{-}({\psi }^{-}(s_i^{\prime }),a_i^{\prime },z_i)){\parallel }^2$.
8.   // Compute orthonormality regularization losses
9.   ${\hat{\mathcal{L}}}_{\mathrm{REG}}(\varphi )=\frac{1}{2B(B-1)}\sum _{i\ne j}(\varphi (s_i{)}^{\mathrm{\top }}\varphi (s_j){)}^2-\frac{1}{B}\sum _i\varphi (s_i{)}^{\mathrm{\top }}\varphi (s_i)$.
10.   ${\hat{\mathcal{L}}}_{\mathrm{REG}}(\psi )=\frac{1}{2B(B-1)}\sum _{i\ne j}(\psi (s_i{)}^{\mathrm{\top }}\psi (s_j){)}^2-\frac{1}{B}\sum _i\psi (s_i{)}^{\mathrm{\top }}\psi (s_i)$.
11.   // Compute actor loss
12.   $\{{\hat{a}}_i{\}}_{i=1}^B\sim \{\pi (\varphi (s_i),z_i){\}}_{i=1}^B$.
13.   ${\hat{\mathcal{L}}}_{\mathrm{actor}}(\pi )=-\frac{1}{B}\sum _{i=1}^B{T}_{\varphi }(\varphi (s_i),{\hat{a}}_i,z_i{)}^{\mathrm{\top }}{z}_i$.
14.   Update $\varphi$, $T_{\varphi }$ to minimize ${\hat{\mathcal{L}}}_{\text{TD-JEPA}}(\varphi ,T_{\varphi },\psi )+\lambda {\hat{\mathcal{L}}}_{\mathrm{REG}}(\varphi )$.
15.   Update $\psi$, $T_{\psi }$ to minimize ${\hat{\mathcal{L}}}_{\text{TD-JEPA}}(\psi ,T_{\psi },\varphi )+\lambda {\hat{\mathcal{L}}}_{\mathrm{REG}}(\psi )$.
16.   Update $\pi$ to minimize ${\hat{\mathcal{L}}}_{\mathrm{actor}}(\pi )$.
17.   Update target networks ${\varphi }^{-}$, $T_{\varphi }^{-}$, ${\psi }^{-}$, $T_{\psi }^{-}$ via EMA of $\varphi$, $T_{\varphi }$, $\psi$, $T_{\psi }$.
18. end while

4. Theoretical Analysis

We now provide some theoretical arguments showing how latent-predictive temporal difference representations capture the long-term dynamics of a given set of policies in a way that makes them amenable to zero-shot RL. Following Tang et al., we consider a simplified tabular setting with linear predictors. We view the representation $\varphi$ (resp. $\psi$) as a $S\times d_{\varphi }$ (resp. $S\times d_{\psi }$) matrix, and consider action-free predictors $T_{\varphi ,z}$ (resp. $T_{\psi ,z}$) as $d_{\varphi }\times d_{\psi }$ (resp. $d_{\psi }\times d_{\varphi }$) matrices for all $z$. The expression $T_{\varphi }(\varphi (s),a,z)$ in Equations (8) and (9) thus reduces to $T_{\varphi ,z}^{\mathrm{\top }}\varphi (s)$, while $M^{\pi }(s^{\prime }\mid s,a)$ and $P(s^{\prime }\mid s,a)$ are replaced by $M^{{\pi }_z}(s^{\prime }\mid s)=M^{{\pi }_z}(s^{\prime }\mid s,{\pi }_z(s))$ and $P^{{\pi }_z}(s^{\prime }\mid s)=P(s^{\prime }\mid s,{\pi }_z(s))$.

下面我们给出一些理论论证，说明潜在预测时序差分表示如何捕获给定策略集合的长期动态，并使其适用于零样本强化学习。 沿用 Tang 等人的研究，我们考虑带有线性预测器的简化表格设置。 我们将表示 $\varphi$（相应地，$\psi$）视为 $S\times d_{\varphi }$（相应地，$S\times d_{\psi }$）矩阵，并把所有 $z$ 对应的无动作预测器 $T_{\varphi ,z}$（相应地，$T_{\psi ,z}$）视为 $d_{\varphi }\times d_{\psi }$（相应地，$d_{\psi }\times d_{\varphi }$）矩阵。 因此，公式（8）和（9）中的表达式 $T_{\varphi }(\varphi (s),a,z)$ 化简为 $T_{\varphi ,z}^{\mathrm{\top }}\varphi (s)$，而 $M^{\pi }(s^{\prime }\mid s,a)$ 和 $P(s^{\prime }\mid s,a)$ 分别替换为 $M^{{\pi }_z}(s^{\prime }\mid s)=M^{{\pi }_z}(s^{\prime }\mid s,{\pi }_z(s))$ 与 $P^{{\pi }_z}(s^{\prime }\mid s)=P(s^{\prime }\mid s,{\pi }_z(s))$。

4.1 Monte-Carlo losses

We define a (non-latent-predictive) successor measure approximation loss

我们定义一个非潜在预测的后继测度近似损失：

$${\mathcal{L}}_{\mathrm{SM}}(\varphi ,\{T_z{\}}_z,\psi ):=\frac{1}{2}{\mathbb{E}}_{z\sim \mathcal{Z}}{\Vert \varphi T_z{\psi }^{\mathrm{\top }}-M^{{\pi }_z}\Vert }_F^2.$$

Minimizing ${\mathcal{L}}_{\mathrm{SM}}$ is equivalent to finding the best multilinear approximation to the successor measures $M^{{\pi }_z}$. We prove the following connection with the Monte Carlo latent-predictive loss of Equation (8).

最小化 ${\mathcal{L}}_{\mathrm{SM}}$ 等价于寻找后继测度 $M^{{\pi }_z}$ 的最佳多线性近似。 我们证明它与公式（8）的蒙特卡洛潜在预测损失存在下列联系。

Theorem 1. For fixed $\varphi$ and $\psi$, let $T_z^{\star }$, $T_{\varphi ,z}^{\star }$, $T_{\psi ,z}^{\star }$ be the optimal predictors for ${\mathcal{L}}_{\mathrm{SM}}(\varphi ,T_z,\psi )$ (Equation (12)), ${\mathcal{L}}_{\text{MC-JEPA}}(\varphi ,T_{\varphi ,z},\psi )$, ${\mathcal{L}}_{\text{MC-JEPA}}(\psi ,T_{\psi ,z},\varphi )$ (Equation (8)), respectively. If (A1) ${\varphi }^{\mathrm{\top }}\varphi ={\psi }^{\mathrm{\top }}\psi =I$, (A2) the state distribution is uniform, and (A3) for all $z\in \mathcal{Z}$, the matrix $P^{{\pi }_z}$ is symmetric, then:

定理1。 固定 $\varphi$ 和 $\psi$，令 $T_z^{\star }$、$T_{\varphi ,z}^{\star }$、$T_{\psi ,z}^{\star }$ 分别为 ${\mathcal{L}}_{\mathrm{SM}}(\varphi ,T_z,\psi )$（公式（12））、${\mathcal{L}}_{\text{MC-JEPA}}(\varphi ,T_{\varphi ,z},\psi )$、${\mathcal{L}}_{\text{MC-JEPA}}(\psi ,T_{\psi ,z},\varphi )$（公式（8））的最优预测器。 若（A1）${\varphi }^{\mathrm{\top }}\varphi ={\psi }^{\mathrm{\top }}\psi =I$，（A2）状态分布均匀，并且（A3）对所有 $z\in \mathcal{Z}$，矩阵 $P^{{\pi }_z}$ 对称，则：

1. For all $z$, $\varphi T_z^{\star }=\varphi T_{\varphi ,z}^{\star }={\mathrm{\Pi }}_{\varphi }{M}^{{\pi }_z}\psi$ and $\psi T_{\psi ,z}^{\star }=\psi (T_z^{\star }{)}^{\mathrm{\top }}={\mathrm{\Pi }}_{\psi }{M}^{{\pi }_z}\varphi$, where ${\mathrm{\Pi }}_{\varphi }$ (resp. ${\mathrm{\Pi }}_{\psi }$) is an orthogonal projection on the span of $\varphi$ (resp. $\psi$).
2. ${\mathrm{\nabla }}_{\varphi }{\mathcal{L}}_{\text{MC-JEPA}}(\varphi ,T_z,\psi )={\mathrm{\nabla }}_{\varphi }{\mathcal{L}}_{\mathrm{SM}}(\varphi ,T_z,\psi )$ and ${\mathrm{\nabla }}_{\psi }{\mathcal{L}}_{\text{MC-JEPA}}(\psi ,T_z,\varphi )={\mathrm{\nabla }}_{\psi }{\mathcal{L}}_{\mathrm{SM}}(\varphi ,T_z^{\mathrm{\top }},\psi )$.

1. 对所有 $z$，$\varphi T_z^{\star }=\varphi T_{\varphi ,z}^{\star }={\mathrm{\Pi }}_{\varphi }{M}^{{\pi }_z}\psi$ 且 $\psi T_{\psi ,z}^{\star }=\psi (T_z^{\star }{)}^{\mathrm{\top }}={\mathrm{\Pi }}_{\psi }{M}^{{\pi }_z}\varphi$，其中 ${\mathrm{\Pi }}_{\varphi }$（相应地，${\mathrm{\Pi }}_{\psi }$）是到 $\varphi$（相应地，$\psi$）张成空间的正交投影。
2. ${\mathrm{\nabla }}_{\varphi }{\mathcal{L}}_{\text{MC-JEPA}}(\varphi ,T_z,\psi )={\mathrm{\nabla }}_{\varphi }{\mathcal{L}}_{\mathrm{SM}}(\varphi ,T_z,\psi )$ 且 ${\mathrm{\nabla }}_{\psi }{\mathcal{L}}_{\text{MC-JEPA}}(\psi ,T_z,\varphi )={\mathrm{\nabla }}_{\psi }{\mathcal{L}}_{\mathrm{SM}}(\varphi ,T_z^{\mathrm{\top }},\psi )$。

This result reveals that 1) the optimal predictors for the successor measure loss ${\mathcal{L}}_{\mathrm{SM}}$ and the latent-predictive loss ${\mathcal{L}}_{\text{MC-JEPA}}$ match, and yield an orthogonal projection of the successor features $M^{{\pi }_z}\psi$ onto the $\varphi$ space; 2) the gradients w.r.t. the representations $\varphi$ and $\psi$, when evaluated at any predictor, match among these two losses, showing that gradient descent on ${\mathcal{L}}_{\text{MC-JEPA}}$ would update representations in the direction that reduces ${\mathcal{L}}_{\mathrm{SM}}$, hence improving the approximation of the successor measures. This result follows as a special case of a novel theorem (see the appendix) generalizing and implying all previous guarantees for latent-predictive representations, which we believe is of independent interest. Finally, we remark that, while the assumptions A1-A3 have been considered in all these related works, they can be relaxed, at the price of more involved proofs and notation, as shown in the appendix.

该结果揭示：1) 后继测度损失 ${\mathcal{L}}_{\mathrm{SM}}$ 与潜在预测损失 ${\mathcal{L}}_{\text{MC-JEPA}}$ 的最优预测器一致，并将后继特征 $M^{{\pi }_z}\psi$ 正交投影到 $\varphi$ 空间；2) 在任意预测器处求值时，这两个损失关于表示 $\varphi$ 和 $\psi$ 的梯度一致，这表明对 ${\mathcal{L}}_{\text{MC-JEPA}}$ 执行梯度下降会沿着降低 ${\mathcal{L}}_{\mathrm{SM}}$ 的方向更新表示，从而改善后继测度近似。 该结果是附录中一个新定理的特殊情况；该定理推广并蕴含潜在预测表示的所有既有保证，我们认为它本身也具有独立价值。 最后，尽管这些相关工作都采用了假设 A1-A3，但正如附录所示，它们可以放宽，只是需要更复杂的证明与记号。

4.2 Temporal-difference losses

We first derive a non-collapse guarantee. While a similar result was originally proved by Tang et al. for the one-step loss (Equation (3)), our case is more complex since TD latent-prediction can be seen as "doubly latent-predictive" (cf. Equation (9)): $T_{\varphi ,z}^{\mathrm{\top }}\varphi (s)$ is optimized to match a representation being learned--$\psi (s^{+})$--plus a bootstrapped version of itself--$T_{\varphi ,z}^{\mathrm{\top }}\varphi (s^{+})$.

我们首先推导一个非坍缩保证。 Tang 等人最初针对单步损失（公式（3））证明了类似结果，但我们的情况更复杂，因为 TD 潜在预测可以视为“双重潜在预测”（参见公式（9））：$T_{\varphi ,z}^{\mathrm{\top }}\varphi (s)$ 被优化为匹配一个正在学习的表示 $\psi (s^{+})$，再加上它自身的自举版本 $T_{\varphi ,z}^{\mathrm{\top }}\varphi (s^{+})$。

Theorem 2. Let ${\varphi }_t$ and ${\psi }_t$ be the representations learned under a continuous-time relaxation of Equation (9) where, at each step $t$, the optimal predictors for $({\varphi }_t,{\psi }_t)$ are first computed and then a gradient step on $({\varphi }_t,{\psi }_t)$ is taken (see the appendix for the explicit formulation). Then, the covariance matrices ${\varphi }_t^{\mathrm{\top }}{\varphi }_t$ and ${\psi }_t^{\mathrm{\top }}{\psi }_t$ are constant over time, i.e., ${\varphi }_t^{\mathrm{\top }}{\varphi }_t={\varphi }_0^{\mathrm{\top }}{\varphi }_0$ and ${\psi }_t^{\mathrm{\top }}{\psi }_t={\psi }_0^{\mathrm{\top }}{\psi }_0$ for all $t\ge 0$.

定理2。 令 ${\varphi }_t$ 和 ${\psi }_t$ 表示在公式（9）的连续时间松弛下学到的表示；在每一步 $t$，先计算 $({\varphi }_t,{\psi }_t)$ 的最优预测器，再对 $({\varphi }_t,{\psi }_t)$ 执行一个梯度步骤（显式形式见附录）。 那么，协方差矩阵 ${\varphi }_t^{\mathrm{\top }}{\varphi }_t$ 和 ${\psi }_t^{\mathrm{\top }}{\psi }_t$ 随时间保持不变，即对所有 $t\ge 0$，都有 ${\varphi }_t^{\mathrm{\top }}{\varphi }_t={\varphi }_0^{\mathrm{\top }}{\varphi }_0$ 且 ${\psi }_t^{\mathrm{\top }}{\psi }_t={\psi }_0^{\mathrm{\top }}{\psi }_0$。

This result suggests that, if predictors are trained at a faster rate than representations, the overall dynamics preserve their covariance, thus preventing $\varphi$ and $\psi$ from collapsing to trivial solutions (e.g., $\varphi =\psi =0$) when properly initialized, e.g., with unitary covariance.

该结果表明，如果预测器的训练速度快于表示，则整体动态会保持其协方差；在适当初始化（例如单位协方差）时，这可以防止 $\varphi$ 和 $\psi$ 坍缩到平凡解（例如 $\varphi =\psi =0$）。

As done for MC objectives (Theorem 1), we now show that the latent-predictive loss of TD-JEPA is related to forward and backward TD losses for approximating the successor measure.

与蒙特卡洛目标的处理方式（定理1）类似，下面我们表明 TD-JEPA 的潜在预测损失与近似后继测度的前向和后向 TD 损失相关。

Theorem 3. Consider the following TD losses for approximating the successor measure:

定理3。 考虑下列用于近似后继测度的 TD 损失：

$$\begin{array}{rl}{\mathcal{L}}_{\mathrm{fw}}(\varphi ,T_z,\psi )& :=\frac{1}{2}{\mathbb{E}}_{z\sim \mathcal{Z}}\left[{\Vert \begin{array}{rl}\varphi T_z{\psi }^{\mathrm{\top }}& -P^{{\pi }_z}\\ & -\gamma \mathrm{sg}(P^{{\pi }_z}\varphi T_z{\psi }^{\mathrm{\top }})\end{array}\Vert }_F^2\right],\\ {\mathcal{L}}_{\mathrm{bw}}(\varphi ,T_z,\psi )& :=\frac{1}{2}{\mathbb{E}}_{z\sim \mathcal{Z}}\left[{\Vert \begin{array}{rl}\psi T_z{\varphi }^{\mathrm{\top }}& -(P^{{\pi }_z}{)}^{\mathrm{\top }}\\ & -\gamma (P^{{\pi }_z}{)}^{\mathrm{\top }}\mathrm{sg}(\psi T_z{\varphi }^{\mathrm{\top }})\end{array}\Vert }_F^2\right].\end{array}$$

For fixed $(\varphi ,\psi )$, let $T_{z,\mathrm{fw}}^{\star }$, $T_{z,\mathrm{bw}}^{\star }$, $T_{\varphi ,z}^{\star }$, $T_{\psi ,z}^{\star }$ respectively be the optimal predictors for ${\mathcal{L}}_{\mathrm{fw}}(\varphi ,T_z,\psi )$, ${\mathcal{L}}_{\mathrm{bw}}(\varphi ,T_z,\psi )$, ${\mathcal{L}}_{\text{TD-JEPA}}(\varphi ,T_z,\psi )$, ${\mathcal{L}}_{\text{TD-JEPA}}(\psi ,T_z,\varphi )$. Under the same assumptions as Theorem 1:

固定 $(\varphi ,\psi )$，令 $T_{z,\mathrm{fw}}^{\star }$、$T_{z,\mathrm{bw}}^{\star }$、$T_{\varphi ,z}^{\star }$、$T_{\psi ,z}^{\star }$ 分别为 ${\mathcal{L}}_{\mathrm{fw}}(\varphi ,T_z,\psi )$、${\mathcal{L}}_{\mathrm{bw}}(\varphi ,T_z,\psi )$、${\mathcal{L}}_{\text{TD-JEPA}}(\varphi ,T_z,\psi )$、${\mathcal{L}}_{\text{TD-JEPA}}(\psi ,T_z,\varphi )$ 的最优预测器。 在与定理1相同的假设下：

1. For all $z$, $\varphi T_{\varphi ,z}^{\star }=\varphi T_{z,\mathrm{fw}}^{\star }={\tilde{\mathrm{\Pi }}}_{\varphi ,z}{M}^{{\pi }_z}\psi$ and $\psi T_{\psi ,z}^{\star }=\psi T_{z,\mathrm{bw}}^{\star }={\tilde{\mathrm{\Pi }}}_{\psi ,z}{M}^{{\pi }_z}\varphi$, where ${\tilde{\mathrm{\Pi }}}_{\varphi ,z}$ (resp. ${\tilde{\mathrm{\Pi }}}_{\psi ,z}$) is an oblique projection on the span of $\varphi$ (resp. $\psi$).
2. ${\mathrm{\nabla }}_{\varphi }{\mathcal{L}}_{\text{TD-JEPA}}(\varphi ,T_z,\psi )={\mathrm{\nabla }}_{\varphi }{\mathcal{L}}_{\mathrm{fw}}(\varphi ,T_z,\psi )$ and ${\mathrm{\nabla }}_{\psi }{\mathcal{L}}_{\text{TD-JEPA}}(\psi ,T_z,\varphi )={\mathrm{\nabla }}_{\psi }{\mathcal{L}}_{\mathrm{fw}}(\varphi ,T_z,\psi )$.

1. 对所有 $z$，$\varphi T_{\varphi ,z}^{\star }=\varphi T_{z,\mathrm{fw}}^{\star }={\tilde{\mathrm{\Pi }}}_{\varphi ,z}{M}^{{\pi }_z}\psi$ 且 $\psi T_{\psi ,z}^{\star }=\psi T_{z,\mathrm{bw}}^{\star }={\tilde{\mathrm{\Pi }}}_{\psi ,z}{M}^{{\pi }_z}\varphi$，其中 ${\tilde{\mathrm{\Pi }}}_{\varphi ,z}$（相应地，${\tilde{\mathrm{\Pi }}}_{\psi ,z}$）是到 $\varphi$（相应地，$\psi$）张成空间的斜投影。
2. ${\mathrm{\nabla }}_{\varphi }{\mathcal{L}}_{\text{TD-JEPA}}(\varphi ,T_z,\psi )={\mathrm{\nabla }}_{\varphi }{\mathcal{L}}_{\mathrm{fw}}(\varphi ,T_z,\psi )$ 且 ${\mathrm{\nabla }}_{\psi }{\mathcal{L}}_{\text{TD-JEPA}}(\psi ,T_z,\varphi )={\mathrm{\nabla }}_{\psi }{\mathcal{L}}_{\mathrm{fw}}(\varphi ,T_z,\psi )$。

Similar to Theorem 1, the optimal predictors and gradients of TD-JEPA match those of the non-latent-predictive TD losses of Equations (13) and (14), which are known to recover an approximation of the successor measure for bilinear parameterizations of the form $F_z^{\mathrm{\top }}B$. Unlike in the Monte Carlo case, here the optimal predictors solve a least-squares TD problem, yielding the fixed point of a projected Bellman operator whose closed-form expression is an oblique projection.

与定理1类似，TD-JEPA 的最优预测器和梯度与公式（13）和（14）的非潜在预测 TD 损失一致；已知后者能够为 $F_z^{\mathrm{\top }}B$ 形式的双线性参数化恢复后继测度近似。 不同于蒙特卡洛情形，这里的最优预测器求解一个最小二乘 TD 问题，得到投影贝尔曼算子的固定点，其闭式表达为斜投影。

4.3 Policy evaluation and zero-shot RL

Finally, the following result motivates the significance of optimizing the successor measure losses of Equations (12), (13), and (14).

最后，下列结果说明了优化公式（12）、（13）和（14）中的后继测度损失的重要性。

Theorem 4. Let $\varphi ,\psi$ have identity covariance matrices. For any reward function $r$, let ${\omega }_r:=({\psi }^{\mathrm{\top }}\psi {)}^{-1}{\psi }^{\mathrm{\top }}r$ be the linear regression weight for representation $\psi$. Then, for any $T_z$,

定理4。 设 $\varphi ,\psi$ 具有单位协方差矩阵。 对任意奖励函数 $r$，令 ${\omega }_r:=({\psi }^{\mathrm{\top }}\psi {)}^{-1}{\psi }^{\mathrm{\top }}r$ 为表示 $\psi$ 的线性回归权重。 那么，对任意 $T_z$，

$$\begin{array}{r}\underset{\begin{array}{c}r\in {\mathbb{R}}^S\\ \parallel r{\parallel }_2\le 1\end{array}}{max}{\mathbb{E}}_{z\in \mathcal{Z}}\left[\sum _{s\in \mathcal{S}}{(V_r^{{\pi }_z}(s)-\varphi (s{)}^{\mathrm{\top }}{T}_z{\omega }_r)}^2\right]\le 2{\mathcal{L}}_{\mathrm{SM}}(\varphi ,T_z,\psi ).\end{array}$$

Moreover, ${\mathcal{L}}_{\mathrm{SM}}(\varphi ,T_z,\psi )\le c{\mathcal{L}}_{\mathrm{fw}}(\varphi ,T_z,\psi )$ and ${\mathcal{L}}_{\mathrm{SM}}(\varphi ,T_z,\psi )\le c{\mathcal{L}}_{\mathrm{bw}}(\varphi ,T_z,\psi )$ for some $c$.

此外，对某个 $c$，有 ${\mathcal{L}}_{\mathrm{SM}}(\varphi ,T_z,\psi )\le c{\mathcal{L}}_{\mathrm{fw}}(\varphi ,T_z,\psi )$ 和 ${\mathcal{L}}_{\mathrm{SM}}(\varphi ,T_z,\psi )\le c{\mathcal{L}}_{\mathrm{bw}}(\varphi ,T_z,\psi )$。

Paraphrasing, the policy evaluation error of the technique in Section 3.3 (i.e., embed $r$ into a vector $\omega$ through linear regression on $\psi$, and compute $T_{\varphi }(\varphi (s),z{)}^{\mathrm{\top }}\omega$) is bounded by the successor measure approximation loss and the corresponding TD errors. Both these quantities are indirectly optimized by TD-JEPA (Theorems 1 and 3), which is thus a sound approach for zero-shot policy evaluation. Moreover, Theorem 4 leads to a zero-shot optimality result analogous to Theorem 2 of Touati et al.: if the approximation of $M^{{\pi }_z}$ is perfect (i.e., $M^{{\pi }_z}=\varphi T_z{\psi }^{\mathrm{\top }}$ for all $z$ or, equivalently, the TD errors in Equations (13) and (14) are zero) and the policies ${\pi }_z$ are optimal for all linear rewards in $\psi$, then the inference procedure above recovers optimal policies for any (even non-linear) reward function.

换言之，第3.3节所述技术的策略评估误差，即通过在 $\psi$ 上进行线性回归把 $r$ 嵌入向量 $\omega$，再计算 $T_{\varphi }(\varphi (s),z{)}^{\mathrm{\top }}\omega$，受到后继测度近似损失及相应 TD 误差的约束。 这两个量都由 TD-JEPA 间接优化（定理1和定理3），因此 TD-JEPA 是一种理论合理的零样本策略评估方法。 此外，定理4导出了一个与 Touati 等人定理2类似的零样本最优性结果：若 $M^{{\pi }_z}$ 的近似完美，即对所有 $z$ 都有 $M^{{\pi }_z}=\varphi T_z{\psi }^{\mathrm{\top }}$，或等价地，公式（13）和（14）中的 TD 误差为零，并且策略 ${\pi }_z$ 对 $\psi$ 中所有线性奖励均最优，则上述推理过程可以为 任意 奖励函数恢复最优策略，包括非线性奖励函数。

5. Experiments

We benchmark zero-shot performance across a diverse set of problems, including 4 locomotion/navigation domains from ExoRL/DMC, as well as 9 navigation/manipulation domains from OGBench. The former suite involves reward-based tasks and high-coverage data, while the latter evaluates goal-reaching and provides low-coverage datasets. We additionally apply BC regularization in OGBench, as detailed in the appendix. We consider both proprioceptive and pixel-based variants of all domains, and report expected returns/success rates across a set of tasks (4-8 depending on the domain) as main evaluation metric. In DMC, we often normalize returns by the maximum achievable ($1000$).

我们在多种问题上评测零样本性能，包括 ExoRL/DMC 的 4 个运动或导航域，以及 OGBench 的 9 个导航或操作域。 前一个套件包含基于奖励的任务和高覆盖率数据，后一个套件评估目标到达，并提供低覆盖率数据集。 我们还在 OGBench 中应用 BC 正则化，具体细节见附录。 我们考虑所有域的本体感知和像素版本，并把一组任务上的期望回报或成功率作为主要评估指标；每个域包含 4 至 8 个任务。 在 DMC 中，我们通常用可达到的最大值（$1000$）对回报进行归一化。

We structure our evaluation in four parts: (i) a comprehensive evaluation of TD-JEPA with respect to existing zero-shot methods; (ii) an ablation over the prediction target, measuring the impact of multi-step, policy-aware dynamics modeling; (iii) a comparison of TD-JEPA to its symmetric variant that learns a shared state-task encoder $\varphi$; and (iv) a demonstration of fast adaptation from pre-trained state representations. Further results are presented in the appendix, and implementation details in the appendix.

我们的评估分为四部分：(i) 将 TD-JEPA 与现有零样本方法进行全面比较；(ii) 对预测目标进行消融，衡量多步、策略感知动态建模的影响；(iii) 将 TD-JEPA 与学习共享状态—任务编码器 $\varphi$ 的对称变体进行比较；(iv) 展示如何基于预训练状态表示快速适应。 更多结果和实现细节见附录。

5.1 How does TD-JEPA compare to zero-shot RL algorithms?

表1：DMC（回报）与 OGBench（成功率）上使用本体感知或 RGB 输入时，各零样本算法的性能。报告跨随机种子的均值和标准误；若置信区间重叠，则表现领先的算法均以粗体标出。

Domain	Laplacian	ICVF*	HILP	FB	RLDP	BYOL*	BYOL-γ*	TD-JEPA
DMCRGB (avg)	293.1 ± 15.1	438.7 ± 14.9	391.2 ± 23.8	456.2 ± 8.6	525.7 ± 13.3	513.8 ± 11.6	582.4 ± 9.8	628.8 ± 5.5
walker	309.4 ± 50.0	534.9 ± 61.3	422.8 ± 32.5	324.4 ± 16.6	576.1 ± 35.3	595.2 ± 9.0	648.3 ± 36.5	738.9 ± 3.5
cheetah	242.4 ± 29.6	394.9 ± 30.1	333.0 ± 86.6	622.4 ± 23.1	605.3 ± 23.5	468.0 ± 46.7	679.8 ± 17.1	706.0 ± 4.1
quadruped	430.1 ± 32.3	583.3 ± 17.2	513.9 ± 10.8	475.4 ± 16.7	551.1 ± 23.4	581.8 ± 16.6	570.0 ± 6.6	626.7 ± 13.6
pointmass	190.4 ± 12.4	241.6 ± 35.6	294.9 ± 33.4	402.8 ± 16.8	370.3 ± 12.0	410.3 ± 8.5	431.6 ± 17.4	443.7 ± 10.9
DMC (avg)	591.1 ± 10.7	619.3 ± 10.3	620.1 ± 8.4	648.2 ± 4.1	610.2 ± 13.5	618.6 ± 10.5	645.4 ± 10.5	661.2 ± 6.3
walker	769.7 ± 4.7	727.0 ± 16.2	796.4 ± 7.7	811.5 ± 5.9	723.9 ± 18.3	746.8 ± 11.0	786.1 ± 9.6	785.2 ± 6.7
cheetah	614.5 ± 18.9	606.3 ± 16.8	618.3 ± 5.8	672.7 ± 4.9	575.6 ± 44.9	622.8 ± 23.9	647.2 ± 9.0	688.7 ± 6.7
quadruped	635.0 ± 38.7	708.5 ± 14.2	694.8 ± 11.0	595.6 ± 9.1	665.0 ± 13.9	611.8 ± 28.1	683.1 ± 26.1	691.4 ± 5.0
pointmass	345.1 ± 22.4	435.5 ± 11.1	371.0 ± 37.1	513.0 ± 20.0	476.3 ± 39.4	493.0 ± 41.3	465.1 ± 17.6	479.3 ± 23.6
OGBenchRGB (avg)	30.58 ± 0.81	25.22 ± 0.55	32.56 ± 0.92	39.89 ± 0.47	39.09 ± 0.59	40.33 ± 0.52	41.58 ± 0.64	41.34 ± 0.45
antmaze-mn	92.20 ± 2.91	85.80 ± 3.02	84.60 ± 3.59	96.80 ± 0.74	97.60 ± 0.50	94.40 ± 1.48	98.00 ± 0.73	96.67 ± 1.11
antmaze-ln	35.40 ± 2.97	42.60 ± 2.84	47.00 ± 4.04	76.80 ± 2.33	63.60 ± 3.89	62.20 ± 3.42	68.80 ± 2.70	74.60 ± 3.35
antmaze-ms	60.20 ± 3.88	46.20 ± 2.74	71.80 ± 2.22	86.20 ± 2.05	90.60 ± 1.91	90.40 ± 1.97	86.00 ± 3.10	84.40 ± 3.85
antmaze-ls	7.20 ± 1.98	7.20 ± 1.20	23.60 ± 1.83	27.40 ± 2.78	21.80 ± 1.01	26.60 ± 2.23	28.60 ± 1.71	28.80 ± 2.50
antmaze-me	0.00 ± 0.00	0.00 ± 0.00	0.20 ± 0.20	1.80 ± 1.09	0.80 ± 0.44	1.20 ± 1.00	3.20 ± 1.98	0.20 ± 0.20
cube-single	73.80 ± 3.53	34.80 ± 7.03	56.40 ± 3.82	62.00 ± 2.27	63.20 ± 3.91	75.40 ± 2.58	76.40 ± 3.24	67.80 ± 3.67
cube-double	1.60 ± 0.72	0.80 ± 0.44	1.60 ± 0.58	1.20 ± 0.61	2.20 ± 1.31	2.40 ± 0.65	1.40 ± 0.67	3.00 ± 0.91
scene	2.80 ± 1.12	8.40 ± 1.45	5.40 ± 1.63	4.20 ± 0.87	9.40 ± 1.33	8.80 ± 1.64	11.20 ± 1.82	14.20 ± 2.22
puzzle-3x3	2.00 ± 1.40	1.20 ± 0.44	2.44 ± 0.99	2.60 ± 0.79	2.60 ± 0.79	1.60 ± 0.40	0.60 ± 0.31	2.40 ± 0.83
OGBench (avg)	14.81 ± 1.32	30.87 ± 0.58	37.98 ± 1.11	39.04 ± 0.66	27.07 ± 0.83	26.42 ± 0.83	30.42 ± 0.94	37.98 ± 0.77
antmaze-mn	50.00 ± 4.94	79.80 ± 2.62	83.60 ± 2.63	73.00 ± 2.72	74.60 ± 4.15	58.40 ± 2.00	51.40 ± 1.55	70.40 ± 3.72
antmaze-ln	21.60 ± 3.90	58.40 ± 1.90	52.60 ± 3.86	36.80 ± 4.28	36.40 ± 4.66	26.60 ± 3.03	21.80 ± 3.57	57.20 ± 4.25
antmaze-ms	21.40 ± 4.32	39.00 ± 3.30	50.60 ± 2.46	70.40 ± 3.95	58.40 ± 3.29	60.60 ± 5.07	45.60 ± 2.84	61.56 ± 4.53
antmaze-ls	11.80 ± 1.47	13.20 ± 1.64	12.20 ± 1.75	49.80 ± 5.64	19.60 ± 2.73	25.80 ± 4.28	20.20 ± 1.80	40.60 ± 2.51
antmaze-me	0.80 ± 0.61	0.00 ± 0.00	2.00 ± 0.84	51.60 ± 2.65	4.80 ± 2.35	11.40 ± 2.29	19.60 ± 2.53	20.20 ± 2.39
cube-single	15.11 ± 1.49	20.40 ± 1.93	74.20 ± 3.53	49.60 ± 3.83	19.80 ± 2.41	22.00 ± 3.16	79.40 ± 2.83	34.20 ± 2.88
cube-double	2.00 ± 0.42	5.00 ± 0.80	20.00 ± 2.72	2.60 ± 0.43	3.80 ± 0.76	4.40 ± 0.72	2.60 ± 0.67	3.60 ± 0.78
scene	7.80 ± 1.28	45.40 ± 2.29	43.80 ± 1.90	12.80 ± 1.61	11.60 ± 1.57	15.40 ± 1.37	14.40 ± 2.32	38.44 ± 1.37
puzzle-3x3	2.80 ± 0.68	16.60 ± 0.73	2.80 ± 0.68	4.80 ± 0.68	14.60 ± 0.90	13.20 ± 1.91	18.80 ± 0.44	15.60 ± 1.11

We first compare TD-JEPA to three groups of successor-feature-based zero-shot RL baselines. Notice that only Laplacian, HILP, FB and RLDP are standard zero-shot unsupervised RL algorithms, while BYOL, BYOL-$\gamma$, and ICVF (henceforth marked with a $\ast$) are representation learning methods: their instantiation in a zero-shot framework is novel and designed to investigate the impact of different representations.

我们首先将 TD-JEPA 与三组基于后继特征的零样本强化学习基线进行比较。 需要注意，只有 Laplacian、HILP、FB 和 RLDP 是标准的零样本无监督强化学习算法；BYOL、BYOL-$\gamma$ 与 ICVF（下文以 $\ast$ 标记）是表示学习方法：我们首次将它们实例化为零样本框架，以研究不同表示产生的影响。

• Laplacian, HILP, and FB are established zero-shot methods that train a task encoder $\psi$, without specific learning objectives for a state encoder.
• BYOL${}^{\star }$, BYOL-${\gamma }^{\star }$ and RLDP learn a state encoder $\varphi$ via latent-predictive learning, which we then use as a task encoder for successor features (learned through a contrastive loss in the case of RLDP).
• ICVF${}^{\star }$ learns a multilinear decomposition of the successor measure via expectile regression, yielding both state and task encoders on top of which we train successor features.

• Laplacian、HILP 和 FB 是已有的零样本方法，它们训练任务编码器 $\psi$，但不为状态编码器设置专门的学习目标。
• BYOL${}^{\star }$、BYOL-${\gamma }^{\star }$ 和 RLDP 通过潜在预测学习状态编码器 $\varphi$，随后将其作为后继特征的任务编码器；对于 RLDP，后继特征通过对比损失学习。
• ICVF${}^{\star }$ 通过期望回归学习后继测度的多线性分解，同时得到状态编码器和任务编码器，我们再在其上训练后继特征。

For a fair comparison, each method is tuned over comparable hyperparameter grids and adopts the same architecture: in particular, the state input is always passed through an explicit state encoder before being fed into, e.g., the successor features estimator $F(s,a;z)$. On average, explicit state encoders actually improve the performance for existing methods. We find that this protocol results in significant improvements in zero-shot performances, even for existing methods (e.g., $1.3\times$ and $2.4\times$ higher than overlapping pixel-based results for HILP and RLDP, respectively), as displayed in Table 1.

为公平比较，每种方法都在规模相当的超参数网格上调优，并采用相同架构；具体而言，状态输入始终先经过显式状态编码器，再送入后继特征估计器 $F(s,a;z)$ 等组件。 平均而言，显式状态编码器实际上也提高了现有方法的性能。 如表1所示，我们发现该协议显著提高了零样本性能，对已有方法亦是如此；例如，HILP 和 RLDP 的结果分别比与既有工作重叠的像素任务结果高 $1.3\times$ 和 $2.4\times$。

When considering suite-aggregated performance, we find that TD-JEPA is on par or better than the best performing baseline in each suite. Given the diverse nature of suites (proprioception vs pixels), domains (locomotion, navigation, manipulation) and datasets (high- vs low-coverage), many algorithms unsurprisingly achieve strong performance in some configurations while under-performing in others. We thus additionally measure how consistently well each algorithm performs by computing the probability of improvement across all domains in Figure 2. We find that TD-JEPA is consistently among the top performing algorithms, whereas most baselines perform well on a narrow subset of problems. For instance, while TD-JEPA is only slightly preferable to FB and HILP from proprioception, it is significantly better than them in visual domains. Similarly, BYOL-$\gamma$ is slightly better than TD-JEPA in OGBench${}_{\text{RGB}}$, but it is significantly worse in DMC${}_{\text{RGB}}$ and OGBench. Finally, we note that latent-predictive methods tend to be generally preferrable in pixel-based domains.

从各套件的聚合性能来看，TD-JEPA 在每个套件中都达到或超过了表现最好的基线。 鉴于套件（本体感知与像素）、领域（运动、导航与操作）以及数据集（高覆盖与低覆盖）都具有多样性，许多算法在部分配置上表现强劲而在其他配置上表现较弱并不意外。 因此，我们进一步计算所有领域上的提升概率，以衡量每种算法持续取得良好表现的程度，如图2所示。 我们发现 TD-JEPA 始终位列表现最好的算法，而多数基线只在一小部分问题上表现良好。 例如，在本体感知输入下，TD-JEPA 仅略优于 FB 和 HILP，但在视觉领域中显著优于二者。 类似地，BYOL-$\gamma$ 在 OGBench${}_{\text{RGB}}$ 上略优于 TD-JEPA，但在 DMC${}_{\text{RGB}}$ 和 OGBench 上明显更差。 最后，我们注意到，潜在预测方法在基于像素的领域中往往总体更具优势。

图2：提升概率：方法 X 在随机领域上优于方法 Y 的可能性。报告对称化的 95% 简单 bootstrap 置信区间；虚线围出的比较表示提升具有统计显著性。

5.2 Which dynamics should latent-predictive zero-shot algorithms model?

The baselines based on BYOL and BYOL-$\gamma$ are algorithmically closest to TD-JEPA, and allow a precise investigation on the dynamics to model. While BYOL${}^{\star }$ and BYOL-${\gamma }^{\star }$ approximate one-step and multi-step transitions of the behavioral policy, respectively, TD-JEPA models multi-step transitions of the zero-shot policies. While approximating the behavioral dynamics can be effective for expert-like data (i.e., in OGBench), we observe a general pattern suggesting that directly modeling policy-conditional successor measures is on average beneficial, as reported in Figure 3 (left).

基于 BYOL 和 BYOL-$\gamma$ 的基线在算法上最接近 TD-JEPA，因此可以精确研究应当建模哪些动态。 BYOL${}^{\star }$ 和 BYOL-${\gamma }^{\star }$ 分别近似行为策略的单步转移和多步转移，而 TD-JEPA 建模的是零样本策略本身的多步转移。 虽然近似行为动态对专家式数据可能有效，例如 OGBench，但我们观察到一个普遍规律：直接建模策略条件后继测度平均而言更有益，如图3左侧所示。

图3：左：潜在预测方法的归一化零样本性能。右：TD-JEPA 与其对称变体之间的归一化性能差异。误差条分别表示归一化性能或性能差异的标准误。

5.3 Should state and task representations differ?

TD-JEPA trains separate state and task encoders: while this may grant a better approximation of successor measures, sharing state and task representations while optimizing a single objective (see Section 3.2) may in practice be more efficient. We measure the difference in per-task normalized performance between TD-JEPA and a symmetric variant in Figure 3 (right): we observe that this variant performs comparatively rather well, while relying on a single predictor-encoder pair. However, using distinct state and task embeddings tends to improve empirical performance more often than not.

TD-JEPA 分别训练状态编码器和任务编码器；这样做可能更好地近似后继测度，但共享状态表示和任务表示并优化单一目标（见第3.2节）在实践中可能更高效。 我们在图3右侧衡量 TD-JEPA 与对称变体在逐任务归一化性能上的差异：该变体只依赖一对预测器与编码器，但相对而言表现相当不错。 不过，使用不同的状态嵌入和任务嵌入通常更有可能提高经验性能。

5.4 Are state representations beneficial for fast adaptation?

While the previous evaluations have focused on aggregated zero-shot performance, we now investigate an additional benefit of explicit state representations: fast adaptation at test-time. Given a pixel-based task, we initialize the agent with the zero-shot policy ${\pi }_z$ and critic learned at pre-training, and we either fine-tune the whole model via TD3 or keep the pre-trained state encoder frozen. We consider two RL adaptation protocols: (i) Offline: a transition-reward dataset is provided ${\mathcal{D}}_{\text{rew}}=\{(s,a,s^{\prime },r)\}$ and TD3 updates are applied offline; (ii) Online: an online buffer is additionally collected over time and batches are sampled by mixing it with the offline buffer mentioned above.

前面的评估聚焦于聚合零样本性能，现在我们研究显式状态表示的另一项优势：测试时快速适应。 给定一个基于像素的任务，我们用预训练阶段学得的零样本策略 ${\pi }_z$ 和评论器初始化智能体，然后选择通过 TD3 微调整个模型，或者将预训练状态编码器保持为冻结状态。 我们考虑两种强化学习适应协议：（i）离线： 提供转移—奖励数据集 ${\mathcal{D}}_{\text{rew}}=\{(s,a,s^{\prime },r)\}$，并离线执行 TD3 更新；（ii）在线： 随时间额外收集在线缓冲区，并将其与上述离线缓冲区混合采样训练批次。

Figure 4 reports results for each DMC domain for the task in which the gap between online and zero-shot algorithms is largest; we consider TD-JEPA and FB as strong, representative algorithms among self-predictive and contrastive methods. We first observe that fine-tuning pre-trained agents leads to large gains in sample efficiency w.r.t. training from scratch, and reaches the asymptotic performance of TD3. More interestingly, frozen representations are often sufficient for downstream learning, and do not need further fine-tuning. We refer to the appendix for further results and details.

图4给出了每个 DMC 领域中在线算法与零样本算法差距最大的任务结果；我们选取 TD-JEPA 和 FB，分别作为自预测方法和对比方法中的强代表算法。 我们首先观察到，与从头训练相比，微调预训练智能体大幅提高了样本效率，并能达到 TD3 的渐近性能。 更值得注意的是，冻结表示通常已足以支持下游学习，无需进一步微调。 更多结果和细节见附录。

图4：零样本策略经离线微调（上）或在线微调（下）后的归一化性能。用零样本解初始化智能体（蓝线和黄线）可实现样本高效学习；冻结表示（虚线）往往已经具有足够的表达能力，可支持快速适应。

6. Conclusion

Through the introduction of a novel temporal-difference latent-predictive loss, we presented a zero-shot unsupervised RL method that operates entirely in latent space and can be shown to recover a factorization of the successor measures of multiple policies. Empirically, we found that TD-JEPA matches the best zero-shot methods when learning from proprioception, and exceeds them when learning from pixels, while also retrieving state representations that allow fast downstream adaptation. As formal guarantees rely on an assumption of symmetry, one exciting direction for future work may study learning objectives that are compatible with asymmetric successor measures, yet remain amenable to practical optimization. On a practical note, we believe that benchmarking latent-predictive zero-shot objectives on large-scale, real robotic dataset can shed further light on opportunities and limitations of this promising framework.

通过引入一种新型时序差分潜在预测损失，我们提出了一种完全在潜在空间中运行的零样本无监督强化学习方法，并且可以证明它能够恢复多个策略的后继测度分解。 在经验上，我们发现，TD-JEPA 使用本体感知输入学习时可以匹配最佳零样本方法，使用像素输入学习时则超过这些方法，同时还能获得支持下游快速适应的状态表示。 由于形式化保证依赖对称性假设，未来一个令人期待的方向是研究既与非对称后继测度兼容、又便于实际优化的学习目标。 从实践角度看，我们认为，在大规模真实机器人数据集上评测潜在预测零样本目标，可以进一步揭示这一有前景框架的机会与局限。