Original: Learning Representations via a Robust Behavioral Metric for Deep Reinforcement Learning
Revised: Learning Representations via a Robust Behavioral Metric for Deep Reinforcement Learning

Original: Abstract
Revised: Abstract

Original: Learning an informative representation with behavioral metrics is able to acceleratedeep reinforcement learning process.
Revised: Learning an informative representation with behavioral metrics is able to accelerate the deep reinforcement learning process.

Original: Two key research issues on behavioralmetric-based representation learning are how to relax the computation of a speciﬁcbehavioral metric, which is difﬁcult or even intractable to compute, and  how toapproximate the relaxed metric by learning an embedding space for states.
Revised: There are two key research issues on behavioral metric-based representation learning: 1) how to relax the computation of a specific behavioral metric, which is difficult or even intractable to compute, and 2) how to approximate the relaxed metric by learning an embedding space for states.

Original: In thispaper, we analyze the potential relaxation and/or approximation gaps for existingbehavioral metric-based representation learning methods.
Revised: In this paper, we analyze the potential relaxation and/or approximation gaps for existing behavioral metric-based representation learning methods.

Original: Based on the analysis,we propose a new behavioral distance, the RAP distance, and develop a practicalrepresentation learning algorithm on top of it. We provide theoretical analysis onthe proposed algorithm.
Revised: Based on the analysis, we propose a new behavioral distance, the RAP distance, and develop a practical representation learning algorithm on top of it with a theoretical analysis.

Original: We conduct extensive experiments on DeepMind Control Suite with distraction, Robosuite, and autonomous driving simulator CARLA todemonstrate new state-of-the-art results.
Revised: We conduct extensive experiments on DeepMind Control Suite with distraction, Robosuite, and autonomous driving simulator CARLA to demonstrate new state-of-the-art results.

Original:  Introduction
Revised:  Introduction

Original: Deep reinforcement learning (RL) aims to interactively learn an optimal policy from high-dimensionalenvironmental observations or states in an end-to-end manner.
Revised: Deep reinforcement learning (RL) aims to interactively learn an optimal policy from high-dimensional environmental observations or states in an end-to-end manner.

Original: In the literature, it has been demonstrated that a robust representation of states, which are task-relevant and invariant to task-irrelevantbackground information, is able to signiﬁcantly speed up the RL process and makes the learned policymore generalizable.
Revised: In the literature, it has been demonstrated that a robust representation of states, which are task-relevant and invariant to task-irrelevant background information, is able to significantly speed up the RL process and makes the learned policy more generalizable.

Original: Therefore, representation learning has played a key role in Deep RL algorithmsand attracts more and more attention in the RL community [4, 23, 17, 20].
Revised: Therefore, representation learning has played a key role in Deep RL algorithms and attracts more and more attention in the RL community [4, 23, 17, 20].

Original: Prior works on representation learning are focused on learning embeddings to represent statesbased on a reconstruction loss [26, 25, 13].
Revised: Prior work on representation learning is focused on learning embeddings to represent states based on a reconstruction loss [26, 25, 13].

Original: Though promising results have been reported on some RL application domains, policies learned with such representations usually fail to generalize wellin a complex environment because minimizing a reconstruction loss may potentially introducelocal (visual) features with task-irrelevant information.
Revised: Though promising results have been reported on some RL application domains, policies learned with such representations usually fail to generalize well in a complex environment because minimizing a reconstruction loss may potentially introduce local (visual) features with task-irrelevant information.

Original: Another research direction is to apply dataaugmentation techniques to learn robust state representations for an RL agent [19, 28, 27].
Revised: Another research direction is to apply data augmentation techniques to learn robust state representations for an RL agent [19, 28, 27].

Original: While dataaugmentation-based approaches may be able to learn more robust feature representations, they donot take the characteristics ofa Markov Decision Process (MDP). e.g., reward signals and dynamicsmodels, into consideration when learning representations.
Revised: While data augmentation-based approaches may be able to learn more robust feature representations, they do not take the characteristics of Markov Decision Processes (MDPs), e.g., reward signals and dynamics models, into consideration when learning representations.

Original: As a result, the learned representations maynot be informative for learning an optimal policy for the task of interest.
Revised: As a result, the learned representations may not be informative for learning an optimal policy for the task of interest.

Original: A third research directionis to construct some auxiliary tasks in addition to the prime RL task and learn state representationsby learning all the tasks simultaneously [14, 2, 21].
Revised: A third research direction is to construct some auxiliary tasks in addition to the prime RL task and learn state representations by learning all the tasks simultaneously [14, 2, 21].

Original: These approaches can be considered as shapingthe state representations implicitly by learning the auxiliary tasks, which do not have a guarantee forlearning a better policy, especially when the auxiliary tasks fail to beneﬁt the learning of the RL task.
Revised: These approaches can be considered as shaping the state representations implicitly by learning the auxiliary tasks, which do not have a guarantee for learning a better policy, especially when the auxiliary tasks fail to benefit the learning of the RL task.

Original: Recently, behavioral metrics, such as the bisimulation metric and its variants [7, 8, 3], have beenexploited in representation learning for deep RL.
Revised: Recently, behavioral metrics, such as the bisimulation metric and its variants [7, 8, 3], have been exploited in representation learning for deep RL.

Original: Behavioral metrics are originally proposed tomeasure behavioral similarity between states in terms of rewards and dynamics models, e.g, state transition probabilities.The high-level idea is to
Revised: Behavioral metrics are originally proposed to measure the behavioral similarity between states in terms of rewards and dynamics models, e.g, state transition probabilities.

Original:  learn state embeddings by preserving the 
Revised: The high-level idea is to learn state embeddings by preserving the behavioral similarity between states based on a specific behavioral metric [29, 1, 16, 4].

Original: behavioral

Original: similarity between states based on a speciﬁc behavioral metric [29, 1, 16, As behavioral metricsprovide a theoretical bound on the difference between the outputs of value function on pair ofstates, the learned representation enjoys a theoretical guarantee to capture behavioral structure in theenvironment for policy learning.
Revised: As behavioral metrics provide a theoretical bound on the difference between the outputs of value function of a pair of states, the learned representation enjoys a theoretical guarantee to capture behavioral structure in the environment for policy learning.

Original: However, as behavioral metrics are expensive or even intractable tocompute, different approximation approaches and learning objectives have been proposed to makebehavioral metric-based representation learning for RL agents more efﬁciently [29, 1, 16, 4].
Revised: However, as behavioral metrics are expensive or even intractable to compute, different approximation approaches and learning objectives have been proposed to make behavioral metric-based representation learning for RL agents more efficient [29, 1, 16, 4].

Original: Though behavioral metric-based representation learning methods have achieved state-of-the-art resultson some benchmark RL problems, they suffer from at least one of the following three issues: lossfunction mismatch, relaxation of dynamics model divergence andlimitation of the L 1 / L 2 norm.
Revised: Though behavioral metric-based representation learning methods have achieved state-of-the-art results on some benchmark RL problems, they suffer from at least one of the following three issues: loss function mismatch, relaxation of dynamics model divergence and the L 1 / L 2 norm limitation.

Original: • Loss function mismatch .
Revised: • Loss function mismatch .

Original: In behavioral metric-based representation learning, the loss function is  the difference between the distance between states on the embedding space and their behavioral metric.
Revised: In behavioral metric-based representation learning, a loss function is to measure the difference between the distance between states on the embedding space and their behavioral metric.

Original: As behavioral metrics are expensive or even intractable to computers, approximation or relaxation are necessary.
Revised: As behavioral metrics are expensive or even intractable to computers, approximation or relaxation are necessary.

Original: However, based on our analysis,state-of-the-art methods adopt improper approximations, which may introduce  bias and makes the bound on value function looser.
Revised: However, based on our analysis, state-of-the-art methods adopt improper approximations, which may introduce a bias and makes the bound of the value function looser.

Original: • Relaxation of dynamics model divergence .
Revised: • Relaxation of dynamics model divergence .

Original: The bisimulation metric or its on-policy invariant is one of the most widely used behavioral metrics, which requires estimating the 1-Wasserstein distance between dynamics models.As the 1-Wasserstein distance is
Revised: The bisimulation metric or its on-policy invariant is one of the most widely used behavioral metrics, which requires estimating the 1-Wasserstein distance between dynamics models.

Original:  usually difﬁcult or intractable to estimate, prior methods propose some relaxations to replace the estimation of the 1-Wasserstein, which may break some theoretical guarantees of the bisimulation metric.
Revised: As the 1-Wasserstein distance is usually difficult or intractable to estimate, prior methods propose some relaxations to replace the estimation of the 1-Wasserstein, which may break some theoretical guarantees of the bisimulation metric.

Original: • Limitation of the L 1 / L 2 norm  .
Revised: • The L 1 / L 2 norm limitation .

Original: The L 1 and the L 2 norms are commonly-used distances with zero self-distance.
Revised: The L 1 and the L 2 norms are commonly-used distances with zero self-distance.

Original: However, due to the two gaps mentioned above, the approximations or relaxations of behavior metrics in prior approaches are potentially non-zero self-distance.
Revised: However, due to the two gaps mentioned above, the approximations or relaxations of behavior metrics in prior approaches are potentially non-zero self-distance.

Original: Using the L 1 or L 2 distance on embedding space is difﬁcult to learn robust representation topreserve the behavioral similarity between states.
Revised: Using the L 1 or L 2 distance on embedding space is difficult to learn robust representation to preserve the behavioral similarity between states.

Original: To  the aforementioned issues or gaps, in this paper, we ﬁrst introduce a new behavior metric namelythe Reducing Approximation Gap (RAP) distance, andwe then develop a practical approximationalgorithm with consistency to its theoretical prototype.
Revised: To address the aforementioned issues or gaps, in this paper, we first introduce a new behavior metric namely the Reducing Approximation Gap (RAP) distance, and then develop a practical approximation algorithm with consistency to its theoretical prototype.

Original: In this way, our algorithm is guaranteed to learna robust state representation to capture  behavioral similarity between states.
Revised: In this way, our algorithm is guaranteed to learn a robust state representation to capture the behavioral similarity between states.

Original: We conduct extensiveexperiments on DeepMind Control Suite (DMC) [24] with background distraction, Robosuite [30]and autonomous driving simulator CARLA [6] to demonstrate new state-of-the-art results comparedwith other behavior metric-based representation learning methods.
Revised: We conduct extensive experiments on DeepMind Control Suite (DMC) [24] with background distraction, Robosuite [30] and autonomous driving simulator CARLA [6] to demonstrate new state-of-the-art results compared with other behavior metric-based representation learning methods.

Original: The contributions of this paper are two-fold: 1) we analyze the potential approximation gaps forexisting behavioral metric-based representation learning approaches, and 2) we introduce a newbehavior distance RAP and develop its practical approximation algorithm with theoretical guarantees.
Revised: The contributions of this paper are two-fold: 1) we analyze the potential approximation gaps for existing behavioral metric-based representation learning approaches, and 2) we introduce a new behavior distance RAP and develop its practical approximation algorithm with theoretical guarantees.

