Title: State Chrono Representation for Enhancing Generalization in Reinforcement Learning
Venue: NeurIPS 2024
Accepted: yes
Keywords: bisimulation metric

deep reinforcement learning

representation learning
Abstract: In reinforcement learning with image-based inputs, it is crucial to establish a robust and generalizable state representation. Recent advancements in metric learning, such as deep bisimulation metric approaches, have shown promising results in learning structured low-dimensional representation space from pixel observations, where the distance between states is measured based on task-relevant features. However, these approaches face challenges in demanding generalization tasks and scenarios with non-informative rewards. This is because they fail to capture sufficient long-term information in the learned representations. To address these challenges, we propose a novel State Chrono Representation (SCR) approach. SCR augments state metric-based representations by incorporating extensive temporal information into the update step of bisimulation metric learning. It learns state distances within a temporal framework that considers both future dynamics and cumulative rewards over current and long-term future states. Our learning strategy effectively incorporates future behavioral information into the representation space without introducing a significant number of additional parameters for modeling dynamics. Extensive experiments conducted in DeepMind Control and Meta-World environments demonstrate that SCR achieves better performance comparing to other recent metric-based methods in demanding generalization tasks. The codes of SCR are available in https://github.com/jianda-chen/SCR.

State Chrono Representation for Enhancing Generalization in Reinforcement Learning

Abstract

In reinforcement learning with image-based inputs, it is crucial to establish a robust and generalizable state representation. Recent advancements in metric learning, such as deep bisimulation metric approaches, have shown promising results in learning structured low-dimensional representation space from pixel observations, where the distance between states is measured based on task-relevant features. However, these approaches face challenges in demanding generalization tasks and scenarios with non-informative rewards. This is because they fail to capture sufficient long-term information in the learned representations. To address these challenges, we propose a novel State Chrono Representation ( SCR ) approach. SCR augments state metric-based representations by incorporating extensive temporal information into the update step of bisimulation metric learning. It learns state distances within a temporal framework that considers both future dynamics and cumulative rewards over current and long-term future states. Our learning strategy effectively incorporates future behavioral information into the representation space without introducing a significant number of additional parameters for modeling dynamics. Extensive experiments conducted in DeepMind Control and Meta-World environments demonstrate that SCR achieves better performance comparing to other recent metric-based methods in demanding generalization tasks. The codes of SCR are available in https://github.com/jianda-chen/SCR.

1 Introduction

In deep reinforcement learning (Deep RL), one of the key challenges is to derive an optimal policy from highly dimensional environmental observations, particularly images [5, 12, 33]. The stream of images received by an RL agent contains temporal relationships and significant spatial redundancy. This redundancy, along with potentially distracting visual inputs, poses difficulties for the agent in learning optimal policies. Numerous studies have highlighted the importance of building state representations that can effectively distinguish task-relevant information from task-irrelevant surroundings [45, 46, 8]. Such representations have the potential to facilitate the RL process and improve the generalizability of learned policies. As a result, representation learning has emerged as a fundamental aspect in the advancement of Deep RL algorithms, gaining increased attention within the RL community [21].

Related Works. The main objective of representation learning in RL is to develop a mapping function that transforms high-dimensional observations into low-dimensional embeddings. This process helps reduce the influence of redundant signals and simplify the policy learning process. Previous research has utilized autoencoder-like reconstruction losses [43, 18], which have shown

impressive results in various visual RL tasks. However, due to the reconstruction of all visual input including noise, these approaches often lead to overfitting on irrelevant environmental signals. Data augmentation methods [42, 22, 23] have demonstrated promise in tasks involving noisy observations. These methods primarily enhance perception models without directly impacting the policy within the Markov Decision Process (MDP) framework. Other approaches involve learning auxiliary tasks [33], where the objective is to predict additional tasks related to the environment using the learned representation as input. However, these auxiliary tasks are often designed independently of the primary RL objective, which can potentially limit their effectiveness in improving the overall RL performance.

In recent advancements, behavioral metrics, such as the bisimulation metric [11, 10, 6, 4, 5, 47] and the MICo distance [7], have been developed to measure the dissimilarity between two states by considering differences in immediate reward signals and the divergence of next-state distributions. Behavioral metrics-based methods establish approximate metrics within the representation space, preserving the behavioral similarities among states. This means that state representations are constrained within a structured metric space, wherein each state is positioned or clustered relative to others based on their behavioral distances. Moreover, behavioral metrics have been proven to set an upper bound on state-value discrepancies between corresponding states. By learning behavioral metrics within representations, these methods selectively retain task-relevant features essential for achieving the optimal policy, which involves maximizing the value function and shaping agent behaviors. At the same time, they filter out noise that is unrelated to state values and behavioral metrics, thereby improving robustness in noisy environments.

However, behavioral metrics face challenges when handling demanding generalizable RL tasks and scenarios with sparse rewards [20]. While they can capture long-term behavioral metrics to some extent through the temporal-difference update mechanism, their reliance on one-step transition data limits their learning efficiency, especially in the case of sparse rewards. This limitation may result in representations that encode non-informative signals [20] and can restrict their effectiveness in capturing long-term reward information. Some model-based approaches attempt to mitigate these issues by learning transition models [14, 16, 27]. However, learning a large transition model with long trajectories requires a large number of parameters and significant computational resources. Alternatively, N -step reinforcement learning methods can be used to address the problem [48]. Nevertheless, these methods introduce higher variance in value estimation compared to one-step methods, which may lead to unstable training.

Contributions. To overcome the above challenges, we introduce the State Chrono Representation ( SCR ) learning framework. This metric-based approach enables the learning of long-term behavioral representations and accumulated rewards that span from present to future states. Within the SCR framework, we propose the training of two distinct state encoders. One encoder focuses on crafting a state representation for individual states, while the other specializes in generating a Chronological Embedding that captures the relationship between a state and its future states. In addition to learning the conventional behavioral metric for state representations, we introduce a novel behavioral metric specifically designed for temporal state pairs. This new metric is approximated within the chronological embedding space. To efficiently approximate this behavioral metric in a lower-dimensional vector space, we propose an alternative distance metric that differs from the typical L p norm. To incorporate long-term rewards information into these representations, we introduce a 'measurement' that quantifies the sum of rewards between the current and future states. Instead of directly regressing this measurement, we impose two constraints to restrict its range and value. Note that SCR is a robust representation learning methodology that can be integrated into any existing RL algorithm.

In summary, our contributions are threefold: 1) We introduce a new framework, SCR , for representation learning with a focus on behavioral metrics involving temporal state pairs. We also provide a practical method for approximating these metrics; 2) We develop a novel measurement specifically adapted for temporal state pairs and propose learning algorithms that incorporate this measurement while enforcing two constraints; 3) Through experiments conducted on the Distracting DeepMind Control Suite [34], we demonstrate that our proposed representation exhibits enhanced generalization and efficiency in challenging generalization tasks.

2 Preliminary

Markov Decision Process: AMarkov Decision Process (MDP) is a mathematical framework defined by the tuple M = ( S , A , P, r, γ ) . Here S represents the state space, which encompasses all possible states. A denotes the action space, comprising all possible actions that an agent can take in each state. The state transition probability function, P , defines the probability of transitioning from the current state s t ∈ S to any subsequent state s t +1 ∈ S at time t . Given the action a t ∈ A taken, the probability is denoted as P ( s t +1 | s t , a t ) . The reward function, r : S × A → R , assigns an immediate reward r ( s t , a t ) to the agent for taking action a t in state s t . The discount factor, γ ∈ [0 , 1] , determines the present value of future rewards.

A policy, π : S × A , is a strategy that specifies the probability π ( a | s ) of taking action a in state s . The objective in an MDP is to find the optimal policy π ∗ that maximizes the expected discounted cumulative reward, π ∗ = arg max π E a ∼ π [ ∑ t γ t r ( s t , a t )] .

Behavioral Metric: In DBC [47], the bisimulation metric defines a pseudometric d : S × S → R to measure the distance between two states. Additionally, a variant of bisimulation metric, known as π -bisimulation metric, is defined with respect to a specific policy π .

Theorem 2.1 ( π -bisimulation metric) . The π -bisimulation metric update operator F bisim : M → M is defined as

where M is the space of d , r π x = ∑ a ∈A π ( a | x ) r a x , P π x = ∑ a ∈A π ( a | x ) P a x , W is the Wasserstein distance, and r a x is r ( x , a ) for short. F bisim has a unique least fixed point d π bisim .

MICo [7] defines an alternative metric that samples the next states instead of directly measuring the intractable Wasserstein distance.

Theorem 2.2 (MICo distance) . The MICo distance update operator F MICo : M → M is defined as

F MICo has a fixed point d π MICo .

