Self-Predictive Representations for Combinatorial Generalization in Behavioral Cloning
Authors: Daniel Lawson
Adriana Hugessen
Charlotte Cloutier
Glen Berseth
Khimya Khetarpal
Venue: ICLR 2026 Poster (2026)
Keywords: Sequential Decision Making
Combinatorial Generalization
Representation Learning
Primary area: reinforcement learning

Abstract: While goal-conditioned behavior cloning (GCBC) methods can perform well on in-distribution training tasks, they do not necessarily generalize zero-shot to tasks that require conditioning on novel state-goal pairs, i.e. combinatorial generalization. In part, this limitation can be attributed to a lack of temporal consistency in the state representation learned by BC; if temporally correlated states are properly encoded to similar latent representations, then the out-of-distribution gap for novel state-goal pairs would be reduced. We formalize this notion by demonstrating how encouraging long-range temporal consistency via successor representations (SR) can facilitate
generalization. We then propose a simple yet effective representation learning objective, $\text{BYOL-}\gamma$ for GCBC, which theoretically approximates the successor representation in the finite MDP case through self-predictive representations, and achieves competitive empirical performance across a suite of challenging tasks requiring combinatorial generalization.

Published as a conference paper at ICLR 2026
SELF-PREDICTIVE REPRESENTATIONS FOR COMBINA- TORIAL GENERALIZATION IN BEHAVIORAL CLONING
Daniel Lawson 1,2,+ Adriana Hugessen 1,2,+
Charlotte Cloutier'
Glen Berseth 1,2,1 Khimya Khetarpal 1.3.1 1 Mila Universiiéé de Montréal
'Google DeepMind
ABSTRACT
While goal-conditioned behavior cloning (GCBC) methods can perform well on m-distribution training tasks, they do not necessarily generalize zero-shot to tasks that require conditioning on novel state-goal pairs, i.e. combinatorial generaliza- tion. In part, this limitation can be attributed to a lack of temporal consistency in the state representation learned by BC; temporally correlated states are properly encoded to similar latent representations, then the out-of-distribution gap for novel state-goal pairs would be reduced. We formalize this notion by demonstrating how encouraging long-range temporal consistency via successor representations (SR) can facilitate generalization. We then propose a simple yet effective representation learning objective, BOOL-y for GCBC, which theoretically approximates the suc- cessor representation in the finite MDP case through self-predictive representations, and achieves competitive empirical performance across suite of challenging tasks requiring combinatorial generalization.
INTRODUCTION
Generalization has been oong-standing goal in machine learning and robotics. Recently, large-scale supervised models for language and vision have demonstrated impressive generalization when trained over vast amounts of data. In robotics, this has motivated large-scale behavior cloning (BC) models trained offline datasets of diverse demonstrations (Ghosh al. 2024 Kim et al. 2024). However, these models still suffer from lack of generalization. In particular, while BC methods can perform well on tasks directly observed in the dataset, they often fail to perform zero-shot transfer to tasks requiring novel combinations of in-distribution behavior, known as combinatorial generalization. In the robotics domain, where demonstration data is time-intensive and costly to produce, simply scaling the dataset is often not possible. Hence, achieving this type of generalization algorithmically will be critical to unlocking the potential for large-scale supervised policy training.
The property of combinatorial generalization has been previously formalized as the ability to "stitch" (Ghugare et al.) 2024). Here, stitching refers to the ability of policy to reach a goal state from start state when trained on dataset, which provides sufficient coverage of the path to the goal, but which does not contain a single complete trajectory of the path. The lack of stitching observed in coal-conditioned behavioral cloning (GCBC) and, more generally, supervised learning, can be understood through the inductive biases of the model. By construction, BC methods do not encode the inductive bias that the observed data are generated from a Markov decision process (MDP). In contrast, reinforcement learning (RL) policies that are trained via temporal difference (TD) learning directly utilize the structure of the MDP to pass information through time using dynamic programming. Offline RL (Levine et al. 2020) has been proposed as method for achieving stitching in policies trained on offline datasets. However, these methods are challenging to scale due to the instability of bootstrapping in TD learning when combined with fully offline training. Scaling has been more successful with supervised methods, such as in robotics, where training robot foundation models with
BC (Ghosh et al. 2024 Kim al. (2024) on large-scale datasets (O Neill et al. 2024 Khazatsky et al. 2024) can lead to more general-purpose policies.
Co first author. f Equal advising. Project page: https: wttpsellfffpll fttp//-eee--rrdd-bb. wit.bb. t Correspondence: [daniel. lwsso, adriana
aarriana.knacchor. www.kcabtttttc...o whwgbessennhhheeeennee @mi quebec

---

Published as a conference paper at ICLR 2026
C
b/b
o(sy) Si
C
representation e
predictor 1
((tt)
(
((a))
((5))
Sh
(a)
(b)
Figure 1: (a) Self-predictive Representations. We consider training on trajectories like, + and S6 > sy, which intersect at w, and then evaluating on task like so + sf. requiring combinatorial generalization. (h) Representation learning with BYOL. We predict future state representations ((*++)) via y ()((s))), a), and also predict backwards with 4/a (+(()+))). The target offset is sampled geometrically; k ~ geom(1 (1 7). Stop-gradients are denoted by //. We provide more details on the
training procedure C in Section 4.2
goal-conditioned policy being general-purpose implies has learned an implicit world model of the enviromment (Richens al. (025). From this intuition, a key desiderata make policy's representations align with the (latent) dynamics of the underlying environment, in order to obtain more robust goal-conditioned policy. However, an open question here which representation learning objective best achieves this property. We begin to investigate this question with Bootstrap Your Own Latent (OL) framework (Grill al. 2020). which in RL, learns a representation space through predicting future latent states (Schwarzer et al. (020),, without requiring negative samples nor TD-learning While the standard BYOL objective has been shown to learn representations capturing spectral information about the one-step transition dynamics (Khetarpal al., 2025), find that a key property is to capture temporally extended information, leading us to (I) propose a novel objective, BYOL- which predicts future states geometrically (Figure Tb), and (2) present unifying framework (Table for understanding objectives related to the successor representation (Blier et al. 202T including contrastive learning, BYOL, BYOL- and a novel application of known bDsed approximation of the SR as an auxiliary loss for BC. Namely, we quantity how these methods uniquely behave when applied to data collected by mixture of policies which is encountered in practical BC settings.
In the finite, single-policy MDP case, we show that BYOL-y approximates the successor repre- sentation. As with other non-TD methods, in the mixture-policy case, BYOL-r corresponds to approximating a mixture of SRs, however, with less pessimism than existing contrastive objectives. Qualitatively, the B BYOL-y objective learns representations that encode long-range temporal distance between states on mixture datasets more faithfully, as compared to TD, than contrastive learning (Figure 2 Empirically. on the challenging OGBench suite (Park et al. 2025), we demonstrate that BYOL-y augmented GCBC outperforms other methods (Table on average, and is robust to combinatorial generalization with increasing horizons (Figure 3 Our representation also be extended to hicrarchical setups (Appendix C), which leads to further improvements in generalization.
2 RELATED WORK
Stitching in Supervised Methods. Outcome (goals or conditiccn behavioral cloning (OCBC) methods (Schmidhuber) 2020 Chen et al. 202T: Emmons et al.) 2022) provide a simple and scalable alternative to traditional offline RL. (Levine et al. 2020) methods. However, these methods do not properly "stitch" and generalize to unseen outcomes (Brandfonbrener et al.) 2022 Ghugare (ct al., 2024). To reduce this problem, various works have proposed augmenting training data used by BC methods. Some work incorporates methodlogy from offline RL to label returns or goals for downstream SL (Char et al. 2022] Yamagata et al. 2023). Other work has considered relabeling goals through clustering states (Ghugare al. (2024). which relies good distance metric, utilized planing Zhou et al. (2024) for goal relabeling, or generative models to synthesize new trajectories (Lu Tee al. 2024b. Rather than using models to generate data, combinatorial generalization
2023
2

---

Published as a conference paper at ICLR 2026
can be achieved by planning with generative models (Luo al. 2025 In this work, neither require explicit Q-learing, generative models, or perform explicit planning.
Representation learning in RL. Our objective is most closely related to approaches using auxiliary
BYOL objectives in online RL (Gelada et al. 2019 Schwarzer et al.] 2020 Ni et al. 2024 Voelcker
et al. 2024 These objectives can help with sample-efficiency. such in challenging, partially observed environments with sparse rewards, or with noisy states. Additionally, self-predictive
dynamics models are used in planning and model-based RL (François-Lavet al. 2019; Ye et al.
|02T] Hansen et al. 2022). Various works have also characterized the dynamics of B YOL objectives
the RL setting, showing that BYOL objectives capture spectral information about the policy's transitions (Tang et al. [023] ( Khetarpal et al. 2025). In the offline setting, how well Joint Embedding Predictive Architecture (JEPA) world models generalize when used for explicit planning has been studied (Sobal et al. (2025), however, not for combinatorial generalization. Additionally, certain representation structures for value functions, namely quasimetrics (Liu et al.) 2023 Wang et al.) 2023]
Wang and Isola 2022; Myers et al. 2024) can also lead to policies that better generalize to longer
horizons Myers et al. 2025a).
Successor Representation (SR) (Dayan 1993) objectives, such as successor features (SF) Barreto
ctt 2017), and the successor measure (SM) (Blier et al., 202I) have been widely used for gener- alization and transfer in reinforcement learning (Carvalho et al. 2024). Similarly to BYOL, these
objectives have been used for representation learning in RI. (Lan et 2022, Farebrother et al. 2023).
While prior BYOL methods either perform 1-step, or relatively short fixed n-step prediction, neither of these choices directly approximate the successor measure. Our setup is most related to temporal representation alignment (TRA) (Myers et al. 2025b), which recently proposed using contrastive learning as an auxiliary objective for BC to improve combinatorial generalization. In this work, we further build on the relationship between the SM and combinatorial generalization, and propose new objectives which can lead to better performance.
BACKGROUND
Controlled Markov Process. We consider goal-conditioned decision-making, with states S, actions A, goals g E S, initial state distribution Po(s). dynamics p(#++1 |st, a), and with policies T (a)ss,)).
Successor Representation (SR) and Successor Measure (SM). In finite MDP, the successor representation (SR) Dayan 1993) of a policy is: (s, := Σi20 1'1 (( 8o = S,
We use the convention of counting from $141, writing in matrix form M = E (PT) (Pr))) The transition matrix for policy π is px with P = Σ. T[a|s i) where P.,a.j St = i,a) The successor representation also satisfies the bellman equation, (F[I->P"-- For a fixed policy, the successor representation de- scribes type of temporal distance between states. The successor measure (SM) (Blier et al. 2021) extends SR to continuous spaces S: (s, == Σizo? 7a P((++1 E X ()) VX V C. We also define the normalized successor representation, or measure M= = (1 - - In the fi- nite case, the normalized successor representation 11 has rows that sum to one like transitions p=. We also define the state occupancy via M M=(s') = Ex-po(=) [M s'). Another quan- tity, successor features (SF) (Barreto al. 2017) are the expected discounted sum of future fea- tures (0s) E <Rd: (s) = E Ei20 Y' ((1()))) So s,
p( M pa P
We can relate SFs to the SM with ab (s) S Sg. (s, ('s') These quantities can also condition an action, e.g. M (s, s")..
3.1 REPRESENTATION LEARNING
We begin with two representation learning methods that approximate the density of the SM.
Contrastive Learning. Temporal contrastive learning used in MDPs (Eysenbach et al. 2022 is related to Monte Carlo (MC) approximation of the (discounted) successor measure. This can be implemented with an InfoNCE (van den Oord et al., 2019) loss that maximizes the similarity of a positive pair consisting of state 8, and future state from the same trajectory 8+, and minimizing
3

---

Published as a conference paper at ICLR 2026
the similarity of negative pairs consisting of s and random states
log e ((((4)))06111)
min IE P.
Se semp((s) Ae-goom( (1 1-)) 2 mp((s)
(
(I)
common choice for the energy function f is the inner product f(v(())((++) (()))) A key aspect to note is that the positive sample 84 comes from an MC sample from ~ M (3 84). The optimal solution to ①1 gives (s, = C (eep((((ss) (o(((((s))((+)))) p(s+).
Temporal-Difference Approximation of SR (TD-SR) We consider Forward-Backee Backward (Touaii and Ollivier 2021)-like loss that approximates the successor measure for a fixed policy using TD learning, discussed by Touati et al. (2023), which we call TD-SR.
min s)mp(s),e(<p(s) P.V s141 wp (544115))
[(v(s))) 6(s') - (((+))))((()))))) 21
stmp(s) #1110p" (ec11l8s)
(ws))) (5(+1))
TD-SR learns an approximation of the successor measure with factorization (M s_, x (s,)" (ss))T((s) o(*+) p((+) using TD learning. Given transitions (8t, 81+1) sampled by policy æ, the second term relates fitting M (st, 8(+1). Given an independently sampled state s', the first term bootstraps estimate of MM (st, s') from M² (sc-1. s'), where é, denote koopggrddient operations. In Appendix D we further elaborate on the relationship between the TD-SR loss, and CL. Particularly, in the limit, n-step version of TD-SR is related to CL.
BYOL We now look at an objective that captures information about single-step transitions instead of the successor measure. In the context of RL, self-predictive models jointly learn a latent space and a dynamics model through predicting future latent representations. Self-predictive models rely on latent bootstrapped targets (B YOL)) (Grill et al. (2220), avoiding reconstruction (generative models), or negative samples (contrastive learning). Self-predictive models are an instance of joint-embedding predictive architectures (JEPAs) (LeCun 2022 Garrido et al. 2024).
Given an encoder which produces a representation 2; = ((*)). and dynamics (fz++1 |z)) for fixed policy n, minimize the difference between our prediction and target representation in latent-space:
min E S (opp [f [f(((()) [(s++1))]
6,47
Where f measures the discrepancy between representations, such as the squared 12 norm, and refers to an EMA target, or stop-gradient. Variants of this BYOL objective have been widely used to learn state abstractions, and work as an auxiliary loss to value-function learning (Gelada et al. 2019 Schwarzer et al. 22000] Ni et al. 20024.. In finite MDPs, this objective captures spectral information about one-step transitions px (Tang al. 2023 Khetarpal et al. 2025). discussed in Appendix
3.2 COMBINATORIAL GENERALIZATION FROM OFFLINE DATA
We now shift focus on how can learn policies from offline data using behavioral cloning, and then introduce a combinatorial generalization gap that arises in this setting.
We consider dataset D = {(+b) ab. s.. ) J== composed of trajectories generated by set of unknown policies {a)(s) Goal Conditioned Behavioral Cloning (GCBC) trains a policy πe with maximum likelihood to reproduce the behaviors from the dataset. After sampling a current state, a goal is sampled as a future state from the same trajectory:
max (Fc = max Te
[log Lognolals.g = s4)] =+]]]
By ())r)( SM (s) an 1), 3 (s.s))
Generalization gap. While this policy can perform well in-distribution, the behavior cloning policy struggles to generalize to reach goals from states that are not in matching training trajectories. We now review more formal definition of this type of generalization gap.
We consider Lemma 3.1 from Ghugare et al. (2024), which says there exists a single Markovian policy ((als) that has the same occupancy the mixture of policies: (s) Bpp(,) (M&(ss) This
4

---

Published as a conference paper at ICLR 2026
policy also has construction: B(a () : E, (a s)p(33 s). where p(β) s) is the distribution over policies in reflected by the dataset.
Using the successor measure of the individual policies, and the mixture policy, we can quantify gap between accomplishing out-of-distribution tasks versus in-distribution training tasks (Ghugare et al. 2024)
E Rq Po (so) Ag (80.*))
10 (800 $g)]
By (o(();)))))))))))))))))) (2), a0 [4" (80, $g)] By 3 (fo.*,)
requiring combinatorial generalization
in-distribution training lasks
Here, is performance metric of the policy тe such the success rate to reach Sy from so- As we perform well on in-distribution tasks due to a correspondence to Equation (4). the BC policy has no guarantees for the first term. This is because after sampling state, the goal is sampled from the successor measure the mixture policy.
4 USING REPRESENTATIONS FOR COMBINATORIAL GENERALIZATION
In this section, we aim to reduce the aforementioned generalization gap. We consider policy trained with the BC objective Tee to made more robust the tasks requiring combinatorial generalization through representation learning. We begin with setup similar to Equation (5). but with a shared initial state so for both the in-distribution and out-of-distribution task. For in-distribution task, we sample goal as before, labeled as 8 However, for the out-of-distribution task, we sample goal sy to be a state that can be reached by the mixture policy after sw- (6):
E Dg (e0 15 (sowew)
E M (swss)) [2 (sg, 8t))]
u Tel (sp Spn)
(6)
extendedd task requiring generalization
in-distribation task
E By-p(()) (oo-pp())))0mMmi))) (P(() eu Ser (sal8w)
MP [u=e (())))
(so, p(sw))
(7)
want invariance with respect to future goals through
Then, in Equation 7 we add a goal representation o that processes the goal before going to policy Te. Intuitively. policy could achieve the out-of-distributin task by first going from to sw finddistibbution. and then completing the remaining task s, to sy. In essence, we want that when conditioning on d(ss), the policy should first go 8w, which can be achieved by learning é where ((sw) is similar Φ(sy) (Myers al. 2025b). More formally, for s M (sw. sy) invariance o(s,) = ((s..). From this observation, we can understand obtaining a representation ø related the successor measure of the mixture policy 3 can be beneficial.
The BYOL framework would be simple way to learn these representations that capture temporal dependencies. However, in Section 5 we demonstrate that simple BYOL objective empirically leads to limited generalization when used as an auxiliary loss. Intuitively, the standard B BOL objective directly approximates the one-step transition dynamics, not the successor measure, and struggles to capture relationships between distant states, separated by several trajectories.
4.1 BYOL-y: CONNECTING SELF-PREDICTIV OBJECTIVES TO THE SUCCESSOR REPRESENTATION
To build better self-predictive representations, we propose BYOL-y which allows us use the BYOL framework to capture temporally extended information, i.e. successor representations. Given a state
8t, BOL objective samples prediction targets from one-step transition as in Equation B However, we make modification to predict empirical samples from the normalized successor measure:
LBYoL- (0, v) Egyr p(s),k>geom([ (aokk-gooll1-)).-ca -= (114))))) [f((()()))) ((s++)))]
(8)
Where f refers to an energy function, ø refers to the encoder, and & the predictor. With y = 0, we have 8t+4 8441 corresponding approximation of the one-step transitions, recovering the base
5