Original:  Related Work
Revised:  Related Work

Original: In RL, early research works have focused on learning state representations by designing and optimizing some auxiliary objectives in addition to the RL task of interest.
Revised: In RL, early research work has focused on learning state representations by designing and optimizing some auxiliary objectives in addition to the RL task of interest.

Original: For instance, Hafner et al.[11] propose to learn a dynamics model to predict future states with  reconstruction loss.
Revised: For instance, Hafner et al. [11] propose to learn a dynamics model to predict future states with a reconstruction loss.

Original: Geladaet al.
Revised: Gelada et al.

Original: [9] aim to learn state representations by predicting  dynamics model and  reward function onembedding space.
Revised: [9] aim to learn state representations by predicting a dynamics model and a reward function on an embedding space.

Original: Laskin et al.
Revised: Laskin et al.

Original: [17] apply contrastive learning with samples generated by momentumencoder.
Revised: [17] apply contrastive learning with samples generated by the momentum encoder.

Original: Hansen et al.
Revised: Hansen et al.

Original: [12] predict an inverse model as  auxiliary task for representation learning.
Revised: [12] predict an inverse model as an auxiliary task for representation learning.

Original: Another type of approaches aims to apply data augmentation techniques to improve representationlearning.
Revised: Another type of approaches aims to apply data augmentation techniques to improve representation learning.

Original: For instance, Laskin et al.
Revised: For instance, Laskin et al.

Original: [18] conduct  extensive study of data augmentation for deep RLwith pixel-based input.
Revised: [18] conduct an extensive study of data augmentation for deep RL with pixel-based input.

Original: Yarats et al.
Revised: Yarats et al.

Original: [28] adopt random crops on pixel-based input and regularizesthe Q-function objectives.
Revised: [28] adopt random crops on pixel-based input and add regularization terms on the Q-function objectives.

Original: Lee et al.
Revised: Lee et al.

Original: [19] introduce convolutional neural network with randomizedparameters.
Revised: [19] introduce convolutional neural networks with randomized parameters.

Original: Stooke et al.
Revised: Stooke et al.

Original: [23] use augmented samples for representation contrastive learning, Incontrast to these works, our proposed method aim to encode state representation into a structuralmetric space based on behavioral metrics.
Revised: [23] use augmented samples for representation contrastive learning, In contrast to these works, our proposed method aim to encode state representations into a structural metric space based on behavioral metrics.

Original: Recently, behavioral metric-based representation learning has attracted more and more attention inthe RL community.
Revised: Recently, behavioral metric-based representation learning has attracted more and more attention in the RL community.

Original: For instance, Zhang et al.
Revised: For instance, Zhang et al.

Original: [29] aim to learn representation by approximating the bisimulation metric [7, 3] on  embedding space.
Revised: [29] aim to learn representations by approximating the bisimulation metric [7, 3] on an embedding space.

Original: Agarwal et al.
Revised: Agarwal et al.

Original: [1] propose a behavioral metricconsidering a distance between action distributions given states for representation learning.
Revised: [1] propose a behavioral metric considering a distance between action distributions given states for representation learning.

Original: Kemertasand Aumentado-Armstrong [16] propose to improve the robustness of the representation learningmethod in [29] by adding norm constraints on the embedding space and intrinsic rewards.
Revised: Kemertas and Aumentado-Armstrong [16] propose to improve the robustness of the representation learning method proposed in [29] by adding norm constraints on the embedding space and intrinsic rewards.

Original: Castroet al.
Revised: Castro et al.

Original: [4] introduce a new behavioral distance and develop a sampling-based approach to preservebehavioral similarity between states on the embedding space.
Revised: [4] introduce a new behavioral distance and develop a sampling-based approach to preserve the behavioral similarity between states on the embedding space.

Original: Chen and Pan [5] propose to learnneural networks to approximate components in the bisimulation metric on embedding state  space.
Revised: Chen and Pan [5] propose to learn neural networks to approximate components in the bisimulation metric on the state embedding space.

Original:  Preliminaries
Revised:  Preliminaries

Original: Reinforcement Learning We consider a Markov Decision Process (MDP) deﬁned by a tuple(cid:104)S , A , R , P, γ (cid:105) , where S is the high-dimensional state space, A is the action space, P ( s (cid:48) | s , a ) is thetransition distribution that captures the probability of entering  next state s (cid:48) ∈ S given  current states ∈ S and  action a ∈ A , R : S ×
Revised: Reinforcement Learning We consider a Markov Decision Process (MDP) defined by a tuple ⟨S , A , R , P, γ ⟩ , where S is the high-dimensional state space, A is the action space, P ( s ′ | s , a ) is the transition distribution that captures the probability of entering a next state s ′ ∈ S given a current state s ∈ S and an action a ∈ A , R : S ×

Original: A → R is the reward function and γ ∈ [0 , 1) is the discount factor.
Revised: A → R is the reward function and γ ∈ [0 , 1) is the discount factor.

Original: In the sequel, we use P a s and r a s to denote P ( ·| s , a ) and R ( s , a ) , respectively.
Revised: In the sequel, we use P a s and r a s to denote P ( ·| s , a ) and R ( s , a ) , respectively.

Original: A policy π ( a | s ) is aprobability distribution over each action a conditioned on  state s .
Revised: A policy π ( a | s ) is a probability distribution over each action a conditioned on a state s .

Original: The value function V π : S → Rfor a given policy π at  state s is deﬁned as the expected sum of discounted future rewards,
Revised: The value function V π : S → R for a given policy π at a state s is defined as the expected sum of discounted future rewards,

Original: [Equation]
Revised: [Equation]

Original: The goal of reinforcement learning is to ﬁnd an optimal policy π ∗ = arg max π V π that maximizesthe expected future rewards.
Revised: The goal of reinforcement learning is to find an optimal policy π ∗ = arg max π V π that maximizes the expected future rewards.

Original: In the scope of representation learning for deep RL, a state encoderφ : S → R n maps a high-dimensional state to a low-dimensional vector, with which a policyπ ( a | φ ( s )) is learned.
Revised: In the scope of representation learning for deep RL, a state encoder ϕ : S → R n maps a high-dimensional state to a low-dimensional vector, with which a policy π ( a | ϕ ( s )) is learned.

Original: Bisimulation Metrics
Revised: Bisimulation Metrics

Original: The bisimulation metric [7, 8] deﬁnes a pseudometric d : S × S
Revised: The bisimulation metric [7, 8] defines a pseudometric d : S × S

Original: → R tomeasure the behavioral distance between states.
Revised: → R to measure the behavioral distance between states.

Original: Recently, a variant of the bisimulation metric, onpolicy bisimulation metric (or π -bisimulation metric) is proposed [3], which is focused on behaviorsrelative to a particular policy π .
Revised: Recently, a variant of the bisimulation metric, onpolicy bisimulation metric (or π -bisimulation metric) is proposed [3], which is focused on behaviors relative to a particular policy π .

Original: π -bisimulation metric consists of a reward difference term and a
Revised: π -bisimulation metric consists of a reward difference term and a

Original: Wasserstein distance in dynamics models between states.
Revised: Wasserstein distance in dynamics models between states.

Original: Theorem 3.1 ( π -bisimulation metric [3]) .
Revised: Theorem 3.1 ( π -bisimulation metric [3]) .

Original: Let M be the set of all pseudometrics on space S .
Revised: Let M be the set of all pseudometrics on space S .

Original: Apseudometric transformation function F πB : M → M is deﬁned as,
Revised: A pseudometric transformation function F πB : M → M is defined as,

Original: [Equation]
Revised: [Equation]

Original: where E a i∼ π r a
Revised: where E a i

Original: i s i =
Revised: [Equation]

Original: E a i ∼ π ( ·| s i )
Revised: E [Equation] E

Original: r a i s i , Pπ s i =
Revised: [Equation] P

Original: E a ∼ π ( ·| s i ) P a s i and W 1 is  1-Wasserstein distance.
Revised: [Equation] and W 1 is the 1-Wasserstein distance.

Original: F πB has116 a least ﬁxed point d πB and d πB is a π -bisimulation metric.
Revised: F πB has a least fixed point d πB and d πB is a π -bisimulation metric.

Original: The following theorem shows that the difference in value function is bounded by d πB .
Revised: The following theorem shows that the difference in value function is bounded by d πB .

Original: Theorem 3.2 (Value difference bound [3]) .
Revised: Theorem 3.2 (Value difference bound [3]) .

Original: Given states s i and s i and policy π , we have
Revised: Given states s i and s j , and policy π , we have

Original: [Equation]
Revised: [Equation]

Original: MICo distance The MICo distance [4] is a variant of the π -bisimulation metric, which measuresthe distribution distance between dynamics models by computing the distance between sampled nextstates from the dynamics models in order to avoid the computation of the Wasserstein distance.
Revised: MICo distance The MICo distance [4] is a variant of the π -bisimulation metric, which measures the distribution distance between dynamics models by computing the distance between sampled next states from the dynamics models in order to avoid the computation of the Wasserstein distance.

Original: Theorem 3.3 (MICo distance [4]) .
Revised: Theorem 3.3 (MICo distance [4]) .

Original: Let M be the space of distance function d : S × S
Revised: Let M be the space of distance function d : S × S

Original: → R , if the
Revised: → R , if the

Original: MICo metric function F πM : M → M is deﬁned as,
Revised: MICo metric function F πM : M → M is defined as,

Original: [Equation]
Revised: [Equation]

Original: then F πM has a unique ﬁxed point d πM .
Revised: then F πM has a unique fixed point d πM .

Original: Theorem 3.4 (Value difference bound [4]) .
Revised: Theorem 3.4 (Value difference bound [4]) .

Original: Given states s i and s i and policy π , we have
Revised: Given states s i and s j , and policy π , we have

Original: [Equation]
Revised: [Equation]

Original:  Approximation Gaps in Behavioral Metric-based Representation Learning
Revised:  Approximation Gaps in Behavioral Metric-based Representation Learning

