TD-JEPA: Latent-predictive Representations for Zero-Shot Reinforcement Learning
Authors: Marco Bagatella
Matteo Pirotta
Ahmed Touati
Alessandro Lazaric
Andrea Tirinzoni
Venue: ICLR 2026 Oral (2026)
Keywords: zero-shot reinforcement learning
unsupervised reinforcement learning
self-predictive representations
joint embedding predictive architecture
Primary area: reinforcement learning
TL;DR: We propose a temporal-difference latent-predictive method for zero-shot unsupervised RL.

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.

Published as a conference paper at ICLR 2026
TD-JEPA: LATENT-PREDICTIVE REPRESENTATIONS FOR ZERO-SHOT REINFORCEMENT LEARNING
Marco Bagatella" 123+ Matteo Pirotta' Ahmed Touati', Alessandro Lazaric', Andrea Tirinzoni! 1 FAIR at Meta, ETH Zürich, 3 Max Planck Institute for Intelligent Systems, Tübingen mbagatella@ethh. ch, {pirotta, touati, lazaric, cirinzoni]@meta...
ABSTRACT
Latent prediction-where agents learn by predicting their own latents-has emerged powerful paradigm for training general representations in machine learning. In reinforcement learning (RL), this approach has been explored to define auxil- iary losses for variety of settings, including reward-based and unsupervised RL, behavior cloning, and world modeling. While existing methods are typically lim- ited to single-task learning, one-step prediction, or on-policy trajectory data, we show that temporal difference (TD) learning enables learning representations pre- dictive of long-term latent dynamics across multiple policies/tasks from offline, reward-free transitions. Building on this, we introduce TD-JEPA, which lever- ages TD-based patent-predictivy representations into unsupervised RL. TD-JEPA trains explicit state and task encoders, policy-conditioned multi-step predictor, and a set of parameterized policies directly in latent space. This enables zero- shot optimization of any reward function test time. Theoretically, we show that an idealized variant of TD-JEPA avoids collapse with proper initialization, and learns encoders that capture low-rank factorization of long-term policy dynam- ics, while the predictor recovers their successor features in latent space. Empiri- cally, TD-JEPA matches or outperforms state-of-the-ar baselines on locomotion, navigation, and manipulation tasks across 13 datasets in ExoRL and OGBench, especially in the challenging setting of zero-shot RI. from pixels!
INTRODUCTION
Learning effective state representations is a core challenge in reinforcement learning (RL). Useful representations should capture the dynamics of the environment a way that supports efficient value estimation and policy optimization across tasks (Watter et al., [0155] Silver et al. 2018: Hafner et al.,
2019 Gelada al.) 2019). A promising line of work latent- predictive (a. self-predictive) representation learning (Schwarzer et al. 2021 Grill et al. 2020 Guo et al.) 2020 Tang et al.) 2023), an instance of the joint-embedding predictive architecture (LeCun) 2022 JEPA) paradigm. These algorithms jointly learn state encoder d(s) and predictor P. i.e., latent dynamics model estimating the representation of a future state s': P()())) = o(s'). Latent- parentt-redictive methods thus perform supervied learning entirely in latent space without any reward reconstruction of (possibly high- dineessinnsiinl) states.
Several RL methods leverage latent prediction as auxiliary loss to improve sample efficiency and generalization in reward-based learning (Schwarzer et al. 2021 Guo et al. 2020 Hansen et al. 2024] behavior cloning (Lawson et al. 2025). ). and curiosity-drive
exploration (Guo et al. (022)_ As latent-predictive losses do not require any reward, they have been recently used for unsupervised
RL: Assran al. (2025a). Zhou et al. (2025) and (oball al. (2025) learn latent world models that can solve goal-reaching tasks via test-time planning, whereas Jaeer et al. (2025) learn a state encoder from trajectory data to define the space of tasks used to optimize zero-shot unsupervised policies.
This paper proposes a novel way to instantiate latent- predditted representations for unsupervised RL. While previous methods have largely focused on either one-step dynamics, single-task/single-
Work done at Meta.
Code available aigihhuh.comffscchokkrrccccchch jepa

---

Published as a conference paper at ICLR 2026
stop grad
FH
Prod
TED-EEPA
Pred (eac
Enc
Eac
-2
Pref
-4
Prod/
Rac()
TD-JEPA trains policies "N parameterized by latents 2. The predictor, conditioned on 2, predicts the representations of future states visited by т. (left). When trained via TD, the predictor (arrows on the right) approximates successor features for each policy, i.e., the weighted barycenter (stars) of representations of visited states (circles).
Figure I:
policy training, relied on on-policy data, we introduce ponicy-conditioned, multi-step formula- tion based on a novel off- pollipllyy iemporal-diffrence 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 as- sociated predictors particularly well-suited for integration with off-policy, approaches to zero-shot unsupervised RL (Touati & Ollivier 2021 Touati et al.) 2023 Park et al., 2024). which have recently emerged as a promising solution for applications such as whole-body humanoid control in simulation (Triinzoni et al. 2025) and on real hardware (Li al. 2025).
seccessor-eatture based
We thus instantiate temporal difference preeit-preeietiv representation learning into TD-JEPA, a zero-shot unsupervised RL algorithm which pre-trains four components: state encoder, 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 D 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.
Theoretically. for an idealized version of TD-JEPA with linear predictors, we show that 1) the rep- resentations 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" argu- ment that extends and generalizes existing theoretical analyses of patent-predicttve representations, and connect TD-JEPA with other unsupervised RL methods such as fowardd-bckwaad (Touati & Ollivier 2021 and intention-conditione value functions (Ghosh et al.) 2023).
Empirically, we evaluate TD-JEPA on 65 tasks across 13 datasets from ExoRL Yarats et al. 2022a) and OGBench (Park et al. (025a)), covering locomotion, navigation, and manipulation with both pro- prioceptive and basel- 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 dimen- sions of the algorithm, demonstrating the importance of learning representations that are predictive of multi-step policy-depender coders. Finally, we show that learned representations can be easily reused for offline or online RL, improving over zero-shot policies and learning from scratch.
dynamics, and the advantage of training distinct state and task en-
2 PRELIMINARIES
We consider a reward-frre Markov Decision Process M PS,,, where and A are state and action spaces, P is the probability measure over next states when taking action in state s P(ds' 8, a), 7 E [0, 1) is a discount factor. Executing a Markov policy x : S - Prob( induces an unnormalized distribution over visited states, which is referred to the successor measure:
M (x (8)a)=
- Pr E X(s,a,T) C
(1)
t=
2

---

Published as a conference paper at ICLR 2026
Given reward function S - R and policy и, the action-value function (T measures the cumulative discounted reward obtained by the policy over an infinite horizon, ie., Q7 (s, a) = 8, 0, T]. Action value functions are connected to successor measures via
E
(s, [[ Md (ds* ss,a)r(*)) = (Ieao
(2)
which shows convenient linear decomposition of O into the reward function and the dy- namics induced by т. Standard RL agents aim at finding reward-maxinn wwwaad-maximizing policies (8) E arg m (s, where Q a) max, Q7 (8)
Latent-predictive representations. In high-dimensional settings, state encoders ø S Rd may be learned ease the estimation of action-value functions. For instance, if encoder is such that Q7 (8, a) = o(s) wa,r for some vector wa,- E Red,, then the RL process reduces to learning vectors in dd rather than high-dimensional functions Q7 (s, a). Latent-predictive learning has been shown to be an effective approach for this problem. In the simplest formulation, preed-predee representations capture the one-step latent dynamics of policy æ by minimizing the loss
Cone-stop (0, T) = (s))
([(s) - o(s")[]),
(3)
where T R: + RRds is a (possibly non-linear) predictor of the latent one-step dynamics induced by é and policy म,, and & denotes soop-gradiet.. Notably, optimizing for this loss does not require any decoding reconstruction, and it only relies instantiations of this approach have been shown both empirically and theoretically to produce rep- resentations that accurately approximate action-value functions policies (Guo et al. 2022 Tang
unsupervised dataset D {(s, a, s')]. Different
et al. 2023 Voelcker et al. 2024 Lawson et al. 2025 Fujimoto et al. 2025).
Successor-features and zero-shot unsupervised RL. Considering state encoder S -> Rd
and the associated space of linear rewards R2 {r(s) v(s) =IzE E R& J. Q-values for any reward function r(s) = whs)] 2, E Rep can be written as
(s, [ Jstes M (ds+ | = (b((s))
F (s, 2r, (4)
where F (s, E R&de captures the successor features of π (Barreto et al.
2017). The majority
of unsupervised zero-shot RL methods (Touati Ollivier 2021: Park et al.
2024; Agarwal et al,
2025 Jajoo al. 2025) learn successor features F(s,, 0; F (s, a) for set of parameterized
policies { T. (=) with Z Rd, that are trained to be optimal for all rewards in Rg. ie.. (.())) Ro arg max, (s, a; 2) 2, where F(s, a; z) z is an approximation of Q a) for r(s) = v0s) z. At vector 2, Rd is first obtained by projecting onto R.p. and the associated policy T., is then returned.
test time, given reward function
Given the role played by & in defining the space of tasks of interest, with abuse of terminology, will refer to as a task encoder. On the other hand, we shall call state encoder map ø S Rd that is used to embed states before feeding them into different networks (e.g., will train successor features Fg (Ø(8), a) and policies π (Ø(s)) in the latent space given by Φ). While the zero-shot methods cited so far train the task encoder & in different ways, and do not train any explicit state encoder o, the next section will show how multi-step policy-dependee antent-predicet learning can be used to train both simultaneously.
3
REPRESENTATIONS
We begin by showing how the latent-predictive loss of Eq. 3 can model multi-step and policy- dependent dynamics, and how temporal difference (TD) learning allows learning from offline tran- sition data. We will then expand this idea to learn separate state and task embeddings, and finally show how it can be instantiated as zero-shot unsupervised RL method.
3.1 MULTI-STEP POLICY-CONDITIONED LATENT PREDICTION
Let {mz)zez be a family of policies parameterized by 2 E Z, and D = ((s, be a dataset of transitions. We train state encoder ò S - Rd and policy-dependent predictor T
3

---

Published as a conference paper at ICLR 2026
Rd to be latent-predictive of the long-term dynamics of the policies {T²) i.., dd A 2
JEA Tτ) =
( [|
- +
(5)
where MC-JEPA stands for Monte-Carlo (MC) JEPA loss, as on-policy samples s M (-1s,a) are needed for all policies of interest. Intuitively, To (d(s), a, z) tries to predict future latent states visited by the policy æ. More formally, predictors trained via minimization of LMC.JERA (0. To) approximate the successor features of & in the latent space induced by e itself.
Proposition For any & and To, we have the following equivalence
LMC-JEPA (0, TT) B(e.))D, (0)(0)))) - F (s, + const.
(6)
Given the connection between Q-functions and successor features (Eq, 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 the span of Φ, as we detail at the end of this section. Since = is the successor features of Φ, with the terminology introduced in Sec. 2 & is used both as a state encoder, i.e., to embed states passed to the predictor, and as task encoder, i.e., defining a space of reward functions.
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 (s, a)
+
y FF (s',a') (Barreto et al. 2017) define temporal-difference version of the previous loss:
CID-JEPA(O, TT) -(sn,))--D.=-2a'=tt (1la)) (IT -)) (To')'))( =)) (7)
Unlike the Monte Carlo loss of Eq. [ LID-JEA only requires sampling one-sted step transitions and actions from the given policies, and it can thus be estimated from off off-colicy offline datasets.
3.2 TRAINING SEPARATE STATE AND TASK REPRESENTATIONS
While Eq. and 7 the same encoder used for both state and task representations, these need not be the same in practice. Consider, for instance, 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 abstract: having flexibility over the dimensionality and content of each representation would be desirable. We thus now introduce an asymmetric variant that trains a distinct encoder v/s : S - Rd define the of reward functions of interest (i.e., task encoder). We first redefine the predictor as To Rdd x A x Z + Rd and the latent-predictive Monte-Carlo loss to train é and T' as
(8) such that To maps states encoded through @ to the long-term dynamics of a policy m, the la- tent space induced, this time, by vb Similar to Prop. T' approximates the successor features F (s, a) of d in the latent space induced by ø. Symmetrically, we train & together with an addi-
LMC-JEPA (d,, Tø, vb) = B(s.a))D. MM= (d(s), a, -++))1]]]
tional predictor Te Rd x A Z Rdo. To do so, follow existing literature according to which joint representations should predictive of each other (Guo et al. 2020 Tang et al. 2023) and train & and TV, through the same latent predictvv loss with the roles of and & inverted, i.c.,
LMC-JEPA(1) Te. ②) As before, can then design off-policy TD variant of this loss,
CID-EPA (0, To, 2/)
E
(s.0,e")D
[[Tg( 2) - (is) - -Toloooolosss
(9)
so that é and To are optimized via LTD-EEPA (ò, To, v)), while vb and Ts via LID-JEPA (v. T).
3.3 TD-JEPA REPRESENTATIONS FOR ZERO-SHIOT RL
The relationship between the learned predictors and successor features suggests a seamless instan- tiation of JD-JEPA zero-shot unsupervised RL algorithm. Redefining the policy parameter
While some existing works use forward-in-time sampling train one representation and backward-in-time for the other, we use two forward-in-time losses. We further discuss this difference in App. C
4

---

Published as a conference paper at ICLR 2026
Algorithm 1 TD-JEPA for zero-shot RL
Inputs: Dataset D, batch size B, regularization coefficient A, networks m, Th. Φ. To. p.
Initialize target networks: T To. Φ t
T To,
xp
while not converged
b Sample training batch
{(5t, 0i <f)]== ~ D, (2i) ~ 2, faahia
{ (6(), 21) )= N=1
b Compute latent-predictivv losses
ETDAEEA (6, To. 1/= 2B E;
To(((s)), ais os, (s() y1's (6 (5
ETDAE (v) To. 6) = 2 Σ
To: (1) ( 21)
x (6]) TTx (v (().
D. Compute orthonormality regularization losses
(REC( 1 Σ (xts(0(ss) ((s;))² B Σ, ((s;) T (sss)
(REG(E)) 28(B-T) (v(ss)) "<:;)("(<))))))))))))>>))))>>))) (x;))) B E, (s:)) (+ss))))
Compute actor loss (a.) (+sss = E.e (s())) as, TT)
Lactor(F)
Update o, Tр minimize ETDJERR (o. (0, To, 1) XL XCREO (Ø)
Update v. T. to minimize LIDAEZA (1), Tw. 6) ACRRG (1)) Update x minimize (actor(T)
Update target networks & To N.T5 TJ via EMA of Φ. To. vb To
space 2 as the task embedding space (i.e., E C Rdv), ) we train latent policies such that (d(s)) = argmax, To (Ø(s), 2, T) for all 2 E 2 Since (o()),,, ~ F (s, a) (Proposition 1). this produces optimal policies for all rewards in the span of th, learned directly from state representa- tions Φ(-). At test time, given an inference dataset of rewarded samples Drwd = {(s,r)) r)). the optimal policy т. can be retrieved by computing 2, through linear regression, e.g. through the closed foom solution 2, argmin, ((s.)))D,,ol [[r (ss) 2) = EgrDaal ())=)).)((()))))) (s) ( (s))))) Alg. | describes TD-JEPA, which combines LTD-JEPA with stabilization strategies, e.g. target net- works and covariance regularization. We remark that latent prediction is not auxiliary: is the core objective that trains encoders and predictors, from which zero-shot policies can be directly distilled.
4 THEORETICAL ANALYSIS
We now provide some theoretical arguments showing how latent- preditttediitive temporal difference rep- resentations capture the long-term dynamics of given set of policies in a way that makes them amenable to zero-shot RL. Following Tang al.) (203),, with linear predictors. We view the representation Φ (resp. 1/) as a S x da (resp. S x dly) matrix, and consider action-free predictors To.z (resp. Te.z) de x de (resp. dæ do) matrices for all z. The expression (Ø((s)) a, 2) in Eq.
consider simplified tabular setting
8 Snnd 9 thus reduces TT ((s), while (s')ss |s, a) and P(s'|s, (s']s,,) are replaced by = (s'[s) M (s' s, (s)) and pr: (s'ls) = P(s'|s ss (s)))).
Monte-Carlo losses.
We define (non-latent- successor measure approximation loss
Lsva (g. (T:).. v)= R.- BI 1
(10)
Minimizing Lsm is equivalent to finding the best multilinear approximation to the successor mea- sures 5 We prove the following connection with the Monte Carlo latent-predictive loss of Eq g
T Theorem 1. For fixed & and w. let T 22, v.2 be the optimal predictors for LsM(Ø, Te, v)) (Eq. 10p, LMC LMC.JEF to),, EMC.JEPA T Ty.2, ②) (Eg. 8. respectively. If (A1) & 16 db = I,
(A2) the state distribution uniform, and (A3) for all z E 2, the matrix p*= is symmetric, then
This decision additionally grounds &b as task encoder, and breaks the symmetry that could arise from the two encoders and di being trained through similar parent-predictive objectives.
5

---

Published as a conference paper at ICLR 2026
1. for all 2, 6T² = =. ILg vp and T. = (T;))² = I where Hlg (resp. I1x))
is an orthogonal projection on the spun of & (resp. v);
2. V &oLMC.JEPA(O T,, 1) = VoLsM(0. Te, 4) and CMC.JEE (2. = V
This result reveals that 1) the optimal predictors for the successor measure loss LsM and the latent- predictive loss JEPNEPA match, and yield an orthogonal projection of the successor features a onto the ē space; 2) the gradients w.n.t. the representations Φ and th, when evaluated at any pre- dictor, match among these two losses, showing that gradient descent on LMC.JEPA would update representations in the direction that reduces LSM. hence improving the approximation of the succes- sor measures. This result follows as a special case of novel theorem (see App./C generalizing and implying all previous guarantees for latent-predien predictive representations (Tang al. 2023 Khetarpal ett al. 2025 Voeleker et al. 2024; Lawson et al.) 2025). which we believe is of independent inter- Finally, we remark that, while the assumptions AT-A3 have been considered in all these related works, they can be relaxed, at the price of more involved proof' and notation, as shown in App. C
Temporal--lifference losses. We first derive non-collapse guarantee. While a similar result was originally proved by Tang et al. (2023) for the one-step loss (Eq. B). our case is more complex since TD iatent-prediction can be seen as "doubly latent- aatent-predictvee (cf. Eq. 9 T ((s) is optimized to
match a representation being learned - (++) ppuu bootstrapped version of itself - T 2.= ((s+).
Theorem 2. Let Φe and the be the representations learned under a continuous-tin time relaxation of Eq. wwere, each step t, the optimal predictors for (òr, 1hx) are first computed and then a gradient step on (ér, 1/2) is taken (see App. Bor the explicit formulation). Then, the covariance matrices &. te and v the are constant over time, i.e., of Φr 60 oo and 1/ 1/2 vo tho for all 2 O.
This result suggests that, if predictors are trained at a faster rate than representations, the overall dynamics preserve their covariance, thus preventing ø and 1é from collapsing trivial solutions (e.g.. Φ = A 0) when properly initialized, e.g., with unitary covariance.
As done for MC objectives (Th.[). we now show that the latent- predittpre loss of TD-JEPA is related to forward and backward TD losses for approximating the successor measure (Blier et al. 2021).
Theorem 3. Consider the following TD losses for approximating the successor measure
[IRTTV - p== > 2
(ib( T2, 1)) 2
Bz-z
(11)
(o, x), = 1
Eznz
[IVT:OT -
"
TT
(12)
2
fixed (0, v). let T 2,w T b.w 20 Th respectively be the optimal predictors for Liw (0, Te, v), T2, 6). Under the same assumptions as Th Th.
Low (0,T,,,) ). LID-JEPA (, Tz, )) LID-EEAA (vb, (v,,T,,
1. for all z, &T (resp. I1p.2) is an oblique projection on the span of à (resp. vb);
2,t M and TT T/ -M=
oo where No.
2. VoLTD-JEPa() Te, v) V.C (0, Te, 4) and V &TD.JEPA (v) Te, Φ) C. Te, v.
Similar to Th. I the optimal predictors and gradients of TD-JEPA match those of the non-latent- predictive TD losses Eq. I and 12, which are known to recover an approximation of the successor measure for bilinear parameterizations of the form F. B (Blier et al. 2021). Unlike in the Monte Carlo case, here the optimal predictors solve a least-squares TD problem (Boyan 1999 Precup et al,] 2001], yielding the fixed point of projected Bellman operator whose closed-form expression is an oblique projection (Scherrer 2010).
Policy evaluation and zero-shot RI. Finally, the following result motivates the significance of optimizing the successor measure losses of Eq. T0 11 and [2]
Theorem 4. Let d, va have identity covariance matrices. For any reward function r, let wr (x))) - be the linear regression weight for representation bb. Then, for any To.
- ((8) -(v)()))oo))TT TTui)) < 2LsM (o, T2, 1)).
max I
EZEZ LEES
Moreover, (o T2, cLtw (0) (H. and Lsx (0. T2.4) CLow(sh T2, 1b) for some c.

---

Published as a conference paper at ICLR 2026
Paraphrasing, the policy evaluation error of the technique in Section 3.3 (i.e., embed into a vector w through linear regression on v. and compute (o((s)) 2)
T)) w) is bounded by the successor measure approximation loss and the corresponding TD errors. Both these quantities are indirectly optimized by TD-JEPA (Th. T 3), which is thus a sound approach for zero-shot policy evaluation. Moreover, Th. + leads to zero-shot optimality result analogous to Theorem 2 of (Touati & Ollivier [021]: if the approximation of _M is perfect (i.e., = = oT,²/" for all a or, equivalently, the TD errors in Eq. and are zero) and the policies a. are optimal for all linear rewards in eb, then the inference
procedure above recovers optimal policies for any (even non-linear) reward function.
RELATED WORK
Zero-shot RL algorithms. Methods that pre-train agents on unsupervised data to enable zero-shot solution of wide range of downstream tasks have achieved impressive results, yielding so-called behavioral foundation models (Pirotta et al. [2024] Tirinzoni et al., (025). The forward-backward al- gorithm (FB. Touati & Ollivier (202 Touati al. (2023)) is established method, and perhaps the most related to TD-JEPA. FB learns a task encoder and estimates its successor features, essentially finding& bilinear decomposition of policy-conditional successor measures (e.g., M*= R F. B") On the other hand, TD-JEPA uses the parameterization M & oT2b", thus training shared (across tasks) state representations. Moreover, FB adopts a contrastive loss, which computes pairwise dot products across each training batch, while TD-JEPA is non non-contrastive at its core. FB has been fur- ther shown capable of zero-shot imitation (Pirotta et al., 2024) and extended to several settings, in- cluding online training regularized by action-free expert data (Tirinzoni et al. 2025), offline training on low-quality data (Jeen et al.) 2024b), training on environments with different dynamics (Bobrin) eet al. 2025), online fine-tuning (Sikchi et al. 2025) and pure exploration (Urpí et al. 2025 Other methods, like HILP, PSM, and RLDP, can also be seen as training a task encoder v plus succes- sor features top. HILP (Park et al. (024) trains th through distance- pewennceeeeeeeiiinggggg "goal-reaching" loss, while PSM (Agarwal et al. (2025) learns an affine decomposition of the successor measure for a discrete codebook of policies, and RLDP (Jajoo al. 2025) trains to using chained multi-step latent prediction (Hansen et al. 2024). Jajoo et al. (2025) also observe that regularizing the representation to be orthonormal is crucial to avoid collapse, which we also observe in TD-JEPA (see Alg. D..
Latent-predictive methods. Latent-predictive methods have mostly been applied to define auxil- iary losses for variety of RL settings. [Schwarzer al. (202T) use a prrent--redictiee loss to enhance state representations learned through a deep Q network. Guo et al. (2020) use latent prediction in POMDPs to encourage two representations (of observations and histories) to be self- -eef-predictive of each other, similarly to the asymmetric variant of TD-JEPA and the method explained in Appendix
Hansen et al. (2024) and [Sobal et al. (2025) train a latent dynamics model that enables test-time planning to improve pre-trained policy or solve goal-reaching tasks, respectively. BYOL- (Law-) son et al., 2025) trains representations by predicting discounted future latent states visited by the behavior policy. BOLL-y may thus be seen as unconditional, Monte Carlo version of TD-JEPA, and off off-policy. The on-policy nature of the algorithm enables Lawson et al. (2025) to implement bi-directional update of asymmetric representations. TD-JEPA can also recover an asymmetric parameterization, but its practical objective is not bi-directional (i.e., it only implements forward TD prediction, cf. Appendix C for formal definition of the bi-directional objective). Crucially, BYOL-y is not proposed as a zero-shot method: the version we evaluate is novel instantiation in successor-feature framework.
which is instead
Theory of latent-predictive representations. The theory of latent-predictive representations has been previously studied in several works (Tang et al.) 2023: Voeleker et al., 2024, Khetarpal et al., 2025 Lawson et al.) (025). with a particular focus on single- single-step prediction (poten- tially, bi-directional). Our analysis of MC-JEPA (Section 4 largely takes place in a multi-policy setting, with generic transition kernels over states; as such, it subsumes and expands on several ex- isting results (see Appendix C.2). On the other hand, representation learning through TD losses, as in TD-JEPA, is largely understudied. The closest studies (Blier et al.) 1021 Lan et al.) 2023) show that, under certain parameterizations and assumptions, TD representation learning can recover low-rank decompositions of the successor measure, (i.e. it optimizes the corresponding approxima- tion loss). While these works rely on having a single policy, we provide a first result connecting latent-predictive TD learning with TD learning over the successor measure for multiple policies.
7

---

Published as a conference paper at ICLR 2026
Laplacian
ICVF
HILP
FB
RLDP
BYOL*
BYOL-
TD JEJA
DMC (avg)
293.1 1511
438.7 149
391.2 = 25.8
456,2 8.
525.7 #: 133
513.8 = 116
582.4 =
628.8 #
walker
3004 200
5349 61.3
4228:43255
=
=
646.3 365
7389 35
cheetal:
24244 6
3949 = N0.1
3300 1: 95.6
6224 = 131
0153 1: 23.5
48.0 46.7
676.5 11.1
7063
quadzuped
323
5833
5129
=
5818 165
$709
6267 136
pointmass
1904 124
2416 158
294. 1: 33A4
4025 = 16.
50033 1: 12.0
4103 S 8.5
431A 17.4
443.7 18.9
(avg)
591.1 183
619.3
620.1
648.2
610.2 ::
618.6 10.5
645.4 10.5
661.2
walker
266.7 1:47
7210
7964 1:
R11A
7219 1: 16.3
7468 1100
1:96
7852 1. K7
cheetal:
6165 18.9
60633
6183
6727
515.6 4 449
622.
602299
658.7
quadruped
6.8.0.3887
7063 14.1
MALB 1. 111
9956 1.461.90 9.1
06500 1: 13.9
6118 281
MMM 161
691.4 1 5.
pointmaas
1451 4:
4153 11.1
3710 37.1
$130 = 30D
676.1 194
a: 413
17.6
479.5 21A
OGBenchaan (avg)
30,58 0.81
25.22 0:55
32.56 0.92
39.89 = 0.47
39.09 +:
40.33 0.52
41.58 0.64
41.34 6.45
antmaze-mn
92:20 1: 2.91
65:90 182
BA.00 4: 3.99
96.80 : 0.74
97.00 1: 0.59
94.40 d: 1.48
98,0 6.73
96.67 1.1
antmaze-In
antmaze-mo
60.30 3.88
46.30 2.74
2.22
96.00 = 2.05
90.10 1.97
N600 7.10
$4.10 1: 385
antma za
225
171
ontmaze--mi
0.0 0 0.00
0.00
0.20 0.30
1,80 L/W
0:00 =: 0.44
1.30 1/0
XI0 d: 1.96
0.20 0.20
cube single
3.5
3480
56.40
62.00 =
(3
3.34
67.50 3.07
cupe-dousa
1/0 =: 0.72
0.80 2: 0.44
1/4 = 0.58
1.30 0.1
2.30 :: 1.31
2.00 :: 6.5
1.40 0.67
3 = 0.91
scene puzzle-3x3
250 = L12
8,40
5
1.92
14.30
2,00 1: 1.4
1.20 0.44
2 9.99
2.00 0.79
2 6.79
1/00 d: 0.40
0.00 d: 0.31
2.49 0.03
OGBench (avg)
14.81 1.32
30.87 0:058
37.98 = 1.11
39.04 :: 0.66
27.07 # D.B
26,42 =: 0,83
30:42 0.94
37.98 6:77
antmaae mn
50.00 + 4.94
79:00 2.42
$3.00 +:
13,00
1.55
70.40
antmeze-In In
21.60 + 3.00
$8.40 1390
$2.00 +: 3.56
36:80 = 4.25
36.40 4.65
26.50 +:
21:90 3.57
57.30 4.25
antmeze-ls
11.50 1.47
13.20 1.64
12:20 # 1.75
49.30 = 5.64
19.00 2.73
+ 4.28
20:20 1.80
40.00 2:1
antmaze-me
cube-aingle le
1511 1 1:10
20:00 183
74.20 1.51
48A0 1.81
1: 7:41
d: 1.16
79:00 2,33
14.30 7:0
cube double
200
5.00
2000
2:60
scene
700 8: 1:24
45.40 1.20
63.00 1: 1.00
17.80 1.61
11 01
15,40 d: 1.37
14:40 1: 7.37
18.44 1: 1.37
puzzle
16:00 003
2,90 0:68
14:00:00.00
13,30
0.44
15.00
Table I: Performance of zero-shot algorithms for DMC (reward) and OGBench (success rate) with either proprioception or RGB inputs. We report means and standard errors across seeds. Numbers are bold for top algorithms if confidence intervals overlap.
EXPERIMENTS
We benchmark zero-shot performance across diverse set of problems, including 4 locomotion/nav-
igation domains from EXoRL/DMC (Tassa et al. 2018 Yarats et (2022a), as well as 9 navigation/- manipulation domains from OGBench (Park et al.) 2025a). The former suite involves reward-baser tasks and high-coverage data, while the latter evaluates goal-reaching and provides low-coverage " We consider both proprioceptive and bixel-based variants of all domains, and report ex- pected returns/success rates across set of tasks (4-8 depending on the domain) as main evaluation metric. In DMC, we often normalize returns by the maximum achievable (1000).
We structure our evaluation in four parts: (i) a comprehensive evaluation of TD-JEPA with respect to existing zero-shot methods; (fi) an ablation over the prediction target, measuring the impact of multi- step. policy-awan dynamics modeling; (iii) comparison of TD-JEPA its symmetric variant that learns a shared state-task encoder Φ; and (iv) a demonstration of fast adaptation from pre-trained state representations. Further results are presented in App. D and implementation details in App. B
How does TD-JEPA compare to zero-shot RL algorithms? We first compare TD-JEPA to three groups of wew.ccessoo--fturee-bee
zero-shot RL baselines
Laplacian (Wu et al. 2019), HILP (Park et al. 2024). and FB (Touati & Ollivier 2021 are es- tablished zero-shot methods that train task encoder e,, without specific learning objectives for state encoder.
BYOL* (Grill al. 2020), BYOL-y* (Lawson et al. 2025 and RLDP (Jajoo al. 2025 learn state encoder & via prrent-rredictive learning, which we then use as a task encoder for successor features (learned through a contrastive loss in the case of RLDP).
ICVF* (Ghosh et al. 2023) learns a multilinear decomposition of the successor measure via expec- tile regression, yielding both state and task encoders on top of which we train successor features.
For 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
We additionally apply BC regularization in OGBench based on Park al. 2025h), detailed in App. E.6 Notice that only Laplacian. HILP, FB and RLDP are standard zero-shot unsupervised RL algorithms, while BYOL, BYOI.-7, and ICVF (henceforth marked with +) are representation learning methods: their instantiation
zero-shot framework is novel and designed to investigate the impact of different representations.
8

---

Published as a conference paper at ICLR 2026
P(X+Y)
RGB
P(X-Y) Proprioception
TDJEPA
50:0:62=7:67
+3:87=4:8221
50:0:65
5:55 =5 (5766665:4
BYOL Y 3827 5924::924:61110
35
50.00
5426
73777 1927 5728 79
BYOL 3325 41 5030
522639723
2345
4626
5020 5317
33:4
2926 4826
7336
RLDP
30=5
36:4 48
50*0
42+7
50+0
37 3845
X
FB 20+4
33+5 43+5 45+5 50:0.64
45 469
52+4
HILP
1923
19 23 2643 3647 5020:5655
4316
71
51
5020
ICVF*
13a4 21 22a6 23 31 64 50 0:
2544 43 8280 5749 43 49 50
Laplacian
18=3 17+5 1845 20x3 27-6 33+6 42+5 50
7=4 21 27-6 25 15a4 18 24 50+0
TDJJPA BYOL-y*
BYOL*
RLDP
FB
HILP
ICVp
Laplacian
TD-JPA
BVOL-y*
BYOL*
RLDP
FB
ICVF HILP
Laplacian
Figure 2: Probabilities of improvement: how likely is method X to outperform method Y on
random domain? We report symmetrized 95% simple bootstrap confidence intervals. Dotted lines surround matches in which the improvement is statistically significant.
Normalized performance
0.6
RGB
Proprioception
antmaze-ls
soene
walker
0.4
pointmass anlimaee my cheetal cube-double
0.2
antmaze-mn
puzzle-3x3
nitmaze-mee
cube-single antmaze-In
0.0
DMCRca
DMC OGBencchaggGGeench
0.0
0.1
-0.1 0.0 0.1
BYOL
R
BYOL- TD-JEPA
Normalized performance difference
Figure 3: Left: normalized zero-shot performance for latent- -reni-pred methods, Right: difference in normalized performance between TD-JEPA and its symmetric variant. Error bars represent stan- dard errors on normalized performance or its differences, respectively.
before being fed into, e.g., the successor features estimator F(s, a; [ We find that this protocol results in significant improvements in phot performances, even for existing methods (e.g.. 1.3x and 2.4x higher than overlapping pixel- based results for the methods presented in Park et al. (2024) and Jajoo et al.] (2025), respectively), as displayed Tab [ 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, naviga- tion, manipulation) and datasets (high- 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 improve- ment (Agarwal et al., 1021) across all domains in Fig. 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 proprio- ception, significantly better than them in visual domains. Similarly, BYOL- slightly better than TD-JEPA in OGBenchRGB, but it is significantly worse in DMCRGB and OGBench. Finally, we note that latent-predictive methods tend be generally preferrable pixel-based domains.
zero-shot algorithms model? The baselines based on BYOL and BYOL-y are algorithmically closest to TD-JEPA, and allow a precise investigation on the dynamics to model. While BYOL* and BYOL-?" approximate one-step and multi-step transi- tions of the hehavioral 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-conditions successor measures is on average beneficial, reported in Fig. 3 (left).
Which dynamics should latent-
Should state and task representations differ? TD-JEPA trains separate state and task encoders: while this may grant better approximation of successor measures, sharing state and task repre- sentations while optimizing a single objective (see Section 3.11 may in practice be more efficient. We measure the difference in per-task normalized performance between TD-JEPA and symmetric
'On average, explicit state encoders actually improve the performance for existing methods, see App. D.I

---

Published as a conference paper at ICLR 2026
Normalized Performance Normalized Pecformance
walker
cheetah
quadruped
pointmass
0.50
0.4
Offline
0.5
0.25
0.2
0.0
0.00
0.0
0K
200K
400K
200K
400K
200K
400K
0K
200K
400K
Steps
Steps
Steps
walker
cheetah
quadruped
pointmass
Online
0.50
0.50
0.4
0.5
0.25
0.2
0.00
0K
200K
400K
200K
400K
0K
200K
400K
0K
200K
400K
Steps
Steps
Steps
Steps
TD-JEPA TD-JEPA (frozen)
FB
FB (frozen)
Scratch
Figure 4: Normalized performance of zero-shot policies when fine-tuned offline (top) online (bot- tom). Initializing the agent to zero-shot solutions (blue and yellow lines) results in sample-efficiee learning; frozen representations (dashed) are often expressive enough to enable fast adaptation.
variant 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 improve empirical performance more often than not.
Are state representations beneficial for fast adaptation? While the previous evaluations have focused on aggregated zero-shot performance, we now investigate an additional benefit explicit state representations: fast adaptation at test-time. Given pixel-based task, initialize the agent with the zero-shot policy π. and critic learned at pre-training and we either fine-tune the whole model via TD3 (Fujimoto al., 2018) keep the pre-trained state encoder frozen. We consider two RL adaptation protocols (i) Offline: a transition-reward dataset is provided Drew = {(x, a, 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 (following the unsupervised-to-one
protocol Kim et al. (004)). Fipure/4reports 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.. training from scratch, and reaches the asymptotic performance of TD3. More interestingly, frozen representations often sufficient for downstream learning, and do not need further fine-tuning. We refer to App. D.3 and App. E.7 for further results and details, respectively.
7 CONCLUSION
Through the introduction of novel lemporal-difference latent- prreettpreeittve loss, we presented zero- shot unsupervised RL method that operates entirely in latent space and can be shown to recover factorization of the successor measures of multiple policies. Our method tackles fundamen- tal representation learning problem for control, and highlights a connection between downstream performance and accurate modeling of the successor measure. We thus suggest that flexible repre- sentations for RL, particularly for value estimation and optimization on downstream tasks, should be predictive of future behaviors, and precisely capture their diverse and long-term nature. 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 rection for future work may study learning objectives that are compatible with asymmetric successor measures, yet remain amenable practical optimization. On practical note, we believe that bench- marking latent- predictve zero-shot objectives on large-scale, real robotic dataset can shed further light on opportunities and limitations of this promising framework.
assumption of symmetry, one exciting di-
10

---

Published as a conference paper at ICLR 2026
REFERENCES
Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron C. Courville, and Mare Bellemare. Deep reinforcement learning the edge of the statistical precipice. NewrIPS, 2021.
Siddhant Agarwal, Harshit Sikchi, Peter Stone, and Amy Zhang. Proto successor measure: Repre- senting the behavior space of an rl agent. ICML, 2025.
Mido Assran, Adrien Bardes, David Fan, Quentin Garrido, Russell Howes, Mojtaba, Komeili, Matthew Muckley, Ammar Rizvi, Claire Roberts, Koustuv Sinha, Artem Zholus, Sergio Arnaud, Abha Gejji, Ada Martin, Francois Robert Hogan, Daniel Dugas, Piotr Bojanowski, Vasil Khali- dov, Patrick Labatut, Francisco Massa, Mare Szafraniec, Kapil Krishnakumar, Yong Li, Xiaodong Ma, Sarath Chandar, Franziska Meier, Yann LeCun, Michael Rabbat, and Nicolas Ballas. V-jepa 2: supervupee video models enable understanding, prediction and planning. arXiv preprint wwXii::5006.0998.. 2025a.
Mido Assran, Adrien Bardes, David Fan, Quentin Garrido, Russell Howes, Matthew Muckley, Am- mar Rizvi, Claire Roberts, Koustuv Sinha, Artem Zholus, al. V-jepa 2: Self-supervised video models enable understanding, prediction and planning. arXiv preprint 2rXiv:2506.00990 2025b.
André Barreto, Will Dabney, Rémi Munos, Jonathan J Hunt, Tom Schaul, Hado van Hasselt, and David Silver. Successor features for transfer in reinforcement learning. NeurIPS, 2017.
Chethan Bhateja, Derek Guo, Dibya Ghosh, Anikait Singh, Manan Tomar, Quan Vuong, Yevgen Chebotar, Sergey Levine, and Aviral Kumar. Robotic offline rl from internet videos via value- function pre-training. arXiv preprint arXiv:2309. 13041, 2023.
Léonard Blier, Corentin Tallec, and Yann Ollivier. Learning successor states and goal-dependent values: A mathematical viewpoint. arXiv preprint arXiv:2101.07123, 2021.
Maksim Bobrin, Ilya Zisman, Alexander Nikulin, Vladislav Kurenkov, and Dmitry Dylov, Zero-shot adaptation of behavioral foundation models to unseen dynamics. arXiv preprint arXiv:2505.13150, 2025.
Justin A Boyan. Least-squares temporal difference learning. ICML, 1999.
Lasse Espeholt, Hubert Soyer, Remi Munos, Karen Simonyan, Vlad Mnih, Tom Ward, Yotam Doron, Vlad Firoiu, Tim Harley, lain Dunning, et al. Impala: Scalable distributed deep-rl with importance weighted actor-learner architectures. ICML, 2018
Scott Fujimoto and Shixiang Shane Gu. A minimalist approach to offline reinforcement learning. Advances in neural information processing systems, 14-20132-20145 2021.
Scott Fujimoto, Herke Hoof, and David Meger. Addressing function approximation error in actor- critic methods. ICML, 2018.
Scott Fujimoto, Pierluca D'Oro, Amy Zhang, Yuandong Tian, and Michael Rabbat. Towards general-purpose model-free reinforcement learning. ICLR, 2025.
Carles Gelada, Saurabh Kumar, Jacob Buckman, Ofir Nachum, and Mare Bellemare. Deepmdp: Learning continuous latent space models for representation learning. ICML, 2019.
Dibya Ghosh, Chethan Anand Bhateja, and Sergey Levine, Reinforcement learning from passive data via latent intentions. ICML, 2023.
Bean-Bastien Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan Guo, Mohammad Gheshlaghi Azar, et al. Bootstrap your own latent-a new approach to self-supervised learning. NeurIPS, 2020.
Zhaohan Guo, Shantanu Thakoor, Miruna Pislar, Bernardo Avila Pires, Florent Altché, Corentin Tallec, Alaa Saade, Daniele Calandriello, Jean-Bastien Grill, Yunhao Tang, et al. Byol-explore: Exploration by bootstrapped prediction. NewrIPS, 2022.

---

Published as a conference paper at ICLR 2026
Zhaohan Daniel Guo, Bernardo Avila Pires, Bilal Piot, -ean-Bastien Grill, Florent Altché, Rémi Munos, and Mohammad Gheshlaghi Azar. Bootstrap preett-predictive representations for multi- task reinforcement learning. ICML, 2020.
Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. ICML, 2019.
Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. ICLR, 2020.
Nicklas Hansen, Hao Su, and Xiaolong Wang. Td-mpe2: Scalable, robust world models for contin- uous control. ICLR, 2024.
Tairan He, Jiawei Gao, Wenli Xiao, Yuanhang Zhang, Zi Wang, Jiashun Wang, Zhengyi Luo, Guanqi He, Nikhil Sobanbab, Chaoyi Pan, et al. Asap: Aligning simulation and real-world physics for learning agile humanoid whole-body skills. arXiv preprint arXiv:2502.01 143, 2025.
Pranaya Jajoo, Harshit Sikchi, Siddhant Agarwal, Amy Zhang, Scott Niekum, and Martha White. Regularized latent dynamics prediction is strong baseline for behavioral foundation models. Workshop on Reinforcement Learning Beyond Rewards @ RLC 2025, 2025.
Scott Jeen, Tom Bewley, and Jonathan Cullen. Zero-shot reinforcement learning from low quality data. Advances in Neural Information Processing Systems, 37: 169441 2024a.
Scott Jeen, Tom Bewley, and Jonathan Cullen. Zero-shot reinforcement learning from low quality data. NeurIPS, 2024b.
Alexander Khazatsky, Karl Pertsch, Suraj Nair, Ashwin Balakrishna, Sudeep Dasari, Siddharth Karamcheti, Soroush Nasiriany, Mohan Kumar Srirama, Lawrence Yunliang Chen, Kirsty El- lis, al. Droid: A large-scale in-the-wild robot manipulation dataset. RSS 2024 Workshop: Data Generation for Robotics, 2025.
Khimya Khetarpal, Zhaohan Daniel Guo, Bernardo Avila Pires, Yunhao Tang. Clare Lyle, Mark Rowland, Nicolas Heess, Diana Borsa, Arthur Guez, and Will Dabney. A unifying framework for action-conditional self selffpredictive reinforcement learning. A/STATS, 2025.
Junsu Kim, Seohong Park, and Sergey Levine. Unsupervised-tt-onlinee reinforcement learning.
arXiv preprint 2rXiv:2408.14780 2024.
Diederik Kingma. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.660 2014.
Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. NeurIPS, 2020.
Charline Le Lan, Stephen Tu, Mark Rowland, Anna Harutyunyan. Rishabh Agarwal, Mare G Belle- mare, and Will Dabney. Bootstrapped representations in reinforcement learning. arXiv preprint arXiv:2306. 10/71, 2023.
Daniel Lawson, Adriana Hugessen, Charlotte Cloutier, Glen Berseth, and Khimya Khetarpal. Self- predictive representations for combinatorial generalization in behavioral cloning. arXiv preprint arXiv:2506. 10137, 2025.
Yann LeCun. A path towards autonomous machine intelligence. Open Review, 2022.
Yitang Li, Zhengyi Luo, Tonghe Zhang, Cunxi Dai, Anssi Kanervisto. Andrea Tirinzoni, Haoyang Weng, Kris Kitani, Mateusz Guzek, Ahmed Touati, Alessandro Lazaric, Matteo Pirotta, and Guanya Shi. Bfm-zero: A promptable behavioral foundation model for humanoid control ing unsupervised reinforcement learning, arXiv preprint arXiv:2511.04133 2025.
Yecheng Jason Ma, Shagun Sodhani, Dinesh Jayaraman, Osbert Bastani, Vikash Kumar, and Amy Zhang. Vip: Towards universal visual reward and representation via value- calue-implicit pre- pre-training. ICLR, 2023.
12

---

Published as a conference paper at ICLR 2026
Sridhar Mahadevan and Mauro Maggioni. Proto-value functions: A laplacian framework for learn- ing representation and control in markov decision processes. JMLR, 2007.
Arjun Majumdar, Karmesh Yadav, Sergio Arnaud, Jason Ma, Claire Chen, Sneha Silwal, Aryan Jain, Vincent-Pierre Berges, Tingfan Wu, Jay Vakil, et al. Where are we in the search for an artificial visual cortex for embodied intelligence? NeurIPS, 2023.
Robert McCarthy, Daniel CH Tan, Dominik Schmidt, Fernando Acero, Nathan Herr, Yilun Du, Thomas G Thuruthel, and Zhibin Li. Towards generalist robot learning from internet video: A survey, Journal of Artificial Intelligence Research, 2025.
Suraj Nair, Aravind Rajeswaran, Vikash Kumar, Chelsea Finn, and Abhinav Gupta. R3m: A univer- sal visual representation for robot manipulation. CORL, 2022.
Abby Open X- X-embodiment collaboration, O O'Neill, Abdul Rehman, Abhiram Maddukuri, Abhishek Gupta, Abhishek Padalkar, Abraham Lee, Acorn Pooley, Agrim Gupta, Ajay Mandlekar, Ajinkya Jain, et al. Open x-embodiment: Robotic learning datasets and rt-x models. ICRA, 2024.
Maxime Oquab, Timothée Darcet, Théo Moutakanni, Huy Vo, Mare Szafranice, Vasil Khalidov, Pierre Fernandez, Daniel Haziza, Francisco Massa, Alaaeldin El-Nouby, et al. Dinov2: Learning robust visual features without supervision. arXiv preprint arX1v:2304.07/93, 2023.
Simone Parisi, Aravind Rajeswaran, Senthil Purushwalkam, and Abhinav Gupta. The unsurprising effectiveness of pre-trained vision models for control. ICML, 2022.
Seohong Park, Tobias Kreiman, and Sergey Levine. Foundation policies with hilbert representations. ICML, 2024.
Seohong Park, Kevin Frans, Benjamin Eysenbach, and Sergey Levine. Ogbench: Benchmarking offline goal-conditioned RL. ICLR, 2025a.
Seohong Park, Qiyang Li, and Sergey Levine. Flow q-learning. ICML, 2025b.
Matteo Pirotta, Andrea Tirinzoni, Ahmed Touati, Alessandro Lazaric, and Yann Ollivier. Fast imi- tation via behavior foundation models. ICLR, 2024.
Doina Precup, Richard S Sutton, and Sanjoy Dasgupta. Off-policy temporal-difference learning with function approximation. ICML, 2001.
Bruno Scherrer. Should one compute the temporal difference fix point or minimize the bellman residual? the unified oblique projection view. ICML, 2010.
Moritz Schneider, Robert Krug, Narunas Vaskevicius, Luigi Palmieri, and Joschka Boedecker. The surprising ineffectiveness of pre-trained visual representations for model-based reinforcement learning. NewrlPS, 2024.
Max Schwarzer, Ankesh Anand, Rishab Goel, R Devon Hjelm, Aaron Courville, and Philip Bach- man. Data-efficient reinforcement learning with self-predictive representations. ICLR, 2021.
Harshit Sikehi, Andrea Tirinzoni, Ahmed Touati, Yingchen Xu, Anssi Kanervisto, Scott Niekum, Amy Zhang, Alessandro Lazaric, and Matteo Pirotta. Fast adaptation with behavioral foundation models. RLC, 2025.
David Silver, Thomas Hubert, Julian Schrittwieser, Ioannis Antonoglou, Matthew Lai, Arthur Guez, Mare Lanctot, Laurent Sifre, Dharshan Kumaran, Thore Graepel, et al. A general reinforcement learning algorithm that masters chess, shogi, and go through self-play. Science, 2018
Sneha Silwal, Karmesh Yadav. Tingfan Wu, Jay Vakil, Arjun Majumdar, Sergio Arnaud, Claire Chen, Vincent-Piere Berges, Dhruv Batra, Aravind Rajeswaran, al. What do we learn from large-scale study of pre-trained visual representations in sim and real environments? ICRA, 2024.
Vlad Sobal, Wancong Zhang, Kyunghyun Cho, Randall Balestriero, Tim GJ Rudner, and Yann LeCun. Learning from reward-free offline data: case for planning with latent dynamics models. arXiv preprint 2502 /4819, 2025.
13

---

Published as a conference paper at ICLR 2026
Yunhao Tang, Zhaohan Daniel Guo, Pierre Harvey Richemond, Bernardo Avila Pires, Yash Chan- dak, Rémi Munos, Mark Rowland, Mohammad Gheshlaghi Azar, Charline Le Lan, Clare Lyle, et al. Understanding self- eeff-predictive learning for reinforeement learning. ICML, 2023.
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Bud- den, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, et al. Deepmind control suite. arXiv preprint arXiv: 00690, 2018.
Gemini Robotics Team, Saminda Abeyruwan, Joshua Ainslie, Jean-Baptiste Alayrac, Montser- rat Gonzalez Arenas, Travis Armstrong, Ashwin Balakrishna, Robert Baruch, Maria Bauza, Michiel Blokzijl, et al. Gemini robotics: Bringing into the physical world. arXiv preprint arXiv:2503.20022, 2025.
Andrea Tirinzoni, Ahmed Touati, Jesse Farebrother, Mateusz Guzek, Anssi Kanervisto, Yingchen Xu, Alessandro Lazaric, and Matteo Pirotta. Zero-shot whole-body humanoid control via behav- ioral foundation models. ICLR, 2025.
Ahmed Touati and Yann Ollivier. Learning one representation to optimize all rewards. NeurIPS, 2021.
Ahmed Touati, Jérémy Rapin, and Yann Ollivier. Does zero-shot reinforcement learning exist? ICLR, 2023.
Nikolaos Tsagkas, Andreas Sochopoulos, Duolikun Danier, Sethu Vijayakumar, Chris Xiaoxuan Lu, and Oisin Mac Aodha. When pre-trained visual representations fall short: Limitations in visuo-motor robot learning. arXiv preprint arXiv:2502.03270 2025.
Núria Armengol Urpí, Marin Vlastelica, Georg Martius, and Stelian Coros. Epistemically-guide forward-backward exploration. RIC,, 2025.
Claas Voelcker, Tyler Kastner, Igor Gilitschenski, and Amir-massoud Farahmand. When does self- prediction help? understanding auxiliary tasks in reinforcement learning. RLC, 2024.
Manuel Watter, Jost Springenberg, Joschka Boedecker, and Martin Riedmiller. Embed to control: A locally linear latent dynamics model for control from raw images. NeurIPS, 2015.
Haoyang Weng, Yitang Li, Nikhil Sobanbabu, Zihan Wang, Zhengyi Luo, Tairan He, Deva Ra- manan, and Guanya Shi. Hdmi: Learning interactive humanoid whole-body control from human videos. arXiv preprint 1rXiv:2509.16755 2025.
Yifan Wu, George Tucker, and Ofir Nachum. The laplacian in rl: Learning representations with efficient approximations. ICLR, 2019.
Denis Yarats, David Brandfonbrener, Hao Liu, Michael Laskin, Pieter Abbeel, Alessandro Lazaric, and Lerrel Pinto. Don't change the algorithm, change the data: Exploratory data for offline reinforcement learning. Generalizable Policy Learning in the Physical World Workshop @ ICLR 2022, 2022a.
Denis Yarats, Rob Fergus, Alessandro Lazaric, and Lerrel Pinto. Mastering visual continuous con- trol: Improved data-augmented reinforcement learning. ICLR, 2022b.
Seonghyeon Ye, Joel Jang, Byeongguk Jeon, Se June Joo, Jianwei Yang, Baolin Peng, Ajay Man- dlekar, Reuben Tan, Yu- Wei Chao, Bill Yuchen Lin, et al. Latent action pretraining from videos. 2025.
Weishuai Zeng, Shunlin Lu, Kangning Yin, Xiaojie Niu, Minyue Dai, Jingbo Wang, and Jiangmiao Pang. Behavior foundation model for humanoid robots. arXiv preprint arXiv:2509. 13780, 2025.
Gaoyue Zhou, Hengkai Pan, Yann LeCun, and Lerrel Pinto. Dino-wm: World models on pre-trained visual features enable zero-shot planning. arXiv preprint arXiv:2411.04983, 2024.
Gaoyue Zhou, Hengkai Pan, Yann LeCun, and Lerrel Pinto. Dino-wm: World models on pre-trained visual features enable zero-shot planning. 2025.
2025.
14

---

Published as a conference paper at ICLR 2026
A EXTENDED RELATED WORK
As TD-JEPA bridges zero-shot reinforcement learning and latent-predictive representation learning, reserve this section to building connections several related works in both areas beyond what was discussed in Sec. [addd
Other representation learning methods for RI. Beyond datent-predictinn methods, ICVF (Ghosh et al. 2023 Bhateja et al., (202)) has also been proposed as multi-poliee multi-step representation learning objective. By relying on an implicit, value-based loss, this method may be applied to off- policy, action-free, cousee data to recover a decomposition of policy-dependent successor measures. This decomposition is multi-linear, but, unlike TD-JEPA, it is restricted to linear pre- dictors; its practical implementation moreover forces the same dimensionality across state and task embeddings. As B ICVF does not natively support zero-shot RL: we thus integrate it into a successor-featuurs--beee based zero-shot policy optimization scheme.
VIP (Ma et al. 2023) also works in a similar setting as ICVF. It casts representation learning as an offline goal-conditioned reinforcement learning problem. This be seen related approximating the successor measure of several goal-reaching policies. On the other hand, TD-JEPA does so for all reward-maximing policies, hence encompassing the goal-reaching ones.
MR.Q (Fujimoto et al. 2025) learns model-based representations that approximately linearize value functions for reward-based RL. This is achieved by combining reward prediction with a single-step latent dynamics loss similar to the one of BYOL, yielding state encoders on top of which value functions and policies are trained. TD-JEPA also aims at linearizing value functions, but does so with multi-step policy-dependent loss and, most importantly, in a reward-free manner.
More broadly, several recent works evaluate visual representations pre-trained from large- scale data (e.g., internet videos) in control problems (Parisi et al. 2022 Majumdar et al., [2023] Silwal et al.) 2024 Schneider al. 2024 Zhou et al. 2024] MMCarthyy et al., 2025 Tsagkas et al. 2025). While there no "best" method overall, performance
problem and data dependent (Majumdar al., 2023), representatiinns pre-trained or fine-tuned with RL-related objectives, like time and task aware- ness (Ma et al. 2023 Bhateja et al. [2023 TTagkas et al.) 2025), ) perform well in general. TD-JEPA aligns with this view, as it shows that visual representations pre-trained with multi-step and policy- dependent objectives are suitable for value estimation and optimization across multiple tasks.
Applications of Unsupervised RL Training on large. reward-free, multi-task offline datasets, which an important motivation for this work, progressively demonstrating its relevance for real-world applications. Current efforts in robotics control, and in particular manipulation, have identified in-the-wild videos as a promising source of information to tap into. These data sourees typically include demonstrations of several tasks, and are largely reward-free (Khazatsky et al.) 2025 Open X-embodiment collaboration et al., 2024). Recent works have pushed this direction sig- nificantly, either by leveraging action annotation (Sobal al. 2025 Assran et al. (025b) or latent actions YYe et al., (025)). Other works in this space focus on large teleoperation datasets, which are again multi-task and reward-free. This is particularly the case of modern VLAs (Team al. 2025
which have achieved impressive applications in the real world.
Important developments have also taken place in whole-body humanoid control. In this area, agents are normally trained on motion-capture datasets, which can be regarded as multi-task unlabeled data, with recent works achieving impressive results on real hardware (Zeng al. 2025 He al. 2025 Weng et al., 2025). The main limitation of these methods is that they focus exclusively on imitating the behaviors in the data (e.g., via motion tracking). Unsupervised zero-shot RL methods have recently stepped in as a promising solution, enabling agents that can quickly solve tasks beyond imitation, like reward optimization and goal reaching. They have been successfully applied both simulation (Tirinzoni et al. 2025) and on real hardware (Li et al.. 2025).
15

---

Published as a conference paper at ICLR 2026
B PROOFS
B.1 PROOF OF PROPOSITION
For any 8, st, the term inside the expectation in Eq. 5 can be rewritten
[Tg(0(s), a.2) o(s+)|2 [TT()(s). = F. (s, Ø(st) )
[[T( a, - F as)
(s. ((s+)( )
- (TT(((s)) F (s,
x(s
Taking the expectation w.r.t. st ~ M (-1s, a). it is easy to see that the last term is zero since F (s,a) = MM- (1s,) ())()))((((()))))))))))))))))))) [0(+)]] while the second term is constant. This concludes the proof.
Remark 1. By replacing the target ((s)) above with ((s)) and (s, a) with F (s, a). this proof generalizes to the Monte Carlo loss with asymmetric representations (Eq. g).
(s.
B.2 PROOF OF THEOREM
We begin by rewriting the MC-JEPA loss of Eq. 1 with the notation of Sec.
LMC-JEPA (0, To 16) = E poup,e (11)( (TI],((s) 4(s
(13)
where p denotes the state distribution. Proposition 1 Implies that LMC-JERA has the same gradients as [ [ [o))(s))
Lme (d, To. 1)
(1t
Let D, SxS be diagonal matrix containing p(s) for all states S on its diagonal. Using that
Dp I by Assumption A2
Luc 2)) =
- Eznz [IID// MD 12 = Bznz [[[ØT= 1
(14)
We now prove all statements for Lme. as gradient equivalence with Eq. 13 1iimplies they also hold for LEC.EEA
Statement 1 Let us first compute the optimal predictors. For any E 2, the gradient of Lime T2, vb w.r.t T,
V To, v) ((2))"(011) MT p(=)(T= A=
where p(2) is the probability to sample - while the second equality that o by Assump- tion A1. This yields The = = T M tb. Moreover, by simply inverting the roles of Φ and vb, we find that Cmc (v, T2, Ø) p(=)(T= - th T Φ) and, thus, T5 + *
The gradient of (0 ( T2, 16) w.r.t is
Ls. Tz, v) = p(2) (OT,U = - MTM
where we used again Assumption A1. Hence, T = TT M Therefore, we clearly have that T T Moreover, since pM: is symmetric by Assumption A3, T symmetric too, and
T - - (MM' TT TTTM". G) = (TT)
Finally, it is easy to see that oT* and TT satisfy the stated expressions for II = TØT and 11 := Th respectively. Moreover, 1I, and 11 are symmetric and idempotent (11,11g = 1lg), hence they are orthogonal projection matrices. This proves the first part the statement.
We ignore the 1/S scaling that only multiplies the loss by a constant.
Without loss of generality, we also assume that p(z) > 0 for all If this is not the case, any with
p(2) can be removed from the loss.
16

---

Published as a conference paper at ICLR 2026
Statement 2 Let now fix any T3 for all E 2 and compute the gradients w.r.. & and sb.
VoLmo (2 1b) Ex~2 [ØT; M = +
V Suc (v, (v, Φ) = Eznz [TT; - = 6)T!]
V a T2, 16) = -2-z [0Tx5* = MoT = E. [ - M V (d T, 16) -Ez= (1 (M**) (MOT)
(v (
(
where used Assumption A1 simplify the last expressions. Given that prs and, thus, M are symmetric by Assumption A3, these gradients match as stated.
B.3 PROOF THEOREM 2
We begin by rewriting the TD-JEPA loss of Eq. 9 with the notation of Sec.
4 as
LTD-JEPA (0, To, vb) Benz. Pa ((s) (++))
7T To (
(15)
Following the proof of Theorem 1 we can put this in matrix form as
1 lod (2, = 2 BE-z
|D (0T= U=)
(16)
where U2 == p > p GIg. As we only brought expectations inside the norm, Eq. 16 has Hence, define continuous-time relaxation of gradient descend dynamics for Eq. 16 (equiv. Eq. 15 by the following ordinary differental equation (ODE) system:
the same gradients as Eq. 15
To.z.t E arg minT, Led (0t; T2, v/2)
To. E arg minl,, Lia (vi, T2, 91)
(17)
di Cod(1, To,z,1 4/t)
1/2 = v. Si,Lia T..
This implicitly assumes that predictors are optimized at much faster rate than representations an important property used in Theorem of Tang et al. (2023) show constant covariance and, thus, no collapse. We now prove this by following similar steps as in the proof of Tang et al. (2023). adapted our setting. particular, we prove
the representations & only. Given the symmetry of the losses, the same result can trivially be proven for 1e well.
We need to prove that d Φt) = O. Since ( 0r) (0) TT) + oJ Sr. it is enough to show that o. d. Simple algebra yields
VoLullo Tz, 1)) = E2 [D, (ØT, - -U,) T.
Te, w) = p(=) D, (ØT, U2).
Therefore,
o. b. = 0. Vo, Lraldr, To. 1br)
= - o E2 DD (KrT U2) TT
= Σ pz) Dp (0i U=) T
3EZ
Σ
Vi..., Gral To,z.11 To,2,144.7T TT 0.2.
2EZ
0,
where the last equation holds since the gradient w.r.t. the predictor is zero at every step (first order optimality conditions from Eq. 17).
17

---

Published as a conference paper at ICLR 2026
B.4 PROOF OF THEOREM 3
We begin by rewriting the TD-JEPA loss of Eq. witth the notation of Sec. as
LTDJEPA (0, 22 v)) := B0 (+1s)
[s)) (1+))
7 0. 00
(18)
Following the proof of Theorem | we can put this matrix form as
(o Tep v0) E ||D// (ØT, Ph 7 PT=
|
(19)
As we only brought expectations inside the norm, Eq. 19 has the same gradients as Eq.
18
SO we
can focus on prove the results. Moreover, can set D, / by Assumption A2.
Statement 1
We start by computing the gradients of Led and Liw w.r.t. Tg. Up to a multiplicative constant p(=) (which doesn't 't change the results), we have
VT, (o, T2, v/ Tt (øT: p y p =Tx T. T gDb p.d cT=
where we used Assumption A1 to set T I. Further using that T I,
VT, Lew T2, 1) Φ0 'OTT247 - p y -pp= ,Tx,v/ )
T, σ p ya0 p.. oT=
Vr, Led (Ø, T)
Therefore, the gradients w.r.t. T. of the two losses match, which means that the stationary points (i.e., optimal predictors) are also the same. Setting these gradients to zero we thus find that
= = F7 = (0 > 06) 10 pras.
Note that matrix TT(T (T--7p ) is positive definite and, thus, invertible. This is because - y pr: is positive definite and Φ =I Using that M (I ? pm=)) 1pT,
OT ST ol0 ol0 =p= )0) 1TT(1 P M M
N..-
where it is easy to verify that az is idempotent (f,,. = Ns,: but not necessarily symmetric, hence an oblique projection as stated.
For the other result in Statement 1, we proceed analogously by first showing that
Vr, (v ()) = T2 1b pm db - -/b p
VT, where the last equality is true since p is symmetric. Then the result follows before after equating the gradients to zero, solving for T a and expressing vT function of I1g..-
= vt (p*a) To mbb (P=x) JTT = V
Statement 2
We show that the gradients w.r.t. & and & match for any predictor T.
V VoLullo, T., 1b) (ØT: - p x - P T.
VoL (0, Te, 3) (0T v pa= - -= of VT. = (o,
where we used that T Φ I. Similarly.
V V Lta Tz, () = -WTT= - s - pp (T)
V Low (g) Tz, 1)) = (v) - (P"=)]T (Pm=) TT T) TT. = V Lia T2, where we used that Φ T = I and (P==)T_ = pxx, This proves the statement.
18

---

Published as a conference paper at ICLR 2026
B.5 PROOF OF THEOREM 4
Let us start from the first inequality. Defining V E IR as a vector containing V (s) for all states E S, we can write the left-hand side for any r IR as
E
Σ LES
d(s) (s)
BBez - 2T22wr11
Since V. r and, by Assumption A1, Wr = Tr,
E E
|V of,,, 12 = BrEz/ 1M 6T,; r
S Bzez \M === OT 1/ 1 |||
(o, Tu, 11111
The inequality is thus obtained by maximizing both sides over rewards with norm bounded by
Let us now prove the bounds of LsM in terms of the Bellman errors Liw and Low. Let us fix E and recall that R admits both a "forward" Bellman equation, M pT +yp M
and a "backward" one, = = p + >M pms (Blier et al. 202T). This implies that = (I >P== - nd, (I- ) -1 Then, for any matrix M RSXS
M - = = 7 (11 7p*= P58) = -1 -1-7 (M p - 7PP=MM)
M - M (M(I - pm=) - p=:) (I - 7 p==)--= (M- - pr= >M (I - - p
Using the first set of equalities with M = oT ( LT, we can easily bound
(0 T2
} E
[| TT M [|:]
2
1
BInz ||[[ - -pp== ) 1221 oT - - o
2
where we used the inequality || - | YI with || 112 denoting the operator norm and || || the frobenius norm. Moreover,
S
[I (
p1= pT
(1 (1-77 7)2 [1[1 71 p 1
(I (1-
where the last inequality holds since (1 - (1 pp=)) )m is a stochastic matrix. Hence,
s
Lo T2,1/) 5 (1 7)2 Ltw (1), Tz, vb),
which proves the first inequality with c [T>)). The second one can be proved analogously.
C THEORETICAL ANALYSIS UNDER RELAXED ASSUMPTIONS
This section describes how the main assumptions used in Section 4 can be removed, namely the uniform state distribution (A1), identity covariances (A2), and symmetric kernels pxz or M (A3). We shall derive similar results Th. and Th. 3] but at the price of more complex proofs and notation. We do so through a novel "gradient matching" argument that reduces a general latent-predictive loss to density approximation. As we shall see, this encompasses not only MC-JEPA and TD-JEPA, but
also existing methods (Tang et al. 2023: Khetarpal et al. 2025
Lawson et al. 2025).
C.I REDUCTION FROM LATENT-PREDICTIVE TO DENSITY- LOSSES
Let Z be a finite set. For 2 E Z, let E, E SxS be a generic kernel. For instance, E, may be the successor measure 9 or the one-step kernel pd: but it is not important this point. Using the
19

---

Published as a conference paper at ICLR 2026
same notation Section 4 and Appendix
consider the following density- -baed loss:
( (0l0) 2 s~z [T
= 27
1
Laens T2: v6)
2 2 aoppptt
p(st)
1
(20)
Minimizing this loss over representations o, v and a collection of predictors { is equivalent finding the best multilinear approximation Note that this is well-defined loss (i.e., it does not involve grggrradeent operations) and the prediction targets =. D are not a function of the representations being learned. Our goal to show that certain latent-predicts dynamics optimize this density-based loss.
the densities of the E, w.r.t. the state distribution
For didactic purpose, let us consider two abstract latent- iatent-predictive losses Lo(Ø, To.z: 1/) and Ce (tb, we. 6). We shall specify what these are later. (o( Te... 15) is optimized over and To for all z, while Ly(s,, T2.2. +) is optimized over v and Ty. for all z. As common in the literature, relaxation of gradient descend dynamics by assuming that predictors
study continuous-timg optimized at a much faster rate than representations an important property to ensure that the latter ones do not collapse (Tang al. 2023). This process be described by the following ordinary differental equation (ODE) system:
{
To, E arg argg,ny, Loldr, Te, the)
To.z. E arg miny, fxp (the, Tz, 0e) d: = -Vo VorCollr (or 2o. di -V. (err, Te,e,e
(21)
We now ask the question: how should L and Ly be defined for the dynamics of Eq. to optimize the density-based loss of Eq.20) 20 The following result provides and answer.
Theorem 5. Suppose that the gradients of Ldens. Lsp and Lu match for all d. v. TV, i.c.
1. Vr, feens (0, Tz, v = Vr, (o, T2,11)
2. Vr, Fdens (0, TT. 1/) Vr, Lae (v), 02)0)
3. V VoLeees T2, (oLo Tb)
4. V Litens T v) = voC (vb D2,
Then, Ldens is Lyapunov function for the ODE of Eq. 21
Proof. Let Te,t E arg milly, Ldens (0r,, f T2, sha) be the optimal predictor for the density-based loss given (ør, the) and define L(t) == Cdena Oc, Te,t 4/c). We verify that a C(t), and thus Ldens is Lyapunov function for the ODE of Eq. 21 First note that Laens is continuous and has continuous first derivates. By the chain rule,
d
df C(t) V VaLee (or, Tz, bh) de + Wdede Ts,s vt) the.
where the operation is the dot product of the vectorized matrices. By Eq. 21
2
Lo To,,,,, vix) and sh: - - Vo Ty.a,t, dt). Moreover, assumptions 1 and 2 directly yield To.z,t = T2, and T -TT, Plugging these into di and di. and using assumptions 3 and 4,
de = - VoLoler T2, we) = - - Co (0r, Te,t, vi),
dh V Lo TI. TI, V Ldens (De,, T2,t, 1/22-
Thus,
((1) = Loenns T2,ts we)l || y (de, T2,25 ve)l 5 0,
Moreover, we clearly have $ C(t) < 0f (dr, vh,) is not a stationary point the ODE, which proves the statement.
20

---

Published as a conference paper at ICLR 2026
Therefore, if the gradients of Leens L.. and Ly match, the gradient dynamics of Eq. 21 monoton- ically improve the density-based loss over time. We remark that this does not imply convergence to the global minimum Lenns is not convex. Th. suggests simple trick for designing the right latent- predicttve dynamics to optimize given density-based loss: just find La and Lup whose gradients match those of Ldens- Latent predictive losses that satisfy this property are derived in the next result.
Proposition 2. Consider the following latent latent-predictive losses
Co(0, 12) B2n2 1D (0T - -.V )13
(22)
1
Lu(w, T2,
(23)
2
where ||X|| ||X Eg == D Ey == T Dpti and E == D' is the
p-adjoint E. Then, Ls and Ly satisfy the gradient matching conditions for Ldeus stated in Th. 5
Proof. This can be directly obtained through simple linear algebra.
Note that, while the preentt-ee losses of Eq. and A are expressed in expectation w.r.t the kernels E, and Ex sample- based estimators for both are possible. In particular, while Eq. 22 involves sampling from the kernel E,, Eq. 23 involves "backward" sampling operation given by the adjoint of E..
C.2 GENERALIZING EXISTING RESULTS
Through the right assumptions and choice of E,, Proposition 2 and Theorem Slicld several existing results. First, with =_ M note that Eq. 22 and 23 are equivalent MC-JEPA (ef. Eq. () with additional covariance weighting and a backward-in-tine sampling in the loss for tb. Assuming that D, (AI), Eg Σn resolved and we recover Th. under its same assumptions. On the other hand, Proposition and Theorem 5 show what happens when such assumptions are relaxed: if only A3 holde then Th. still holds up to covariance transformations, while if A3 is further removed then MC-JEPA needs be modified with backward sampling in the loss v to guarantee optimization of the successor measure approximation loss.
(A2), and pS is symmetric for all z (A3), this mismatch
Several analyses in related works, which assume A1-A3, are also recovered by proper choice of E. For instance, for single policy æ (i.e., |2| 1), we obtain Theorem 6 of Tang et al. (2023)
by setting E = pr and Theorem 4.1 of Lawson et al. (2025) by setting E, M By letting
Z correspond to the action space A, we recover Theorem 2 of Khetarpal et al. (0055) by setting
Eg = =P for all actions a, with Pa the matrix containing ((s" (s a.) for all (s,s").
C.3 TD-JEPA WITH FORWARD-BACKWER
SAMPLING
In the previous section, we have seen that MC-JEPA requires a backward-in -n-timme sampling operation in the loss for 1a to provably optimize the successor measure approximation loss. In this section, we derive the TD dynamics corresponding to this process.
We thus start from Letess La. and L. with E. M and replace the latter with the one-step kernel p*= plus the bootstrapped parameterization M z TT= For Laens similarly to Th. 3 we derive density- based forward and backward TD losses, respectively based on the application of the forward and backward Bellman operator. For the first one, we use that T pr, 7 M Fe pss + 7P ΦT. u D, and write
Lew (0, T, v) ENnN
1
(OT p== PP== DP.DD
(24)
To be precise, here we need the density 4 DJ, be symmetric, which implies that (M*.)* M
21

---

Published as a conference paper at ICLR 2026
For the second one, rewrite the backward Bellman equation in terms of the adjoint of M as
(M1 = (P*x)* + p=M
(P=)) ((M D = (Pxx)* +> (P=x)
( D = (PTi)* + t (P=.) (D
= (Pxx)' + ,(p(.)) D'
This yields
Cow (6, Te, v) = 1 E
VT, D |1
(25)
We then use the same bootstrapping for the patent-predictive losses of Eq. 22 and
23
thus obtaining
Co(Ø, T2, 1)) BInz |D/ (ØT= PT E pp oT=
(26)
1
ColV, 0) 6) = 2 1 (PT. ) -
(27)
As for MC-JEPA, this the counterpart of TD-JEPA with additional covariance weighting and backward-in-time sampling in the loss of tb. Through simple algebra, it is not difficult to check that Lg and Cg. have matching gradients (as in Th. D w.r.. Liw and Low, respectively. This implies that Th. holds for this modified algorithm without assumptions A1-A3.
Unfortunately, differently from the original TD-JEPA, this variant is not easy to be optimized off- policy. In fact, while for 26 we can simply replace the on-policy kernel pt, with P(s'|s, a) and condition the predictor on actions, the same trick cannot be used for the adjoint (P**)* 1 in Eq. 27 This is because there is action-conditioned backward Bellman equation. We leave the study of practical learning dynamics for this theoretically sound variant of TD-JEPA for future work.
D ADDITIONAL RESULTS
This section reports additional experiments, as well as detailed numerical results for plots in the main part of the paper.
D.1 EXPLICIT STATE ENCODERS FOR ZERO-SHOT BASELINES
As mentioned in Section 6 existing zero-shot methods do not necessarily learn explicit state em- beddings Wu al. (019) Touati Ollivier (202T) Park et al. (024)) Nevertheless, we find that introducing a state encoder tends to improve average zero-shot performance, even in the absence of a specific representation learning objective. We consider three established zero-shot baselines (Laplacian (Wu al.) 2019), FB (Touati & Ollivier (2021) and HILP (Park et al.) (025aa), and com- pare their performance without an explicit state encoder (i.e., successor features are trained on raw states), to those attained when introducing a state encoder (either shallow or deep). Note that, in this case, the state encoder is trained through gradient flowing from the original loss, and is not coupled to ad-hoc objective. Results mains as a state encoder is necessary when learning from pixels. We observe that, while having an explicit state encoder may be detrimental in few specific domains, it remains beneficial average, and crucial to obtain better performance in some domains. The optimal depth of the encoder is how- ever domain-specific: OGBench domains generally prefer deeper encoder, while DMC domains can be solved with a shallow encoder, or no explicit encoder at all. In order to maximize perfor- mance of the baselines, we thus evaluate them coupled with deep encoder in OGBench, and with soallow one in DMC (e.g., in Tab. Tb. We additionally summarize performance differences induced by this choice of encoders in Fig. (eft).
summarized in Tab. which only considers proprioceptive do-
D.2 CONTRASTIVE VARIANT OF SYMMETRIC TD-JEPA
In order to further isolate the effect of different objectives zero-shot performance, instantiate the symmetric variant of TD-JEPA described by Eq. | where a single representation is used both state and task encoder, as well as its contrastive counterpart, which objective similar to
22

---

Published as a conference paper at ICLR 2026
H1I
HIT Faop
Eaplacion
Laglaciimaater
Laplacioneeo
DMC (avg)
648.2
659.4
603.2
585.6
508,8
valker
8112
18361
1.13
788
7011
cheatah
6123
quedzuped
965
565
pointmass
ST88 23
5130 310
1339 960
1943 218
14677
3451
19.07 1:000
21.96 1:05
39.04 104
30.51 11
35.02 2:100
37.98 1:111
9.07 2:871
9:29 1:05
14.81 1.10
mntmate-mi
20.30
2:13
7189 2:17
0231 11139
M333 2.37
1130
15 13911-50
16:1
sataeze-in
1520 342
12.00 254
36.99 1-43
3-66 2.5
27.4 4:40
52,8 1.0
13,02 2:30
10.91 246
21.00 350
wntine my
1.6
03:11:115 313
1516 0 1.0
1348 5.10
S0:8 1.36
14.38 123
13.91 3.1
21.31-433
anlineee 1s
15:0 1/5
4850 1:14
133
11.00 L.
13:39::11 175
14:5:250
600 151
1189
antmeze-ne
23002:10
23912:10 1/1
2:214
8:912:100 1.30
1/::3445
caba-single
2100
2680 242
300
2.4
1928
11.11:214
1511 1311--14
cabe double
2.4
135
2380
23.99 229
835
2.88
03.8
1.
580
puzzle-Jac3
Table 2: Ablation over encoder depth on proprioceptive DMC and OGBench domains for zero-shot baselines. Each method is evaluated in three variants: a baseline without explicit state encoder, one with a shallow (i.e., linear) state encoder, one with deep one (2 or 4 hidden layers for DMC and OGBench, respectively, see Tab. 5)) We report mean performance and standard error. Results for each variant are bold if their confidence intervals overlap with the best variant of the same method.
Algorithm Symmetric TD-JEPA for zero-shot RL (latent-predich) and contrastive variants)
Inputs:
Dataset D, batch size B, regularization coefficient A, networks x, Tg.
Initialize target networks: T5 To.
while not converged do
D. Sample training batch
((s,, ai; <()},=, D, [=) 2, fai ~ { (6)),:))
Compute lncent-rrrdictiiv//eenttssssssss loss
E( Tw) 28 Σ, (o((s)),
2
[(0, Tb) 28(8-1) Eej (To( ai (0(0)) - T(()))
- Σ, (o(ss). dt, 2;) ((11)
D Compute orthonormality regularization loss
[REC(Ø) 2BTB-1] Σ (o(s)) ((;;))² 1/2 E, ((s)))
D Compute actor loss
(a:) (0)()))) = Lactor (T) El. (o()
Update é, To minimize L(bb, T + (o)
Update æ to minimize Luctor (T)
Update target networks & T5 via EMA d. T.
that Touati & Ollivier (202I). These algorithms are described Alg. 2 where blue and yellow lines are exclusive to the latent-predictive and contrastive variants, respectively. Intuitively, the contrastive variant can be seen as a specific, symmetric instantiation of FB (Touati & Ollivier 202T), and the comparison to the symmetric variant parallels that between TD-JEPA and FB. We report performance differences between the symmetric variant of TD-JEPA and the symmetric-contrastiv
variant in Fig. (right). Similarly results from Fig. 2 we observe that gaps in performance between self-predictive and contrastive methods grow larger when learning directly from pixels. For numerical comparison, see Tab. 3
D.3 ADDITIONAL FAST ADAPTATION RESULTS
We extend the empirical evaluation on fast adaptation in Section Joction[]throve 6thhough Fig. 6 In the top part, we repeat the evaluation of Fig. 4 while only initializing encoders to pre-trained weights. The remain- ing components of actor and critic need to be learned from scratch. These experiments thus recall the standard setting in which pre- pre-trained trained visual representations are evaluated (Nair et al.] 2022: Ma et al. 2023; Majumdar et al. 2023). We observe that, while initial performance near-zero for all methods, pre-trained representations maintain their effectiveness in terms of sample efficiency. In the middle row, again repeat the evaluation from Fig. 4 but only freeze convolutional weights
23

---

Published as a conference paper at ICLR 2026
Laplacian
HILP
FB
RGB
Proprioception
antmaze-ms
walker
antmaze-ms
antmaze-1s
quadruped
antmaze-In
cherlah
antmaze-me
antmaze-1s
cube-single
antmaze-In
quadruped
cube-single scene
walker
puzzle-3x3
poin tmass antmaze-mn
cheetah
H
puzzle-30 1e-zlee-3 3x3 cube antmaze-me
cube-double
pointmass
0.2
-0.2 0.2
0.0
0.5
0.0
0.2
0.0
0.2
Normalized performance difference
Normalized performance difference
Figure 5: Difference in normalized performance between zero-shot baselines with and without an explicit encoder (left); normalized performance difference between symmetric TD-JEPA and its con- trastive variant (right). Error bars represent standard errors normalized performance differences.
for the variants represented by a dashed line. We observe that this causes the performance gap to full fine-tuning to shrink significantly, suggesting that convolutional filters extract suitable represen- tations for both TD-JEPA and FB.
Finally, we extend our empirical evaluation OGBench in the bottom part of Figure 6 in which we present results for the most challenging tasks in three representative domains. This evaluation differs from the previous ones, as we found it to require strong BC regularization, even during fine-tuning. We confirm that pre-trained representations remain beneficial in terms of sample efficient adaptation. Interestingly, we find that frozen representations pre-trained may outperform full fine-tuning in antmaze 1s, while they remained bottleneck in DMC.
D.4 ARCHITECTURAL ABLATIONS
Architectural choices are often crucial for self- eeff-predictive learning, and capacity is usually carefully distributed between encoder and predictor (Guo al.) 2022). This section ablates the width of the three main components in TD-JEPA in a controlled setting, namely in two visual OGBench tasks (antmaze- -n and antmaze ze 1s), which were chosen due their complexity, and the fact that they often reward different approaches (see Tab. D We measure zero-shot performance as widths change in Fig. 7 In general, we observe that the state encoder (on the y axis) should be as deep as possible (although this trend is not present in DMC, see Table Instead, the task encoder (on the x axis) should only be as deep as needed: at least layer in antmaze-11 and closer to 4 for antma ls- Finally, we observe that the predictor may be shallow, as long as the encoders have sufficient capacity. This result matches the general conjecture that predictrede representations should capture the key aspects of the input, to the extent that prediction problem, e.g. predicting successor features, may be solved with limited capacity.
D.5 VISUALIZATION OF TD-JEPA REPRESENTATIONS
Given state-action pair (s, a) and further state g, one may easily evaluate the successor measure in of the policy that tries to reach from s: for instance, in the case of TD-JEPA, Ta (g 8, a) RJ (d(s), (,,,),,)) =g) (gg). where 2g DEw (tb (()())())))))))))) s) - Intuitively, this connected to how quickly the policy may reach the goal, and should reflect the dynamics of the MDP, As often done the literature (Lawson et al. 2025). we consider the methods evaluated in Tab. D and visualize these estimates for visual antmaze-iee in Fig. g highlighting how representations reflect tempo- ral distances in the environment. While all plots reveal similar pattern (which is not surprising representations by definition), some methods show
successor features inconsistent temporal distances. For instance, BYOL, RLDP, and HILP have latents of few states far from the goal with higher similarity to the latter than closer states. This is sign of poor perfor- mance, as the predictor would be optimistic in predicting the number of steps to reach those goals.
A slightly different visualization of similarities is also helpful in highlighting the differences in the representation learning objectives of TD-JEPA, BYOL" and BYOL-y*. Figure D.4 displays simi- larities between the task embedding of given goal (g)) again depicted star, and predictions
24

---

Published as a conference paper at ICLR 2026
Normalized Petformance Normalized Petormance
walker
cheetah
quadruped
pointmass
0.4
Offine
0.5
0.2
0.2
0.0
0.0
0.0
0.0
200K
200K
400K
400K
Steps
Steps
Steps
Steps
walker
cheetah
quadruped
pointmass
Online
0.50
0.50
0.4
0.5
0.25
0.2
0.00
0.00
0.0
0.0
200K
400K
Steps
Steps
Steps
Steps
Normalized Performance Normalized Performance
walker
cheetah
quadruped
pointmass
0.4
0.50
0.4
Offine
0.5
0.2
0.25
0.2
0.0
0.00
0.0
0.0
0K
200K
400K
200K
400K
0K
200K
400K
0K
200K Steps
400K
Steps
Steps
Steps
walker
cheetah
quadruped
pointmass
0.4
Online
0.5
0.25
0.25
0.00
0.00
0.0
0.0
200K
200K
400K
Steps
Steps
Steps
Steps
Normaltzed Performance Normalized Porformancs
antmaze-In
antmaze-ls
cube-single
Offline
0.5
0.2
0.0
0K
200K
400K
0K
200K
400K
DK
200K
400K
Steps
Steps
Steps
antmaze-In
antmaze-ls
cube-single
Online
0.5
0.25
0.0
0.0
0.00
0K
200K
400K
0K
200K
100K
0K
200X
100K
Steps
Steps
Staps
TD-JEPA
TD-JEPA (frozen)
FB
FB (frozen)
Scratch
Figure 6: Additional fast adaptation results, obtained when only loading encoder's weights (top). when only freezing convolutional layers (middle) and in OGBench (bottom).
25

---

Published as a conference paper at ICLR 2026
Predictor depth:
Predictor depth:
Predictor depth:
4=2
17#5
18:1
antmaze- -In State encoder depth
2=1
18+5
22
30
1=0
24
17*4
20*2
2+1
26
12*3
17#2
3
25a?
32.14
27
3a:1
26.55
22
22:4
1s:1
32.a3
10a3
21a2
2=1
44
38
38+10
4
0+00
37
32#5
32#11
4=1
46#3
22
30
0=:
14#3
2=1
7+2
1=1
12#2
6#2
11#3
0
0=00
13#5
3=2
14#2
antmaze-15 encoder
2
0=:
12#3
20
24#2
2
0=0
18#1
15#2
24:4
2
0=00
17#0
12#3
22#1
2
Task encoder depth
Task encodar depth
Task encoder depth
Figure 7: Zero-shot performance of TD-JEPA in antmaze-in (top) and antmaze-ls (bottom) as the number of hidden layers in the encoders and predictors varies (from and from 4, respectively).
TD. JEPA BYOL
BYOL"
RLDP
RLDP
GEE HIP
FB
ICVP"
Losiacian
Figure 8: Cosine similarities between successor features and features (e.g., (To(6()), ag), v(g) for TD-JEPA) projected over zy position of the agent's center of mass. The red star marks g. Similarities reflect shortest-path distances in the MDP.
TD-JEPA C BYOL
BOO DYOL:
Figure 9: Cosine similarities between predictions and features (i.e., (Ta()))) zg), (gg)) for TD- JEPA and (T((()) a), ((g)) for BYOL-based methods) projected over the zy position of the agent's center of mass. The red star marks g. Similarities reflect alignment between predictions from a given state and the goal BYOL-based approaches model dynamics induced by the (unconditional) behavioral policy, and their predictions are not directed towards specific tasks.
26

---

Published as a conference paper at ICLR 2026
Figure 10: Top: a trajecctory from the visual sinel dataset of OGBench. Middle: visual features learned by TD-JEPA, computed as a PCA projection of the channels returned by the convolutional layers in é down to dimensions. Bottom: the same PCA projection applied to the visual features computed by DINO-v2.
from several state-action pairs in the dataset. In the case of TD-JEPA, this prediction is condi- tioned on the goal reaching task 2y (v))) (s)) (g), and we thus again visualize (o((s) a, 2g) (g)).. BYOL-based methods extract their representations through unconditional prediction of the very next state (BYOL"") or of future state (BYOL-?"). For these two algorithms, we thus may only compute predictions that are not conditioned on any particular task, and visualize their similarities (()(s), a) ((()). As expected, the predictions of BYOL* align well with the goal g only near g itself due to the one-step nature the loss, while those of BOOL->* are more "aver- aged out" since the algorithm models the long-term occupancy of very stochastic behavior policy. Only TD-JEPA yield predictions that model the directed behavior of a task-aware policy.
We finally compare the visual features learned by TD-JEPA with those of DINO-v2 (Oquab al. 2023). chosen as a representative method for general purpose (i.e., not control/dynamics hawoo//dyamic-awaarre vision encoders. As shown in Figure D.5 for the OGBench cube-si meiinlle le task, while DINO-v2 nicely segments the robotic arm (i.e., the main object in the scene), TD-JEPA focuses more on the end-effector and the cube (i.e., the objects that are most descriptive of the agent's 's behaviors). This showcases one of the advantages of (long-term) dynamics prediction, as focusing on those features that are most relevant for control.
D.6 NUMERICAL RESULTS FOR PERFORMANCE DIFFERENCE PLOTS
We supplement the performance difference plot detailing numerical results.
Fig. (right) and Fig. 5 (right) with Tab.
D.7 COMPUTATIONAL COST
We report the training speed (in terms of gradient steps per second) of all tested algorithms in Table As expected, the symmetric variant of TD-JEPA (see Section(.11 is generally the fastest algorithm, since it trains a single encoder and predictor through a simple TD loss. The asymmetric version of TD-JEPA is slower, requiring two encoders and two predictors, but is notably as fast, faster than successor-feature- based methods like BYOL and B BOOL-Y.
D.8 PERFORMANCE LOW-QUALITY LOW-COVERAGE DATA
The zero-shot RL experiments of Section focus on two types of datasets: low-quality high- coverage (ExoRL), where trajectories cover the state-action space well, but do not contain any "purposeful" behavior, and high-quality low-coverage (OGBench), where trajectories lipert-like, but cover a narrow portion of the state space. Here we analyze TD-JEPA's perfor- mance under a more extreme setting of low-quality low-coverage data. To do so, follow Jeen
mostly
27

---

Published as a conference paper at ICLR 2026
C-TD-JEPA
TD-JEPA gym
TD-JEPA
DMCRGB (avg)
437.2 9.8
598. 1 5.9
628.8 5.5
walker
413.2
685.9
738.9 + 3.5
cheetah
517.9 1: 31.3
688.0 +: 7.2
706.0 1: 4.1
quadruped
428.3
pointmass
389.5 - 17.2
411.6 13.9
443.7 + 10.9
DMC (avg)
586.3 + 14.6
657.5 + 3.6
661.2 + 6.3
walker
757.8 12.4
800.7 4.7
785.2 + 6.7
cheetah
583.3 23.3
618.1 11.3
688.7 1: 6.7
quadruped
565.5 1:
731.7 =:: 17.3
691.4 1: 5.0
pointmass
438.8
479.6
479.3 23,6
OGBenchgas (avg)
33.93
39. 74 1:
41.34
antmaze-mn
93.80 1: 1.05
95.80 1.09
96.67 1.11
antmaze-1n
antmaze-ms
61.20 + 3.00
82.20 1,80
84,40 + 3.85
antmaze-1s
8.80 = 0.90
20.60 1.03
28,80 2.50
antmaze-me
1.40 = 0.52
0.60 + 0.31
0.20 + 0.20
cube-single
64.00 1: 5.70
69.00 2.62
67.80 1: 3.67
cube-double
2,00 0.67
1,40 1 0.60
3,00 :: 0.91
scene
4.00 = 0.79
8.00 + 1.94
14.20 + 2.22
puzzle-3x3
2.20 0.76
2.22 + 0.78
=
OGBench (avg)
1: 35.58
35.20 +:
37.98
antmaze-mi
82.20 + 2.87
76,40 2.65
70.40 +: 3,72
antmaze-1n
antmaze-ms
74,00 + 2.27
62.80 3.59
61.56 4.53
antmaze-1s
43,20
41.40 3.82
40.60 2.51
antmaze-me
21.80 1: 3.30
17.00 2.65
20.20 2.39
cube-single
16.80 1: 2.15
20.00 2.17
34.20 2.88
cube-double
3.60 1.07
2.40 1 0.78
3.60 1: 0.78
scene
16.80 + 1.31
39.40 = 2.17
38.44 + 1.37
puzzle-3x3
4,60 1.19
13.80
15.60 1.11
Table 3: Performance of TD-JEPA and symmetric variants (contrastive and latent-predictive) in DMC (returns) and OGBench (success rate) with either proprioception RGB inputs. We report means and standard errors across seeds. Numbers are bold for top algorithms if confidence intervals overlap.
DMCRGB
DMC
OGBenchRGB
OGBench
TD-JEPA
19
52
35
132
TD-JEPA (sym)
34
84
61
161
FB
23
72
45
146
BYOL
19
55
36
112
BYOL-Y
18
55
34
110
HILP
16
62
29
111
Laplacian
16
59
31
112
ICVF
11
55
26
99
RLDP
15
55
137
Table 4: Average number of iterations (gradient steps) per second of each algorithm for the ex- periments Table [ All methods are trained on the same hardware (a single V100 GPU) and similarly sized architectures (see Appendix E One training iteration samples batch from the dataset and updates all networks once.
28

---

Published as a conference paper at ICLR 2026
100 episades
500 episades
Reward
Reward
100
100
200
200
PB
FB
FB PQL
TD TD. JEPA
TD-
PB
FB
FB PQL
TDJEPA TD JEPA
JEA
BC
+ BC
+ FQL
BBC
+ BC
+ FOL
1000 episodes
5000 episodes
600
800
Reward
Reward
400
400
200
200
FB
FB
FR FOL
TDJEPA TD.JEPA JEPA TD-JEPA
FB
FB
FB FQL
TDJJFA TD.
TD-JEPA
BC
+ BC
FOL
BC
+ BC
FOL
1M
3M
Figure 11: Average performance over all DMC wa ker and quadruped tasks for varying dataset sizes and number of training steps (IM, 2M, or 3M).
700
Reward
600
500
400
100
500
1000
5000
Number of episodes
TD-JEPA
TD-JEPA BC
TD-JEPA + FQL
+
+ FQL
Figure 12:
Average performance over all DMC walker and quadruped tasks function of the dataset size. All algorithms are trained for 3M steps.
et al. (2024a) and test TD-JEPA on two representative domains from ExoRL (walker and quadruped)
while decreasing the dataset size. It well known that offline training low-coverage data requires some form of constraints to prevent extrapolation beyond the data distribution. We thus further re- port the performance of TD-JEPA combined with plain behavior-cloning (BC/Fujmmto Gu 2021) and FQL-style (FQL Park et al. 2025b) regularization, the latter of which is used across OGBecn experiments. The results, which also include FB-based variants, are shown in Figure D.8 and D.8 We can make the following observations: (1) as expected, the performance of all algorithms degrade when decreasing the data size, with drops consistent to those observed by Jeen et al.) (0024a). (2) Both BC and FQL regularization are effective in mitigating the drops. For instance, the performance of TD-JEPA on the smallest dataset increases between 10% and 20% when adding regularization. (3) On small datasets, performance tend to decrease an expected sign of overfitting. While behavior regularization does not mitigate this decrease, other techniques (e.g., early stopping) likely to be beneficial. Finally, remark that TD-JEPA could
increase the number of training steps,
29

---

Published as a conference paper at ICLR 2026
Task Encoder
Predictor (%)
Actor
Convolutional Encoder
Input Embedding
Linear Norm
[K*
Lincare
State Encoder
Fiee RelU
v g
LayarNorm lanH
Liner ReLU
Linear ReEU
Hatton
LayerNorm TanH
Lieer Norm
Liner
Ingun
Co RelU
Connoittionll
Corvelutional
Figure 13: Overview of the architectures used by TD-JEPA in DMCRGE. We refer to Tab. 5 for the different instantiations in other domains.
be easily combined with other conservative methods, such as advantage weighting CQL (Jeen) jet al., 2024a). Evaluating these techniques goes beyond the scope of our work and represents and interesting avenue for future investigations.
E IMPLEMENTATION DETAILS
We organize the discussion of implementation details in several subsections.
E.1 ENVIRONMENTS
Zero-shot results in Section consider the standard Deepmind Control Suite (DMC, Tassa et al. 20188) evaluation over four domains from (Touati et al. (202)): walker, cheetah, quadruped and maze. The first three define several locomotion tasks (e.g. walk, run, flip), while the latter evaluates both goal- and reward-based tasks. We additionally extend the zero-shot evalu- ation to the more recent OGBench suite (Park et al.] 2024), which focuses exclusively on goal reaching. For computational reasons, we consider nine representative domains. We thus evaluate navigation across five antmaze datasets (medium -um---nvvgat,, stitch, explore) and Large (naiie-nnaviatte,, stitch}) and manipulation through cube cube-(sinng) (single, double), }, scene and puwzle-uaae 3x33. In each domain we consider the five default tasks; for consistency with DMC, carry zero-shot evaluation through the standard reward inference procedure, and thus define each task through its reward function. During inference, we shift each reward by +1, which found significantly improve zero-shot performance.
E.2 LEARNING FROM PIXELS
In visual experiments, environment return states as 64 x 64 RGB images. In order to alleviate non-Markovianity, states composed by stacking frames. Fach image is scaled to [-0.5, 0.5] and undergoes random shifts with maximum intensity of 2 pixels, which is slightly milder than common strategies (Yarats al. 2022b). The images are passed through convolutional encoders before being processed by further networks. Namely, we instantiate two convolutional encoders: one for the state encoder, and one for the task encoder (except for the symmetric variant of TD-JEPA, which only uses one). We however found that using single, shared convolutional encoder does not significantly impact performance, and do not exclude that further specialization of the encoder might improve it. Each convolutional encoder uses the DrQ-v2 architecture introduced in Yarats ctt al. (2022b); we briefly experimented with IMPALA- (Espcholt et al., 2018)) and Dreamer- like (Hafner al. 2020) architectures, and we found them to achieve similar results when well-tuned.
E.3 ARCHITECTURES
All algorithms rely on three types of networks, each of which is instantiated to the standard ar- chitectures from (Tuuaii et al. 2023): (i) successor feature estimators, predictors and F-networks are MLPs with two aooerizor embedding layers (ii) state/task encoders and B-networks standard MLPs with 12-normalized output and (iii) actor networks Gaussian MLPs with
For compactness, mn, ms, me, 1 dataset abbreviations, respectively.
30

---

Published as a conference paper at ICLR 2026
DMCRGB
DMC
OGBenchRGB
OGBench
SFs: To. To., hidden layers
4
SFs: Tø. Tg. F hidden width
1024
1024
512
512
State Encoder: Φ hidden layers
0
4
State Encoder: hidden width
256
256
512
512
State Encoder: Φ output dimension dæ
256
256
256
256
Task Encoder: hidden layers
Task Encoder: v hidden witch
256
256
512
512
Task Encoder: d output dimension di
50
50
50
Actor: hidden layers
3
3
4
4
Actor: π hidden witch
256
256
512
512
Discount factor y
0.98
0.98
0.99
0.99
Total gradient steps
2M
2M
1M
1M
Batch size B
512
1024
256
256
Default learning rate
10
10
10
10
Default regularization coefficient À
1.0
1.0
1.0
1.0
EMA coefficient
0.001
0.001
0.005
0.005
Pgoal
0.5
0.5
0.5
0.5
Table 5: Architectural (top) and training (bottom) hyperparameters.
similar architecture to predictors, and fixed standard deviation of 0.2. All networks use ReLU ac- tivations, except for embedding layers, which use TanH. The number and width of layers in DMC closely follow those described in Touati et al. (2023), while we utilize deeper, narrower networks in OGBench, in order better align with the implementation of benchmarked methods in Park et (2025a). An overview of the architectures used for TD-JEPA in DMCRGB is depicted in Fig. 13 further architectural hyperparameters describing depth and width of these networks are reported in the first half of [ We remark that the state encoder's ' depth in proprioception was found to be quite impactful, and was tuned according baseline performance, as shown Table
E.4 TRAINING HYPERPARAMETER
The second part of Table describes hyperparameters used for training, and is complemented by the following discussion of further details. We again closely follow the default hyperparameters from Touati et al. (2023) and Park al. (2025a) whenever possible; the batch size is reduced in visual domains to meet computational limitations. All networks are optimized through Adam (Kingma 2014). TD-targets and value estimates for SVG-like policy updates computed the mean of twin networks; latents E Z are representations of random uniform states from the dataset (1()) with probability Pgoal, and are sampled from the hypersphere (i.e., Z) otherwise (Touati al. 2023
E.5 METHOD-SPECINE DETAILS
The baselines' implementation closely follows the public codebases, when available.
Our implementation of FB builds upon the one released in Tirinzoni et al. (0025), which is in turn aligned with the code released by Touati et al. (2023).
As described Jajoo et al. (0025), RLDPP is implemented in the same framework, but the latent dynamics model and the task encoder trained in parallel with the remaining components; gradi- ents from the contrastive FB objective are not backpropagated through the task encoder. As in the original work, we found that multi-step prediction results in better performance, and we similarly adopt prediction horizon of H 5. Given batch of trajectories ((16) of.... .fg....,ff encoder v/ and latent predictor TRLDP Rd x A Rdv the loss for training task representations thus:
.., B- } task
B
ÉRLDR (TRLDP v) =
Σ
Σ
+ (s))))
(28)
2B
T=0
31

---

Published as a conference paper at ICLR 2026
where há +0))). h; TRLDE ( and 6= is target network. The latent predictor is instantiated with the hyperparameters presented for the SFs architecture described in Tab. 5
Laplacian (Mahadevan & Maggioni 2007), HILP Park al. 2024), BYOL" (Grill et al. 2020
BYOL-7" (Lawson et al. 2025) y and ICVF* GGosh et al. 2023 are all implemented in successor- feature framework: the losses proposed in the original publications are optimized to retrieve task encoder 1.. Considering batch of transitions {(s, a, s)) B- the feature learning objectives for Laplacian, BYOL* and B are, respectively:
B-1
Lappaciin (v)) = 2B
Σ
-
(29)
=00
B
CBYOL (TBYOL, v)
T (1b(s)), 2 (s])
(30)
2B
==
LavoL-;(TuooL. 2B
Σ
(M + (1
(31)
==0
(32)
where TBYOL: Rdd x A - Ted is jointly trained latent predictor, St M (-Is,, a.), ag is the
behavioral policy, and the minus sign (-) denotes target networks. HILP and ICVF instead train representations through
[H))(
(- 1 1itk - - -- (1b(()) - (g))))),
(33)
2B
k=0
B-1
CICVF(TicVF, 2,0)
Σ
T cV.
11 - 7V- (V..
(34)
2B
where = - is an expectile regression loss with expectile т. 9, and y; are goals (or intents) sampled from mixture of future and random states described in Ghosh et al. (2023), (II <(x) Ticvp( (b (z)) TicvF Rd + Rd is a matrix predictor and A;
1s:=9: + (5 M/, m1i) V(si, yr, yt). In an unified way across baselines, universal successor
feature estimators Fw : Rd x A Z Red (and state encoders & S - IRds) are then trained
through standard TD-learning over features vé by optimizing:
B Σ =
(F = 2B
Fo(((ss) o;;) -
FF. 10
a: zi)
(35)
24). The latent dynamics model in BYOL* and BYOL-y*, as well as the intent- conditioned predictor from ICVF* (Ghosh al., 2023). share the architecture of SFs, described in Table[5] potentially dropping the subnetworks that receive states, actions or latents as necessary. The state encoder in ICVF* receives gradients from both the ICVF loss, and the successor feature prediction, which we found to slightly improve performance compared to only propagating gradients in either direction.
where ~ (6(()))
We found all zero-shot algorithms to be sensitive to certain hyperparameters, such as the strength of orthonormal regularization X. For fair comparison. we evaluate all algorithms over small hyperparameter grid (6 configurations in DMC. and 4 in OGBench), and report performances for the best performing configuration for each domain. For all algorithms, we sweep over two values for the learning rate of the task encoder th: [10-4, 10-5]. A second important hyperparameter is the orthonormal regularization coefficient X, for which optimal ranges strongly differ across algorithms. In order to avoid an excessively large sweep, the ranges were thus tuned for each algorithm on rep- resentative subset of domains, and are reported in Table 6 In general, we observe that contrastive methods prefer very strong regularization, while self- eeffpreddietive methods can learn with weaker regu- larization certain domains. For algorithms that do not leverage orthonormal regularization (HILP, ICVF*), we instead sweep over the likelihood of sampling random goals/intents in [0.375, 0.5).
32

---

Published as a conference paper at ICLR 2026
Laplacian
FB
RLDP
BYOL
BYOL-
TD-JEPA
DMCRGB
[0.01. 0.1,
[0.01. 0.1, 1]
[0.01. 0.1. 1]
[0.01. 0.1. 1]
[0.01. 0.1.
[0.01. 0.1. 1]
DMC
[0.01. 0.1, 1]
10.01. 0.1. 1]
0.101.01.11 1]
00.01.0...1
[0.01.0.1.1]
[0.01. 1]
OGBenchaca (nav)
[0, 1]
[100, 1000]
[100. 1000]
[0.001, 0.011
[0.001, 0.011
10.
OGBench (nav)
[0. 1]
[100. 1000]
[100. 1000]
[0.001, 0.01]
00.001,0011
[0.1.
BGBencgRGD (man)
[0. 1]
[100. 1000]
[100. 1000]
[0.01. 0.1]
[0.01. 0.1]
[1. 10)
OGBench (man)
[0. 1]
[100, 1000]
[100, 1000]
[0.01, 0.11
[0.01, 0.11
[1, 10]
Table Orthonormal regularization ranges for each algorithm, (nav) and (man) indicate navigation and manipulation domains, respectively.
E.6 OFFLINE CORRECTION
The standard zero-shot evaluation pipeline (Touati et al.] 20233 has traditionally relied on high- coverage datasets (Yarats et al., 2022a).. OGBench (Park et al. 2025a) represents an interesting challenge in this sense, most datasets are expert-like, and the incomplete support over actions induces well-known offline issues (Kumar et al., 2020). While these problems have been previously studied in the context of zero-shot RL (Jeen et al.) 2024b), regularization techniques for single-task RL works well this setting. In particular, we rely on FlowQ-like regularization scheme (Park al.] 2025b): we train flow model to estimate the behavioral policy, and replace the Gaussian policy with a noise-conditioned, one-step policy. This policy is regularized to samples generated by the flow model through 10 integration steps. The resulting behavior cloning loss term is normalized by the mean absolute Q-value in the batch, and scaled by a regularization coefficient at we use a = 3 and a 0.3 for manipulation and navigation tasks, respectively. Both networks are instantiated with the architectural hyperparameters described for ı in Table []] Finally, in order be able to track BC targets, the policy is trained directly over the state space (or on the output of convolutional encoders in visual domain), instead of acting over state representations produced by state encoders ((())
find that novel instantiation of
E.7 EVALUATION: ZERO- SHOT AND FAST ADAPTATION
While zero-shot evaluation is averaged over all tasks (4-7, depending the domain), fast adaptation is only evaluated on the hardest task per domain, estimated through average zero-shot performance of TD-JEPA and FB. This corresponds run walker, cheetah, quadruped and square in maze; we evaluate task 2 in ant domains and task 4 in cube-singge Fast adaptation methods that rely on pre-trained weights are initialized from the best performing zero-shot run. Behavior cloning regularization is carried over from pre-training if present, but orthonormal regularization is no longer applied. For online adaptation, the Update-to-Data ratio is fixed 1, and 5000 initial transitions are collected by the policy before fine-tuning, Batches are sampled from 50-50 mixture of the relabeled pre-training dataset, and a replay buffer containing the most recent 2 1 transitions.
The procedure to extract an actor and critic that are trainable through TD3 (Fujimoto et al., 2018 from zero-shot agents is rather direct. First, the relabeled dataset is used to infer the optimal latent 2, through linear regression. Then, the critic network may then be derived from successor feature estimates: in the case of FB (Touati & Ollivier 2021h. Q(s. a) F(o(s),a, 2,) TTzT in the case of TD-JEPA the same holds as soon as F is replaced by Ta. Although adaptation schemes over Z are possible (Sikchi et al. 2025), ) 2, can be kept frozen. In order to match the output of the critic the scale of rewards, z, is scaled such that its squared L2 norm matches the maximum reward over the dataset. Finally, the actor π (s) can be directly extracted from zero-shot policies as soon as they are conditioned on žy.
All evaluation metrics are averaged over 10 episodes in OGBench, and 20 in DMC, and computed over 10 random seeds for OGBench and 5 for DMC.
33

---

Published as a conference paper at ICLR 2026
E.8 PSEUDOCODE FOR TD-JEPA
def train (salf):
sample training batch
action, next_oba self.replay_buffer hww.ay..uuffrofoooooo semple
self. sell.samlle_ (obs)
6
# compute targets
next_phi hww.-ttrrgt__hh_ccceeee
8
next_psi self www.:aaree__pieeeooeee tnext_obna
texi_ection self.ectoz tnext laext__th)t___e___________ .eam___
target_phi next_pai discount self.target_phi_preee (ne::__rnee__pppp 2, next_action
11
parget_pai pext_phi discount
predicitor (nex___cppo (nex._pai, 2, next_action)
compute predictions
13
phi = self-ghi_encoder pst
15
pred_phi self,phi prepiippeeeetortggeee iphi, action!
pred_ps: self.gei_predictorr ipsi, 2, action!
jepa_losa (pred_phi target_phi.detachh]
(pred_ps1 target_pei.detant
18
regulazize
phi_cov torch.matmul (phi, phi.1)
20
phi_orthe_loss = oni_eav.diagll.mcan + 0.5 phi_cor.off_diog/ .pow .mean
psi_cov torch.matmul (ps1, pai.
22
psi_orthe_lsse psi_cov.diag mean 0.5 pei_cov.oft_diatt
mean
23
corpute
loss
25
acter_actien =
(phi. decach 2) sample
acter_pred pelf..ph_predictor. (phi.detaccl (),, 2, actor_action)
27
actor_loss tactor_pred
-um(-)).mman()
29
aggregate losses and optimize
30
loss jepa_ioss self, sand, + ohi + self.lembda.psi + psi_ortho_loss
aator_losa
32
loss. ban.bbackward)
34
imi.o
35
update tazgel networks
update_target (self. tar..ttteee__ieeeeeeeee self.phi_ennoder/
37
update_target war.fl.trret.psccnnooe.
celf.psii_nccode
update_target teelff.trret_pht.. cto", pellf.phi_predicttte
39
update_target ttelf.target.ssi_rrede.
self..ss__peedicter
Listing I: Python-like pscudocode for TD-JEPA (simplified).
Pseudocode for the default variant of TD-JEPA is shown in Listing ! assumed to be averaged across twin networks.
the output of predictors is
34


Reviews:
[{"review_id": "8iQzosHaWT", "reviewer": "Reviewer_WTAu", "summary": "The paper proposes TD-JEPA, a latent-predictive representation learning method for zero-shot reinforcement learning (RL) that uses temporal difference (TD) learning to capture long-term multi-policy dynamics from offline, reward-free data. By training separate state/task encoders, policy-conditioned predictors, and parameterized policies in latent space, TD-JEPA enables zero-shot optimization of arbitrary reward functions—even for challenging pixel-based inputs. Theoretical guarantees and extensive experiments across 13 datasets (ExoRL/OGBench) validate its effectiveness.", "questions": "In terms of methodology, could you further compare the approach proposed in this paper with that proposed by Motivo [1]?\n\n[1] Andrea Tirinzoni, Ahmed Touati, Jesse Farebrother, Mateusz Guzek, Anssi Kanervisto, Yingchen Xu, Alessandro Lazaric, and Matteo Pirotta. Zero-shot whole-body humanoid control via behavioral foundation models. ICLR, 2025.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Sound Motivation and Novel Method: The work addresses a key limitation of existing latent-predictive methods (single-task/one-step prediction, on-policy dependence) by leveraging TD learning for offline, multi-step, multi-policy dynamics modeling. Its design—separating state (low-level dynamics) and task (high-level context) encoders, plus TD-based losses and regularization—balances innovation and practicality, avoiding representation collapse while enabling offline training.\n2. Comprehensive and Convincing Experiments: TD-JEPA is evaluated across 65 tasks (locomotion, navigation, manipulation) with proprioceptive/pixel inputs. It matches or outperforms SOTA baselines. Ablations confirm the value of multi-step prediction and separate encoders, while fast adaptation results show pre-trained representations boost sample efficiency for fine-tuning.\n3. Rigorous Theoretical Foundations:\nThe paper provides solid theoretical support.", "weaknesses": "TD-JEPA relies on FlowQ-like behavioral cloning regularization to handle low-coverage datasets in OGBench, but its performance under more extreme data scarcity (e.g., critical action gaps, sparse trajectories) is not fully validated. It would help to add analysis on how TD-JEPA’s performance decays as data coverage decreases, and compare it to methods specifically designed for low-quality offline data. This would better demonstrate its practical applicability to real-world scenarios where data is often incomplete.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "dT0XUdXS2N", "reviewer": "Reviewer_u6m8", "summary": "TD-JEPA targets learning latent policy dynamics models through successor features from a collection of reward-free off-policy data. Unlike prior work often limited to one-step predictions, this TD loss enables the model to learn representations predictive of long-term, policy-conditioned latent dynamics. The system trains four components directly in latent space: a state encoder ($\\phi$), a task encoder ($\\psi$), a policy-conditioned multi-step predictor ($T_\\phi$), and a set of policies ($\\pi_z$). The predictor learns to approximate successor features, which allows the agent to perform zero-shot optimization of any new reward function at test time.", "questions": "In the abstract, it is written “(TD) learning enables learning representations predictive of long-term latent dynamics across multiple policies…” . Would it be more accurate to say “long-term latent task/policy dynamics” instead?\n\nFor me, the notational assignments were difficult to maintain as I was reading through the relatively terse section 3. I wonder if there is a more distinguishing notation between \\psi and \\phi?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "TD-JEPA introduces a method for learning long-term transition dynamics rather than single-step dynamics. There is extensive experimentation in simulated benchmarks across DMC and OGBench showing general improvement over other zero-shot RL baselines.", "weaknesses": "Perhaps it is because I am not in the immediate area, but it was challenging for me to determine what the real-world impact of this method is. It certainly makes sense in relation to recent latent prediction models, but the introduction may benefit from some framing that takes a step back. Is there more application-facing work that has called for learning from large reward-free multi-task offline datasets that could be cited? And for zero-shot RL? If this doesn’t exist, why not yet?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "4vuNlVgxsA", "reviewer": "Reviewer_FzbF", "summary": "The authors propose TD-JEPA, a method for learning policy-dependent representations of long-term dynamics that can be used for zero-shot RL. They learn the representations by combining temporal difference methods with the self-predictive framework of JEPA. Across a number of experimental benchmarks, they show that TD-JEPA matches or outperforms existing baselines, especially on pixel-based environments. The authors also discuss a number of theoretical properties of the proposed algorithm.", "questions": "Please see above.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "The paper is well written and clearly motivates the proposed algorithm. The experiments are comprehensive and show convincing performance. The proofs justify the algorithm theoretically and build on techniques used in previous works.", "weaknesses": "I do not have any major weaknesses for this work, but I point out things I found interesting which could benefit from more detailed discussion (although I understand some of these points might be outside of the scope of this work).\n\n1. What is the computational cost of the different proposed methods? I know many of the considered methods train multiple function approximators, but it would be nice to have at least a brief discussion of the training speed for each method.\n2. I found the comparison between BYOL and TD-JEPA interesting. In particular, the contrast between modeling expert policies vs. policy conditional measures. I wonder if there are other simple analyses/visualizations that can pinpoint this down further, beyond performance on benchmarks. For example, is one representation more robust to noise / generalizes better (perhaps since TD-JEPA is better on pixel environments, this would be the case)?  \n3. For pixel-based environments, how do the visual features compare from using a purely visual pre-training strategy (MAE for example) to learning the encoders with TD-JEPA (or another RL method)?\n\nIn general, there seem to be a number of different \"representation learning\" methods that are useful for RL in the community. The authors already briefly do this in the discussion section, but continued conversation about the long-term desiderata for a representation, beyond performance on zero-shot benchmarks would be interesting for the community.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "8ULd0AqA1O", "reviewer": "Reviewer_ghUc", "summary": "This paper proposes a new state representation learning method called TD-JEPA. The main idea of TD-JEPA is to use TD learning to learn BYOL/JEPA-like self-predictive representations, where positives are sampled from geometrically distributed future states (in the MC case). Importantly, instead of learning representations w.r.t. the behavioral policy (like BYOL-$\\gamma$), the authors simultaneously learn a latent task embedding $z$ and train a $z$-conditioned policy and the corresponding representations, somewhat similarly to FB representations. The authors show that TD-JEPA representations enable zero-shot RL, outperforming previous zero-shot RL approaches on ExORL and OGBench. They also demonstrate that these representations can be effectively fine-tuned for offline-to-online RL.", "questions": "- In Table 4, why do some of the settings have zero hidden layers?\n- While I understood that the TD-JEPA representation approximates the successor measure of $\\pi^z$ in a low-rank manner ($F_z^\\top B$) for a **fixed** $z$, can the authors explain what kind of behavior **set** this method would learn? In other words, are there any explicit descriptions of learned skills? For example, for HILP, one can describe its skills as those that maximally span the isometric latent embedding space. Do the authors have a similar intuitive (yet \"correct\") description of skills learned by TD-JEPA? (I guess this question directly applies to FB representations as well.)", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "This is a well-written paper with solid theoretical and empirical results. The proposed method is (to my knowledge) novel, even though there are a number of closely related (but different) works, such as FB and BYOL-$\\gamma$. The authors compare their method with previous methods across diverse categories, and convincingly demonstrate its effectiveness on a wide array of tasks and settings.\n\nAnother strength is that the authors provide a solid theoretical analysis of their method. They theoretically show that their method (more or less) low-rank approximates successor measures. With this connection, they are able to relate downstream performance to their loss. While I didn't exhaustively check the correctness of the proofs in Appendix, they appear to be solid and have some new, intriguing aspects of their own.", "weaknesses": "I don't see any major weaknesses in this work. While I believe the current results are already enough for an ICLR publication, I think they could be further strengthened by demonstrating the performance on even more complex, \"new\" environments that go beyond the standard benchmarks used in previous work (Motivo is a good example of this, and I think this could also be done in a separate follow-up work). Another nitpick is that the related work section is placed in Appendix. Especially given that this area is (relatively) dense, I believe discussions about related work are essential to understanding this method, and would like to encourage the authors to move this section to the main paper in the final version.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Oral)