3 State Chrono Representation

Although the bisimulation metric [47] and MICo [7] have their strengths, they do not adequately capture future information. This limitation can degrade the effectiveness of state representations in policy learning. To address this issue and incorporate future details, we propose a novel approach, State Chrono Representation ( SCR ). The detailed architecture of SCR is illustrated in Figure 1.

Γ3ΓRΓOΓLΓFΓ'ΓΛΓ1ΓHΓWΓZΓRΓUΓN

SCR consists of two key components: a state representation ϕ ( x ) ∈ R n for a given state x , and a chronological embedding ψ ( x i , x j ) ∈ R n that captures the relationship between a current state x i and its future state x j within the same trajectory. The state representation, ϕ ( x ) , is developed using a behavioral metric d that quantifies the reward and dynamic divergence between two states. On the other hand, the chronological embedding, ψ ( x i , x j ) , fuses these two state representations using deep learning, with a focus on capturing the long-term behavioral correlation between the current state x i and the future state x j . To compute the distance between any two chronological embeddings, we propose a 'chronological' behavioral metric, which is further improved through a Bellman operator-like MSE loss [36]. Moreover, ϕ ( x ) incorporates a novel temporal measurement m to evaluate the transition from the current state to a future state, effectively capturing sequential reward data from the optimal behavior. Due to the complexity of the regression target, this temporal measurement operates within the defined lower and upper constraints, providing a stable prediction target for both the measurement and state representation during the learning process.

Γ4ΓΛΓ1ΓHΓW ZΓRΓUΓN

3.1 Metric Learning for State Representation

The state representation encoder ϕ is trained by approximating a behavioral metric, such as the MICo distance. In our model, we utilize a MICo-based metric transformation operator. Instead of using sampling-based prediction in MICo, we employ latent dynamics-based modeling to assess the divergence between two subsequent state distributions. This approach draws inspiration from the methodology in SimSR [45]. The metric update operator for latent dynamics, denoted by F , is defined as follows.

Theorem 3.1. Let ˆ d : R n × R n → R be a metric in the latent state representation space, d ϕ ( x i , y i ′ ) := ˆ d ( ϕ ( x i ) , ϕ ( y i ′ )) be a metric in the state domain. The metric update operator F is defined as,

where a x i and a y i ′ are the actions in states x i and y i ′ , respectively, and ˆ P is the learned latent dynamics model. F has a fixed point d π ϕ .

To learn the approximation for d π ϕ in the embedding space, one needs to specify the form of distance ˆ d for latent vectors. Castro et al. [7] showed that a behavioral metric with a sample-based next state distribution divergence is a diffuse metric, as the divergence of the next state distribution corresponds to the Łukaszyk-Karmowski distance [49] that measures the expected distance between two samples drawn from two distributions, respectively (detailed definition is in Appendix A.4.1).

Definition 3.2 (Diffuse metric [7]) . A function d : X × X → R based on a vector space X is a diffuse metric if the following axioms hold: 1) d ( a , b ) ≥ 0 for any a , b ∈ X ; 2) d ( a , b ) = d ( b , a ) for any a , b ∈ X ; 3) d ( a , b ) + d ( b , c ) ≥ d ( a , c ) for any a , b , c ∈ X .

MICo provides an approximation of the behavioral metric using an angular distance: ˆ d MICo ( a , b ) = ∥ a ∥ 2 2 + ∥ b ∥ 2 2 2 + βθ ( a , b ) , where a , b ∈ R n , θ ( a , b ) represents the angle between vectors a and b , and β = 0 . 1 is a hyperparameter. This distance calculation includes a non-zero self-distance, which makes it compatible with expressing the Łukaszyk-Karmowski distance [49]. However, the angle function θ ( a , b ) only considers the angle between a and b , which requires computations involving cosine similarity and the arccos function. This can sometimes lead to numerical discrepancies. DBC [47] recommends using the L 1 norm with zero self-distance, which is only suitable for the Wasserstein distance. On the other hand, SimSR [45] utilizes the cosine distance, derived from cosine similarity, but it does not satisfy the triangle inequality and does not have a non-zero self-distance.

To address the challenges mentioned above, we propose a revised distance, ˆ d ( a , b ) , in the embedding space. It is characterized as a diffuse metric and formulated as follows:

Definition 3.3. We define ˆ d : R n × R n → R as a distance function, where ˆ d ( a , b ) = √ ∥ a ∥ 2 2 + ∥ b ∥ 2 2 -a ⊤ b , for any a , b ∈ R n .

Theorem 3.4. ˆ d is a diffuse metric.

Proof. Refer to Appendix A.3.2 for the detailed proof.

Lemma 3.5 (Non-zero self-distance) . The self-distance of ˆ d is not strict to zero, i.e., ˆ d ( a , a ) = ∥ a ∥ 2 ≥ 0 . It becomes zero if and only if every element in vector a is zero.

Theorem 3.4 validates that ˆ d is a diffuse metric satisfying triangle inequality, and Lemma 3.5 shows ˆ d has non-zero self-distance capable of approximating dynamic divergence, which is a ŁukaszykKarmowski distance. Moreover, the structure of ˆ d closely resembles the L 2 norm, with the only difference being that the factor of a ⊤ b is -1 instead of -2 . This construction, which involves only vector inner product and square root computations without divisions or trigonometric functions, helps avoid numerical computational issues and simplify the implementation.

To learn the representation function ϕ , a common approach is to minimize MSE loss between the two ends of (1). The loss, which incorporates ˆ d and is dependent on ϕ , can be expressed as follows:

where D represents the replay buffer or the sampled rollouts used in the RL learning process.

3.1.1 Long-term Temporal Information

The loss L ϕ ( ϕ ) in Eqn. (2) heavily relies on the temporaldifference update mechanism, utilizing only one-step transition information. However, this approach fails to effectively capture and encode long-term information from trajectories. Incorporating temporal information into the representation while maintaining the structure and adhering to behavioral metrics presents a complex challenge. To address this challenge, we propose two distinct methods: Chronological Embedding (Section 3.2) and Temporal Measurement (Section 3.3). Each technique is designed

Figure 2: An example with two rollouts. to capture the temporal essence of a rollout denoted as τ ( x i ; x j ) , which represents a sequence starting from state x i and reaching a future state x j .

As illustrated in Figure 2, the goal of the chronological embedding is to create a novel paired state embedding, denoted as ψ ( x i , x j ) , capturing the temporal representation of the rollout τ ( x i ; x j ) . This is achieved by transforming a pair of state representations ϕ ( x i ) and ϕ ( x j ) to a vector: ψ ( x i , x j ) := ψ ( ϕ ( x i ) , ϕ ( x j )) ∈ R n . The objective of learning ψ is to develop a 'chronological' behavioral metric d ψ that quantifies the distance between rollouts τ ( x i ; x j ) and τ ( y i ′ ; y j ′ ) . On the other hand, temporal measurement aims to compute a specific 'distance' m between states x i and x j , which provides insights into the cumulative rewards accumulated during the rollout τ ( x i ; x j ) .

Note that learning the state representation solely based on either d ψ or m poses challenges. While d ψ improves the representation expressiveness, capturing both rewards and dynamics may be less efficient in the learning process. On the other hand, learning m is efficient but may not adequately capture the dynamics. Therefore, we propose a unified approach that leverages the strengths of both d ψ and m to overcome the challenges and improve the overall state representation. Empirical results in Section 4.3 support the effectiveness of our proposal.

3.2 Chronological Embedding

The chronological embedding , denoted by ψ ( x i , x j ) ∈ R n , is designed to capture the relationship between a given state x i and any of its future states x j over the same trajectory. To capture the comprehensive behavioral knowledge, we introduce a distance function d ψ : ( S×S ) × ( S×S ) → R , which is intended to reflect the behavioral metric and enable the encoder ψ to incorporate behavioral information. Building upon the MICo distance in Theorem 2.2, we specify the metric update operator F Chrono for d ψ in the following theorem. Proof is detailed in Appendix A.3.3.

Theorem 3.6. Let M ψ be the space of d ψ . The metric update operator F Chrono : M ψ → M ψ is defined as follows,

ψ

Here d π ψ represents the 'chronological' behavioral metric. Our objective is to closely approximate d π ψ , which measures the distance between two sets of states ( x i , x j ) and ( y i ′ , y j ′ ) , taking into account the difference in immediate rewards and dynamics divergence. To co-learn the encoder ψ with d π ψ , we represent d π ψ as d π ψ ( x i , x j , y i ′ , y j ′ ) := ˆ d ( ψ (( x i , x j )) , ψ ( y i ′ , y j ′ )) , where ˆ d is defined in Definition 3.3. Similar to Eqn. (1), we construct d π ψ to compute the distance in the embedding space.

