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 accelerate the deep reinforcement learning process.
Revised: Learning an informative representation with behavioral metrics is able to accelerate the deep reinforcement learning process.

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

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

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

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

Original:  Introduction
Revised:  Introduction

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

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

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

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

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

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

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

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

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

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

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

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

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

Original: As behavioral metrics provide a theoretical bound on the difference between the outputs of value function of a pair of states, the learned representation enjoys a theoretical guarantee to capture behavioral structure in the environment for policy learning.
Revised: As behavioral metrics provide a theoretical bound on the difference between the outputs of value function of a pair of states, the learned representation enjoys a theoretical guarantee to capture behavioral structure in the environment for policy learning.

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

Original:  Related Work
Revised:  Related Work

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

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

Original: [11] propose to learn a dynamics model to predict future states with a reconstruction loss.
Revised: [11] propose to learn a dynamics model to predict future states with a reconstruction loss.

Original: Gelada et al.
Revised: Gelada et al.

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

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

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

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

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

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

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

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

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

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

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

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

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

Original: [23] use augmented samples for representation contrastive learning,
Revised: [23] use augmented samples for representation contrastive learning.

Original: In contrast to these works, our proposed method aim to encode state representations into a structural metric space based on behavioral metrics.
Revised: In contrast to these works, our proposed method aims to encode state representations into a structural metric space based on behavioral metrics.

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

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

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

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

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

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

Original: Castro et al.
Revised: Castro et al.

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

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

Original:  Preliminaries
Revised:  Preliminaries

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

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

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

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

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

Original: [Equation]
Revised: [Equation]

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

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

Original: Bisimulation Metrics
Revised: Bisimulation Metrics

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

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

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

Original:  π -bisimulation metric consists of a reward difference term and a
Revised: The π -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: A pseudometric transformation function F πB :
Revised: A pseudometric transformation function F πB :

Original: M → M is defined as,
Revised: M → M is defined as,

Original: [Equation]
Revised: [Equation]

Original: where E a i
Revised: where E a i

Original: [Equation] 
Revised: [Equation] =

Original: E
Revised: E

Original: [Equation] 
Revised: [Equation] =

Original: E
Revised: E

Original: [Equation]
Revised: [Equation]

Original: P
Revised: P

Original: [Equation]
Revised: [Equation]

Original: and W 1 is the 1-Wasserstein distance.
Revised: and W 1 is the 1-Wasserstein distance.

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

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

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

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

Original: [Equation]
Revised: [Equation]

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

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

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

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

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

Original: [Equation]
Revised: [Equation]

Original: then F πM has a unique fixed 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 j , and policy π , we have
Revised: Given states s i and s j , and policy π , we have

Original: [Equation]
Revised: [Equation]

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

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

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

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

Original: [Equation]
Revised: [Equation]

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

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

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

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

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

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

Original: By specifying the term d π ( s i , s j ) in (5) as the MICo distance defined 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 ¯ ω is a copy of parameters for the target network.
Revised: where ¯ ω is a copy of parameters for the target network.

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

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

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

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

Original: [Equation]
Revised: [Equation]

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

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

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

Original: However, such a revision introduces a gap between solutions of minimizing (7) and (6), because in (7) the reference behavioral metric is no longer the
Revised: However, such a revision introduces a gap between solutions of minimizing (7) and (6), because in (7) the reference behavioral metric is no longer the

Original: MICo distance but the “shift” MICo distance defined as follows.
Revised: MICo distance but the “shift” MICo distance defined as follows.

Original: Definition 4.1 (Shift MICo distance) .
Revised: Definition 4.1 (Shift MICo distance) .

Original: The shift MICo distance function ˜ F πM is defined 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 fixed-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 fixed-point theorem.
Revised: This can be proved by following the proof of Theorem 3.3 by using Banach’s fixed-point theorem.

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

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

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

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

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

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

Original: [Equation]
Revised: [Equation]

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

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

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

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

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

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

Original: [Equation]
Revised: [Equation]

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

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

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