Original: The high-level idea of behavioral metric-based representation learning is to learn an embedding spacesuch that after mapping states onto the embedding space, the behavioral similarity can be preserved.
Revised: The high-level idea of behavioral metric-based representation learning is to learn an embedding space such that after mapping states onto the embedding space, the behavioral similarity can be preserved.

Original: We denote the state encoder by φ ω : S → R n with parameter ω and the distance between states onthe embedding space R n by ˆ d ( φ ω ( s i ) , φ ω ( s j )) , e.g. ˆ d can be the L 2 norm distance.
Revised: We denote the state encoder by ϕ ω : S → R n with parameters ω and the distance between states on the embedding space R n by ˆ d ( ϕ ω ( s i ) , ϕ ω ( s j )) , e.g., ˆ d can be the L 2 norm distance.

Original: The problem ofrepresentation learning in terms of ω can be cast as a minimization problem of the following expectedsquared difference or loss between the distance on the embedding space, ˆ d ( φ ω ( s i ) , φ ω ( s j )) , and thecorresponding behavior metric, d π ( s i , s j ) , between any pair of states s i and s j :
Revised: The problem of representation learning in terms of ω can be cast as a minimization problem of the following expected squared difference or loss between the distance on the embedding space, ˆ d ( ϕ ω ( s i ) , ϕ ω ( s j )) , and the corresponding behavior metric, d π ( s i , s j ) , between any pair of states s i and s j :

Original: [Equation]
Revised: [Equation]

Original: To develop a practical algorithm to minimize the above objective, prior approaches adopt differentapproximation or relaxation strategies to make the resultant optimization problem computationallytractable.
Revised: To develop a practical algorithm to minimize the above objective, prior approaches adopt different approximation or relaxation strategies to make the resultant optimization problem computationally tractable.

Original: As discussed above, there are three common issues underlying most prior approaches,which introduce gaps between the practical algorithms and their theoretical prototype, i.e., (5), namelyloss function mismatch, relaxation of dynamics model divergence, andlimitation of the L 1 or the L 2norm.
Revised: As discussed above, there are three common issues underlying most prior approaches, which introduce gaps between the practical algorithms and their theoretical prototypes, i.e., (5), namely loss function mismatch, relaxation of dynamics model divergence, and the L 1 or the L 2 norm limitation.

Original: We discuss these 3 gaps in detail in the following sections.
Revised: We discuss these 3 gaps in detail in the following sections.

Original: 4.1 Loss Function Mismatch
Revised: 4.1 Loss Function Mismatch

Original: In general, to learn the encoder φ ω by minimizing (5) is computationally intractable or expensivebecause of the computation of the behavior metric d π .
Revised: In general, to learn the encoder ϕ ω by minimizing (5) is computationally intractable or expensive because of the computation of the behavior metric d π .

Original: We take MICo (MICo-based representationlearning [4]) as an example.
Revised: We take MICo (MICo-based representation learning [4]) as an example.

Original: By specifying the term d π ( s i , s j ) in (5) as the MICo distance deﬁned in(3), the loss of MICo becomes,
Revised: By specifying the term d π ( s i , s j ) in (5) as the MICo distance defined in (3), the loss of MICo becomes,

Original: [Equation]
Revised: [Equation]

Original: where ¯ ω  a copy of parameter for the target network.
Revised: where ¯ ω is a copy of parameters for the target network.

Original: However, the reward expectations E a i ∼ π r a
Revised: However, the reward expectations E a i ∼ π r a

Original: i s i and E a j ∼ π r a j s j in the 2nd term in (6) are computationally intractable and also difﬁcult to estimate evenbased on sampling.
Revised: i s i and E a j ∼ π r a j s j in the 2nd term in (6) are computationally intractable and also difficult to estimate even based on sampling.

Original: Thus, Castro et al.
Revised: Thus, Castro et al.

Original: [4] propose to approximate the loss in (6) with the followingalternative loss,
Revised: [4] propose to approximate the loss in (6) with the following alternative loss,

Original: [Equation]
Revised: [Equation]

Original: Note that in (7), the expectation operator on rewards is moved out of the absolute value operator ofthe difference between rewards to avoid the estimation of expectation over rewards, E a i ∼ π r a
Revised: Note that in (7), the expectation operator on rewards is moved out of the absolute value operator of the difference between rewards to avoid the estimation of expectation over rewards, E a i ∼ π r a

Original: i s i and
Revised: i s i and

Original: E a j ∼ π r a j s j  and enable sampling more efﬁcient.
Revised: E a j ∼ π r a j s j , and enable sampling more efficient.

Original: However, such a revision introduces a gap betweensolutions of minimizing (7) and (6),respectively, because in (7) the reference behavioral metric is nolonger the MICo distance but the “shift” MICo distance deﬁned as follows.
Revised: However, such a revision introduces a gap between solutions of minimizing (7) and (6), because in (7) the reference behavioral metric is no longer the MICo distance but the “shift” MICo distance defined as follows.

Original: Deﬁnition 4.1 (Shift MICo distance) .
Revised: Definition 4.1 (Shift MICo distance) .

Original: The shift MICo distance function ˜ F πM is deﬁned as
Revised: The shift MICo distance function ˜ F πM is defined as

Original: [Equation]
Revised: [Equation]

Original: Lemma 4.2 (Fixed-point) .
Revised: Lemma 4.2 (Fixed-point) .

Original: The shift MICo distance function ˜ F πM has a unique ﬁxed-point ˜ d πM .
Revised: The shift MICo distance function ˜ F πM has a unique fixed-point ˜ d πM .

Original: Proof.
Revised: Proof.

Original: (Sketch)
Revised: (Sketch)

Original: This can be proved by following the proof of Theorem 3.3 by using Banach’sﬁxed-point theorem.
Revised: This can be proved by following the proof of Theorem 3.3 by using Banach’s fixed-point theorem.

Original: A detailed proof can be found in Appendix.
Revised: A detailed proof can be found in Appendix A.1.

Original: The above lemma shows that the approximated distance on the embedding space ˆ d ( φ ω ( s i ) , φ ω ( s j ))still converges to the Shift MICo distance, ˜ d π M , by minimizing (7).
Revised: The above lemma shows that the approximated distance on the embedding space ˆ d ( ϕ ω ( s i ) , ϕ ω ( s j )) still converges to the Shift MICo distance, ˜ d πM , by minimizing (7).

Original: However, as E a i ∼ π a j ∼ π | r a i s i − r a j s j | ≥| E a i ∼ π r a
Revised: However, as E a i ∼ π a j ∼ π | r a i s i − r a j s j | ≥ | E a i ∼ π r a

Original: i s i − E a j ∼ π r a j s j | , we have the following theorem, whose proof can be found in Appendix.
Revised: i s i − E a j ∼ π r a j s j | , we have the following theorem, whose proof can be found in Appendix A.2.

Original: Theorem 4.3 (Looser value difference bound) .
Revised: Theorem 4.3 (Looser value difference bound) .

Original: Given states s i and s j , and a policy π , we have
Revised: Given states s i and s j , and a policy π , we have

Original: [Equation]
Revised: [Equation]

Original: Based on the above theorem, as the Shift MICo distance has looser value difference bound, it maybe less relevant to  value function.
Revised: Based on the above theorem, as the Shift MICo distance has looser value difference bound, it may be less relevant to the value function.

Original: As a result, the learned representation may not be able to encodebehavior similarity between states accurately.
Revised: As a result, the learned representation may not be able to encode the behavior similarity between states accurately.

Original: Apart from MICo, other behavioral metric-basedmethods, which consist of the on-policy reward difference term, such as DBC [29, 16] and AMBS [5],also suffer from a similar approximation gap.
Revised: Apart from MICo, other behavioral metric-based methods, which consist of the on-policy reward difference term, such as DBC [29, 16] and AMBS [5], also suffer from a similar approximation gap.

Original: 4.2 Relaxation of Dynamics Models Divergence
Revised: 4.2 Relaxation of Dynamics Models Divergence

Original: The π -bisimulation metric needs to compute  1-Wasserstein distance W 1 between dynamics modelsto measure the distribution distance.
Revised: The π -bisimulation metric needs to compute the 1-Wasserstein distance W 1 between dynamics models to measure the distribution distance.

Original: In π -bisimulation metric-based representation learning, one canlearn the encoder φ ω by minimizing the following loss,
Revised: In π -bisimulation metric-based representation learning, one can learn the encoder ϕ ω by minimizing the following loss,

Original: [Equation]
Revised: [Equation]

Original: However, the 1-Wasserstein distance is computationally expensive or intractable.
Revised: However, the 1-Wasserstein distance is computationally expensive or intractable.

Original: In DBC [29] the2-Wasserstein distance W 2 is proposed to replace W 1 , as W 2 has a convenient closed-form of a
Revised: In DBC [29] the 2-Wasserstein distance W 2 is proposed to replace W 1 , as W 2 has a convenient closed-form of a

Original: Gaussian distribution with respect to the L 2 distance.
Revised: Gaussian distribution with respect to the L 2 distance.

Original: Speciﬁcally, the loss function of DBC withbatched sampled transitions is deﬁned as,
Revised: Specifically, the loss function of DBC with batched sampled transitions is defined as,

Original: [Equation]
Revised: [Equation]

Original: (cid:21) , where (cid:107) · (cid:107) 2 is the L 2 norm, (cid:98) P is a dynamics model on the representation space, and ¯ π is the expectedpolicy output.
Revised: where ∥ · ∥ 2 is the L 2 norm, (cid:98) P is a dynamics model on the representation space, and ¯ π is the expected policy output.

Original: The use of W 2 almost breaks all theoretical guarantees for the bisimulation metric.
Revised: The use of W 2 almost breaks all theoretical guarantees for the bisimulation metric.

Original: Theexistence of unique ﬁxed-point in the bisimulation metric requires the continuity and monotonicityof W 1 with respect to d [7].
Revised: The existence of unique fixed-point in the bisimulation metric requires the continuity and monotonicity of W 1 with respect to d [7].

Original: The properties of continuity and monotonicity do not hold with W 2 .
Revised: The properties of continuity and monotonicity do not hold with W 2 .

Original: Therefore there is no more guarantee about the ﬁxed-point existence in DBC except that both thedynamics model and the policy π are deterministic, in which case W 2 ( d ) degenerates to d and
Revised: Therefore there is no more guarantee about the fixed-point existence in DBC except that both the dynamics model and the policy π are deterministic, in which case W 2 ( d ) degenerates to d and

Original: Banach’s ﬁxed-point exists [16].
Revised: Banach’s fixed-point exists [16].

Original: However, this assumption may be too strong to hold in practice.
Revised: However, this assumption may be too strong to hold in practice.

