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 acceleratedeep 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: Two key research issues on behavioralmetric-based representation learning are how to relax the computation of a specificbehavioral metric, which is difficult or even intractable to compute, and how toapproximate 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 thispaper, we analyze the potential relaxation and/or approximation gaps for existingbehavioral 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.
Revised: Based on the analysis,we propose a new behavioral distance, the RAP distance and develop a practicalrepresentation learning algorithm on top of it.

Original: We provide theoretical analysis onthe proposed algorithm.
Revised: We provide theoretical analysis onthe proposed algorithm.

Original: We conduct extensive experiments on DeepMind Control
Revised: We conduct extensive experiments on DeepMind Control

Original: Suite with distraction, Robosuite, and autonomous driving simulator CARLA todemonstrate new state-of-the-art results.
Revised: Suite with distraction, Robosuite, and autonomous driving simulator CARLA todemonstrate 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-dimensionalenvironmental 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-irrelevantbackground information, is able to significantly speed up the RL process and makes the learned policymore 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 algorithmsand attracts more and more attraction 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 works on representation learning are focused on learning embeddings to represent statesbased on a reconstruction loss [26, 25, 13].

Original: Though promising results have been reported on some
Revised: Though promising results have been reported on some

Original: 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: 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.

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 dataaugmentation 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 of a Markov Decision Process (MDP).
Revised: While dataaugmentation-based approaches may be able to learn more robust feature representations, they donot take the characteristics of a Markov Decision Process (MDP).

Original: e.g., reward signals and dynamicsmodels, into consideration when learning representations.
Revised: e.g., reward signals and dynamicsmodels, 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 maynot 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 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].

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 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 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 beenexploited 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.
Revised: Behavioral metrics are originally proposed tomeasure behavioral similarity between states in terms of rewards and dynamics models, e.g, state transition probabilities

Original: The high-level
Revised: The high-

Original:  idea is to learn stateembeddings by
Revised: level idea is to learn state

Revised: embeddings

Original:  preserving the behavioral
Revised: by preserving the behavioral

Original: similarity between states based ona
Revised: similarity between states based on

Original: speciﬁc
Revised: a specific

