Return-Based Contrastive Representation Learning for Reinforcement  Learning
Authors: Guoqing Liu
Chuheng Zhang
Li Zhao
Tao Qin
Jinhua Zhu
Li Jian
Nenghai Yu
Tie-Yan Liu
Venue: ICLR 2021 Poster (2021)
Keywords: reinforcement learning
auxiliary task
representation learning
contrastive learning

Abstract: Recently, various auxiliary tasks have been proposed to accelerate representation learning and improve sample efficiency in deep reinforcement learning (RL). However, existing auxiliary tasks do not take the characteristics of RL problems into consideration and are unsupervised. By leveraging returns, the most important feedback signals in RL, we propose a novel auxiliary task that forces the learnt representations to discriminate state-action pairs with different returns. Our auxiliary loss is theoretically justified to learn representations that capture the structure of a new form of state-action abstraction, under which state-action pairs with similar return distributions are aggregated together. Empirically, our algorithm outperforms strong baselines on complex tasks in Atari games and DeepMind Control suite, and achieves even better performance when combined with existing auxiliary tasks.

Published as a conference paper at ICLR 2021
RETURN-BASED CONTRASTIVE REPRESENTATION LEARNING FOR REINFORCEMENT LEARNING
Guoqing Liu' Chuheng Zhang Li Zhao Tao Qin', Jinhua Zhu', Jian , Nenghai Yu', Tie- Yan Liu Univeesitty of Science and Technology of China
IISS, Tsinghua University
3 Microsoot Research
1gq1001 @maillmmil..... uste ustc.edu. en, zhangchuhngg223@1vee com
[lizo, taoqin, (mlluu@microooi
teslazhnee aaslazuummaii ustc.. edu cn
1ijian83email1 hammail.ssinghhie tssinguu.. edu. cn, ynh@uste. edu. cn
ABSTRACT
Recently, various auxiliary tasks have been proposed to accelerate representa- tion learning and improve sample efficiency deep reinforcement learning (RL). However, existing auxiliary tasks do not take the characteristics of RL problems into consideration and are unsupervised. By leveraging returns, the most impor- tant feedback signals in RL, we propose a novel auxiliary task that forces the learnt representations discriminate state-action pairs with different returns. Our auxil- iary loss is theoretically justified to learn representations that capture the structure a new form of state-action abstraction, under which state-action pairs with sim- ilar return distributions outperforms strong baselines on complex tasks in Atari games and DeepMind Control suite, and achieves even better performance when combined with existing auxiliary tasks.
aggregated together. In low data regime, algorithm
INTRODUCTION
Deep reinforcement learning (RL) algorithms can learn representations from high-dimensional in- puts, as well as learn policies based on such representations to maximize long-term returns simul- taneously. However, deep RL algorithms typically require large numbers of samples, which can be quite expensive obtain (Mnih al. 2015). In contrast, is usually much more sample efficient to learn policies with learned reprrssentationseextraaee
features (Srinivas et al. 2020). To this end, various auxiliary tasks have been proposed to accelerate representation learning in aid of the main RL task (Suddarth and Kergosien 1990; [Sutton et al.) [001 Gelada et al.) 2201 Bellemare et al. François-Lavet et al. 2019 Shen et 2000 et al. 20220 Dabney et al. [0200] Srinivas 2020). Representative examples of auxiliary tasks include predicting the future in either the space or the latent space with reconstruction based losses (e.g.. Jaderberg et al. 2016 Hafner
2019
al.
pixel
al.
2019alb)
Recently, contrastive learning has been introduced to construct auxiliary tasks and achieves het- ter performance compared reconstruction based methods in accelerating RI. algorithms (Oord) ct al. 2018; Srrinvaa et al. 2020). Without the need to reconstruct inputs such as raw pixels, con- trastive learning based methods can ignore irrelevant features such as static background in games and learn more compact representations. Oord et al. (2018) propose a contrastive representation learning method based on the temporal structure of state sequence. (Srinivas al.] (2020) propose to leverage the prior knowledge from computer vision, learning representations that are invariant to image augmentation. However, existing works mainly construct contrastive auxiliary losses in an unsupervised manner, without considering feedback signals in RI. problems as supervision.
In this paper, we take a further step to leverage the return feedback design a contrastive auxil- iary loss to accelerate RL algorithms. Specifically, we propose novel method, called Return-based
Thii work is conducted Microsoft Research Asia. The first two authors contributed equally to this work.

---

Published as a conference paper at ICLR 2021
Contrastive representation learning for Reinforcement Learning (RCRL). In method, given an anchor state-action pair, we choose a state-action pair with the same or similar return as the pos- itive sample, and state-action pair with different return the negative sample. Then, train a discriminator to classify between positive and negative samples given the anchor based on their representations the auxiliary task. The intuition here is to learn state-action representations that capture return-relevan features while ignoring cttrrn-ireeeeant features.
From a theoretical perspective, RCRL is supported by a novel state-action abstraction, called Z"- irrelevance. 2 -iieeeeeevence abstraction aggregates state-action pairs with similar return distributions under certain policy и. We show that 2 -rreeeeeeeence abstraction can reduce the size of the state- action space (ef. Appendix(A) as well as approximate the Q values arbitrarily accurately (cf. Section We further propose method called Z-learning that can calculate Z*-irrelevance abstraction with sampled returns rather than the return distribution, which is hardly available in practice. Z-- learning can learn Z -irrelevance abstraction provably efficiently. Our algorithm RCRL can be seen as the empirical version of Z-learning by making few approximations such as integrating with deep RL algorithms, and collecting positive pairs within a consecutive segment in trajectory of the anchors.
We conduct experiments on Atari games (Bellemare et al., 2013) and Deep.indd Control suite (Tassa ct al. 2018) in low data regime. The experiment results show that our auxiliary task combined with Rainbow (Hessel et al. 2017) for discrete control tasks or SAC (Haamnjj et al.] (018) for continuous other state-of-the-art baselines for this regime. Our method can be further combined with existing unsupervised contrastive learning methods to achieve even better performance. We also perform detailed analysis how the representation changes during training with/without our auxiliary loss. We find that a good embedding network assigns similar/dissimilar representations to state-action pairs with similar/dissimilar return distributions, and our algorithm can boost such generalization and speed up training.
control tasks achieves superior performance
Our contributions are summarized follows:
We introduce a novel contrastive loss based on return, to learn eeturntrelevant representa- tions and speed up deep RL algorithms.
We theoretically build the connection between the contrastive loss and new form of state- action abstraction, which can reduce the size of the state-action space as well as approxi- mate the Q values arbitrarily accurately.
Our algorithm achieves superior performance against strong baselines in Atari games and DeepMind Control suite in low data regime. Besides, the performance can be further en- hanced when combined with existing auxiliary tasks.
2 RELATED WORK
2.1 AUXILIARY TASK
In reinforcement learning. the auxiliary task can be used for both the model-hased setting and the model-free setting. In the model-based settings, world models can be used as auxiliary tasks and lead better performance, such as CRAR (François-Lavet et al., 2019), Dreamer (Hafner et al.) (019a)), and PlaNet (Hafner et al. 2019b). Due to the complex components (c.g., the latent transition or reward module) in the world model, such methods are empirically unstable to train and relies on different regularizations to converge. In the model-free settings, many algorithms construct various auxiliary tasks improve performance, such as predicting the future (Jaderberg al., 2016 Shel-
hamer et al. 2016 Guo et al.) 2020 Lee et al. (2020: Mazoure et al. (0220). learning value functions with different rewards or under different policies (Veeriah et al., 2019 (Schaul et al. 2015 Borsa ct al. 2018; [Bllemaaee et al., 2019 Dabney et al.) (020). learning from many-goal Veeriah et al.
2018
or the combination different auxiliary objectives (de Bruin et al. (0188. Moreover, aux-
Iliary tasks can be designed based on the prior knowledge about the environment MMrowski al.
2016 Shen et al.) 2020 van der Pol et al. (020) or the raw state representation
(Srinivas et al.
2020
Hessel al. (2019) also apply auxiliary task to the multi-task RL setting.
Contrastive learning has seen dramatic progress recently, and been introduced to learn state repre-
sentation (Ord et al. 2018 Sermanet et al. 2018 Dwibedi al. 2018 Aytar al. 2018 Anand
2

---

Published as a conference paper at ICLR 2021
et al. 2019 Srinivas et al., 2020). Temporal structure (Sermanet et al. 2018 Aytar et al. 2018 and local spatial structure (Anand et al., 2019) has been leveraged for state representation learning via contrastive losses. CPC (Oord et al. 2018) and CURL (Srinivas et al. 2020 adopt contrastive auxiliary tasks to accelerate representation learning and speed up main RI. tasks, by leveraging the temporal structure and image augmentation respectively. To the best of our knowledge, we are the first to leverage return to construct a contrastive auxiliary task for speeding up the main RL task.
2.2 ABSTRACTION
State abstraction (or state aggregation) aggregates states by ignoring irrelevant state information. By reducing state space, state abstraction can enable efficient policy learning. Different types of abstraction are proposed in literature, ranging from fine-grained to coarse-grained abstraction, each reducing state space to a different extent. Bisimulation or model irrelevance (Dean and Givan 1997; Givan et al. 2003) define state abstraction under which both transition and reward function are kept invariant. By contrast, other types of state abstraction that are coarser than bisimulation such as Q" irrelevance or Q" irrelevance (see e.g.. Li et al. 2006), which keep the Q function invariant under any policy the optimal policy respectively. There also some works on state-action abstractions, e.g. MDP MDP homomorphism (Ravindran [2003] Ravindran and Barto 2004a) and approximate MDP homomorphism (Ravindran and Barto] 2004b; Taylor et al. (2009). which are similar to bisimulation in keeping reward and transition invariant, but extending bisimulation from state abstraction to state- action abstraction.
In this paper, we consider a new form of state- aate-action abstraction z"-irrelevance, which aggregates state-action pairs with the same return distribution and is coarser than bisimulation homomor- phism which are frequently used as auxiliary tasks (e.g.. Biza and Platt 2018 Gelada et al.) 2019] Zhang et al.) 2020). However, it is worth noting that Z=-irrelevance only used build the theo- retical foundation of our algorithm, and show that our proposed auxiliary task is well-aligned with the main RL task. Representation learning in deep RL is in general very different from aggregating states in tabular case, though the latter may build nice theoretical foundation for the former. Here we focus on how design auxiliary tasks to accelerate representation learning using contrastive learning techniques, and we propose novel return- contrastive method based on our proposed 7" -irreeevance abstraction.
3 PRELIMINARY
We consider a Markov Decision Process (MDP) which is a tuple (S, (S, P, R. nI, 7) specifying the state space S, the action space A. the state transition probability P(s at), the reward R(r:|s;,a, a.), the initial state distribution j E A and the discount factor 7. Also, we denote = () E X == S x A to be the state-action pair. A (stationary) policy π S - A specifies the action selection probability on each state. Following the policy π, the discounted sum of future rewards (or return) is denoted by the random variable Z a) Lee ([,s,, az). where so 8, 40 = a, ~ P(- P(-1 0z 1;-1), and a; ~ π (-(+)). We divide the range of return into K equal bins {Ro = Rminn R2,
Rx = Rmax) such that Rg R (Rmax Rmmin// /K, Ak E [K], where Rmai (resp. Rmax) is the minimum (reps. maximum) possible return, and [K] := {1,2,
K).
We use b(R) = k [K] to denote the event that R falls into the kth bin, i.c., RR-1 < R RR- Hence, b( R)) can be viewed as the discretized version of the return, and the distribution of discretized return can be represented by K -ddimensional vector Zd (x) A where the k-th element equals to Pr[ < (x) < Rg.. The Q function is defined as Q ) = [I])(x)] and the state value function is defined as (T(s) Ba- [2 (s, a)). The objective for RL is find policy æ that maximizes the expected cumulative reward J(x) E.- (V*(s)] We denote the optimal policy as
T and the corresponding optimal Q function as Q" := ==
4 METHODOLOOGY
In this section, present method, from both theoretical and empirical perspectives. First, propose 2 -irreleevance a new form of state-action abstraction based on return distribution. We show that the Q functions for any policy (and therefore the optimal policy) can be represented under
3

---

Published as a conference paper at ICLR 2021
Algorithm 1: Z-leeaming
Given the policy n, the number of bins for the return K, constant N 2 Nx the encoder
class NN the regressor class Wy, and a distribution d A with supp(d) X
2: D D=0
3: for =ai==1....
do
4: 21, ~ d
5: R. ~ (x1)), R2 (x)2)
6: D=DU [111, [[b(R1) + ((2))))
7: end for
8: (Ø, 2) arg min L(o, w; D), where L(o, w; D) is defined in !
9: return the encoder Φ
-"-irrelevance abstraction. Then we consider an algorithm, Z-learning, that enables learn -"-irrelevance abstraction from the samples collected using π. Z-learning is simple and learns the abstraction by only minimizing a contrastive loss. We show that Z-leamnng can learn 2" -rrieeeleveece abstraction provably efficiently. After that, we introduce return-based contrastive representation learning for RL (RCRL) that incorporates standard RL algorithms with an auxiliary task adapted from Z-learning. At last, we present our network structure for learning state-action embedding, upon which RCRL built.
4.1 -T IRIREEEVAN ABSTRACTION
A state-action abstraction aggregates the state-action pairs with similar properties, resulting in an abstract state-action space denoted as [N], where N is the size of abstract state-action space. In this paper, we consider new form of abstraction, Z"-irrelevance, defined as follows: Given policy x, -rrireelevance abstraction is denoted as Φ X - [N] such that, for any I1, E X with (xx1) ((22), we have Zd (a1) Z= (I22). Given a policy и and the parameter for return discretization K, use Nr.K denote the minimum N such that Z*-irrelevance exists. It is true that Nr, S No S ( AA| for any π and K, where ( is the number of abstract states for the coarsest bisimulation (ef. Appendix A..
Proposition 4.1. Given policy π and any Z" Z"-irrelevance o X + [N], there exists a function Q: [N] + R such that (Q(((x)) (x:) Remnn Vx E X.
We provide a proof in Appendix A Note that controls the coarseness of the abstraction. When accurately represent the value function and therefore the optimal policy when π + ~² When using an auxiliary task to learn such abstraction, this proposition indicates that the auxiliary task (to learn Z -=-irrelevance) well-aligned with the main RL task (to approximate Q*). However, large K results in a fine-grained abstraction which requires to use a large N and more samples to learn the abstraction (cf. Theorem [.J].. In practice, this may not be a problem since we learn state-action representation in low-dimensional space Rd instead of [N] and reuse the samples collected by the base RL algorithm. Also, we do not need choose explicitly in the practical algorithm (cf. Section 4.3)
K - co, 7 -i-irrelevance
4.2 Z-LEARNING
We propose Z-lleamiig to learn -"-irrelevance based on a dataset D with a contrastive loss (see Algorithm[]] D. Each tuple the dataset collected as follows: First, two state-action pairs are drawn i..dd from a distribution d E A with supp(d) = X' (cf. Line 4 in Algorithm D). In practice, we can sample state-action pairs from the rollouts generated by the policy æ. In this case, a stochastic policy (e.g., using c-greedy) with a standard ergodic assumption on MDP ensures supp(d) = X. Then, we obtain returns for the two state-action pairs (i.e., the discounted sum the rewards after I and 222) which can be obtained by rolling out using the policy π (cf. Line 5 in Algorithm The binary label for this inte-action pair indicates whether the two returns belong to the same bin (cf. Line 6 in Algorithm
D. The contrastive loss is defined follows:
min weWy
L(o, D) == (e =DD [(w((())) ((22)) - y)²]
(1)
4

---

Given the encoder è returned by Algorithm 1 the following inequality holds with Published as a conference paper at ICLR 2021
where the class of encoders that map the state-action pairs to N discrete abstractions is defined as
@N := { + [M]], and the class of tabular regressors is defined as Wy := {[N] x [N] - [0,1]]- Notice that choose N 2 Nr, to ensure that Z d'-irrelevance & X' + [N] exists. Also. to aggregate the state-action pairs, should be smaller than | (otherwise we will obtain an identity mapping). In this case, mapping two state-action pairs with different return distributions to the same abstraction will increase the loss and therefore is avoided. The following theorem shows that Z-learning can learn - provably efficiently.
Theorem 4.1.
probability 1 8 and for any x' X:
B [Tδ(1)) (((22) 2 (x') (z" (x) 2d (12)) |
(2)
SN 3 33+422 Inn + 4ln | + 4 41 2)
where ( the cardinality of encoder function class and the size of the dataset.
We provide the proof in Appendix Although N is upper bounded by N/ it is generally much smaller for deep encoders that generalize over the state-action space. The theorem shows that whenever maps two state-actions 1 32 the same abstraction, 2 (x1x)) ) ~ - (x22 up to an error proportional to 1/vii (ignoring the logarithm factor). The following corollary shows that becomes a 2d -rireeeleeecce when +
Corollary 4.1.1. The encoder è returned by Algorithm with n - oo is a Z -r-irrelevance. i.e. any 21, 22 E X, Zd (x1) = z" (x22 if ((x1) = 2(x2).
4.3 RETURN- CONTRASTIVE LEARNING FOR RL (RCRL)
We adapt Z-learning as an auxiliary task that helps the agent to learn a representation with mean- ingful semantics. The auxiliary task based RL algorithm is called RCRL and shown in Algorithm 2 Here, we use Rainbow (for discrete control) and SAC (for continuous control) as the base RL algorithm for RCRL. However, RCRL can also be easily incorporated with other model-free RL algorithms. While Z-learning relies on a dataset sampled by rolling out the current policy, RCRL constructs such a dataset using the samples collected by the base RL algorithm and therefore does not require additional samples, e.g., directly using the replay buffer in Rainbow or SAC (see Line 7 and in Algorithm Dp. Compared with Z-lleaning, we use the state-action embedding network that shared with the base RL algorithm Φ X' + Rd as the encoder, and use an additional discriminator trained by the auxiliary task w. IRd x IRd + [0, 1] as the regressor.
the auxiliary task, the labels in the dataset may be unbalanced. Although this does not cause problems in the theoretical analysis since we assume the Bayes optimizer be obtained for the contrastive loss, it may prevent the discriminator from learning properly in practice (cf. Line 8 in Algorithm Ip. To solve this problem, instead of drawing samples independently from the replay buffer B (analogous in Z-learning), we sample the pairs for D as follows: As a preparation, we cut the trajectories in B into segments, where each segment contains state-action pairs with the same similar returns. Specifically, in Atari games, we create a new segment once the agent receives a non-zero reward. In DMControl tasks, we first prescribe a threshold and then create new segment once the cumulative reward within the current segment exceeds this threshold. For each sample in D, we first draw an anchor state-action pair from randomly. Afterwards, we generate a positive sample by drawing state-action pair from the same segment of the anchor state-action pair. Then, we draw another state-action pair randomly from B and use it the negative sample.
However, when implementing Z-learning
sampling from the distribution d
We believe our auxiliary task may boost the learning, due to better return-inducen representations that facilitate generalization across different state-action pairs. Learning on one state-action pair will affect the value of all state-action pairs that share similar representations with this pair. When the embedding network assigns similar representations to similar state-action pairs (e.g., sharing similar distribution over returns), the update for state-action pair is representative for the updates for other similar state-action pairs, which improves sample efficiency. However, such generalization may not be achieved by the base RL algorithm since, when trained by the algorithm with only
5

---

Published as a conference paper at ICLR 2021
Algorithm 2: Return based Contrastive learning for RL (RCRL)
Initialize the embedding Φ0 X' + Rd and discriminator wa R x R - 00
Initialize the parameters &e for the base RL algorithm that uses the learned embedding do
Given batch of samples D, the loss function for the base RL algorithm is LRL (0o, ppi D)
A replay buffer B
foreach iteration do
Rollout the current policy and store the samples to the replay buffer B
Draw a batch of samples D from the replay buffer B
Update the parameters with the loss function L(0o; wa; D) + LRI yp; D)
end
return The learned policy
return-based loss, similar state-action pairs may have similar Q values but very different represen- tations. One may argue that, we can adopt several temporal difference updates to propagate the values for state-action pairs with same sampled return, and finally all such pairs are assigned with similar Q values. However, since adopt a learning algorithm with bootstrapping/lemprrd differ- ence learning and frozen target network in deep RL, it could take longer time to propagate the value across different state-action pairs, compared with direct generalization over state-action pairs with contrastive learning.
Mcanwhile, since construct auxiliary tasks based return, which is very different structure from image augmentation or temporal structure, our method could be combined with existing meth- ods to achieve further improvement.
4.4 NETWORK STRUCTURE FOR STATE-ACTION EMBEDDING
In our algorithm, the auxiliary task is based on the state-action embedding, instead of the state embedding that is frequently used in the previous work (e.g., [Siinvas al. 2020). To facilitate our algorithm, we design two new structures for Atari 2600 Games (discrete action) and DMControl Suite (continuous action) respectively. We show the structure in Figure [ For Atari, we learn an action embedding for each action and use the element-wise product of the state embedding and action embedding as the state-action embedding. For DMControl, the action embedding is a real- valued vector and we use the concatenation of the action embedding and the state embedding.
Rainbow Loss ((,aa
Auxiliary
SAC Loss
Auxiliary
Input
(6)
Q(s.a)
Output
Trainable
Value net
Discriminator
Actor
Critic
Discriminator
Variable
Legend
(s, embedding
(s', (s'aa)) embedding
(5.4) embedding
(s'a') embedding
Element wise product
Concatenation
State embedding
Action embedding
State embedding
Flatten
Flatten
Conv
Conv
State
Action
RCAL for Atari
State
Action a
RCAL for Control
Figure 1: The network structure of RCRL for Atari (left) and DMControl Suite (right).
EXPERIMENT
In this section, we conduct the following experiments: 1) We implement RCRL on Atari 2600 Games (Bellemare et al. 2013) and DMControl Suite (Tassa et al. 2020). and compare with state- of-the-art model-free methods and strong model-based methods. In particular, we compare with CURL (Srinivas et al. 2020), top performing auxiliary task based RL algorithm for baxel-based control tasks that also uses a contrastive loss. addition, we also combine RCRL with CURL to study whether our auxiliary task further boosts the learning when combined with other auxiliary

---

Published as a conference paper at ICLR 2021
tasks. 2) To further study the reason why our algorithm works, we analyze the generalization of the learned representation. Specifically, we compare the cosine similarity between the representations of different state-action pairs. We provide the implementation details in Appendix
5.1 EVALUATION ON ATART AND DMCONTROI.
Experiments on Atari. Our experiments on Atari are conducted in low data regime of 100k in- teractions between the agent and the environment, which corresponds to two hours of human play. We show the performance of different algorithns/baselines including the scores for average human (Human), SimPLe (Kaiser et al. 2019) which a strong model-based baseline for this regime, orig- inal Rainbow (Hessel et al. 2017),
Rainbow (ERainbow, van Hasselt et al. 2019), ERainbow with state-action embeddings (ERainbow-sa, cf. Figure Left), CURL (Srinivas et al,, 2020) that is based on ERainbow, RCRL which is based FRainbow-sa, and the algorithm that combines the auxiliary loss for CURL and RCRL to FRaainbo-sa (RCRL+CURL). We show the evaluation results of our algorithm and the baselines on Atari games in Table First, we observe that using state-action embedding instead of state embedding in ERainbow does not lead to significant performance change by comparing ERainbow with ERainbow-sa. Second, built upon ERainbow-sa, the auxiliary task in RCRL leads to better performance compared with not only the base RL algo- rithm but also SimPLe and CURL in terms of the median human normalized score (HNS). Third, we can see that RCRL further boosts the learning when combined with CURL and achieves the best per- formance for 7 out of 26 games, which shows that our auxiliary task can be successfully combined with other auxiliary tasks that embed different sources of information to learn the representation.
Human
SimPL.c
Rainbow
ERainhow
ERainbow-st
CURL
RCRL
RCRLICURI
ALIEN
7127.7
616.9
318.7
739.9
813.8
558.2
K54.2
912.2
AMIDAR
1719.5
88.0
32.5
188.6
154.2
142.1
157.7
125.
ASSAULT
742.0
527.2
231.0
431.2
576.2
600.6
569.6
583.4
ASTERIX
8503.3
1128.3
243.6
470.8
697.0
734.5
799.0
683.0
HEIST
753.1
131.6
107.2
99.0
BATTLE ZONE
37187.5
5184.4
2360.0
10124.6
13920.0
14870.0
14280.0
17380.0
BONING
2.2
2.7
BREAKOUT
30.5
16.4
1.2
1.9
3.4
4,9
4.3
4.0
CHOPPER COMMAND
7387,8
1246.9
120.0
861.8
106-10
1058.5
1262.0
1008.0
CRAZY CLIMBER
35829.4
62583.6
2254.5
16185.3
21840.
12146.5
15120.0
15032.0
DEMON ATTACK
1971.0
208.1
163.6
508.0
768.0
817.6
790.4
618.3
HREEWAY
29.6
0.0
26.7
2
PROSTBITE
4334.7
254.7
60.2
866.8
1472.0
1181.3
1337.6
1516.6
GOPHER
2412.5
771.0
431.2
349.5
384.8
669.3
429.6
458.8
HERO
30826.4
2656.6
487.0
6857.0
4787.9
6279.3
6454.1
7647.4
JAMESBOND
302.8
125.3
47.4
301.6
308.0
471.0
314.0
503.0
KANGAROO
3035.0
323.1
732.0
KRULL
2665.5
4539.9
1468.D
2851.5
2740.0
4229.6 6
2997.5
3905 8
KUNG MASTER
22736.3
17257.2
0.0
11914.0
11856.0
MS PACMAN
6951.6
1480,0
67.0
1204.1
1384.5
1465.5
1555.2
1336.8
PONG
14,6
12,8
-20.6
-19.3
-18.3
-16.5
-16.9
-18.72
PRIVATE
69571.3
0.0
97.8
80.0
102.6
QBERT
13455.0
1288.8
123.5
1152.9
893.5
1042.4
1121.0
942.0
ROAD RUNNER
7845.D
5640.6
9600.0
5661.0
61.8800
$392.0
SEAQUEST
42054.7
683.3
131.7
354.1
402.0
384.5
375.6
489.5
UP DOWN
11693.2
3350.3
504.6
2877.4
3235,2
2955.2
4210.2
3127,8
Median HNS
14.4%
0.006
16.1%
16.7%
17.5%
18.5%
19.6%
Table 1: Scores of different algoribhmsbaselines on 26 games for Atari-100k benchmark. We show the mean score averaged over five random seeds.
Experiments on DMControl. For DMControl, we compare our algorithm with the following base- lines: Pixel SAC (Haarnoja et al.] 2018) which is the base RL algorithm that receives images as the input; SLAC Lee et al. 2019) that learns a latent variable model and then updates the actor and critic based on it; SAC+AE Yarats et al. (019) that uses a regularized autoencoder for recon- struction in the auxiliary task; PlaNet (Hafner et al. 2019b) and Dreamer (Hafner et al.) (019a) that learn latent space world model and explicitly plan with the learned model. We also compare with a skyline, State SAC, the receives the dimensioai state representation instead of the image as the input. Different from Atari games, tasks in DMControl yield dense reward. Consequently, we split the trajectories into segments using threshold such that the difference of returns within each segment does not exceed this threshold. Similarly, we test the algorithms in low data regime of 500k interactions. We show the evaluation results in Figure 2 We can see that auxiliary task not only brings performance improvement over the base RL algorithm but also outperforms CURL and other state-of-the. -frtt-rrrtt baseline algorithms in different tasks. Moreover, we observe that algorithm is
7

---

Published as a conference paper at ICLR 2021
Finger, Spin
Caripole, Swingup
Easy
850
800
000
620
400
402
400
200
100
200
300
400
500
303
203
300
403
500
100
200
300
400
800
1000
620
800
ICO
Score
600
402
400
100
200
300
400
500
303
200
300
403
500
100
200
300
400
500
Koteps
Ksteps
Ksteps
Pood SAC
RDLL
CURL
Planet
Dreamer
SACIAE
SLAC
State SAG (skyline)
Figure 2: Scores achieved by RCRL and other baseline algorithms during the training for differ- ent tasks in DMControl suite. The line and the shaded area indicate the average and the standard deviation over random seeds respectively.
more robust across runs with different seeds compared with Pixel SAC and CURL (e.g., for the task Ball in cup, Catch).
5.2 ANALYSIS ON THE LEARNED REPRESENTATION
We analyze on the learned representation our model to demonstrate that our auxiliary task at- tains representation with better generalization, which may explain why our algorithm succeeds. We use cosine similarity to measure the generalization from one statc-action pair to another in the deep learning model. Given two state-action pairs 21, 22 X, cosine similarity is defined as (0(21)" da(a)) [[e[[=)]]] [=0(2]]]] where ()(()) is the learnt embedding network.
We show the cosine similarity of the representations between positive pairs (that are sampled within the same segment and therefore likely to share similar return distributions) and negative pairs (i.e., randomly sampled state-action pairs) during the training on the game Alien in Figure 4 First, we observe that when a good policy is learned, the representations of positive pairs are similar while those of negative pairs are dissimilar. This indicates that good representation (or the representation that supports good policy) aggregates the state-action pairs with similar return distributions. Then, find that our auxiliary loss can accelerate such generalization for the representation, which makes RCRL learn faster.
(b)
())
1330
Rainbow RCRL
Cosine simitarty
Score
Rainbow (pos)
Reinbow fneal
RCRL (pos)
RCRL (neg)
Training steps (K)
Training steps (K)
Figure 3: Analysis of the learned representation Alien. (a) The cosine similarity between the rep- resentations of the positive/negative state-action pair and the anchor during the training of Rainbow and RCRL. (b) The scores of the two algorithms during the training.
8

---

Published as a conference paper at ICLR 2021
6 CONCLUSION
In this paper, propose return-based contrastive representation learning for RL (RCRL), which introduces a return-based auxiliary task to facilitate policy training with standard RL algorithms. Our auxiliary task is theoretically justified to learn representations that capture the structure of Z irrelevance, which can reduce the size the state-action space as well as approximate the Q values arbitrarily accurately. Experiments Atari games and the DMControl suite in low data regime demonstrate that our algorithm achieves superior performance not only when using our auxiliary task alone but also when combined with other auxiliary tasks..
As for future work, we are interested in how to combine different auxiliary tasks in a more so- phisticated way, perhaps with meta-controller. Another potential direction would be providing theoretical analysis for auxiliary tasks and justifying why existing auxiliary tasks can speed up deep RL algorithms.
7
ACKNOWLEDGMENT
Guoqing Liu and Nenghai Yu supported in part by the Natural Science Foundation of China under Grant U20B2047, Exploration Fund Project of University of Science and Technology of China under Grant YD3480002001. Chuheng Zhang and Jian Li supported in part by the National Natural Science Foundation of China Grant 61 8222230, 61772297, 61632016 and the Zhongguancun Haihua Institute for Frontier Information Technology, Turing AI Institute Nanjing and Xi an Institute for Interdisciplinary Information Core Technology.
REFERENCES
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Marc G Belle- mare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning, Nature, 518(7540):529, 2015.
Aravind Srinivas, Michael Laskin, and Pieter Abbeel. Curl: Contrastive unsupervised representa- tions for reinforcement learning. arXiv preprint arXiv:2004. 023220044.4133 04136, 2020.
Steven C Suddarth and YL Kergosien. Rule-injection hints as a means of improving network per- formance and learning time. In European Association for Signal Processing Workshop, pages 120-129. Springer, 1990.
Richard S Sutton, Joseph Modayil, Michael Delp. Thomas Degris, Patrick Pilarski, Adam White, and Doina Precup. Horde: A scalable real-time architecture for learning knowledge from unsu- pervised sensorimotor interaction. The 10th International Conference on Autonomous Agents and Multiagent Systems- Volume 2, pages 761-768, 2011.
Carles Gelada, Saurabh Kumar, Jacob Buckman, Ofir Nachum, and Mare G Bellemare. Deep- mdp: Learning continuous latent space models for representation learning. arXiv preprint arXiv:1906. 02736, 2019.
Mare Bellemare, Will Dabney, Robert Dadashi, Adrien Ali Taiga, Pablo Samuel Castro, Nicolas Le Roux, Dale Schuurmans, Tor Lattimore, and Clare Lyle. A geometric perspective on opti- mal representations for reinforcement learning. Advances Neural Information Processing Systems, pages 4358 4369, 2019.
Vincent François-Lavet. Yoshua Bengio, Doina Precup, and Joelle Pineau. Combined reinforce- ment learning via abstract representations. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pages 3582-3589, 2019.
Wei Shen, Xiaonan He, Chuheng Zhang, Qiang Ni, Wanchu Dou, and Yan Wang. Auxiliary- based deep reinforcement learning for participant selection problem in mobile crowdsourcing. arXiv preprint arXiv:2008.1 /087, 2020.

---

Published as a conference paper at ICLR 2021
Amy Zhang, Rowan McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learning invariant representations for reinforcement learning without reconstruction. arXiv preprint 2rXh::2006.10742, 2020.
Will Dabney, André Barreto, Mark Rowland, Robert Dadashi, John Quan, Mare Bellemare, and David Silver. The celee-imprvvement path: Towards better representations for reinforcement learning. arXiv preprint arXiv:2006.02243 2020.
Max Jaderberg, Volodymyr Mnih, Wojciech Marian Czarnecki, Tom Schaul, Joel Z. Leibo, David Silver, and Koray Kavukcuoglu. Reinforcement learning with unsupervised auxiliary tasks. arXiv preprint /61 1.05397, 2016.
Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. arXiv preprint 1912v 2019a.
Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In International Conference on Machine Learning, pages 2555-2565, 2019b.
Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predic- tive coding, arXiv preprint 1rXXiv:807.037488 2018.
Mare G Bellemare, Yavar Naddaf, Joel Veness, and Michael Bowling. The arcade learning environ- ment: An evaluation platform for general agents. Journal of Artificial Intelligence Research, 47: 253-279, 2013.
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Bud- den, Abbas Abdolmalcki, Josh Merel, Andrew Lefrancq, al. Deepmind control suite, arXiv preprint
2018.
Matteo Hessel, Joseph Modayil, Hado Van Hasselt, Tom Schaul, Georg Ostrovski, Will Dabney, Dan Horgan, Bilal Piot, Mohammad Azar, and David Silver. Rainbow: Combining improvements in deep reinforcement learning. arXiv preprint arXiv: 2017.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In Proceedings of the 35th International Conference Machine Learning, volume 80, pages 1861-1870. JMLR, 2018.
Evan Shelhamer, Parsa Mahmoudieh, Max Argus, and Trevor Darrell. Loss its own reward: Self- supervision for reinforcement learning. arXiv preprint arXiv: 161.070722.0707,,00 2016.
Daniel Guo, Bernardo Avila Pires, Bilal Piot, Jean-bastien Grill, Florent Altché, Rémi Munos, and Mohammad Gheshlaghi Azar. Bootstrap latent-predictive representations for multitask reinforce- ment learning. arXiv preprint arXiv:2004. 14646, 2020.
Kuang-Huei Lee, Ian Fischer, Anthony Liu, Yijie Guo, Honglak Lee, John Canny, and Sergio Guadarrama. Predictive information accelerates learning in rl. arXiv preprint urXiv:2007. 12401, 2020.
Bogdan Mazoure, Remi Tachet des Combes, Thang Doan, Philip Bachman, and R Devon Hjelm. Deep reinforeement and infomax learning. arXiv preprint arXiv:2006 07200007727..... 2020.
Vivek Veeriah, Matteo Hessel, Zhongwen Xu, Janarthanan Rajendran, Richard L Lewis, Junhyuk Oh, Hado P van Hasselt, David Silver, and Satinder Singh. Discovery of useful questions as auxiliary tasks. In Advances in Neural Information Processing Systems, pages 9310-9321, 2019.
Tom Schaul, Daniel Horgan, Karol Gregor, and David Silver. Universal value function approxima- tors. In International conference on machine learning, pages 1312-1320, 2015.
Diana Borsa, André Barreto, John Quan, Daniel Mankowitz, Rémi Munos, Hado Hasselt, David Silver, and Tom Schaul. Universal successor features approximators. arXiv preprint arXiv: 18/2. 2.06666, 2018.
10

---

Published as a conference paper at ICLR 2021
Vivek Veeriah, Junhyuk Oh, and Satinder Singh. Many-goale reinforcement learning. arXiv preprint arXiv:1806.09605, 2018.
Tim de Bruin, Jens Kober, Karl Tuyls, and Robert Babuška. Integrating state representation learning into deep reinforcement learning. IEEE Robolics and Automation Letters, 3(3):194-1401,, 2018.
Piotr Mirowski, Razvan Pascanu, Fabio Viola, Hubert Soyer, Andrew Ballard, Andrea Banino, Misha Denil, Ross Goroshin, Laurent Sifre, Koray Kavukcuoglu, et al. Learning to navigate in complex environments. arXiv preprint arXiv: 0eto6/6111..367, 2016.
Elise der Pol, Daniel E Worrall, Herke van Hoof, Frans A Oliehoek, and Max Welling- Mdp homomorphic networks: Group symmetries in reinforcement learning. arXiv preprint arXiv:2006.16900 16908, 2020.
Matteo Hessel, Hubert Soyer, Lasse Espeholt, Wojciech Czarnecki, Simon Schmitt, and Hado van Hasselt. Multi-task deep reinforcement learning with popart. In Proceedings of the AAAI Confer- ence on Artificial Intelligence, volume 33, pages -196-, 2019.
Pierre Sermanet, Corey Lynch, Yevgen Chebotar, Jasmine Hsu, Eric Jang, Stefan Schaal, Sergey Levine, and Google Brain. Time-contrastiv networks: seff-supervised learning from video. In 2018 IEEE International Conference Robotics and Automation (ICRA), pages 1134-1141. IEEE, 2018
Debidatta Dwibedi, Jonathan Tompson, Corey Lynch, and Pierre Sermanet. Learning actionable rep- resentations from visual observations. In 2018 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS), pages 1577-1584. IEEE, 2018.
Yusuf Aytar, Tobias Pfaff, David Budden, Thomas Paine, Ziyu Wang, and Nando de Freitas. Playing hard exploration games by watching youtube. In Advances in Neural Information Processing Systems, pages 2930-2941, 2018.
Ankesh Anand, Evan Racah, Sherjil Ozair, Yoshua Bengio, Marc-Alexandre Côté, and R Devon Hjelm. Unsupervised state representation learning in atari. In Advances in Neural Information Processing Systems, pages 0769-8782, 2019.
Thomas Dean and Robert Givan. Model minimization in markov decision processes. In AAAI/IAAI, pages 106-111, 1997.
Robert Givan, Thomas Dean, and Matthew Greig. Equivalence notions and model minimization in markov decision processes. Artificial Intelligence, 147(1-2):163-223, 2003.
Lihong Li, Thomas J Walsh, and Michael Littman. Towards a unified theory of state abstraction for mdps. In ISAIM, 2006.
Balaraman Ravindran. Smdp homomorphisms: An algebraic approach to abstraction in semi markov decision processes. 2003.
Balaraman Ravindran and Andrew G Barto. An algebraic approach abstraction reinforcement learning. PhD thesis, University of Massachusetts at Amherst, 2004a.
Balaraman Ravindran and Andrew G Barto. Approximate homomorphisms: A framework for non- exact minimization in markov decision processes. 2004b.
Jonathan Taylor, Doina Precup, and Prakash Panagaden. Bounding performance loss in approximate mdp homomorphisms. Advances in Neural Information Processing Systems, pages 1649-1656, 2009.
Ondrej Biza and Robert Platt. Online abstraction with mdp homomorphisms for deep learning. arXiv preprint arXiv:1811.12929 2018
Yuval Tassa, Saran Tunyasuvunakool, Alistair Muldal, Yotam Doron, Siqi Liu, Steven Bohez, Josh Merel, Tom Erez, Timothy Lillicrap, and Nicolas Heess. dmcontrol: Software and tasks for continuous control, 2020.
11

---

Published as a conference paper at ICLR 2021
Lukasz Kaiser, Mohammad Babaeizadeh, Piotr Milos, Blazej Osinski, Roy H Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, et al. Model-based reinforcement learning for atari. arXiv preprint arXiv:1903.000
2019.
Hado P Hasselt, Matteo Hessel, and John Aslanides. When use parametric models in rein- forcement learning? In Advances in Neural Information Processing Systems, pages 14322-14333, 2019.
Alex X Lee, Anusha Nagabandi, Pieter Abbeel, and Sergey Levine. Stochastic latent actor-critic:
Deep reinforcement learning with latent variable model. arXiv preprint urXiv:197.000955 2019.
Denis Yarats, Amy Zhang, Ilya Kostrikov, Brandon Amos, Joelle Pineau, and Rob Fergus. Im- proving sample efficiency in model-free reinforcement learning from images. arXiv preprint 1rXiv:1910.01741, 2019.
Pablo Samuel Castro. Scalable methods for computing state similarity in deterministic markov decision processes. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pages 10069-10076, 2020.
Dipendra Misra, Mikael Henaff, Akshay Krishnamurthy, and John Langford. abstraction and provably efficient rich-observation reinforcement learning. arXiv:1911.05815. 2019.
Kinematic state arXiv preprint
Marc G Bellemare, Will Dabney, and Rémi Munos. distributional perspective reinforcement learning, arXiv preprint arXiv:1707.06887 2017.
12

---

Published as a conference paper at ICLR 2021
A Z" -.-RREELEVNNCE
A.1 COMPARISON WITH BISIMULATION.
We consider a bisimulation and denote the set of abstract states as 2 = Φ(S).. The bisimulation ΦB is defined as follows:
Definition A.1 (Bisimulation Givan et al. 2003)) @B bisimulation if 48i, 82 E S where
(8($1) = (B()2), Va E A 2,
R(81, a) = R(82, a)
Σ
())
Σ
P(8182,)))
N'E (=)
d'E
Notice that the coarseness of 7" -r-iirelevance is dependent on K, the number of bins for the return.
When K - cc, two acatinctioon pairs z and are aggregated only when (x)x) 2 2dd (x'), which is strict condition resulting in fine-grained abstraction. Here, provide the following proposition to illustrate that 2 -rrreleevvnec abstraction is coarser than bisimulation even when K + oo.
Proposition A.1. Given OB to be the coarsest bisimulation, ΦB induces Z -r-irrelevance abstrac- tion for any policy π defined over Z. Specifically, V8 $2 E where dg((1) = (B(3)), then 7" (81, a) 2 7" (82, a), Va E Afr any policy a defined over Z.
Consider state-action abstraction ΦB that is augmented from the coarsest bisimulation Φb: ΦB(S1, a1) (B( a2) if and only if (B(s)) 6(s2) and aj a2 The proposition indi- cates that (0)
(X) 2 Nz.x0 2 Na.K for any and for any defined over Z, i.e., bisimulation is no coarser than Z"-irrelevance. Note that there exists optimal policy that is defined over (Li et al. 2006). practice, Z" -rreleelveeeee should far more coarser than bisimu- lation when only consider specific policy Therefore learning a 2"-irrelevance should be casier than learning a bisimulation.
Proof. First, for a fixed policy x : Z + A we prove that two state distributions projected to 2 are identical, the corresponding reward distributions or the next state distributions projected to Z will be identical. Then, we use such invariance property to prove the proposition.
The proof for the invariance property goes follows: Consider two identical state distributions Z on the t-th step, Pt, P2 E such that Pt = P2. Notice that we only require that the two state distributions are identical when projected to Z, therefore they may be different S. Specifically, we denote the state distribution over S as P(s) = P(=)q(s|z) where Φg(+) = z, the distribution q for the two state distributions may different. However, we will show that this is sufficient to ensure that the corresponding state distributions on Z are identical on the next step (and therefore the subsequent steps).
We denote R1 and R2 as the reward on the f-th step, which are random variables. The reward distribution on the t-th step specified as follows:
P(R r(P:)
Σ xE2,0EA,EES
P(rls, Pr(ss.0g((5))7(())3))
Σ
((1l2_0)((al2)
zEZ,0EA
(3)
Σ
P(rl2,0)x(dl))333)
zE2,0EA
= =
for any IR, where we use the property of bisimulation, R|s,a D the second equation. This indicates that Rí P. 2 R2|P2.
Rls', a, Vob(s) = db(('), in
13

---

Published as a conference paper at ICLR 2021
Similarly, denote Pf and P/z as the state distribution over Z on the (t + 1)-th step. We have
P((=)|
E
Σ
P(d))s,a) ((((())(()))))(((3
aE
(EE
Σ
zEZ,aEA
Σ
(((3)(
=EZ,nEA
P/
for any z' E 2, where we also use the property of bisimulation in the second equation, i.e., the quantity in the bracket equals for different such that g(s|2) with fixed This indicates that PiI P3
At last, consider two state 81, 82 S with (81) = ( = on the first step, and we take E A for both states this step and later take the actions following æ We denote the subsequent reward and state distribution over as _ _ .P(()..... and R PPP)) R!! p!3)
respectively, where the superscripts indicate time steps. Therefore, we can deduce that p(U)
p( (1) D (1) p(()) =P?2) n!(2) P n!? , and therefore (=1 2 nd (52, a), Va E A.
p(2)
A.2 COMPARISON WITH BISSIMULAAIION
Similar to 2 -irrllevacee a recently proposed state abstraction r-bisimulation (Castro 2020) is also tied behavioral policy It is beneficial compare the coarseness of 7 -iiireleevance and T- r-bisimulation. For completeness, we restate the definition r-bisimulation.
Definition A.2 (m-bisimulation (Castro 2020 ) Given a policy n, ΦB,a is a r-bisimulation if V81,82 82 E S where (B,r(s1) = 0B,
Σ T(als))
Σ
Σ ((all)))
Σ ((
Σ T(als2)
Σ
P(s'(s2,a),
Vz' E 2,
('E)
('Emml(aa))
where 2 = (S) is the set of abstract states.
However, r-bisimulation does not consider the state-action pair that is not visited under the policy и (c.g., a state-action pair (s, a) when s{ass) = 0), whereas 2 ---irrelevance is defined on all the state-action pairs. Therefore, it is hard to build connection between them unless we also define Z -rrreeleeeacce on the state space (instead of the state-action space) in a similar way.
Definition A.3 (Z -r-eielevvance the state space). Given S, denote (²)ss E. π (a)s) ([ass]Z( (s a). Given a policy n, o is a Z" -irrelevance if V'ss, 82 E where ((81) = ((82), 2"(x)) (x2).
Based the above definitions, the following proposition indicates that such nd -i-irrelevance is coarser than m-bisimulation.
Proposition A.2. Given a policy π and ΦB,a a to be the coarsest r-bisimulation, if 82 8 where (31) ΦB.a (82), then 2 (s1) = (82).
14

---

Published as a conference paper at ICLR 2021
Proof. Starting from $1 and 82 and following the policy x, the reward distribution and the state distribution over Z on each step are identical, which can be proved by induction. Then, we can conclude that 2 (s1) _ zn ($2) and thus -dd (=1))=2
A.3 BOUND ON THE REPRESENTATION ERROR
For completeness, restate Proposition 4.1
Proposition A.3. Given a policy and any Z" -rrrelevaee Φ X + [N], there exists a function Q: [M] + R such that [Q((())) - (x) s Rmux - Raain Vx E X.
K
Proof. Given a policy π and a Z irrelevance Φ, we can construct a Q such that Q()(z)) = (²)x), Vx X in the following way: For all X. one by one, we assign Q(z) t Q (x),
where 2 = ((x). In this way, for any x E X with z = <(x), Q(z) = (x') for some x' E X such
that Z (x²) nx (ir). This implies that (Q()) - Q")x) = (x') (x(x)= Rman Rmin This
also applies the optimal policy n².
B PROOF
Notice that Corollary 4.1.1 is the asymptotic case (n 20) for Theorem 4.1 We first provide sketch proof for Corollary 4.1. which ignores sampling issues and thus more illustrative. Later, we provide the proof for Theorem 4.1 which mainly follows the techniques used Misra et al. (2019).
B. PROOF OF COROLLARY
Recall that Z-learning aims solve the following optimization problem:
min L(o, w: D) == =2))=))))> [(w( ((x2) - v) "EW
(5)
which can also be regarded as finding compound predictor f(:,)) .) w(O()), o()) over the func- tion class FN = [(11, x22) + w(((x1), Φ(x2)) w Wx. @ E 2
For the first step, it is helpful to find out the Bayes optimal predictor f* when size of the dataset is infinite. We notice the fact that the Bayes optimal predictor for a square loss is the conditional mean, ie., given distribution D, the Bayes optimal predictor f arg min (() y) satisfies f (a') E E(ry)>Dly x r'). Using this property, can obtain the Bayes optimal predictor over all the functions {X' x X + [0, 1]} for our contrastive loss:
(2) = E(,4, a 21, x22 = 22]
ER,~Z (x.) (. [b(())) + b(R2)
(6)
-1 (ZF (122)
where we use D denote the distribution from which each tuple in the dataset D is drawn.
To establish the theorem, we require that such an optimal predictor f to be in the function class Fx. Following similar argument Proposition 10 in Misra et al. (2019), it is not hard to show that using > Na.K is sufficient for this realizability condition to hold.
Corollary 4.1.1. The encoder Φ returned by Algorithm / with n + 0o is a Z"-irrelevance, i.e., for any F2 E X, 7 (x1) = (x2) if ((x)) = Φ(x2).
Proof of Corollary 4.1.1 Considering the asymptotic case (i.e., n - oc), we have f where If (((()) = Ö(x2), we have for any
i(:, = (((()), (-)) and and à is returned by Algorithm
PEX X.
1 ndd (x1) T ndd (x) ff (x1 f(z,, = f(2, f=2221=)))
(x2, = 2d (x2) T Z (x).
We obtain 27 (2)) Zz (x²2) by letting 31 472-
15

---

Published as a conference paper at ICLR 2021
B.2 PROOF OF THEOREM 4.1
Theorem 4.1. Given the encoder à returned by Algorithm | the following inequality holds with probability 8 and for any x' Ex:
Bagndaaz (10ó(0)) = 2(x22) 2 2" (x') (2" (21 -21)
8N
- (3+ Inn + 4hn N + 4 Inl 2
where N is the cardinality of encoder function class and is the size of the dataset.
The theorem shows that 1) whenever à maps two state-actions J1, to the same abstraction, Z" (x1) R Zd (122) up to an error proportional to 1/vn (ignoring the logarithm factor), and 2) when the difference of the return distributions (e.g.. 2 (21) 2 (22)) is large, the chance that two state-action pairs (e.g., Ii and X2) are assigned with the same encoding is small. However, since state-action pairs sampled i.idd from the distribution d, the bound holds in average sense instead of the worse-case sense.
The overview of the proof is follows: Note that the theorem builds connection between the opti- mization problem defined in equation 5 and the learned encoder é. To prove the theorem, first find out the minimizer f of the optimization problem when the number of samples is infinite (which was calculated in equation the previous subsection). Then, bound the difference between the learned predictor f and f Lemma B.2 (which utilizes Lemma B.1 when the number of samples finite. At last, utilizing the compound structure of f. we can relate it to the encoder S. Specifically, given æ1, D2 E X with d(x1) = δ(x2). we have f(z1, = f(z2, x'), Vxl E X.
Definition B.1 (Pointwise covering number). Given a function class G : X - R. the pointwise covering number scale e, (N(G,e), e), the size of the set V X R such that Vg G SuPzEx [g(x) v(z)] <.
Lemma B.1. The logarithm of the pointwise covering number of the function class FN satisfies (nA N. S In(?) + xl.
Proof. Recall that FN (x1, + w(o(z1), ó(x2))
WN 2x). where Wy [[]] x [N] + [0, 1]} and VN = fx + [M]]. Given e > O, we discretize [0, 1] to Y (E, fe..). Then, we define WN == [[N] x [N]+ Y} and is easy to see that W is covering [
of W' with |Wxl S N Next, we observe that FN = ((11, + (((((1)) 6(x2)) : w e WN, ò n is covering of FN and [Fv 12 [w] We complete the proof by combining the above results.
TR. and n samples ((x;, == drawn from D, where E X and E 10, The Bayes op- timal function is g" arg mingeg B(za)>-D (g(x) - y) and the empirical risk minimizer is arg mingeg E [(((x.) -v)² With probability at least 1 8 for a fixed 8 E (0, 1).
Proposition B.1 (Proposition 12 Misra et al.
2019)) Consider a function class G
6c (Im(2N(G.e)/0)
((DN))D [(((x) - o (x)) < inf
e>0
Lemma B.2. Given the empirical risk minimizer in Algorithm have
8
In( ))
E E (f(21, - (211 S Ang
+
N2 In |n |n
(8)
n
with probability at least 6.
Proof. This directly follows from the combination Lemma B.I and Proposition B.I by letting
16

---

Published as a conference paper at ICLR 2021
Proof of Theorem 4.1 First, we denote & := [[((x1) = = [(22) and b, == Eawd [[((()) = =]
to be the prior probability over the i-th abstraction. Then, for all x' E A, we have
E a2d
Ma(xa) = i-= ((12))
(Z" (Z"())
2d =z (12))
En aad
[ (f1, f (z2,,
& (11))1, f( x')l 1f (2) 2²) - f(02, 2 xz)) (x), (x)1,(0)
n add
& (11 2') (20
2 nd [0(21) f (x1,2x) f(22,
22vbiAn
The first step obtained using Equation O The second step the triangle inequality. The third step uses the fact that under event &,, we have ((x1) - )(x²) and therefore xx11,²") = f(22, x²). In the fourth step, we drop the dependence on the other variable, and marginalize over Deoup The last step is the Cauchy-Schware inequality.
At last, we complete the proof by summing over all the i E [M] and using the fact Σ vo, S VN and Lemma B.2
C IMPLEMENTATION DETAILS
C.I IMPLEMENTATION FOR ATARI GAMES
Codebase, For Atari games, we use the codehase from htttps://githu/ com Ka xhin Rainbow which an implementation for data-efficient Rainbow (van Hasselt et al., 2019).
Network architecture. To implement our algorithm, we modify the architecture for the value net- work as shown in Figure Lef In data-efficient Rainbow, the state embedding has a dimension of 576. We maintain an action embedding for each action, which is a vector of the same dimension and treated as trainable parameters. Then, we generate the state-action embedding by conducting an element- wesnent-wise product to the state embedding and the action embedding. This state-action embed- ding is shared with the auxiliary task and the main RL task. Afterwards, the value network outputs the return distribution for this state-action pair (noting that Rainbow uses distributional RL algo- rithm, C51 1 (Bellemare et al. 2017)) instead of the return distributions of all actions for the input state is in the original implementation.
Hyperparameters. We use exactly the same hyperparameters as those used in van Hasselt et al. (2199) and CURL (Srinivas et al. (020)) to quantify the gain brought by our auxiliary task and compare with CURL. We refer the readers to their paper for the detailed list of the hyperparameters.
Balancing the auxiliary loss and the main RL loss. Unlike CURL (or other previous work such as JJderberg et al.] (2016); Yarats et al. (2019)) that uses differrent/earned coefficients/learnngg rates for different games balance the auxiliary task and the RL updates, our algorithm uses equal weight and learning rate for both the auxiliary task and the main RL task. This demonstrates the our auxiliary task is robust and does not need careful tuning for these hyperparameters compared with the previous work.
Auxiliary loss. Since the rewards in Atari games are sparse, we divide the segments such that all the state-action pairs within the same segment have the same return. This corresponds to the setting of Z-learning with K + do where the positive sample has exactly the same return with that of the anchor. Then, the auxiliary loss for each update is calculated follows: First, we sample batch of 64 anchor state-action pairs from the prioritized replay memory. Then, for each state-action pair, sample the corresponding positive pair (i.e., the state-action pair within the same segment as the anchor state-action pair) and the corresponding negative pair (randomly selected from the replay memory). The auxiliary loss is calculated these samples with effectively |D] 128.
17

---

Published as a conference paper at ICLR 2021
C.2 IMPLEMENTATION FOR DMCONTROL SUITE
Codebase. We use SAC as the base RL algorithm and build our algorithm on top of the publicly released implementation from CURL (Srinivas al. 2020).
Network architecture. Similarly, we modify the architecture for the critic network SAC. In SAC, the state embedding has dimension of 50. Since the actions are continuous vectors of dimension d in the continuous control tasks of DMControl suite, we directly concatenate the action to the state embedding, resulting in state-action embedding of size 50 + d. Then, the critic network receives the sttt-actioo embedding as the input and outputs the Q value. The actor network receives the state embedding as the input and output the action selection distribution on the corresponding state. Note that, although our auxiliary loss is based on the state-action embedding, the state embedding used by the actor network is also trained by the auxiliary loss through propppropagttion of the gradients.
Hyperparameters. We set the threshold for dividing the segments to 1.0, i.e., when appending transitions to the replay buffer, we start a new segment when the cumulative reward within the last segment exceeds this threshold. The auxiliary loss and the hyperparameters to balance the auxiliary loss and the main RL loss are the same as those used for Atari games. Other hyperparameters we exactly the same those in CURL implementation and we refer the readers to their paper for the details.
D ADDITIONAL EXPERIMENT RESULTS
D.1 MORE SEEDS RESULTS OF REPRESENTATION ANALYSIS
For clearness, we only show the result of representation analysis with single seed in the main text. We add the results for multiple seeds here. The detailed description of analysis task can be found in the first paragraph in Section 5.2.
18

---

Published as a conference paper at ICLR 2021
Bardon
INTIL soei
Raihio spod
INCTL Oregi
Samning Aeun
Bambon
11)
Raribon
(irelyyl (ellty Jppd
PCTIL Aumoli4
(7/3)
Figure 4: Analysis the learned representation on Alien with five seeds. (a) The cosine similar- ity between the representations of the positive/negative state-action pair and the anchor during the training of Rainbow and RCRL. (b) The game scores of the two algorithms during the training.
19

---

Published as a conference paper at ICLR 2021
D.2 HIGH DATA REGIME RESULTS
To empirically study how applicable our model to higher data regimes, we run the experiments the first five Atari games (of Table I for 1.5 millon agent interactions. We show the evaluation results of both our algorithm and the rainbow baseline in T We can see that RCRL outperforms the ERainbow-sa baseline for 4 out of games, which may imply that our auxiliary task has the potential to improve performance in the high-data regime.
Game
ERainbow-sa (100k)
RCRL (100k)
ERainbow-sa (1.SM)
RCRL (1.5M)
Alien
813.8
854.2
1721
1824
Amidar
154.2
157.7
398.8
454.5
Assault
576.2
569.6
572.5
757.9
Asterix
697
799
1370
1306.7
Bank Heist
96
107.2
257.3
550.7
Table 2: Scores of RCRL and ERainbow-sa the first five Atari games.
20


Reviews:
[{"review_id": "Zo7QIaT_Xkd", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "**Summary**\n\nThe authors propose the inclusion of an auxiliary task for training an RL model, where the auxiliary task objective is to learn an abstraction of the state-action space that clusters (s,a) pairs according to their expected return.  The authors first describe a basic abstraction learning framework (Z-learning) followed by the extension to Deep RL as an auxiliary task (RCRL). The authors present results in Atari (discrete action) building on Rainbow, showing an improvement compared to baselines on median HNS in the low-data regime, and results on DMControl (continuous action) building on SAC, showing similar or improved performance compared to baselines.\n\n**Quality**\n\nOverall I found the approach and results to be interesting and moderately compelling. At first glance the improvement is surprising, given that model-free Deep RL already needs to abstract the state space on the basis of returns even without an auxiliary task. The key appears to be the focus on sample efficiency in the low-data regime, where the task seems to improve non-local value signal propagation compared to a bootstrapped algorithm (particularly on Atari, note that in the 100k regime the model-free algorithms have not yet learned to play Pong). Since it is not clear that the algorithm will generalize to more data (it's easy to imagine that the abstraction task will hinder performance when the base algorithms become more finely tuned), I would like to see more clarification of the goal throughout the paper (e.g. \"In the low-data regime, our algorithm outperforms strong baselines on complex tasks in the DeepMind Control suite and Atari games\" in the Abstract), as well as a reference to the focus in the Conclusions.\n\nIn the low-data regime, it's also critical to justify this approach compared to a model-based alternative. On Atari the authors compare to SimPLe, but MuZero would be a stronger baseline.\n\nBesides the empirical results, the authors also nicely provide a description of the Z_\\pi abstraction and an error bound.\n\n**Clarity**\n\nI was confused by the description of the positive/negative sampling procedure in 4.3 paragraph 2. Are segments temporally consecutive within a trajectory? If so, is it primarily a heuristic that they will \"contain state-action pairs with similar returns\" (i.e. couldn't a reward achieved mid-segment make this statement incorrect)? As I understand it, segmenting avoids the problem of determining bins on the return distribution a-priori, however it also seems like it will limit the agent's ability to cluster non-local (s,a) pairs with the same returns. It might also mean that the agent is learning to cluster temporally adjacent states in the underlying state space rather than similar returns.\n\n**Originality**\n\nThe paper builds on existing work in the abstraction literature and auxiliary tasks for deep RL. The primary novel component is using a return-based auxiliary task. \n\nThe Z_\\pi abstraction framework also appears to be novel, although its closely related to existing abstractions like Q_\\pi abstraction.\n\n**Significance**\n\nThe RCRL model itself improves on existing model-free approaches and can be easily incorporated into many model-free architectures, although it seems unlikely to beat a strong model-based baseline like MuZero in the low-data regime. \n\nThe description of the Z_\\pi abstraction, and the exploration of return-based auxiliary tasks in general, could prove more significant in the long term.\n\n**Pros:**\n- The model improves performance in the low-data regime over existing model-free baselines\n- The model can be easily added to many existing architectures\n- Description and theoretical results on a new type of abstraction\n\n**Cons:**\n- The paper needs some more clarity around the focus on low-data / sample efficiency and how applicable the model is to higher data regimes\n- Unclear if the segment-based sampling strategy is clustering (s,a) pairs with similar returns or just states that are nearby in the underlying state space\n- The model seems unlikely to improve on a stronger model-based baseline in the low-data regime\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "WsjdqyLjdl", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "## Return-Based Contrastive Representation Learning for Reinforcement Learning\n### Summary\nThe authors propose Return-based Contrastive Representation Learning (RCRL), a contrastive auxiliary learning task that guides the feature network to encode representation relevant to the task rewards. The experiment results show that RCRL helps improve two commonly used RL algorithm (RAINBOW and SAC) in low data regime. Additionally, RCRL can also be used in combination with other auxiliary tasks to boost performance.\n\nOverall, the paper is well-written, the topic is relevant to the field and the approach is novel.\n\n### Strength\n- Theoretically-backed.\n- The topic of representation learning is pretty relevant to the field now.\n- The learned representation is \"task-relevant\", and therefore can achieve higher performance compared to other representation learning methods.\n\n### Weakness\n- The reliance on environment returns make the approach pretty susceptible to poorly or sparsely defined rewards. Specifically:\n  1. The auxiliary loss does not work in sparse reward environments.\n  2. In the task with dense but deceptive rewards, the representation may be biased toward representation that is not helpful in the long run.\n- The improvement in continuous control tasks seem to be really marginal. Why is that?\n- The learned representation may not be very general due to its reliance on return signals. Certainly, it can help achieve better performance when we only focus on a single task with a well-defined reward function. Yet, the representation may not be as useful when we considered some practical real-world settings that require policy adaptation and transfer learning.", "extra_scores": {}, "extra_text": {}}, {"review_id": "H6ZRjZPWYdP", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "\nSummary:\n\nThe authors present a contrastive auxiliary loss based upon state-action returns.\n\nThey introduce an abstraction over state-action pairs and divide the space of state-action returns into K bins over which the Z function is defined where Z(s,a) is distributed over K-dimensional vectors.  Given an encoding function, phi, and an input x, Z-irrelevance is defined as phi(x_1) = phi(x_2) when Z(x_1) = Z(x_2) which motivates the objective for Z-Learning: to classify state-action pairs with similar returns (within bounds) to be similar.  From this a contrastive loss can be defined (Return-based Contrastive RL, RCRL) where class labels are determined by Z-irrelevance sets encouraging state-action encodings to be similar when the returns are.  In the limit Z becomes the RL state-action value function Q.\n\nThe authors evaluate their approach on Atari (discrete actions) and the DeepMind Control Suite (continuous actions) across both model-free and model based RL algorithms against and in combination with other auxilliary losses including CURL (Srinivas et al. 2020). \n\nStrengths & Weaknesses:\n\nAuxiliary losses have become an important component in RL for developing stable agents that can generalize well and form good representations.  In particular, contrastive losses have come into increasing use with growing literature around these methods and so I believe the domain area of this paper is relevant and of interest.  The authors do a good job of covering the recent developments of background literature in their related work section and grounding their approach with recent efforts undertaken in RL auxiliary losses, contrastive learning approaches and state abstraction/representation learning literature.\n\nThe approach is overall novel as many contrastive learning methods are defined against input data or downstream representations, whereas this work derives it's data from RL returns and creates a link between the representational landscape of the observations and actions and broad outcomes as they are valuable to an agent.  As the author's have framed the problem, I believe this approach is more powerful and also more tractable than something like reward prediction.  Intuitively the formulation seems solid to me since we often would like to understand not only when we're in a good state and taking a useful action but also, in general, what kind of properties state-action pairs with similar returns should have.  The authors do note that this may be learnable by temporal difference updates alone however, this approach aims to directly encourage the learning of this relationship and decouple it from the RL algorithm (where perhaps other things may be focused on such as planning etc.). \n\nOne shortfall of this approach could be the available data itself as you'd rely on the policy to provide you with good samples for RCRL.  The authors indicate that they segment trajectories to ensure better quality positive and negative samples for learning however, it could be made clearer how much of a problem this can be.    This approach could possibly be combined with a self-supervised approach to alleviate these types of concerns.   It would also be nice to know the additional computational burden of RCRL and how this compares to other auxiliary losses.\n\nThe experiments on Atari & Control look solid and demonstrate that this method attained a stronger score both alone and when combined with CURL and good top performance on DeepMind control suite tasks when compared again to CURL and pixelSAC.  It might have been nice to see more comparisons or combinations with other contrastive methods that have had some success in  learning visual representations (SimCLR: Chen et al. 2020, BYOL: Grill et al. 2020).  The similarity analysis also provided some nice insight into the inductive bias induced by RCRL.\n\nOverall, the paper is well written and has a clear layout.  The authors provide clear algorithms and figures and the content flows well from section to section.\n\n\nRecommendation:\n\nI believe that this is a promising and very active area of research and that this work makes the case for a solid new approach and a set of encouraging results to back it up. ", "extra_scores": {}, "extra_text": {}}, {"review_id": "1xdBB97REc_", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This is an interesting paper that proposes abstractions based on return distribution similarities to be used as auxiliary tasks to aid in learning. The idea is quite promising and I think could open up new avenues for future research, but it does not appear to me to be ready for publication yet. In particular, although the authors claim to be distinguishing based on returns, many of the design decisions seem to implicitly assume determinism, or near-determinism (in particular, see point 3 and 10 below).\nThe theoretical results seem interesting, but I have some questions on their validity and clarity (see point 7 and 8 below).\nThe practical algorithm has a few issues that need clarification (see points 3, 4, 9, 11, and 12 below).\n\n1. At the bottom of page 2, the authors write \"we are the first to leverage return to construct a contrastive auxiliary task for speeding up the main RL task.\" This is not quite true, see [Zhang, A.; McAllister, R.; Calandra, R.; Gal, Y.; and Levine, S. 2020. Learning Invariant Representations for Reinforcement Learning without Reconstruction\". arXiv preprint arXiv:2006.10742 .].\n2. In section 2.2 you should cite [Taylor, J.; Precup, D.; and Panagaden, P. 2009. Bounding performance loss in approximate MDP homomorphisms. In Advances in Neural Information Processing Systems, 1649–1656.]\n3. In line 6 of Algorithm 1, $y = \\mathbb{I}[b(R_1) \\ne b(R_2)]$ seems problematic for stochastic returns. In particular, if the number of bins is very large and the returns have wide variance, $y$ will almost always be zero.\n4. In line 8 of Algorithm 1 is $\\hat{w}$ also learned?\n5. In the last sentence of the first paragraph of section 4.1, rather than comparing against regular bisimulation it seems more appropriate to compare to $\\pi$-bisimulation from [Castro, P. S. 2020. Scalable methods for computing state similarity in deterministic Markov Decision Processes. In Proceedings of the AAAI Conference on Artificial Intelligence.].\n6. In equation (1), the minimization appears over $f$, but $f$ doesn't appear on the RHS. Shouldn't the minimization be over $w$?\n7. In Theorem 4.1 it's not clear what role $x_1$ plays in the result. Why is this third state necessary? It does not seem to show up in the proof in the appendix. Further, the proof in the appendix could do with some elaboration, as it's not completely clear how lemmas B.1 and B.2 result in the proof of Theorem 4.1. It would be better if the authors restate the theorem statement in the appendix and were more explicit about the connections. There are no page limits for the appendix.\n8. How do we get a sense for how big $|\\Phi_N|$ is? Couldn't it be as large as $N^{\\mathcal{|X|}}$?\n9. Typically auxiliary losses are combined with the main loss into a single loss, but in Algorithm 2 you seem to be updating them sequentially. Why? Does the order of update matter?\n10. In Figure 1, on the left, the purple rectangle says \"Value net\". If so, are you really learning a distribution?\n11. In Figure 2, why do some algorithms only have triangles and not learning curves? Also, it seems from that figure that \"State SAC (skyline)\" seems to outperform all others, including RCRL.\n12. Is Figure 3 over a single seed?\n13. In Definition A.1, bisimulation was _not_ introduced in (Jiang, 2018), it was introduced in [Givan, R., Dean, T., & Greig, M. (2003). Equivalence notions and model minimization in markov decision processes. Artificial Intelligence, 147, 163–223].\n\n\nMinor comments:\n1. It would be helpful if the authors specify how to pronounce RCRL. While reading the paper I was pronouncing it [like this](https://youtu.be/dQw4w9WgXcQ).\n2. In the second-to-last sentence on page 1, should read \"while ignor**ing** return-irrelevant features.\"", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: RCRL is return-based contrastive learning for reinforcement learning, where the label is whether two samples belong to the same return bin. The reviewers found this to be a well executed paper with good theoretical and experimental results.
Decision: Accept (Poster)