Original: Instead of using the family of Wasserstein distances, in MICo as shown in (3) the sample-baseddistribution divergence, E s (cid:48)
Revised: Instead of using the family of Wasserstein distances, in MICo as shown in (3) the sample-based distribution divergence, E s ′

Original: i ∼ P
Revised: i ∼ P

Original: π s i , s (cid:48)
Revised: π s i , s ′

Original: j ∼ P
Revised: j ∼ P

Original: π s j d( s (cid:48) i , s (cid:48) j ) , is used to measure the difference between dynamicsmodels.
Revised: π s j d s ′ i , s ′ j ) , is used to measure the difference between dynamics models.

Original: This sample-based distribution divergence can be considered as a Łukaszyk-Karmowskimetric.
Revised: This sample-based distribution divergence can be considered as a Łukaszyk-Karmowski metric.

Original: While a Wasserstein distance has zero self-distance, a Łukaszyk-Karmowski metric does notsatisfy the identity of indiscernibles.
Revised: While a Wasserstein distance has zero self-distance, a Łukaszyk-Karmowski metric does not satisfy the identity of indiscernibles.

Original: As a result, the approximated distance on the learned embeddingspace based on the MICo distance, which involves a Łukaszyk-Karmowski metric to measure distancebetween dynamics models, may also suffer from the violation issue of the identity of indiscernibles.
Revised: As a result, the approximated distance on the learned embedding space based on the MICo distance, which involves a Łukaszyk-Karmowski metric to measure distance between dynamics models, may also suffer from the violation issue of the identity of indiscernibles.

Original: 4.3 Limitation of Using the L 1 or the L 2 Norm on the Embedding Space
Revised: 4.3 Limitation of Using the L 1 or the L 2 Norm on the Embedding Space

Original: As mentioned in Section 4.1, to avoid the expensive or intractable computation of the expectationover rewards, prior approaches, such as MICo, and DBC, use an alternative term to measure thereward difference between states , E a i ∼ π a j ∼ π | r a
Revised: As mentioned in Section 4.1, to avoid the expensive or intractable computation of the expectation over rewards, prior approaches, such as MICo, and DBC, use an alternative term to measure the reward difference between states, E a i ∼ π a j ∼ π | r a

Original: i s i − r a j s j | .
Revised: i s i − r a j s j | .

Original: In the following, we discuss the case that states i and s j are identical.
Revised: In the following, we discuss the case that state s i and s j are identical.

Original: Lemma 4.4.
Revised: Lemma 4.4.

Original: If s i
Revised: If s i

Original: =
Revised: =

Original: s j , then E a i ∼ π a j ∼ π | r a i s i − r a j s j | ≥ 0 .
Revised: s j , then E a i ∼ π a j ∼ π | r a i s i − r a j s j | ≥ 0 .

Original: The equality holds only if the reward function ris constant w.r.t.
Revised: The equality holds only if the reward function r is constant w.r.t.

Original: action a
Revised: action a

Original: or the policy π is deterministic.
Revised: or the policy π is deterministic.

Original: Note that in most RL tasks, a stochastic policy is widely used for exploration and a reward functionis rarely constant w.r.t.
Revised: Note that in most RL tasks, a stochastic policy is widely used for exploration and a reward function is rarely constant w.r.t.

Original: actions.
Revised: actions.

Original: Therefore, in practice, E a i ∼ π a j ∼ π | r a i s i − r a j s j | > 0 for most RL tasks.
Revised: Therefore, in practice, E a i ∼ π a j ∼ π | r a i s i − r a j s j | > 0 for most RL tasks.

Original: Besides, as mentioned in Section 4.2, a Łukaszyk-Karmowski metric measuring distance betweendynamics models does not satisfy the identity of indiscernibles.
Revised: Besides, as mentioned in Section 4.2, a Łukaszyk-Karmowski metric measuring distance between dynamics models does not satisfy the identity of indiscernibles.

Original: Therefore, the behavioral metric,which is a sum of the reward difference term and the dynamics model distance term (no matter a
Revised: Therefore, the behavioral metric, which is a sum of the reward difference term and the dynamics model distance term (no matter a

Original: Wasserstein distance or a Łukaszyk-Karmowski metric), can be greater than zero on a pair of identicalstates.
Revised: Wasserstein distance or a Łukaszyk-Karmowski metric), can be greater than zero on a pair of identical states.

Original: Note that both  L 1 and  L 2 norms satisfy the identity of indiscernibles, i.e. when x i = x j , (cid:107) x
Revised: Note that both the L 1 and the L 2 norms satisfy the identity of indiscernibles, i.e. when x i = x j , ∥ x

Original: i − x j (cid:107) 1 =(cid:107) x i − x j (cid:107) 2 = 0 .
Revised: i − x j ∥ 1 = ∥ x i − x j ∥ 2 = 0 .

Original: If we leverage the L 1 or the L 2 norm as the form of distance ˆ d on the embeddingspace to approximate the behavioral metric, then the identical states pairs has zero distance on embedding space.
Revised: If we leverage the L 1 or the L 2 norm as the form of distance ˆ d on the embedding space to approximate the behavioral metric, then the identical states pairs has zero distance on embedding space.

Original: However, the regression target, i.e., the behavioral metric, isusually greater thanzero on identical states.
Revised: However, the regression target, i.e., the behavioral metric, is greater than zero on identical states.

Original: As a result, when minimizing regression loss between the L 1 / L 2 norm andthe behavioral metric, the representations for similar states, especially identical ones, will be pushedapart in the embedding space.
Revised: As a result, when minimizing regression loss between the L 1 / L 2 norm and the behavioral metric, the representations for similar states, especially identical ones, will be pushed apart in the embedding space.

Original: Recently, AMBS [5] proposes to use neural networks to measure distance on state embedding spacerather than using the L 1 or the L 2 norm.
Revised: Recently, AMBS [5] proposes to use neural networks to measure distance on state embedding space rather than using the L 1 or the L 2 norm.

Original: However, in this case, the learned “behavior metric” does nothave theoretical support such as the ﬁxed-point convergence guarantee.
Revised: However, in this case, the learned “behavior metric” does not have theoretical support such as the fixed-point convergence guarantee.

Original: Besides, MICo proposes a newform of distance on the embedding space, which is a sum of angular distance between embeddingsand L 2 norm of the embeddings.
Revised: Besides, MICo proposes a new form of distance on the embedding space, which is a sum of angular distance between embeddings and the L 2 norm of the embeddings.

Original: Note that the proposed  distance has non-zero self-distance.
Revised: Note that the proposed wo distance has non-zero self-distance.

Original:  The Proposed RAP Distance
Revised:  The Proposed RAP Distance

Original: We ﬁrstly propose a new behavioral metric to measure the state similarity without computing the
Revised: We firstly propose a new behavioral metric to measure the state similarity without computing the

Original: Wasserstein distance between dynamics models in Section 5.1.
Revised: Wasserstein distance between dynamics models in Section 5.1.

Original: The proposed behavioral metricnamely the RAP distance enjoys theoretical properties such as ﬁxed-point existence and a valuedifference bound.
Revised: The proposed behavioral metric namely the RAP distance enjoys theoretical properties such as fixed-point existence and a value difference bound.

Original: We then present a practical algorithm to learn the state encoder by approximatingthe RAP distance on the state embedding space in Section 5.2.
Revised: We then present a practical algorithm to learn the state encoder by approximating the RAP distance on the state embedding space in Section 5.2.

Original: Our algorithm uses the learnedestimation of reward functions and dynamics models to provide distance approximation which isconsistent to the behavioral metric.
Revised: Our algorithm uses the learned estimation of reward functions and dynamics models to provide distance approximation which is consistent to the behavioral metric.

Original: It addresses all the aforementioned approximation gaps andpreserves the theoretical guarantee about the value function difference bound.
Revised: It addresses all the aforementioned approximation gaps and preserves the theoretical guarantee about the value function difference bound.

Original: Particularly, theapproximation gap of loss function mismatch is addressed in Section 5.2, the relaxation of dynamicsmodel divergence is addressed in Section 5.1 and the limitation of the L 1 or the L 2 norm on theembedding space is addressed in Section 5.3.
Revised: Particularly, the approximation gap of loss function mismatch is addressed in Section 5.2, the relaxation of dynamics model divergence is addressed in Section 5.1 and the limitation of the L 1 or the L 2 norm on the embedding space is addressed in Section 5.3.

Original: 5.1 Deﬁnition and Properties of the RAP Distance
Revised: 5.1 Definition and Properties of the RAP Distance

Original: In order to avoid the high computational cost of W 1 or the approximation gap introduced by relaxationto W 2 as described in Section 4.2, we consider a distance measure between dynamics models
Revised: In order to avoid the high computational cost of W 1 or the approximation gap introduced by relaxation to W 2 as described in Section 4.2, we consider a distance measure between dynamics models

Original: P π s i and
Revised: P π s i and

Original: P π s j with sampling.
Revised: P π s j with sampling.

Original: To be speciﬁc, our on-policy behavioral distance is deﬁned as follows.
Revised: To be specific, our on-policy behavioral distance is defined as follows.

Original: Deﬁnition 5.1 (the RAP distance) .
Revised: Definition 5.1 (the RAP distance) .

Original: Let M be the space of distance function d : S × S
Revised: Let M be the space of distance function d : S × S

Original: → R , the RAPdistance function F πG : M → M is deﬁned as,
Revised: → R , the RAP distance function F πG : M → M is defined as,

Original: [Equation]
Revised: [Equation]

Original: 9) where E [ s (cid:48) i ]
Revised: where E [ s ′ i ]

Original: =
Revised: =

Original: E s (cid:48)
Revised: E s ′

Original: i ∼ P a i s i[ s (cid:48)
Revised: i ∼ P a i s i

Original:  i ] is the expectation value of next state over the dynamics model P a
Revised: [ s ′ i ] is the expectation value of next state over the dynamics model P a

Original: i s i .
Revised: i s i .

Original: We design the behavioral distance by measuring the expected states over dynamics models recursively,which removes the requirement of sampling on P
Revised: We design the behavioral distance by measuring the expected states over dynamics models recursively, which removes the requirement of sampling on P

Original: a i s i and P
Revised: a i s i and P

Original: a j s j but only performs sampling on thepolicy π .
Revised: a j s j but only performs sampling on the policy π .

Original: In practice, the expected next states are generated by a learnable approximated dynamicsmodel as described in Section.
Revised: In practice, the expected next states are generated by a learnable approximated dynamics model as described in Section.

Original: Theorem 5.2.
Revised: Theorem 5.2.