To ensure computational efficiency, the parameters of the encoders ϕ and ψ are shared. The encoder ψ extracts outputs from ϕ , and the distance measure is appropriately adapted as d π ψ ( x i , x j , y i ′ , y j ′ ) := ˆ d ( ψ ( ϕ ( x i ) , ϕ ( x j )) , ψ ( ϕ ( y i ′ ) , ϕ ( y j ′ ))) . The objective of learning the chronological embedding is to minimize the MSE loss between both sides of Eqn. (4) w.r.t. ϕ and ψ ,

Minimizing Eqn. (5) is to encourage the embeddings of similar state sequences to become closer in the embedded space, thereby strengthening the classification of similar behaviors.

3.3 Temporal Measurement

To enhance the ability of SCR to gain future insights, we introduce the concept of temporal measurement , which quantifies the disparities between a current state x i and a future state x j . This measurement, denoted by m ( x i , x j ) , aims to measure the differences in state values or the cumulative rewards obtained between the current and future states. We construct an approximate measurement based on the state representation ϕ , denoted by ˆ m ϕ ( x i , x j ) := ˆ m ( ϕ ( x i ) , ϕ ( x j )) . Here, ˆ m : R n × R n → R is a non-parametric asymmetric metric function (details are presented at the end of this section). This approach structures the representation space of ϕ ( x ) around the 'distance' m , ensuring that ϕ ( x ) contains sufficient information for future state planning.

We propose using m to represent the expected discounted accumulative rewards obtained by an optimal policy π ∗ from state x i to state x j : m ( x i , x j ) = E π ∗ [ ∑ j -i t =0 γ t r s t ∣ ∣ ∣ s 0 = x i , s j -i = x j ] . However, obtaining the exact value of m ( x i , x j ) is challenging because the optimal policy π ∗ is unknown and is, in fact, the primary goal of the RL task. Instead of directly approximating m ( x i , x j ) , we learn the approximation ˆ m ( x i , x j ) in an alternative manner that ensures it falls within a feasible range covering the true m ( x i , x j ) .

To construct this range, we introduce two constraints. The first constraint, serving as a lower boundary, states that the expected discounted cumulative reward obtained by any policy π , whether optimal or not, cannot exceed m :

This constraint is based on a fact that any sub-optimal policy is inferior to the optimal policy. Based on this constraint, we propose the following objective to learn the approximation ˆ m ϕ :

where ReLU ( x ) = max (0 , x ) . This objective is non-zero when the constraint in Eqn. (6) is not satisfied. Hence, it increases the value of m ( ϕ ( x i ) , ϕ ( x j )) until it exceeds the sampled reward sum.

The second constraint serving as the upper boundary is proposed based on Fig. 2. To ensure that the absolute value | m ( x i , x j ) | remains within certain limits, we propose the following inequality:

where d is the behavioral metric introduced in Section 3.1. The right-hand side of the inequality represents the longer path from x i to x j . This inequality demonstrates that the absolute temporal measurement | m ( x i , x j ) | should not exceed the sum of the behavioral metrics at the initial states ( x i and y i ′ ), i.e., d ( x i , y i ′ ) , the final states pair ( x j and y j ′ ), i.e. d ( x j , y j ′ ) , and the measurement | m ( y i , y j ) | . This constraint leads to the following objective for training ˆ m ϕ :

where sg means 'stop gradient'. This objective aims to decrease the value of ˆ m ϕ when the constraint in Eqn. (8) is not satisfied. By simultaneously optimizing both constraints in a unified manner,

we guarantee that the approximated temporal measurement, m , is confined within a specific range, defined by the lower and upper constraints. The overall objective for ˆ m is formulated as follows:

Asymmetric Metric Function for ˆ m . The measurement ˆ m ( x i , x j ) , designed to measure the distance based on the rewards, is intended to be asymmetric with respect to x i and x j . This is because we assume that state x i comes before x j , resulting in a distinct relationship compared to the progression from x j to x i . Recent research has focused on studying quasimetrics in deep learning [28, 39, 38] and developed various methodologies to compute asymmetric distances. In our method, we choose to utilize Interval Quasimetric Embedding (IQE) [38] to implement ˆ m (details are in Appendix A.5.1).

3.4 Overall Objective

As shown in Figure 1, the encoders are designed to predict state representations ϕ ( x ) for individual states and chrono embedding ψ ( x i , x j ) to capture the relationship between states x i and x j . The measurement ˆ m is then computed based on ϕ ( x i ) and ϕ ( x j ) to incorporate the accumulated rewards between these states. The components ψ and ˆ m work together to enhance the state representations ϕ and capture temporal information, thereby improving their predictive capabilities for future insights. The necessity of ψ and ˆ m is demonstrated in the ablation study in Section 4.3. Therefore, a comprehensive objective to minimize is formulated in a unified manner as follows,

Our method, SCR , can be seamlessly integrated with a wide range of deep RL algorithms, which can effectively utilize the representation ϕ ( x ) as a crucial input component. In our implementation, we specifically employ Soft Actor-Critic (SAC) [13] as our foundation RL algorithm. The state representation serves as the input state for both the policy network and Q-value network in SAC. Other implementation details can be found in Appendix A.5.

4 Experiments

4.1 Configurations

DeepMind Control Suite. The primary objective of our proposed SCR is to develop a robust and generalizable representation for deep RL when dealing with high-dimensional observations. To evaluate its effectiveness, we conduct experiments using the DeepMind Control Suite (DM_Control) environment, which involves rendered pixel observations [37] and a distraction setting called Distracting Control Suite [34]. This environment utilizes the MuJoCo engine, offering pixel observations for continuous control tasks. The DM_Control environment provides diverse testing scenarios, including background and camera pose distractions, simulating real-world complexities of camera inputs. (1) Default setting evaluates the effectiveness of each RL approach. Frames are rendered at a resolution of 84 × 84 RGB pixels as shown in Figure 3(a). We stack three frames as states for the RL agents. (2) Distraction setting evaluates both the generalization and effectiveness of each RL approach. The distraction [34] consists of 3 components, as shown in Figure 3(b): 1) background video distraction, where the clean and simple background is replaced with a natural video; 2) object color distraction, which involves slight changes in the color of the robots; and 3) camera pose distraction, where the camera's position and angle are randomized during rendering from the simulator. We observe that tasks become significantly more challenging when the camera pose distraction is applied.

Baselines. We compare our method against several prominent algorithms in the field, including: 1) SAC [13], a baseline deep RL method for continuous control; 2) DrQ [42], a data augmentation

(a)

(b)

Median

IQM

Mean

OptimalityGap

SAC

DrQ

DBC

MICo

SimSR

SCR

0.15

0.30

0.1

0.2

0.3

0.1

0.2

0.3

0.4

0.7

0.8

0.9

normalizedscores

method using random crop; 3) DBC [47], representation learning with the bisimulation metric; 4) MICo [7], representation learning with MICo distance; and 5) SimSR [45], representation learning with the behavioral metric approximated by the cosine distance.

4.2 Main Results

Default Setting. To verify the sample efficiency of our method, we compare it with others on eight tasks in DM_Control and report the experimental results in Table 1. We trained each method on each task for 500K environment steps. All results are averaged over 10 runs. We observe that SCR achieves comparable results to the augmentation method DrQ and the state-of-the-art behavioral metric approach SimSR. It is worth noting that for a DM_Control task, the maximum achievable scores are 1000, and a policy that collects scores around 900 is considered nearly optimal. These findings highlight the effectiveness of our SCR in mastering standard RL control tasks.

Distraction Setting. To further evaluate the generalization ability of our method, we perform comparison experiments on DM_Control with Distraction using the same training configuration as in the default setting. We evaluate 10 different seeds, and conduct 100 episodes per run. The scores from each run are summarized in Table 2. The std is computed across the mean scores of these 10 runs. The results show that our method outperforms the others in all eight tasks, notably in the sparse reward task cartpole-swing_up_sparse , where other methods achieve nearly zero scores. The camera-pose distraction poses a significant challenge for metric-based methods like DBC, MICo, and SimSR, as it leads to substantial distortion of the robot's shape and position in the image state. DrQ outperforms other behavioral metric methods, possibly due to its random cropping technique, which facilitates better alignment of the robot's position. The aggregate metrics [2] are shown in Figure 4, where scores are normalized by dividing by 1000. Figure 5 presents the training curves, while additional curves can be found in Appendix B.3. These results demonstrate the superior performance of our method in handling distractions and highlight our strong generalization capabilities.

4.3 Ablation Study

Impact of Different Components. To evaluate the impact of each component in SCR , we perform an ablation study where certain components are selectively removed or substituted. Figure 6 shows the training curves on cheetah-run and walker-walk in the distraction setting. SCR denotes the full