---

Published as a conference paper at ICLR 2026
Method
Approx. 7 T~
Approx. 18,}
Batch
TRA (CL)
84), (St, 8,) )88 (8t, 86+1 ) 83) 8
TD-SR
(s,
BYOL
p² (5(41155)))
(#), 8841
BYOL- (ours)
M (8,84)
(81, 84)
Table 1: Auxiliary Representation Objectives. We provide an overview of the representation objectives we consider. In the first two columns, we label the quantities which representations are approximating in finite MDPs, either from datasets with trajectories collected from single policy T 3, or a mixture of policies T ~ {3,}. We provide additional derivations for mixture datasets in Appendix F In the last column, we list samples used for each objective, where the superscript denotes the number of loss terms for a pair of samples.
BYOL objective. Figure 1b depicts our overall representation learning objective. We can view this objective as iteratively minimizing an upper-bound on the error between (b(((s)) and target the true successor features of the policy uhth with changing basis features é. With convex f. by Jensen's s inequality we have:
LBVOL 2 Ee, [(((((())) 1,, (s.. ((5+)) = E., [(()((()))) (1 (s) (1
(9)
Specifically, we precisely show the relationship of BYOL-Y to the SR with the following result: Theorem 4.1. Given a finite MDP with linear representations & E Bllixd, and predictor y E Rdddd under assumptions of orthogonal initialization for & (Ass. E. D. a uniform initial state distribution Po(s) (Ass. E.2) and symmetric transition dynamics (Ass. E.3 minimizing the seff-predictive learning objective LAVOL- (ob b approximates spectral decomposition of the successor representation M xe T,कT, corresponding to successor features (1 W5 as tT
Proof in Appendix F.2 where show that existing theory (Khetarpal al. 2025 also translates to the proposed BYO objective. Finally, we can see the relation between this objective and CL ( with the most striking difference being the removal of the denominator involving negative samples. Surprisingly, we reveal that this simplified system still captures similar information and can also lead empirical generalization Section 5.1). while relying neither on TD learning nor negative samples.
BYOL-> Variants. We discuss few variants on our base objective, namely, we evaluate bidirec- tional prediction (Guo et al., 2020 Tang et al. 2023) where we add additional backwards predictor vb which predicts past representation from the future. We also utilize an action-conditioned variant of the forward predictor (0xt), az), which can be interpreted as a temporally extended latent dynamics model, or capturing information about (M a, giving:
LBYOL ) E g ~M ( ae), Ø(++)) f(((s)) (+(()))) (10)
For f. we choose cross cross-entropy loss between (softmax) normalized representations, similar to DINO (Caron et al. 2021): fcr(a, b) softmax(1) (b) log softmax(a). However, we find that a normalized
loss, Jtg |ল beह I1ž. commonly used BYOL setups (Grill al. 2020 Schwarzer et al.
2020) also works, which we ablate in Section 5.4
4.2 TRAINING A POLICY WITH AUXILIARY REPRESENTATION
We consider BYOL-- and other objectives as auxiliary losses for BC policies (alss, to improve their generalization. We label all parameters θ = (0, o, )), with the parameters of the encoder and predictor corresponding é, vh, respectively. The pheicy-head (0) transforms representations to actions via an MLP: m (a) g) MLP, (concat( (())) With this policy, we train with the objective:
Eg E3,~0(8,),70,, [Lnc(O) + aLaux (1) $]
(11)
The term Lac updates the parameters of both the policy head 9 and its inputs, i.c. the encoder 6, while Laxx updates 1h, & but not 0. With affected by both terms, the BC loss ensures that the

---

Published as a conference paper at ICLR 2026
representation is sufficient for action prediction, preventing collapse, which can be an issue under certain representation learning objectives such as BYL.. Additionally, the auxiliary loss prevents overfitting and help generalization for the policy. We provide additional details about the architecture in Appendix A
We wish learn representations related the successor measure of the mixture policy motivated by Equation 7 However, there are trade-offs with the representation learning objectives in terms of the quantities they approximate and the data they use, as shown in Table When we directly have full MC samples from a single policy (T ~ 9), TRA, TD-SR, and B YOL-y each capture information related to its SM. However, in practice, we only have MC samples from individual policies {3,}- rather than the mixture.
First, we consider using TD-SR as an auxiliary loss for BC, as we can see it still approximates the correct quantity. To our knowledge, we are the first to study this objective as an auxiliary loss for BC. TD-SR explicitly can "stitch" across policies, i.e. ((s))ssss)) via (8(())1))( for
St, Se+1 v² (r)) (s141 18{)) and 5'~ (s) However, we wish to understand if we can obtain representations that help with generalization without TD, as this may be more scalable when applied with policy learning. This leads us to quantify how MC methods behave when trained on mixture datasets.
While 1-step MC methods like BYOL are consistent across dataset composition, we can see that TRA and BYOL-7 approximate different quantities than TD-SR. Namely, rather than approximating the SM of the mixture policy, MC methods capture a mixture of SRs as shown in Table T and Appendix [F] In practice, MC methods still learn relationships between states encountered in different policies by effectively approximating many SRs in single representation space, which qualitatively shown in Figure? 2 Surprisingly, we show that B BOL--y can lead to representations that are similar, or even better than TD-SR without utilizing TD. However, we find that CL as used in TRA leads to pessimism in the relationship between states sampled by different policies. Namely, for states that are not in the same trajectory, they will only be paired as negative examples, whose representations are pushed apart. This also shows up in the denominator of its approximation, with normalization from p (s+) in E, ( (s, s+)/p²(( (8+). On the other hand, this pessimism is not encountered with BYOL-7, which does not utilize negative examples, giving an approximation of Σ3 P(33(s) M (s) (s, by simply predicting latents. Finally, we highlight that BOO--y only computes O(B) loss terms, while CL compute (B²) negatives (st, Sy), and we utilize O(B2) bootstrap terms with TD-SR.
EXPERIMENTS
Now that we have shown a theoretical basis for studying choices of representations, including CL. TD-SR, BYOL, and our new objective (BYOL-7), we study how these methods behave empirically. We compare representation learning algorithms across three axes: (1) First, we compare qualitatively whether the representations appear capture temporal relationships (2) Second, assess representations quantitatively by measuring zero-shot generalization performance on unseen tasks that require combinatorial generalization (3) Third, we assess generalization performance over an increasing generalization horizon. Finally, we perform ablations on the various components of our proposed method to demonstrate the relative importance of each algorithmic choice.
Environments. We empirically evaluate how well our approach can help with combinatorial general- ization on offline goal-reaching tasks on OGBench (Park et al., 2025), which contains both navigation and manipulation tasks, across low-dimensions and visual observations. We focus on navigation environments, where OGBench provides st stitch datasets, that assess combinatorial generalization by training on trajectories that span at most 4 maze cells, while evaluating on tasks that are longer, requiring combining information from multiple smaller trajectories.
Baselines. We benchmark against non-hierarchical methods that perform control from state to low-level actions (e.g. joint-contro)). In addition to BYOL-y used as an auxiliary loss for BC, we evaluate several baselines: GCBC is the standard BC baseline, which we aim to improve upon with representation learning. Offline RL from OGBench, including implicit [V.))-learning (IVL, IQL) Kostrikov et al.) 2022). Quasimetric RL (QRL)/Wang et al. 2023). and Contrastive RL (CRL) (Eysenbach et al. 2022). BYOL is a minimal version of our setup with 1 1-step prediction and loss f1g- TRA
(7 0), only forwards prediction (4yy) without action-conditioning
7

---

Published as a conference paper at ICLR 2026
BYOL-y (ours)
TD-SR
BYOL
TRA
Figure 2: Visualization of the Learned Representation: depicts the similarity between the prediction of the current state representation to the goal representation. For BYOL-Y and TD-SR, we visualize the cosine similarity between 1b(((s), or (s,,)) to o(g) Vs E D for fixed goal g which is indicated by the star marked in red.
(Myers et al. 2025b) is an auxiliary representation objective using contrastive learning related to an MC approximation of the SM. TD-SR is a TD-based approximation of the SM in Equation (2)
used as an auxiliary objective for BC. We also compare to an n-step version of BYOL in Appendix
B.2 and the Norward-Backward (FB) (Touati and Ollivier 2021 Touati et al. 2023) Appendix B.3
Experimental Setup. We match the training details of OGBench, and consider similar representa- tion learning setup to TRA. We found was beneficial to add action conditioning to TD-SR, did not an overall improvement for TRA,
the original setup without action conditionn..... While we use policy ग (Ø((s, s), o(g) and train with action-conditioning for BYOL-Y and TD-SR, TRA originally uses parameterization xt (<b(s)) (())) and does not condition actions. We provide full comparison for changing (s)) to ((s) and action-conditee conditionng in Appendix B.1 for for TRA. Howeve,r obtain similar performance on average with the original setup. For clarity, Table 2 we utilize superscript to denote methods with action-condittoning. Notably, we find that the weight of the auxiliary representation learning objectives (a) can be sensitive both the embodiment, and size of environment (medi um vs large). For each method, we perform a hyperparameter sweep over 4 a values, and report the best result for each environment Table 2] We hold other hyperparameters constant, except with variation between non-visual and visual noted in Appendix A
5.1 QUALITATIVE ANALYSIS OF REPRESENTATIONS
we display qualitative analysis of the representations. We visualize the similarity between the future prediction v for each state to ((g) for fixed goal g. We can see that BYOL-y seems to learn representation that encodes reachability between states, and has similar structure to TD-SR, which is known to approximate the successor measure. TRA and base BYOL seem to both capture similar structure and leam less well-defined latent space. However. BYOL-y and TD-SR have more distinct similarity. and have visible "paths" of similar states. BYOL- also appears capture the most similarity among more distant pairs of states. Compared TRA, hypothesis here is that BOOL-y has more optimistic similarity between distant states due to the lack of negative term in the loss, pushing representations apart. We show additional environments in Appendix H. We also check the correlation of distance in representation space with shortest paths in the maze in Appendix H.2 showing that BYOL-y best captures the structure of the environment.
In Figure 2
5.2 ZERO-SHOP PERFORMANCE COMBINATORIAL GENERALIZATION TASKS
In Table we provide the performance results across all methods. Overall, our proposed method BYOL--, shows improved performance vs. GCBC across most environments, and is either competi- tive with or outperforms TD-SR and TRA. Importantly, we find that minimal BYOL setup does not confer significant benefit over the base GCBC except in non-visual antmaz environments. Generally, auxiliary representation learning with GCBC outperforms existing offline RL methods.
Within the auxillary loss methods, we find that TD-SR and BYOL-y tend to outperform TRA on most environments. While we find that TD-SR outperforms BYOL-Y on environments with smaller state spaces (antmaze- medi um, large) ), find that BYOL-Y's simpler training pro- cedure is beneficial in environments with larger state spaces (humano i dm ze /medium, large), vi sual 1 antma -med i um and vi visull scene- scwne-pleayyy.eee -play).
8

---

Published as a conference paper at ICLR 2026
Dataset
BYOL-y"
BYOL
TRA
TD-SR"
GCBC
GCIVL
GCIQL
QRL
CRL
antmaze-median meee---eeuu-eeee st-uumm-sttttttee ståtch
58
+9+4
6
64
45+11
44+6
29 + 6
+
53 +
antmare- peee--arrg--se tic
19
1746
11
23
313
1842
7
1812
1142
humanci iennnnidmzzeee deezo-medi un-sittch
23+33
2
121.
12
1
364
humanci memancidmsze--ee setsoocar-orrn--ttii
13+33
3+1
5+4
11 3
6+3
1+1
0+0
3+1
4+1
25 2:5
12
14
22
10
2
21
2
1
1
vi satuall-stttaeeeere--tteeee visual scene play
68 a
8
62
= 3
49 2
67 4
6
2
0
2
visuel visuel-aatmaze
26 5
26 5
17
1
29 +
24 4-3
1+1
0+0
1+3 1
11+33
17:
13
25
10
average-state
33
23
26
32
21
19
21
average-visual avecage all
37
32
28
31
34
11
30
35
26
27
32
26
16
14
25
Table 2: OGBench: We find that BYOL-Y performs better overall compared to prior methods. We report mean and standard deviation over 10 training seeds non-visual environments, and seeds in visual environments. We match the OGBench evaluation setup of 5 evaluation (state.goal) tasks, and 50 episodes per task. The success rate is then averaged over the last 3 checkpoints. We color the best non-RL method, and bold values within 95% of its value in the same row. We use superscript a to denote methods utilizing action-conditioning.
Interestingly, in visual-antmant TRA and TD-SR actually seem to hurt performance in com- parison to base GCBC. On the other hand, with BYOL-Y we see no performance degradation over GCBC on the visual environments, considerable improvement over other methods. In Ap- pendix[[ we extend BYOL-y to the hierarchical setting (HBYOL.-7), where we obtain significant improvement over BC baselines, including visual maze environments.
We also find relationship between success and representation quality (Section 5. Namely, in Table 12 we calculate the correlation of representations to shortest path distances and success rate over these same checkpoints. We see that the ranking of methods in terms of average correlation to shortest path (average maze correlation ion) in representation space matches the ordering of methods in terms of average empirical policy success (average maze success).
5.3 EVALUATING GENERALIZATION WITH INCREASING HORIZON
We conduct experiments to understand how suc- cess rate changes as an agent has to reach more challenging goals further away from its starting position. For each ma ze environment, we con- sider the same base 5 evaluation tasks used in Table 2 but construct intermediate waypoints along the shortest path to the final goal deter- mined by breadth-first search. We also include an additional maze environment, giant which all methods have zero success rates to reach distant goals. This gives more holistic view on an agent's performance.
antmaze-giant
BYOL-y (ours))
1RA
GCBC
We display results in Figure Appendix G where we can see how performance drops off for all methods after generalization threshold denoted by the red bar. While all methods can- not fully reach distant goals on giant, we see that BYOL-y has the slowest drop-off in performance. We note that this is a challenging task, that requires stitching up to approximately 8 different trajectories.
Figure 3: Evaluating Generalization with In- creasing Horizons: shows that BYOL-y not only performs well goals in the near horizon, but also, helps to generalize well to goals requiring stitching, after the red bar (> 4).
5.4 COMPONENTS AFFECTING GENERALIZATION
We ablate key components the BYOL-y objective in Table This includes removing action conditioning for forward predictor / (-a), swapping the loss from cross-entropy normalized squared l2 norm (f(.), removing backwards predictor tha, and predicting the representation of the adjacent state (r 0). Both removing action-conditioning. and backwards prediction overall lead to

---

Published as a conference paper at ICLR 2026
Table BYOL-7 ablations. For each ablation, we perform sweep over a, and report the best result per-environment. For environments, we re- port results over seeds (for BYOL-Y, we use the first 4 of 10 in Table 2).
Dataset
BYOL-Y"
ft2
1/b
antmaze-medt um-st um-st itch
6
63 1
56 1 4
67
2
5
anLmaze large sLitch
2
27
6
numanoidmare www.ieddmaedaaieeeeeeeeeee -emmaze-eeiu---ao it
48 5
6
1
humancidmi humancidmaze- ancsocces
wwwdddaaaaaaeeeeeeeeeeeaeee large stitch
14+2
12+6
15+1
2
3+1
stitch
41
20
11
27
25 =
viaual sesal-ande --antma um Mlaual-antm::ee
66
3
03
5
61 4
54
2
25
27
28
average-a:1
33
31
33
24
similar results, but variability per-enviromment. For fr2, we obtain slightly worse average performance, and for O, we see the largest drop-off, especially on humanoli dmaze.
DISCUSSION
Limitations. While we demonstrate that BYOL-Y and other representation learning objectives offer promising recipe for obtaining combinatorial generalization, find that there still exists generalization gap, especially on challenging navigation environments e.g. giant. We also find less significant improvement over BC visual environments, which may motivate additional investigation. Additionally, we may anticipate more benefit from representation learning when applied to larger visual datasets, which has been fruitful other domains.
Conclusion. In this work, we provide a stronger understanding of the relationship between quantities related to successor representations and the generalization of policies trained with behavioral cloning through unified understanding of objectives. We propose new self- ceff-predictive representation learning objective, BYOL-Y, and show that captures information related the successor measure, resulting in competitive choice of auxiliary loss for better generalization. We demonstrate that augmenting behavior cloning with meaningful representations results in new capabilities such as improved combinatorial generalization, especially in larger and more complex environments.
ACKNOWLEDGMENTS
The authors thank Zhaohan Daniel Guo, Faisal Mohamed, and Özgür Aslan for feedback on earlier drafts of the work.
We want to acknowledge funding support from Natural Sciences and Engineering Research Council of Canada, Samsung AI Lab, Google Research, Fonds de recherche du Québec and The Canadian Institute for Advanced Research (CIFAR) and compute support from Digital Research Alliance of Canada, Mila IDT, and NVidia.
REFERENCES
André Barreto, Will Dabney, Rémi Munos, Jonathan J Hunt, Tom Schaul, Hado P van Hasselt, and David Silver. Successor features transfer in reinforcement learning. Advances in neural information processing systems, 30, 2017. URI.
/ee..///666665333 05312
Léonard Blier, Corentin Tallec, and Yann Ollivier. Learning successoo states and goal-dependent values: mathematical viewpoint, 2021. URL
David Brandfonbrener, Alberto Bietti, Jacob Buckman, Romain Laroche, and Joan Bruna. When does return-conditioned supervised learning work for offline reinforcement learning? In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances Neural Information Processing Systems, 2022. URL
net hwtp:/eeet/fouuneeeee
forum 7 www.iiddxXyyotttyyeaeeyo
Mathilde Caron, Hugo Touvron, Ishan Misra, Hervé Jegou, Julien Mairal, Piotr Bojanowski, and Armand Joulin. Emerging properties in self-supervised vision transformers. In 2021 IEEE/CVF International Conference on Computer Vision (ICCV), pages 9630-9640, 2021. doi: 10.1/09/
ICCV48922.2021..00 10051.
10

---

Published as a conference paper at ICLR 2026
Wilka Carvalho, Momchil S. Tomov, William de Cothi, Caswell Barry, and Samuel J. Gershman. Predictive representations: Building blocks of intelligence. Neural Computation, 36(11)22225 2298, 10 2024. ISSN 0899-7667. 1162/neco_a Lateaco____1______ 015
doi: 10. 10.1nec/ncco... 01705. URL [RRl/attss////..ooo
Yash Chandak, Shantanu Thakoor, Zhaohan Daniel Guo, Yunhao Tang, Remi Munos, Will Dabney, and Diana L Borsa. Representations and exploration for deep reinforcement learning using singular value decomposition. In International Conference on Machine Learning, pages 4009-4034. PMLR, 2023. URL /ttttps://arxee http//aaxi...or//bb.../ 2200/33550000 00654
Ian Char, Viraj Mchta, Adam Villaftor, John M. Dolan, and Jeff Schneider. Bats: Best action trajectory stitching, 2022. URL
12026
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Michael Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, 2021. URL
www.wwwwweeeeee.e.ee...ee...eeee nww. net.forum?
id LddaaAA Ld=aaAAPM4BSS
Peter Dayan. Improving generalization for temporal difference learning: The successor representation. Neural Compulation, 5(4):613 1993. doi: 10. 1616/nccoo. 19122/eco..193..4...2 1993.5.4.6133.
Scott Emmons, Benjamin Eysenbach, Ilya Kostrikov, and Sergey Levine. Rvs: What is essential for offline RL via supervised learning? In International Conference Learning Representations, 2022.
S874XAIpKR
Benjamin Eysenbach, Tianjun Zhang, Sergey Levine, and Russ R Salakhutdinov. Contrastive learning as goal-conditioned reinforcement learning. Advances in Neural Information Processing Systems, 35:35603-35620, 2022.
Jesse Farebrother, Joshua Greaves, Rishabh Agarwal, Charline Le Lan, Ross Goroshin, Pablo Samuel Castro, and Mare Bellemare. Proto-value networks: Scaling representation learning with auxiliary tasks. In The Eleventh International Conference on Learning Representations, 2023. URL httt:::/ wwt.:wiwwwwweeeeee..eemm..m.eeee /
Vincent François-Lavee Yoshua Bengio, Doina Precup, and Joelle Pineau. Combined reinforcement learning via abstract representations. Proceedings of the AAAI Conference on Artificial Intelli- gence, volume 33, pages 3582-3589, 2019. URL httpst / http/aaxiiv..gr///eee
1809. 04506
Kevin Frans, Seohong Park, Pieter Abbeel, and Sergey Levine. Diffusion guidance is controllable policy improvement operator. arXiv preprint 2rXiv:2505.23458. 2025.
Quentin Garrido, Mahmoud Assran, Nicolas Ballas, Adrien Bardes, Laurent Najman, and Yann
LeCun. Learning and leveraging world models in visual representation learning, 2024. URL
httts::/arxiv.cor http://wwwww..rg///////oaa//////
abs /403 00504
Carles Gelada, Saurabh Kumar, Jacob Buckman, Ofir Nachum, and Mare G Bellemare. Deepmdp: Learning continuous latent space models for representation learning. In International conference on machine learning, pages 2170-2179. PMLR, 2019. URL
1906. 02736
Dibya Ghosh, Homer Rich Walke, Karl Pertsch, Kevin Black, Oier Mees, Sudeep Dasari, Joey Hejna, Tobias Kreiman, Charles Xu, Jianlan Luo, You Liang Tan, Lawrence Yunliang Chen, Quan Vuong, Ted Xiao, Pannag R. Sanketi, Dorsa Sadigh, Chelsea Finn, and Sergey Levine. Octo: An open-source generalist robot policy. Robotics: Science and Systems, 2024. URL https://do/ itt.oo..rr11..156000000
2024 XX. 090
Raj Ghugare, Matthieu Geist, Glen Berseth, and Benjamin Eysenbach. Closing the gap between TD learning and supervised learning generalisation point of view. In The Twelfth International Conference on Learning Representations, 2024. URL https:
:tttts::/arxa ///rxxiv. wwg/abs...............e........e / 2401
11237
11

---

Published as a conference paper at ICLR 2026
Jean-Bastiens 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 approach to sulf-super learning. Advances in neural information processing systems, 33:21271-21284, 2020. URL https: htt///aaxvvvrrr/// abs/
2006. 07733
Zhaohan Daniel Guo, Bernardo Avila Pires, Bilal Piot, Yean-Bastien Grill, Florent Altché, Remi Munos, and Mohammad Gheshlaghi Azar. Bootstrap antent-predicee representations for multitask reinforcement learning. In Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 3875-3886. PMLR, 13-18 Jul 2020. URL https: / http::://pprceedinggs mlr press/v119/guo20g..to
Nicklas Hansen, Xiaolong Wang, and Hao Su. Temporal difference learning for model predictive control. In International Conference on Machine Learning (ICML), 2022. URL https://
arxiv. bbbrg/bbs/iie http:/abbb22000.04333
Alexander Khazatsky, Karl Pertsch, Suraj Nair, Ashwin Balakrishna, Sudeep Dasari, Siddharth Karamcheti, Soroush Nasiriany, Mohan Kumar Srirama, Lawrence Yunliang Chen, Kirsty Ellis, Peter David Fagan, Joey Hejna, Masha Itkina, Marion Lepert, Yecheng Jason Ma, Patrick Tree Miller, Jimmy Wu, Suneel Belkhale, Shivin Dass, Huy Ha, Arhan Jain, Abraham Lee, Youngwoon Lee, Marius Memmel, Sungjae Park, Iija Radosavovic, Kaiyuan Wang, Albert Zhan, Kevin Black, Cheng Chi, Kyle Beltran Hatch, Shan Lin, Jingpei Lu, Jean Mercat, Abdul Rehman, Pannag R Sanketi, Archit Sharma, Cody Simpson, Quan Vuong, Homer Rich Walke, Blake Wulfe, Ted Xiao, Jonathan Heewon Yang, Arefeh Yavary, Tony Z. Zhao, Christopher Agia, Rohan Baijal, Mateo Guaman Castro, Daphne Chen, Qiuyu Chen, Trinity Chung. Jaimyn Drake, Ethan Paul Foster, Jensen Gao, David Antonio Herrera, Minho Heo, Kyle Hsu, Jiaheng Hu, Donovon Jackson, Charlotte Le, Yunshuang Li, Xinyu Lin, Zehan Ma, Abhiram Maddukuri, Suvir Mirchandani, Daniel Morton, Tony Khuong Nguyen, Abigail O'Neill, Rosario Scalise, Derick Seale, Victor Son, Stephen Tian, Emi Tran, Andrew E. Wang. Yilin Wu, Annie Xie, Jingyun Yang, Patrick Yin, Yunchu Zhang, Osbert Bastani, Glen Berseth, Jeannette Bohg, Ken Goldberg, Abhinav Gupta, Abhishek Gupta, Dinesh Jayaraman, Joseph Lim, Jitendra Malik, Roberto Martín- Martin, Subramanian Ramamoorthy, Dorsa Sadigh, Shuran Song, Jiajun Wu, Michael C. Yip, Yuke Zhu, Thomas Kollar, Sergey Levine, and Chelsea Finn. DROID; A large-scale in in-the-wild robot manipulation dataset. In RSS 2024 Workshop: Data Generation for Robotics, 2024. URL http:::///peenreveww net
Khimya Khetarpal, Zhaohan Daniel Guo, Bernardo Avila Pires, Yunhao Tang, Clare Lyle, Mark Rowland, Nicolas Heess, Diana L Borsa, Arthur Guez, and Will Dabney, A unifying framework for action-conditional self- predictive reinforcement learning. The 28th International Conference on Artificial Intelligence and Staristics, 2025. URL https: / wttpiiiiiiii///ii//aaaaaaaaaaava org/ abs 2406 02035
Moo Jin Kim, Karl Pertsch, Siddharth Karamcheti, Ted Xiao, Ashwin Balakrishna, Suraj Nair, Rafael Rafailov, Ethan P Foster, Pannag R Sanketi, Quan Vuong, Thomas Kollar, Benjamin Burchfiel, Russ Tedrake, Dorsa Sadigh, Sergey Levine, Percy Liang, and Chelsea Finn. OpenVLA: An open-source vision-
model. In 8th Annual Conference on Robot Learning. 2024.
URL httpttts/////aiii/a wwwp:/ii.oooorreeee www.oorrg/gaaaooo.ooo.oooosssess
09246
Ilya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit q-learning. In International Conference on Learning Representations, 2022. URL https
//arxiv. wwwg/aass....................... 20111.111550 066669
Charline Le Lan, Stephen Tu, Adam Oberman, Rishabh Agarwal, and Marc G. Bellemare. On the generalization of representations in reinforcement learning, 2022. URL https: orx. 0t5//2203.00544
Yann LeCun. A path towards autonomous machine intelligence version, 2022. URL https
//openreview www.eeeoooroiddatteeeo
Jaewoo Lee, Sujin Yun, Taeyoung Yun, and Jinkyoo Park. GTA: Generative trajectory augmentation with guidance for offline reinforcement learning. In The Thirty-eighth Annual Conference on
12

---

Published as a conference paper at ICLR 2026
Neural Information Processing Systems, 2024. UtL.attpss//ppeneeeeeeeeee/////s id=k2pNDb2rzy
Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline deffiinnereinnfrrceeetttttamii leamnng Tutorial, review, and perspectives on open problems, 2020. URL htt.tttss//aaiii..rr/aab///////a 01643
Timothy P Lillicrap, Jonathan J Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learning. arXiv preprint urXiv:1509002977 2015.
Bo Liu, Yihao Feng. Qiang Liu, and Peter Stone. Metric residual network for sample efficient goal-conditioned reinforcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 37, pages 8799- 8806, 2023. httpstt::: arxiv 08133
http//wwww.....aaaaa///a org/ 2208
Cong Lu, Philip J. Ball, Yee Whye Tch, and Jack Parker-Holder. Synthetic experience replay. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL https
//openreview net / forummiduum JNQ1AYIUE]
Yunhao Luo, Utkarsh A. Mishra, Yilun Du, and Danfei Xu. Generative trajectory stitching through diffusion composition, 2025. URL
wtt.:org/abb....../////// 05153
Vivek Myers, Chongyi Zheng, Anca Dragan, Sergey Levine, and Benjamin Eysenbach. Learning temporal distances: Contrastive successor features can provide metric structure for decision- making. In Forty-forst International Conference on Machine Learning, 2024. URL https //openreview.oo net.:/iwwwweeeeeeeeeeeeeeeeeeeee forum? forp///oummmmm YCmDr
Vivek Myers, Catherine Ji, and Benjamin Eysenbach. Horizon Generalization in Reinforcement
Learning. In International Conference on Learning Representations, January 2025a. URL
Vivek Myers, Bill Chunyuan Zheng, Anca Dragan, Kuan Fang, and Sergey Levine. Temporal representation alignment: Successor features enable emergent compositionality in robot instruction following, 2025b. URL httttttttt//xiiiiia org/ 2502. 05454
Tianwei Ni, Benjamin Eysenbach, Erfan Seyedsalehi, Michel Ma, Clement Gehring, Aditya Mahajan, and Pierre-Luc Bacon. Bridging state and history representations: Understanding self-predictive rl. In The Twelfth International Conference on Learning Representations, 2024. (ttps:
7/arxiv. org/ abs 2401. 08898
Abby O'Neill, Abdul Rehman, Abhiram Maddukuri, Abhishek Gupta, Abhishek Padalkar, Abraham Lee, Acor Pooley, Agrim Gupta, Ajay Mandlekar, Ajinkya Jain, Albert Tung, Alex Bewley, Alex Herzog, Alex Irpan, Alexander Khazatsky, Anant Rai, Anchit Gupta, Andrew Wang, Anikait Singh, Animesh Garg, Aniruddha Kembhavi, Annie Xie, Anthony Brohan, Antonin Raffin, Archit Sharma, Arefeh Yavary, Arhan Jain, Ashwin Balakrishna, Ayzaan Wahid, Ben Burgess-Limerick, Beomjoon Kim, Bernhard Schölkopf. Blake Wulfe, Brian Ichter, Cewu Lu, Charles Xu, Charlotte Le, Chelsea Finn, Chen Wang, Chenfeng Xu, Cheng Chi, Chenguang Huang, Christine Chan, Christopher Agia, Chuer Pan, Chuyuan Fu, Coline Devin, Danfei Xu, Daniel Morton, Danny Driess, Daphne Chen, Deepak Pathak, Dhruv Shah, Dieter Büchler, Dinesh Jayaraman, Dmitry Kalashnikov, Dorsa Sadigh, Edward Johns, Ethan Foster, Fangchen Liu, Federico Ceola, Fei Xia, Feiyu Zhao, Freek Stulp. Gaoyue Zhou, Gaurav S. Sukhatme, Gautam Salhotra, Ge Yan, Gilbert Feng, Giulio Schiavi, Glen Berseth, Gregory Kahn, Guanzhi Wang. Hao Su. Hao-Shu Fang. Haochen Shi. Henghui Bao. Heni Ben Amor, Henrik I Christensen, Hiroki Furuta, Homer Walke, Hongjic Fang. Huy Ha, Igor Mordatch, Ilija Radosavovic, Isabel Leal, Jacky Liang, Jad Abou-Chakra, Jaehyung Kim, Jaimyn Drake, Jan Peters, Jan Schneider, Jasmine Hsu, Jeannette Bohg. Jeffrey Bingham, Jeffrey Wu, Jensen Gao, Jiaheng Hu, Jiajun Wu, Jialin Wu, Jiankai Sun, Jianlan Luo, Jiayuan Gu, Jie Tan, Jihoon Oh, Jimmy Wu, Jingpei Lu, Jingyun Yang. Jitendra Malik, João Silvério, Joey Hejna, Jonathan Booher, Jonathan Tompson, Jonathan Yang. Jordi Salvador, Joseph J. Lim, Junhyek Han, Kaiyuan Wang, Kanishka Rao, Karl Pertsch, Karol Hausman, Keegan Go. Keerthana Gopalakrishnan, Ken Goldberg, Kendra Byrne, Kenneth Oslund, Kento Kawaharazuka, Kevin Black, Kevin Lin, Kevin Zhang, Kiana Ehsani, Kiran Lekkala, Kirsty Ellis, Krishan Rana, Krishnan
13

---

Published as a conference paper at ICLR 2026
Srinivasan, Kuan Fang, Kunal Pratap Singh, Kuo-Hao Zeng, Kyle Hatch, Kyle Hsu, Laurent Itti, Lawrence Yunliang Chen, Lerrel Pinto, Li Fei-Fei, Liam Tan, Linxi Jim Fan, Lionel Ott, Lisa Lee, Luca Weihs, Magnum Chen, Marion Lepert, Marius Memmel, Masayoshi Tomizuka, Masha Itkina, Mateo Guaman Castro, Max Spero, Maximilian Du Michael Ahn, Michael C. Yip. Mingtong Zhang, Mingyu Ding, Minho Heo, Mohan Kumar Srirama, Mohit Sharma, Moo Jin Kim, Naoaki Kanazawa, Nicklas Hansen, Nicolas Heess, Nikhil J Joshi, Niko Suenderhauf, Ning Liu, Norman Di Palo, Nur Muhammad Mahi Shafiullah, Oier Mees, Oliver Kroemer, Osbert Bastani, Pannag R Sanketi, Patrick Tree Miller, Patrick Yin, Paul Wohlhart, Peng Xu, Peter David Fagan, Peter Mitrano, Pierre Sermanet, Pieter Abbeel, Priya Sundaresan, Qiuyu Chen, Quan Vuong, Rafael Rafailov, Ran Tian, Ria Doshi, Roberto Martín-Martín, Rohan Baijal, Rosario Scalise, Rose Hendrix, Roy Lin, Runjia Qian, Ruohan Zhang. Russell Mendonca, Rutav Shah, Ryan Hoque, Ryan Julian, Samuel Bustamante, Sean Kirmani, Sergey Levine, Shan Lin, Sherry Moore, Shikhar Bahl, Shivin Dass, Shubham Sonawani, Shuran Song, Sichun Xu. Siddhant Haldar, Siddharth Karamcheti, Simcon Adebola, Simon Guist, Soroush Nasiriany, Stefan Schaal, Stefan Welker, Stephen Tian, Subramanian Ramamoorthy, Sudeep Dasari, Suneel Belkhale, Sungjae Park, Suraj Nair, Suvir Mirchandani, Takayuki Osa, Tanmay Gupta, Tatsuya Harada, Tatsuya Matsushima, Ted Xiao, Thomas Kollar, Tianhe Yu, Tianli Ding, Todor Davchev, Tony Z. Zhao, Travis Armstrong, Trevor Darrell, Trinity Chung, Vidhi Jain, Vincent Vanhoucke, Wei Zhan, Wenxuan Zhou, Wolfram Burgard, Xi Chen, Xiaolong Wang, Xinghao Zhu, Xinyang Geng, Xiyuan Liu, Xu Liangwei, Xuanlin Li. Yao Lu, Yecheng Jason Ma, Yejin Kim, Yevgen Chebotar, Yifan Zhou, Yifeng Zhu, Yilin Wu, Ying Xu, Yixuan Wang, Yonatan Bisk, Yoonyoung Cho, Youngwoon Lee, Yuchen Cui, Yue Cao, Yueh-Hua Wu, Yujin Tang, Yuke Zhu, Yunchu Zhang, Yunfan Jiang, Yunshuang Li, Yunzhu Li, Yusuke Iwasawa, Yutaka Matsuo, Zehan Ma, Zhuo Xu, Zichen Jeff Cui, Zichen Zhang. and Zipeng Lin. Open x-embodiment: Robotic learning datasets and rt-x models : Open x-embodiment collaborationO. In 2024 IEEE Internationnal Conference Robotics and Automation (ICRA), pages 6892-6900 6903, 2024. doi: 10.1 10.110///CAA77147..0000 106114777.
Seohong Park, Dibya Ghosh, Benjamin Eysenbach, and Sergey Levine. Hiql: Offline goal-conditioned rl with latent states as actions. Advances in Neural Information Processing Systems, 36:34866 34891, 2023.
Scohong Park, Kevin Frans, Benjamin Eysenbach, and Sergey Levine. Ogbench: Benchmarking offline goal-conditioned rl. In International Conference on Learning Representations (ICLR), 2025.
URL htttt::://ttttttt www.:i.iioooooooooooooo.ooooooo abs 2200. 05929
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, Gretchen Krueger, and Ilya Sutskever. Learning transferable visual models from natural language supervision. In International Conference on Machine Learning, 2021. URL https: /ttp:/rrr.....rg/////0......
Jonathan Richens, David Abel, Alexis Bellot, and Tom Everitt. General agents contain world models, 2025. /ttp://:///aall.ooo//////////
01622
Juergen Schmidhuber. Reinforcement leanr upside down: Don't t predict rewards just map them to actions, 2020.
abs/ /at.s./911.... 02875)
Max Schwarzer, Ankesh Anand, Rishab Goel, R. Devon Hjelm, Aaron C. Courville, and Philip Bachman. Data-efficient reinforcement learning with self-predictive representations. In Inter- national Conference on Learning Representations, 2020. URI.
2007. 05929
Vlad Sobal, Wancong Zhang, Kynghyun Cho, Randall Balestriero, Tim G. J. Rudner, and Yann LeCun. Learning from reward-free offline data: A case for planning with latent dynamics models, 2025. URL /ttttt:///aaxittii
org/ abs 25022552220 14819
Yunhao Tang, Zhaohan Daniel Guo, Pierre Harvey Richemond, Bernardo Avila Pires, Yash Chandak, Rémi Munos, Mark Rowland, Mohammad Gheshlaghi Azar, Charline Le Lan, Clare Lyle, et al. Understanding self- predictive learning for reinforcement learning. In International Conference on Machine Learning, pages 33632-33656 PMLR, 2023. URL https: / /axxiviiivvo//ii org/ abs/
221 2. 03339 9
14

---

Published as a conference paper at ICLR 2026
Andrea Tirinzoni, Ahmed Touati, Jesse Farebrother, Mateusz Guzek, Anssi Kanervisto, Yingchen Xu, Alessandro Lazaric, and Matteo Pirotta. Zero-shot shole-body humanoid control via behavioral foundation models. arXiv preprint arXhv:2504.11054 2025.
Ahmed Touati and Yann Ollivier. Learning one representation optimize all rewards. Advances in Neural Information Processing Systems, 34:13-23. 2021.
Ahmed Touati, Jérémy Rapin, and Yann Ollivier. Does zero-shot reinforcement learning exist? In The Eleventh International Conference on Learning Representations, 2023. URL https
Aaron den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding, 2019. URL UL/_Bt/pe////xxx..orrr//////1//0
Claas A. Voelcker, Tyler Kastner, Igor Gilitschenski, and Amir-massout Farahmand. When does self-prediction help? understanding auxiliary tasks in reinforcement learning. Reinforcement Learning Conference, August 2024. URL httt::://arxdv..og////////21////
Tongzhou Wang and Phillip Isola. Improved representation of asymmetrical distances with interval quasimetric embeddings. NewrIPS 2022 Workshop on Symmetry and Geometry in Neural Representations, 2022. URL /ttpps:///xxivv org/ abs 15120
Tongzhou Wang, Antonio Torralba, Phillip Isola, and Amy Zhang. Optimal goal-reaching reinforce- ment learning via quasimetric learning. In International Conference on Machine Learning. PMIR, 2023. httttt:////xxii.tttp/a
abs 22004/3011.1 01203
Taku Yamagata, Ahmed Khalil, and Raúl Lantos-Rodriguez Q-learning decision transformer: leveraging dynamic programming for conditional sequence modelling in offline rl. In Proceedings of the 40th International Conference on Machine Learning, ICMI.'3. JML.Rorgg 2023.
Weirui Ye, Shaohuai Liu, Thanard Kurutach, Pieter Abbeel, and Yang Gao. Mastering atari games 2021.
with limited data. Advances in neural information processing systems, URL /attps://aax..
org/ abs/ 00210
Zhaoyi Zhou, Chuning Zhu, Runlong Zhou, Qiwen Cui, Abhishek Gupta, and Simon Shaolei Du. Free from bellman completeness: Trajectory stitching via model-based return-conditioned supervised learning. In The Twelfth International Conference on Learning Representations, 2024. URL
https://arae org/ abs/ 19308
15

---

Published as a conference paper at ICLR 2026
A EXPERIMENTAL SETUP
Table 4: Hyperparameters for BYOL-y
Hyperparameter
Shared
actor head
MLP (512.512,512)
representation encoder (d)
MLP (64,64,64)
predictor (+h)
MLP (64,64,64)
encoder ensemble
learning rate
3 x 10
optimizer
Adam
Non-visual
Visual
Gradient steps
1000k
500k
Batch size (EMA)
1024
256
1.0
0.99
0.99
(0.66, 0.99)
(alignment) additional encoder
1 1,6,400,100 1000
[1,6,10,20)
n/a [
impala small
encoder output dimension
64
A.
IMPLEMENTATION DETAILS
In this section we provide more training details for BYOL-Y, and representation learning baselines. We match the training details of OGBench, including gradient steps, batch size, learning rate.
a
Network Architecture. We follow the same gen- eral setup as TRA, where we utilize MLP-based encoders, and action head, For the output dimen- sion of the encoder, we use the state dimension for non-visual experiments, and 64 for visual experi- ments. For the predictor th, we utilize an MLP of the same architecture the encoder. For image- based tasks, there is an additional CNN, which then passes output to the MLP encoder.
π
π
((st)
(($g)
(8g)
<b(31
81)
((sg)
(a)
(b)
Figure 4: Encoder Variation. When training with BYOL, BYOL-Y and TD-SR, we utilize policies with architecture (a) which uses o to process states and goals. We utilize architecture (b) for TRA to match prior implementation, however Appendix B. we train TRA with architecture (a) and action-conditioning.
Representation Ensemble. We follow the setup of TRA which utilizes representation ensembling, such that two copies of the encoder Φ1, Φ2 are in parallel. We also have two distinct predictors 4/1, 4/2 for each ensemble. As input to the policy head, we average the representations, ž (x ( Each representation is trained independently for the BYOL loss, but the BC loss differentiates through both ós.
Alignment. We find that the choice of weight of the auxiliary loss for the representation learning objective is sensitive both the robot embodiment and the environment size. For comparison, perform hyperparameter search over four alignment values for BYOL-y, TRA, and TD-SR, and then report the best value for each environment in Table 2
16

---

Published as a conference paper at ICLR 2026
Discount. For sampling the next-state, utilize discount factor of y 0.99 for all non-visual environments. For visual environments, we perform a hyperparameter search over {0.66, 0.99), however all representation learning methods performed better at 0.66.
A.2 BYOL-7
Target network. For BYO,, we find that exponential moving average (EMA) target networks for the encoder are not necessary for non-visual environments (T 1), but for visual environments, we find that a fast target stabilizes training (T = 0.99):
Orarget
Toolline (1 T) Ptarget
A.3 TRA
In practice, TRA uses symmetric version (Radford et al. 2021 of the InfoNCE objective discussed
in Equation
We write this in batch form, {$ 8 s+.1))+
M rather than in expectation:
B
((
18
LTRA Es
log
log
(12)
B
B
B
カー e(((())))(()1.)))
Additionally, TRA minimizes the squared norm of representations ming, AE, with A = 10-6. For TRA, we search over = {10, 40, 60, 100}.
|M d + M
A.4 TD-SR
Prior work similar to TD-SR, which trains FB for zero-shot policy optimization (Touati al. 2023 typically normalizes with additional loss term that E [øo"] 2 Ia However, we found that adding this loss term was not beneficial to performance our setting and hence not include
TD-SR uses an EMA target network described in A.2 with 0.005. For TD-SR, we search over a = {0.01, 0.05, 0.001, 0.005).
A.5 CODE.
We utilize the OGBench PPak al. 2025) codebase and benchmark, and its extensions in the TRA codebase (Myers et al. 2025b) for equal comparison.
A.6 COMPUTE REQUIREMENTS
We perform all experiments utilizing single GPUs, predominately NVIDIA RTXA8000 and L40S. We utilize 6 CPU cores, 24G of RAM for non-visual environments, and 64G for visual experiments. Experiments take hours for non-visual and 6 to hours for visual environments.
B ABLATIONS.
B.1 ACTION-CONDITTONIIG
In this section, we ablate the component of performing ciniionnnniiioninn for the predictor (sx) vs (st,,, a.) for TRA and TD-SR. We consider similar comparison for BYOL-> in Table 3 For this comparison, when we perform action-condittiningg, we utilize a policy representation 7 o(s),9 g ((g)) we have predictor (8(8, a), and otherwise a(s ( g &(g)) as in the original TRA implementation. We find that results can be environment specific. On average, results improved for TRA, but find improvement for TD-SR, hence in our main Table 2 we include the aetion-conditional results for TD-SR and the action-fice results for TRA to match the original implementation.
17

---

Published as a conference paper at ICLR 2026
Dataset
TRA
TRA"
TD-SR
TD-SR"
anzmaze
5446
57
64 10
64
mrrmaze-large-stiee
1148
7+7
1
23 14
45
36
42 14
5+4
944
642
11 3
an:soccoer-arrne-sttee
14
25
17
22
52
33 4
47 +
49 2
wisoal-aattaaxe-laree-ssii
17:11
22
1.3
29
visual-scene-pla.
16 3
18 2
12 4: 2
14
average-all
27 27
28
32
Table 5: Action-conditioning ablations. We ablate the choice to condition on the first action for predictor v for TRA and FB over 10 seeds for non-visual and 4 seeds for visual environments.
Dataset
BYOL-7"
-46
0
BYOL,
BYOL P
BYOL
anLmaze medlum sLitch
6
67 22
59 + 5
60 =0=0 5
60 =
60 =
58 # 3 3
antia ze- Larg--are---ee Laagg--ttiie sttthhh sumanoidmas medium stitch
2145
19 7
844
13
19 4
8 4
54+55
33
sumanolcmace lazge stlich
14+2
13 22
3
2
3
21
27 =
26 1
23
25 12
11
13 1
averoge-al1
36
28
23
20
Table 6: N-step BYOL ablations. We ablate BYOL-y with an n-step BYOL baseline, where we report results over 4 seeds.
B.2 N-STEP BYOL
We provide an additional BYOL-based baseline that utilizes n-step next-representation recurrent prediction, while BYOL-y uses non-recurrent prediction. We utilize the same BYOL-y architecture with forward prediction, but with the following objective, computing loss with terms, where v5 is shorthand for n recurrent calls:
C ==((W at), (($(+1)) +
a.), a4+1), 4441), ē(s142) + ... ((s)))
Theoretically, in a finite MDP, we can interpret this objective of capturing information up to n- step transitions (Tang al. 2023), i.e. information related to P? ,..) is captured by loss terms V7.
vf) respectively. As BYOL-? captures information related M (1-7) Erz0 7 ''p these two objectives match in theory at 7 0. practice, as n-step operates recurrently, we are constrained shorter n which limits the ability for learning long-horizon information.
Empirically, we report comparison of n-step with n {1,3, 3, 5) BYOL-y in Table We validate our n-step implementation in the base case (n = 1) with the ablation {-3b.. 0f to BYOL- making them equivalent. With increased multi-step prediction (as we increase n), we find worse performance on average.
B.3 FORWARD-BACKWAN ALGORITHM
As an alternative to GCBC, we could instead consider the full -aakww-- (FB) algorithm for zero-shot goal- roal-reaching, proposed in Touati al. (2023). Here, instead of conditioning the policy on a goal representation o(g), we instead condition v on a vector such that jointing learning o and produces policy-dependent successor representation where
M (8,0, (888 (,(s,,,2)) (((4)) p(s+), and Tz (a 8) = argmax F(s, a, T2,
(13)
a and Φ can be learned through a TD relationship analogous to Equation 2 additionally sampling vectors according to some distribution. In the discrete setting, the policy be derived directly from Equation] 13 In the continuous setting, Touati and Ollivier (202T) additionally learn a policy network (s, =), trained to maximize F(s, 2) 2, in DDPG-sttye (Lillicrap al. 2015) algorithm.
18

---

Published as a conference paper at ICLR 2026
At inference time, policy for goal state g be obtained by first encoding the goal state to the z-representation space using the relationship & = Eg [r(s)((s]], which implies & = ((g) for goal-reaching tasks.
For these experiments, we follow the 2 sampling method from Touati and Ollivier (2021) by using a 50-50 mixture of states sampled from / and encoded
((s) and vectors sampled uniformly on sphere of radius v where is the latent dimension. We use network architectures for & and & matching those used in the implementation of FB provided in Tirinzoni et al. (2025), however we keep the number and size of the hidden layers as well the latent dimension consistent with our implementations of other methods. Additionally, we add a BC-loss to the policy loss as regularization, with coefficient We sweep the learning rate over three values: (10-1, 10-5, 10-9) and selected the best performing, averaged over four seeds, for each environment.
Table we compare our proposed BC with auxiliary loss methods (TD-SR" and BYOL-?"), which use successor measure learning use a goal-conditioned value function (GCIQL) or successor measure (FB) to learn a policy through RL. We find that auxiliary loss methods significantly outperform across almost all environments.
auxiliary loss for BC, to value-based methods which instead
Dataset
BYOL-
TD-SR
FB
GCIQL
antmaze-mede um-stitch
61
64
36
2
antmaze sarre-stttthh
21+5
22 3-3
5+ 4
7+2
humanol deeeeldddzeeeedeeee stchh
54
41
26
1243
humanoi mamazeelarge-stieeee
1412
12
2
0=0
antsoccer-arena-sttie
2114
18 12
19 10:14
4
2=0
average-all
34
31
17
10
Table 7: Heward--ackwaar Algorithm. We compare our proposed GCBC with auxiliary loss methods (TD-SR" and BYOL-7") " to an implementation of the Berward-Backwerr algorithm (FB) and an offline RL method GCIQL, both of which learn actor which maximizes goal-conditioned value function. We report best results from a hyperparameter sweep, averaged over four seeds
B.4 CONSTANT ENCODER OUTPUT DIMENSION
Our main experimental setup utilized in Table 2 and other experiments utilize an encoder output dimension of size equal the state dimension, corresponding to ant |s| 29, and humanoia |s| = 69 for non-visual environments. We perform an additional comparison in Table S usiig fixed size latent dimension 64, matching the latent dimension used for visual environments. We can see that the larger latent dimension helps performance for each method on antmaze. Generally, we see similar trends to our prior experiments, such as BYOL-Y performing stronger in humano dma ze experiments, while FB performs stronger on antmaze
Dataset
BYOL-y"
TD-SR"
TRA
TRA (Myers et al. 2025b
antma in
64
7348
67
=
antmaze-largg--titee
18
24
15 10
13
48
4113
41
46
2
sumarooldhaaze-arrgee
12
10+2
4+2
9+1
entsoccer-arent itch
21
1244 4
17
average-all
33
32
29
29
Table 8: Constant Encoder Output Dimension. We conduct an ablation repeating our experimental
setup for representation learning methods with a constant encoder output dimension at 64, For reference, we also report results from Myers al. 2025b).
B.5 GCBC ENCODER ABLATION
We perform ablation where the same architecture as representation learning methods for GCBC. Standard GCBC learns a shared state-goal encoder, o(s, g), while representation learning methods pass inputs through encoder separately, Φ(s), Φ(g) and representations are concatenated
19

---

Published as a conference paper at ICLR 2026
and fed to action head. In our main results, report OGBench GCBC results with shared state- goal encoder, as this is a stronger baseline. However, to better illustrate the impact that auxiliary loss learning has on GCBC performance, in Table g report GCBC results for architecture matching representation learning methods (GCBC-Ø). We especially see a difference in visual environments, where state, goal stacked (64 x 64 6) before going through the CNN.
Dataset
GCBC
GCBC-
artmmse--med.um-sttttee
45
uatmaze-large-stice
3+3
544
29
324
humanciioazz--lageeettei
64:33
4+3
67 4:4
37 +
24
4
visual-sceme-plas
12::2
10
average-state
21
19
average-visual
34
17
average-all
27
Table 9: GCBC Encoder Ablation.
C HIERARCHICAL POLICIES
Although we focus on the impact representation learning in "flat" learning methods, hicrarchical policies also effective orthogonal direction for improving generalization to longer horizon tasks. We demonstrate that BYOL-y also improves on a hierarchical GCBC setup (HGCBC) used in Frans et al. (2025). With HGCBC, we train a high-level policy ="(1 [s,g] that predicts sub-goals 1, and a low-level policy conditioned on sub-goals <"(a 1). which are both trained with BC. Using BYOL-y, we implement its hierarchical version, HBYOL-y, as follows: (1) perform our standard BYOL-> setup, which produces π (a 6(s), 6()) and (2) train hicrarchical policy in the existing latent space of the low-level policy T (o(l) o(s). ((g)). where ő denotes that the representation is fixed for the high-level policy. HB BOLY- allows for both policies to operate in shared representation space, and avoids state reconstruction performed by standard HGCBC. Prior work does not implement HGCBC in visual settings which would require predicting in pixel space. As a baseline, we implement HGCBC-0, which avoids pixel prediction by using GCBC-6 (Appendix B.5). This matches our HBYOL architecture, and two-sta two-stage setup but without representation learning on 6. For HGCBC-0, we also found it was better to only use representations & for the output space of Way and to train shared input encoder from scratch.
Dataset
GCBC
HGCBC HGCBC-
BYOL-
H
HIQL
antmaze um-st s tcc secmazz-large-sse icc.
45
60
61 6
76
94
3
11=8
21
29 9
67 1:
humanai http:n:::mmaa-aaa--ettttte
it
29
35=4
54
2
88 4
humanoicmas
623
420
1412
21 21
28 3
Hisuaa-aatttaze-mmdimn--s- Waaual-acttazz-larg---teee
67
7416
68
84
87 1:2
+4
1941
26
31-3
visual-scene
12
8
17
14
49 +:4
average-state average-visual average-all
21
28
38
47
69
34
34
37
43
55
27
45
Table 10: Hierarchical BC with BYOL--. We report performance averaged over seeds, for HGCBC, HGCBC-0, and (H)BYOL-y. We report GCBC and HIQL results from OGBench. We highlight the the best performing BC methods, bold for methods within 95% of the BC, and darker highlight for BC methods which are within 95% better than HIQL.
In Table 10 we compare between these BC setups and also report results of hierarchical implicit Q-learning (HIQL) (Park et al. (002)). For HGCBC methods, we use a sub-goal (I) step of 25, listing other hyperparamters in Table 1I We that HB is the strongest BC setup, outperforming non non-hierarchical BYOL-y and HGCBC. HBYOL-> is also competitive with HIQL, especially on vi sual 1 -anmmaae environments.
20

---

Published as a conference paper at ICLR 2026
a1
a1
π
π h
hh
24
Zt
2
Φ(st)
$1) ((s1)
((st)
お $ g)
(($t)
st
((sp)
(a) Training
(b) Evaluation
Figure 5: Architecture for HBYOL-7. During training (a), first train low-level policy with BYOL-y. Then, we train high-level policy with similar procedure to T, in HGCBC, but using the representation space Φ To train wh, freeze ø. labeled Φ. and use it for encoding inputs, and the output space, where the predicts the representation of the sub-goal: (b), Ty first predicts sub-goal representation, which then passed the The where both policies utilize a common state representation.
Φ(s)). During evaluation
Table 11: Additional hyperparameters for HGCBC, HGCBC-0, HBYOL-7, For other hyperparam- eters we match those in Table 4 For high-level policies иy that predict in representation space (HGCBC-6, HBYOL-?), find it is better to smaller learning rate.
Hyperparameter
Value
Hierarchical head
MLP (512, 512, 512)
Low-level head
MLP (512, 512, 512)
Sub-goal steps
25
Learning rate
3 x 10-a (HGCBC), 10-1 (HGCBC-0, HBYOL-7)
D CL TO TD-SR
Here, illustrate that connection between CL and TD-SR, showing that in the limit an n-step version of TD-SR becomes similar to CL.
We can rewrite Equation (T) to see the connection between TD-SR and CL (MC). Under assumptions that f is the dot product between á and vh, and Φi, t centered, if apply second-order Taylor expansion to the denominator of the CL loss (Touati et al. 2023) we have:
CLinfONCE 2o Earp.s'>p [(b(s) [)()))) 2E
k- -7)
( (((++))
(14)
P (s. (kls))
Next, consider n-step variant of the TD-SR loss (Blier et al. 2021 which refer as (n):
min Eapop (()))))T(())) ("(((s)))"((s)) - Σ Exppayrrr ( 6/6 s'-p
(15)
21

---

Published as a conference paper at ICLR 2026
We make the full connection CL with infinite horizon n2:
12
TD-SR(n) = Exyo ((b(s))*((s')))
Σ
scop
(x(())) (((())))
(16)
1-000
mp
i=1
27
16
Expo (v)(())) ( s'mp
(1 11-77
E
rp (814150)
soup
(1-2)) ((>))))>>(>>)) (((4)))))))))))))))))))))))))))
(17)
27
= Egop (<((s)))(ss'))) d'mp
(1 (1-7)
E
kongoamm| -77)
(s)(s)) ((s)))((s)))))))))))))))))))))))
(18)
sempleeth xpp (8t+klee)
Thus, we can see that in the infinite horizon form of TD-SR(n), it related to the form of CLi in but with the positive contrastive term weighted by factor 1-7
E FINITE MDP
E.1 BYOL
BYOL as an Ordinary Differential Equation (ODE) In finite MDPs, can characterize the BOL objective which gives intuition about what information is captured in o, v) and conditions that may be useful for stability (Tang al.) 2023] Khetarpal et al., [025]. Consider a finite MDP with transition ps, linear d-dimensional encoder + ERiSTxd and linear action-free latent-dynamies
E Rddd In a finite MDP, Equation (3) becomes:
min B V) = min EeN popP) 4.V
[[] "
T" 1
(19)
4 V
property to prevent this objective from collapsing is that Yi is updated more quickly than t. In practice, this is commonly realized as the dynamics generally smaller network than the encoder. This system can be analyzed in an ideal setup, where we first find the optimal Vr each time before taking a gradient step for 4, which leads to the ODE for representations & (Tang al. 2023
V arg min ( V), Φ (,OYOL VV
(20)
We are able to analyze this ODE with the following assumptions (Tag et al. 2023):
Assumption F.1 (Orthogonal initialization). Φ $-
Assumption E.2 (Uniform state distribution). po(s) [
Assumption E.3 (Symmetric dynamics). pr = (P*)T
Under these three assumptions, Khetarpal et al. (2025) prove that the BYOL. ODE equivalent to monotonically minimizing the surrogate objective:
min |P" oV6" |F +C
(21)
Where || || is the Frobenius matrix norm. Thus, we can understand that the BYOL objective as learning d-rank decomposition of the underlying dynamics pa, Additionally, the top d eigenvectors of pa match those of (I - - = A (Chandak et al. (2023). However, we will highlight that there are key differences when learning a low-rank decomposition between p' and M=, This is described by Touati et al. (2023), where we can consider that in real-world problem with underlying continuous- time dynamics, actions may have little effect, and p close to the identity, i.e. close to full-rank. However, M which takes powers of (P*)', has a "sharpening effect" on the difference between eigenvalues, which gives a clearer learning signal. This is intuitive on a real-world problem like robotics, even with discrete-time dynamics, where 8141 R3 8t, but we have larger differences between s, and $1+k-
E.2 BYOL-Y
In the finite MDP, we now verify theorem 4.1 where BYOL-> approximates the successor repre- sentation with matrix decomposition A TVΦT.
22

---

Published as a conference paper at ICLR 2026
We consider the same objective 19 where need to update the expectation of the sampling distribution:
min BYOL->(2, v) == min E 4.4
po(s),44~a17)
[[ o" 31 1
(22)
V.V
Assuming that this objective is optimized under the ODE (2)). We have that our objective monotoni- cally minimizes:
min V
|M -1
|| +
(23)
This directly translates as we can consider M pa as simply valid transition matrix for new, temporally abstract, version of the original MDP. We maintain the original assumptions E. E.2 and
E.3
We do not need an additional assumption for M", as assumption E.3 for symmetric pr implies
a symmetric M (1 7) E120 P
Under this setup, also have that T/ E Baxd relates to the successor feature matrix, where each row (D contains the vector (1 7b² (8:):
(1 - v/ (si) = Σ M (81, 85)(()))
(24)
(M))));
(25)
= (क p);
(26)
V)));
(27)
In other words, in the restricted finite MDP, where minimize successor features uith ~ 4r and basis features +
23
we simultaneously learning
F MIXTURE DATASETS
we describe a practical setting where we have an offline dataset generated by a set of policies }. While we previously describe that BYOL-y approximates M7 when we have MC samples directly an arbitrary n. we now describe the behavior of BYOL-y when trained jointly on MC samples from multiple { first in the finite MDP, and how this relates to approximating the SR of the unknown mixture policy M
In Section 3.2.
SR of Mixture Policy, We begin by obtaining the SR for the mixture policy 3(als)
ΣS (a(ss)))(33)11 (3) first defining 1-step transitions:
P. pd (s411 1|st 1)
Σ 3,(ss = i)p(s141 =1 st = i.
(28)
P. ΣΣ ( = = i,0) Σ 1(88)l==1)Pp
(29)
Using (i) = p(By|s = i). W,, = diag( wy (1) wy (1))), we can see the transitions of the mixture policy 3 as simply -eeeedde))ndent) weighted average of the transitions of {3,}-
ps Σ W.,p%
(30)
M = (1 7)
W, p8
(31)
120
Approximated SR of BYOL-y. Using samples from set of unknown policies {8)} the BYOL-?
objective corresponds to:
min E 0.0 sLm p0),),)-0(((()))))) 9 [ -0.TTT... "pTs+113
(32)
min E
4,~ p()) 17
T St
क.. 84 s11]
(33)
0,W
23