Original: F πG is a contraction mapping w.r.t.
Revised: F πG is a contraction mapping w.r.t.

Original: the L ∞ norm and has a unique ﬁxed-point d πG .
Revised: the L ∞ norm and has a unique fixed-point d πG .

Original: Proof.
Revised: Proof.

Original: Let d, d (cid:48) ∈ M .
Revised: Let d, d ′ ∈ M .

Original: We have
Revised: We have

Original: [Equation]
Revised: [Equation]

Original: Therefore, F πG is a contraction mapping w.r.t.
Revised: Therefore, F πG is a contraction mapping w.r.t.

Original: the L ∞ norm and there exists a unique ﬁxed-point for
Revised: the L ∞ norm and there exists a unique fixed-point for

Original: F πG due to Banach’s ﬁxed-point theorem.
Revised: F πG due to Banach’s fixed-point theorem.

Original: This completes the proof.
Revised: This completes the proof.

Original: Theorem 5.2 provides  convergence guarantee for the RAP distance that by iterating F πG , distance dwill converge to the ﬁxed-point d πG .
Revised: Theorem 5.2 provides a convergence guarantee for the RAP distance that by iterating F πG , distance d will converge to the fixed-point d πG .

Original: Theorem 5.3 (Value function difference bound) .
Revised: Theorem 5.3 (Value function difference bound) .

Original: Given states s i and s j  and a policy π , we have
Revised: Given states s i and s j , and a policy π , we have

Original: [Equation]
Revised: [Equation]

Original: The proof can be found in Appendix.
Revised: The proof can be found in Appendix A.3.

Original: Theorem 5.3 demonstrates that the RAP distance between statesupper-bounds the difference of their states values.
Revised: Theorem 5.3 demonstrates that the RAP distance between states upper-bounds the difference of their states values.

Original: 5.2 Approximation of RAP
Revised: 5.2 Approximation of RAP

Original: A straightforward way to learn a representation encoder to approximate the RAP distance on theembedding space is to minimize the loss in (5).
Revised: A straightforward way to learn a representation encoder to approximate the RAP distance on the embedding space is to minimize the loss in (5).

Original: However, the approximation gap of loss functionmismatch as mentioned in Section 4.1 still exists if we relax the ﬁrst term in (9) as E a i ∼ π a j ∼ π | r a i s i − r a j s j | .
Revised: However, the approximation gap of loss function mismatch as mentioned in Section 4.1 still exists if we relax the first term in (9) as E a i ∼ π a j ∼ π | r a i s i − r a j s j | .

Original: This leads to the learned metric having a looser value difference bound than the original behavioralmetric as proved in Section 4.1.
Revised: This leads to the learned metric having a looser value difference bound than the original behavioral metric as proved in Section 4.1.

Original: Here, we propose another alternative relaxation of | E a i ∼ π r a
Revised: Here, we propose another alternative relaxation of | E a i ∼ π r a

Original: i s i −
Revised: i s i −

Original: E a j ∼ π r a j s j | to address this issue.
Revised: E a j ∼ π r a j s j | to address this issue.

Original: Let r s be a random variable over the action distribution deﬁnedby p ( r s = r as )
Revised: Let r s be a random variable over the action distribution defined by p ( r s = r as )

Original: =
Revised: =

Original: π ( a | s ) .
Revised: π ( a | s ) .

Original: We ﬁrst analyze the difference between E a i ∼ π a j ∼ π (cid:104)(cid:12)(cid:12) r ai s i − r a j s j (cid:12)(cid:12) 2 (cid:105)
Revised: We first analyze the difference between E a i ∼ π a j ∼ π (cid:104)(cid:12)(cid:12) r a

Original: and| E a i
Revised: [Equation]

Original: ∼ π r a i s i −E a j ∼ π r a j s j | 2 , E
Revised: s i − r a j s j (cid:12)(cid:12) 2 (cid:105) and

Original: [Equation]
Revised: [Equation]

Original: = 
Revised: = V ar [ r s i ] + V ar [ r s j ] − 2 Cov

Original: [Equation]
Revised: [ r s i , r s j ] .

Original: Since r s i and r s j are independent, Cov
Revised: Since r s i and r s j are independent

Revised: [Equation]

Original:  [ r s i , r s j ] = 0 .
Revised: Cov [ r s i , r s j ] = 0 .

Original: Therefore, we have | E a i ∼ π
Revised: Therefore [Equation]

Original: r a i s i − E a j ∼ π r
Revised: we have the reward difference term

Original: a j s j | = 257 (cid:114) E a i ∼ π a j ∼ π (cid:104)(cid:12)(cid:12) r a i s i − r a j s j (cid:12)(cid:12) 2 (cid:105) − V ar [ r s i ] − V ar
Revised: [Equation]

Original: [ r s j ] and the revised RAP distance at ﬁxed-point as
Revised: and the revised RAP distance at fixed-point as

Original: [Equation]
Revised: [Equation]

Original: In (12), we successfully move the expectation operator on rewards out of the absolute value operatorwithout introducing any approximation gap caused by loss function mismatch.
Revised: In (12), we successfully move the expectation operator on rewards out of the absolute value operator without introducing any approximation gap caused by loss function mismatch.

Original: However, thereare three issues that need to be further addressed: 1) the square root introduces new bias undersampling, 2) the variances V ar [ r s i ] and V ar [ r s j ] are intractable to compute, and 3) how to estimatethe expected next states E s (cid:48)
Revised: However, there are three issues that need to be further addressed: 1) the square root introduces new bias under sampling, 2) the variances V ar [ r s i ] and V ar [ r s j ] are intractable to compute, and 3) how to estimate the expected next states E s ′

Original: i ∼ P
Revised: i ∼ P

Original: [Equation]
Revised: a i s i [ s ′ i ] and E s ′ i ∼ P

Original:  [ s [Equation]
Revised: a j s j [ s ′ j ] .

Original: i ] and E s
Revised: In order to reduce the bias issue introduced by the square root, we try to remove the square root in the loss of learning the RAP distance. We move the dynamics term in (12) to the left-hand side, then take square on both sides and get

Original: [Equation]
Revised: [Equation]

Original:  s 
Revised: Let ˆ d ( ϕ ω ( s i ) , ϕ ω ( s j )) be the approximated RAP distance between s i and s j parameterized by ω . The loss for learning ˆ d ( ϕ ω ( s i ) , ϕ ω ( s j )) is to minimize the mean squared error between the left-hand side and the right-hand side in (13):

Original: [Equation]
Revised: [Equation]

Revised: Such a loss consists of a pair of squared difference terms as the regression target, which approximate the squared difference between distances of current states and distances of dynamics models, respectively.

Original: Though the reward variance V ar [ r s ] is computationally intractable,  we can learn a neural networkapproximator to estimate it by assuming that the reward r s on state s is Gaussian distributed.
Revised: The reward variance V ar [ r s ] is computationally intractable, but we can learn a neural network approximator to estimate it by assuming that the reward r s on state s is Gaussian distributed.

Original: Letr ψ be the learned reward function  parameterized by ψ , which outputs a Gaussian distribution,r ψ ( s ) = { (cid:98) µ ( r s ) , (cid:98) σ ( r s ) } , where (cid:98) µ ( r s ) and (cid:98) σ ( r s ) are the mean and the standard deviation, respectively.
Revised: Let r ψ be the learned reward function approximation parameterized by ψ , which outputs a Gaussian distribution, r ψ ( s ) = { (cid:98) µ ( r s ) , (cid:98) σ ( r s ) } , where (cid:98) µ ( r s ) and (cid:98) σ ( r s ) are the mean and the standard deviation, respectively.

Original: Similarly, in order to estimate the expected next states E s (cid:48) ∼ P as [ s (cid:48) ] with a neural network approximator on the embedding space, we learn a dynamics model (cid:98) P taking input of state embedding φ ( s ) and action a and outputs a Gaussian distribution over the next state embedding,(cid:98) P ( φ ω ( s ) , a ) = { (cid:98) µ ( (cid:98) P a φ ω ( s ) ) , (cid:98) σ ( (cid:98) P a φ ω ( s ) ) } , where (cid:98) µ ( (cid:98) P a φ ω ( s ) ) and (cid:98) σ ( (cid:98) P a φ ω ( s ) ) are the mean vector andthe standard deviation vector for predictive  next state embedding, respectively.Let ˆ d ( φ ω ( s i ) , φ ω ( s j )) be the approximated RAP distance between s i and s j parameterized byω .
Revised: Similarly, in order to estimate the expected next states E s ′ ∼ P as [ s ′ ] with a neural network approximator on the embedding space, we learn a dynamics model (cid:98) P taking input of state embedding ϕ ( s ) and action a and outputs a Gaussian distribution over the next state embedding, (cid:98) P ( ϕ ω ( s ) , a ) = { (cid:98) µ ( (cid:98) P a ϕ ω ( s ) ) , (cid:98) σ ( (cid:98) P a ϕ ω ( s ) ) } , where (cid:98) µ ( (cid:98) P a ϕ ω ( s ) ) and (cid:98) σ ( (cid:98) P a ϕ ω ( s ) ) are the mean vector and the standard deviation vector for predictive embeddings of the next state, respectively.

Original: The RAP distance between expected next states, d πG ( E [ s (cid:48) i ] , E [ s (cid:48) j ]) , can be approximated byˆ d ( (cid:98) µ ( (cid:98) P a 
Revised: The RAP distance between expected next states, d πG ( E [ s ′ i ] , E [ s ′ j ]) , can be approximated by ˆ d ( (cid:98) µ ( (cid:98) P a i ϕ ω

Original: i φ ω ( s i ) ) , (cid:98) µ ( (cid:98) P a j φ ω ( s j ) )) .
Revised: ( s i ) ) , (cid:98) µ ( (cid:98) P a j ϕ ω ( s j ) )) .

Original: The variance V ar [ r s ] can be replaced by the learned reward functionoutput as ( (cid:98) σ ( r s )) 2 .
Revised: The variance V ar [ r s ] can be replaced by the learned reward function output as ( (cid:98) σ ( r s )) 2 .

Original: In order to reduce the biased issue introduced by the square root, we propose anovel regression loss deﬁned as,
Revised: By replacing the dynamics term and reward variance terms in (14), we propose the RAP loss defined as

Original: [Equation]
Revised: [Equation]

Original: where ¯ ω is a copy of parameter with stop gradient and φ ¯ ω is the encoder in the target network. The lossin (13) consists of a pair of squared difference terms as the regression target, which approximate thesquared difference between distances of current states and distances of dynamics models, respectively.
Revised: where ¯ ω is a copy of parameter with stop gradient and ϕ ¯ ω is the encoder in the target network and D is the replay buffer or the set of data that RL algorithm, e.g. SAC [10], learns from. The loss (15) is trained over the transitions sampled from D .