cartpole_swingup_sparse

cheetah_run

hopper_stand

walker_walk

250

200

300

400

600

500

300

400

return

150

200

200

300

100

50

0

SAC

DrQ

SimSR

SCR

DBC

MICo

100

0

SimSR

SCR

SAC

DrQ

DBC

MICo

100

0

SCR

SAC

DrQ

DBC

MICo

SimSR

200

100

0

SimSR

SCR

SAC

DrQ

DBC

MICo

0

1

environment steps

2

3

4

1e5

5

Q

1

environment steps

2

3

4

1e5

5

0

1

environment steps

2

3

4

1e5

5

0

1

environment steps

2

3

4

1e5

5

cheetah_run

walker_walk

600

300

500

400

return

200

300

100

SCR

SCR wo

SCR w/o mh

200

SCRw/o

SCR w/o m

SCR

SCR  only

SCR wf cos

SCR w/ MICo

100

SCR  only

SCR w/ cos

SCR w/MICo

0

SCR w/L1

0

SCR w/ L1

Q

1

environment steps

2

3

4

1e5

5

Q

1

environment steps

2

E

4

1e5

5

cheetah_run

walker_walk

600

300

500

return

200

return

400

300

100

n[1.10]

n[1,50]

200

n[1.10]

n~[1,50]

n[1.100](SCR)

n~[1,150]

n=50

100