Original: Specifically, the loss function of DBC with batched sampled transitions is defined as,
Revised: Specifically, the loss function of DBC with batched sampled transitions is defined as,

Original: [Equation]
Revised: [Equation]

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

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

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

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

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

Original: Banach’s fixed-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-based distribution divergence, E s ′
Revised: Instead of using the family of Wasserstein distances, in MICo as shown in (3) the sample-based distribution divergence, E s ′

Original: i ∼ P
Revised: i ∼ P

Original: π s i , s ′
Revised: π s i , s ′

Original: j ∼ P
Revised: j ∼ P

Original: π s j d
Revised: π s j d

Original: s ′ i , s ′ j ) , is used to measure the difference between dynamics models.
Revised: s ′ i , s ′ j ) , is used to measure the difference between dynamics models.

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

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

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

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

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

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

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

Original: Lemma 4.4.
Revised: Lemma 4.4.

Original: If s i
Revised: If s i

Original: =
Revised: =

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

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

Original: action a
Revised: action a

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

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

Original: actions.
Revised: actions.

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

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

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

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

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

Original: i − x j ∥ 1 = ∥ x i − x j ∥ 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 embedding space to approximate the behavioral metric, then the identical states pairs has zero distance on embedding space.
Revised: If we leverage the L 1 or the L 2 norm as the form of distance ˆ d on the embedding space to approximate the behavioral metric, then the identical states pairs has zero distance on embedding space.

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

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

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

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

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

Original: Note that the proposedwo 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 firstly 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 metric namely the RAP distance enjoys theoretical properties such as fixed-point existence and a value difference bound.
Revised: The proposed behavioral metric namely the RAP distance enjoys theoretical properties such as fixed-point existence and a value difference bound.

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

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

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

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

Original: 5.1 Definition 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 relaxation to W 2 as described in Section 4.2, we consider a distance measure between dynamics models
Revised: In order to avoid the high computational cost of W 1 or the approximation gap introduced by relaxation to W 2 as described in Section 4.2, we consider a distance measure between dynamics models

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

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

Original: To be specific, our on-policy behavioral distance is defined as follows.
Revised: To be specific, our on-policy behavioral distance is defined as follows.

Original: Definition 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 RAP distance function F πG : M → M is defined as,
Revised: → R , the RAP distance function F πG : M → M is defined as,

Original: [Equation]
Revised: [Equation]

Original: where E [ s ′ i ]
Revised: where E [ s ′ i ]

Original: =
Revised: =

Original: E s ′
Revised: E s ′

Original: i ∼ P
Revised: i ∼ P

Original: a
Revised: a

Original: i s i
Revised: i s i

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

Original: a
Revised: 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 the policy π .
Revised: a j s j but only performs sampling on the policy π .

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

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 fixed-point d πG .
Revised: the L ∞ norm and has a unique fixed-point d πG .

Original: Proof.
Revised: Proof.

Original: Let d, d ′ ∈ 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 fixed-point for
Revised: the L ∞ norm and there exists a unique fixed-point for

Original: F πG due to Banach’s fixed-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 a convergence guarantee for the RAP distance that by iterating F πG , distance d will converge to the fixed-point d πG .
Revised: Theorem 5.2 provides a convergence guarantee for the RAP distance that by iterating F πG , distance d will converge to the fixed-point d πG .

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

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

Original: [Equation]
Revised: [Equation]

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

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

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

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

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

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

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

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

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

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

Original: =
Revised: =

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

Original: We first 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 i s i − r a j s j (cid:12)(cid:12) 2 (cid:105) and

Original: [Equation]
Revised: [Equation]

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

Original: [Equation]
Revised: [Equation]

Original: = V ar
Revised: = V ar

Original: [ r s i ]
Revised: [ r s i ]

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

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

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

Original: [Equation]
Revised: [Equation]

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

Original: Therefore
Revised: Therefore

Original: [Equation]
Revised: [Equation]

Original: we have the reward difference term
Revised: we have the reward difference term

Original: [Equation]
Revised: [Equation]

Original: and the revised RAP distance at fixed-point as
Revised: and the revised RAP distance at fixed-point as

