Contrastive Difference Predictive Coding
Authors: Chongyi Zheng
Ruslan Salakhutdinov
Benjamin Eysenbach
Venue: ICLR 2024 poster (2024)
Keywords: contrastive learning
reinforcement learning
goal-reaching
goal-conditioned RL
temporal difference
Primary area: reinforcement learning
TL;DR: a temporal difference version of contrastive predictive coding

Abstract: Predicting and reasoning about the future lie at the heart of many time-series questions. For example, goal-conditioned reinforcement learning can be viewed as learning representations to predict which states are likely to be visited in the future. While prior methods have used contrastive predictive coding to model time series data, learning representations that encode long-term dependencies usually requires large amounts of data. In this paper, we introduce a temporal difference version of contrastive predictive coding that stitches together pieces of different time series data to decrease the amount of data required to learn predictions of future events. We apply this representation learning method to derive an off-policy algorithm for goal-conditioned RL. Experiments demonstrate that, compared with prior RL methods, ours achieves $2 \times$ median improvement in success rates and can better cope with stochastic environments. In tabular settings, we show that our method is about $20\times$ more sample efficient than the successor representation and $1500 \times$ more sample efficient than the standard (Monte Carlo) version of contrastive predictive coding.

Published as a conference paper at ICLR 2024
CONTRASTIVE DIFFERENCE PREDICTIVE CODING
Chongyi Zheng
Ruslan Salakhutdinoy Carnegie Mellon University
Benjamin Eysenbach Princeton University
Carnegie Mellon University chongyi 28 andrew. cmu. edu
ABSTRACT
Predicting and reasoning about the future lie at the heart of many time-series ques- tions. For example, goal-conditioned reinforcement learning can be viewed learning representations to predict which states are likely to be visited in the fu- ture. While prior methods have used contrastive predictive coding to model time series data, learning representations that encode long-term dependencies usually requires large amounts of data. In this paper, we introduce a temporal difference version of contrastive predictive coding that stitches together pieces of different time series data to decrease the amount of data required to learn predictions of fu- events. We apply this representation learning method to derive off-policy algorithm for goal-conditioned RL. Experiments demonstrate that, compared with prior RL methods, ours achieves 2x median improvement in success rates and can better cope with stochastic environments. In tabular settings, we show that our method is about 20 more sample efficient than the successor representation and 1500 more sample efficient than the standard (Monte Carlo) version of con- trastive predictive coding.
Code: https:://hoogyii whwnriiihhhhggg...... ootoo//// Besiddinne infonce Website:
infonce
INTRODUCTION
Learning representations is important for mod- cling nigh-dimensional time series data. Many applications of time-series modeling require representations that not only contain informa- tion about the contents of particular observa- tion, but also about how one observation relates to others that co-occur in time. Acquiring rep- resentations that encode temporal information challenging, especially when attempting to capture long-term temporal dynamics: the fre- quency of long-term events may decrease with the time scale, meaning that learning longer- horizon dependencies requires larger quantities of data.
random states
trajectory
current state
S
S'
state
86' of -11..))
7)
Figure 1: InfoNCE nonparametrie version of the successor representation. (Top) The distances tween learned representations indicate the prohability of transitioning to the next state and set of randomly- sampled states. (Bottom) We update these representa- tions so they assign high likelihood to (a) the next state and (b) states likely to be visited after the next state. See Sec. for details.
In this paper, we study contrastive represen- tation learning on time series data positive examples co-occur nearby in time, so the dis- tances between learned representations should encode the likelihood of transiting from one representation to another. Building on prior work that uses the InfoNCE (Sohn) 2016 Oord et al.) 2018) loss to learn representations of time- series data effectively, we will aim to build a temporal difference version of this loss. Doing so allows us to optimize this objective with fewer samples, enables us to stitch together pieces of dif- ferent time series data, and enables us to perform counterfactual reasoning we should be able to estimate which representations would have learned, if had collected data in different way.

---