Original: behavioral metric [29, 1, 16,
Revised: behavioral metric [29, 1, 16,

Original: 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 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.

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 tocompute, different approximation approaches and learning objectives have been proposed to makebehavioral metric-based representation learning for RL agents more efficiently [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 and limitation of the L 1 / L 2 norm.
Revised: 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 and limitation of the L 1 / L 2 norm.

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

Original: In behavioral metric-based representation learning, the loss
Revised: In behavioral metric-based representation learning, the loss

Original: function is the difference between the distance between states on the embedding space
Revised: function is the difference between the distance between states on the embedding space

Original: and their behavioral metric.
Revised: and their behavioral metric.

Original: As behavioral metrics are expensive or even intractable to
Revised: As behavioral metrics are expensive or even intractable to

Original: computers, approximation or relaxation are necessary.
Revised: 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
Revised: However, based on our analysis,state-of-the-art methods adopt improper approximations, which may introduce bias and

Original: makes the bound on value function looser. • Relaxation of dynamics model divergence .
Revised: makes the bound on value function looser.• Relaxation of dynamics model divergence .

Original: The bisimulation metric or its on-policy
Revised: The bisimulation metric or its on-policy

Original: invariant is one of the most widely used behavioral metrics, which requires estimating
Revised: invariant is one of the most widely used behavioral metrics, which requires estimating

Original: the 1-Wasserstein distance between dynamics models.
Revised: the 1-Wasserstein distance between dynamics models.

Original: As the 1-Wasserstein distance is
Revised: As the 1-Wasserstein distance is

Original: usually difﬁcult or intractable to estimate, prior methods propose some relaxations to replace
Revised: usually difficult or intractable to estimate, prior methods propose some relaxations to replace

Original: the estimation of the 1-Wasserstein, which may break some theoretical guarantees of the
Revised: the estimation of the 1-Wasserstein, which may break some theoretical guarantees of the

Original: bisimulation metric.
Revised: bisimulation metric.

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

Original: The L 1 and the L 2 norms are commonly-used distances
Revised: The L 1 and the L 2 norms are commonly-used distances

Original: with zero self-distance.
Revised: with zero self-distance.

Original: However, due to the two gaps mentioned above, the approximations
Revised: However, due to the two gaps mentioned above, the approximations

Original: or relaxations of behavior metrics in prior approaches are potentially non-zero self-distance.
Revised: 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 topreserve 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, and we then develop a practical approximationalgorithm with consistency to its theoretical prototype.
Revised: To the aforementioned issues or gaps, in this paper, we first introduce a new behavior metric namelythe Reducing Approximation Gap (RAP) distance, and we then develop a practical approximationalgorithm 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 learna robust state representation to capture 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 expensiveexperiments 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.

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 forexisting behavioral metric-based representation learning approaches, and 2) we introduce a newbehavior 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 works have 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 reconstruction loss.

Original: Geladaet al.
Revised: Geladaet 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 dynamics model and reward function onembedding 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 momentumencoder.

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 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 representationlearning.

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 extensive study of data augmentation for deep RLwith 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 regularizesthe Q-function objectives.

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

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

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, Incontrast to these works, our proposed method aim to encode state representation into a structuralmetric 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 inthe 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 representation by approximating the bisimulation metric [7, 3] on 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 metricconsidering 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: 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.

Original: Castroet al.
Revised: Castroet 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 preservebehavioral 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 learnneural networks to approximate components in the bisimulation metric on embedding state 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 thetransition distribution that captures the probability of entering next state s ′ ∈ S given current states ∈ S and 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 aprobability distribution over each action a conditioned on 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 → Rfor a given policy π at 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 maximizesthe 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 tomeasure 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 behaviorsrelative 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: Apseudometric 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 ∼ π r a

Original: i s i
Revised: i s i

Original: =
Revised: =

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

Original: r a i s i , P
Revised: r a i s i , P

Original: π s i
Revised: π s i

Original: =
Revised: =

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

Original: F πB has 116 a least ﬁxed point d πB and d πB is a π -bisimulation metric.
Revised: F πB has 116 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 i 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 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.

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 i and policy π , we have

Revised: [Equation]

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 spacesuch 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 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.

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 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 :

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 differentapproximation or relaxation strategies to make the resultant optimization problem computationallytractable.

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, and limitation 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 prototype, i.e., (5), namelyloss function mismatch, relaxation of dynamics model divergence, and limitation of the L 1 or the L 2norm.

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 expensivebecause 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 representationlearning [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 ¯ ω a copy of parameter 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
Revised: i s i and

Original: 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: E a j ∼ π r a j s j in the 2nd term in (6) are computationally intractable and also difficult to estimate evenbased 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 followingalternative 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 ofthe 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 betweensolutions of minimizing (7) and (6), respectively, because in (7) the reference behavioral metric is nolonger 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’sfixed-point theorem.

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

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 | ≥| Ea i ∼ π r a
Revised: However, as E a i ∼ π a j ∼ π | r a i s i − r a j s j | ≥| E

Revised: [Equation]

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 ∼ π r a i s i − E a j ∼ π r a j s j | , we have the following theorem, whose proof can be found in Appendix.

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 maybe less relevant to 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 encodebehavior 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-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.

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 1-Wasserstein distance W 1 between dynamics modelsto 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 canlearn 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] the2-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 withbatched 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 expectedpolicy 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: Theexistence of unique fixed-point in the bisimulation metric requires the continuity and monotonicityof 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 thedynamics 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-baseddistribution 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 dynamicsmodels.

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-Karmowskimetric.

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 notsatisfy 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 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.

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 overrewards, prior approaches, such as MICo and DBC, use an alternative term to measure the rewarddifference 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 ands 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 ris 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 functionis 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 betweendynamics 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 identicalstates.

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 L 1 and 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 embeddingspace 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, is usually greater thanzero on identical states.
Revised: However, the regression target, i.e., the behavioral metric, is usually greater thanzero 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 andthe behavioral metric, the representations for similar states, especially identical ones, will be pushedapart 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 spacerather 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 nothave 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 newform of distance on the embedding space, which is a sum of angular distance between embeddingsand L 2 norm of the embeddings.

Original: Note that the proposed distance has non-zero self-distance.
Revised: Note that the proposed 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 metricnamely the RAP distance enjoys theoretical properties such as fixed-point existence and a valuedifference 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 approximatingthe 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 learnedestimation of reward functions and dynamics models to provide distance approximation which isconsistent 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 andpreserves 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, 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.

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 relaxationto 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 RAPdistance 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 thepolicy π .

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 dynamicsmodel 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 convergence guarantee for the RAP distance that by iterating F πG , distance dwill 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.

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 statesupper-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 theembedding 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 functionmismatch 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 behavioralmetric 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 definedby 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 a
Revised: We first analyze the difference between E a i ∼ π a j ∼ π (cid:104)(cid:12)(cid:12) r a

Original: i s i − r a j s j (cid:12)(cid:12) 2 (cid:105)
Revised: i s i − r a j s j (cid:12)(cid:12) 2 (cid:105)

Original: and| E a i
Revised: and| E a i

Original: ∼ π ra
Revised: ∼ π r

Original: i s i − E a j ∼ π r a j s j | 2 , E
Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: =
Revised: =

Original: [Equation]
Revised: [Equation]

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

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

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

Original: r
Revised: r

Original: a
Revised: a

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

Original: a j s j | = 257 (cid:114)
Revised: a j s j | = 257 (cid:114)

Original: E a i ∼ π a j ∼ π (cid:104)(cid:12)(cid:12) r a
Revised: E a i ∼ π a j ∼ π (cid:104)(cid:12)(cid:12) r a

Original: i s i − r a j s j (cid:12)(cid:12) 2 (cid:105) − V ar [ r s i ] − V ar
Revised: i s i − r a j s j (cid:12)(cid:12) 2 (cid:105) − V ar [ r s i ] − V ar

Original: [ r s j ] and the revised RAP distance at ﬁxed-point as
Revised: [ r s j ] 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 operatorwithout 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, 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 ′

Original: i ∼ P
Revised: i ∼ P

Original: [Equation]

Original: [ s
Revised: a

Original: [Equation] i ] and E s [Equation]
Revised: i s i [ s ′ i ] and E s ′

Original:  s [Equation]
Revised: i ∼ P a j s j [ s ′ j ] .

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: 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.

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: 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.

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.
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 andthe standard deviation vector for predictive next state embedding, respectively.

Original: Let ˆ d ( φ ω ( s i ) , φ ω ( s j )) be the approximated RAP distance between s i and s j parameterized byω .
Revised: Let ˆ d ( ϕ ω ( s i ) , ϕ ω ( s j )) be the approximated RAP distance between s i and s j parameterized byω .

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 functionoutput 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: In order to reduce the biased issue introduced by the square root, we propose anovel regression loss defined as, [Equation] ω [Equation]

Revised: − γ

Original: [Equation]
Revised: [Equation]

Original:  where ¯ ω is a copy of parameter with stop gradient and φ ¯ ω is the encoder in the target network.
Revised: ω [Equation] , (13) where ¯ ω is a copy of parameter with stop gradient and ϕ ¯ ω is the encoder in the target network.

Original: 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: 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.

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, thedistance 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 distanceon 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.

Original: [Figure]
Revised: [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: Ourmethod is built upon SAC [10].

Original: Actor and critic networks take inputof state representation φ ω ( s ) .
Revised: Actor and critic networks take inputof 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 implementationdetails can be found in Appendix.

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)
Revised: In this section, we evaluate the efficiency, robustness and generalization ability of the representationlearned by our method on three RL benchmarks: 1) Distracting DeepMind Control Suite (DMC) [24],2)