(805001']u

n~[1,150]

n=50

0

n=100

n=100

0

1

environment steps

2

m

4

1e5

5

1

environment steps

2

3

4

1e5

5

model. SCR w/o ψ removes the chronological embedding ψ . SCR w/o ˆ m refers to the exclusion of the approximation ˆ m . SCR w/ ϕ only removes losses L ψ ( ϕ, ψ ) and L ˆ m ( ϕ ) but keep L ϕ ( ϕ ) . SCR w/ cos replaces the distance function ˆ d for computing metrics on the representation space with cosine distance, akin to SimSR does. SCR w/ MICo replaces ˆ d with MICo's angular distance. SCR w/ L1 replaces ˆ d with the L 1 distance as adopted by DBC. The results demonstrate the superior performance of the full model and the importance of ψ and ˆ m . The absence of these components can lead to worse performance and unstable training. The results also demonstrate that the proposed ˆ d surpasses existing distance functions.

Impact of Sampling Steps. In our previous experiments, we uniformly sample the number of steps between i and j in the range [1 , 100] . In this experiment, we investigate the impact of this hyper-parameter on SCR . We conduct experiments on the cheetah-run and walker-walk tasks in the distraction setting with various step sampling ranges: [1 , 10] , [1 , 50] , [1 , 100] , and [1 , 150] . We also make comparisons with fixed step counts: 50 and 100. The results presented in Figure 7 show that sampling step counts in the range [1 , 100] yields the optimal results and remains stable across tasks. Therefore, we set this hyper-parameter to uniformly sample in the range [1 , 100] for all experiments.

4.4 Experiments on Meta-World

We present experimental investigations conducted in Meta-World [44], a comprehensive simulated benchmark that includes distinct robotic manipulation tasks. Our focus is on six specific tasks: basketball-v2 , coffee-button-v2 , door-open-v2 , drawer-open-v2 , pick-place-v2 , and window-open-v2 , chosen for their diverse objects in environments. The observations are rendered as 84 × 84 RGB pixels with a frame stack of 3, following the DM_Control format. Table 3 displays the average success rates over all tasks and five seeds, while Figure 8 shows the training curves for a subset of these tasks. Additional training curves for all tasks can be found in Appendix B.5. Our proposed SCR consistently achieves optimal performance in all tasks, while other baseline methods, except for DrQ, fail to converge to such high levels of performance. Even though DrQ demonstrates proficiency in achieving optimal performance, it shows less sampling efficiency than SCR , highlighting the effectiveness of SCR in the applied setting.

5 Conclusion

Our proposed State Chrono Representation learning framework ( SCR ) effectively captures information from complex, high-dimensional, and noisy inputs in deep RL. SCR improves upon previous metricbased approaches by incorporating a long-term perspective and quantifying state distances across

basketball-v2

door-open-v2

pick-place-v2

window-open-v2

1.2

1.2

1.2

1.2

1.0

1.0

1.0

1.0

rate

0.8

rate

0.8

rate

0.8

rate

0.8

success

0.6

Success

0.6

Success

0.6

success

0.6

0.4

0.4

0.4

0.4

0.2

SAC

SCR

0.2

SCR

SAC

0.2

SCR

SAC

0.2

SCR

SAC

0.2

0.0

DrQ

DBC

MICo

SimSR

0.2

0.0

MICo

DBC

DrQ

SimSR

0.2

0.0

DrQ

DBC

MICo

SimSR

0.2

0.0

DrQ

DBC

MICo

SimSR

0

1

2

3

4

5

0

1

2

3

4

5

0.0

0.2

0.4

0.6

0.8

1.0

0

1

2

3

4

5

environment steps

1e5

environment steps

1e5

environment steps

1e6

environment steps

1e5

temporal trajectories. Our extensive experiments demonstrate its effectiveness compared with several metric-based baselines in complex environments with distractions, making a promising avenue for future research on generalization in representation learning.

Limitations. This work does not address truly Partially Observable MDPs (POMDPs), which are common in real-world applications. In future work, we plan to integrate SCR with POMDP methods to solve real-world problems, or to design unified solutions that better accommodate these complexities.

References

The successor representation in human reinforcement learning. Nature human behaviour , 1(9): 680-692, 2017.

Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron Courville, and Marc G Bellemare. Deep reinforcement learning at the edge of the statistical precipice. Advances in Neural Information Processing Systems , 2021.

André Barreto, Will Dabney, Rémi Munos, Jonathan J Hunt, Tom Schaul, Hado P van Hasselt, and David Silver. Successor features for transfer in reinforcement learning. Advances in neural information processing systems , 30, 2017.

Pablo S. Castro and Doina Precup. Using bisimulation for policy transfer in mdps. In Proceedings of the 9th International Conference on Autonomous Agents and Multiagent Systems: Volume 1 - Volume 1 , AAMAS '10, Richland, SC, 2010. International Foundation for Autonomous Agents and Multiagent Systems. ISBN 9780982657119.

Pablo Samuel Castro. Scalable methods for computing state similarity in deterministic markov decision processes. Proceedings of the AAAI Conference on Artificial Intelligence , 34(06): 10069-10076, Apr. 2020. doi: 10.1609/aaai.v34i06.6564. URL https://ojs.aaai.org/ index.php/AAAI/article/view/6564 .

Pablo Samuel Castro, Prakash Panangaden, and Doina Precup. Equivalence relations in fully and partially observable markov decision processes. In Proceedings of the 21st International Joint Conference on Artificial Intelligence , IJCAI'09, San Francisco, CA, USA, 2009. Morgan Kaufmann Publishers Inc.

Pablo Samuel Castro, Tyler Kastner, Prakash Panangaden, and Mark Rowland. MICo: Improved representations via sampling-based state similarity for markov decision processes. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems , 2021. URL https://openreview.net/forum?id=wFp6kmQELgu .

Jianda Chen and Sinno Pan. Learning representations via a robust behavioral metric for deep reinforcement learning. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems , volume 35, pages 36654-36666. Curran Associates, Inc., 2022. URL https://proceedings.neurips.cc/paper_files/ paper/2022/file/eda9523faa5e7191aee1c2eaff669716-Paper-Conference.pdf .

Maxime Chevalier-Boisvert, Bolun Dai, Mark Towers, Rodrigo Perez-Vicente, Lucas Willems, Salem Lahlou, Suman Pal, Pablo Samuel Castro, and Jordan Terry. Minigrid & miniworld: Modular & customizable reinforcement learning environments for goal-oriented tasks. In Advances in Neural Information Processing Systems 36, New Orleans, LA, USA , December 2023.

Norm Ferns and Doina Precup. Bisimulation metrics are optimal value functions. UAI'14, pages 210-219, Arlington, Virginia, USA, 2014. AUAI Press. ISBN 9780974903910.

Norm Ferns, Prakash Panangaden, and Doina Precup. Metrics for finite markov decision processes. UAI '04, pages 162-169, Arlington, Virginia, USA, 2004. AUAI Press. ISBN 0974903906.

Carles Gelada, Saurabh Kumar, Jacob Buckman, Ofir Nachum, and Marc G. Bellemare. DeepMDP: Learning continuous latent space models for representation learning. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning , volume 97 of Proceedings of Machine Learning Research , pages 2170-2179. PMLR, 09-15 Jun 2019. URL https://proceedings.mlr.press/v97/gelada19a.html .

Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Offpolicy maximum entropy deep reinforcement learning with a stochastic actor. In Jennifer Dy and Andreas Krause, editors, Proceedings of the 35th International Conference on Machine Learning , volume 80 of Proceedings of Machine Learning Research , pages 1861-1870. PMLR, 10-15 Jul 2018. URL https://proceedings.mlr.press/v80/haarnoja18b.html .

Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. In International Conference on Learning Representations , 2019.

Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In International Conference on Machine Learning , pages 2555-2565, 2019.

Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models. arXiv preprint arXiv:2301.04104 , 2023.

Philippe Hansen-Estruch, Amy Zhang, Ashvin Nair, Patrick Yin, and Sergey Levine. Bisimulation makes analogies in goal-conditioned reinforcement learning. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato, editors, Proceedings of the 39th International Conference on Machine Learning , volume 162 of Proceedings of Machine Learning Research , pages 8407-8426. PMLR, 17-23 Jul 2022. URL https://proceedings.mlr.press/v162/hansen-estruch22a.html .

Irina Higgins, Arka Pal, Andrei Rusu, Loic Matthey, Christopher Burgess, Alexander Pritzel, Matthew Botvinick, Charles Blundell, and Alexander Lerchner. DARLA: Improving zero-shot transfer in reinforcement learning. In Doina Precup and Yee Whye Teh, editors, Proceedings of the 34th International Conference on Machine Learning , volume 70 of Proceedings of Machine Learning Research , pages 1480-1490. PMLR, 06-11 Aug 2017. URL https://proceedings. mlr.press/v70/higgins17a.html .

Michael Janner, Igor Mordatch, and Sergey Levine. Gamma-models: Generative temporal difference learning for infinite-horizon prediction. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems , volume 33, pages 1724-1735. Curran Associates, Inc., 2020. URL https://proceedings.neurips.cc/ paper_files/paper/2020/file/12ffb0968f2f56e51a59a6beb37b2859-Paper.pdf .

Mete Kemertas and Tristan Ty Aumentado-Armstrong. Towards robust bisimulation metric learning. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems , 2021. URL https://openreview.net/forum? id=ySFGlFjgIfN .

Robert Kirk, Amy Zhang, Edward Grefenstette, and Tim Rocktäschel. A survey of zero-shot generalisation in deep reinforcement learning. J. Artif. Int. Res. , 76, feb 2023. ISSN 1076-9757. doi: 10.1613/jair.1.14174. URL https://doi.org/10.1613/jair.1.14174 .

Michael Laskin, Aravind Srinivas, and Pieter Abbeel. CURL: Contrastive unsupervised representations for reinforcement learning. In Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning , volume 119 of Proceedings of Machine Learning Research , pages 5639-5650. PMLR, 13-18 Jul 2020. URL https://proceedings.mlr.press/v119/laskin20a.html .

Misha Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, Pieter Abbeel, and Aravind Srinivas. Reinforcement learning with augmented data. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems , volume 33, pages 19884-19895. Curran Associates, Inc., 2020. URL https://proceedings.neurips.cc/ paper_files/paper/2020/file/e615c82aba461681ade82da2da38004a-Paper.pdf .

Alex X. Lee, Anusha Nagabandi, Pieter Abbeel, and Sergey Levine. Stochastic latent actor-critic: Deep reinforcement learning with a latent variable model. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems , volume 33, pages 741-752. Curran Associates, Inc., 2020. URL https://proceedings.neurips.cc/paper_files/paper/2020/file/ 08058bf500242562c0d031ff830ad094-Paper.pdf .

Kuang-Huei Lee, Ian Fischer, Anthony Liu, Yijie Guo, Honglak Lee, John Canny, and Sergio Guadarrama. Predictive information accelerates learning in rl. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems , volume 33, pages 11890-11901. Curran Associates, Inc., 2020. URL https://proceedings.neurips.cc/paper_files/paper/2020/file/ 89b9e0a6f6d1505fe13dea0f18a2dcfa-Paper.pdf .

Qiyuan Liu, Qi Zhou, Rui Yang, and Jie Wang. Robust representation learning by clustering with bisimulation metrics for visual reinforcement learning with distractions. In Proceedings of the Thirty-Seventh AAAI Conference on Artificial Intelligence and Thirty-Fifth Conference on Innovative Applications of Artificial Intelligence and Thirteenth Symposium on Educational Advances in Artificial Intelligence , AAAI'23/IAAI'23/EAAI'23. AAAI Press, 2023. ISBN 978-1-57735-880-0. doi: 10.1609/aaai.v37i7.26063. URL https://doi.org/10.1609/ aaai.v37i7.26063 .

Tung D Nguyen, Rui Shu, Tuan Pham, Hung Bui, and Stefano Ermon. Temporal predictive coding for model-based planning in latent space. In International Conference on Machine Learning , pages 8130-8139. PMLR, 2021.

Silviu Pitis, Harris Chan, Kiarash Jamali, and Jimmy Ba. An inductive bias for distances: Neural nets that respect the triangle inequality. In International Conference on Learning Representations , 2020. URL https://openreview.net/forum?id=HJeiDpVFPr .

Jordi Pont-Tuset, Federico Perazzi, Sergi Caelles, Pablo Arbeláez, Alexander Sorkine-Hornung, and Luc Van Gool. The 2017 davis challenge on video object segmentation. arXiv:1704.00675 , 2017.

Chris Reinke and Xavier Alameda-Pineda. Successor feature representations. Transactions on Machine Learning Research , 2023. ISSN 2835-8856. URL https://openreview.net/ forum?id=MTFf1rDDEI .

John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347 , 2017.

Max Schwarzer, Ankesh Anand, Rishab Goel, R Devon Hjelm, Aaron Courville, and Philip Bachman. Data-efficient reinforcement learning with self-predictive representations. In International Conference on Learning Representations , 2021. URL https://openreview.net/ forum?id=uCQfPZwRaUu .

Younggyo Seo, Danijar Hafner, Hao Liu, Fangchen Liu, Stephen James, Kimin Lee, and Pieter Abbeel. Masked world models for visual control. In 6th Annual Conference on Robot Learning , 2022. URL https://openreview.net/forum?id=Bf6on28H0Jv .

Austin Stone, Oscar Ramirez, Kurt Konolige, and Rico Jonschkowski. The distracting control suite - a challenging benchmark for reinforcement learning from pixels. arXiv preprint arXiv:2101.02722 , 2021.

Adam Stooke, Kimin Lee, Pieter Abbeel, and Michael Laskin. Decoupling representation learning from reinforcement learning. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning , volume 139 of Proceedings of Machine Learning Research , pages 9870-9879. PMLR, 18-24 Jul 2021. URL https://proceedings. mlr.press/v139/stooke21a.html .

Richard S Sutton. Reinforcement learning: An introduction. A Bradford Book , 2018.

Saran Tunyasuvunakool, Alistair Muldal, Yotam Doron, Siqi Liu, Steven Bohez, Josh Merel, Tom Erez, Timothy Lillicrap, Nicolas Heess, and Yuval Tassa. dm_control: Software and tasks for continuous control. Software Impacts , 6:100022, 2020. ISSN 2665-9638. doi: https:// doi.org/10.1016/j.simpa.2020.100022. URL https://www.sciencedirect.com/science/ article/pii/S2665963820300099 .

Tongzhou Wang and Phillip Isola. Improved representation of asymmetrical distances with interval quasimetric embeddings. In NeurIPS 2022 Workshop on Symmetry and Geometry in Neural Representations , 2022. URL https://openreview.net/forum?id=KRiST_rzkGl .

Tongzhou Wang and Phillip Isola. On the learning and learnability of quasimetrics. In International Conference on Learning Representations , 2022. URL https://openreview. net/forum?id=y0VvIg25yk .

Tongzhou Wang, Antonio Torralba, Phillip Isola, and Amy Zhang. Optimal goal-reaching reinforcement learning via quasimetric learning. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett, editors, Proceedings of the 40th International Conference on Machine Learning , volume 202 of Proceedings of Machine Learning Research , pages 36411-36430. PMLR, 23-29 Jul 2023. URL https://proceedings.mlr.press/v202/wang23al.html .

Denis Yarats, Rob Fergus, Alessandro Lazaric, and Lerrel Pinto. Reinforcement learning with prototypical representations. 2021.

Denis Yarats, Ilya Kostrikov, and Rob Fergus. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. In International Conference on Learning Representations , 2021. URL https://openreview.net/forum?id=GY6-6sTvGaf .

Denis Yarats, Amy Zhang, Ilya Kostrikov, Brandon Amos, Joelle Pineau, and Rob Fergus. Improving sample efficiency in model-free reinforcement learning from images. Proceedings of the AAAI Conference on Artificial Intelligence , 35(12):10674-10681, May 2021. doi: 10.1609/aaai.v35i12.17276. URL https://ojs.aaai.org/index.php/AAAI/article/ view/17276 .

Tianhe Yu, Deirdre Quillen, Zhanpeng He, Ryan Julian, Karol Hausman, Chelsea Finn, and Sergey Levine. Meta-world: A benchmark and evaluation for multi-task and meta reinforcement learning. In Conference on Robot Learning (CoRL) , 2019. URL https://arxiv.org/abs/ 1910.10897 .

Hongyu Zang, Xin Li, and Mingzhong Wang. Simsr: Simple distance-based state representations for deep reinforcement learning. Proceedings of the AAAI Conference on Artificial Intelligence , 36(8):8997-9005, Jun. 2022. doi: 10.1609/aaai.v36i8.20883. URL https://ojs.aaai.org/ index.php/AAAI/article/view/20883 .

Xianghua Zeng, Hao Peng, Angsheng Li, Chunyang Liu, Lifang He, and Philip S Yu. Hierarchical state abstraction based on structural information principles. In Proceedings of the Thirty-Second International Joint Conference on Artificial Intelligence , pages 4549-4557, 2023.

Amy Zhang, Rowan Thomas McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learning invariant representations for reinforcement learning without reconstruction. In International Conference on Learning Representations , 2021. URL https://openreview.net/ forum?id=-2FCwDKRREu .

Junmin Zhong, Ruofan Wu, and Jennie Si. A long n-step surrogate stage reward for deep reinforcement learning. In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine, editors, Advances in Neural Information Processing Systems , volume 36, pages 12733-12745. Curran Associates, Inc., 2023. URL https://proceedings.neurips.cc/paper_files/ paper/2023/file/29ef811e72b2b97cf18dd5d866b0f472-Paper-Conference.pdf .

Szymon Łukaszyk. A new concept of probability metric and its applications in approximation of scattered data sets. Computational mechanics , 33:299-304, 2004.

A Appendix

A.1 Broader Impact

This paper is a research on improving the representation of deep reinforcement learning. We think that there are no potential societal consequences that need to be highlighted.

A.2 Additional Related Work

Recent studies have extensively explored representation learning in reinforcement learning (RL). Some previous works [18, 24, 43] employ autoencoders to encode image states into low-dimensional latent embeddings, resulting in improved visual perception and accelerated policy learning. Other approaches [42, 22, 23, 41, 35] utilize data augmentation techniques such as random crop or noise injection, combined with contrastive loss, to learn more generalizable state representations. Another set of approaches [25, 33, 15] focus on learning representation by predicting auxiliary tasks to extract more information from the environment. Successor representations [19, 30] tackles the state occupancy distribution to learn a more informative state representation, showing generalization capabilities between tasks.

In recent research, metric learning methods have been employed in RL to measure distances between state representations. Some approaches aim to approximate bisimulation metrics [47, 20] while others focus on learning sample-based distances [45, 7]. Bisimulation metrics have also been utilized in goal-based RL to enhance state representation [17]. A more recent work introduces quasimetrics learning as a novel RL objective for cost MDPs [40] but it is not for general MDPs.

Self-predictive representation (SPR) [32] learn representations regarding a future state but it focuses on predicting future states, enforcing the representation of predicted future states to be close to the true future states. SPR's representation focuses on dynamics learning, without directly considering the reward or value. In contrast, SCR focuses on predicting a metric/distance between current state and future state and the metrics are related to rewards or values regarding policy learning.

Successor Representations (SRs) [1, 3, 30] are a class of methods designed to learn representations that facilitate generalization. SRs achieve this by focusing on the occupancy of states, enabling generalization across various tasks and reward functions. In contrast, SCR measures the distance between states specifically to handle observable distractions.

A.3 Proofs

A.3.1 Proof of Theorem 3.1

Theorem 3. Let ˆ d : R n × R n → R be a metric in the latent state representation space, d ϕ ( x i , y i ′ ) := ˆ d ( ϕ ( x i ) , ϕ ( y i ′ )) be a metric in the state domain. The metric update operator F is defined as,

where ˆ M is the space of d , with a x i and a y i ′ being the actions at states x i and y i ′ , respectively, and ˆ P is the learned latent dynamics model. F has a fixed point d π ϕ .

Proof. We follow the proof techniques from [7] and [45]. By substituting ˆ d ( ϕ ( x i +1 ) , ϕ ( y i ′ +1 )) with d ϕ ( x i +1 , y i ′ +1 ) , we have

The operator F d ϕ is a contraction mapping with respect to the L ∞ norm because,

By Banach's fixed point theorem, operator F has a fixed point d π ϕ .

A.3.2 Proof of Theorem 3.4

To prove the distance ˆ d is a diffuse metric, we need to prove that ˆ d satisfies all of three axioms in the Definition 3.2 (definition of diffuse metric).

Lemma A.1. ∥ a ∥ 2 2 + ∥ b ∥ 2 2 -a ⊤ b ≥ 0 for any a , b ∈ R .

Proof. Because a ⊤ b ≤ ∥ a ∥∥ b ∥ , we have,

This lemma indicates that the term under the square root of ˆ d is always non-negative. ˆ d is able to measure any two vectors a , b ∈ R n .

Lemma A.2 (Non-negative) . ˆ d ( a , b ) ≥ 0 for any a , b ∈ R .

Proof. By definition, the square root is non-negative.

Lemma A.3 (Symmetric) . ˆ d ( a , b ) = ˆ d ( b , a )

Lemma A.4 (Triangle inequality) . ˆ d ( a , b ) + ˆ d ( b , c ) ≥ ˆ d ( a , c ) , for any a , b , c ∈ R .

Proof. To prove this lemma, it is equivalent to prove the following inequality by definition of ˆ d ,

Because -∥ x ∥∥ y ∥ ≤ x ⊤ y ≤ ∥ x ∥∥ y ∥ , ∀ x , y , we have

and

If

is true, then inequality (15) is true and Lemma A.4 is proven. To prove inequality (18), we can take squares on both sides without sign changing because both sides are non-negative. Then we have,

To prove inequality (18), it is equivalent to prove inequality (19). Expand and simplify inequality (19), we have

The left-hand side of inequality (20) is non-negative.

1) if right hand side -2 ∥ b ∥ 2 + ∥ a ∥∥ b ∥ + ∥ b ∥∥ c ∥ + ∥ a ∥∥ c ∥ < 0 , then inequality (20) is proven and backtrace to Lemma A.4 is proven.