Original: 5.3 Explicit Form of Distance on the Embedding Space
Revised: 5.3 Explicit Form of Distance on the Embedding Space

Original: As discussed in Section 4.3, the behavioral metric is usually with non-zero self-distance.
Revised: As discussed in Section 4.3, the behavioral metric is usually with non-zero self-distance.

Original: Besides, thedistance measured between the next-state distributions in the RAP distance is a Łukaszyk-Karmowski metric [22] which also has non-zero self-distance.
Revised: Besides, the distance measured between the next-state distributions in the RAP distance is a Łukaszyk-Karmowski metric [22], which also has non-zero self-distance.

Original: Here, we adopt the approximation form of distanceon the embedding space as proposed in MICo [4].
Revised: Here, we adopt the approximation form of distance on the embedding space as proposed in MICo [4].

Original: Deﬁnition 5.4 (MICo approximation) .
Revised: Definition 5.4 (MICo approximation) .

Original: Let ˆ d G : R n × R n → R be a distance function on representation space R n , x , y ∈ R n and k > 0 .
Revised: Let ˆ d G : R n × R n → R be a distance function on representation space R n , x , y ∈ R n and k > 0 .

Original: ˆ d G is deﬁned as ˆ d G ( x , y ) 
Revised: ˆ d G is defined as ˆ d G ( x , y ) =

Original: = (cid:107) x (cid:107) 22 + (cid:107) y [-(cid:107) 
-]
Revised: ∥ x ∥ 22 + ∥ y [-(cid:107) 
-]∥

Original: + kθ ( x , y ) , where θis angular function evaluating the absolute angle between vectors x and y , and k is a hyperparameter.
Revised: + kθ ( x , y ) , where θ is angular function evaluating the absolute angle between vectors x and y , and k is a hyperparameter.

Original: In practice, we set k = 10 − 5 for our method.
Revised: In practice, we set k = 10 − 5 for our method.

Original: As the above form of distance produces non-zero self-distance, with ˆ d G the approximation of the
Revised: As the above form of distance produces non-zero self-distance, with ˆ d G the approximation of the

Original: RAP distance will not push apart similar states on the embedding space.
Revised: RAP distance will not push apart similar states on the embedding space.

Original: Lemma 5.5 (Non-zero self-distance) .
Revised: Lemma 5.5 (Non-zero self-distance) .

Original: The self-distance of ˆ d G is not restrict to zero:
Revised: The self-distance of ˆ d G is not restrict to zero:

Original: ˆ d G ( x , x ) =2 (cid:107){+= 
 ∥+} x (cid:107) 22 ≥ 0 .
Revised: ˆ d G ( x , x ) {+= 
 ∥+} x ∥ 22 ≥ 0 .

Original: The equality holds only if all the elements of x are zero.
Revised: The equality holds only if all the elements of x are zero.

Revised: 5.4 Implementation

Original: [Figure]
Revised: [Figure]

Original: The network architecture of our method is shown in Figure 1.
Revised: The network architecture of our method is shown in Figure 1.

Original: Ourmethod is built upon SAC [10].
Revised: Our method is built upon SAC [10].

Original: Actor and critic networks take inputof state representation φ ω ( s ) .
Revised: Actor and critic networks take input of state representation ϕ ω ( s ) .

Original: The SAC objectives and the RAP regression loss L RAP ( φ ω ) are optimized jointly.
Revised: The SAC objectives and the RAP regression loss L RAP ( ϕ ω ) are optimized jointly.

Original: More implementationdetails can be found in Appendix.
Revised: More implementation details can be found in Appendix B.

Revised: The source code is available at https://github.com/jianda-chen/RAP_distance .

Original:  Experiment
Revised:  Experiment

Original: [Figure] In this section, we evaluate the efﬁciency, robustness and generalization ability of the representationlearned by our method on three RL benchmarks: 1) Distracting DeepMind Control Suite (DMC) [24],2) Robosuite [30], and 3) CARLA [6].
Revised: In this section, we evaluate the efficiency, robustness and generalization ability of the representation learned by our method on three RL benchmarks: 1) Distracting DeepMind Control Suite (DMC) [24], 2) Robosuite [30], and 3) CARLA [6].

Original: Thesethree domains are  continuous action space controltask with visual input.
Revised: These are all control tasks in continuous action spaces with visual input.

Original: We compare our methods with several baselines and state-of-the-art methodsincluding metric-based representation learning and data augmentation: 1) SAC [10], a baseline RLmethod for continuous control, 2) MICo [4], a sample-based behavioral metric-based representationlearning method for RL, 3) DBC [29], a representation learning method by approximating thebisimulation metric with the L 1 norm, 4) RobustDBC [16], a DBC-styled method with intrinsicrewards and inverse dynamics, and 5) DrQ [28], an image augmentation method on pixel inputs RL.
Revised: We compare our method with several baselines and state-of-the-art methods including metric-based representation learning and data augmentation: 1) SAC [10], a baseline RL method for continuous control, 2) MICo [4], a sample-based behavioral metric-based representation learning method for RL, 3) DBC [29], a representation learning method by approximating the bisimulation metric with the L 1 norm, 4) RobustDBC [16], a DBC-styled method with intrinsic rewards and inverse dynamics, and 5) DrQ [28], an image augmentation method on pixel inputs RL.

Original: [Figure]
Revised: [Figure]

Original: Distracting DeepMind Control Suite To evaluate the generalization ability and robustness of ourmethod, we perform experiments on 2 settings, original background and natural video backgroundsettings, in DMC [24].
Revised: Distracting DeepMind Control Suite To evaluate the generalization ability and robustness of our method, we perform experiments on 2 settings, original background and natural video background settings, in DMC [24].

Original: We render 84 × 84 pixels as observation at each time step and stack framesas state.
Revised: We render 84 × 84 pixels as observation at each time step and stack frames as state.

Original: In the original background setting, we use the default background provided by DMC(as shown in Figure 2a).
Revised: In the original background setting, we use the default background provided by DMC (as shown in Figure 2a).

Original: For the natural video background distracting setting, we follow [29] toreplace the background with natural video sampled from Kinetics dataset [15] (as shown in Figure 2b).
Revised: For the natural video background distracting setting, we follow [29] to replace the background with natural video sampled from Kinetics dataset [15] (as shown in Figure 2b).

Original: We sample 1000 continuous frames from the video dataset as background for training RL agentsand evaluate agents on another 1000 continuous frames.
Revised: We sample 1000 continuous frames from the video dataset as background for training RL agents and evaluate agents on another 1000 continuous frames.

Original: The video background is considered asdistraction to RL algorithm.
Revised: The video background is considered as distraction to RL algorithm.

Original: For each setting, we train and evaluate on 4 tasks: cartpole-swingup,cheetah-run, ﬁnger-spin, and walker-walk.
Revised: For each setting, we train and evaluate on 4 tasks: cartpole-swingup, cheetah-run, finger-spin, and walker-walk.

Original: Each task is trained for 1 million environment steps.
Revised: Each task is trained using 1 million environment steps.

Revised: curve is average on 3 runs.

Original: Figure 3 is the training curves in original background settings.
Revised: Figure 3 shows the training curves in original background settings.

Original: It conﬁrms that our method RAP canaccelerate the training and has comparable results to data augmentation method DrQ. Figure 4 isthe training curves in natural video background settings.
Revised: It confirms that our method RAP can accelerate the training and has comparable results to data augmentation method DrQ. Figureshows the training curves in natural video background settings.

Original: Our method RAP outperforms DBC and MICo on all 4 tasks.
Revised: Our method RAP outperforms DBC and MICo on all 4 tasks.

Original: It also converges to higher score than DrQ in 1 task and learns much fasterin 3 tasks.
Revised: It also converges to higher score than DrQ in 1 task and learns much faster in 3 tasks.

Original: The new SOTA results verify that our method of reducing the gaps in approximation ofbehavioral distance can improve the efﬁciency and robustness for deep RL and they also conﬁrm thatour method is able to learn generalizable state representation in complex environments.
Revised: The new SOTA results verify that our method of reducing the gaps in approximation of behavioral distance can improve the efficiency and robustness for deep RL and they also confirm that our method is able to learn generalizable state representation in complex environments.

Original: Robosuite Robosuite [30] is a robotic simulatorwith various manipulation tasks based on MuJoCoengine.
Revised: Robosuite Robosuite [30] is a robotic simulator with various manipulation tasks based on MuJoCo engine.

Original: To evaluate the robustness of methods,we use the distraction settings for Robosuite byrandomizing the colors of robot arms and table,the lighting source and luminance, and the cameraposition.
Revised: To evaluate the robustness of methods, we use the distraction settings for Robosuite by randomizing the colors of robot arms and table, the lighting source and luminance, and the camera position.

Original: The color, lighting and camera positionare changed at the beginning of an episode but are consistent within the episode.
Revised: The color, lighting and camera position are changed at the beginning of an episode but are consistent within the episode.

Original: Our experiment isperformed by controlling the Panda arms.
Revised: Our experiment is performed by controlling the Panda arms.

Original: The control rate is 20Hz and the episode length is 500 steps.
Revised: The control rate is 20Hz and the episode length is 500 steps.

Original: We evaluate on two tasks: Door Opening and Two Arm Peg-In-Hole.
Revised: We evaluate on two tasks: Door Opening and Two Arm Peg-In-Hole.

Original: Figure 5 shows illustrations ofobservations of two tasks.
Revised: Figure 5 shows illustrations of observations of two tasks.

Original: We compare our method with behavioral metric methods MICo and DBC,and data augmentation method DrQ. Table 1 shows the training scores in 500k environment steps.
Revised: We compare our method with behavioral metric methods MICo and DBC, and data augmentation method DrQ. Table 1 shows the training scores in 500k environment steps.

Original: Our method RAP outperforms the others.
Revised: Our method RAP outperforms the others.

Original: The converge scores are around 7 × and 1.4 × comparedto the second-best method in Door Opening and Two Arm Peg-In-Hole, respectively.
Revised: The converge scores are around 7 × and 1.4 × compared to the second-best method in Door Opening and Two Arm Peg-In-Hole, respectively.

Original: It shows therobustness of our RAP in environments with random distractions.
Revised: It shows the robustness of our RAP in environments with random distractions.

Original: [Table]
Revised: [Table]