Original: Robosuite [30], and 3) CARLA [6].
Revised: Robosuite [30], and 3) CARLA [6].

Original: These three domains are continuous action space controltask with visual input.
Revised: These three domains are continuous action space controltask 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 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.

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 ourmethod, we perform experiments on 2 settings, original background and natural video backgroundsettings, 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 framesas 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] toreplace 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 agentsand evaluate agents on another 1000 continuous frames.

Original: The video background is considered asdistraction to RL algorithm.
Revised: The video background is considered asdistraction 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 for 1 million environment steps.

Original: Figure 3 is the training curves in original background settings.
Revised: Figure 3 is 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 canaccelerate the training and has comparable results to data augmentation method DrQ. Figure 4 isthe training curves in natural video background settings.

Original: Our method RAP outperforms DBC and
Revised: Our method RAP outperforms DBC and

Original: MICo on all 4 tasks.
Revised: 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 fasterin 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 ofbehavioral distance can improve the efficiency and robustness for deep RL and they also confirm thatour 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 simulatorwith various manipulation tasks based on MuJoCoengine.

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 byrandomizing the colors of robot arms and table,the lighting source and luminance, and the cameraposition.

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 positionare 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 isperformed 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 ofobservations 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 × comparedto 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 therobustness 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 autonomousdriving simulator CARLA [6], which provides 3D realistic on-worldscenarios.

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 aspossible 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 rewardfunction followed [29] is designed to encourage driving far and avoidcollisions 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 × 84pixels of a 300 degrees ego-centric view, constructed by concatenatingfive cameras.

