Learning Generalizable Representations for Reinforcement Learning via Adaptive Meta-learner of Behavioral Similarities
Authors: Jianda Chen
Sinno Pan
Venue: ICLR 2022 Poster (2022)
Keywords: deep reinforcement learning
deep learning
representation learning

Abstract: How to learn an effective reinforcement learning-based model for control tasks from high-level visual observations is a practical and challenging problem. A key to solving this problem is to learn low-dimensional state representations from observations, from which an effective policy can be learned. In order to boost the learning of state encoding, recent works are focused on capturing behavioral similarities between state representations or applying data augmentation on visual observations. In this paper, we propose a novel meta-learner-based framework for representation learning regarding behavioral similarities for reinforcement learning. Specifically, our framework encodes the high-dimensional observations into two decomposed embeddings regarding reward and dynamics in a Markov Decision Process (MDP). A pair of meta-learners are developed, one of which quantifies the reward similarity and the other quantifies dynamics similarity over the correspondingly decomposed embeddings. The meta-learners are self-learned to update the state embeddings by approximating two disjoint terms in on-policy bisimulation metric. To incorporate the reward and dynamics terms, we further develop a strategy to adaptively balance their impacts based on different tasks or environments. We empirically demonstrate that our proposed framework outperforms state-of-the-art baselines on several benchmarks, including conventional DM Control Suite, Distracting DM Control Suite and a self-driving task CARLA.

Published as a conference paper at ICLR 2022
LEARNING GENERALIZABLE REPRESENTATIONS FOR REINFORCEMENT LEARNING VIA ADAPTIVE META- LEARNER OF BEHAVIORAL SIMILARITIES
Jianda Chen & Sinno Jialin Pan
Nanyang Technological University. Singapore
i jinda0011ee nt utu.u edu maa.opanentu.edu..
ABSTRACT
How to learn an effective reinforcement learning-based model for control tasks from high-level visual observations is practical and challenging problem. A key solving this problem to learn oow-dimesiional state representations from observa- tions, from which an effective policy can be learned. In order to boost the learning of state encoding, recent works are focused on capturing behavioral similarities between state representations or applying data augmentation visual observations. In this paper, we propose novel meta-learner-baat learning regarding behavioral similarities for reinforcement learning. Specifically, our framework encodes the high-dimenn -igh-dimensional observations into two decomposed embeddings regarding reward and dynamics in a Markov Decision Process (MDP). A pair of meta- meta-learners are developed, one of which quantifies the reward similarity and the other quantifies dynamics similarity over the correspondingly decomposed embeddings. The meta-learners are self-learned to update the state embeddings by approximating two disjoint terms in on-policy bisimulation metric. To incorporate the reward and dynamics terms, further develop a strategy adaptively balance their impacts based on different tasks or environments. We empirically demonstrate that proposed framework outperforms state-of-the-art baselines several benchmarks, including conventional DM Control Suite, Distracting DM Control Suite and self-driving task CARLA.
framework for representation
INTRODUCTION
Designing effective reinforcement learning algorithms for learning to control from nigh-dimensional visual observations is crucial and has attracted more and more attention Yaratt et al. [202 Laskin ct al. 2020a Schwarzer et al. 202T Lesort et al. 2018). To learn a policy efficiently from high- dimensional observations, prior approaches first learn encoder map dimensionl obser- vations, c.g., images, to low-dimensional representations, and subsequently train a policy from low-dimensional representations to actions based various RI. algorithms. Therefore, how to learn low-dimensional representations, which are able to provide semantic abstraction for high-dimensional observations, plays key role.
Early works on deep reinforcement learning train encoders based on various reconstruction losses (Lange & Riedmiller 2010; Watter et al. 2015 Wahlström et al. 2015 Higgins al. 2017), which aim to enforce the learned low-dimensional representations to reconstruct the original high-dimensional observations after decoding. Promising results have been achieved in some ap- plication domains, such as playing video games, simulated tasks, etc. However, the policy learned on state representations trained with reconstruction loss may not generalize well complex en- vironments, which are even though semantically similar to the source environment. The reason that a reconstruction loss is computed over all the pixels, which results in all the details of the high-dimensional images tending be preserved in the low-dimensional representations. However, some of the observed details, such as complex background objects in an image, are task-irrelevant Encoding such details in representation learning makes the
and highly anviromment-depender learned representation less effective for downstream reinforcement learning task of interest and less generalizable to environments. To address the aforementioned issue, some data augmentation

---

Published as a conference paper at ICLR 2022
techniques have been proposed to make the learned representations more robust (Yarats et al. Lee et al., 2020b Laskin et al. 2020b). However, these approaches rarely consider the properties of a Markov Decision Process (MDP), such as conditional transition probabilities between states given an action, in the representation learning procedure.
202
Recent research (Zhang al. 2021 Agarwal et al. 2021 Castro 2020
Ferns et al. 2004) has
shown that the bisimulation metric and its variants are potentially effective to be exploited to learn more generalizable reinforcement learning agent across semantically similar environments. The bisimulation metric measures the "behavioral similarity" between two states based on two terms: 1) reward difference that considers the difference in immediate task-specific reward signals between two states, and 2) dynamics difference that considers the similarity of the long- tengttrmm behaviors between two states. A general idea of these approaches is to learn an encoder to map observations to a latent space such that the distance or similarity between two states in the latent space approximates their bisimulation metric. In this way, the learned representations (in the latent space) are task-relevant and invariant to environmental details. However, manually specifying form of distance, c.g., the L norm as used in (Zhng et al. 202T), in the latent space may limit the approximation precision for the bisimulation metric and potentially discard some state information that is useful for policy learning Moreover, existing approaches rarely explore how to learn an adaptive combination of reward and dynamics differences in the bisimulation metric, which may vary in different tasks environments.
We propose a novel framework for learning generalizable state representations for RI., named Adap- tive Meta-leamner of Behavioral Similarities (AMBS). In this framework, we design network with two encoders that map the nigh-dimmessionl observations to two decomposed representations regard- ing rewards and dynamies. For the purpose of learning behavioral similarity on state representations, we introduce a pair of meta-learners that learn similarities in order to measure the reward and the dynamics similarity between two states over the corresponding decomposed state representations, respectively. The meta-learners are self-learned by approximating the reward difference and the dynamies difference in the bisimulation metric. Then the meta-learners update the state representa- tions according to their behavioral distance to the other state representations. Previous approaches with a hand-crafted form of distance/similarity evaluating state encoding in the L1 space are diffi- cult to minimize the approximation error for the bisimulation metric, which may lead to important side information being discarded, e.g., information regarding to Q-value but not relevant to states distance/similarrty. Instead, our learned similarities measure two states representations via neural ar- chitecture, where side information can be preserved for policy learning. Our experiments also showed that a smaller approximation loss for similarity learning can be obtained by using the meta-learners. This demonstrates that the proposed meta-learners can overcome the approximation precision issue introduced by the L1 distance in previous approaches and provide more stable gradients for robust learning of state representations for deep RL. Moreover, we explore the impact between the reward and the dynamics terms in the bisimulation metric. We propose learning- based adaptive strategy to balance the effect between reward and dynamics in different tasks or environments by introducing learnable importance parameter, which is jointly learned with the state-action value function. Finally, we use a simple but effective data augmentation strategy to accelerate the RL procedure and learn more robust state representations.
The main contributions of our work are 3-fold: 1) we propose meta-learner-based framework to learn task-relevant and architecture that decomposes each state into two different types of representations for measuring the similarities in terms of reward and dynamics, respectively, and design a learnable adaptive strategy to balance them to estimate the bisimulation metric between states; 3) we verify our proposed framework on extensive experiments and demonstrate new state-of-the-art results on background-
state representations; 2) we propose network
distraction DeepMind control suite (Tassa al. 2018 Zhang al., 2018 Stone et al. 202 1 and other visual- based RL tasks.
2 RELATED WORK
Representation learning for reinforcement learning from pixels Various existing deep RL methods have been proposed to address the sample-effcciency and the generalization problems for conventional RL from pixel observation. In end-to-end deep RL, neural networks learn representations implicitly by optimizing some RI. objective (Mnih al. 2015 Espeholt et al., 2018). Watter et al.
2

---

Published as a conference paper at ICLR 2022
(2015) and Wahlstrom et al. (2015) proposed to learn an encoder with training dynamics model jointly to produce observation representations. Lee et al. (2020a). Hafner et al. 2211)) Hafner et al. (2020) and Zhangg et al. (2119) aimed to learn environment dynamics model with the reconstruction loss to compact pixels to latent representations. Gelada et al. 2019) proposed to learn representations by predicting the dynamies model along with the reward, and analyzed its theoretical connection to the bisimulation metric. Agarwal et al. (2021) proposed to learn representations by state distances based on policy distribution. Kemertas & Soimenttdo-Armstroonn bisimulation metric with intrinsic rewards and inverse dynamics regularization. Castro et al. 20021) replaced the Wasserstein distance in bisimulation metric with a parametrized metric.
(202I) modified the learning of
Data Augmentation in reinforcement learning Laskin et al. (2020b) and Ye al. (2020b) explored various data augmentation techniques for deep RI., which are limited to transformations on input images. Cobbe et al. (2019) applied data augmentation to domain transfer. Ye et al. (2020a) studied data augmentation on game environments for zero-shot generalization. RandConv (Lee et al. 2020b) proposed randomized convolutional neural network generate randomized observations in order to perform data augmentation. A recent work DrQ (Yarats et al. 202I) performs random crop on image observations and provides regularized formulation for updating Q-function. RRaleanu et al. (2021) proposed to use data augmentation for regularizing on-policy actor-critic RI. objectivs Piiii et al. (2020) proposed counterfactual data augmentation technique by swapping observed trajectory pairs.
PRELIMINARIES
We define an environment as Markow Decision Process (MDP) described by tuple M (S, AA, P, R, 7), where S is the high-dimensional state space (e.g. images), A is the action space, P(s' ss. a) is the transition dynamics model that captures the probability of transitioning to next state s' E given current state 8 S and action a E A. R is the reward function yielding reward signal r R(s, a) E R, and 7 [0, 1) is the discounting factor. The goal of rein- forcement learning is to learn a policy т T[ass) that maximizes the expected cumulative rewards: max, 7' partial observability, and use stacked consecutive frames as the fully observed states.
7'r(St, a.) a:)aa; (-]ss),, Sy ~ P(s'|s. (s. a)]. In the scope of this paper, we do not consider
Soft Actor-Critic (SAC) is widely used off-policy model-free reinforcement learning al- gorithm (Haarnoja et al., 2018). SAC aims to maximize the entropy objective Ziebart 2010) which is the reinforcement learning objective augmented with an entropy term J(T) E, E [r(s,, a.) + a7 (T[-|s; In order to maximize the objective, SAC learns a state-value func- tion (s, a), stochastic policy Ty (a(s) and the temperature a, where and 1a are the parameters of Qo and Ty. respectively. The Q-function is trained by minimizing the squared soft Bellman residual
(0) a
(Qo(s,, a.) (r(s,, ae) + y IPP a..).) (V (+++)))))
(1)
where D is the dataset or replay buffer storing the transitions, and Vj is the value function parameter- ized by 0. The parameters zo of policy is learned by maximizing the following objective
Ja (16) EAa-D [LaI~T [ar log a (als)) - Qo(s,, a:]]]
The Bisimulation Metric defines pseudometric d behavioral similarity of two discrete states (Ferns et a (004). An extension to both continuous and discrete state spaces has also been developed (Ferns et al.) 2011). A variant of the bisimulation metric proposed in Castro (2020) defines metric w.r.t. a policy ग,, which is known as the on-policy bisimulation metrie. It removes the requirement of matching actions in the dynamics model but focuses on the policy x, which is able to better capture behavioral similarity for a specific task. Because of this property, in this paper, we focus on the π r-bisimulation metric (Castro) 2020):
S R where d quantifies the
FT (d)(s;, S5) (1 - c) - RT, + cWi PP)))
The mapping F met + met, where met denotes the set of all pseudometrics in state space S and Wi (i(d)(--) denotes the 1- Wasserstein distance given the pseudometric d. Then fer has least fixed point denoted by d. Deep
where R5 := E. (()s)) a) and P. : (aals))(-iss,)))
If the pseudometric d of states 0, then the two states helong to an equivalence relation.
3