Original: [Equation]
Revised: [Equation]

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

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

Original: i ∼ P
Revised: i ∼ P

Original: a
Revised: a

Original: i s i
Revised: i s i

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

Original: i ∼ P
Revised: j ∼ P

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

Original: In order to reduce the bias issue introduced by the square root, we try to remove the square root in the loss of learning the RAP distance.
Revised: In order to reduce the bias issue introduced by the square root, we try to remove the square root in the loss of learning the RAP distance.

Original: We move the dynamics term in (12) to the left-hand side, then take square on both sides and get
Revised: We move the dynamics term in (12) to the left-hand side, then take square on both sides and get

Original: [Equation]
Revised: [Equation]

Original: 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 loss for learning ˆ d ( ϕ ω ( s i ) , ϕ ω ( s j )) is to minimize the mean squared error between the left-hand side and the right-hand side in (13):
Revised: The loss for learning ˆ d ( ϕ ω ( s i ) , ϕ ω ( s j )) is to minimize the mean squared error between the left-hand side and the right-hand side in (13):

Original: [Equation]
Revised: [Equation]

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

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

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

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

Original: The RAP distance between expected next states, d πG ( E [ s ′ i ] , E [ s ′ j ]) , can be approximated by ˆ d ( (cid:98) µ ( (cid:98) P a i ϕ ω
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: ( 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 function output as ( (cid:98) σ ( r s )) 2 .
Revised: The variance V ar [ r s ] can be replaced by the learned reward function output as ( (cid:98) σ ( r s )) 2 .

Original: By replacing the dynamics term and reward variance terms in (14), we propose the RAP loss defined as
Revised: By replacing the dynamics term and reward variance terms in (14), we propose the RAP loss defined as

Original: [Equation]
Revised: [Equation]

Original: where ¯ ω is a copy of parameter with stop gradient and ϕ ¯ ω is the encoder in the target network and D is the replay buffer or the set of data that RL algorithm, e.g. SAC [10], learns from.
Revised: where ¯ ω is a copy of parameter with stop gradient and ϕ ¯ ω is the encoder in the target network and D is the replay buffer or the set of data that RL algorithm, e.g. SAC [10], learns from.

Original: The loss (15) is trained over the transitions sampled from D .
Revised: The loss (15) is trained over the transitions sampled from D .

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

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

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

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

Original: Definition 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 defined as ˆ d G ( x , y ) =
Revised: ˆ d G is defined as ˆ d G ( x , y ) =

Original: ∥ x ∥ 22 + ∥
Revised: ∥ x ∥ 22 + ∥

Original: y
Revised: y

Original: ∥
Revised: ∥

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

Revised: ˆ d G ( x , x ) =

Original: ∥ x ∥ 22 ≥ 0 .
Revised: ∥ 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: 5.4 Implementation
Revised: 5.4 Implementation

Original: [Figure]
Revised: [Figure]

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

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

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

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

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

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

Original:  Experiment
Revised:  Experiment

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

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

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

Original: [Figure]
Revised: [Figure]

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

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

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

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

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

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

Original: For each setting, we train and evaluate on 4 tasks: cartpole-swingup, cheetah-run, finger-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 using 1 million environment steps.
Revised: Each task is trained using 1 million environment steps.

Original: curve is average on 3 runs.
Revised: curve is average on 3 runs.

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

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

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

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

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

Original: Robosuite Robosuite [30] is a robotic simulator with various manipulation tasks based on
Revised: Robosuite Robosuite [30] is a robotic simulator with various manipulation tasks based on

Original: MuJoCo engine.
Revised: MuJoCo engine.

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

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

Original: Our experiment is performed by controlling the
Revised: Our experiment is performed by controlling the

Original: Panda arms.
Revised: 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:
Revised: We evaluate on two tasks:

Original: Door Opening and Two Arm Peg-In-Hole.
Revised: Door Opening and Two Arm Peg-In-Hole.

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

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

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

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

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

Original: [Table]
Revised: [Table]

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

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

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

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

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

Original: The weather (sunlight, rain, etc.), which affects the visual observation, can be considered as real-world distraction to RL agents.
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 method achieves comparable scores with data augmentation method DrQ, and its score is higher than behavioral metric DBC and MICo.
Revised: Our method achieves comparable scores with data augmentation method DrQ, and its score is higher than behavioral metric DBC and MICo.

Original: It shows that our RAP is able to generalize well in tasks with real-world scenarios.
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 this paper, we propose a new behavioral metric and a practical representation learning algorithm on top of the new behavioral metric for deep RL.
Revised: In this paper, we propose a new behavioral metric and a practical representation learning algorithm on top of the new behavioral metric for deep RL.

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

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

Original: Acknowledgement and Funding Disclosure
Revised: Acknowledgement and Funding Disclosure

Revised: This work is supported by the 2020 Microsoft Research Asia collaborative research grant.

Revised: A Proofs A.1 Proof of Lemma 4.

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

Revised: The “shift” MICo distance function ˜ F πM has a unique fixed-point ˜ d πM .

Revised: Proof.

Revised: The shift MICo distance function ˜ F πM is defined as,

Revised: [Equation]

Revised: Let d, d ′ ∈ M .

Revised: We have

Revised: [Equation]

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

Revised: the L ∞ norm and there exists a unique fixed-point ˜ d πM for ˜ F πM due to Banach’s fixed-point theorem.

Revised: This completes the proof.

Revised: A.2 Proof of Theorem 4.

Revised: Lemma

Revised: [Equation]

Revised: Proof.

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

Revised: a j s j − r a i s i by symmetric.

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

Revised: [Equation]

Revised: 16) and

Revised: [Equation]

Revised: By combining (16) and (17), we have

Original: This work issupported by the 2020 Microsoft Research Asia collaborative research grant.
Revised: [Equation] Lemma A.2 (Lifted MDP of MICo [4]) . The MICo metric function F πM is the Bellman operator for a Lifted MDP. Proof.

Revised: Define the MDP for RL by a tuple ⟨S , A , R , P, γ ⟩ .

Revised: We consider a lifted MDP constructed by a tuple ⟨ ˆ S , ˆ A , ˆ R , ˆ P, γ ⟩ , where state space ˆ S = S × S , action space ˆ A =

Revised: A × A , transition distribution ˆ P ( a i , a j ) ( s i , s j ) =

Revised: P a

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

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

Revised: | E a i ∼ π r a

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

Revised: a j s j | .

Revised: The

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

Revised: π ( a i | s i ) π

Revised: [Equation]

Revised: a j

Revised: [Equation]

Revised: s j

Revised: [Equation]

Revised: is:

Revised: [Equation]

Revised: ˆ

Revised: [Equation]

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

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

Revised: Proof.

Revised: We construct the lifted MDP of shift MICo by a tuple ⟨ ˆ S , ˆ A , ˜ˆ R , ˆ P, γ ⟩ which is the same as the lifted MDP in Lemma A.2 except for the reward function ˜ˆ R (( s i , s j ) , ( a i , a j )) =

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

Revised: The

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

Revised: [Equation]

Revised: Lemma A.4.

Revised: Consider an auxiliary MDP specified by a tuple ⟨ ˆ S , ˆ A , ˆ R ∆ , ˆ P, γ ⟩ where ˆ R ∆ (( s i , s j ) , ( a i , a j ))

Revised: =

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

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

Revised: a j s j | .

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

Revised: [Equation]

Revised: Proof.

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

Revised: [Table]

Revised: , therefore,

Revised: [Equation]

Revised: Theorem 4.3 (Looser value difference bound) .

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

Revised: [Equation]

Revised: Proof.

Revised: Due to Theorem 3.

Revised: [Equation]

Revised: which has been proven by Castro et al.

Revised: [4] in Proposition 4.8 of their paper

Revised: [Equation]

Revised: we already have the first inequality,

Revised: [Equation]

Revised: Lemma A.2 shows that d πM is the value function of lifted MDP ⟨ ˆ S , ˆ A , ˆ R , ˆ P, γ ⟩ .

Revised: d πM can be expanded as the sum of discounted future rewards,

Revised: [Equation]

Revised: Similarly

Revised: [Equation]

Revised: due to Lemma A.

Revised: [Equation]

Revised: ˜ d πM can be expanded as,

Revised: [Equation]

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

Revised: Since ˜ d πM and d πM are defined on the same policy ˆ π and dynamics model ˆ P , we have

Revised: [Equation]

Revised: V ˆ π ∆

Revised: [Equation]

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

Revised: Therefore,

Revised: [Equation]

Revised: By combining (18) and (19), we finally prove that,

Revised: [Equation]

Revised: A.3 Proof of Theorem 5.

Revised: Theorem 5.3 (Value function difference bound) .

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

Revised: [Equation]

Revised: Proof.

Revised: We follow Castro et al.

Revised: [4] (the proof of Proposition 4.8) to prove the value function difference bound.

Revised: 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 ) | ≤

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

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

Revised: ( s j ) | ≤ d ( s i , s j ) holds

Revised: [Equation]

Revised: then

Revised: [Equation]

Revised: We make an assumption that (cid:80) s ′ P as ( s ′ )

Revised: V π ( s ′ ) =

Revised: V π ( E s ′ ∼ P as [ s ′ ]) .

Revised: We make this assumption because we use the learned dynamics model to predict the next states

Revised: [Equation]

Revised: and the learned dynamics model is assumed as Gaussian distribution with small standard deviation.

Revised: If this assumption holds, then we have

Revised: [Equation]

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

Revised: By symmetric,

Revised: V

Revised: [Equation]

Revised: s

Revised: [Equation]

Revised: Therefore, | V π ( s i ) − V π ( s j ) | ≤ F πG ( d )( s i , s j ) .

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

Revised: = 2 max s , a | r as | / (1 − γ ) that satisfies | V π

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

Revised: By repeatedly applying F πG on d , d will eventually converge to the fixed-point d πG .

Revised: Because for each iteration F πG ( d ) satisfies | V π ( s j ) − V π ( s i ) | ≤

Revised: F πG ( d )( s i , s j ) , the fixed-point d πG satisfies

Revised: [Equation]

Revised: B Implementation Details B.1 Objectives B.1.1 Soft Actor-critic

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

Revised: SAC is an actor-critic algorithm that optimizes a stochastic policy in an off-policy manner.

Revised: We use the version of SAC that incorporates clipped double-Q trick.

Revised: Let Q θ

Revised: i denote the Q-networks where i = 1 , 2 , Q ¯ θ i be the target Q-networks, and π θ be the actor network.

Revised: All Q-networks, target Q-networks and actor network take input of state representation ϕ ω ( s ) instead of state s .

Revised: The loss functions for the Q-networks are

Revised: [Equation]

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

Revised: We stop gradients for V target and V target is defined as

Revised: [Equation]

Revised: where a ′ ∼ π θ ( ·| ϕ ω ( s ′ )) is the action at next state.

Revised: The policy improvement is minimizing the following objective for actor network,

Revised: [Equation]

Revised: The loss function for α of SAC is,

Revised: [Equation]

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

Revised: B.1.2 RAP Distance

Revised: We learn the reward function and dynamics model by neural network approximators, in order to approximate RAP distance.

Revised: Let r ψ 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: The loss function for learning r ψ is

Revised: [Equation]

Revised: We learn a dynamics model (cid:98) P η parameterized by η to take input of state embedding ϕ ( s ) and action 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 the predictive next state embedding, respectively.

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

Revised: [Equation]

Revised: B.2 Learning Algorithm

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

Revised: B.3 Networks and Hyperparameters

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

Revised: The encoder takes the input of a state, where 3 frames are stacked, and outputs a vector, the state representation ϕ ω ( s ) .

Revised: The detailed dimensions of each convolutional layer and fully-connected layer are described in Table 2.

Revised: The Q-network consists of 3 stacked fully-connected layers with 1024 hidden dimensions, and takes the input of state representation ϕ ω ( s ) and action a .

Revised: The actor network takes the input of ϕ ω ( s ) and feeds it into three stacked fully-connected layers with 1024 hidden dimensions to output policy π .

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

Revised: ReLU activation is used for every layer.

Revised: We train the whole model on one NVIDIA A100 GPU.

Revised: Other hyperparameters are listed in Table 2.

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: C Additional Details of Experiments C.1 DeepMind Control Suite

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

Revised: [Table]

Revised: C.2 CARLA

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

Revised: The observation is 420 × 84 pixels constructed by concatenating five cameras.

Revised: The action repeat for CARLA is 8.

Revised: D Additional Experiments D.1 Ablation study on reward correction and learned dynamics

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

Revised: Gaussian dynamics model.

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

Revised: DMC with natural video background setting.

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

Revised: The curve RAP w/o learned dynamics is our method but measuring distances with next states that are sampled from MDP instead of a learned dynamics model mean/expectation.

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

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

Revised: Without a learned dynamics model mean it will meet a larger performance drop but it is still better than MICo.

Revised: D.2 Experiment on Simple 2D Highway

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

Revised: The 2D highway task contains a straight road with 4 lanes and 50 other vehicles running on.

Revised: Figure 9 shows an illustration of 2D highway and Figure 10 shows the observation which is cropped to a square.

Revised: The green vehicle is the vehicle that RL agent controls and blue vehicles are controlled by the environment.

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

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

Revised: DBC, DrQ, and the proposed method RAP achieve similar scores.

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

Revised: This 2D highway is a simple task so that they converge at around 100K to 200K steps.

Revised: E Visualization of Representation Space

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

Revised: We collect data from cheetah-run with different natural video backgrounds.

Revised: Figure 12 shows the latent representation space where the data points are collected from four different video backgrounds and marked with four different colors.

Revised: RAP and MICo can mix different video background which means they can prevent overfit to the background features, while DBC grouping the same color points into several clusters potentially captures the background task-irrelevant features.

Revised: RAP maps the observations to a structural latent space but MICo maps to a circle latent space which may lose the structure.

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

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

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

Revised: MICo seems to have random Q-values with respect to the representation structure.

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

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

Revised: RAP can map similar observations to the nearby locations, while MICo and DBC map them far away from each other.

Revised: [Figure]

Revised: [Figure]

Revised: F Explanation of Policy Dependence of RAP Loss

Revised: The loss (15) in Section 5.2 should be dependent on the policy.

Revised: More precisely, the expectation is taken over the transitions sampled from a replay buffer D .

Revised: The loss (15) should be written as

Revised: [Equation]

Revised: The rewards r a

Revised: 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 rewards depend on are sampled from the policy when the rewards are generated from the MDP.

Revised: Therefore this

Revised: [Table]

Revised: [Table]

Revised: [Table]

Revised: [Table]

Revised: [Table]

Revised: [Table]

Revised: [Table]

Revised: [Table]

Revised: [Table]

Revised: loss function is dependent on the policy.

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

Revised: This indicates that the learned reward function is also dependent on the policy.

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

Revised: However, off-policy RL algorithms have better sample efficiency, especially in control tasks.

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

Revised: The underlying policy in the replay buffer may not be identical to the current policy.

Revised: We want to verify that even though the learning of RAP is not dependent on the "truly" on-policy, our approximation is not far away from on-policy and is enough to learn good representations.

Revised: We try to reduce the replay buffer capacity.

Revised: If the replay buffer is small enough, the RL algorithm would become on-policy.

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

Revised: We reduce the replay buffer capacity from 1 million to 500K, 100K, 50K, and 10K. Figure 15 shows the learning curves of different replay buffer capacities.

Revised: It shows that smaller replay buffer sizes lead to lower performance and slower convergence.

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

Revised: Then we measure the RAP distance in between the 1K observations by using different models learned by different replay buffer sizes, i.e. 500K, 100K, 50K and 10K.

Revised: We compare the distances among different replay buffer sizes and show the average differences in Table 4.

Revised: The differences of distances stay on the same scale.

Revised: The 1M version which is the proposed method is not far away from the models learned with other buffer sizes.

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