---

Published as a conference paper at ICLR 2026
Le. by theorem 4.1 we are approximating .M P(s(ssss))M which compare to Equation
31 via:
M = (1 r)
Σe +1 Σ (P
(34)
120
Intuitively, while M corresponds to the SR of the average policy, BYOL-Y approximates M, an average of policy SRs.
General Case. We rewrite the inequality from Equation 9 but with mixture data:
LBVOL- = E
(35)
2 E B ( (c))
[(++))
(8t)
F.1 CL ON MIXTURE DATA
We diseuss the behavior CL on mixture datasets. First, we write Equation when practically applied to mixture data as implemented in TRA:
36 we rewrite Equation
(3 (f(((s))) ((s+)) -E BUN log (((())))0))
CIRA E
(36)
(x.
We note that a mismatch occurs between the numerator (attractive), and the denominator (repulsive) terms. Namely, we attract two representations only when they are sampled from the same policy B5. but minimize similarly for states sampled under the occupancy of the mixture policy B.
We could consider other forms where both terms sample from the same distributions. Namely, in the ideal case if could get MC samples from 8+
M Ms, s_), the loss has the form:
Las Rc
(e) (s,,44
(f(((s)) E E (s)
log
(37)
In practice, we only can take MC samples from, Bj's, so we could view the loss as an expectation over these policies:
f Rs R. By ppppp ) s [1(b())) ((++)) - T. (
log
e (((C')))(()))
84 ()),)+)
1
N
Eg, E
(($()) M
(s)
E, ap
(5)
Dog
e ((
(38)
This objective similarly corresponds to capturing information related to mixture of different policies, similarly to the BYOL-r objective. We can see the positive term of LIRA matches L while the negative term matches f In other words, LTRA is under-optimistic compared to L and over-pessimistic compared to f We can see how this may discourage stitching. For example, have trajectory a b, as negative term.
Although want relation from a, c, v(())((c) is only sampled
SVD approximation of TRA. In the single policy case, Touati 2023) demonstrates that CL with single policy (Lcr,)) relates to an SVD of (a,a")
(=()")
27
M (s.
Le. 2 supp
(5s) ((()))
+
p (s'')
24

---

Published as a conference paper at ICLR 2026
We now show that the mixture policy case corresponds to SVD of 23 plg, (s)
(=1)
CTRA se E Bagp((t))) ( (f(1((e)) o(5 Baups [log up I [e([U(t))06631]
o' s'm M (oL,0')
E, M f(b(s), Eaupe log By [
= an of mpp
(5))
(39)
Under assumptions that f is the dot product between ái and s/, and Φ, sb are centered, if we apply second-order Taylor expansion to second term:
Σ, p(3))))M (8, (())
sugg a
(5('))
2 E a [()())
(40)
E 2(B (5()) o(5)) pp()
27
((())(((8)
-Bangg sampp
(41)
25

---

Published as a conference paper at ICLR 2026
G
ADDITIONAL RESULTS FOR HORIZON GENERALIZATION
Giant Maze
Large Maze
Medium Maze
THC
2H
humanoidmaze-giant
1.0
Avrrgge Suitis Nate
0.6
0.4
0.2
0.0
3
9-10
11-12
13-14
15-16
17-18
19.20
21.22
23-24
25-26
27-28
29-30
(Muza Cella)
humanoidnazz-large
antmaze-large
1.0
1.0
Beeage Sucoee Rete
8.10
11.12
13:14
15:16
17:18
34
7-8
8.10
11-12
13-14
15:46
17.18
10
Doteaze Godl Date Otl.
Dinasox to Coul (Mace Coly,
humanoidmazz-mmedium
antmaze-medium
Avrrage Sectess Hute
0.8
0.8
0.6
0.6
0.2
0.2
0.0
0.0
2
9-10
Distanoe Goal (Maze Gello
Distance Goal (Maze Calle)
3YOL. (ours)
1BA
OOSIC
Figure 6: Evaluating Generalization with Increasing Horizons: The distances to the right of the red dotted line require combinatorial generalization. The maze maps show examples of how intermediate goals are selected along the optimal path.
We include additional results matching the setup in Section (humanoidman dmaze )- (medium, large, giant) in Figure We can observe that BYOL-y leads in performance as the distance between the start and goal grows when compared to other methods.
5.3 for antmaze-med um, and
26

---

Published as a conference paper at ICLR 2026
H
REPRESENTATIONS
H.1
ADDITIONAL VISUALIZATIONS
BYOL-Y
TD-SR
BYOL
TRA
JF
LA
antmaze-large
x2
humanoi i dmaze mediuu
UP 2H
UP 26
humanoidnaaze-aagge
Figure 7: Additional Visualization of the Learned Representation: depicts the similarity between the prediction of the current state representation to the goal representation. Brighter color indicates higher similarity.
H.2 CORRELATION TO SHORTEST PATH
We conduct quantitative comparison between representations through alignment with shortest- path distance in the environment. Namely, we compute the correlation between similarity in the representation space, [[[s.:]]]
((s.)) F(g)
the shortest path distance in the between sampled start and goal cells (ry space). While the shortest path distance does not measure ground truth temporal distance, as it does not account for robot dynamics, it still provides simple reference for the general structure we expect to see in representations. We can see that in Table 12 on average BYOL-y's representations
to most strongly correlate with shortest path distance. We also compute the success rate the same checkpoints used for correlation, and notice relationships between these correlations and empirical success rate. We see that the ranking of methods in terms of average correlation in representation space matches the ordering of methods in terms of average empirical policy success.
27

---

Published as a conference paper at ICLR 2026
Setemaze-medinee
antmaze- large
TD-5R
A.75
BYOLY
BYOL-Y
BYOL
6.00
13:35
Sterios Parrouradriongh
5 Longth
Sharaes
Shattert Part
humanoi www.iiiddaeeeeeeeee
humano www.e-
TRA
BYOLYY
BYOLLY
BYOL
Sharta 7Ash Langth
Shorfaat Tuth Langth
19
Charteet Farth
Figure 8: Scatter plot of negative cosine similarity between randomly sampled (state.goal) pairs in representation spaces and true shortest path, aggregated over 4 model seeds, each sampled at 100 (saate.goal) pairs.
Dataset
BYOL-y"
BYOL
TRA
TDSR#
aetmazz--medinny
0.71 0.01
0.59 0.05
0.05
0.72 0.03
antmaze Selmmaze-aarge-gte
0.66 0.02
0.10+0.02
0.62
0.67 0.02
humamooidrazz---rree---- 0.02
0.18
0.02
0.03
0.62=0.00
0.00
0.15
0.08
average correlation
0.66
0.27
0.32
0.53
average maze success
26
31
36
Table 12: Correlation of representation space with shortest path distance. For each method, we use 10, 000 (state, goal) pairs to compute correlation, and then compute the average and standard deviation of the correlation over 4 model seeds, and the success rate over these same checkpoints.
28


Reviews:
[{"review_id": "9QtVJHa2ac", "reviewer": "Reviewer_6mmR", "summary": "This paper addresses the failure of Goal-Conditioned Behavioral Cloning (GCBC) on tasks requiring combinatorial generalization (\"stitching\"). The authors propose BYOL-γ, a self-predictive auxiliary loss. The core idea is that this loss encourages the learning of representations that approximate the Successor Representation (SR) without the instability of TD-learning or the pessimism of contrastive learning.", "questions": "None", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- simple, intuitive design.\n- well-suited for offline setting (as shown in TD-SR < BYOL-γ).\n- good empirical performance.", "weaknesses": "- The authors explicitly benchmark only against non-hierarchical methods. However, the current state-of-the-art on the OGBench 'stitch' tasks (e.g., HIQL) is hierarchical. By omitting this comparison, the paper fails to demonstrate how its performance stacks up against the actual SOTA, making its practical significance unclear. It also remains unclear whether BYOL-γ's benefits could be combined with hierarchical methods.\n\n\n- The empirical improvements on visual-based tasks (visual-antmaze, visual-scene-play) are marginal at best (Table 2). In some cases (e.g., visual-antmaze-large), the proposed BYOL-γ (26.0) performs worse than the standard GCBC baseline (29.2), questioning the method's applicability and benefits for high-dimensional visual inputs.", "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": "atnMSAL0YP", "reviewer": "Reviewer_o9Vq", "summary": "The authors focus on Goal-Conditioned Behavior Cloning and its inability to “stitch” or as they refer, perform combinatorial generalization. They attribute this inability to the lack of temporal consistency in state representations. The authors look into dynamics aware state representations and connections to successor measures to instill representations that are close for the states visited by can be sampled by successor measure of the policy collecting the data. In other words, the representations of all the states collected by a policy should be close. To ensure scalability, the authors look into BYOL style contrastive methods rather than TD based methods to learn representations that can predict “future” state representations from the policy. The authors further create a few variants of their method using a bidirectional loss. The policy is trained using BC with this BYOL-$\\gamma$ being an auxiliary loss.", "questions": "(1) The authors say (on page 5, line 238-243) that they want an invariance $\\phi(s_f) \\approx \\phi(s_w)$. This can be fine for the BC policy that was conditioned for $\\phi(s_f)$ but would feel is going towards $\\phi(s_w)$. But what about the BC policy from $\\phi(s_w)$ to $\\phi(s_f)$? Wont it assume that it's already at the goal?\n\n(2) How is a collapse avoided? It looks like all states of a trajectory would be mapped close to each other so for the same example, $\\phi(s_0) \\approx \\phi(s_w)$ as they belong to some policy and $\\phi(s_f) \\approx \\phi(s_w)$ which would also imply, $\\phi(s_f) \\approx \\phi(s_0)$. How is a collapse prevented?\n\n(3) Suppose there are four trajectories, $s_0$ to $s_w$, $s_w$ to $s_f$, $s_0$ to $s_y$ and $s_y$ to $s_f$. The trajectory $s_0 - s_y - s_f$ is optimal. Would this method be able to find out the optimal as there is no notion of reward/cost/value?\n\n(4) Are the results statistically significant?", "limitations": "", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "(1) The authors formalize the combinatorial generalization gap or the stitching error using a mixture policy. This formulation is interesting and can be used to study several algorithms and future works. \n\n(2) The authors expand on self-predictive representations from single step or n-step to a general geometric distribution over long horizons. The authors also link this representation to successor measures and show how these representations can be used to estimate successor measures (and features).", "weaknesses": "(1) Empirically there is a minor improvement over TD-SR. In visual setting, the performance is similar to GCBC. First of all, how is GCBC stitching? Wasn’t the argument that GCBC cant stitch?\n\n(2) To goal to enforce stitching in GCBC was to come up with an offline RL method that can scale (as the authors claim the TD methods dont). But in visual environments, their method does not scale well too (similar performance to GCBC)\n\n(3) $\\psi$ is a network and $\\psi^\\pi$ is successor features which is confusing, a different symbol should be used.\n\n(4) While the motivation for adding stitching ability to GCBC makes some sense, the motivation for the representation learning objective is not strong. The representation sure can be used to estimate successor measures but does that make them better for BC or does that allow stitching? Are these representations helping with stitching in a theoretical way or its an empirical hypothesis/observation?\n\n(5) Does adding a bidirectional prediction affect the theoretical properties? As the loss changed so the solution would change too.", "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": "ws6v13Zvv0", "reviewer": "Reviewer_JKTQ", "summary": "This paper investigates the question of representation learning to improve combinatorial generalization in behavior cloning: what learning objective best recovers the latent dynamics of the environment? They focus on the self-predictive Bring Your Own Latent framework, introducing a modified objective BYOL-$\\gamma$ that displays better ability to stitch together long trajectories, enabling better generalization in longer-horizon tasks.", "questions": "1. The topic I am most unclear on is that equations (7) and (8) seem that they would be prone to collapse. How is this avoided? Is it related to the use of `stop gradient`? (the below points should be interpreted as both parts of this question and notes/suggestions)\n    - The invariance $\\phi(s_f) \\approx \\phi(s_w)$ seems that, without any \"counter-balancing\" from a term that differentiates states, it would coerce $\\phi(\\cdot)$ to encode all states similarly. How is this avoided?\n    - The theoretical results in Appendix D mention stability, but it could be helpful to also discuss this in the main text.\n2. This is not a question, but I wanted to note that the referenced paper (Richens et al., 2025) reminded me of the paper \"Resolving Causal Confusion in Reinforcement Learning via Robust Exploration\" (Lyle et al., 2021) that similarly discusses the importance of learning about the environment in order to generalize. This is not to say that it should be cited here; I just thought it might be interesting to the authors.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The writing of the paper is generally good (easy to follow, good motivation of the problem, good coverage of related work)\n    - In particular, the paragraph that discusses successor representations at the beginning of section 3 was very clear\n2. The introduction to the experiments section (section 5 on page 7) is very nice and thorough\n3. The various kinds of experimental results reported (eg, Figures 2 and 3) are interesting\n    - The decision to \"bold values within 95% of [the best] value in the same row\" (table 2) is a great idea, very helpful for quickly parsing/comparing methods within each row\n    - The latent state representation visualizations are very interesting\n4. The inclusion of a limitations section at the end of the paper is appreciated\n5. The BYOL-$\\gamma$ method is motivated both intuitively and theoretically\n    - Because of this, as well as because the paper seeks to study the broader topic of representation-learning objectives, I would say that the simplicity of the method (i.e., the similarity to the prior BYOL objective, eq. 3) is not really a weakness.", "weaknesses": "1. The theoretical assumption of symmetric transition dynamics does not seem realistic.\n2. The qualitative results of the representation visualizations aren't convincing: the representation similarities for BYOL-$\\gamma$ do not look clearly better than, e.g., BYOL and TRA.\n    - In mazes, presumably (based on the concepts in this paper) the embedding similarity should decrease roughly monotonically with distance from the goal; thus, along with the provided visuals in Figures 2 and 6, it could be helpful to provide a one-dimensional scatterplot of \"similarity (e.g., dot product) vs. true shortest-path distance from goal\" for each sampled point in the space; this could facilitate easier (and less subjective) comparison of the methods.\n3. The results in Figure 5 (Appendix F) seem to show that BYOL-$\\gamma$ is not really better than the other methods (i.e., the plot for Figure 3 in the main text seems cherry-picked).\n4. It could be useful to study some more interesting environments of different types, since all the experiments are just done in various mazes.", "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": "lUlvy6geMc", "reviewer": "Reviewer_c6Ng", "summary": "This paper proposed BYOL-$\\gamma$ for learning representations in goal-conditioned reinforcement learning (GCRL) problems. The core idea of BYOL-$\\gamma$ is simple: instead of learning representations by aligning one-step latent transitions as in prior RL methods utilizing BYOL, BYOL-$\\gamma$ learns representations by aligning multi-step latent transitions following the successor measures. Empirically, BYOL-$\\gamma$ achieves competitive empirical performance across a suite of challenging tasks from the OGBench benchmarks.", "questions": "- In Sec. 3.1, what are the connections between CL and TD-SR to BYOL-$\\gamma$? Does the connection between CL and TD-SR, line 174: \"an n-step version of TD-SR is related to CL\", help us to understand BYOL-$\\gamma$?\n\n- Sec. 3.2 and the beginning of Sec. 4 reviewed prior definition of generalization gap in GCBC. What are the connections between these definitions and the BYOL-$\\gamma$ objective in Sec. 4.1?\n\n- The motivation of developing the objective in Eq. 8 is to enable combinatorial generalization. But why does this objective enable combinatorial generalization? What is the policy $\\pi$ in this equation? Is it the behavioral policy or the learned policy?\n\n- Eq. 10 includes a bidirectional prediction term. Any explanation for the reason for including this bidirectional prediction term? Results in Table 3 seems to suggest the effect of this bidirectional prediction term is minor.\n\n- Sec 4.2 introduces another auxiliary loss function called TD-SR. What’s the relationship between TD-SR and BYOL-$\\gamma$? Does the paper propose TD-SR or BYOL-$\\gamma$ as the main representation learning objective?\n\n- The second and third columns in Table 1 are confusing. From line 205, it looks like $\\tilde{M}^{\\beta}(s, s_{+})$ and $\\sum_j p(\\beta_j \\mid s) \\tilde{M}^{\\beta_j}(s, s_{+})$ are equivalent to each other. What’s the difference between these two columns in the context of different methods?\n\n- Line 335: More loss terms with the same batch of dataset typically results in higher learning efficiency. Is this sentence saying that BYOL-$\\gamma$ is less efficient than CL and TD-SR?\n\n- Sec 5.1 visualizes the representation learned by different methods. From Fig 2, it looks like the cosine similarity of BYOL-$\\gamma$ and TD-SR are more concentrated than BYOL and TRA. How can we relate these representation heatmaps to the reachability of different states?", "limitations": "", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "- Learning representations for control is an important topic in RL. The paper motivates learning temporal representation for the GCBC problem by using the successor measure. Instead of bootstrapping the representations of the next state as in BYOL, BYOL-$\\gamma$ proposes to bootstrap the representations of a future state sampled from the successor measure, which encodes temporal information in the dataset. The learned representations are then used to drive the learning of the goal-conditioned policy via behavioral cloning.\n\n- The author conducted experiments on the challenging OGBench benchmarks, showing that BYOL-$\\gamma$ achieves competitive empirical performance, compared to prior representation learning methods for GCBC and prior goal-conditioned reinforcement learning methods in the offline setting. Additional experiments ablate the key components of BYOL-$\\gamma$ and demonstrate that BYOL-$\\gamma$ can generalize over longer horizons than prior methods.", "weaknesses": "-  The paper reviewed different representation learning objectives for RL in Sec. 3.1 and reviewed the generalization gap from prior work in Sec. 3.2 and the beginning of Sec. 4 (line 223 - line 249). However, Sec. 4.1 proposed the BYOL-$\\gamma$ objective directly without explaining its connection with the preliminaries. It is difficult to understand the roles of Sec. 3.1, Sec. 3.2, and the beginning of Sec. 4 in the paper.\n\n- One of the motivations of developing BYOL-$\\gamma$ is to learn representations that enable combinatorial generalization in GCRL. The objective function in Eq. 8, the following Eq. 9 and Theorem 4.1 do not provide enough explanations for why the objective function can enable combinatorial generalization. \n\n- Sec 4.2 introduces another auxiliary loss function called TD-SR for learning representations. The relationships between TD-SR and BYOL-$\\gamma$ are not clear from the current draft, making it unclear whether the paper proposed BYOL-$\\gamma$ or TD-SR as its main representation learning objective.\n\n- Results in Table 3 seem to suggest some components of BYOL-$\\gamma$ has minor effects.\n\nOverall, the motivations, the algorithms, and the theory are somewhat inconsistent.", "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 (Poster)