---

Published as a conference paper at ICLR 2022
Dynamice Model
(((p)(())
(8)(5)
Encoder
(s)(s)
e
(s,)
(1
[S
Enooder
(x.(88
f-
[rs - 7yl
S
W
(8,
8j
Eneoder
Φ
(8,(8))
Acto
* (-1a))))
Dynamies Nodel P(-(Au(sy). a.)
Figure I: Architecture of our AMBS framework. The dotted arrow represents the regression target and the dash arrow means stop gradient. Left: the learning process of meta-learner. Right: the model architecture for SAC with adaptive weight which is jointly learned with SAC objective.
Bisimulation for Control (DBC) (Zhang al. 2021 is learn latent representations such that the L distances in the latent space are equal to the x r-bisimulation metric in the state space:
arg min (1(b():) - *(s,))l - |R_, R, -W2 P pT)))
where @ is the state encoder, and (-) denotes the 2-Wasserstein distance. DBC is combined with the reinforcement learning algorithm SAC, where d(s) is the input for SAC.
ADAPTIVE META-LEARNER OF BEHAVIORAL SIMILARITIES
In this section, we propose framework named Adaptive Meta-learner of Behavioral Similarities (AMBS) to learn generalizable states representation regarding the т n-bisimulation metric. The learning procedure is demonstrated in Figure [I Observe that the r-bisimulation metric is composed of two terms: [R_ Rgy, which computes the difference of rewards between states, and (2(d))Ps, 8.) )s,)) which computes the difference of the outputs of dynamics model between states. We propose a network architecture which contains encoders to transform the digh-dimensiona visual observations to two decomposed encodings regarding rewards and dynamics. We develop two meta-learners, of which quantifies the reward difference dynamics distance (Section 4. Each meta-learner is self-learned to update the corresponding state representations by learning to approximate a term in the r-bisimulation metric, respectively. Rather than enforcing the L1 distance between embedded states to be equal to the r-bisimulation metric, our meta- meta-learners are able to use more flexible form of similarity, i.e., well-designed non-linear neural network, with each similarity evaluates the reward difference or dynamics difference between states beyond the original Euclidean space. Moreover, we propose a strategy for learning to combine the outputs of the two learned similarities in a specific environment (Section 4.2) We introduce learnable weight for the combination and such weight adaptively learned together with the policy learning procedure (in this work use SAC the base reinforcement learning algorithm). In addition, we also use a data augmentation technique to make the learned representations more robust (Section 4.3). The whole learning procedure is trained in end-to-end manner.
reward representations and the other captures
4.1 META-LEARNERS FOR DECOMPOSED REPRESENTATIONS
We design a network architecture with two encoders, d, and Ca, to encode decomposed features from visual observations as shown in the left side of Figure Each encoder maps high-dimensionn observation to owwdimensiooal representation, i.c., Φ, S + Z, capturing wwward-relevan features features. We design two meta-learners, fr and fa-
and du : + Za capturing dynamics-relevan where fr learns to measure reward difference and fa aims to measure dynamics difference between two state representations. Specifically, each meta-learner takes two embedded states as input and outputs the corresponding similarity. The procedure is summarized as follows.
S; S oh (s:), (s,) - (s:), (s;)), where E (r,
Note that the aggregation of the outputs of the two encoders or and due of each observation (i.e., state)is also fed into SAC to learn a policy (the right side of Figure []. Each self-learned meta- learner leams capture similarity by approximating distance term the r-bisimulation metric, respectively, and provides stable gradients for updating o, or de according to the distance to the other state representations. The details of network architecture be found in Appendix B.2
4

---

Published as a conference paper at ICLR 2022
As discussed in Section restricting the form of distance be the L1/Lz norm in the latent space may limit the approximation precision. As shown in Figure 2 the L1 and the Lg norm for measuring the distance of two latent representations ((s;) and o(s,) lead to large regression losses during RL training, which destabilize the representation learning and consequently decrease final performance. Besides, such hand-crafted distances make semantically similar states encoding close to each other in terms of the L1 norm in Euclidean space, which however may lose part of useful information, e.g., policy-relevant information but immediately regarding the distance.
To overcome the aforementioned limitations, we propose to exploit meta-learners f. and fa to learn similarities. By learning with regression target, the similarity learned by f. is easier converge to the target and leads to faster descent tendency of regression loss (as shown in Figure 2 where y-axis log-scale). Such a property provides more stable gradients for state representation comparing to the Ln/L.2 norm, and therefore the meta-learner f. is able to guide the process of updating the state encoder. Besides, f. a non-linear transformation that evaluates the state similarity more flexible space instead of the original Euclidean space. Consequently, it is able to preserve more task-relevant information in state representation that required for further SAC policy learning.
As the goal of meta-learners is approximate different types of simi- larities in the x-bisimulation metric, we design the loss functions for the meta-learners using the mean squared error:
cortooli sne ngup, regressin bos
E(rr
(s)( (r(s,) IR -R,, l)
(4)
101
{(Sa, da) = (Ja(Ou(Si), oa(s)) - W2(Ps, s. p.,
(5)
To make the resultant optimization problems compatible with SGD up- dates with transitions sampled from replay buffer, we replace the reward metric term |R_, R2,1 by the distance of rewards in the sampled transi-
Figure 2:
The regres- sion loss among different forms of distances.
tions in (4). and use a learned parametric dynamics model P to approximate the true dynamics model PT in (5) We also use 2-Wasserstein distance W2 in (5) (as in DBC) because it has closed form for Gaussian distributions. Denote by (s., a,, rs. s)) and (s). a,. Ty. s',) two transitions sampled from the replay buffer. The loss functions (4) and (5) can be revised as follows,
((Fr, = (s. ((fa.0a) = (Fllda(s)) (2((P(lou(s)) a.), (()l)))))) a.)))
Ir: - ry
(6)
where P(.l@a(s.), a.) is a learned probabilistic dynamics model, which is implemented as a neural network that takes the dynamics representation (s(s.) and an action a. as input, and predicts the distribution of dynamics representation of next state s'.- The details of P can be found Appendix
4.2 BALANCING IMPACT OF REWARD AND DYNAMICS
The π r-bisimulation metric 3 is defined as a linear combination of the reward and the dynamics difference. Rather than considering the combination weight hyper- apper-paramete, which needs be tuned in advance, we introduce a learnable parameter e E (0, 1) to adaptively balance the impact of the reward and the dynamics in different environments and tasks.
As the impact factor should be automatically determined when learning in a specific task or environment, we integrate the learning of e into the update of the Q-function SAC such that the value of is learned from the state-action value which is highly relevant to specific task and environment. To specific, we concatenate the low-dimensional representations (s)(s) and (a(s) weighted by - and e, where 1 - and output of softmax to ensure c (0, 1). The loss for the Q-function is revised as follows, which takes or (s), (s(s), weight c and action a as input:
JQ Dr. Φa)= = E(s.a..''~D
(Qo(Or (s(s), c, a, (r(s, a) + +V((s)))
(8)
where Qu is the Q-fuccioon parameterized by 0 and V--) is the target value function.
Moreover, to balance the learning of the two meta- neta-learners f. and fa. we jointly minimize the regression losses of fr and ʃa (66) and (7). respectively) with the balance factor Thus, the loss
5

---

Published as a conference paper at ICLR 2022
function is modified follows,
(()) = (1 - c))(s, + cl(fa, where θ = Ja, Dr. 02).
Note that c is already learned along with the Q-function, we stop gradient of when minimizing the loss (9). This is because we perform gradient descent w.r.t may only capture which loss (e(frs Φr) or [fa. 0a)) is larger/smaller and fail to learn the quantities of their importance. We also stop gradient for the dynamics model P since P only predicts one-step dynamics.
4.3 OVERALL OBJECTIVE WITH DATA AUGMENTATION
In this section, we propose to integrate data augmentation strategy into our proposed framework to learn more robust state representations. We follow the notations of an existing work on data augmentation for reinforcement learning, DrQ (Yarats et al. (021), and define state transformation x T + S that maps state to a data-augmented state, where T is the space of parameters of h. In practice, possible crop positions and v T is one crop position drawn from T. We apply DrQ to regularize the objective of the Q-function (see Appendix for the full objective of Q-function).
random crop the transformation h in this work. Then T contains all
Besides, by using the augmentation transformation h, we rewrite the loss of representation learning and similarity approximation in (9) as follows,
(()) -0 - (f. (.) - -Ir: -Ini-
).
+ fa ( )e(ss)))) W/ (P(-lo((!")), P(Iou(()) ),))
(10)
(s(()))) + (1 c) fr 2 (r(ss) ())
- -Ir -1r--10
ry1) -m:-ry))
+ + (f (d(s)) (f
)ols! W2 (P(loo(s(")),
a.), a.), P()1o(()))) ),)))
where s*x) Mss,, vh)) is the transformed
Algorithm AMBS+SAC
state with parameters (1)) Specifically, s,
1: Input: Replay Butfer D. initialized e (Sr, Sa, da) netwoek Qo. actor pig. larget Q network Qū.
s(")) s!(²) and sf²) are transformed states with
2: Sample a batch with size B: [(s,, as, Ti, T 10 D.
parameters of (1) (1)) v(2) and (2) respec-
3: Shuflle batch [(s,, au, my v()) to [(5), my.
tively, which are all drawn from o indepen-
4: Update Q network by g with DrQ augmentation.
dently. The first two terms (TOp are similar
5: update actor network by
6: update alpha
to the two terms in (9, while the observation
7: update encoder by 10
s. in (9) is transformed s!"). For the last two terms in IO). the observation is transformed by different parameter v! except for the last term, where the parameter !)) does not change. The reason why the parameter v! used in the last term that we expect the meta-learner sa is able to make a consensus prediction on s and s. as they are transformed from same state s., Another major difference between the first two terms and the last two terms the order of the subscripts and j: is followed by in first two terms while is followed by i in last two terms. The reason behind this is that we aim to make fr and fa to be symmetric.
8: update dynamics model
The loss of the actor of SAC, which is based on the output of the encoder, is
J (+) (s.))=D aa lolog(T (a)(())(()((s)) (alō(s))) Qo (d))s))
(1I)
where å denotes the convolutional layers of ó and (o(((s), a) is a convenient form of the Q-function with c-weighted state representations input. We also stop gradients for & and & (Φ, and Φa) regarding the actor loss. In other words, the loss in (IT) only calculates the gradients w.r.t. tb.
The overall learning algorithm is summarized in Algorithm Note that due limit of space, more details can be found in Appendix[A) A In summary, as illustrated in Figure in AMBS, the convnet is shared by the encoders o, and da, which consists of 4 convolutional layers. The final output of

---

Published as a conference paper at ICLR 2022
the convnet is fed into two branches: 1) one is cully-connected layer to form reward representation (s),,, and 2) the other is also a fully-connected layer but forms dynamics representation (a(s). The meta-learners fr and fa are both MLPs with two layers. More implementation details provided at Appendix B Source code is available at wtttttts///// ptteps:://gttii wtg.tthuubb.c cindd
5 EXPERIMENTS
The major goal of our proposed AMBS is to learn a generalizable representation for RL from high-dimensional raw data. We evaluate AMBS on environments: 1) the DeepMind Control (DMC) suite (Tassa et al., 2018) with background distraction; and 2) autonomous driving task on CARLA (Dosovitskiy et al. 2011) simulator. Several baselines methods are compared with AMBS: 1)
Deep Bisimulation Control (DBC) (Zhang et al., (02T) which is recent research on RL representation learning by using the L, distance to approximate bisimulation metric; 2) PSEs (Agarwal al. (021) which proposes new state-pair metric called policy similarity metric for representation learning and learns state embedding by incorporating a contrastive learning method SimCLR (Chen et al. 2020); ;
3) DrQ Yarats al., 2021), state-of-th offthe-artt art data augmentation method for deep RL; 4) Stochastic Latent Actor-Critic (SLAC) (Lee et al. 2020a), state-of-the--ee approach for partial observability on controlling by learning sequential latent model with reconstruction loss; 5) Soft Actor actor-Critic (SAC) (Haarnoja et al.] (2018), a commonly used off-policy deep RL method, upon which above baselines and our approach are built.
5.1 DMC SUITE WITH BACKGROUND DISTRACTION
DeepMind Control (DMC) suite (Tassa et al. 2018) is environment that provides high dimensional pixel observations for RL tasks. DMC suite with background distraction (Zhang et al., 2018) Stone et al.) 2021) replaces the background with natural videos which are task-irrelevant the RL tasks. We perform the comparison experiments in two settings: original background and nature video background. For each setting evaluate AMBS and baselines in 4 environments, cartpole- swingup, finger-spin, cheetah-run and walker-walk, Results of additional environments are shown in Appendix D
(b)
Figure 3:
(a) Pixel observations of DMC suite with original background. (b) Pixel observations of DMC suite with natural video background. Videos sampled from Kinetics (Kay et al. 2017).
Original Background. Figure 3a shows the DMC observation of original background which static and clean. The experiment result is shown in Figure Our method AMBS has comparable learning efficiency and converge scores comparing to state-of of-the-rrt pixel-level RL methods DrQ and SLAC. Note that DBC performs worse in original background setting.
300
100
200
DO
05
05
0G
Figure 4: Training curves of AMBS and comparison methods. Each curve is average 5 runs.
7