Original: In order to evaluate the generalization ability, we randomlysample akind of weather in each episode starts.
Revised: In order to evaluate the generalization ability, we randomlysample a weather at 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 achievescomparable 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.
Revised: In thispaper, we propose a new behavioral metric and a practical representation learning algorithm on top ofthe new behavioral metric for deep RL.

Original: We provide theoretical analysis for our proposed metric aswell as the representation learning algorithm.
Revised: We provide theoretical analysis for our proposed metric aswell as the representation learning algorithm.

Original: We conduct empirical studies on multiple RL domainsto verify the effectiveness of our proposed method.
Revised: We conduct empirical studies on multiple RL domainsto verify the effectiveness of our proposed method.

Original: Checklist
Revised: Checklist

Original: The checklist follows the references.
Revised: The checklist follows the references.

Original: Please read the checklist guidelines carefully for information onhow to answer these questions.
Revised: 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] .
Revised: 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.
Revised: You are strongly encouraged to include a justification to your answer , either by referencingthe appropriate section of your paper or providing a brief inline description.

Original: For example:
Revised: For example:

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

Original: [Yes] See Section ?? .•
Revised: [Yes] See Section ?? .•

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

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

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

Original: [N/A]
Revised: [N/A]

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

Original: Note that the
Revised: Note that the

Original: Checklist section does not count towards the page limit.
Revised: 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.
Revised: In your paper, please delete this instructionsblock and only keep the Checklist section heading above along with the questions/answers below.

Original: For all authors...
Revised: For all authors...