Original: CARLA In order to validate the generalization ability and learning efﬁciency in natural scenarios, we construct experiments on an autonomousdriving simulator CARLA [6], which provides 3D realistic on-worldscenarios.
Revised: CARLA In order to validate the generalization ability and learning efficiency in natural scenarios, we construct experiments on an autonomous driving simulator CARLA [6], which provides 3D realistic onworld scenarios.

Original: The goal of this task is to control a vehicle driving as far aspossible on a high-way map (Town 4) in 1000 time steps.
Revised: The goal of this task is to control a vehicle driving as far as possible on a high-way map (Town 4) in 1000 time steps.

Original: The rewardfunction followed [29] is designed to encourage driving far and avoidcollisions with other vehicles.
Revised: The reward function followed [29] is designed to encourage driving far and avoid collisions with other vehicles.

Original: The observation is formed as 420 × 84pixels of a 300 degrees ego-centric view, constructed by concatenatingﬁve cameras.
Revised: The observation is formed as 420 × 84 pixels of a 300 degrees ego-centric view, constructed by concatenating five cameras.

Original: In order to evaluate the generalization ability, we randomlysample a kind of weather in each episode starts.
Revised: In order to evaluate the generalization ability, we randomly sample a kind of weather in each episode starts.

Original: The weather (sunlight,rain, etc.), which affects the visual observation, can be considered as real-world distraction to RLagents.
Revised: The weather (sunlight, rain, etc.), which affects the visual observation, can be considered as real-world distraction to RL agents.

Original: Figure 6 shows the comparison results of our RAP and other SOTA methods.
Revised: Figure 6 shows the comparison results of our RAP and other SOTA methods.

Original: Our methodachieves comparable scores with data augmentation method DrQ, and its score is higher than behavioral metric DBC and MICo.
Revised: Our method achieves comparable scores with data augmentation method DrQ, and its score is higher than behavioral metric DBC and MICo.

Original: It shows that our RAP is able to generalize well in tasks with real-worldscenarios.
Revised: It shows that our RAP is able to generalize well in tasks with real-world scenarios.

Original:  Conclusion
Revised:  Conclusion

Original: Representation learning is one of the most critical problems in high-dimensional deep RL.
Revised: Representation learning is one of the most critical problems in high-dimensional deep RL.

Original: In thispaper, we propose a new behavioral metric and a practical representation learning algorithm on top ofthe new behavioral metric for deep RL.We provide theoretical analysis for our proposed metric aswell as the representation learning algorithm. We conduct empirical studies on multiple RL domainsto verify the effectiveness of our proposed method. Checklist
Revised: In this paper, we propose a new behavioral metric and a practical representation learning algorithm on top of the new behavioral metric for deep RL.

Original: The checklist follows the references.

Original: Please read the checklist guidelines carefully for information onhow to answer these questions.

Original: For each question, change the default [TODO] to [Yes] , [No] or[N/A] .

Original: You are strongly encouraged to include a justiﬁcation to your answer , either by referencingthe appropriate section of your paper or providing a brief inline description.

Original: For example:

Original: • Did you include the license to the code and datasets?

Original: [Yes] See Section ?? .•

Original: Did you include the license to the code and datasets?

Original: [No] The code and the data areproprietary.•

Original: Did you include the license to the code and datasets?

Original: [N/A]

Original: Please do not modify the questions and only use the provided macros for your answers.

Original: Note that the

Original: Checklist section does not count towards the page limit.

Original: In your paper, please delete this instructionsblock and only keep the Checklist section heading above along with the questions/answers below.

Original: For

Original: all

Original: authors...

Original: (a) Do the main claims made in the abstract

Original: and introduction

Original: accurately reﬂect the paper’s

Original: contributions and scope?

Original: [Yes]

Original: (b) Did you describe the limitations of your work?

Original: [No]

Original: (c) Did you discuss any potential negative societal impacts of your work?

Original: [N/A]

Original: (d) Have you read the ethics review guidelines and ensured that your paper conforms to

Original: them?

Original: [Yes]

Original: If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results? [Yes] (b) Did you include complete proofs of all theoretical results?
Revised: We provide theoretical analysis for our proposed metric as well as the representation learning algorithm.

Original: [Yes] In Appendix.

Original: If you ran experiments...

Original: (a) Did you include the code, data, and instructions needed to reproduce the main experi-

Original: mental results (either in the supplemental material or as a URL)?

Original: [No] We will publishcode once paper is accepted.

Original: (b) Did you specify all the training details (e.g., data splits, hyperparameters, how theywere chosen)?

Original: [Yes] In Appendix.

Original: (c) Did you report error bars (e

Original: g., with respect to the random seed after running experiments multiple times)?

Original: [Yes]

Original: (d) Did you include the total amount of compute and the type of resources used (e

Original: g., typeof GPUs, internal cluster, or cloud provider)?

Original: [Yes] In Appendix.4.

Original: If you are using existing assets (e

Original: g., code, data, models) or curating/releasing new assets...

Original: (a) If your work uses existing assets, did you cite the creators?

Original: [Yes]

Original: (b) Did you mention the license of the assets?

Original: [N/A]

Original: (c) Did you include any new assets either in the supplemental material or as a URL?

Original: [No]

Original: (d) Did you discuss whether and how consent was obtained from people whose data you’re

Original: using/curating?

Original: [N/A]

Original: (e) Did you discuss whether the data you are using/curating contains personally identiﬁable

Original: information or offensive content?

Original: [N/A]

Original: If you used crowdsourcing or conducted research with human subjects...

Original: (a) Did you include the full text of instructions given to participants and screenshots, ifapplicable?

Original: [No]

Original: (b) Did you describe any potential participant risks, with links to Institutional Review

Original: Board (IRB) approvals, if applicable?

Original: [N/A]

Original: (c) Did you include the estimated hourly wage paid to participants and the total amount

Original: spent on participant compensation?

Original: [N/A] A ProofsA.1 Proof of Lemma 4.  Lemma 4.2 (Fixed-point of shift MICo) . The “shift” MICo distance function ˜ F πM has a uniqueﬁxed-point ˜ d πM .
Revised: We conduct empirical studies on multiple RL domains to verify the effectiveness of our proposed method.

Original: Proof.

Original: The shift MICo distance function ˜ F πM is deﬁned as,

Original: [Equation]

Original: Let d, d (cid:48) ∈ M .

Original: We have

Original: [Equation]

Original: ˜ F πM is a contraction mapping w.r.t.

Original: the L ∞ norm and there exists a unique ﬁxed-point ˜ d πM for ˜ F πMdue to Banach’s ﬁxed-point theorem.

Original: This completes the proof.

Original: A.2 Proof of Theorem 4.

Original: Lemma A.

Original: [Equation]

Original: Since r a i s i and r a j s j are rewards which are scalars, we have | r a i s i − r a j s j | ≥ r a i s i − r a j s j , and | r a i s i − r a j s j | ≥r a j s j − r a

Original: i s i by symmetric.

Original: By taking expectation over a i and a j , we get,

Original: [Equation]

Original: ) and

Original: [Equation]

Original: Combine (14) and (15), we have

Original: [Equation]

Original: Lemma A.2 (Lifted MDP of MICo [4]) .

Original: The MICo metric function F πM is the Bellman operator fora Lifted MDP.

Original: Proof.

Original: Let the MDP for RL deﬁned by a tuple (cid:104)S , A , R , P, γ (cid:105) , we consider a lifted MDP constructedby a tuple (cid:104) ˆ S , ˆ A , ˆ R , ˆ P, γ (cid:105) , where state space ˆ S = S × S , action space ˆ A

Original: = A × A , transitiondistribution ˆ P ( a i , a j ) ( s i , s j ) = P a i s i P a j s j , and reward function
Revised: Acknowledgement and Funding Disclosure

Original: ˆ R (( s i , s j )) =

Original: | E a i ∼ π r a

Original: i s i − E a j ∼ π r

Original: a j s j | .

Original: The

Original: Bellman operator T ˆ πM under policy ˆ π ( a i , a j | s i , s j ) =

Original: π ( a i | s i ) π ( a j | s j ) is:

Original: T ˆ πM ( d πM )(( s i , s j ))

Original: = (cid:88)

Original: [Equation]

Original: = F πM ( d πM )( s i

Original: [Equation]

Original: s j ) .

Original: Lemma A.3 (Lifted MDP of shift MICo) .

Original: The shift MICo metric function ˜ F πM is the Bellmanoperator for a lifted MDP but with different reward function.

Original: Proof.

Original: We construct the lifted MDP of shift MICo by a tuple (cid:104) ˆ S , ˆ A , ˜ˆ R , ˆ P, γ (cid:105) which is same asthe lifted MDP in Lemma A.2 except for reward function

Original: ˜ˆ R (( s i , s j ) , ( a i , a j )) =

Original: | r a i s i − r a j s j | .

Original: The

Original: Bellman operator ˜ T ˆ πM under policy ˆ π is:

Original: [Equation]

Original: Lemma A.4.

Original: Consider an auxiliary MDP speciﬁed by a tuple (cid:104)

Original: ˆ S , ˆ A , ˆ R ∆ , ˆ P, γ (cid:105) whereˆ R ∆ (( s i , s j ) , ( a i , a j )) =

Original: | r a i s i − r a j s j | − | E a i ∼ π r a

Original: i s i − E a j ∼ π r

Original: a j s j | .

Original: Let V ˆ π ∆ denote the value function of such MDP over policy ˆ π , then

Original: [Equation]

Original: s i , s

Original: [Equation]

Original: [Equation]

Original: Proof.

Original: The value function V ˆ π ∆ can be expanded as,

Original: [Equation]

Original: Due to Lemma A.1, E

Original: [Table]

Original: [Table]

Original: ) i ∼ π

Original: [Equation]

Original: t )

Original: [Equation]

Original: ≥ 0 , therefore,

Original: [Equation]

Original: Theorem 4.3 (Looser value difference bound) .

Original: Given states s i and s j , and a policy π , we have| V π

Original: ( s i ) − V π ( s j ) | ≤ d πM ( s i , s j ) ≤ ˜ d πM ( s i , s j ) .

Original: Proof.

Original: Due to Theorem 3.4, which has been proven by Castro et al.

Original: [4] in Proposition 4.8 of theirpaper, we already have the leftmost two terms,

Original: [Equation]

Original: [Equation]

Original: Lemma A.2 shows that d πM is the value function of lifted MDP (cid:104)

Original: ˆ S , ˆ A

Original: [Equation]

Original: ˆ R , ˆ P, γ (cid:105) .

Original: d πM can be expandedas sum of discounted future rewards,

Original: [Equation]