Published as a conference paper at ICLR 2024
After a careful derivation, our resulting method can be interpreted non-parametric form of the successor representation (Dayan 1993), as shown in Fig.
The main contribution of this paper is a temporal difference estimator for InfoNCE. We then apply this estimator to develop a new algorithm for coal-conditioned RL. Experiments on both state-based and image-based benchmarks show that algorithm outperforms prior methods, especially on the most challenging tasks. Additional experiments demonstrate that our method can handle stochastic- ity in the environment more effectively than prior methods. We also demonstrate that our algorithm can be effectively applied in the offline setting. Additional tabular experiments demonstrate that TD InfoNCE is up to x more sample efficient than the standard Monte Carlo version of the loss and that it can effectively stitch together pieces of data.
RELATED WORK
This paper will study the problem of self-supervised RL, building upon prior methods on goal- condition RL., contrastive representation learning, and methods for predicting future state visitations. Our analysis will draw connection between these prior methods, a connection which will ultimately result in new algorithm for goal-conditioned RL. We discuss connections with unsupervised skill learning and mutual information in Appendix B
Goal-conditioned reinforcement learning. Prior work has proposed many frameworks for learn-
ing goal-conditioned policies, including conditional supervised learning (Ding et al. 2019 Ghosh
et al.
2020 Gupta et al. 2020 Emmons et al.] 202T; [Lncch et al. [2200] Oh et al. 2018 Sun et al.
Sun
2019
actor-critic methods Andrychowicz et al. (2017); Nachum et al. (2018): ; Chane-Sam et al.
(2021
semi-parametric planning (Pertsch et al. 2020: Fang et al. 2022 2023 Eysenbach et al.
2019
Nair & Finn 2019 Gupta et al. 2020), and distance metric learning (Wang et al. [0233] Tian
et al. 2020a Nair et al. 2020b] Durug kar et al. (002). These methods have demonstrated impres
sive results on a range of tasks, including real-world robotic tasks (Ma al. 2022 Shah al. 2022
Zheng et al., 2023). While some methods require manually-specified reward functions or distance functions, our work builds upon a self-supervised interpretation of goal-conditioned RL that casts this problem as predicting which states are likely to be visited in the future (Eysenbach et al. 2020
2022 Blier et al. 202T).
Contrastive representation learning. Contrastive learning methods have become key tool for learning representations in computer vision and NLP (Chopra et al.] [2055] Schroff et al. 2015
Sohn, 2016 Oh Song et al. 2016)] Wang Isola 2020 Oord et al. 2018 [ Tschannen et al. 2019
Weinberger & Saul 2009 He et al. 2022, Radford et al. 2021: Chen et al. 2020: Tian et al.
2020b Gao et al. 2021). These methods assign similar representations to positive examples and dissimilar representations to negative examples or outdated embeddings (Grill et al. 2020). The two main contrastive losses are based on binary classification ("NCE") ranking loss ("InfoNCE") (Ma &
Collins 2018). Modern contrastive learning methods typically employ the ranking- based objective
to learn representations of images (Chen et al. 2020 Tian et al.] 20206] Henaff, 2020; Wu et al,
2018 text (Logeswaran & Leel 2018 et al.] 2021 Radford et al. [022 and sequential data (Nair
ct al. 2022 Sermanet et al. 2018). Prior works have also provided theoretical analysis for these methods from the perspective of mutual information maximization (Linsker 1988 Poole et al.
2019), noise contrastive estimation (Gutmann & Hyvärinen, 2010 [a & Collins [018 Tsai et al.
2020 Arora et al.] 2019). and the geometry the learned representations (Wang & Isola 2020).
In the realm of RL,, prior works have demonstrated that contrastive methods can provide effective reward functions and auxiliary learning objectives (Laskin et al., [020a |Hansen-Estruch et al., 2022 Choi et al. 202 Nair et al., 2020a 2018), and can also be used to formulate the goal-reaching problem in an entirely self-supervised manner (Ma et al. 2022 Durugkar et al., 202 Eysenbach et al. 2020 2022). Our method will extend these results by building temporal difference version of
the "ranking" based contrastive loss; this loss will enable us to use data from one policy to estimate which states different policy will visit.
Temporal difference learning and successor representation. Another line of work studies us- ing temporal difference learning to predict states visited in the future, building upon successor rep- resentations and successor features (Dayan 1993. Barreto al. 2017 2019 Blier et al.) 2021
2

---

Published as a conference paper at ICLR 2024
While learning successor representation using temporal difference bears similarity the typical Q-Learning algorithm (Watkins & Dayan 1992] Fu et al. 2019; Mnih et al,, 2013) in the tabular setting, directly estimating this quantity difficult with continuous states and actions (Janner al. 2020 Barreto et al. 2017; Touati & Ollivier 2021 Blier et al. (021). J. To lift this limitation, we will follow prior work (Eysenbach al. 2022 2020 Touati & Ollivier 202T) in predicting the successor representation indirectly: rather than learning a representation w whose coordinates correspond to vis- itation probabilities, we will learn state representations such that their inner product corresponds to visitation probability. Unlike prior methods, we will show how the common InfoNCE objective can be estimated temporal difference fashion, opening the door to off ooflcy reasoning and enabling our method to reuse historical data to improve data efficiency.
METHOD
We start by introducing notation and prior approaches to the contrastive representation learning and the goal-conditioned RL problems. We then propose new self-supervised actor-critic algorithm that we will use in our analysis.
3.1 PRELIMINARIES
We first review prior work in contrastive representation learning and goal-conditioned RL. Our method will use ideas from both.
Contrastive representation via InfoNCE. Contrastive representation learning aims to learn representation space, pushing representations of positive examples together and pushing represen- tations of negative examples away. InfoNCE (also known contrastive predictive coding) (Sohn) 2016 Jozefowicz et al. [016)] Oord et al., 2018 Henaff 2020) is a widely used contrastive loss, which builds upon noise contrastive estimation (NCE) (Gutmann & Hyvärinen 2010 Ma Collins) 2018). Given the distribution of data (x( py(v) over data x E X, E y and the conditional dis- tribution of positive pairs (y)x(() over X' x V, we sample px(x), (1) ~ pyiæ(1 æ)) and
(2 ,", ..(,()))) py(y)- The InfoNCE loss is defined
LInONCE (f) A smpx(x), "0 log Ei=1 (f(yyt)))
N
(
where f x + IR is parametric function. Following prior work (Eysenbach et al. 2022 Wang Touati & Ollivier 2021). choose parameterize f(:.)) via the inner product ((x) (yy), where Φ0-) and (b(.) map data to & normalized
& Isola 2020
of representations of data [xx vectors of dimension d. We will call f the critic function and and 1 the contrastive representations. The Bayes-optimal critic for the InfoNCE loss satisfies (Poole et al. 2019 & Collins) 2018 Oord
al. 2018
ply x)) p(y)c(x))
exp (f* (2.
where e(-) is an arbitrary function. We can estimate this arbitrary function using the optimal critic f by sampling multiple negative pairs from the data distribution:
plu x)) dy ploic(z)
c(x) / p(y x)dd c(x)
Ep(n /exp (f* y))l
plot
Reinforcement learning and goal-conditioned RL. We will consider a Markov decision process defined by states E S, actions E A. rewards : S x A x S + IR Using A(-) denotes the probability simplex, we define an initial state distribution Po
+ A(S), discount factor y E (0, 1], and dynamics p S x A + A(S). Given policy pF (s++ 8, a) to denote the probability density of reaching state $t after exactly steps, starting at state and action and then following the policy aaa s). We can then define the discounted state
S + A(A), we will use
3

---

Published as a conference paper at ICLR 2024
occupancy measure (Ho & Ermon 2016 Zhang al. 2020 Eysenbach al. 2020 et al. )023) starting from state s and action a as
2022
Zheng
p" (sLt s. a (1
(s+ s,).)
(3)
Prior work (Dayan 1993) have shown that this discounted state occupancy measure follows recur- sive relationship between the density at the current time step and the future time steps:
p= s, a) (1 ())(s' = se+ a, + TlE. (pppppplsss) [p" (sc+ ss,a')]
(4)
d'>xxa''1
For goal-conditioned RL, we define goals g E S in the same space as states and consider a goal- conditioned policy (a 9,g) and the corresponding goal-conditioned discounted state occupancy measure p" (st+ 8, a. g). For evaluation, we will sample goals from a distribution Pg S + A(S). Following prior work (Eysenbach et al.] 20200: Rudner et al., (02T). we define the objective of the goal-reaching policy maximizing the probability of reaching desired goals under its discounted state occupancy measure while commanding the same goals:
max Epp(s) ((ls)) [v" (5++ = g 8, 8,a,9)] =(1++)
(5)
In tabular settings, this objective is the same as maximizing expected returns using a sparse reward function r(s, a, 8', g) = (1 7)8(s' g) (Eysenbach et al. 2022). Below, we review two strate- gies for estimating the discounted state occupancy measure. Our proposed method (Sec. 3.2) will
combine the strengths of these methods while lifting their respective limitations.
Contrastive RL and C-Learning. Our focus will be on using contrastive representation learning to build
conditiodd RL algorithm, following template set in prior work (Eysenbach) jet al. 2022: 2020 These contrastive RI. methods are closely related to the successor representa- tion (Dayan 1993): they aim learn representations whose inner products correspond the like-
lihoods of reaching future states. Like the successor representation, representations from these con- trastive RL methods can then be used represent the Q function for any reward function (Mazoure et al. 2022). Prior work (Eysenbach et al. 2022) has shown how both NCE and the InfoNCE losses can used to derive Monte Carlo algorithms for estimating the discounted state occupancy mea- sure. We review the Monte Carlo InfoNCE loss below. Given a policy T(a | s). consider learning contrastive representations for state and action pair x = (s, a) and potential future state y = 8t4- We define the data distribution to be the joint distribution of state-action pairs pxxxx) = p(s, a) and the marginal distribution of future states py (y) = p(8++), represening either the distribution of re- play buffer (online) or the distribution of dataset (offline). The conditional distribution of positive
pairs is set to the discounted state occupancy measure for policy म. PyJæ (3) x) = p= (s | s, resulting in Monte Carlo (MC) estimator
((aa log
LMC InfoNCE (f) E.
(6)
(s.a)))(saa))) ~p (m.0) (2:M) -p(s(-)
1== e (2)
and an optimal critic function satisfying
a, (8,) p" (sL+ 8.a)
p(5)))c(s,))
This loss estimates the discounted state occupancy measure in a Monte Carlo manner. Computing this estimator usually requires sampling future states from the discounted state occupancy measure of the policy π. i.e., on-policy data. While, in theory. Monte Carlo estimator can be used in an off-policy manner by applying importance weights to correct actions, this estimator usually suffers from high variance and is potentially sample inefficient than temporal difference methods (Precup et al. 2000 2001).
In the same way that temporal difference (TD) algorithms tend to be more sample efficient than Monte Carlo algorithms for reward maximization (Sutton & Barto 2018), we expect that TD con- trastive methods are more sample efficient at estimating probability ratios than their Monte Carlo counterparts. Given that the InfoNCE tends to outperform the NCE objective in other machine learning disciplines, we conjecture that our TD InfoNCE objective will outperform the TD NCE objective (Eysenbach et al. 2020) (see experiments in Appendix.
4

---

Published as a conference paper at ICLR 2024
3.2 TEMPORAL DIFFERENCE INFONCE
In this section, we derive a new loss for estimating the discounted state occupancy measure for fixed policy. This loss will be temporal difference variant of the InfoNCE loss. We will temporal difference InfoNCE (TD InfoNCE) to refer to our loss function.
In the oof-po setting, we aim to estimate the discounted state occupancy measure of the policy given a dataset of transitions D {(s, a, } collected by another behavioral policy 3(a s). This setting is challenging because we do not obtain samples from the discounted state occupancy measure of the target policy æ. Addressing this challenge involves two steps: (i) expanding the MC estimator (Eq. ( via the recursive relationship of the discounted state occupancy measure (Eq. + and (ii) estimating the expectation over the discounted state occupancy measure via importance sampling. We first the identity from Eq. express the MC InfoNCE loss the sum of next-state term and a stture-state term:
log e Ms, (2)
E (s.a)~((sa) (2:N) xp(xet)
(1->) E (--p(())s.))
Σ 1= (
ftl
Li(f)
ef(rass !1) D (eaa
+ E aaaa
log
C2(s)
While this estimate is similar to TD target for Q-Learning Watkins & Dayan 1992, Fu et al., 2019), the second term requires sampling from the discounted state occupancy measure of policy To avoid this sampling, we next replace the expectation over p" (s(+ s',a)) in (2( by an importance weight,
( s' a') p(8
ef (1)
C2(l) = E.' "p('' ( a'a.)
log
Li=l ((
x(=+)
we could estimate the importance weight, then we could easily estimate this term by sampling from p(st.). ). We will estimate this importance weight by rearranging the expression for the optimal
critic (Eq. 7 and substituting our estimate for the normalizing constant c(s, a) (Eq, 2):
n" (5 a,a)
ef'(s,;o
c(s
exp exp (ft(s, = Eppree) [ef (1,
(8)
p(s
We will use w(s, stt to denote our estimate of this, using in place of f* and using a finite- sample estimate of the expectation in the denominator:
(I:N)
(
wis, a, (1:N)
N EN (9,),199))
This weight accounts for the effect of the discounted state occupancy measure of the target policy. Additionally, corresponds to the categorical classifier that InfoNCE, produces (without constant N). Taken together, we can now substitute the importance weight in L2(f) with our estimate in Eq. yielding a temporal difference (TD) InfoNCE estimator
LTD InfaNCE (f) 2 (s.0)~((s.a) (2: mp(ee=)
(1 (1- (1-7)BB 7) o2 (o(
log
Ei=1 ed ((
el(sianet))
+y tyBe'n(''lsad d'va(a'le") ott xpp(4))
[w(s',a', o4, (1:N)) (s log
(10)
ΣN @ (ssano)
5

---

Published as a conference paper at ICLR 2024
Algorithm 1 'emmporl Difference InfoNCE. We CE to denote the cross entropy loss, taken across the rows of a matrix of logits and labels. We use F as a matrix of logits, where Fli,j] = (6) (6) ( T (x(2)) See Appendix D.1 for details.
(( at (g())
1: Input contrastive representations Φ9 and </g. target representations oy and Vo. and goal- conditioned policy Tu.
2: for each iteration do
3: Sample ((5)" (0) .a.)) s41, ,
(1) gld)
replay buffer / dataset, al) T(a s g)))
4: Compute Frext, P' Phauree P gog using @u and vho.
5:
Compute Fw using og and v/g-
6: W <- N stop.grad (SOFTMAX(Fu))
7:
C(0) + (1 7)EE
FRem. labels IN) + CE Litture labels W)
8:
[(w) + (logits = F gest labels IN)
9:
Update f.w by taking gradients of [(w).
10:
Update using an exponential moving average.
11:
Return @o. Vo. and ww..
where L'Jsg indicates the gradient of the importance weight should not affect the gradient of the entire objective. As shown in Fig. we can interpret the first term as pulling together the repre- sentations of the current state-action pair ((s, a)) and the next state (ss'); the second term pulls the representations at the current step o(s, a) similar to the (weighted) predictions from the future state (s++)). Importantly, the TD InfoNCE estimator is equivalent to the MC InfoNCE estimator for the optimal critic function: LTD InfoNCE (f*) = LMC (J
Convergence and connections. Appendix A we prove that optimizing a variant of the TD In- foNCE objective is equivalent perform one step policy evaluation with a new Bellman operator; thus, repeatedly optimizing this objective yields the correct discounted state occupancy measure. This analysis considers the tabular setting and assumes that the denominators of the softmax func- tions and w in Eq. are TD InfoNCE and C-learning (Eysenbach al. 2020) (a temporal difference estimator of the NCE objective) in Appendix F..3 Appendix discussses how TD InfoNCE corresponds to nonparamet- ric variant of the successor representation.
are computed using an exact expectation. We discuss the differences between
3.3 GOAL-CONDITIONED POLICY LEARNING
The TD InfoNCE method provides a way for estimating the discounted state occupancy measure. This section shows how this estimator can be used to derive algorithm for goal-conditioned RL. This algorithm will alternate between (/) estimating the occupancy measure using the TD In- foNCE objective and (2) optimizing the policy maximize the likelihood of the desired goal under the estimated occupancy measure, Pseudo-code is shown in Algorithm [ and additional details are in Appendix D.I and code is available online
While our TD InfoNCE loss in Sec. 3.2 estimates the discounted state occupancy measure for policy Taa s). we can extend to the goal-conditioned setting by replacing aTaa s) with x(a s,g) and f(s, $14) with f(, a. 9. $14). resulting in goal-conditioned TD InfoNCE estimator. This goal-conditioned TD InfoNCE objective estimates the discounted state occupancy measure of any future state for goal-conditioned policy commanding any goal. Recalling that the discounted state occupancy measure corresponds the Q function (Eysenbach al., 2022), the policy objective to select actions that maximize the likelihood of the commanded goal:
ef ((000.00044))))))
(p)(s)po)()) Tog 8, 8, = E 7( s.9)
g-pp(a)))-po(o)) (al_ral)).)).) (1:N) rpp(ss))
log
DN ('
(9)
(11)
In practice, we optimize both the critic function and the policy for one gradient step iteratively, using estimated ff in place of f*
https::/githhb.... www..ooww.oooo.oooooooooooooco.c httriiidhhhnn_____e_e fonce

---

Published as a conference paper at ICLR 2024
1.0
TD IntoNCE
1.0
TD IntoNCE
Quesimetric IL
Qruelmmnnic RL.
success rate
GEBE
DDFG + HER
0.5
0.5
0.0
0.0
push
push
pick place (image) (image)
pash
(state)
(state)
(state)
(image)
(a) Fetch robotics benchmark from (Plappert al.
2018
(b) Stochastic tasks.
Figure 2: Evaluation online GCRL benchmarks. (Left) TD InfoNCE performs similarly or outper-
forms all baselines on both state-based and image-based tasks. (Right) On stochastic versions of the state-based tasks, TD InfoNCE outperforms the strongest baseline (QRL). Appendix Fig. 6 shows the learning curves.
4 EXPERIMENTS
Our experiments start with comparing goal-conditioned TD InfoNCE to prior goal-conditioned RL approaches on both online and offline goal-conditioned RI (GCRL) benchmarks. We then analyze the properties of the critic function and the policy learned by this method. Visualizing the represen- tations learned by TD InfoNCE reveals that linear interpolation corresponds to a form of planning. Appendix E.3 ablates the difference between TD InfoNCE and a prior temporal difference method based on NCE. All experiments show means and standard deviations over five random seeds.
4.1 COMPARING TO PRIOR GOAL-CONDITIONEDD RL METHODS
We compare TD InfoNCE to four baselines containing four manipulation tasks for the Fetch robot. The observations and goals of those tasks be either state of the robot and objects 64 x 64 RGB image. We will evaluate using both versions. The first baseline, Quasimetric Reinforcement Learning (QRL) (Wang et al. 2023).
online GCRL benchmark (Plappert et al. 2018)
state-of-the-ard approach that uses quasimetric models learn the optimal goal-conditioned value functions and the corresponding policies. The second baseline is contrastive RL (Eysenbach al., 2022]. which estimates the discounted state occupancy measure using EMCLLIONCE (Eq. (6).
Our third baseline is variant of contrastive RI. (Eysenbach et al., 2022) using binary NCE loss.
We call this method contrastive RL (NCE). The fourth baseline is the goal-conditioned behavioral
cloning (GCBC) (Ding et al., 2019: |Emmons et al., 2021; Ghosh et al.
[220 Lynch et al.) 2020
Sun et al. 2019 Srivastava et al. 2019). We also include a comparison with an off-the-shelf
actor-
critic al Igorithm augmented with hindsight relabeling (Andrychowicz et al.
2017 Levy et al.
2018
Riedmiller et al. 2018 Schaul et al. 2015) to learn a goal-conddtioned policy (DDPG HER).
We report results in Fig. 2a and defer the full learning curves to Appendix Fig. 6 These results show that TD InfoNCE matches or outperforms other baselines on all tasks, both for state and image observations. On those more challenging tasks (pick & pl place (state / image) and s ide (state image) ), TD InfoNCE achieves a 2x median improvement relative to the strongest baseline (Appendix Fig. (). On the most challenging tasks, image-based pick place and s slide, TD InfoNCE is the only method achieving non-negligible success rates. For those tasks where the success rate fails to separate different methods significantly (e.g., slide (state) and push (image)), ), we include comparisons using minimum distances of the gripper or the object to the goal over an episode Appendix Fig. / We speculate this observation because TD InfoNCE estimates the discounted state occupancy measure more accurately, a hypothesis we will investigate in Sec.
Among those baselines, QRL is the strongest one. Unlike TD InfoNCE, the derivation of QRL assumes the dynamics are deterministic. This difference motivates us to study whether TD InfoNCE continues achieving high success rates in environments with stochastic noise. To study this, we compare TD InfoNCE to QRL on a variant of the Fetch benchmark where observations are corrupted with probability 0.1. As shown in Fig. 2b TD InfoNCE maintains high success rates while the performance of QRL decreases significantly. suggesting that TD InfoNCE can better cope with stochasticity in the environment.
7

---

Published as a conference paper at ICLR 2024
Table Evaluation on offline D4RL AntMaze benchmarks.
TD InfoNCE
QRL
Contrastive RI.
GCBC
DT
IQL
TD3 + BC
umaze-v2
84.9 + 1.2
76.8 + 2.3
79.8
1.6
65.4
65.6
87.5
78.6
umaze-diverse-v2
91.7 + 1.3
80.1 = 1.3
77.6 +
2.8
2.8
60.9
51.2
62.2
71.4
mediuum-play-v2
86.8 = 1.7
76.5 16.5211
72.6
58.1
1.0
71.2
10.6
medium-diivese--22
82.0 1: 3.4
1.94
71.5
1.3
67.3
0.6
70.0
3.0
large-play-v2
17.0 2.5
+: 2.8
48.6
4.4
32.4
39.6
0.2
large-diverse-v2
55.6 + 3.6
51.5 3.8
54.1
+: 5.5
36.9
0.2
17.5
0.0
average absolute error
Monte Carlo method
10-4
10-1
temporal difference methods
10-1
10
10a
10a
gradient steps
dataset size
TD InfoNCE
Successor Representation
C-Learning
Contrastive RL
Estimating the discounted state occupancy measure in tabular setting. (Left) Temporal difference methods have lower errors as fast as the best baseline (successor representation). (Right) TD InfoNCE is more data efficient than other methods. Using dataset of size 10M, TD InfoNCE achieves error rate 25% lower than the best baseline; TD InfoNCE also matches the performance of C-learning with x less data.
Figure 3:
Monte Carlo method. Also note our TD InfoNCE converges
4.2 EVALUATION ON OFFLINE GOAL REACHING
We next study whether the good performance of TD InfoNCE transfers to the setting without any interaction with the environment (i.e., offline RL). We evaluate on AntMaze tasks from the D4RL. benchmark (Fu et al. 2020). The results in Table [ show that TD InfoNCE outperforms most baselines on most tasks. See Appendix D.3 for details.
4.3 ACCURACY OF THE ESTIMATED DISCOUNTED STATE OCCUPANCY MEASURE
This section tests the hypothesis that our TD InfoNCE loss will be more accurate and sample ficient than alternative Monte Carlo methods (namely, contrastive RL (Eysenbach et al.) (022)) in predicting discounted state occupancy measure. We will use the tabular setting so that we can ground truth estimate. We compare TD InfoNCE to three baselines. Successor representa- tions (Dayan 1993) can also learned in a TD manner, though can be challenging to apply beyond tabular settings. C-learning similar to TD InfoNCE in that uses temporal difference method to optimize a contrastive loss, but differs in using binary cross entropy loss instead a softmax cross entropy loss. Contrastive RL is the MC counterpart of TD InfoNCE. We design x 5 gridworld with 125 states and 5 actions (up, down, left, right, and no-op) and collect 100K transitions using UNIF(AA). We evaluate each method by measuring the absolute
uniform random policy, μla error between the predicted probability p and the ground truth probability p", averaging over all pairs of (s, a, ):
1
Σ
1p(s ss,) - -ppppp"aaa (s s.a)).
S S|
8,0,8
For the three TD methods, compute the TD target in SARSA manner (Sutton & Barto 2018). For those methods estimating probability ratio, we convert the prediction to a probability by mult- plying by the empirical state marginal. Results in Fig. show that TD methods achieve lower emror than the Monte Carlo method, while TD InfoNCE converges faster than C- C-Learnngg Appendix E.2 discusses why all methods plateau above zero
8

---

Published as a conference paper at ICLR 2024
Dataset
TD InfoNCE
Contrastive RI.
Dataset
TD InfoNCE
Contrastive RL
Figure 4: Stitching trajectories in a dataset. The behavioral policy collects "7." style trajectories. Un- like the Monte Carlo method (contrastive RI.) our TD InfoNCE successfully "stitches" these trajecto- ries together, navigating between pairs (start *, goal *) states unseen in the training trajectories. Ap- pendix Fig. shows additional examples.
Figure Searching for shortcuts skewed datasets. (Leftt) Conditioned on different initial states * and goals *, collect datasets with 95% long paths (dark) and 5% short paths (light). (Center) TD InfoNCE infers shortest path, (Right) while con- trastive RL fails to find this path. Appendix Fig. shows additional examples.
Our next experiments studies sample efficiency. We hypothesize that the softmax in the TD InfoNCE loss may provide more learning signal than alternative methods, allowing it to achieve lower error on fixed budget of data. To test this hypothesis,
experiments with dataset sizes from IK to 10M on the same gridworld, comparing TD InfoNCE to the same set of baselines. We report results in Fig. 3 with errors showing one standard deviation after training for 50K gradient steps for each approach. These results suggest that methods based on temporal difference learning predict more accurately than Monte Carlo method when provided with the same amount of data. Compared with its Monte Carlo counterpart, TD InfoNCE is 1500 x more sample efficient (6.5 x 10x vs 10" transi- tions). Compared with the only other TD method applicable in continuous settings (C-learning), TD InfoNCE can achieve a comparable loss with 130x less data (7.7 x 10* vs 10 transitions). Even compared with the strongest baseline (successor representations), which makes assumptions (tabular MDPs) that our method avoids, TD InfoNCE can achieve comparable error rate with almost 20 fewer samples (5.2 x 10a vs 107 transitions).
4.4 DOES TD INFONCE ENABLE OFF-POLICY REASONING?
The explicit temporal difference update (Eq. 10 in TD InfoNCE is similar to the standard Bellman backup, motivating us to study whether the resulting goal-conditioned policy is capable of perform- ing dynamic programming with offline data. To answer these questions, conduct experiments on the same gridworld environment as in Sec. 4.3 comparing TD InfoNCE to contrastive RL (i.e., Monte Carlo InfoNCE). Fig. 4 shows that TD InfoNCE successfully stitches together pieces of dif- ferent trajectories to find a route between unseen (state, goal) pairs. Fig. S shows that TD InfoNCE can perform off- off-policy reasoning, finding path that shorter than the average path demonstrated in the dataset. See Appendix D.4 for details.
5 CONCLUSION
This paper introduced a temporal difference estimator for the InfoNCE loss. Our goal-conditioned RL algorithm based this estimator outperforms prior methods in both online and offline settings, and is capable of handling stochasticity in the environment dynamics. While we focused on spe- eific type RL problem (goal-conditioned RL), in principle the TD InfoNCE estimator can be used to drive policy evaluation for arbitrary reward functions. One area for future work is to determine how compares to prior off- off-policy evaluation techniques.
While we focused on evaluating the TD InfoNCE estimator on control tasks, it is worth noting that the MC InfoNCE objective has been previously applied to NLP, audio, video settings; one intriguing and important question is whether the benefits of TD learning seen on these control tasks translate into better representations in these other domains.
Limitations. One limitation of TD InfoNCE is complexity: compared with its Monte Carlo coun- terpart, ours is more complex and requires more hyperparameters. It is worth noting that even TD InfoNCE struggles to solve the most challenging control tasks with image observations. On the the- oretical front, our convergence proof uses a slightly modified version of our loss (replacing a sum with expectation), which would be good to resolve in future work.

---

Published as a conference paper at ICLR 2024
Acknowledgements We thank Ravi Tej and Wenzhe Li for discussions about this work, and anonymous reviewers for providing feedback on early versions of this work. We thank Tongzhou Wang for providing performance of baselines in online GCRL experiments and thank Raj Ghugare for sharing code of environment inplementation. We thank Vivek Myers for finding issues in the code. BE is supported by the Fannie and John Hertz Foundation.
REFERENCES
Alekh Agarwal, Nan Jiang. Sham M Kakade, and Wen Sun. Reinforcement learning: Theory and algorithms. CS Dept., UW Seattle, Seattle, WA, USA, Tech. Rep, 32, 2019.
Marcin Andrychowicz, Filip Wolski, Alex Ray, Jonas Schneider, Rachel Fong, Peter Welinder, Bob McGrew, Josh Tobin, OpenAl Pieter Abheel, and Wojciech Zaremba. Hindsight experience re- play. Advances in neural information processing systems, 30, 2017.
Sanjeev Arora, Hrishikesh Khandeparkar, Mikhail Khodak, Orestis Plevrakis, and Nikunj Saunshi. A theoretical analysis of contrastive unsupervised representation learning. In 36th International Conference on Machine Learning, ICML 2019, pp. 9904-9923, International Machine Learning Society (IMLS), 2019.
André Barreto, Will Dabney, Rémi Munos, Jonathan Hunt, Tom Schaul, Hado van Hasselt, and David Silver. Successor features for transfer in reinforcement learning. Advances in neural information processing systems, 30, 2017.
André Barreto, Diana Borsa, Shaobo Hou, Gheorghe Comanici, Eser Aygün, Philippe Hamel, Daniel Toyama, Shibl Mourad, David Silver, Doina Precup, et al. The option keyboard: Combining skills in reinforcement learning. Advances in Neural Information Processing Systems, 32, 2019.
Dimiuri P Bertsekas and John N Tsitsiklis. Neuro-dynamic programming: overview. In Pro- ceedings of 1995 34th IEEE conference on decision and control, volume pp. 560-564. IEEE, 1995.
Léonard Blier, Corentin Tallec, and Yann Ollivier. Learning successor states and goal-dependent values: A mathematical viewpoint. arXiv preprint arXiv:2101.07/11 2021.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, et al. Jax: composable transformations of python+ numpy programs. 2018.
Víctor Campos, Alexander Trott, Caiming Xiong, Richard Socher, Xavier Giró-i Nieto, and Jordi Torres. Explore, discover and learn: Unsupervised discovery of state-covering skills. In Interna- tional Conference on Machine Learning, pp. 1317-1327. PMLR, 2020.
Elliot Chane-Sane SSnee Cordelia Schmid, and Ivan Laptev. Goal-conditioned reinforcement learning with imagined subgoals. In International Conference on Machine Learning, pp. 1430-1440. PMLR, 2021.
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems, 34:150844155999 2021.
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In International conference machine learning, pp. 1597-1607. PMLR, 2020.
Jongwook Choi, Archit Sharma, Honglak Lee, Sergey Levine, and Shixiang Shane Gu. Variational empowerment as representation learning for goal-conditioned reinforcement learning. In Interna- tional Conference on Machine Learning, pp. 1953-1963, PMLR, 2021.
Sumit Chopra, Raia Hadsell, and Yann LeCun. Learning similarity metric discriminatively, with application to face verification. 2005 IEEE computer society conference on computer vision and pattern recognition (CVPR'05), volume pp. 539-546. IEFF, 2005.
10

---

Published as a conference paper at ICLR 2024
Peter Dayan. Improving generalization for temporal difference learning: The successor representa- tion. Neural computation, 5(4):613-624) 1993.
Yiming Ding, Carlos Florensa, Pieter Abbeel, and Mariano Phielipp. Goal-conditioned imitation learning. Advances neural information processing systems, 32, 2019.
Dubi and YS Horowitz. The interpretation of conditional monte carlo as a form of importance sampling. SIAM Journal on Applied Mathematics, 36(1):115-122) 1979.
Ishan Durugkar, Mauricio Tec, Scott Nickum, and Peter Stone. Adversarial intrinsic motivation for reinforcement learning. Advances in Neural Information Processing Systems, 34:8622-8636, 2021.
Scott Emmons, Benjamin Eysenbach, Ilya Kostrikov, and Sergey Levine. Rvs: What is essential for offline rl via supervised learning? In International Conference on Learning Representations, 2021.
Damien Ernst, Pierre Geurts, and Louis Wehenkel. Tree-based batch mode reinforcement learning. Journal of Machine Learning Research, 6, 2005.
Benjamin Eysenbach, Abhishek Gupta, Julian Ibarz, and Sergey Levine. Diversity all you need: Learning skills without a reward function. arXiv preprint arXiv: arXiv: 1802.06070, 2018.
Benjamin Eysenbach, Ruslan Salakhutdinov, and Sergey Levine. Search on the replay buffer: Bridg- ing planning and reinforcement learning. Advances Neural Information Processing Systems, 32, 2019.
Benjamin Eysenbach, Ruslan Salakhutdinov, and Sergey Levine. C-learning: Learning achieve goals via recursive classification. arXiv preprint arXiv:201 0et:20011.08999 2020.
Benjamin Eysenbach, Tianjun Zhang, Sergey Levine, and Russ R Salakhutdinov. Contrastive learn- ing as goal-conditioned reinforcement learning. Advances in Neural Information Processing Sys- tems, 1-.35603-356200 2022.
Kuan Fang, Patrick Yin, Ashvin Nair, and Sergey Levine. Planning to practice: Efficient online fine-tuning by composing goals in latent space. In 2022 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS), pp. 4076 4083. IEEE, 2022.
Kuan Fang, Patrick Yin, Ashvin Nair, Homer Rich Walke, Gengchen Yan, and Sergey Levine. Gen- cralization with lossy affordances: Leveraging broad offline data for learning visuomotor tasks. In Conference on Robot Learning. pp. 106-117. PMLR, 2023.
Justin Fu, Aviral Kumar, Matthew Soh, and Sergey Levine. Diagnosing bottlenecks in deep q- learning algorithms. In International Conference Machine Learning, pp. 2021-2030. PMLR, 2019.
Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning. arXiv preprini arXiv:2004.07219, 2020.
Scott Fujimoto and Shixiang Shane Gu. minimalist approach offline reinforcement learning. Advances in neural information processing systems, 1-20132-221443 2021.
Tianyu Gao, Xingcheng Yao, and Danqi Chen. Simese: Simple contrastive learning of sentence em- beddings. In 202/ Conference on Empirical Methods in Natural Language Processing, EMNLP 2021, pp.
Association for Computational Linguistics (ACL), 2021.
Samuel Gershman. The successor representation: its computational logic and neural substrates. Journal of Neuroscience, 38(33);7193-7200, 2018.
Dibya Ghosh, Abhishek Gupta, Ashwin Reddy, Justin Fu, Coline Manon Devin, Benjamin Eysen- bach, and Sergey Levine. Learning to reach goals via iterated supervised learning. International Conference on Learning Representations, 2020.
Michael B Giles. Multilevel monte carlo methods. Acta numerica, 24:259-328. 2015.
11

---

Published as a conference paper at ICLR 2024
Karol Gregor, Danilo Jimenez Rezende, and Daan Wierstra. Variational intrinsic control. arXiv preprint arXiv:1611.07000 2016.
Jean- Bant-Batieen Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan Guo, Mohammad Gheshlaghi Azar, et al. Bootstrap your own latent-a new approach to self-supervised learning. Advances in neural information processing systems, 33:21271-21284, 2020.
Abhishek Gupta, Vikash Kumar, Corey Lynch, Sergey Levine, and Karol Hausman. Relay policy learning: Solving long-horizon tasks via imitation and reinforcement learning. In Conference on Robot Learning. pp. 1025-1022 1037. PMLR, 2020.
Michael Gutmann and Aapo Hyvärinen. Noise-contrastive estimation: A new estimation principle for unnormalized statistical models. Proceedings of the thirteenth international conference on artificial intelligence and statistics, pp. 297-304. JMLR Workshop and Conference Proceedings, 2010.
JM Hammersley. Conditional monte carlo. Journal of the ACM (JACM), 3(2):33-66, 1956.
Philippe Haansnn-Estruch Amy Zhang, Ashvin Nair, Patrick Yin, and Sergey Levine. Bisimula- tion makes analogies goal-conditioned reinforcement learning. In International Conference on Machine Learning, pp. 8407-8426. PMLR, 2022.
Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. Masked au- toencoders are scalable vision learners. In Proceedings of the IEEE/CVF conference computer vision and pattern recognition, pp. 16000-16009, 2022.
Olivier Henaff. Data-efficient image recognition with contrastive predictive coding, In International conference on machine learning, pp. 1182-4192, PMLR, 2020.
Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning. Advances in neural information processing systems, 29, 2016.
Michael Janner, Igor Mordatch, and Sergey Levine. numma-modelss Generative temporal difference learning for infinite-horizon prediction. Advances in Neural Information Processing Systems, 33: 7724-1735. 2020.
Chao Jia, Yinfei Yang, Ye Xia, Yi-Ting Chen, Zarana Parekh, Hieu Pham, Quoe Le, Yun-Hsuan Sung, Zhen Li, and Tom Duerig. Scaling up visual and vision-language representation learning with noisy text supervision. In International conference on machine learning, pp. 4904-4916. PMLR, 2021.
Rafal Jozefowicz, Oriol Vinyals, Mike Schuster, Noam Shazeer, and Yonghui Wu. Exploring the limits of language modeling. arXiv preprint arXiv:1602.02410. 2016.
Ilya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit q- learning. In International Conference Learning Representations, 2021.
Aviral Kumar, Justin Fu, Matthew Soh, George Tucker, and Sergey Levine. Stabilizing off- tif-policy q- learning via bootstrapping error reduction. Advances in Neural Information Processing Systems, 32, 2019.
Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning offline reinforcement learning. Advances in Neural Information Processing Systems, 33:1179-111 179-11911 2020.
Michael Laskin, Aravind Srinivas, and Pieter Abbeel. Curl: Contrastive unsupervised representa- tions for reinforcement learning. In International Conference on Machine Learning, pp. 5639- 5650. PMLR, 2020a.
Misha Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, Pieter Abbeel, and Aravind Srinivas. Rein- forcement learning with augmented data. Advances in neural information processing systems, 33: 19884-19895. 2020b.
12

---

Published as a conference paper at ICLR 2024
Andrew Levy, George Konidaris, Robert Platt, and Kate Saenko. Learning multi-level hierarchies with hindsight. In International Conference on Learning Representations, 2018.
Ralph Linsker. Self-organization in perceptual network. Computer, 21(3):005-177, 1988.
Lajanugen Logeswaran and Honglak Lee. An efficient framework for learning sentence representa- tions. arXiv preprint arXiv: 02893, 2018
Corey Lynch, Mohi Khansari, Ted Xiao, Vikash Kumar, Jonathan Tompson, Sergey Levine, and Pierre Sermanet. Learning latent plans from play. In Conference on robot learning, pp. 1113- 1132. PMLR, 2020.
Yecheng Jason Ma, Shagun Sodhani, Dinesh Jayaraman, Osbert Bastani, Vikash Kumar, and Amy Zhang. Vip: Towards universal visual reward and representation via value-implicit pre-training. In The Eleventh International Conference on Learning Representations, 2022.
Zhuang Ma and Michael Collins. Noise contrastive estimation and negative sampling for conditional models: Consistency and statistical efficiency. arXiv preprint arXiv:1809.0188/ 2018.
Bogdan Mazoure, Benjamin Eysenbach, Ofir Nachum, and Jonathan Tompson. Contrastive value learning: Implicit models for simple offline rl. arXiv preprint arXiv:2211. 1tt.:22211.0220000 2022.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Mare Belle- mare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. nature, 51 (77554)::5225533 2015.
Ofir Nachum, Shixiang Shane Gu, Honglak Lee, and Sergey Levine. Data-efficient hicrarchical reinforcement learning. Advances in neural information processing systems, 31, 2018.
Ashvin Nair, Shikhar Bahl, Alexander Khazatsky, Vitchyr Pong. Glen Berseth, and Sergey Levine. Contextual imagined goals for self-supervised robotic learning, In Conference on Robot Learning, pp. 530-539. PMLR, 2020a.
Ashvin V Nair, Vitchyr Pong, Murtaza Dalal, Shikhar Bahl, Steven Lin, and Sergey Levine. Visual reinforcement learning with imagined goals. Advances in neural information processing systems, 31, 2018
Suraj Nair and Chelsea Finn. Hierarchical foresight: Self-supervised learning of long-horizon tasks via visual subgoal generation. In International Conference Learning Representations, 2019.
Suraj Nair, Silvio Savarese, and Chelsea Finn. Goal-aware prediction: Learning model what matters. In International Conference Machine Learning, pp. 7207-7219. 2020b.
Suraj Nair, Aravind Rajeswaran, Vikash Kumar, Chelsea Finn, and Abhinav Gupta. R3m: A univer- sal visual representation for robot manipulation. arXiv preprint arXiv:2203.12601, 2022.
Junhyuk Oh, Yijie Guo, Satinder Singh, and Honglak Lee. Self-imitation learning. In International
Conference on Machine Learning, pp. 3878-3887 PMLR, 2018
Hyun Oh Song, Yu Xiang, Stefanie Jegelka, and Silvio Savarese. Deep metric learning via lifted structured feature embedding. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 4004-4012, 2016.
Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predic- tive coding. arXiv preprint arXiv: 2018
Karl Pertsch, Oleh Rybkin, Frederik Ebert, Shenghao Zhou, Dinesh Jayaraman, Chelsea Finn, and Sergey Levine. Long-horizon visual planning with goal-conditionss hierarchical predictors. Ad- vances Neural Information Processing Systems, 33:1 10217322-77333 2020.
Matthias Plappert, Marcin Andrychowicz, Alex Ray, Bob McGrew, Bowen Baker, Glenn Pow- ell, Jonas Schneider, Josh Tobin, Maciek Chociej, Peter Welinder, al. Multi-goal reinforce- ment learning: Challenging robotics environments and request for research. arXiv preprint arXiv: 18020094444 2018.
13

---

Published as a conference paper at ICLR 2024
Ben Poole, Sherjil Ozair, Aaron Van Den Oord, Alex Alemi, and George Tucker. On variational bounds mutual information. In International Conference on Machine Learning, pp. 5171- 5180. PMLR. 2019.
Doina Precup, Richard S Sutton, and Satinder Singh. Eligibility traces for oof-policy policy eval- uation. In Proceedings of the Seventeenth International Conference on Machine Learning, pp. 759-766, 2000.
Doina Precup, Richard S Sutton, and Sanjoy Dasgupta. Off-policy temporal difference learning with function approximation. In Proceedings of the Eighteenth International Conference Machine Learning, pp. 417-424, 2001.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In International conference on machine learning, pp- 8748 8763. PMLR, 2021.
Tom Rainforth, Rob Cornish, Hongseok Yang, Andrew Warrington, and Frank Wood. On nest- ing monte carlo estimators. In International Conference Machine Learning, pp. 4267-4276. PMLR, 2018.
Martin Riedmiller, Roland Hafner, Thomas Lampe, Michael Neunert, Jonas Degrave, Tom Wiele, Vlad Mnih, Nicolas Heess, and Jost Tobias Springenberg, Learning by playing solving sparse reward tasks from scratch. International conference on machine learning, pp. 4344 4353.
PMLR, 2018
Tim GJ Rudner, Vitchyr Pong, Rowan McAllister, Yarin Gal, and Sergey Levine. Outcome-driven reinforcement learning via variational inference. Advances in Neural Information Processing Systems, 34:13045-1300 2021.
Tom Schaul, Daniel Horgan, Karol Gregor, and David Silver. Universal value function approxima- tors. In International conference on machine learning, pp. 1312-1320 PMLR, 2015.
Florian Schroff, Dmitry Kalenichenko, and James Philbin. Facenet: A unified embedding for face recognition and clustering. In Proceedings the IEEE conference computer vision and pattern recognition, pp. 815-883, 2015.
2015.
Pierre Sermanet, Corey Lynch, Yevgen Chebotar, Jasmine Hsu, Eric Jang, Stefan Schaal, Sergey Levine, and Google Brain. Time-contrastive networks: Self-supervised learning from video. In 2018 IEEE international conference on robotics and automation (ICRA), pp. 1134-1141. IEEE, 2018.
Dhruv Shah, Benjamin Eysenbach, Nicholas Rhinchart, and Sergey Levine. Rapid exploration for spen-world navigation with latent goal models. In Conference on Robot Learning, pp. 674-684. PMLR, 2022.
Kihyuk Sohn. Improved deep metric learning with multi-class n-pair loss objective. Advances in neural information processing systems, 29, 2016.
Rupesh Kumar Srivastava, Pranav Shyam, Filipe Mutz, Wojeiech Jaskowski, and Jürgen Schmidhu- ber. Training agents using upside-down reinforcement learning. arXiv preprint arXiv: urXiv:191..0287. 2019.
Hao Sun, Zhizhong Li, Xiaotong Liu, Bolei Zhou, and Dahua Lin. Policy continuation with hind- sight inverse dynamics. Advances in Neural Information Processing Systems, 32, 2019.
Richard Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
Stephen Tian, Suraj Nair, Frederik Ebert, Sudeep Dasari, Benjamin Eysenbach, Chelsea Finn, and Sergey Levine. Model-based visual planning with self-supervised functional distances. In Inter- national Conference Learning Representations, 2020a.
14

---

Published as a conference paper at ICLR 2024
Yonglong Tian, Dilip Krishnan, and Phillip Isola. Contrastive multiview coding. In Computer 2020: 16th European Conference, Glasgow, UK, August 23-28, 2020, Proceedings, Part XI 16. pp- 776-794. Springer, 2020b.
Ahmed Touati and Yann Ollivier. Learning one representation optimize all rewards. Advances in Neural Information Processing Systems, 34:13-23, 2021.
Yao-Hung Hubert Tsai, Han Zhao, Makoto Yamada, Louis-Philippe Morency, and Russ R Salakhut- dinov. Neural methods for point-wise dependency estimation. Advances in Neural Information Processing Systems, 33:62-72, 2020.
Michael Tschannen, Josip Djolonga, Paul K Rubenstein, Sylvain Gelly, and Mario Lucie. On mutual information maximization for representation learning. arXiv preprint arXiv: 13625, 2019.
Tongzhou Wang and Phillip Isola. Understanding contrastive representation learning through align- ment and uniformity on the hypersphere. In International Conference on Machine Learning, pp- 9929-9939. PMLR, 2020.
Tongzhou Wang, Antonio Torralba, Phillip Isola, and Amy Zhang. Optimal goal-reaching rein- forcement learning via quasimetric learning. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Proceedings of the 40th In- ternational Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp.
PMLR, 23-29 Jul 2023. URL press/v202/wane wwtp:/22wwwaagg2aaa.m html
ings mmlr.
David Warde-Farley, Tom Van de Wiele, Tejas Kulkarni, Catalin Ionescu, Steven Hansen, and Volodymyr Mnih. Unsupervised control through non-parametric discriminative rewards. arXiv preprint
2018.
Christopher JCH Watkins and Peter Dayan. Q-learning. Machine learning, 8:279-292, 1992.
Kilian Q Weinberger and Lawrence K Saul. Distance metric learning for large margin nearest neigh- bor classification. Journal of machine learning research, 10(2), 2009.
Zhirong Wu, Yuanjun Xiong, Stella X Yu, and Dahua Lin. Unsupervised feature learning via non- parametric instance discrimination. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 3733-3742, 2018.
Shangtong Zhang, Bo Liu, and Shimon Whiteson. Gradientdice: Rethinking generalized offline estimation stationary values. International Conference on Machine Learning, pp. 11194- 11203. PMLR, 2020.
Chongyi Zheng, Benjamin Eysenbach, Homer Walke, Patrick Yin, Kuan Fang, Ruslan Salakhutdi- nov, and Sergey Levine. Stabilizing contrastive rl: Techniques for offline goal reaching. arXiv preprint arXiv:2306.0m 03346, 2023.
15

---

Published as a conference paper at ICLR 2024
A THEORETICAL ANALYSIS
Our convergence proof will focus on the tabular setting with known p(s' 8, a) and p(*++). ). and follows the fitted Q-iteration strategy (Fu al.] 2019 Ernst al.. 2005| Bertsekas & Tsitsiklis, 1995]: at each iteration, an optimization problem will be solved exactly to yield the next estimate of the discounted state occupancy measure. One key step in the proof is to employ a preserved invariant; we will define the classifier derived from the TD InfoNCE objective (Eq. (TO] and show that this classifier always represents valid probability distribution (over future states). We then construct variant of the TD InfoNCE objective using this classifier and prove that optimizing this objective exactly equivalent to perform policy evaluation, resulting in the convergence to the discounted state occupancy measure.
Definition of the classifier.
We start by defining the classifier derived from the TD InfoNCE as a(s,a, p(s(+) ((sa,,0))) ((laa)
p(s(+)e f(sa.ee)
(12)
Edi, Es ((+++ )e (ora.e44)
Ep(stt)
suggesting that C(s, .) is a valid distribution over future states: a(s,a,)) E A(S).
variant of TD InfONCE. Our definition of the classifier (Eq, 12 allows us to rewrite the impor- tance weight w(s, sr+) and softmax functions in LTD InfoNCE (Eq. 10 as Monte Carlo estimates of (1:N)
the classifier using samples of $1
((s (1)
C(s,a, p($1+)
w(s a, $44
(1:N)
En
Thus, we construct variant of the TD InfoNCE objective using C::
LTD InfoNCE(C) A Epis.a)
[[1 B) (s,a)
[log C(s, a, (t+))
[C(s', stt))ss log C(s, a, 844))
ttl (1)0))))) p(==)
p(sc+)
This objective is similar to LID InfeNCEs but differs in that (a) softmax functions are replaced by C's,a, ) up constant N.p(..j] and (b) w(s', a', sf, (1-N)) is replaced by C( act) Formally,
pla
LTD InfeNCE (C) a nested Monte Carlo estimator of ÊTD InfeNCe (Rainforth et al. 2018 Giles 2015)
and we leave the analysis of the gap between them as future works. We now find the solution of LTD InfoNCE (C) analytically by rewriting it using the cross entropy and ignore the stop gradient operator to reduce clutter: LTD InloNce(CC)
sc (8.10) [log C(s, Se+ 1B [sa) [log C(s,a, a, St+ Ep(s.a) (1 - 1))p) -
c(d'd'
-Eplss,) 1(1 -)C (p(s' = (s,a),
+xCE (Ep) (C(s',cl a', -)), C(s,a,)))
= -Epp(+.a) [CE ()177)p(s' = 88 + "B
(13)
where the cross entropy for p.q E ( is defined as CE(p()).g(
(og ((x)] - Σ p(x) log ((x),
BEX
with the minimizer q² arg min CE(p(-), p. Note that p(s' = a) A(S) and Eppo'ls.e) (a'(e)))) (C(s', al, )) ASS) in Eq. 13 indicate that their convex combination is also distribution in ASS). This objective suggests a update for the classifier given any (s, a, st4):
C(s, a, 824) + (1 - r)p(s' = stt s. a) +
s()),
(14)
which bears resemblance to the standard Bellman equation.
16

---

Published as a conference paper at ICLR 2024
InfoNCE Bellman operator. We define the InfoNCE Bellman operator for any function Q a. : AxSxAXS-2X with policy T(a s) as
TINONCE Q( sc-) - p)p(s' sit s. + (( d' se+ }],
(15)
and write the update of the classifier as C(s, a, s(+) <- TINONNE Cs, 8t+)). Like the standard Bellman operator, this InfoNCE Bellman operator a 7-contraction. Unlike the standard Bellman operator, TInONCE replaces the reward function with the discounted probability of the future state being the next state (1 7)p(s' = sc+ s, a) and applies to a function depending on a state-action pair and future state (s, a, 81+)-
Proof of convergence. Using the same proof of convergence for policy evaluation with the stan- dard Bellman equation (Sutton & Barto 2018 Agarwal et al. 2019), we conclude that repeatedly applying IICONCE C' results convergence
unique O regardless of initialization,
C' (s,, = = (1-7)p(s' Sc- s, a) + (app)))s)))))))))) ( (s' st))].
Since C a s2+) and pm (8++ 8, a) satisfy the same identity (Eq. 4 have C* a, 8(+) p" (st+ 8, a), i.e., the classifier of the TD InfoNCE estimator converges to the discounted state occupancy measure. To recover f from C*, we note that f satisfies
(8,00888)) log (* 0, 86+) logp(st+) + log R lexp(/ (4,)4) =log (st+ 8, 3, log + log Ep (s,,
by definition. Since the (expected) softmax function is invariant to translation, we can write f (s, a, 8t4.) logp"(sr++ s, p(g+())+)) log c(s, a), where c(s, a) is an arbitrary func- tion that does not depend on 5t+ Thus, we conclude that TD InfoNCE objective converges to the same solution as that of MC InfoNCE (Eq. 7 i.e. LID inIoNCe(f") LMC InfoNCE(f*)
It is worth noting that the same proof applies to the goal-conditioned TD InfoNCE objective. After finding an exact estimate of the discounted state occupancy measure of a goal-conditioned policy 11
(a 8, g). maximizing the policy objective 11] is equivalent to doing policy improvement. We can apply the same proof as the Lemma of (Eysenbach al. 2020 to conclude that n (a 8, g) converges the optimal goal-conditioned policy πT" [a] 8,
B CONNECTION WITH MUTUAL INFORMATION AND SKILL LEARNING.
The theoretical analysis in Appendix A has shown that the TD InfoNCE: estimator has the same effect as the MC InfoNCE estimator. As the (MC) InfoNCE objective corresponds to a lower bound on mutual information (Poole et al. 2019), we can interpret our goal-conditioned RL method as having both the actor and the critic jointly optimize lower bound on mutual information. This perspective highlights the close connection between unsupervised skill learning algorithms (Eysenbach et al.) 2018 Campos et al., 2020; Warde-Farley et al.) 2018 Gregor et al.) 20166), and goal-conditioned RL, connection previously noted in Choi et al. (2021). Seen as unsupervised skill learning algorithm, goal-conditioned RL lifts one of the primary limitations of prior methods: can be unclear which skill will produce which behavior. In contrast, goal-conditioned RL methods learn skills that are defined as optimizing the likelihood of reaching particular goal states.
C CONNECTION WITH SUCCESSOR REPRESENTATIONS
In settings with tabular states, the successor representation (Daynn 1993 canonical method for estimating the discounted state occupancy measure (Fq. 3). The successor representation has strong ties to cognitive science (Gershman) 2018) and has been used to accelerate modern RI. meth-
ods (Barreto et al. 2017 Touati & Ollivier 2021
"Technically, f should be a set functions satisfying
C* (5,0,844) p(ae=)
17

---

Published as a conference paper at ICLR 2024
Successor representation M S x A + A(S) is long-horizon, policy dependent model that estimates the discounted state occupancy measure for every s E S via the recursive relationship (Eq. 4). Given policy a(a s), the successor representation satisfies
(16) where s p(s' 8,a) and a' vaa' s¹). Comparing this update the TD InfoNCE update shown in Fig. T and Eq. [4] we see that this successor representation update is a special case of TD InfoNCE where (a) every state is used instead of andomll-sampile the states, and (b) the probabilities are encoded directed in a matrix M, rather than encoding the probabilities as the inner product between two learned vectors.
(=,s,)) + (1 HOTEEOTT (8') + 7M",,'')
This connection is useful because it highlights how and why the learned representations can be used solve
reinforcement learning tasks. In the same way that the successor representation can be used to express the value function of a reward (M" (s, a) (r()).. the representations learned by TD InfoNCE can be used to recover value functions:
((
(*(,,))= = r(), +
Y
N E ellsaioit) r( (2)
r( aet)
1-7 *1l
See Mazoure et al. (0022) for details on this construction.
D EXPERIMENTAL DETAILS
D.1 THE COMPLETE ALGORITHM FOR GOAI-CONDITIONED RL
The complete algorithm of TD InfoNCE alters between estimating the discounted state occupancy measure of the current goal-conditioned policy via contrastive learning (Eq. [Tp and updating the policy using the actor loss (Eq. TT), while collecting more data. Given batch of N transitions of
{(st di (2) ( s ( ) sampled from p(st, ae: (t:g), 9). p(s $t. 0e). and p(s++), we can first compute the critic function for different combinations of goal-conditioned state-action pairs and future states by computing their contrastive representations Φ(8t (r,g), (/(++), and (h(8+)). ), and then construct two critic matrices Foext, Fiuture E ARNXA N using those representations:
Frealli, = ((s!") a..)) (1))
( (6((2) [i ((s) (T) T))
Note that the inner product parameterization of the critic function f(sr, ae d,,g, 814) ((2,, as .)) T (vb(())) helps compute these matrices efficiently. Using these critic matrices, we rewrite the TD InfoNCE estimate
of two cross entropy losses. The first cross entropy loss involves
(1:N) the correct next state for the corresponding goal- predicting which of the N next states 8/41 conditioned state and action pair:
(1 - C-E))C = F Fnexi, labels = IN),
IN [i]i.]] log So SOFTMAX(-) denotes row-wise softmax normalization, and In is N dimensional identity ma- trix. For the second cross entropy term, we first sample batch of N actions from the target policy at the next time step, (I:N) ~ T((++1 S141, and then estimate the importance weight matrix W E N N that serves labels as
where CE = Frest, labels
= IN) = -
Fu J = ( d.41, g.f1,g(()) ( (ssl)) W = N SOFTMAX(F.))
Thus, the second cross entropy loss takes as inputs the critic Frucure and the importance weight W:
Clogitss = Frunure labels = W).
(17)
Regarding the policy objective (Eq. can also rewritten as the cross entropy between critic matrix Fg with [g = o(s) (t) (1)) (( (ggا)))
where a( T[a st ()),,, and the
identity matrix Iy:
CE Fgn labels IN)
In practice, neural networks with parameters 0 {0.. 0 to parameterize (normalized) contrastive representations ó and & and use a neural network with parameters w to parameterize the goal-conditioned policy æ and optimize them using gradient descent.
18

---

Published as a conference paper at ICLR 2024
reach
push
pick & place
slide
1.0
1.0
1.0
success rate
0.5
0.5
0.5
0.5
0.0
0.0
0.0
environment
environment stops
environment steps
environment stops
TD InfoNCE
Quasimetric AL.
Contrastive RL
Contractive RL. (NCE)
GCBC
DDPG HER
(a) State-based tasks
reach
push
pick & place
slide
success rate
0.8
environment steps
environment steps
environment steps
environment steps
TD InfoNCE
Quasimetric RL
Contrastive RL
Contractive RL (NCE)
GCBC
DDPG HER
Image- based tasks
Figure 6: Evaluation on online GCRL benchmarks. TD InfoNCE matches or outperforms all baselines both state-based and image-based tasks.
0.5
TD InfoNCE Quasimetric RL
minimum distance
Contrastive RL
0.4
0.3
0.2
0.1
0.0
reach (state)
push (state)
pick & place (state)
slide (state)
reach (image)
push (image)
pick & place (image)
slide (image)
Figure 7: We also compare different methods using the minimum distance of the gripper or the object the goal over episode. Note that lower minimum distance indicates better performance. TD InfoNCE achieves competitive minimum distances on online GCRL benchmarks.
D.2 ONLINE GOAI-CONDITIONED RL EXPERIMENTS
We report complete success rates for online GCRL experiments in Fig. showing the mean success rate and standard deviation (shaded regions) across five random seeds. TD InfoNCE outperforms or achieves similar performance on all tasks, compared with other baselines. For those tasks where the success rate fails to separate different methods significantly (e.g., slide (state) and push (image)) ). we include comparisons using minimum distances of the gripper the object to the goal over an episode in Fig. 7 selecting the strongest baselines QRL and contrastive RL. Note that lower minimum distance indicates a better performance. These results suggest that TD InfoNCE able to emerge a goal-conditioned policy by estimating the discounted state occupancy measure, serving competitive goal-conditioned RI. algorithm.
19

---

Published as a conference paper at ICLR 2024
TD InfoNCE
Contrastive RL
Dataset
Figure 8: Stitching trajectories in dataset. We show additional (start, goal) pairs for experi- ment in Fig.
D.3 OFFLINE GOAL-CONDITIONED RL EXPERIMENTS
Similar to prior works (Eysenbach et al. 2022 Wang et al. 2023), we adopt an additional goal-conditioned behavioral cloning regularization prevent the policy from sampling out-of-
distribution actions (Fujimoto Gu 2021 (5q
Kumar et al. 2020 2019) during policy optimization
e mg)
arg max E (5. -pp..omm.ps +(-b--) awvr(als.g)).
(1 (1- A) log
A ||a aorig 1
ΣN (eagg.))
(1:N) (1:N) mp((++)
where À is the coefficient for regularization. Note that squared error instead of the maximum likelihood estimate of Gorig under policy π used in prior works. We compare TD InfoNCE to the state-of-the-art QRL (Wang et al. (023) and its Monte Carlo counterpart (contrastive RL (Eysenbach et al., (022 We also compare to the pure goal- conditioned behavioral cloning implemented (Emmons al. 2021) as well as recent baseline that predicts optimal actions via sequence modeling using a transformer (DT (Chen et al.) (20TT) Our two baselines are offline actor-critic methods trained via TD learning: TD3 + BC (Fujimoto & Gu 2021) and IQL. (Kostrikov et al. (021), not involving goal-conditioned relabeling. We use the result for baselines except QRL from [Eseebbach et al., 2022).
supervised loss based on the mean
As shown in Table [ TD InfoNCE matches outperforms all baselines 5 tasks. On tasks (medium sredium-playe paaippllyyy-a v2 v2 and medi um diverse-vee 22), TD InfoNCE achieves a +33% improvement the Monte Carlo
contrastive RL, showing the advantage of temporal difference learning approach with a fixed dataset. We conjecture that this benefit comes from the dynamic programming property of the TD method and will investigate this property further in later experiments (Sec. 4.4 Additionally, TD InfoNCE performs 1.4x better than GCBC and retains a 3.8x higher scores than DT on average, where these baselines (autoregressive) supervised losses instead of TD learning. These results suggest that TD InfoNCE is also a competitive goal-co coal-conditioned RL algorithm in the offline setting.
D.4 OFF-POLICY REASONING EXPERIMENTS
Stitching trajectories. The first set of experiments investigate whether TD InfoNCE successfully stitches pieces trajectories in a dataset to find complete paths between (start, goal) pairs unseen together in the dataset. We collect a dataset with size 20K consisting of "Z" style trajectories moving diagonal and off-diagonal directions (Fig.8) 8). while evaluating the learned policy on reaching goals on the same edge as starting states after training both methods for 50K gradient steps. Figure shows that TD InfoNCE succeeds in stitching parts of trajectory the dataset, moving along "C style paths towards goals, while contrastive RL fails to do These results justify hypothesis that TD InfoNCE performs dynamic programming and contrastive RL instead naively follows the behavior defined by the data.
20

---

Published as a conference paper at ICLR 2024
Dataset
TD InfoNCE
Contrastive RL
Figure 9: Searching for shortcuts in skewed datasets. We show additional (start, goal) pairs for the experiment in Fig.
21

---

Published as a conference paper at ICLR 2024
Table Hyperparameters for InfoNCE.
Hyperparameters
Values
actor learning rate
5 10
critic learning rate
10
using 2 normalized representations
yes
hidden layers sizes (for both actor and representations)
(512, 512, 512, 512)
contrastive representation dimensions
16
Table Changes to hyperparameters for offline RL experiments. (Table
Hyperparameters
Values
batch size (on large tasks)
256 1024
hidden layers sizes (for both actor and rep- resentations on -arge- tasks)
(512, 512, 512, 512) (1024, 1024, 1024, 1024)
behavioral cloning regularizer coefficient À
0.1
goals for actor loss
random states - future states
Searching for shorteuts. Our second of experiments aim compare the performance of TD InfoNCE against contrastive RL. on searching shorteuts in skewed datasets. To study this, we col- lect different datasets of size 20K with trajectories conditioned on the same pair of initial state and goal, with 95% of the time along a long path and 5% of the time along a short path. Using these skewed datasets, we again train both methods for 50K gradient steps and then evaluate the policy performance on reaching the same goal starting from the same state. We show the goal-conditioned policies learned by the two approaches in Fig. The observation that TD InfoNCE learns to take shortcuts even though those data are rarely seen, while contrastive RI. follows the long paths domi- nating the entire dataset, demonstrates the advantage of temporal difference learning over its Monte Carlo counterpart in improving data efficiency.
D.5 IMPLEMENTATION AND HYPERPARAMETERS
We implement TD InfoNCE, contrastive RI., and C-I C-Learning using JAX (Bradbury et al. 2018)
building upon the official codebase of contrastive RI Fl For baselines QRI,, GCBC, and DDPG HER, we use implementation provided by the author of QRLi We summarize hyperparameters for TD InfoNCE in Table 2 Whenever possible, we used the same hyperparameters
contrastive RL (Eysenbach et al. 2022) Since TD InfoNCE computes the loss with 22 negative examples, we increase the capacity of the goal-conditioned stateacction encoder and the future state encoder layers MLP with 512 units in each layer applying ReLU activations. For fair comparisons, we also increased the neural network capacity in baselines to the same number and used fixed batch size 256 for all methods. Appendix [.. includes ablations studying the effect of differet hyperparamters in Table 2
For offline RL experiments, we make some changes to hyperparameters (Table
E ADDITIONAL EXPERIMENTS
E. HYPERPARAMETER ABLATIONS
We conduct ablation experiments to study the effect of different hyperparameters in Table aiming to find the best hyperparameters for TD InfoNCE. For each hyperparameter, we selected a set of
22

---

Published as a conference paper at ICLR 2024
100
10
1.0
10
success ret
success
0.5
05
03
05
0.0
0.0
0.0
0.0
Grvnrmment St8ps
305
anviromment stops
105
cnvtronment Stpp
100
environment st0pS
305
5x 10 (Dun)
1x
10
10 (Duns)
(a) Actor learning rate
(b) Critic learning rate
slide
1.0
10
success rele
success
0.5
05
0.5
0.0
0.0
0.2
0.5
3.6
1.5
100
0.2
0.2
0.5
0.6
0.0
0.2
orvtrmmmnt staps
300
anviromment staps
105
enstronment steps
enviromment steps
2nd
W/1 normalized repr (Ours)
W/o & normallzed ropr
[Ours]
256
(c) Representation normalization
(d) MLP hidden layer size
<lide
slide
suicces rate
0.5
0.2
0.2
0.6
3.0
0.2
0.5
0.9
3.5
0.0
0.2
2.4
anviromment Staps
166
anvirommnt sTEDS
anstronmant staps
ensiromment staps
205
0.90
(c) Representation dimension
(1) Discount factor Y
Figure 10: Hyperparameter ablation. We conduct ablations to study the effect of different hyper- paramters listed in Table and the discout factor 7 on state-based push and slide.
different values and conducted experiments on push (state) and slide (state), one easy task and one challenging task, for five random seeds. We report and standard deviation of success rates in Fig. 10 These results suggest that while some values of the hyperparameter have similar effects, e.g. actor learning rate- 10-5 1 x 10-1, our choice of combination is optimal for TD InfoNCE.
E.2 PREDICTING THE DISCOUNTED STATE OCCUPANCY MEASURE
Our experiments estimating the discounted state occupancy measure in the tabular setting (Sec. 4.3) observed small "irreducible" error. To test the correctness of our implementation, we applied the successor representation with a known model (Fig. [T]. finding that its error does go to zero. This gives confidence that our implementation of the successor representation baseline is correct, and suggests that the error observed Fig. 3 arises from sampling the transitions (rather than having a known model).
E.3 UNDERSTANDING THE DIFFERENCES BETWEEN TD INFONCE AND C-LEARNING
While conceptually similar, our method is a temporal difference estimator building upon InfoNCE whereas C-learning instead bases on the NCE objective (Gutmann & Hyvärinen 2010). There are
We use y denote y layers ML.P with x units in each layer.
23

---

Published as a conference paper at ICLR 2024
1,0
0.32
Successor Representation
absolute
C-learning
Successor Impresertation w/ Dynamics
suecees
0.23
0.0
0.00
INOOCE
Cusaming tofmnn
Cueaming mottingge
CLeaming lisss
Clearning
gradient steps
lel
Figure 11: Errors discounted state occupancy
Figure 12: Differences between TD InfoNCE and C-Learning.
measure estimation in a tabular setting.
goal state
goal state
1.0
1.0
0.7
0.9
0.5
0.3
0.5
0.9
0.1
0.1
0.0
Interpolation
0.0
Interpolation coefficients
coefficients
initial state
initial state
(a) Parametric interpolation
(b) Non-parametric interpolation
Figure 13: Visualizing representation interpolation. Using spherical interpolation of representa-
tions (Left) or linear interpolation of softmax features (Right), TD InfoNCE learns representations that capture not only the content of states, but also the causal relationships.
mainly three distinctions between TD InfoNCE and C-Learning: (a) C-Learning uses binary cross entropy loss, while TD InfoNCE uses a categorical cross entropy loss. (b) C-Learning uses impor- tance weights of the form exp(f(s, a. g)); if these weights self-normalized (Dubi & Horowitz/ 1979 Hammersley 1956). they corresponds to the softmax importance weights in our objectives (Eq.) (c) For the same batch of N transitions, TD InfoNCE updates representations of N neg- ative examples (Eq. (7). while C-Learning only involves N negative examples. We ablate these decisions in Fig. [2] finding that differences (b) and (c) have little effect. Thus, attribute the better performance of TD InfoNCE to its use of the categorical cross entropy loss.
E.4 REPRESENTATION INTERPOLATION
Prior works have shown that representations learned by self-supervised learning incorporate struc- ture of the data (Wang & Isola 2020 Arora et al.) 2019), motivating us to study whether the repre- sentations acquired by TD InfoNCE contain task-specific information. To answer this question, we visualize representations learned by TD InfoNCE via interpolating in the latent space following prior work (Zheng et al.) 2023). We choose to interpolate representations learned on the offline AntMaze medi um-play-v2 task and compare parametric interpolation method against non-parametric variant. Importantly, the states and goals this task are 29 dimensions and we visualize them 2D from top-down view.
24

---

Published as a conference paper at ICLR 2024
Parametric interpolation. Given pair of start state and goal (su, 9). compute the normalized representations @(Su. ape-op: 9) and Φgg. Ono-op- g), where Ono-op is an action taking no operation. Applying spherical linear interpolation to both of them results in blended representations,
sin - a)m
(880, Uno-ops g) + sin ad ( Ano-op:
sin
sin
where E [0, 1] is the interpolation coefficient and 77 is the angle subtended by the are between (Su, ano-op: 9) and Φ(g. an--op: These interpolated representations can be used find the spher- ical nearest neighbors in set of representations of validation states { vval Ono-op: g) } and we call this method parametric interpolation.
Non- non-parametric interpolation. We can also sample another set of random states and using (6)
g) as anchors to construct softmax feature for state their representations ( ono-op; (no-op: feat(s; 9, {)rand)):
SOFTMAX ([o(s, op g)_ (s (1)
T rand ono 9), 0(s,, Ono op: g)
(s op 9)l)
We compute the softmax features for representations of start and goal states and then construct the linear interpolated features,
afeat( so: 9, {Srand)) (1 fe)attg g. {saand)
Those softmax features of interpolated representations used to find the &g nearest neighbors in softmax feature validation set. We call this method non-pan con-parmettric interpolation.
Results in Fig. 1 suggest that when interpolating the representations using both methods, the in- termediate representations correspond to sequences of states that the optimal policy should visit when reaching desired goals. Therefore, conjecture that TD InfoNCE encodes causality in its representations while the policy learns to arrange them in a temporally correct order.
25


Reviews:
[{"review_id": "7ITt49PRjA", "reviewer": "Reviewer_yDjK", "summary": "Predicting future states is crucial for many time-series tasks, including goal-conditioned reinforcement learning (RL). While contrastive predictive coding (CPC) has been used to model time series data, learning representations that capture long-term dependencies often requires large datasets. This paper introduces a temporal difference (TD) version of CPC that combines segments of different time series, reducing the data needed to predict future events. This representation learning method is applied to derive an off-policy algorithm for goal-conditioned RL. Experiments show that the proposed method achieves higher success rates with less data and better handles stochastic environments compared to previous RL methods.", "questions": "Can you explain how you selected the hyperparameters for the proposed algorithm?\n\nCan you provide more details about the observation that TD InfoNCE learns on image-based pick & place and slide, while baselines fail to make any progress?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper proposes a new temporal difference (TD) estimator for the InfoNCE loss, which is shown to be more efficient than the standard (Monte Carlo) estimator.\n\n- The proposed goal-conditioned reinforcement learning (RL) algorithm outperforms prior methods in both online and offline settings.\n\n- The proposed algorithm is capable of handling stochasticity in the environment dynamics.\n\n- In stochastic tasks, there is an excellent improvement in performance versus the baseline of Quasimetric RL, with some healthy gains on non stochastic tasks versus other baselines, although this is not the primary target of the paper.\n\n- The paper provides a clear and concise explanation of the proposed algorithm.\n\n- The paper is well-written and easy to understand.\n\n- The paper is well-supported by experiments.\n\n- The proposed algorithm is evaluated on a variety of tasks, including the Fetch robotics benchmark.\n\n- TD InfoNCE learns on image-based pick & place and slide, while baselines fail to make any progress.\n\n- TD InfoNCE maintains high success rates on more challenging tasks with observation corruption, while the performance of QRL decreases significantly.", "weaknesses": "- The paper focuses on fairly trivial environments, it would be nice to see these methods working on more challenging higher dimensional goal conditioned RL tasks, as its not a given that these gains will carry over to tasks that matter a lot more.\n\n- The proposed TD estimator is more complex than the standard (Monte Carlo) estimator and its implementation requires more hyperparameters.\n\n- The performance of the proposed goal-conditioned RL algorithm on the most challenging tasks was less than 50%.\n\n- QRL assumes deterministic dynamic of the environment, while TD InfoNCE learns without such assumption.\n\nLoss Function Composition: The loss function L(θ) is composed of two cross-entropy (CE) loss terms, one for predicting the next state and one for predicting the future distribution of states. The γ hyperparameter is used to weight these two terms, but the choice of γ and its impact on the algorithm's performance are not discussed in detail.", "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": "cIhT3oTCV1", "reviewer": "Reviewer_2Mg8", "summary": "This paper derives a TD variant of the InfoNCE objective function, relating this to some of the work in reinforcement learning using future distributions of state as an objective (i.e. successor representations/features). This algorithm is then applied to goal-conditioned reinforcement learning, showing competitive performance among several baselines.", "questions": "See above.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The derived method fits nicely within the literature and seems to fill a nice gap between contrastive objectives from self-supervised objectives and more online focused temporal-difference updates.", "weaknesses": "After the rebuttal, I'm raising my score. While I believe there are issues with empirical section still, these are issues the rest of the literature are also facing. I don't think rejecting this paper is a way to a solution. I also appreciate the fix to some of the inaccurate statements that were overlooked! \n\nGreat job authors!\n\n--------before edit-------\n\nThis paper struggles with clarity and accuracy in some of the ancillary statements made about the literature surrounding the paper and in the main content of the paper itself. There are also several issues with the experimental section that should be addressed.\n\n## Accuracy and Clarity:\n1. On page 1, in the paragraph starting with “the key aim…”, the language on motivating why a TD version of InfoNCE is oddly phrased. I think a fix should be easy here by removing phrases such as “may allow”, “may enable” and be more actionable in your language. This could also be replaced by actual hypotheses of what you expect to see from your objective and stated more formally.\n2. In the same paragraph, and further throughout, you make a statement which suggests TD can do counter factual reasoning (or in the parlance of RL make use of off-policy updates) while Monte Carlo estimates cannot.  This is not true as Monte Carlo estimates can be made with off-policy corrections (using Importance Sampling like in TD). While this induces a more variant estimate as compared to TD (there is a classic bias-variance trade-off between MC and TD updates) and a TD update enables the use of incomplete trajectories (because you are using an estimate to inform your update rather than a full trajectory), I came out of the paper with the feeling the paper was suggesting Monte Carlo estimates couldn’t be off-policy. \n    - This comes up very strongly on page 5 (the first paragraph) “This is, we cannot share…”. We should be able to derive an off-policy version of the monte-carlo update for InfoNCE. This doesn’t mean it would be an estimator we would want to use in this setting, but it should be definable. If this is not the case, the paper should show that this can’t be done using importance sampling or cite a reference which shows monte carlo estimates can’t be off-policy.\n\n3. **Notation clarity issues:** \n   - In your expectations above equation 7, I’m not sure what s’, a’ are here. Are you using these instead of s_{t+1} as used in equation 4? This notation should be unified. \n   - Equation (1) and following uses, I’m not sure you explain what the superscript is signifying. I think it is time, but it is not clear from the writing.\n   - Shouldn’t the expectation in the middle of page 4 on the RHS (i.e. after applying the importance weight) be selecting actions a’ from the behavior policy? Or is the importance weight only correcting for the state distribution? Shouldn’t we also correct for the action distribution as well?\n\n\n\n## Empirical Results:\n\nThere are two major issues with the empirical results as presented. \n\n### Major issues:\n\n4. 3 seeds is too few to get any statistical confidence, especially without doing independent hyperparameter sweeps for each baseline. While in the past this has been standard, as a field we continually have shown that the statistical power of our experiments are laughably poor, even if the bounds of our results show statistical significance. This continually misleads the community, and needs to be addressed. This doesn’t include the issue with not running hyperparameter studies on the methods independently. See https://sites.ualberta.ca/~amw8/cookbook.pdf for a reference. While this paper is a draft it does a good job going through these issues and providing context from the literature.\n\n5. In the appendix it is mentioned “We increase the capacity of the goal-conditioned state-action encoder…”. This suggests the model you are using may have more parameters than your counterparts. Is this true? Also did you use a larger batch size for all baselines or just your algorithm? If this was done for your method, this makes the results difficult to trust. If not, it likely means the hyperparameters of your baselines are now invalid. If both methods have the same hyperparameters then I’m usually ok with re-using old hyperparameter studies, unfortunately when any of the hypers change OR the new method has additional hypers this begins to weaken the validity of re-using the same hypers.\n\n### Minor Issues:\n- You should include all hyperparameters you used, even for baselines. \n- You include only success rate as a metric to compare. While this is reasonable, I think there is a lot that could be learned from more traditional metrics (i.e. return or something similar). This is especially the case when the success metric doesn’t clearly separate the methods (for instance reach, Push, slide (state) in firgure 2a). \n\n\n\n# Edits/Suggestions \n- I don’t like the notation $s_{t+}$. I think it could be replaced with something that is more understandable on first read (and looks less like a mistake). This is a preference though.\n- Page 4: “Our proposed method (Sec 3)…” Should say Sec 3.2.\n- Page 4: “we conjecture that our…” -> You should state that you test this in the empirical section (I think you do at least).\n- Page 7: “We will to evaluate” -> “we will evaluate”\n- If you reference a result in the main paper you should include this in the main paper. “TD InfoNCE achieves a 2x median improvement etc…” references figure 6 (I believe).", "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": "Tn58GxBzxM", "reviewer": "Reviewer_a2XQ", "summary": "This paper extends previous works on contrastive RL by using a temporal difference estimator. Similar to contrastive RL, the paper proposes to use contrastive learning, InfoNCE in particular, for estimating the discounted state occupancy measure. Unlike the previous contrastive RL approach, which averages over the goal distribution, the proposed method uses a temporal difference estimator and results in a Bellman-like update rule (TDInfoNCE). Comparing the Bellman update for value function, TDInfoNCE requires taking expectation over the future states from a potentially different goal. It turns out that this can be estimated using importance weight. Then, the paper shows how the estimated state occupancy measure can be used in conjuction with goal-conditioned RL to form a full-fledged RL agent. Experimental results are shown on both online and offline settings and compared to a range of existing methods. In both settings, TDInfoNCE outperforms in most of the environments compared to previous approaches.", "questions": "1. Questions regarding the algorithm:\n - It's not clear to me what's the goal distribution is used? Is it a random goal? Does the different goal distribution affect data efficiency?\n- How is $s_{t+}$ sampled? Is it the same as previous approaches - sample t from a geometric distribution?\n\n2. The notation in Figure 1 is not very clear to me. Is it suppose to visualize Eq. 4?\n\n3. In Eq. 7 and the one above, shoud it be $p^{\\pi}(s_{t+}^{(1)}|s^{'}, a^{'}, g)$ instead of  $p^{\\pi}(s_{t+}^{(1)}|s^{'}, a^{'})$?\n\n4. In Table 1, result for contrastive RL on large-diverse-v2 should not be bold; result for TDInfoNCE on unmaze-v2 should not be bold.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The paper is mostly well written, apart from some details (see questions section). \n\nThe derivations are sound. \n\nExperimental results show strong performance comparing to previous methods. The paper also presents some analysis and insights to explain the performance.", "weaknesses": "The novelty is slightly limited. The idea of using InfoNCE to estimate the state occupancy measure has been presented in contrastive RL; the Bellman-like update and the use of importance weight has been presented in C-Learning.", "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": "lJwlZm0trs", "reviewer": "Reviewer_Jmo8", "summary": "This paper studies the representation learning for goal-conditioned reinforcement learning problems. Built upon InfoNCE objective, this paper proposes a temporal difference estimator of InfoNCE objective and applied it to goal-conditioned RL algorithm.\n\nIn the experiment section, in the online RL setting, the proposed method is compared with prior goal-conditioned RL algorithms, including quasimetric RL, contrastive RL (Monte Carlo estimator of NCE objective), and hindsight experience relabelling. The proposed method achieved a higher average return in the comparison. Also, in the offline RL setting, The proposed method is compared with quasimetric RL, contrastive RL, and SOTA offline RL algorithms. The proposed algorithm generally outperforms baselines.", "questions": "How is the proposed method sensitive to hyper-parameters? Do we need careful hyper-parameter tuning to make it work? Is there any intuitive guidance about how to adjust the hyper-parameters?\n\nIn Algorithm 1, many notations are introduced for the first time without any definition. Could you please clarify them?\n\nOne important baseline, contrastive RL, is the Monte Carlo estimator of the NCE loss. Could you please also compare with the algorithm using Monte Carlo estimator of InfoNCE loss, since it is already introduced in the prior work Eysenbach et al., 2022", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "This paper is well-written and easy to read. The proposed method is explained and presented clearly.\n\nThis paper provides clear derivation and a solid theoretical foundation for the proposed method in Section 3. I \n\nThe proposed method is supported by extensive experiments comparing with many baseline approaches.\n\nThe analysis in Section 4.3 and 4.4 validate the advantages of the proposed method in comparison with other representation learning approaches. It is impressive to see that the proposed method can stitch together pieces of data.", "weaknesses": "Some statements, especially some in the introduction part, seem not fully supported by evidence provided in the paper.\n\nFor example, it is claimed that the proposed method \"enables us to perform counterfactual reasoning\". However, this point is not clear in the following section. Could you please explain it in detail?", "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)