Original: (a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s
Revised: (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s

Original: contributions and scope?
Revised: contributions and scope?

Original: [Yes]
Revised: [Yes]

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

Original: [No]
Revised: [No]

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

Original: [N/A]
Revised: [N/A]

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

Original: them?
Revised: them?

Original: [Yes]
Revised: [Yes]

Original: If you are including theoretical results...
Revised: If you are including theoretical results...

Original: (a) Did you state the full set of assumptions of all theoretical results?
Revised: (a) Did you state the full set of assumptions of all theoretical results?

Original: [Yes]
Revised: [Yes]

Original: (b) Did you include complete proofs of all theoretical results?
Revised: (b) Did you include complete proofs of all theoretical results?

Original: [Yes] In Appendix.
Revised: [Yes] In Appendix.

Original: If you ran experiments...
Revised: If you ran experiments...

Original: (a) Did you include the code, data, and instructions needed to reproduce the main experi-
Revised: (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)?
Revised: mental results (either in the supplemental material or as a URL)?

Original: [No] We will publishcode once paper is accepted.
Revised: [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)?
Revised: (b) Did you specify all the training details (e g., data splits, hyperparameters, how theywere chosen)?

Original: [Yes] In Appendix.
Revised: [Yes] In Appendix.

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

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

Original: [Yes]
Revised: [Yes]

Original: (d) Did you include the total amount of compute and the type of resources used (e
Revised: (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)?
Revised: g., typeof GPUs, internal cluster, or cloud provider)?

Original: [Yes] In Appendix.4.
Revised: [Yes] In Appendix.

Original: If you are using existing assets (e
Revised: If you are using existing assets (e

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

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

Original: [Yes]
Revised: [Yes]

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

Original: [N/A]
Revised: [N/A]

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

Original: [No]
Revised: [No]

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

Original: using/curating?
Revised: using/curating?

Original: [N/A]
Revised: [N/A]

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

Original: information or offensive content?
Revised: information or offensive content?

Original: [N/A]
Revised: [N/A]

Original: If you used crowdsourcing or conducted research with human subjects...
Revised: 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?
Revised: (a) Did you include the full text of instructions given to participants and screenshots, if applicable?

Original: [No]
Revised: [No]

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

Original: Board (IRB) approvals, if applicable?
Revised: Board (IRB) approvals, if applicable?

Original: [N/A]
Revised: [N/A]

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

Original: spent on participant compensation?
Revised: spent on participant compensation?

Original: [N/A]
Revised: [N/A]

Original: A ProofsA.1 Proof of Lemma 4.

Original: Lemma 4.2 (Fixed-point of shift MICo) .

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

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. A.2 Proof of Theorem 4.  Lemma A. [Equation]
Revised: A Appendix

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: =

Original: A × A , transitiondistribution ˆ P ( a i , a j ) ( s i , s j ) =

Original: P a

Original: i s i P a j s j , and reward function

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. [Table] Discount [Table]
Revised: Appendix is uploaded in supplemental material

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]

Original: C Additional Details of Experiments

Original: C.1 DeepMind Control Suite

Original: Table 3 shows the action repeat for various DMC tasks.

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]

Original: F Explanation of Policy Dependence of RAP Loss

Original: The loss (13) in Section 5.2 should be dependent on the policy.

Original: More precisely, the expectation istaken over the transitions sampled from a replay buffer D .

Original: The loss (13) should be written as

Original: [Equation]

Original: The rewards r a

Original: i s i and r a j s j are sampled from the replay buffer, and the actions a i and a i that the rewardsdepend on are sampled from the policy when the rewards are generated from the MDP.

Original: Therefore thisloss function is dependent on the policy.

Original: The learning of the neural network approximated rewardvariance (cid:98) σ is dependent on the data sampled from replay buffer, as described in Appendix B.1.2.

Original: Thisindicates that the learned reward function is also dependent on the policy.

Original: If we use on-policy RL algorithms, then D would be the set of transitions sampled from the currentpolicy.

Original: However, off-policy RL algorithms have better sample efﬁciency, especially in control tasks.

Original: In this paper, we build our method upon SAC which is also an off-policy algorithm.

Original: The underlinepolicy in the replay buffer may not be identical to the current policy.

Original: We want to verify that eventhough the learning of RAP is not dependent on the "truly" on-policy, our approximation is not faraway from on-policy and enough to learn good representations.

Original: We try to reduce the replay buffer capacity.

Original: If replay buffer is small enough, the RL algorithm wouldbecome on-policy.

Original: However, smaller replay buffer usually leads to worse performance.

Original: We reducethe replay buffer capacity from 1 million to 500K, 100K, 50K, and 10K. Figure 15 shows the learningcurves of different replay buffer capacities.

Original: It shows that smaller replay buffer sizes lead to lowerperformance and slower to converge.

Original: 1M 500K

Original: [Table]

Original: 1M \

Original: [Table]

Original: [Table]

Original: [Table]

Original: [Table]

Original: [Table]

Original: [Table]

Original: [Table]

Original: [Table]

Original: [Table]

Original: [Table]

Original: [Table]

Original: [Table]

Original: 10K

Original: [Table]

Original: [Table]

Original: [Table]

Original: [Table]

Original: We generate trajectories (in total 1K steps) by the optimal policy which is our RAP policy.

Original: Thenwe measure the RAP distance in between the 1K observations by using different models learnedby different replay buffer size, i.e. 500K, 100K, 50K and 10K.

Original: We compare the distances amongdifferent replay buffer sizes and show the average differences in Table 4.

Original: The differences of thedistances stay in the same scale.

Original: The 1M version which is the proposed method is not far awayfrom the models learned with other buffer sizes.

Original: Such difference/variance may be generated by therandomly initialized weights of the neural networks.