---

Published as a conference paper at ICLR 2022
certpole awtrgup
Tonger spen
Cheedan P20
Sour waX
430
100
200
DBC
DMC
S6C
300
54C
boonert staxs
Training curves of AMBS and comparison methods on natural video background setting. Videos are sampled from Kinetics Kay et al. 2017) dataset. Each curve is average on runs.
Figure 5:
Natural Video Background. Figure[3bjshowe [3bsshww the DMC observation with natural video background. The background video is sampled from the Kinetics (Kay et al., 2017) dataset under label "driving car". By following the experimental configuration of DBC, we sample 1,000 continuous frames as background video for training, and sample another 1,000 continuous frames for evaluation. The experimental result is shown in Figure [5] Our method AMBS outperforms other baselines in terms of both efficiency and scores. It converges to the scores that are similar the scores on original background setting within million steps. DrQ can only converge to the same scores on cartpole- swingup and finger-spin but the learning is slightly slower than AMBS, The sequential latent model method SLAC performs badly on this setting. This experiment demonstrates that AMBS is robust to the background task-irrelevant information and the AMBS objectives improve the performance when comparing to other RI. representation learning methods, e.g., DBC.
5.2 ABLATION STUDY
Our approach AMBS consists of several major components: meta- meta-leamerr, factor to balance the impact between reward and dynamics, and data augmentation on representation learning. To evaluate the importance of each component, We conduct experiments on cheetah-run and walker-walk under the natural video background setting. Figure 6 shows the performance of AMBS, variants of AMBS and baseline methods. AMBS w/o data aug is constructed by removing data augmentation in AMBS. AMBS w/o f is referred to as replacing the meta-learners fr and fa by the L1 distance, as in DBC. AMBS w/o e removes the component of balancing impact between reward and dynamics but encodes the observation into a single embedding.
eliminate replace each component by baseline methods.
uses single meta-learner to approximate the sum of reward and dynamics distance. DBC DrQ is DBC with data augmentation applied on representation learning and SAC. AMBS = 0.5 denotes AMBS with fixed weight 0.5. AMBS w/ (1,7) replaces e in AMBS by freezing the reward weight to 1 and the dynamics weight to y. Such a setting of weights is used in DBC. Figure [ demonstrates that AMBS performs better than any of its variant. Comparing to DBC that does not utilize data augmentation, better. This comparison shows that using meta-learners fr and fa to learn reward- and dynamics- relevant representations is able to improve RL performance compared with using the Lı distance to approximate the whole r-bisimulation metric.
variant AMBS w/o data aug still performs
5.3 TRANSFER OVER REWARD FUNCTIONS
The environments walker-walk, walker-run and walker-stand share the same dynamics but have different reward functions according to the moving speed. To evaluate the transfer ability of AMBS over reward functions, we transfer the learned model from walker-walk to walker-run and walker- stand. We train agents on walker-run and walker-stam stand with frozen convolutional layers of encoders that well trained in walker-walk. The experiments are done under the natural video setting
Figure 6: Ablation Study cheetah-run (left) Figure 7: Transfer from walk to (left)
and walker-walk (right) in natural video setting. walker-run and (right) walker-stand.
8

---

Published as a conference paper at ICLR 2022
compared with DBC as shown in Figure 7 It shows that the transferring encoder converges faster than training from scratch. Besides our approach has a better transfer ability than DBC.
5.4 GENERALIZATION OVER BACKGROUND VIDEO
We evaluate the generalization ability of our method when background video is changed. We follow the experiment setting of PSE (Agarwal et al. 202 n,. We pick videos, 'bear' and "bmx-bumps" from the training set of DAVIS 2017 (Pont-TTuee one video and a frame at each episode start, and then play the video forward and backward until episode ends. The RL agent is evaluated on validation environment where the background video is sampled from validation set of DAVIS 2017. Those validation videos are unseen during the training. The evaluation scores at 500K environment interaction steps is shown in Table T. The comparison method DrQ+PSEs is data augmented by DrQ. AMBS outperforms DrQ+PSEs all the tasks. It demonstrates that AMBS is generalizable on unseen background videos.
al., 2018) for training. We randomly sample
Methods
C-swingup
F-spin
C-run
W-walk
DrQ PSEs
749 + 19
779 + 49 308 + 12 789 = 28
AMBS
807 1
933 96 332 27 893 885
Table 1: Generalization unseen background videos sampled from DAVIS 2017 (Pont-Tuset et al. 2018 Scores of DrQ+PSEs are reported from Agarwal et al. (00TT)
(a)
(h)
Figure 8: (a) Illustration of a third- serrdpperson view in "Town4" scenario of CARLA. (b) first-person observation for RL agent. It concatenates five cameras for 300 degrees view.
5.5 AUTONOMOUS DRIVING TASK ON CARLA
CARLA is an autonomous driving simulator that provides 3D simulation for realistic on-road scenario. In real world cases or in realistic simulation, the learning of RL agents may suffer from complex background that contains task-irrelevant details. To argue that AMBS can address this issue, we perform exper- iments with CARLA. We follow the setting of DBC to create an autonomous driving task on map "Town04" for controlling a car drive as far as possible in 1,000 frames. The environment contains a highway with 20 moving vehicles. The reward function prompts the driving distance and penalizes collisions. The observation shape is 84 x 420 pixels which consists of five 84 x 84 cameras. Each camera has 60 degrees view, together they produce 300 degrees first-person view, as shown in Figure S The weather, clouds and brightness may change slowly during the simulation. Figure 9 shows the training curves of AMBS and other comparison methods. AMBS performs the best, which provides empirical evidence that AMBS can potentially generalize over fask-irrelevant details, e.g., weather and clouds, in real world scenarios.
100
Figure 9: CARLA simulation.
6 CONCLUSION
This paper presents a novel framework AMBS with meta-learners for learning task-relevant and environmentt-detallinvveee
state representation for deep RL. In AMBS, we propose state encoders that decompose an observation into rewards and dynamics representations for measuring behavioral similarities by two self-learned meta-leamnrrs, and design a learnable strategy to balance the two types of representations. AMBS archives new state-of-the-a results in various environments and demonstrates its transfer ability on RL tasks. Experimental results verify that our framework is able to effectively learn a generalizable state representation.

---

Published as a conference paper at ICLR 2022
ACKNOWLEDGEMENT
This work is supported by Microsoft Research Asia collaborative research grant 2020.
REFERENCES
Rishabh Agarwal, Marlos C. Machado, Pablo Samuel Castro, and Mare G Bellemare. Contrastive behavioral similarity embeddings for generalization reinforcement learning. In International Conference Learning Representations, 2021. URL
id=oda7-sVgB4
Pablo Samuel Castro. Scalable methods for computing state similarity in deterministic markov decision processes. In Proceedings of the AAAI Conference on Artificial Intelligence, pp. 10069- 10076, 2020.
Pablo Samuel Castro, Tyler Kastner, Prakash Panangaden, and Mark Rowland. Mico: Improved rep- resentations via sampling-based state similarity for markov decision processes. In A. Beygelzimer, Y. Dauphin, Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, 2021. URL http::///pppeeeeeee.eeeee. net Fwrumm www..:.wwwwwwww...pgp.
Ting Chen, Simon Korbblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In Hal Daumé III and Aarti Singh (eds.), Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings Machine Learning Research, pp. 1597-1607. PMLR, 13-18 Jul 2020. URL http::/proceeddingt.
m1 wtt.ppeessl/l11/eeeeee
html
Karl Cobbe, Oleg Klimov, Chris Hesse, Tachoon Kim, and John Schulman. Quantifying generalization
in reinforcement learning. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pp. 1282-1288 1289. PMLR, 09-15 Jun 2019. URL http: pettp://proceed ings ngg
mlr. w.r.press//. wtt
9a. html
Alexey Dosovitskiy, German Ros, Felipe Codevilla, Antonio Lopez, and Vladlen Koltun. CARLA: An open urban driving simulator. In Proceedings of the 1st Annual Conference on Robot Learning, pp. 1-16, 2017.
Lasse Espeholt, Hubert Soyer, Remi Munos, Karen Simonyan, Vlad Mnih, Tom Ward, Yotam Doron, Vlad Firoiu, Tim Harley, lain Dunning, Shane Legg, and Koray Kavukcuoglu. IMPALA: Scalable distributed deep-RL with importance weighted actor-learner architectures. In Jennifer Dy and Andreas Krause (eds.), Proceedings of the 35th International Conference Machine Learning, volume 80 of Proceedings of Machine Learning Research, pp. 1407-1416. PMLR, 10-15 Jul 2018.
Norm Ferns, Prakash Panangaden, and Doina Precup. Metries for finite markov decision processes. In Proceedings of the 20th Conference on Uncertainty in Artificial Intelligence, UAI '04, pp. 162-169, Arlington, Virginia, USA, 2004. AUAI Press. ISBN 0974903906.
Norm Ferns, Prakash Panangaden, and Doina Precup. Bisimulation metrics for continuous markov decision processes. SIAM J. Comput., 40(6); 1662-1114, December 2011. I. ISSN 0097-5397. doi: 10.1137/100804444X URL /ttp::/t/////o.otg/// http:///oo..rgg11....11000000000
Carles Gelada, Saurabh Kumar, Jacob Buckman, Ofir Nachum, and Marc G. Bellemare. DecpMDP: Learning continuous latent space models for representation learning. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pp. 2170-2179. PMLR, 9a html
09-15 Jun 2019. URL htthhht::////ceeettiiit www.ngssmml..... www.:mmerr.p..s...mmmmmmmmmmmm.e /taess///9/
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In Jennifer Dy and Andreas Krause (eds.), Proceedings of the 35th International Conference of Proceedings of Machine Learning Research, pp. 1861-1888888 1870. PMLR, 10-15 Jul 2018. URL http::/proceedii igs ml www.pres//vee httpssss////hhaaaaaaa
Machine Learning, volume 80
8b. html
10

---

Published as a conference paper at ICLR 2022
Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson, Learning latent dynamics for planning from pixels. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pp. 2565. PMLR, 09-15 Jun 2019. htttttttt///occeeee
ml www.pppwsssv...c.. http://////aafffeeeemmm 19a.199..... 9a.html
Danijar Hafner, Timothy Lillicrap. Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. International Conference on Learning Representations, 2020.
URL https::/
net Fore://ormmm2m wit.mmmnmsss1111111111
Irina Higgins, Arka Pal, Andrei Rusu, Loic Matthey, Christopher Burgess, Alexander Pritzel, Matthew Botvinick, Charles Blundell, and Alexander Lerchner. DARLA: Improving zero-shot transfer in reinforcement learning. In Doina Precup and Yee Whye Teh (eds.), Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pp. 1480-1490 PMLR, 06-11 Aug 2017. URL httt::/pprcceedine. mlr
press/v70/higee ins1 ns
ht.ml
Will Kay, Joao Carreira, Karen Simonyan, Brian Zhang, Chloe Hillier, Sudheendra Vijayanarasimhan, Fabio Viola, Tim Green, Trevor Back, Paul Natsev, Mustafa Suleyman, and Andrew Zisserman. The kinetics human action video dataset. CoRR, abs/1705.06950, 2017.
Mete Kemertas and Tristan Ty Homentado-Armstrongg Towards robust bisimulation metric learning. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, 2021. URL httpp:://openevvii net fww.m?rrmmmmm
ySFG1F jgINN
Sascha Lange and Martin Riedmiller. Deep auto-encoder neural networks in reinforcement learning. In The 2010 International Joint Conference on Neural Networks (IJCNN), pp. 1-8, 2010.
Michael Laskin, Aravind Srinivas, and Pieter Abbeel. CURL: Contrastive unsupervised repre- sentations for reinforcement learning. In Hal Daumé IIII and Aarti Singh (eds.), Proceed- ings of the 37th International Conference on Machine Learning, volume 119 of Proceed- ings of Machine Learning Research, pp. 5639-566666999 56500555.005 PMLR, 13-18 Jul 2020a. URL http: wwppoceedinggsmmr..ppeo.. wtw.ep/sswsm....reee//eeee..eeee
1askin2 html
Misha Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, Pieter Abbeel, and Aravind Srinivas. Re- inforcement learning with augmented data. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 19884-19895, Curran Associates, Inc., 2020b. URL 2eae6aa1aae22222222222222
neurips
2t//ppperr/22///11//////. (t/ppppe//022///1eeeeeee
Alex X. Lee, Anusha Nagabandi, Pieter Abbeel, and Sergey Levine. Stochastic latent actor-critie: Deep reinforcement learning with a latent variable model. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 741-752, Curran Associates, Inc., 2020a. URL UtL.:tttpp:://rcceeiiig neurips. cc/
Kimin Lee, Kibok Lee, Jinwoo Shin, and Honglak Lee. Network randomization: A simple technique for generalization in deep reinforcement learning. In International Conference on Learning Representations, 2020b. URL 20200.
Timothee Lesort, Natalia Díaz-Rodríguez Jean-Franois Goudou, and David Filliat. State representa- tion learning for control: An overview. Neural Networks, 108:379-392, 2018
Volodymyr Mnih, Koray Kavukeuoglu, David Silver, Andrei A. Rusu, Joel Veness, Marc G. Bellemare, Alex Graves, Martin Riedmiller, Andreas K. Fidjeland, Georg Ostrovski, Stig Pe- tersen, Charles Beattie, Amir Sadik, Ioannis Antonoglou, Helen King, Dharshan Kumaran, Daan Wierstra, Shane Legg, and Demis Hassabis. Human-level control through deep rein- forcement learning. Nature, 518(7540):529 533, February 2015. ISSN 00280836. URL
http::/dx. doi wr.. 1ww1/00tt/neeeeee naturel 4236

---

Published as a conference paper at ICLR 2022
Marek Petrik and Bruno Scherrer. Biasing approximate dynamic programming with lower discount factor. In D. Koller, D. Schuurmans, Y. Bengio, and L. Bottou (eds.), Advances in Neural Information Processing Systems, volume 21. Curran Asso- ciates, Inc., 2009. URL htttps::/prrcceediimg neurips......e.. httsscccpaaerr//08////o/
Silviu Pitis, Elliot Creager, and Animesh Garg. Counterfactual data augmentation using locally factored dynamics. In H. Larochelle. M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. -996-39900.. Curran As- sociates, Inc., 2020. URL URL neurin wwtprrrs..cccppppeeem 29460028166677766cc1117777888888
Jordi Pont-Tuset, Federico Perazzi, Sergi Caelles, Pablo Arbeláez, Alex Sorkine- Hornung, and Lue Van Gool. The 2017 davis challenge on video object segmentation. CoRR, 0oss/704.006755 2018.
Roberta Raileanu, Max Goldstein, Denis Yarats, Ilya Kostrikov, and Rob Fergus. Automatic data augmentation for generalization in deep reinforcement learning, 2021.
Max Schwarzer, Ankesh Anand, Rishab Goel, R Devon Hjelm, Aaron Courville, and Philip Bach- man. Data-efficien reinforcement learning with self- celf-predictive representations, In International Conference on Learning Representations, 2021.
Austin Stone, Oscar Ramirez, Kurt Konolige, and Rico Jonschkowski. The distracting control suite- challenging benchmark for reinforcement learning from pixels. CoRR, abs/2101.022 .02722, 2021.
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, Timothy Lillicrap, and Martin Riedmiller. Deepmind control suite. CoRR, abs/ 8011 1801.006900 2018.
Niklas Wahlström, Thomas B. Schön, and Marc Peter Deisenroth. From pixels to torques: Policy learning with deep dynamical models. CoRR, abs/1502. 502.0225, 2015.
Manuel Watter, Jost Springenberg. Joschka Boedecker, and Martin Riedmiller. Embed control: A locally linear latent dynamics model for control from raw images. In C. Cortes, N. Lawrence, D. Lee, M. Sugiyama, and Garnett (eds.), Advances in Neural Information Processing Systems, volume 28. Curran Associates, Inc., 2015. URL htt.t:s::/pppcceeinngnoo newrips cc/ paper/20151 www.:////ile/aaaaaeeee.ooee.
954 0d05 www.00000729999300
pdf]
Denis Yarats, Ilya Kostrikov, and Rob Fergus. Image augmentation all you need: Regularizing deep reinforcement learning from pixels. In International Conference on Learning Representations, net wttp:/eeeeeeeettonneeeeeeeeeeeee
2021. URL httttaat////peeeeeeeee
6sTvaff
Chang Ye, Ahmed Khalifa, Philip Bontrager, and Julian Togelius. Rotation, translation, and cropping for zero-shot generalization. In 2020 IEEE Conference on Games (CoG), pp. 57-64, 2020a. doi: 10.1109966647356.220000
Chang Ye, Ahmed Khalifa, Philip Bontrager, and Julian Togelius. Rotation, translation, and cropping for zero-shot generalization. In 2020 IEEE Conference on Games (CoG), pp. 57-64, 2020b. doi:
10.1109CCC4444355220099299999
Amy Zhang, Yuxin Wu, and Joelle Pineau. Natural environment benchmarks for reinforcement learning. CoRR, abs/181 81 106032, 2018.
Amy Zhang, Rowan Thomas McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learning invariant representations for reinforcement learning without reconstruction. In International Conference on Learning Representations, 2021. URL
1d==2FCwDKRREU
Marvin Zhang, Sharad Vikram, Laura Smith, Pieter Abbeel, Matthew Johnson, and Sergey Levine. SO- Deep structured representations for model-basee based reinforcement learning. In Kamalika Chaud- huri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pp. 1444-7453, PMLR, 09-15 Jun 2019. URL http: /tppo/e// ngs. ml www.ppressvvv. / zhang 9m. .tml]
LAR:
12

---

Published as a conference paper at ICLR 2022
Brian D. Ziebart. Modeling Purposeful Adaptive Behavior with the Principle of Maximum Causal Entropy. PhD thesis, USA, 2010. AA13438449.
13

---

Published as a conference paper at ICLR 2022
A OBJECTIVES OF ADAPTIVE META-LEARNER OF BEHAVIORAL SIMILARITIES (AMBS)
We denote two encoders o, and dd. where each encoder maps tigh-dimensional observation to low-dimensional representation. The encoder or : S - Z, captures seward-relevare features and the encoder dd : S - Za captures dynamics-relevant features, where 2, and Za are representation spaces. We design two meta-learners fr : 2, x 2, - IR and fa Zn x Zn - R to output reward difference and dynamics difference, respectively. We follow DrQ (Yarats al. (021) to define the state transformation h : S x T + S that maps state to data-augmented state, where T is the space of parameters of h. In practice, we use random crop as transformation h.
With augmentation transformation h, the loss of representation learning is,
((0)=(11 -1 - () (z (5) (s) - -Ir; r31)
+ ( (1))
W2(P(-10a )() + (1 c) (1) (5)2), (s. - Ir: rsl)
(12)
+ to (faldull)) (2)). 2s(s ) W2 (2(P((lu(s a.), P(-lou(s))) ),
where θ (Jr. 6r ma). (1) M(s,, v.
) is the transformed state with parameters !! Specif-
(1) (2) and 5) transformed states with parameters of (1) (1) v v.() and v ically, s!") 8) 8
respectively, which are all drawn from T independently. P(-løu(s,) a. is a learned probabilistic dynamies model, which outputs Gaussian distribution over 2a for predicting dynamics represen- tation of next state sy. W2 denotes the 2- 2-Wasserstein distance which has closed-form for Gaussian distributions:
(2V W Σ,)) = ||+ 2yl12 + || ー Σ || where μl, R are the mean vectors of Gaussian distribution, Σ. E RR. Zal are the diagonals of covariance matrices, and || || is Frobenius norm. Note that we stop gradients for as mentioned in Section 4.2
The adaptive weight e and (1 c) are the softmax output of corresponding learnable parameter Me I It can be formulated as
[e: (1- so sof (maa((1))))
(13)
where the operation : denotes concatenation of two vectors/scalars,
We integrate the learning of into the update of Q-function of SAC. The objective of the Q-function with DrQ augmentation and clipped double trick is
(0 ne) B.
2 ((1--))))) cdb - (r(ss + +V((''))
(r(s, + + " (14)
+ c(11- cQ.
where s!!) = h(s, v((),, v()) are all drawn from T independently. and (Qo, 12_1 the double Q networks. The target value function with DrQ augmentation
V(s') = 2/2
min Qa (ICL - ((1=)))((s(1)) (oa(s"())) all1) + a log n. (aa0))
1=12
+ min Qs (111- - ((1- (2)) c54 x + a (a (21))
(15)
1=1
where s((1) = his', v()), a no (s(())()))) 0. is the set of parameters of the target network. S, and Φ. are target encoders specifically for the target Q network and ē is the adaptive weight for target encoders. The parameters 0x, Sr. du and îe are softly updated.
The loss of actor of SAC is.
JR (1b) = E Banay laloo((T log(T (al((()() min Qs. ()(s()),a))
(16)
k=1,
14

---

Published as a conference paper at ICLR 2022
The loss of of SAC is,
J(a) = (s)-D Baaxy la log( (al((((())) - -
(17)
where H. E R is the target entropy hyper-parameter which in practice is = -AA
The loss for updating of dynamics model p is
( - ( ( 20 (P(((loo)() a))
27
((P) (...s))))
(18)
where ((P(-)) and (P(())) are mean and standard deviation of P output Gaussian distribution, respectively.
Algorithm
shows the algorithm at each learning step.
Algorithm 2 AMBS+SAC
Input: Replay Buffer D, initialized @ 0r far Φat Q network QØ actor pig, target Q network Qo
2: Sample a batch with size B: {ss; Ti, (/)s D.
3: Shuffle batch {(s,, a, s 1b
to a
4: Update Q network by equation 14 with DrQ augmentation.
5: Update actor network by equation 16
6: Update alpha a by equation [7
7: Update encoder @ by equation
8: Update dynamics model P by equation
18
9: Softly update target Q network: 0₂ = TQ + (1 TQ)O,
10: Softly update target encoder: Φ = T&t + T6)Ø.
11: Softly update target č: A7 TØ (1 7
B IMPLEMENTATION DETAILS
B.1 PIXELS PROCESSING
We stack 3 consecutive frames observation, where each frame is 84 x 84 RGB images in DeepMind control suite. Each pixel divided by 255 to down scale [0, 1]. We consider the stacked frames as fully observed state.
B.2 NETWORK ARCHITECTURE
We share the convnet for encoder o, and encoder da. The shared convnet consists of four convolu- tional layers with 3 kernlls and 32 output channels. We set stride everywhere, except for the first convolutional layer, which has stride We use ReI.U activation for each convolutional layer output. The final output of convnet fed into two branches: 1) one is fully-connected layer with 50 dimensions to form reward representation (s),,, and 2) the other one is also fully-connected layer with 50 dimensions but form dynamics representation (s)... The meta-learners fr and fa are both MLPs with two layers with 50 hidden dimensions.
The critic network takes o, (s), da(s) and action as input, and feeds them into three stacked fully- connected layers with 1024 hidden dimensions. The actor network takes the shared convnet output as input, and feeds it into four fully-connected layers where the first layer has 100 hidden dimensions and other layers have 1024 hidden dimensions. The dynamics model is MIPss with two layers with 512 hidden dimensions. ReLU activations are used in every hidden layer.
B.3 HYPERPARAMETERS
The hyperparameters used in algorithm are listed in Table
15

---

Published as a conference paper at ICLR 2022
Parameter name
Value
Replay buffer capacity
10c
Discount factor 7 Minibatch size Optimizer
0.99
128
Adam
Learning rate for
10
Learning rate except a
5 x 10-4
Target update frequency
2
Actor update frequency
2
Actor log stddev bounds
[-10, 2]
TQ
0.01
T
0.05
Init temperature
0.1
Table 2: Hyperparameters used in algorithm.
B.4 ACTION REPEAT FOR DMC
We different action repeat hyper-parameters for each task in DeepMind control suite, which listed in Table
Task name
Action repeat
Cartpole Swingup
8
Finger Spin
2
Cheetah Run
4
Walker Walk
2
Reacher Easy
4
Ball In Cup Catch
4
Table Action repeat used for each task in DeepMind control suite.
B.5 SOURCE CODES
Source code of AMBS is available at htttttss///tth/ii wtwpchhuu.cmmmiaaaaaa www../owmmaaaaaaiee Chena/MM
C ANALYSIS
C.I REGRESSION LOSSES OF APPROXIMATING BISIMULATION METRIC
We compare the regression losses of approximating the bisimulation metric over RL environment steps among DBC and our AMBS. DBC uses the L1 norm to calculate the distances between two states embeddings. AMBS performs meta-learners on state representations to measure the distance with learned similarities. Figure 10 demonstrates that AMBS has smaller regression loss and also faster descent tendency compared to DBC. Meta-leamers AMBS are able to provide more stable gradients update the state representations.
16

---

Published as a conference paper at ICLR 2022
10P1
AMB5
10"1
0.3
0.2
2.4
3.0
0.0
01.4
2.6
0.0
0.2
0.5
02
5.4
19
Figure 10: The regression losses over RL environment steps among DBC and AMBS.
C.2 VISUALIZATIONS OF COMBINATION WEIGHT C
Figure I1 shows the values of combination weight e AMBS trained on DMC with origin background and Figure iis on DMC with video background setting. The values of at 1M steps vary in different tasks. The common trend that they increase at the beginning of training. In original background setting changes slowly after the beginning. In natural video background setting, has fast drop after the beginning. In the beginning, agent learns a large weight for dynamics feature that may boost the learning of RI. agent. Then it drops for natural video background setting because it learns that the video background is distracting the agent. Lower weight is learned for the dynamics features.
original background
video background
0.65
0.65
0.60
0.60
5
segght
0.55
0.55
0.50
0.50
cartpore awingup
cartpore swingup
finger spin cheetaturun
finger spin
0.45
cheetan run
walter
walker
0.4
1e6
Figure 11: The value of combination weight over environment steps. AMBS is trained DMC on Original background.
Figure 12: The value of combination weight c over environment steps. AMBS is trained on DMC on video background.
C.3 NORMS OF STATE EMBEDDINGS
We record the L1 norms of reward representation (s) and dynamics representation (s(s) during training. We record 1 n. (s)]l for reward representation where n, the number of dimensions of (s),, and n [|b((s))|| for dynamics representation where n,, is the number of dimensions of da(s). Figure 13 shows the values averaged on each sampled training batch. Although the T, norm of dba(s) decreases during training, it converges around 0.4-0.5. The L1 norm (x( also decreases in the training procedure. It can verify that dynamics representation ds(s) will not converge to all zeros.
embedding norm cheetah run
0.9
dyn
0.0
0.00
0.8
envirnment
165
Figure 13: Norms of state embeddings:
1 (8) and 1 [+)a())
17

---

Published as a conference paper at ICLR 2022
C.4 SHARING ENCODERS BETWEEN ACTOR AND CRITIC
In AMBS, we share only the convolutional layers ő encoders between actor and critic network. We compare to sharing the whole encoders 6, and @u in this experiment. Figure shows the raining curves on Walker- Walk with natural video background. Sharing full encoder Ør and da is worse than sharing only CNN Ø.
walker walk shared net
800
600
400
#14MBS)
share and &
environment
Figure 14: Training curves of different ways of sharing encoder. Experiments on Walker-War with natural video background.
D ADDITIONAL EXPERIMENTAL RESULTS
D.1 ADDITIONAL RESULT OF DMC SUITE WITH AND w/o BACKGROUND DISTRACTION
Figure 1. shows the training curves on original background setting. Figure 16 shows the training
curves on Kinetics (Kyy al. 2017) video background setting.
cartpole
finger
cheetah
1000
1200
800
800
500022
600
400
DBC
D/O
DrO
200
DrO
SAC
SAC
SAC
SUNC
SLAC
SLAC
envirommet
105
envimmment
106
endmrnmmnt
106
walker walk
reacher easy
hall in cup cup.caten catcn
1200
1000
1000
800
800
600
AMBS
400
AMBS
400
AMB5
DBC
200
DBC
0-
SLAC
SLAC
-200
SLAC
0.4
0.6
0.8
1.0
0.0
0.2
1.0
0.0
0.2
0.4
0.6
anvironment steps
106
Insolrooment sleps
1o0
environment steps
106
Figure 15: Training curves of AMBS and comparison methods. Each curve is average runs.
18

---

Published as a conference paper at ICLR 2022
caripole swingup
finger spin
cheetah
000
500
400
AMBS
AMB5
300
AMB5
DBC
DBC
200
DBC
D10
DHQ
SAC
SAC
100
0.0
0.4
0.6
0.8
1.0
0.0
0.4
0.5
0.8
1.0
0.0
0.2
0.4
0.6
0.8
snvirormmnnt slacs
enverormenn
walker walk
ball
1000
800
800
600
600
DDC
200
DOC
200
DIQ
DIQ
DrD
SAC
SAC
SAC
-200
200
SEAC
166
Figure 16: Training curves of AMBS and comparison methods on natural video background setting. Videos are sampled from Kinetics (Kay et al. 2017) dataset. Each curve is average 3 runs.
D.2 ADDITIONAL. RESULT OF GENERALIZZTION OVER BACKGROUND VIDRO
Table 4 shows the scores DAVIS video background setting comparing PSEs (Agarwal et al.)
2021
Methods
C-swingup
F-spin
C-run
W-walk
R-easy
BiC-catch
DrQ PSEs
749 + 19
779 + 49 308 + 12 789 = 28 955 + 10 821 + 17
AMBS
807 41
933
96 332
27
893
85 980 + 944 +
Table 4: Generalization with unseen background videos sampled from DAVIS 2017 (Pont-Tuset et al. 2018 dataset. We evaluate the agents at 500K environment steps. Scores of DrQ + PSEs are reported from Agarwal et al. (221))
D.3 ADDITIONAL RESULT ON CARLA
The CARLA experiment is to learn an autonomous driving RL agent to avoid collisions with other moving vehicles and drive as far as possible on highway within 1000 frames. Those vehicles are randomly generated, including vehicles' types, colors and initial positions. We raise the number of generated vehicles from to 40 so that the road becomes more crowded for RL agent. Besides, we also extend the training steps to 2e5 for AMBS and DBC. Figure 17 shows the training curve CARLA. AMBS and DBC run on the precious setting of CARLA, while AMBS40 and DBC40 run on the highway with 40 generated vehicles. When increasing the number of vehicles from 20 to 40, DBC becomes more difficult to learn but AMBS is just slightly slower. DBC still doesn't 't achieve high scores after le5 steps.
19

---

Published as a conference paper at ICLR 2022
255
AMDS
OOC
352
35
0.82
0.75
3.00
1.25
1.50
Figure 17: Training curve CARLA environment.
E PROOFS
E.1 VALUE-FUNCTION BOUND
Theorem E.1 (Value function difference bound for different discounting factors Petrik & Scherrer
202T)). Consider two identical MDPs except for different discounting factors "/1 and 722 where "/1 S 7/2 and reward r E [0, Let V denotes the value function for MDP with discounting factor ? given policy π.
2009 Kemertas Aumentado-Armstrone
72 71
V (s)
2/2 [s) E (1 11-mi 71 72)
(19)
Definition E.1 (c-weighted on policy bisimulation metric). Given a policy n, and E (0, 1). the on-policy bisimulation metric exists,
d(s:, 8,) = (1 (1-c) c) R2,, + cWi P.,))
(20)
Theorem E.2 (Value difference bound). For any c I 1), given two state s; and Sj.
(1 (v=)) d(s;,
(21)
proof. Lemma 6 of Kemertas & Aumentado-Armstrong (202I) Theorem E.3 (Generalized value difference bound). For any e E (0, 1), given two states S; and Sy.
(1-c) - V" (s,)) S s;) + 2(1 - c) - min(c,
(22)
(1 - (1--c)
proof. We follow the proof of Theorem in Kemertas & Aumentado Armstrong (2021)
Suppose
another MDP with discounting factor 3 exists. From Theorem H.2 we have
(1 (17c) - V2(s))] - s.).
(23)
Then,
(1 (s:) (T(s))
=(1 (s,) - + -P() -V5())
<(1 (1V35(s)) ([7 +1) (V - (7((),
<d(s;, 8,) + +(1 (==))((((s)) V75(s))
(24)
V,((s))
2(1 -0) miin(()))()
sd( +
(1 (1- (1-))
Replace 7 by c, finally we have
2(1 2(1-c)(r - min (c, 7))
(1 (1-c) (s=)) - v*( ( < d(s, +
(25)
(1 (1-1) (1- c)
20


Reviews:
[{"review_id": "rOHaH7tjfwx", "reviewer": "Reviewer_vhfp", "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": "## Strengths\n- Overall, the approach proposed in this work seems technically sound and well-motivated. Learning a state representation based on a bisimulation metric seems quite reasonable, and doing so appears to yield strong empirical results.\n- The experimental evaluation is quite thorough and illustrates several interesting phenomena. I appreciate that this work both considers environments where most of the state observation includes only distracting information in the DM Control tasks, as well as environments where a large part of the state observation is potentially useful, as in the CARLA tasks. The ablation studies and results on transferring to different reward functions are also interesting, and present fairly promising results.\n- The paper is generally fairly clear.\n\n## Weaknesses\n- There are no critical or significant weaknesses that I am aware of, though potentially slightly incremental on top of DBC.\n- There are numerous minor typos and small errors throughout the paper. For example, the domain and range of the function $F^\\pi$ in Equation (3) appear to be some undefined \"met\" space. At the end of the preliminaries, \"partial observations\" is used, where I believe \"partial observability\" is meant. The equations (4) and (5) are defined as \"least squares error,\" where I believe \"mean squared error\" is meant, or this could refer to a form of least squares regression. I believe that these are generally easy to address, but it would be useful to clean this up.\n\n## Additional Questions / Comments\n- In Equations (6) and (7), shouldn't this be restricted to samples where $a_i = a_j$? If the two actions don't match, I don't think you can necessarily expect the rewards and dynamics to be similar, even in similar states, e.g., even in exactly the same state. This seems potentially problematic to me.\n- Why does the actor condition only on the CNN layers of $\\phi$ in Equation (11) and not on the full learned state representation like the critic?\n- Do the representations $\\phi_r$ and $\\phi_s$ also receive gradients through updating the critic in $J_Q$? Or are they learned exclusively via the $\\ell(\\Theta)$ loss?\n\n-------\n\n## Post-rebuttal comments\n\nI appreciate the rebuttal, which has alleviated the concerns I've raised. I still find the proposed approach technically sound and well-motivated, and find that this work provides a reasonable empirical contribution. Several of the other reviewers also raised concerns about whether the approach is empirically justified. I find these concerns to be reasonable, although they are partially alleviated via the rebuttal and I believe they are outweighed by the empirical contribution of the work. Therefore, I continue to recommend acceptance, though I do not feel strongly enough to raise my score to the next possible option, which is an 8.", "summary_of_the_paper": "This work considers the problem of learning representations of high-dimensional pixel observations for RL. High-dimensional pixel observations often include many task-irrelevant details and ideally, an effective representation of such observations should only encode the task-relevant details. To implement this intuition, this work proposes to learn a two-part representation based off of a bisimulation metric. Two states that have similar first parts of the representation should exhibit similar reward structure, and two states that have similar second parts of the representation should exhibit similar dynamics structure. This work evaluates this learned representation on several continuous control tasks with distracting backgrounds and on CARLA, and shows that the proposed learned representation improves over baselines.", "summary_of_the_review": "Overall, this work offers a principled and empirically strong approach. There are some weaknesses, but they seem relatively easy to address, and the contributions of the work outweigh the cons. There were also some points of confusion for me in the paper that seem potentially serious, which I included above. I initially recommend a score of 6, but am willing to raise my score if the points of confusion are clarified.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "SE2me4-6FWx", "reviewer": "Reviewer_Np3q", "summary": "", "questions": "", "limitations": "", "rating": 5, "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": "*Strengths* \n\n- The paper studies an important problem in learning visual representations for RL robust to distractions. \n- The experimental results are pretty strong, outperforming the relevant baselines (DBC, DrQ) and including ablations which show all components of the proposed approach are useful.\n- Some components of the method, like combining bisimulation with data augmentation are well motivated.\n\n*Weaknesses*\n\nDespite the ablation which shows that each component of the method is important to performance, I still have some concerns/questions on the correctness of the proposed method.\n\nFirst, the idea of learning 2 separate state embeddings, one for reward and one for dynamics. The ablation suggests this is very important for good performance, however it seems to produce a metric that is no longer reflective of bisimulation. The bisim metric should capture states which are functionally similar (hence similar in both immediate reward and future state distribution). But by splitting the two up into separate embeddings, and training one only with the reward and one only with dynamics, it seems this functional similarity is no longer captured? The reward representation will only capture immediate reward, while its not clear what the dynamics representation will capture since its only objective is future state similarity in the same dynamics representation space. Is it not possible that the dynamics model would collapse to simple predicting 0 everywhere to minimize its loss?\n\nSecond, the intuition given for using a learned distance instead of L1 or L2 is unclear. While it may be easier to optimize, the point of using L1 or L2 is to explicitly structure the embedding space to make RL easier. When instead using an MLP its not clear how the embedding space itself will be shaped, or why it would be shaped in a way that is good for RL. Also, this is an important part of the method and should be clearly described in the main text, instead of in the appendix. Also, the framing of this learned distance as a \"meta-learner\" seems incorrect. From what I can understand it is simply using an MLP taking in the two embeddings and predicting distance, but not doing any meta learning. The confusing description + the fact that all implementation details about this are only in the appendix makes this component of the method difficult to understand.\n\nI hope the authors can clarify the above. Besides these aspects of the method which seem incorrect, the results appear strong, and if the authors can address the above concerns I'd be open to raising my score.\n", "summary_of_the_paper": "The paper studies the problem of learning invariant visual representations for reinforcement learning. Specifically, it builds on prior work Deep Bisimulation for Control (DBC) (Zhang et al), which learns an image representation predictive of rewards and dynamics and uses it with SAC. This paper makes a few modifications on top of DBC, specifically (1) instead of a single state embedding, it learns two separate ones (one for rewards and one for dynamics), (2) instead of using L1 distance in the embedding space to measure distance it instead uses a learned MLP which predicts distance, (3) using gradients from SAC it learns to dynamically adjust the balance between the dynamics weight and reward weight, and (4) adds image augmentation to the training. Experiments indicate on Distracting Control that on 2/4 tasks it matches DrQ and DBC and on the other 2 tasks outperforms them. Ablations suggest that all components are important, particularly (1) and (2).", "summary_of_the_review": "The paper studies an important problem and has some solid results. However in their current form the major components of the method are not clearly motivated and there are some questions I have about the correctness/justification of the approach. \n\n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "LKoq985axWg", "reviewer": "Reviewer_bfCF", "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": "\nStrengths: \n- The paper’s improvements over prior work are generally pretty reasonable and backed by convincing experimental results\n- The experimental results are generally pretty thorough with relevant comparisons to prior works and ablations\n- The writing is clear and exposition of the literature is thorough \n\nWeaknesses:\n- This paper generally feels like a “DBC with more bells and whistles” that enable it to work better; I think the general novelty is on the weaker side. The authors can consider a more thorough investigation into where existing representation learning methods fall short and address them more fundamentally\n- The experiments with natural video backgrounds generally feel pretty contrived (granted this is what’s in the literature). The only realistic setting is CARLA and I would have liked to see more environments of that nature or more detailed experimentation there.\n\nDetailed questions:\n- If reward and dynamics have different encoders, to what extent is this still learning a bisimulation, or mapping to an equivalent state space?\n- What do the experimental evaluations look like with more environment steps? Are the asymptotic performances of all methods similar or is AMBS just more sample efficient in the beginning? e.g. DBC uses 1e6 steps on CARLA\n- nit: it’s not really clear to me the naming of “meta-learner”, what is the meta train/test tasks here? I think something like learned distance metric would be clearer\n\n", "summary_of_the_paper": "\nThis paper proposes a representation learning method for reinforcement learning via the bisimulation metric. It’s largely an extension of the work in Deep Bisimulation for Control (DBC) (Zhang et al., 2020) with some additional improvements:\n- Learned distance metrics (“meta-learners”) instead of L1 distance on the learned representations\n- Separate feature encodings for state and dynamics\n- Data augmentation from DrQ (Yarats et al., 2021) \n- Learned c tradeoff between reward and dynamics\n\n", "summary_of_the_review": "Overall this paper’s contributions are generally reasonable and empirically backed, however its novelty and experimental evaluation leave much to be desired. To this end, I am leaning towards acceptance but would not be upset if this paper is rejected.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "1V0dC4F5yGP", "reviewer": "Reviewer_rJ9b", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 5, "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": "While an interesting extension of DBC from Zhang et al. (2021), I have 3 main concerns with this paper:\n1. **Theoretical issues:** My main concern with this paper is in terms of the theoretical justification for learning $c$.\n    * The $c$ term in equation (3) (which originally came from Ferns et al., [2004]) is _directly_ tied to $\\gamma$. Indeed, this is how the authors of that paper were able to prove the value-function bounds. Castro [2020] only uses the second $c$ term (and sets it equal to $\\gamma$), which still yields the value-function bounds. By making the $c$ a learnable parameter, it seems like the connections to the theoretical properties of bisimulation are completely lost.\n    * Furthermore, by having a varying $c$ term, it is not at all clear that the underlying metrics even converge to a fixed point!\n    * **Post-rebuttal note:** The authors have mostly addressed this in their rebuttal. There are some minor issues with their proofs and I've provided some suggestions.\n1.  **Implementation design choices:**\n    * Why is $f_r$ even necessary? Rewards are typically fully observed so it's not clear why this needs to be learned at all.\n    * The use of $V_*^{(1)}$ in the last term of equation (1) is a strange design decision, and I'm not sure I follow the justification. What is meant by \"make a consensus prediction\"?\n    * **Post-rebuttal note:** The authors have addressed these concerns in their rebuttal.\n1. **Statistical significance of results:*** 3 runs is on the lower-side of what should be used for these types of experiments. Further, the authors do not specify what the shaded areas represent in their figures.\n    * **Post-rebuttal note:** The authors have promised to run more seeds. I have provided some suggestions in my high-level comment.\n\n\n---\n\nSome questions/comments for the authors:\n1. In the third paragraph of page 2, the authors say \"their behavioral distance to the other state representations\". What other state representations are being referred to? It is not clear.\n1. In the third paragraph of page 2 the authors say \"more side information can be preserved\". What does \"side information\" mean, and how is it being preserved?\n1. In the third paragraph of page 2 it says \"observe that a smaller loss can be obtained\". What loss? A smaller loss with respect to what?\n1. In the third paragraph of page 2 it says \"the approximation precision issue\", which issue are you referring to, specifically?\n1. Can you clarify what you mean by \"least fixed point\" below equation 3?\n1. In figure 1, are all the $c$s in the figure (there are three of them) all the same?\n1. Below Figure 1 the authors say \"where $\\phi$ is the state encoder\", but no state encoder has been introduced.\n1. In Section 4.1, do $\\phi_r$ and $\\phi_d$ share parameters? Figure 1 suggests they do.\n1. In page 5 it says \"lead to large regression losses\". What regression loss is this, specifically?\n1. In page 5 it says \"which destabilze the representation learning\". In what way is it destabilized?\n1. In page 5 it says \"which however may loss[sic] part of useful information\". What useful information is being referred to, specifically?\n1. In the second paragraph of page 5 it says \"able to preserve more task-relevant information\". What do the authors mean by this, specifically? What task-relevant information?\n1. The sentence immediately above equation (11) ends with \"comprises of encoder is\". What is meant by this?\n1. In the line below equation (11) what is meant by \"a convention form of Q-function\"?\n1. Below equation (11), if Q depends on $c$, shouldn't $c$ also be a parameter of the Q function?\n1. In the learning curves, what do the shaded areas represent?\n1. In section 5.2 it says \"**AMBS + Drq** is DBC with data augmentation\", do the authors mean \"AMBS with data augmentation\"?\n1. In section 5.2 the authors tried fixing $c$ to $0.5$, but a more natural choice would have been $\\gamma$, given the main point made above. Further, neither DBC nor the original $\\pi$-bisimulation from Castro (2020) use the $(1-c)$ term in front of the reward differences. Fixing $c=\\gamma$ and removing the $(1-c)$ term would have yielded a more direct comparison to DBC.\n\n---\n\nSome minor comments:\n1.  In the first paragraph of the introduction, consider replacing \"based on some RL algorithms\" with \"based on various RL algorithms\"\n1.  In the first line of page 2, should be \"of **a** Markov Decision Process\"\n1.  A paper from early in the year is quite relevant and should be included in the related work section for representation learning:\n    * [Agarwal et al., Contrastive Behavioral Similarity Embeddings for Generalization in Reinforcement Learning; ICLR 2021](https://arxiv.org/abs/2101.05265).\n1.  In the second paragraph of page 2, consider including two recent papers that have been accepted to NeurIPS:\n     *  [Castro et al., MICo: Improved representations via sampling-based state similarity for Markov decision processes](https://arxiv.org/abs/2106.08229)\n     *  [Kemertas and Aumentado-Armstrong, Towards Robust Bisimulation Metric Learning](https://arxiv.org/abs/2110.14096)\n1.  In the second paragraph in page 2, should be \"and potentially ***lose*** parts of ***the*** state features\"\n1.  In the third paragraph in page 2 they say \"pair of meta-learners that learn similarities\". Please specify which similarities are being referred to.\n1.  In the third paragraph in page 2 it should say \"hand-craft***ed*** form\"\n1.  Right above section 3, use `\\citet` for \"Pitis et al\".\n1. In section 3, specify the range of $\\gamma$. e.g. $\\gamma\\in [0, 1)$\n1. Right below equation 2, should say \"where $d$ quantif***ies*** the behavioral...\"\n1.  In the same paragraph, better to say \"defines a metric ***with respect*** to a policy $\\pi$.\n1. In the caption of Figure 1, should say \"which ***is*** jointly learned\"\n1. Below Figure 1 caption, $\\gamma$ is referred to \"the hyper-parameter\", but I think the authors mean discount factor.\n1. Below Figure 1 caption, should say \"combined with ***the*** reinforcement learning...\"\n1. In the first paragraph of section 4, remove the \"the\" before Figure 1 (i.e. \"is demonstrated in Figure 1\")\n1. In the first paragraph of section 4, should say \"two learned similarit***ies*** in a specific...\"\n1. In section 4.1 the network architecture is mentioned but has not been introduced. I'd suggest referencing the appendix where it is introduced.\n1. In the second paragraph of page 5 it should say \"and therefore ***the*** meta-learner...\"\n1. In the second paragraph of page 5 it should say \"the process of udpating ***the*** state encoder\"\n1. In the second paragraph of page 5 it should say \"Besides, $f_*$ is ***a*** non-linear transformation\"\n1. In equation (5) specify that you are using the closed form of the $W_2$ metric (as in DBC). This is mentioned in the appendix, but should be clarified in the main paper as well.\n1. Above equation (6) you mention the \"learned parameteric dynamics model\". Please add a reference to the appendix where it is defined.\n1. Above equation (6) should say \"two transitions sampled from ***the*** replay buffer\"\n1. In the first paragraph of 4.2 you should specify that $c\\in (0, 1)$ is enforced by a softmax. This is specified in the appendix but should be clarified in the main paper as well.\n1. In the paragraph above equation (11) should say \"we aim to make $f_r$ and $f_d$ ***symmetric***\"\n1. In the line above 5.1, should say \"which is a common***ly*** use***d*** off-policy...\"\n1. In the first sentence in 5.1, should say \"is a***n*** environment ***that*** provides high dimensional pixel observation***s*** for RL tasks\"\n1. In the first sentence of 5.3 it should say \"but have differen***t*** reward functions\"\n1. In section 5.5 should say \"learning of RL agent***s*** may suffer from...\"\n1. In section 5.5 should say \"DBC to create a***n*** autonomous driving...\"", "summary_of_the_paper": "This paper leverages some recent work on bisimulation metrics to develop a pair of meta-learners to capture the two parts of a bisimulation metric: reward similarities and dynamics similarities. The authors evaluate their method on the environments used by the recent DBC (Zhang et al., 2021) paper, and compare against this and other related methods.", "summary_of_the_review": "See main concerns above.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper proposes to learn a state-representation using bi-simulation in an RL setting. The approach is thoroughly evaluated on several benchmarks. In its current form the paper is mainly an empirical contribution, with now some theoretical contributions tucked away in the appendices. Nevertheless, an interesting approach with promising results.

The reviewers appreciated the revised paper and the discussion. The replies and discussions successfully addressed all serious concerns of the reviewers. Please also clarify the discussed points in the next iteration of the paper, and run the experiments with more seeds, as promised.
Decision: Accept (Poster)