2) if right hand side -2 ∥ b ∥ 2 + ∥ a ∥∥ b ∥ + ∥ b ∥∥ c ∥ + ∥ a ∥∥ c ∥ ≥ 0 , we take square on both sides of inequality (20) and have,

To prove inequality (21), we let left-hand side subtract right-hand side,

(22)

Therefore, inequality (21) is proven. Consequently, inequality (20) in the case of -2 ∥ b ∥ 2 + ∥ a ∥∥ b ∥ + ∥ b ∥∥ c ∥ + ∥ a ∥∥ c ∥ ≥ 0 is proven. Summarize with the case of -2 ∥ b ∥ 2 + ∥ a ∥∥ b ∥ + ∥ b ∥∥ c ∥ + ∥ a ∥∥ c ∥ < 0 , inequality (20) is proven and Lemma A.4 is proven.

Theorem A.5. ˆ d is a diffuse metric.

Proof. By summarizing Lemma A.2, A.3, and A.4, function ˆ d holds the three axioms of diffuse metric. Therefore, function ˆ d is a diffuse metric.

A.3.3 Proof of Theorem 3.6

Theorem 5. Let M ψ be the space of d ψ . The metric update operator F Chrono : M ψ → M ψ is defined as,

Proof. In practice, we have assumed that the sampled x j and y j ′ are future states of x i and y i ′ , respectively. To be precise, we extend the definition of F Chrono to cover the cases that i > j or i ′ > j ′ .

We follow the proof in Section A.3.1. F Chrono is contraction mapping because.

By Banach's fixed point theorem, operator F Chrono has a fixed point d π ψ .

A.4 Theoretical Details

A.4.1 Łukaszyk-Karmowski Distance

Definition A.6 (Łukaszyk-Karmowski Distance) . Given a metric space ( X,d ) where d is a metric defined on X × X , the Łukaszyk-Karmowski distance [49] is a distance between two probability measures µ and ν on X using metric d , and is defined as follows:

Intuitively, d LK measures the expected distance between two samples drawn from µ and ν , respectively. If we measure the d LK ( d ) distance between two identical distributions, the distance is not restricted to zero, i.e., d LK ( d )( µ, µ ) ≥ 0 .