Original: Similarly, due to Lemma A.3, ˜ d πM can be expanded as,

Original: [Equation]

Original: r a ( t ) i s ( t )

Original: [Equation]

Original: − r a ( t )

Original: j s ( t )

Original: [Equation]

Original: (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12) s (0)

Original: i =

Original: s i , s (0)

Original: j = s j   .

Original: Then we analyze the difference between ˜ d πM and d πM .

Original: Since ˜ d πM and d πM are deﬁned on the samepolicy ˆ π and dynamics model ˆ P , we have

Original: [Equation]

Original: t

Original: [Equation]

Original: (cid:12) (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12) s (0)

Original: i =

Original: s i , s (

Original: [Equation]

Original: j = s j 

Original:  = V ˆ π ∆ (( s i , s j )) .

Original: where V ˆ π ∆ (( s i , s j )) is the value function in Lemma A.4 and V ˆ π ∆ (( s i , s j )) ≥ 0 .

Original: Therefore,

Original: [Equation]

Original: Combine (16) and (17), ﬁnally prove that,

Original: [Equation]

Original: A.3 Proof of Theorem 5.

Original: Theorem 5.3 (Value function difference bound) .

Original: Given states s i and s j , and a policy π , we have

Original: [Equation]

Original: Proof.

Original: We follow Castro et al.

Original: [4] (the proof of Proposition 4.8) to prove the value function differencebound.

Original: We will show that if | V π ( s i ) − V π ( s j ) | ≤ d ( s i , s j ) , ∀ s i , s j ∈ S , then | V π ( s i ) − V π ( s j ) | ≤

Original: F πG ( d )( s i , s j ) .

Original: Suppose | V π ( s i ) − V π

Original: ( s j ) | ≤ d ( s i , s j ) holds, then

Original: V

Original: [Equation]

Original: s i

Original: [Equation]

Original: − V

Original: [Equation]

Original: ( s j

Original: [Equation]

Original: We make an assumption that (cid:80) s (cid:48) P as ( s (cid:48) )

Original: V π ( s (cid:48) )

Original: =

Original: V π ( E s (cid:48) ∼ P as [ s (cid:48) ]) .

Original: We make this assumptionbecause we use learned dynamics model to predict the next states and the learned dynamics model isassumed as Gaussian distribution with small standard deviation.

Original: If this assumption holds, then

Original: [Equation]

Original: ∼ π r a j s j (cid:12)(cid:12) +

Original: [Equation]

Original: a i ∼ π a j ∼

Original: [Equation]

Original: d (

Original: [Equation]

Original: s (cid:48)

Original: i ∼ P a i s i [ s (cid:48)

Original: i ] , E s (cid:48)

Original: j ∼ P

Original: a j s j [ s (cid:48) j ])

Original: =

Original: F πG ( d )( s i , s j ) .

Original: By symmetric,

Original: V

Original: [Equation]

Original: Therefore,

Original: [Equation]

Original: We have an initial distance d 0 ( s i , s j )

Original: = 2 max s , a | r as | / (1 − γ ) that satisﬁes | V π

Original: ( s i ) − V π ( s j ) | ≤d ( s i , s j ) , and F πG is contraction mapping on d .

Original: By repeatedly applying F πG on d , d will eventuallyconverge to the ﬁxed-point d πG .

Original: Because for each iteration F πG ( d ) satisﬁes | V π ( s j ) − V π ( s i ) | ≤

Original: F πG ( d )( s i , s j ) , the ﬁxed-point d πG satisﬁes

Original: [Equation]

Original: B Implementation Details

Original: B.1 Objectives

Original: B.1.1 Soft Actor-critic

Original: Our approach is built upon Soft Actor-Critic (SAC) [10].

Original: SAC is an actor-critic algorithm thatoptimizes a stochastic policy in off-policy manner.

Original: We use the version of SAC that incorporatesclipped double-Q trick.

Original: Let Q θ i denotes the Q-networks where i = 1 , 2 and Q ¯ θ i be the target

Original: Q-networks.

Original: Let π θ be the actor network.

Original: All Q-networks, target Q-networks and actor network takeinput of state representation φ ω ( s ) instead of state s .

Original: The loss functions for the Q-networks are

Original: [Equation]

Original: where D is the replay buffer, and V target is the target value function.

Original: We stop gradients for V targetand V target is deﬁned as

Original: [Equation]

Original: where a (cid:48) ∼ π θ ( ·| φ ω ( s (cid:48) )) is the action at next state.

Original: The policy improvement is minimizing thefollowing objective for actor network,

Original: [Equation]

Original: The loss function for α of SAC is,

Original: [Equation]

Original: where ¯ H ∈ R is the target entropy and ¯ H = −|A| in practice.

Original: B.1.2 RAP Distance

Original: We learn the reward function and dynamics model by neural network approximators, in order toapproximate RAP distance.

Original: Let r ψ be the learned reward function parameterized by ψ , whichoutputs a Gaussian distribution, r ψ ( s ) = { (cid:98) µ ( r s ) , (cid:98) σ ( r s ) } , where (cid:98) µ ( r s ) and (cid:98) σ ( r s ) are the mean andthe standard deviation, respectively.

Original: The loss function for learning r ψ is

Original: [Equation]

Original: We learn a dynamics model (cid:98) P η parameterized by η to take input of state embedding φ ( s ) andaction a , and output a Gaussian distribution over the next state embedding, (cid:98) P η ( φ ω ( s ) , a ) ={ (cid:98) µ ( (cid:98) P a φ ω ( s ) ) , (cid:98) σ ( (cid:98) P a φ ω ( s ) ) } , where (cid:98) µ ( (cid:98) P a φ ω ( s ) ) and (cid:98) σ ( (cid:98) P a φ ω ( s ) ) are the mean vector and the standard deviation vector for predictive next state embedding, respectively.

Original: The loss function for (cid:98) P η is

Original: [Equation]

Original: B.2 Learning Algorithm

Original: Algorithm shows the learning steps for learning SAC and RAP jointly.

Original: B.3 Networks and Hyperparameters

Original: We stack convolutional layers followed by a fully-connected layer as state encoder φ ω .

Original: Encoder takesinput of state, where 3 frames are stacked, and output a vector, the state representation φ ω ( s ) .

Original: Thedetail dimensions of each convolutional layer and fully-connected layer are described in Table 2.

Original: The Q-network consists of 3 stacked fully-connected layers with 1024 hidden dimensions, and takesinput of state representation φ ω ( s ) and action a .

Original: The actor network takes input of φ ω ( s ) and feed itinto three stacked fully-connected layers with 1024 hidden dimensions to output policy π .

Original: Both theapproximated dynamics model (cid:98) P and the approximated reward function r ψ are two-layers MLPswith 512 hidden dimensions.

Original: ReLU activation is used for every layers.

Original: We train the whole model onone NVIDIA A100 GPU.

Original: Other hyperparameters are listed in Table 2.

Original: [Table]

Original: Discount

Original: [Table]

Original: Encoder conv kernels [3,3,3,3] [3,3,3,3] [5, 3, 3, 3]

Original: Encoder conv channels [32,32,32,32] [32,32,32,32] [64,64,64,64]

Original: [Table]

Original: conv strides [2,1,1,1] [2,1,1,1] [2,2,2,2]

Original: [Table] C Additional Details of Experiments C.1 DeepMind Control Suite Table 3 shows the action repeat for various DMC tasks.
Revised: This work is supported by the 2020 Microsoft Research Asia collaborative research grant.

Original: [Table]

Original: C.2 CARLA

Original: [Figure]

Original: Figure 7 shows an illustration of observation in CARLA.

Original: The observation is 420 × 84 pixelsconstructed by concatenating ﬁve cameras.

Original: The action repeat for CARLA is 8.

Original: [Figure]

Original: D Additional Experiments

Original: [Figure]

Original: D.1 Ablation study on reward correction and learned dynamics

Original: [Figure]

Original: We perform this ablation study in order to evaluate the effects of the reward correction and the learned

Original: Gaussian dynamics model.

Original: Figure 8 shows the experiment results of cheetah-run and walker-walk of

Original: DMC with natural video background setting.

Original: The curve RAP w/o rew correction is our method RAPwith removing the reward variance correction.

Original: The curve RAP w/o learned dynamics is our methodbut measuring distances with next states that are sampled from MDP instead of a learned dynamicsmodel mean/expectation.

Original: This curve can be considered as MICo with reward variance correction.

Original: Figure 8 shows that without reward variance correction our method becomes more unstable and getsworse results in cheetah-run.

Original: Without learned dynamics model mean it will meet a larger performancedrop but it is still better than MICo.

Original: D.2 Experiment on Simple 2D Highway

Original: We add one more domain which is a 2D highway driving task to further evaluate our method and thebaseline methods.

Original: The 2D highway task contains a straight road with 4 lanes and 50 other vehiclesrunning on.

Original: Figure 9 shows an illustration of 2D highway and Figure 10 shows the observation whichis cropped to a square.

Original: The green vehicle is the vehicle that RL agent controls and blue vehicles arecontrolled by the environment.

Original: The RL agent learns to run as far as possible within 600 time steps.

Original: Figure 11 shows the learning curves at 500k steps.

Original: DBC, DrQ, and the proposed method RAP achievesimilar scores.

Original: SAC and MICo are slightly worse than the other three.

Original: This 2D highway is a simpletask so that they converge at around 100K to 200K steps.

Original: [Figure]

Original: E Visualization of Representation Space

Original: We show the t-SNE of the latent representation space of learned RAP, MICo and DBC.

Original: We collect datafrom cheetah-run with different natural video backgrounds.

Original: Figure 12 show the latent representationspace where the data points are collected from four different video backgrounds and marked withfour different colors.

Original: RAP and MICo can mix different video background which means they canprevent overﬁt to the background features, while DBC grouping the same color points into severalclusters potentially capture the background task-irrelevant features.

Original: RAP maps the observations to astructural latent space but MICo map to a circle latent space which may lose the structure.

Original: [Figure]

Original: Figure 13 shows the same representation spaces but each data point is colored according to thepredicted Q-value.

Original: The darker blue indicates higher Q-value and lighter blue means lower Q-value.

Original: The representation space of RAP preserves the information for predicting Q-values as it gets darkblue in the bottom-left and light blue in the upper-right.

Original: MICo seems to have random Q-values withrespect to the representation structure.

Original: DBC is likely to predict more extreme high/low Q-values.

Original: Figure 14 shows some representations of the corresponding observations in the latent space.

Original: RAPcan map similar observations to the nearby locations, while MICo and DBC map them far away fromeach other.

Original: [Figure]