A.5 Implementation Details

A.5.1 Implementation of ˆ m

We adopt IQE [38] to implement ˆ m . Given two vectors a , b ∈ R n , reshaping to R k × l where k × l = n , IQE first computes the union of the interval for each component:

where [ · , · ] is the interval on the real line. Then it computes the distance among all components d i as

where α ∈ R is an adaptive weight to balance the 'max' and 'mean' terms. In the scope of our method, we adopt d IQE ( a , b ) to implement ˆ m .

A.5.2 Network Architecture

The encoder ϕ takes an input of the states and consists of 4 convolutional layers followed by 1 fullyconnected layer. The output dimension of ϕ is 256. The encoder ψ takes an input of 512-dimensional vector (concatenated with ϕ ( x i ) and ϕ ( x j ) ), feed it into two layer MLPs with 512 hidden units, and output a 256 dim embedding. Q network and policy network are 3-layer MLPs with 1024 hidden units.

A.5.3 Hyperparameters

A.5.4 Algorithm

Algorithm 1 A learning step in jointly learning SCR and SAC.

A.6 Computational Resources 1

The experiment is done on servers with an 128-cores CPU, 512 GB RAM and NVIDIA A100 GPUs. 2

Two instances are running in one GPU at the same time. 3

B Experiments 4

B.1 Additional Information for Distraction Settings in DeepMind Control Suite 5

For the distraction setting in Section 4.2, we utilized the Distracting Control Suite [34] with the 6 setting "difficulty=easy". This involves mixing the background with videos from the DA VIS2017 [29] 7 dataset. Specifically, the training environment samples videos from the DA VIS2017 train set, while 8 the evaluation environment uses videos from the validation set. Each episode reset triggers the 9 sampling of a new video. Additionally, it introduces variability in each episode by applying a 10 uniformly sampled RGB color shift to the robot's body color and randomly selecting the camera 11 pose. The specifics of the RGB color shift range and camera pose variations are in line with the 12 Distracting Control Suite paper [34]. Different random seeds are used for the training and evaluating 13 environments at the start of training to ensure diverse environments. 14

B.2 Name of Tasks in DeepMind Control Suite 15

Due to limited space in main text, we use short name for tasks in DeepMind Control Suite. The full 16 name of tasks are listed in below Table 5:

17

B.3 Additional Results of Distracting Control Suite in Section 4.2 18

19

We extend the experiments in Section 4.2 with compared with RAP [8], showing results in Table 6.

Figure 9 shows the training curves of SCR and baseline methods in distraction setting of DM_Control. 20

21

ball_in_cup_catch

cartpole_swingup

cartpole_swingup_sparse

600

400

250

500

200

300

400

return

150

100

SCR

100

200

50

100

RAP

RAP

environment steps

2

le5

5

0

environment steps

2

3

1e5

0

environment steps

2

3

1e5

5

cheetah_run

finger_spin

hopper_stand

800

400

300

600

300

return

400

100

200

100

0

RAP

RAP

SimSR

0

1

environment steps

2

3

4

1e5

5

0

environment steps

2

3

1e5

0

1

environment steps

2

3

1e5

reacher_easy

walker_walk

700

600

600

500

500

400

#40

return

300

300

200

200

100

100

0

RAP

0

0

2

3

1

2

3

environment steps

1e5

environment steps

1e5

B.4 Additional Comparison Experiments with Data Augmentation Methods

Data augmentation is a category of methods efficient on reinforcement learning, toward sample 22 efficiency and noise invariance. Additional baseline CBM [26] is bisimulation metric method combing 23 with data augmentation method DrQ-v2, and it achieves significant performance in distracting 24 DM_Control. To compare with CBM, we embed data augmentation method DrQ into our proposed 25 method SCR. Table 7 indicates that the performance of SCR combined with DrQ surpass baseline 26 CBM. 27

B.5 Training Curves Meta-World 28

Figure 10 shows the training curves of SCR and baseline methods in Meta-World in Section 4.4. 29

basketball-v2

coffee-button-v2

door-open-v2

1.2

1.2

1.2

1.0

1.0

1.0

success rate

80

success rate

0.8

success rate

80

0.6

0.6

0.6

0.4

0.4

0.4

0.2

0.2

0.0

SAC

DrQ

DBC

MICo

SimSR

SCR

0.2

0.0

0.2

SAC

Dro

DBC

MICo

SimSR

SCR

0.2

0.2

0.0

MICo

SimSR

SCR

0

environment steps

2

3

4

1e5

5

0

1

environment steps

2

3

4

1e5

5

0

1

environment steps

2

3

1e5

5

drawer_open-v2

pick-place-v2

window-open-v2

1.2

1.2

1.2

1.0

1.0

1.0

success rate

6

0.8

successrate

0.6

0.8

success rate

0.8

0.4

0.2

0.0

SAC

Dro

DBC

MICo

SimSR

SCR

0.2

0.2

00

SimSR

SCR

0.2

0.0

0.2

SimSR

SCR

SAC

DrQ

DBC

MICo

0

1

environment steps

2

3

4

1e5

5

0.0

0.2

environment steps

0.4

0.6

80

1e6

1.0

0

1

environment steps

2

3

4

1e5

5

B.6 Additional Experiments on MiniGrid 30

31

We provide additional experiments on MiniGrid-FourRooms [9] environment. MiniGrid is under a

sparse reward setting that the agent receives a reward only when it successfully reaches a specific 32

goal, without rewards at previous steps. This experiment verifies the motivation of SCR that is able 33

to handle non-informative rewards. We compare various metric-based representation methods that 34

combine with backbone PPO [31] algorithms. We train each run for 5M environment steps. Table 35 shows the results with PPO, PPO+DBC, PPO+MICo, PPO+SimSR, and PPO+SCR.

36

PPO

PPO+DBC

0.5

PPO+MICo

PPO+SCR

0.4

F E°O

0.2

0.1

0.0

0

5

1e6

37

38

39

40

41

42

43

44

45

46

47

48

49

50

51

52

53

54

55

56

57

58

59

60

61

62

63

64

65

66

67

68

69

70

71

72

73

74

75

76

77

78

79

80

81

82

83

84

85

86

87

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: We have included the claims of contributions and scope.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: We discuss in the Conclusion Section.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: We provide complete proof in Appendix.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: We provide implementation details in Appendix.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

164

165

166

167

168

169

170

171

172

173

174

175

176

177

178

179

180

181

182

183

184

185

186

187

188

189

190

191

192

Answer: [No]

Justification: We provide pseudo codes in Appendix. We will release source code after published.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: We provide details in Appendix.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: We report the standard deviation of the experimental results.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: We provide details in Appendix.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: We follow NeurIPS Code of Ethics.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [Yes]

Justification: This paper is a research on improving the representation of deep reinforcement learning. We think that there are no potential societal consequences that need to be highlighted.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: Our paper poses no such risks.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: We have cited their original paper.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

296

297

298

299

300

301

302

303

304

305

306

307

308

309

310

311

312

313

314

315

316

317

318

319

320

321

322

323

324

325

326

327

328

329

330

331

332

333

334

335

336

337

338

339

340

341

342

343

344

345

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: We do not release new assets.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: This paper does not involve crowdsourcing nor research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: This paper does not involve crowdsourcing nor research with human subjects. Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes an algorithm to improve upon the recently proposed class of deep bisimulation-based methods by accounting for long-term consequences (in terms of future states) instead of relying one-step bootstrapping based on reward differences to be more robust in settings such as sparse reward cases.\n\nSTRENGTHS:\n1. I think the idea is valuable. The sparse reward setting is challenging for bisimulation-based methods and this type of idea of factoring future states is a reasonable attempt to make these algorithms better.\n2. The algorithm is evaluated on a variety of challenging pixel-based environments, which is good when evaluating representation learning algorithms.\n\nWEAKNESSES:\n1. It is a bit concerning that the main results are over only 5 seeds. Personally, it makes it hard for me to reach a conclusion on the performance because of this. I’d refer authors to [1].\n2. Related to above, the variance of the reported results are quite high. For example, in Table 1, 2 and Figure 3, the variance is so high that it is difficult to conclusively say which algorithm is better than another.\n3. Table 3 does not report the number of trials nor any measure of variance.\n4. While the approach proposed for bisimulation methods is indeed unique, I believe there are strong conceptual relations to self-predictive representation algorithms and successor representations [2, 3]. It would be good if these were discussed in the related works section and if similarities/differences can be drawn between SCR and these related works to better shape the work and help readers.\n5. It appears like Section 3.1 is basically prior work, but the paper includes it as though it is a contribution. I think it would be better to move this to the Appendix since the actual contributions are the sections after Section 3.1.\n\nIf the above issues are adequately addressed, I will consider increasing the score.\n\n[1] Empirical Design in Reinforcement Learning. Patterson et al. 2023.\n\n[2] Data-Efficient Reinforcement Learning with Self-Predictive Representations. Schwarzer et al. 2020.\n\n[3] Understanding Self-Predictive Learning for Reinforcement Learning. Tang et al. 2022\n\nQUESTIONS:\n1. A common issue in self-predictive representation algorithms (see Weaknesses above) is that the training procedure may lead to representation collapse [1]. Is SCR susceptible to this as well? It appears like it may be so. If not, how is that avoided here? \n2. I am curious how much of a burden it is to track future states in a given trajectory? A benefit of the other deep bismulation methods is that they can operate on a dataset of the form $\\{(s,a’,s’)\\}$, which is very useful in off-policy learning. SCR seems to require trajectory knowledge which may be a memory burden. In general, this is one of the advantages of one-step bootstrapping over n-step returns. If this is a limitation, it should be mentioned in the limitations section.\n3. Related to the above, how far into the future does a state $y$ have to be from previous state $x$ when computing the loss function. That is, using the terminology of n-step returns, how many steps in the future does $y$ have to be when computing these loss functions?\n\n\n[1] Understanding Self-Predictive Learning for Reinforcement Learning. Tang et al. 2022\n\nLIMITATIONS:\nYes, the paper addresses this in Section 5.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces SCR that extends state metric-based representations by embedding rich time-related information into the bisimulation metric learning process. SCR is designed to calculate state distances by contextually framing them in a temporal framework that considers both future dynamics and cumulative rewards in immediate and distant future scenarios. This technique fosters the development of long-term behavioral representations, capturing the progressively accumulated rewards from present to future states. Two distinct encoders are employed: one for generating representations of individual states and another dedicated to forming chronological embeddings that represent the interconnections between a state and its subsequent states. The framework also introduces a novel metric within the chronological embedding space, alongside an unconventional distance metric that diverges from the standard $L_p$ norm. Additionally, constraints are applied to manage the magnitude of the cumulative rewards. The methodology was tested within DeepMind Control and Meta-World settings, where SCR demonstrated state-of-the-art results in tasks requiring advanced generalization capabilities. Furthermore, an ablation study was conducted to assess the individual effects of the components comprising SCR.\n\nSTRENGTHS:\nThe paper addresses a noteworthy challenge in RL: capturing long-term temporal state representations within a well-defined metric space. To this end, SCR is proposed for learning representations centered on behavioral metrics that consider pairs of states over time. Moreover, it presents an effective technique for approximating the behavioral metrics. It develops a novel metric tailored for assessing pairs of states across temporal sequences and introduces new algorithms for this learning process. The introduced representation significantly outperforms others, particularly in the distraction setting and Meta-World.\n\nWEAKNESSES:\nIn Table 1, the outcomes for the default setting do not demonstrate a strong level of significance, with the possible exception of noteworthy results for Ch-Run. Sometimes, the performance appears to be worse than the compared methods. There is a section of the text discussing the distraction setting, and the results presented in Table 2 are quite impressive. These should be considered the primary findings of the study. Nonetheless, the paper does not sufficiently convey the rationale or motivation behind these findings.\n\nQUESTIONS:\n1) What factors contribute to SCR's occasional underperformance in the default setting?\n2) Is the distraction setting providing SCR with an advantage, and would the findings be consistent in an alternative setting demanding generalization abilities?\n\nLIMITATIONS:\nYes. The limitations and societal impact are adequately addressed.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper argues that metric learning for RL from one-step reward signal faces challenges in non-informative reward / sparse reward settings. The authors thus propose SCR, attempting to incorporate multi-step information in metric learning. The key components are as follows:\n1. The authors provide a MICo-like metric with a new distance function, which is also a diffuse metric and more numerically stable.\n2. Then they delineate a MICo-like chronological behavioral metric, measuring distance between two pairs of states. States in each pair are taken from the same trajectory.\n3. Then they propose a temporal measurement which can be approximated by upper-bounding and lower-bounding the measure, and thus shape the representation.\n\nThe paper experiments on the DMC environment with default setting and distraction setting, and multiple settings in Meta-World domain and show positive results.\n\nSTRENGTHS:\nThe issue of learning metric from sparse rewards is crucial and reported by previous work such as robust DBC. \n\nThe idea incorporating multi-step information in metric learning is interesting and novel.\n\nResults on DMC benchmarks are clearly better than chosen baselines. \n\nThe ablation study shows the effectiveness of the moving components of this approach.\n\nWEAKNESSES:\nMethodology:\n1. Eq.3 is not well-defined. What is the base case (e.g., what if $i=j$ but $i^\\prime\\neq j^\\prime$, thus $i+1>j$ and the recursive step is not well defined?)\n2. Eq.6 may not be true, as the optimal policy might not achieve the optimal expected discounted reward in any segment. In line 239, do you mean an optimal policy only within a given segment $(x_i,x_j)$? \n3.  It is unclear that Eq. 8 will co-hold with the classic metric bound that d upper bounds value difference, especially given that m measures local property while d measures global one. \nThe overall method is complicated which adds two more temporal embeddings based on state representations.\n\nExperiments: \n- Since the motivation of the proposed method is to handle non-informative reward scenarios, I think the experiment section should stress more importance on the sparse reward settings.  However, most experiments are done in environments with dense rewards except carepole_sparse and ball_in_cup. Neither do the authors give a comprehensive discussion on the challenges and results of those environments. \n- To align with the main claim, it should be expected to mainly work on Minigrid/Crafter-like sparse-reward environments. \n- For ablation study, the tasks of cheetah_run and walker_walk are dense rewards which cannot effectively support the claim on non-informative rewards, either. \n- Some other strong baselines are missing and worth comparing with, e.g., RAP, robust DBC (DBC-normed). RAP seems to outperform the SCR in some DMC tasks. \n- For experiments on Meta-World, please also report std in Table 3.\n\nQUESTIONS:\n- Eq. 2 and Eq. 5: is $\\phi$ in the target also optimized? \n- Line 83: Should Meta-World be better  to prove the generalizability of your proposed method, rather than distracting DMC?\n- Line 144: \"a_x and a_y are the actions in states x and y\" is not precise. Do the actions come from policy \\pi?\n- Line 267: it's better to succinctly talk about what IQE is and how you use it for completeness.\n- Appendix A.2: not finished writing\n- Can you talk more about the connection to GCB (Hansen-Estruch et al.)?\n- According to the SimSR and RAP paper, in some environments it requires more than 1e6 steps to converge. Is 5e5 steps enough for convergence?\n\nLIMITATIONS:\nIt is unclear that the proposed method can be seamlessly applied to intrinsically stochastic environments (i.e., the underlying abstract MDP has stochastic dynamics), where the temporal distance between two latent states may not be fixed.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper presents the State Chrono Representation (SCR), a novel approach to enhancing generalization in reinforcement learning (RL) with image-based inputs. SCR introduces a temporal perspective to bisimulation metric learning. The authors propose a learning framework that includes two encoders to capture individual state representations and the relationship between current and future states. The method is evaluated extensively in DeepMind Control and Meta-World environments, demonstrating SOTA performance in generalization tasks. The paper is well-structured and provides a solid theoretical foundation for the proposed approach.\n\nSTRENGTHS:\n1. The paper offers a robust theoretical foundation and a well-reasoned motivation for the SCR framework.\n2. The incorporation of future behavioral information into the representation space is a logical and effective strategy for improving generalization in RL.\n3. The experimental results are compelling, showing SCR's superiority over existing methods, particularly in challenging generalization tasks.\n\nWEAKNESSES:\n1. The discussion on limitations lacks depth. The paper could benefit from an insightful analysis of scenarios where SCR might underperform, such as the noted inferior performance on the R-Easy task compared to DrQ and DBC.\n2. While the paper mentions the commonality of future information prediction/regularization in RL representation learning, it does not provide a discussion on how SCR's approach differs from prior works like SPR or PlayVirtual and what new insights it offers. This would be valuable for the community.\n\nQUESTIONS:\n1. How does SCR handle data augmentation, and what impact does this have on its performance relative to other methods? As far as I know, data augmentation is very effective to generalization.\n2. Could the authors elaborate on the role of future information prediction/regularization in SCR and how it compares to existing techniques in the literature?\n\nLIMITATIONS:\n1. The paper could benefit from a more thorough exploration of its limitations and a discussion of situations where SCR might not be the best-performing approach.\n2. There is a need for a more comprehensive discussion on the relationship between SCR and established practices in RL representation learning, particularly concerning the use of future information for regularization.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-13"}