State Chrono Representation for Enhancing Generalization in Reinforcement Learning

Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-13 10:11:30

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.

Main paper: 
 introduction: 
 In deep reinforcement learning (Deep RL), deriving an optimal policy from highly dimensional environmental observations, particularly images, is a critical challenge (Castro, 2020; Gelada et al., 2019; Seo et al., 2022) . An RL agent continually receives images that display temporal relationships and substantial spatial redundancy. Redundant and potentially distracting visual inputs make it difficult for the agent to formulate optimal policies. Numerous studies have emphasized the importance of crafting state representations capable of discerning task-relevant information amidst taskirrelevant surroundings. Such representations hold the potential to greatly facilitate the RL process and enhance the generalizability of the learned policies. Consequently, representation learning has been recognized as a cornerstone in the advancement of Deep RL algorithms, garnering increased attention within the RL community (Kirk et al., 2023) .The primary focus of representation learning in reinforcement learning (RL) lies in developing a mapping function to transform high-dimensional observations into low-dimensional embeddings, which mitigates the influence of irrelevant signals to simplify the process of policy learning. Previous research in this area has employed autoencoder-like reconstruction losses (Yarats et al., 2021c; Higgins et al., 2017) , yielding impressive outcomes across various visual RL tasks. However, these approaches do not fully account for noise in visual features, which can be vital for accurately reconstructing input images. Data augmentation (Yarats et al., 2021b; Laskin et al., 2020a; b) methods have shown promise in tasks involving noisy observations, primarily enhancing perception models without directly impacting the policy within the context of the Markov Decision Process (MDP). The methods by learning auxiliary tasks (Seo et al., 2022) aim to predict additional tasks related to the environments using the learned representation as input. Nonetheless, these auxiliary tasks are often designed independently of the primary RL objective, potentially limiting their effectiveness.In recent advancements, behavioral metrics (Ferns et al., 2004; Ferns & Precup, 2014) , such as the bisimulation metric (Castro, 2020; Zhang et al., 2021) and MICo (Castro et al., 2021) , have emerged to quantify the dissimilarity between two states by considering differences in immediate reward signals and the divergence of next-state distributions. These metric learning methods establish approximate metrics within the representation space, preserving the behavioral similarities among states. 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 final RL goal, which involves maximizing the value function and shaping agent behaviors. Conversely, they filter out noise unrelated to state values and behavioral metrics. Behavioral metric approaches have demonstrated remarkable performance in various RL tasks, including control tasks from images, particularly in the presence of noisy images.However, behavioral metrics encounter challenges in handling demanding generalizable RL tasks and scenarios with sparse reward (Kemertas & Aumentado-Armstrong, 2021) . While behavioral metrics can somehow capture long-term behavioral metrics by temporal-difference update mechanism, their reliance on one-step transition data, where the information is limited in the case of sparse reward, significantly hampers learning efficiency. As a result, the representations learned using behavioral metrics may suffer from encoding with non-informative signals, such as sparse rewards (Kemertas & Aumentado-Armstrong, 2021) . Some model-based approaches attempt to mitigate these issues by learning transition models, but the task of learning a large transition model with long trajectories places increased demands on computational resources and parameters.To address the aforementioned challenges, we introduce the State Chrono Representation (SCR) framework, a metric-based approach to learning long-term behavioral representation and accumulated rewards spanning from present to future states. Within the SCR framework, we advocate for the training of two distinct state encoders. An encoder specializes in crafting a state representation for individual states, while the other focuses on generating a Chronological Embedding, which encapsulates the relationship between a state and one of its future states. In addition to learning the conventional behavioral metric for state representations, we introduce a novel behavioral metric tailored to temporal state pairs. This new metric is approximated within the chronological embedding space. We also propose an alternative distance metric, distinct from the typical L p norm, to efficiently approximate this behavioral metric in a lower-dimensional vector space. To infuse long-term rewards information into these representations, we present a "measurement" that quantifies the sum of rewards between the current and future states. Instead of directly regressing this measurement, we impose two constraints on it to restrict its range and value. Note that SCR is a versatile representation learning methodology that can be integrated into any existing RL algorithm.In summary, our contributions are threefold: 1) We introduce the SCR framework for representation learning with a focus on behavioral metrics involving temporal state pairs. Additionally, we provide a practical method for approximating these metrics.; 2) We develop a novel measurement specifically tailored for temporal state pairs and propose learning algorithms that incorporate this measurement while enforcing two constraints; 3) Our proposed representation demonstrates enhanced generalization and efficiency in challenging generalization tasks, as exemplified by experiments conducted on the Distracting DeepMind Control Suite (Stone et al., 2021; Tunyasuvunakool et al., 2020) . 
preliminary: 
 Markov Decision Process: A Markov Decision Process (MDP) is defined as a tuple M = (S, A, P, r, γ), where S represents the state space, consisting of all possible states, A indicates the action space, consisting of all possible actions which an agent can take in each state. The term P stands for the state transition probability function. Given a current state s t ∈ S and an action a t ∈ A taken, P (s t+1 |s t , a t ) gives the probability of transitioning to any state s t+1 ∈ S. r : S ×A → R denotes the reward function, which gives the immediate reward r(s t , a t ) received for taking an action a t in a state s t . The discount factor, γ ∈ [0, 1], determines the present value of future rewards.A policy π : S → A is a mapping function that determines the action that an agent will take in each state. The goal in MDP is to determine the optimal policy π * that maximizes the expected discounted cumulated reward, π * = arg max π E[ t γ t r(s t , π(s t ))].  
state chrono representation: 
 Despite their capabilities, the bisimulation metric (Zhang et al., 2021) and MICo (Castro et al., 2021) fall short when encoding future information. This limitation can impede the effectiveness of state representations in policy learning. To overcome this shortcoming and integrate future details, we present State Chrono Representation (SCR). Figure 1 shows the detailed architecture of SCR.SCR encompasses two representations: a state representation ϕ(x) ∈ R n for a state x and a chronological embedding ψ(x i , x j ) ∈ R n for a state x i and its future state x j . The state representation, ϕ(x), is developed through a behavioral metric d, which discerns the reward and dynamic divergence between two states. In contrast, the chronological embedding, ψ(x i , x j ), fuses these two state representations using deep learning, highlighting the long-term behavioral correlation between the current state x i and future state x j . A "chronological" behavioral metric is proposed to learn and compute the distance between any two chronological embeddings and is further refined through the Bellman operator-like MSE loss. Moreover, ϕ(x) employs an innovative temporal measurement, m, to assess the transition from the current state to a future one, effectively capturing sequential reward data. This devised temporal measurement operates within defined lower and upper constraints, directing the learning trajectory of both the measurement and state representation.The state representation encoder ϕ is trained by approximating a behavioral metric like the MICo distance. In our model, we adopt a MICo-based metric transformation operator, swapping the samplingbased prediction in MICo with latent dynamics-based modeling to determine the divergence between two subsequent states distribution, drawing parallels with the methodology in SimSR (Zang et al., 2022) . The metric update operator for latent dynamics, denoted as F, is defined below.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,EQUATIONwhere M is the space of d, with a xi 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 π ϕ .To learn the approximation for d π ϕ in the representation space, the form of distance d for lowdimensional vectors must be specified. Castro et al. (2021) demonstrated that a behavioral metric with a sample-based next state distribution divergence is a diffuse metric due to the divergence of the next state distribution being the Łukaszyk-Karmowski distance.Definition 3.1 (diffuse metric (Castro et al., 2021)). A function d : X × X → R based on the set 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 offers an approximation of the behavioral metric through an angular distance: b) represents the angle between vectors a and b, and β is a hyperparameter pre-determined to be 0.1. This distance calculation features a non-zero selfdistance, rendering it compatible with expressing the Łukaszyk-Karmowski distance. However, the angle function θ(a, b) exclusively considers the angle between a and b, necessitating computations involving the cosine similarity and arccos function, which can lead to numerical discrepancies. DBC recommends employing the L 1 norm with zero self-distance, suitable exclusively for the Wasserstein distance. Meanwhile, SimSR utilizes the cosine distance, derived from the cosine similarity, albeit without fulfilling the triangle inequality and the non-zero self-distance.dMICo (a, b) = ∥a∥ 2 2 +∥b∥ 2 2 2 + βθ(a, b), where a, b ∈ R n , θ(a,To mitigate the aforementioned challenges, we propose a revised distance, d(a, b), in the embedding space, characterized as a diffuse metric. This is mathematically formulated as:Definition 3.2. 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 4. d is a diffuse metric.Proof. Refer to Appendix for proofs of Theorem 4.Lemma 1 (Non-zero self-distance). The self-distance of d is not stricter to zero, i.e., d(a, a) = ∥a∥ 2 ≥ 0. This becomes zero if and only if every element in vector a is zero.Theorem 4 validates that d is a diffuse metric satisfying triangle inequality, and Lemma 1 shows d has non-zero self-distance capable of approximating dynamic divergence that is a Łukaszyk-Karmowski distance. Moreover, The structure of d resembles the L 2 norm with the exception that the weight before a ⊤ b is -1 instead of -2. Its construction, which only includes vector inner product and square root computations exempting divisions and trigonometric functions, prevents numerical computational issues and simplifies implementation.To learn the representation function ϕ, a prevalent approach is minimizing the MSE loss between both ends of Equation 13. The loss, combined with d and relative to ϕ, is expressed as:L ϕ (ϕ) = E xi,y i ′ ,rx i ,ry i ′ ∼D ϕ(xi+1),ϕ(y i ′ +1 )∼ P d(ϕ(x i ), ϕ(y i ′ )) -|r xi -r y i ′ | -γ d(ϕ(x i+1 ), ϕ(y i ′ +1 )) 2 , (2)where D represents the replay buffer or the sampled rollout for the RL learning process.The loss L ϕ (ϕ) in Equation 2 leans heavily on temporal-difference update mechanism with one-step transitions information. Consequently, it lacks the capacity and efficiency to grasp and encode rich long-term information from trajectories. Encoding temporal information within the representation, while ensuring it remains structured and adheres to behavioral metrics, poses an intricate challenge. To surmount this challenge, we propose two distinct methods, Chronological Embedding (discussed in Section 3.2) and Temporal Measurement (discussed in Section 3.3). Each technique is tailored to harness the temporal essence of a rollout, denoted as τ (x i ; x j ), which represents a sequence that originates from state x i and reach its future state x j . As illustrated in Figure 2 , the chronological embedding seeks to craft an innovative paired state embedding, symbolized as ψ(x i , x j ). This is achieved by transforming state representations ϕ(x i ) and ϕ(x j ) with deep network. The objective for learning ψ is to learn a novel "chronological" behavioral metric, one that measures the distance between rollouts τ (x i ; x j ) and τ (y i ′ ; y j ′ ). On the other hand, temporal measurement aspires to compute a special "distance" between states x i and x j . This measurement offers insights into the cumulative rewards amassed throughout the rollout τ (x i ; x j ). However, the learning temporal distance is formidable, and neither single method can claim full mastery over it. Consequently, we synergize both techniques to fortify and elevate the quality of the state representation.The chronological embedding, denoted as ψ(x i , x j ) ∈ R n , is tailored to capture the relationship between a given state x i and its long-term future states x j . It is premised on the assumption that both states, x i and x j , originate from the same trajectory. With a focus on capturing extended behavioral knowledge, we introduce a distance function d ψ : (S × S) × (S × S) → R, which is envisaged to mirror the behavioral metric, and allows the encoder ψ to integrate the behavioral information.Building upon the MICo distance in Theorem 2, we specify the metric update operator F Chrono for d ψ . Theorem 5. Let M ψ be the space of d ψ . The metric update operator F Chrono : M ψ → M ψ is defined as,F Chrono d ψ (x i , x j , y i ′ , y j ′ ) = |r xi -r y i ′ | + γE xi+1∼P π x ,y i ′ +1 ∼P π y d ψ (x i+1 , x j , y i ′ +1 , y j ′ ). (3) F Chrono has a fixed point d π ψ .Here d π ψ denotes the "chronological" behavioral metric. Our goal is to closely approximate d π ψ . It quantifies the distance between two sets of states, (x i , x j ) and (y i ′ , y j ′ ), considering the immediate reward difference and dynamics divergence. In our strategy to co-learn the encoder ψ withd π ψ , we represent d π ψ in terms of 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.2. Similar to Equation 13, we construct d π ψ to compute the distance in the embedding space.d π ψ (x i , x j , y i ′ , y j ′ ) = |r xi -r y i ′ | + γE xi+1∼P π x ,y i ′ +1 ∼P π y d(ψ(x i+1 , x j ), ψ(y i ′ +1 , y j ′ )). (4)For computational efficiency, the parameters between encoders ϕ and ψ are shared. The encoder ψ extracts outputs from ϕ, and the distance measure is suitably adapted as d π ψ (x i , x j , y i ′ , y j ′ ) := d(ψ((ϕ(x i ), ϕ(x j ))), ψ(ϕ(y i ′ ), ϕ(y j ′ ))). The objective for learning the chronological embedding is formulated as minimizing the MSE loss between both sides of Equation 4 w.r.t ϕ and ψ,L ψ (ψ, ϕ) = E xi,xj ,y i ′ ,y j ′ ,rx i , ry i ′ ,xi+1,y i ′ +1 ∼D d(ψ((ϕ(x i ), ϕ(x j ))), ψ(ϕ(y i ′ ), ϕ(y j ′ ))) -|r xi -r y i ′ | -γ d(ψ((ϕ(x i+1 ), ϕ(x j ))), ψ(ϕ(y i ′ +1 ), ϕ(y j ′ ))) 2 .(5)The goal of this objective is to nudge embeddings of analogous state sequences closer in the embedded space, bolstering the categorization of congruent behaviors.To facilitate SCR in acquiring future insights, we introduce temporal measurement, a conceptual "distance" to quantify the discrepancies between the current state x i and future state x j . This measurement, m(x i , x j ), aims to measure the differences in state value or the sum of rewards received within the current and future states. We build the approximated measurement upon the state representation ϕ, i.e., mϕ (x i , x j ) := m(ϕ(x i ), ϕ(x j )). where m : R n × R n → R is a non-parametric asymmetrical metric function further detailed in Section 3.3.1. This approach organizes the representation space of ϕ(x) around the "distance" m, enabling the structured representation of ϕ(x) to hold sufficient information to plan for future states.We propose m to represent the expected discounted accumulated rewards, obtained by an optimal policy π * from state x i to state x j :EQUATIONHowever, it is non-trivial to obtain m(x i , x j ) because the optimal policy π * is unknown and is, de facto, 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 way that ensures it is located in a feasible range covering the true m(x i , x j ). To construct this range, we introduce two constraints.The first constraint, which is considered as a lower boundary, asserts that the expected discounted cumulative reward collected by any policy π, optimal or otherwise, cannot surpass m:EQUATIONThis constraint is constructed based on an assumption that any sub-optimal policy is inferior to the optimal policy. Based on the constraint in Equation 7, we propose the first objective for learning the approximation mϕ :EQUATION)where ReLU (x) = x + = max(0, x). This objective becomes non-zero when the constraint in Equation 7 is not satisfied, pushing the value of m(ϕ(x i ), ϕ(x j )) to be larger until it becomes larger than the sampled reward sum.The second constraint, corresponding to upper boundary, is proposed according to inspiration from the triangle inequality. The absolute value |m(x i , x j )| is limited by the following inequality,EQUATION)where d is the behavioral metric introduced in Section 3.1. The right-hand side represents the longer path from x i to x j . This inequality demonstrates that the absolute temporal measurement |m(x i , x j )| is no greater than the sum of behavioral metrics at the beginning states (x i and y i ′ ), i.e. d(x i , y i ′ ), and end states pair (x j and y j ′ ), i.e. d(x j , y j ′ ), respectively, plus the measurement |m(y i , y j )|. This constraint leads to the following formulation of the 2nd objective for training mϕ :L up (ϕ) = ReLU | m(ϕ(x i ), ϕ(x j ))| -stop grad d((ϕ(x i ), ϕ(y i ′ ))) + d((ϕ(x j ), ϕ(y j ′ ))) + m(ϕ(y i ′ ), ϕ(y j ′ )) 2 (10)This objective pushes down the value of mϕ when the constraint in Equation 9 is not satisfied.By optimizing both constraints in a unified manner, we ensure the approximated temporal measurement, m, is anchored within a specific range, bounded by the lower and upper constraints. The overall objective for m is formulated as:EQUATIONThe measurement m(x i , x j ) designed to measure the distance regarding the rewards should be asymmetrical with respect to x i and x j . This is predicated on our assumption that state x i precedes x j , making its relation distinctly different from the progression from x j to x i . Recently research studies the quasimetric in deep learning (Pitis et al., 2020; Wang & Isola, 2022b ;a) and develop various methodologies to compute asymmetrical distances. In our method, we opt to leverage Interval Quasimetric Embedding (IQE) (Wang & Isola, 2022a ) to implement m.As shown in Figure 1 , the encoders are designed to predict state representations ϕ(x) for individual states and chrono embedding ψ(x i , x j ) for capturing the relationship existing between that states x i and x j . The measurement m is subsequently computed based on ϕ(x i ) and ϕ(x j ) to account for the accumulated rewards in between these states. The components ψ and m collaboratively enhance the state representations ϕ to capture the temporal information and its predictive capabilities for future insight. We will show the necessity of ψ and m in the ablation study which is detailed in Section 4.3. Therefore, a comprehensive objective is formulated in a unified manner:EQUATION)Our method, denoted as SCR, possesses the flexibility to be integrated with a broad spectrum of deep RL algorithms. These algorithms can effectively utilize the representation ϕ(x) as an integral input component. In our implementation, we employ Soft Actor-Critic (SAC) (Haarnoja et al., 2018) as our foundational RL algorithm. The state representation serves as the input state for the policy network and Q-value network in SAC. Other implementation details are referred to Appendix A.2. 
experiments: 
 Benchmarks. The primary objective of our proposed SCR is to cultivate a versatile and generalizable representation for deep RL when dealing with high-dimensional observation. To assess its efficiency, we conduct experiments utilizing the DeepMind Control Suite (DM Control) environment with rendered pixels observations (Tunyasuvunakool et al., 2020 ) and a distraction setting, Distracting Control Suite (Stone et al., 2021) . This environment utilizes the MuJoCo physics engine, providing pixel observations for a set of continuous control tasks. It includes rich testing scenarios, especially given its inclusion of background distractions and camera pose distractions, which simulate real-world complexities with camera inputs. Specifically,• Default setting. We evaluate our SCR on seven tasks in DM Control compared with other RL approaches. Each frame is rendered 3 × 84 × 84 pixels, as demonstrated in Figure 3 . We stack three frames as states and feed them into the RL agents. • Distraction setting. To evaluate to generalizability of SCR, we perform experiments on DDM Control with distraction. The distraction (as shown in Figure 4 ) includes 1) background video distraction, replacing the clean and simple background with a natural video; 2) object color distraction, slightly changing the color of the bodies of the robot; and 3) camera pose distraction, randomizing the camera pose of position and angle for rendering from the simulator. We observe that tasks become very hard if camera pose distraction is applied. Baselines. For a comprehensive understanding of our method's performance, we benchmark it against prominent algorithms in the domain, including: 1) SAC (Haarnoja et al., 2018) , a baseline deep RL method for continuous control; 2) DrQ (Yarats et al., 2021b) , a data augmentation method using random crop; 3) DBC (Zhang et al., 2021) , representation learning with the bisimulation metric; and 4) SimSR (Zang et al., 2022) , representation learning with the behavioral metric approximated by the cosine distance. In order to verify the sample efficiency of our method, we compare it with other methods on seven tasks in DM Control: ball in cup-catch, cartpole-swing up, cartpole-swing up sparse, cheetah-run, finger-spin, reacher-easy and walker-walk. We train each method in each task for 500K steps. Table 1 reports the scores evaluated at the end of training. All experimental results are averaged over 3 runs. We can observe that SCR has comparable results with the augmentation method DrQ and the state-of-the-art behavioral metric approach SimSR. Given that the maximum achievable returns for a DM Control task stands at 1000, a policy that collects scores around 900 is nearly optimal. These outcomes underscore the potency of SCR in mastering standard RL control tasks.To further evaluate the generalization ability of our method, we perform comparison experiments on DM Control with Distraction. We use the same training configuration as with the default setting.The camera-pose distraction presents a challenge for metric-centric methods like DBC and SimSR, primarily due to the significant distortion of the robot shape and position in the image state. Table 2 shows the results scores. Our method outperforms all other methods, including sparse reward task cartpole-swing up sparse, where other methods receive almost zero scores. DrQ outperforms other behavioral metric methods as its random cropping facilitates better alignment of robot position.Figure 5 : Ablation study on cheetah-run (left) and walkerwalk (right) in distracting setting. Curves are evaluation scores average on 3 runs and shadow shapes are std.To evaluate the impact of each component in the proposed SCR, we perform an ablation study where certain components were selectively removed or substituted. We provide another version of DM Control with distraction which includes only background distraction by replacing with greyscale videos. This experiment setup follows DBC and SimSR. Table 3 shows the experiment results. It is crucial to note that tasks limited to only background distractions are considerably simpler compared to those with object color and camera pose distractions.We perform these experiments to fairly compare with DBC and SimSR. The results show that SimSR performs quite well under this setting and our SCR achieves comparable results with SimSR. 
related work: 
 Recent studies have investigated representation learning in RL in many ways. Previous works (Higgins et al., 2017; Lee et al., 2020a; Yarats et al., 2021c) train autoencoder to encode image states into low-dimensional latent embeddings which improve the visual perception and accelerate policy learning. Approaches (Yarats et al., 2021b; Laskin et al., 2020a; b; Yarats et al., 2021a; Stooke et al., 2021) utilize data augmentation, e.g., random crop or noise injection, accompanied with contrastive loss to learn better generalizable state representations. Auxiliary tasks approaches (Lee et al., 2020b; Seo et al., 2022; Hafner et al., 2019) learn representation by predicting auxiliary tasks to extract more information from environments.Recent research with metric learning method for RL learns to measure distance on state representations. Some approaches learn to approximate bisimulation metric (Zhang et al., 2021; Kemertas & Aumentado-Armstrong, 2021) while other approaches learn sample-based distance (Zang et al., 2022; Castro et al., 2021) . Goal-based RL have harnessed bisimulation metrics for state representation (Hansen-Estruch et al., 2022) . Additionally, a recent work introduces quasimetrics learning as a fresh RL objective for cost MDPs (Wang et al., 2023) . 
conclusion: 
 The primary challenge of deep RL lies in cultivating an optimal policy from intricate, highdimensional noisy observations, such as images. In this work, we propose a novel metric-based representation framework State Chrono Representation (SCR) that leverages temporal dynamics in RL. SCR stands out by fusing the foundational principles of behavioral metrics with a holistic appreciation for long-term state dynamics. Our proposed method, while acknowledging past strides made using behavioral metrics, accentuates the need for a long-term vision in state representation, plugging gaps left by the one-step-focused models. Furthermore, by innovatively measuring rewards over temporal trajectories and weaving this measurement through representations, SCR pushes the frontier of efficient and adaptive RL representation. Our comprehensive experiments underscore the potency of this approach, especially in intricate environments laden with distractions. As the field advances, SCR serves as a testament to the blend of temporal relevance and behavioral metrics, hinting at the trajectory future representation learning endeavors might take.EQUATION)where M is the space of d, with a xi 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 (Castro et al., 2021) and (Zang et al., 2022) . By substituting d(ϕ(x i+1 ), ϕ(y i ′ +1 )) with d ϕ (x i+1 , y i ′ +1 ), we haveEQUATIONThe operator Fd ϕ is a contraction mapping with respect to the L ∞ norm because,EQUATIONBy Banach's fixed point theorem, operator F has a fixed point d π ϕ . 
 
 Appendix: To prove the distance d is a diffuse metric, we need to prove that d satisfies all of three axioms in the Definition 3.1 (definition of diffuse metric). Lemma 2. ∥a∥ 2 2 + ∥b∥ 2 2 -a ⊤ b ≥ 0 for any a, b ∈ R.Proof. Because a ⊤ b ≤ ∥a∥∥b∥, we have,EQUATIONThis 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 3 (Non-negative). d(a, b) ≥ 0 for any a, b ∈ R.Proof. By definition, the square root is non-negative. Proof. To prove this lemma, it is equivalent to prove the following inequality by definition of d,Lemma 4 (Symmetric). d(a, b) = d(b, a) Proof. d(a, b) = ∥a∥ 2 2 + ∥b∥ 2 2 -a ⊤ b = ∥b∥ 2 2 + ∥a∥ 2 2 -b ⊤ a = dEQUATIONBecause -∥x∥∥y∥ ≤ x ⊤ y ≤ ∥x∥∥y∥, ∀x, y , we haveEQUATIONandEQUATIONIf∥a∥ 2 + ∥b∥ 2 -∥a∥∥b∥ + ∥b∥ 2 + ∥c∥ 2 -∥b∥∥c∥ ≥ ∥a∥ 2 + ∥c∥ 2 + ∥a∥∥c∥ (20)is true, then inequality ( 17) is true and Lemma 5 is proven. To prove inequality (20), we can take squares on both sides without sign changing because both sides are non-negative. Then we have,EQUATION)To prove inequality (20), it is equivalent to prove inequality (21). Expand and simplify inequality (21), we haveEQUATIONThe left-hand side of inequality ( 22) is non-negative.1) if right hand side -2∥b∥ 2 + ∥a∥∥b∥ + ∥b∥∥c∥ + ∥a∥∥c∥ < 0, then inequality ( 22) is proven and backtrace to Lemma 5 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 ( 22) and have,4(∥a∥ 2 + ∥b∥ 2 -∥a∥∥b∥)(∥b∥ 2 + ∥c∥ 2 -∥b∥∥c∥) ≥ (-2∥b∥ 2 + ∥a∥∥b∥ + ∥b∥∥c∥ + ∥a∥∥c∥) 2 .(23) To prove inequality (23), we let left-hand side subtract right-hand side,EQUATIONTherefore, inequality ( 23) is proven. Consequently, inequality ( 22) 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 ( 22) is proven and Lemma 5 is proven.Theorem 6. d is a diffuse metric.Proof. By summarizing Lemma 3, 4, and 5, function d holds the three axioms of diffuse metric. Therefore, function d is a diffuse metric.A.1.3 PROOF OF THEOREM 5Theorem 5. Let M ψ be the space of d ψ . The metric update operator F Chrono : M ψ → M ψ is defined as,F Chrono d ψ (x i , x j , y i ′ , y j ′ ) = |r xi -r y i ′ | + γE xi+1∼P πx ,y i ′ +1 ∼P π y d ψ (x i+1 , x j , y i ′ +1 , y j ′ ). (25) F Chrono has a fixed point.Proof. We follow the proof in Section A.1.1. F Chrono is contraction mapping because EQUATIONEQUATIONwhere [•, •] is the interval on the real line. Then it computes the distance among all components d i asEQUATIONwhere α ∈ 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.The encoder ϕ takes input of the states and consists of 4 convolutional layers followed by 1 fullyconnected layer. The output dimension of ϕ is 256. The encoder ψ takes 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 EQUATION3: Compute loss L ϕ (ϕ) according Equation 24: Compute loss L ψ (ψ, ϕ) according Equation 55: Compute loss L m(ϕ) according Equation 116: Compute loss L(ϕ, ψ) = L ϕ (ϕ) + L ψ (ψ, ϕ) + L m(ϕ) according Equation 127: Update ϕ and ψ by minimizing loss L(ϕ, ψ) 8: Compute RL loss L RL according to SAC objectives 9: Update ϕ, Q and π by minimizing loss L RL 10: Soft update target Q network: For the distraction setting in Section 4.3, we follow the approach in DBC (Zhang et al., 2021) to setup experiments focusing solely on background video distraction. The background videos for this setting are sampled from the Kinetics (Kay et al., 2017) dataset (Kay et al., 2017) . We use 1000 consecutive frames for the training environment and a different set of 1000 consecutive frames for the evaluation environment, providing a varied visual context between training and evaluation phases.EQUATIONIn previous experiments, we set the number of steps between i and j to 50. To demonstrate the impact of the hyper-parameter, we include additional experiments with varying step counts: 1, 5, 10, 50, and 100 steps. We evaluate on cheetah-run and walker-walk tasks with the distracting setting in Section 4.2. The results are shown in Figure 6 . We observe that that 50 steps yield optimal results in these tasks. The outcomes of these evaluations are depicted in Figure 8 . Notably, our proposed SCR outperforms existing baseline methodologies across all evaluated tasks. While the DrQ algorithm demonstrates proficiency in achieving optimal performance levels, it is observed that SCR maintains superior sample efficiency, underscoring its effectiveness in the applied setting.  


Decision: Nearly all reviewers recommend acceptance of this paper. The sole exception is reviewer EHt5, who provided valuable feedback and raised valid concerns. The authors have addressed some of these concerns, but additional experiments are necessary to fully resolve these issues. However, I believe the current set of experiments is sufficient to demonstrate the capabilities and some limitations of the model.

That being said, I agree with reviewer Eht5 that the current set of experiments does not fully support the claim of "SOTA performance in demanding generalization tasks." This claim should be rephrased to be less assertive. However, I don't think this is enough grounds for rejecting the paper. The paper provides valuable theoretical and empirical contributions to the field.  

I encourage the authors to incorporate the results from the rebuttal, make the SOTA claim more moderate, and overall recommend this paper for acceptance.

Reviews: review: paper_summary: paper_summary: The 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.
strength_weakness: strength_weakness: 1. 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.
2. 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.
questions: questions: 1. 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.
2. Could the authors elaborate on the role of future information prediction/regularization in SCR and how it compares to existing techniques in the literature?
limitations: limitations: 1. 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.
2. 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.
paper_summary: paper_summary: This 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:
1. The authors provide a MICo-like metric with a new distance function, which is also a diffuse metric and more numerically stable.
2. 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.
3. Then they propose a temporal measurement which can be approximated by upper-bounding and lower-bounding the measure, and thus shape the representation.

The paper experiments on the DMC environment with default setting and distraction setting, and multiple settings in Meta-World domain and show positive results.
strength_weakness: strength_weakness: Methodology:
1. 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?)
2. 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)$? 
3.  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. 
The overall method is complicated which adds two more temporal embeddings based on state representations.

Experiments: 
- 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. 
- To align with the main claim, it should be expected to mainly work on Minigrid/Crafter-like sparse-reward environments. 
- 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. 
- 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. 
- For experiments on Meta-World, please also report std in Table 3.
questions: questions: - Eq. 2 and Eq. 5: is $\phi$ in the target also optimized? 
- Line 83: Should Meta-World be better  to prove the generalizability of your proposed method, rather than distracting DMC?
- 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?
- Line 267: it's better to succinctly talk about what IQE is and how you use it for completeness.
- Appendix A.2: not finished writing
- Can you talk more about the connection to GCB (Hansen-Estruch et al.)?
- According to the SimSR and RAP paper, in some environments it requires more than 1e6 steps to converge. Is 5e5 steps enough for convergence?
limitations: limitations: It 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.
paper_summary: paper_summary: This 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.
strength_weakness: strength_weakness: In 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.
questions: questions: 1) What factors contribute to SCR's occasional underperformance in the default setting?
2) Is the distraction setting providing SCR with an advantage, and would the findings be consistent in an alternative setting demanding generalization abilities?
limitations: limitations: Yes. The limitations and societal impact are adequately addressed.
paper_summary: paper_summary: The 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.
strength_weakness: strength_weakness: 1. 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].
2. 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.
3. Table 3 does not report the number of trials nor any measure of variance.
4. 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.
5. 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.

If the above issues are adequately addressed, I will consider increasing the score.

[1] Empirical Design in Reinforcement Learning. Patterson et al. 2023.

[2] Data-Efficient Reinforcement Learning with Self-Predictive Representations. Schwarzer et al. 2020.

[3] Understanding Self-Predictive Learning for Reinforcement Learning. Tang et al. 2022
questions: questions: 1. 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? 
2. 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.
3. 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?


[1] Understanding Self-Predictive Learning for Reinforcement Learning. Tang et al. 2022
limitations: limitations: Yes, the paper addresses this in Section 5.

Comments: title: Response to Reviewer EHt5 [2/2]
Response to Reviewer EHt5 [1/2]
[Gentle Reminder] Discussion Deadline is Approaching
Reply to the rebuttal
Thank you
Gentle Reminder to Review Our Rebuttal
Gentle Reminder to Review Our Rebuttal
Gentle Reminder to Review Our Rebuttal
Response to Reviewer UVpA (Continued)
Response to Reviewer UVpA
Rebuttal by Authors (Continued)
Rebuttal by Authors [3/3]
Rebuttal by Authors [2/3]
Rebuttal by Authors [1/3]
comment: We sincerely appreciate your effort in reviewing our work.
Thank you for your response. I don’t have any further questions at this time. Since the remaining concerns have not yet been addressed, though I believe the authors is on the right track to resolving them, I will maintain my current rating.
Thank you for your insightful comments. We address each point as follows:

>  theoretical convergence 

Thanks for your clarification. The convergency of SCR operator may hold true if $j$ is sufficient large. If the time steps between $i$ and $j$ are sufficient large, $\gamma$ will discount the $d$ of very future states to a negligible value and the convergency of $d(x_i,x_j, y_{i'}, y_{j'})$ may hold by the fixed-point theorem. We will provide the theoretical analysis under this condition in the revised version.

> Eq.8 looks weird

Your concern regarding Eq.8 is understandable, as it was intuitively proposed without direct theoretical support. We introduced Eq.8 based on perspective of implementation. The intuition behind Eq.8 is derived from Figure 2. Eq.8 serves as a cap to limit the increasing value of $\hat{m}$.

> criteria for selecting baselines


We would like to clarify the selection of baselines. SCR aims to set a new SOTA among metric-based representation methods. However, removing DrQ from baselines is not reasonable, because it is a strong baseline in DM_Control. We considered DrQ as an improved version of SAC, categorized as non-metric-based baselines. 

We agree with your suggestion and will split the experiments between metric-based methods and those incorporating data augmentation in our expanded experimental setup.


> results not convincingly demonstrate 

We will extend our experiments to integrate SCR with data augmentation across all tasks in revised version. 
Currently, the improvement seen in the cheetah-run and walker-walk tasks preliminarily indicate the potential effectiveness of DrQ+SCR.

> tuned hyperparameter settings of baseline methods

We acknowledge the importance of transparency in hyperparameter settings for baseline. We will provide the hyperparameters settings of all baselines in the appendix of the revised version.
Thank you for the clarification and for sharing the quick experiment results demonstrating the adaptivity of SCR with data augmentation.

I have some remaining comments. 

(1) The paper should clearly outline the criteria for selecting baselines. For instance, DrQ, which also uses data augmentation, is considered a proper baseline, while DrQ-v2 + CBM is deemed an unfair comparison in the rebuttal due to their use of data augmentation. This reasoning appears inconsistent. To address the effect of data augmentation fairly, the paper could consider including two sets of baselines and experiments: one without data augmentation and one with it. This approach would allow you to claim SOTA performance in either or both scenarios.

(2) The current results do not convincingly demonstrate that the method achieves SOTA performance across all tasks. While SOTA performance might be achieved on DCS tasks like cheetah-run and walker-walk, other tasks require additional empirical evidence to support similar claims (from my table, BiC-Catch and R-Easy have big gaps). 

(3) Since the authors have invested significant effort in tuning the baseline methods, it would be valuable to share the tuned hyperparameter settings in the appendix. This would benefit the research community by sparing other researchers the need to retune the baselines on the DCS benchmark.
Thank you for your quick and detailed response. I have a few points for clarification:

> In general, you can always add one to the index infinitely.

I believe there might be a misunderstanding regarding my concern. My focus is on the theoretical convergence of the operator, rather than the practical convergence in its implementation.

> This is a common problem across all bisimulation metric methods, including DBC, MICo, SimSR and RAP, where index addition is potentially infinite and it will exceed the feasible state space.

I don't agree with the statement that "index addition is potentially infinite and will exceed the feasible state space" in all bisimulation metric methods. These algorithms apply the operator similarly to the standard Bellman operator in an infinite-horizon MDP, where there is no terminal state. If you believe that these operators share a common issue, then the same concern would also apply to the standard Bellman operator, which is proven to converge by the fixed-point theorem.

The SCR operator is different from these operators. The monotonic increase in the index makes $i=j$ ($i'=j'$) possible for $d(x_i,x_j, y_{i'}, y_{j'})$, regardless of $i$ being infinite or not.  I believe further clarification is needed in the proof to address this distinction.

>  especially given that m measures local property while d measures global one

I get your clarification that $m$ is local optimal policy on the segment. My original concern is that a local optimal policy $m$ on a segment is not necessarily the global optimal policy in an MDP, as $m$ is myopic. The ideal condition of Eq. 8, which connects local and global properties, therefore, looks weird. 

Thank you for the results to show SCR is better than baselines in terms of learning curve.
> benchmark, baselines, and SOTA


We use the Distracting Control Suite (DCS) [Stone et al., 2021], as detailed in experiment configure section (***Section 4.1,lines 286-287***) of our paper. We acknowledge the differences between DCS and the DBC benchmark, which is the changing of body colors and camera poses. Such difference is the reason for the degrading performance of RAP in DCS. To ensure a fair comparison, we have  **tried our best to tune** the baseline methods to optimize their performance within the DCS environment. The enhanced level of distractions in DCS, particularly variations in **camera poses**, **decreases performance** more significantly compared to the DBC benchmark. We selected DCS specifically because of its increased challenges. 

We appreciate your suggestion to include additional baselines like CBM from [Liu et al., 2023]. However, CBM integrates bisimulation metrics with **data augmentation**, a technique known for its efficiency. This combination leading to a **unfair** direct comparisons with **solely metrics-based representation** method. Our focus in this paper is to improve metric-based representation learning **independently of other techniques**, particularly data augmentation, to clearly understand the contributions of the metric approach itself.
Another limitation with applying DrQ is its specificity to image observation spaces. For instance, in environments like Minigrid, where the observation space consists of a 7x7 matrix, DrQ’s random cropping technique would likely eliminate essential information. 
However, metrics-based representation methods do not depend on the observation space type. 

To ensure a fair comparison with CBM, we **incorporated DrQ into SCR** framework. Given the time constraints of less than 24 hours, we managed to complete training only for the cheetah-run and walker-walk tasks in the DCS. We will include full tasks set for DrQ+SCR and CBM in revised version. The results are as follows:

method | cheetah-run | walker-walk
---|---|---
SCR | 330.4  | 530.9
DrQv2+CBM | ~410 | ~640
DrQ+SCR | **463.7** $\pm$ 28.0 | **877.3** $\pm$ 40.2

These outcomes demonstrate that **DrQ+SCR outperforms DrQv2+CBM** in both tasks, supporting our claim of **achieving SOTA** performance. Additionally, this highlights the **adaptability of metric-based representation learning** to integrate effectively with data augmentation techniques.
Thank you for your feedback. We will address your concerns as detailed below:

> base case for Eq.3,4,5

In this paper, $i$ and $j$ are discrete time steps where $i < j$ implies $i + 1 \leq j$, with the maximum possible value for $i + 1$ being $j$. This ensures the feasibility of calculating the metric between $x_{i+1}$ and $x_j$, which are sampled from the replay buffer. A natural base case occurs when $i+1 = j$.

> In general, you can always add one to the index infinitely.

This is a **common problem** across all bisimulation metric methods, including DBC, MICo, SimSR and RAP, where index addition is potentially infinite and it will exceed the feasible state space. For example, consider a bisimulation metric (in simplified form):

$d(x_i, y_{i'}) = |r_i - r_{i'}| + \gamma d(x_{i+1}, y_{i'+1})$.

If $x_i$ is a **terminal state**, then $x_{i+1}$ becomes **infeasible**. In practice, it can be managed by ensuring states $x_{i+1}$ sampled  only from feasible states space, which is the replay buffer. 

The **base case** for bisimulation metric, $d(x_{i+1}, x_{i'+1})=0$, is **unlikely to be sampled**. if $x_{i+1}=x_{i'+1}$, they likely originate from the same transition, i.e. $(x_i, r_i, x_{i+1})=(x_{i'}, r_{i'}, x_{i'+1})$. In implementation, the transition pairs are uniformly sampled from buffer. The probability of sampling same transitions with a pair is approximately $1 / (replay\\_buffer\\_size)^2$, inversely proportional to the square of the replay buffer size. The probability decreases significantly as the replay buffer size increases.

While the absence of base cases can lead to overestimation of metrics, this effect **may not be harmful** to representation learning. Many research [1][2] in contrastive learning studies indicates that benefit of **pushing away** dissimilar samples often outweighs the need to align similar ones for **effectively learning representations**. The overestimation of metrics will rapidly distinguish state representations and RL get efficiency benefits from representation learning.

> For Eq.6

We wish to clarify the definition and role of $m$ in this context. We **define** $m$ as the optimal accumulated rewards between $x_i$ and $x_j$. Therefore, by definition, $m$ satisfies Eq.6 as the maximum potential accumulated rewards. In other words, $m$ is defined by upper bounding the LHS of Eq.6 that is the accumulated rewards of any policy. Policy that can achieve $m$ is regarded as optimal policy.

Rather than explicitly seeking this $m$, which is intractable, our approach focuses on optimizing an estimation $\hat{m}$. We optimize $\hat{m}$ to upper bond the LHS via Eq.7. Importantly, our objectives **do not require** the optimal policy and the extract $m$. 


> learning curves

We acknowledge your point regarding the importance of presenting learning curves in RL. We provide a detailed table showing the performance of PPO+SCR vs baselines at 500K steps intervals in Minigrid. Scores are computed using the latest 50K steps within each interval. The results highlight PPO+SCR’s sample efficiency compared to PPO.

Step   | 500K              | 1M                | 1.5M              | 2M                | 2.5M              | 3M                | 3.5M              | 4M                | 4.5M              | 5M                |
--- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
PPO     | 0.141 $\pm$ 0.004     | 0.224 $\pm$ 0.018     | 0.269 $\pm$ 0.026     | 0.324 $\pm$ 0.042     | 0.381 $\pm$ 0.058     | 0.431 $\pm$ 0.036     | 0.472 $\pm$ 0.037     | 0.494 $\pm$ 0.034     | 0.513 $\pm$ 0.033| 0.518 $\pm$ 0.031
PPO+SCR     | **0.146** $\pm$ 0.015     | **0.248** $\pm$ 0.014     | **0.322** $\pm$ 0.063     | **0.418** $\pm$ 0.034     | **0.467** $\pm$ 0.016     | **0.481** $\pm$ 0.023     | **0.499** $\pm$ 0.026     | **0.523** $\pm$ 0.017     | **0.539** $\pm$ 0.017| **0.551** $\pm$ 0.013
PPO+DBC      | 0.141 $\pm$ 0.014     | 0.249 $\pm$ 0.019     | 0.331 $\pm$ 0.061     | 0.399 $\pm$ 0.057     | 0.442 $\pm$ 0.039     | 0.470 $\pm$ 0.019     | 0.485 $\pm$ 0.018     | 0.499 $\pm$ 0.020     | 0.513 $\pm$ 0.012| 0.515 $\pm$ 0.017
PPO+MICo    | 0.090 $\pm$ 0.009     | 0.140 $\pm$ 0.056     | 0.194 $\pm$ 0.059     | 0.216 $\pm$ 0.067     | 0.235 $\pm$ 0.074     | 0.270 $\pm$ 0.068     | 0.309 $\pm$ 0.064     | 0.321 $\pm$ 0.076     | 0.339 $\pm$ 0.083| 0.356 $\pm$ 0.067
PPO+SimSR   | 0.084 $\pm$ 0.007     | 0.145 $\pm$ 0.027     | 0.168 $\pm$ 0.036     | 0.197 $\pm$ 0.037     | 0.223 $\pm$ 0.039     | 0.243 $\pm$ 0.033     | 0.273 $\pm$ 0.051     | 0.303 $\pm$ 0.058     | 0.316 $\pm$ 0.052| 0.327 $\pm$ 0.051

For the **Meta-World** experiments, the learning curves can be found in ***Figure 10 of Appendix B.4***.
Thank you for your response and I appreciate your running more seeds. I will stick to the updated score as I already made the update.
Thank you for the rebuttal. While it addresses some of my concerns, my primary issues remain unresolved.

For Eq.3,4,5: Even if you add the constraint, it is still necessary to define a base case, as the metric is recursively defined. For example, even if you have a constraint $i<j$, but $i+1<j$ may not hold, and the metric may still be ill-defined. In general, you can always add one to the index infinitely.

For Eq.6: it seems regardless of stochastic/deterministic policies (even if you assume the support of policy always covers all action and the support of transition function covers the state space), as you said, there exists a $\pi$ that $m_\pi> m$. When training, the agent’s policy can be any policy in the policy space. Thus, the LHS of Eq.6 can exceed m, even by a large amount in theory. Why do you think it has something to do with stochastic/deterministic policies and how the LHS can be a valid lower bound?

> generalization and sparse reward

Thanks for the clarification that the main focus is on generalization.

> Additional results in table

I appreciate the new results that demonstrate your approach outperforms the baselines in Minigrid. However, both this table and the Meta-World table lack the learning curves (OpenReview should allow you to upload PDF for the rebuttal), which are commonly presented in DRL research due to the known instability during training. The results between PPO and PPO+SCR in Minigrid are quite close, and the learning curves would help to differentiate the two methods more clearly.

> benchmark, baselines, and SOTA

In the RAP paper, the reported result in many distracting DMC tasks is much better than yours. Are you using different benchmarks? As far as I know, there are two related benchmarks -- the one DBC [Zhang et al., 2021] created (referred to as the DBC benchmark) and the distracting control suite (DCS) [Stone et al., 2021]. The baselines adopted in this work, including DBC, robust DBC, MICo, RAP, SimSR, were evaluated in the DBC benchmark in their respective papers. Since your paper cites DCS [Stone et al., 2017] for your benchmark, it seems you are using the DCS benchmark, which differs from those baselines. If this is the case, the baselines may be untuned for this new benchmark, which is not best practice (refer to the discussion on untuned baselines in new tasks in Section 4.1 of [Patterson et al., 2023, Empirical Design in Reinforcement Learning]). 

Moreover, if this is the case, the authors might consider including more appropriate baselines evaluated on DCS with strong performance, such as [Liu et al., 2023, Robust Representation Learning by Clustering with Bisimulation Metrics for Visual Reinforcement Learning with Distractions]. Based on Liu et al.'s paper (their Figure 3) and using the same sample budget (500k steps), I have summarized the comparison on 6 tasks in DCS (average only, as Liu et al. does not provide detailed table results):

|task  | SCR | DrQ-v2 + CBM|
| -------- | ------- |------- |
|BiC-Catch | 138.5 | ~780|
|C-Swingup | **566.1** |~630|
|Ch-Run | **330.4**|~410|
|F-Spin | **780.0**|~810|
|R-Easy| 158.7|~490|
|W-Walk| **530.9**| ~640|

The numbers in bold indicate where your paper claims that SCR achieves SOTA among the baselines. However, SCR performs worse than DrQ-v2 + CBM on all these tasks, which raises significant concerns about the validity of the SOTA claim. Quoting from your conclusion:

> Our extensive experiments demonstrate its effectiveness compared with several SOTA baselines in complex environments with distractions.



While the authors' response has addressed some concerns, I believe the paper still requires further refinement before it can be accepted. I recommend a thorough revision to ensure the technical correctness of the method and to better organize the experimental design, which will help more convincingly support the claims made.
Dear Reviewer UVpA,
 
Thank you for your thoughtful feedback and for updating your score.
 
We appreciate your agreement with our presentation of the experimental results using $N \times num\\_episodes$. To provide a more comprehensive comparison, we will include Tables A and B in the revised manuscript.
 
Regarding the number of seeds, we recognize the importance of increasing the robustness of our findings. We plan to run 5 **additional seeds** for the final version of the manuscript to further strengthen our results. Although the current results are based on 5 seeds, the data in Table 2, Table B, and Figure 4 show a **distinct separation** in the performance ranges. These results **convincingly** demonstrate the significant improvement of SCR over the baselines in the distraction setting.
 
Best Regards,
 
The Authors
Dear Reviewer EHt5,
 
I hope this message finds you well.
 
As the discussion deadline of **August 13 AoE** is approximately **30 hours away**, we noticed that we have not yet received your feedback on our rebuttal. All other reviewers have kindly provided their responses, and your insights are crucial for the final evaluation of our submission.
 
Could you please take a moment to review our rebuttal and share your feedback within this time frame? We greatly appreciate your expertise and are looking forward to your valuable comments.
 
Thank you very much for your attention to this matter.
 
Best regards,
 
The Authors
Dear Reviewer R7iQ,

Thank you for your valuable feedback. We are pleased that our responses addressed your concerns. 

Best regards,

The Authors
Dear Reviewer b4Fc,

Thank you for your thorough review and for adjusting the score based on our responses. We are pleased that we could successfully address your questions.

Best regards,

The Authors
Thanks for the response which has addressed many concerns. Good luck!
I appreciate the author for answering my questions. I raise my score to 6.
Dear Reviewer b4Fc,

We sincerely appreciate the time and effort you have invested in reviewing our manuscript. We have provided responses that address the issues raised in your initial feedback. We believe that further discussion would be highly valuable in improving the quality and impact of our work.

We would like to gently remind you of the upcoming deadline for the discussion period on **August 13 AoE**. Your expertise and insights are vital to helping us refine our manuscript.

Thank you for your continued support and consideration.

Best regards,

The Authors
Dear Reviewer EHt5,

We sincerely appreciate the time and effort you have invested in reviewing our manuscript. We have provided responses that address the issues raised in your initial feedback. We believe that further discussion would be highly valuable in improving the quality and impact of our work.

We would like to gently remind you of the upcoming deadline for the discussion period on **August 13 AoE**. Your expertise and insights are vital to helping us refine our manuscript.

Thank you for your continued support and consideration.

Best regards,

The Authors
Dear Reviewer R7iQ,

We sincerely appreciate the time and effort you have invested in reviewing our manuscript. We have provided responses that address the issues raised in your initial feedback. We believe that further discussion would be highly valuable in improving the quality and impact of our work.

We would like to gently remind you of the upcoming deadline for the discussion period on **August 13 AoE**. Your expertise and insights are vital to helping us refine our manuscript.

Thank you for your continued support and consideration.

Best regards,

The Authors
Thank you for updating with the results and showing the std of the runs. I think reporting over N * num_episodes seems reasonable. Even if there is no ideal metric for high variance estimates, discarding them is not correct in my opinion, so thank you for including them in your new result. I understand the computational cost issue. I will say, however, that citing prior work for having done 5 seeds is not necessarily a good argument to continue the trend though. I've updated my score.
Methods | BiC-Catch | C-SwingUp | C-SwingUpSparse | Ch-Run | F-Spin | H-Stand | R-Easy | W-Walk
---|---|---|---|---|---|---|---|---
SAC | 465.9 $\pm$ 12.7 | 730.2 $\pm$ 101.9 | 21.1 $\pm$ 10.6 | 357.5 $\pm$ 3.4 | 492.1 $\pm$ 35.6 | 753.8 $\pm$ 34.6 | 333.8 $\pm$ 19.9 | 398.1 $\pm$ 98.4
DrQ | 968.3 $\pm$ 4.8 | 834.5 $\pm$ 13.3 | 739.4 $\pm$ 26.8 | 477.8 $\pm$ 14.8 | 962.3 $\pm$ 4.1 | 856.4 $\pm$ 16.8 | 970.8 $\pm$ 10.8 | 924.1 $\pm$ 10.7
DBC | 80.6 $\pm$ 44.8 | 300.8 $\pm$ 58.8 | 122.7 $\pm$ 15.1 | 397.7 $\pm$ 32.3 | 633.5 $\pm$ 28.2 | 82.8 $\pm$ 26.7 | 199.1 $\pm$ 17.8 | 359.5 $\pm$ 38.1
MICo | 206.6 $\pm$ 34.3 | 811.7 $\pm$ 6.8 | 0.0 $\pm$ 0.0 | 3.6 $\pm$ 0.1 | 2.0 $\pm$ 0.2 | 815.7 $\pm$ 57.3 | 199.8 $\pm$ 18.5 | 27.3 $\pm$ 2.5
SimSR | 949.3 $\pm$ 16.7 | 862.9 $\pm$ 4.8 | 435.3 $\pm$ 35.3 | 486.8 $\pm$ 25.3 | 975.8 $\pm$ 2.8 | 5.0 $\pm$ 1.2 | 94.3 $\pm$ 7.2 | 929.1 $\pm$ 5.9
SCR | 951.0 $\pm$ 6.7| 847.9 $\pm$ 12.5 | 799.3 $\pm$ 22.4 | 799.5 $\pm$ 26.4 | 983.0 $\pm$ 3.1 | 830.5 $\pm$ 14.5 | 691.8 $\pm$ 29.4 | 923.2 $\pm$ 5.6

**Table A**: Mean and std over 5 averaged scores on DM_Control with default setting.

Methods | BiC-Catch | C-SwingUp | C-SwingUpSparse | Ch-Run | F-Spin | H-Stand | R-Easy | W-Walk
---|---|---|---|---|---|---|---|---
SAC | 60.0 $\pm$ 31.7 | 223.3 $\pm$ 6.2 | 1.4 $\pm$ 0.8 | 168.7 $\pm$ 9.9 | 53.4 $\pm$ 10.7 | 4.5 $\pm$ 1.2 | 102.1 $\pm$ 22.0 | 175.7 $\pm$ 6.8
DrQ | 194.4 $\pm$ 76.8 | 239.4 $\pm$ 37.6 | 4.0 $\pm$ 4.4 | 104.6 $\pm$ 55.5 | 514.5 $\pm$ 99.5 | 22.8 $\pm$ 3.1 | 102.8 $\pm$ 59.7 | 31.1 $\pm$ 1.7
DBC | 32.0 $\pm$ 8.9 | 125.5 $\pm$ 9.3 | 0.0 $\pm$ 0.0 | 7.0 $\pm$ 1.1 | 3.3 $\pm$ 2.7 | 5.8 $\pm$ 0.5 | 191.7 $\pm$ 21.1 | 26.3 $\pm$ 1.1
MICo | 115.7 $\pm$ 25.7 | 206.6 $\pm$ 23.6 | 0.0 $\pm$ 0.1 | 3.8 $\pm$ 0.1 | 113.4 $\pm$ 20.4 | 5.5 $\pm$ 1.0 | 162.4 $\pm$ 54.4 | 20.2 $\pm$ 1.0
SimSR | 121.9 $\pm$ 61.3 | 134.4 $\pm$ 35.7 | 0.0 $\pm$ 0.0 | 7.3 $\pm$ 3.3 | 0.3 $\pm$ 0.4 | 5.5 $\pm$ 1.5 | 97.8 $\pm$ 23.0 | 26.0 $\pm$ 2.6
SCR | 138.5 $\pm$ 34.8 | 566.1 $\pm$ 37.7 | 166.7 $\pm$ 62.0 | 330.4 $\pm$ 23.5 | 780.0 $\pm$ 39.4 | 88.5 $\pm$ 12.9 | 158.7 $\pm$ 40.4 | 530.9 $\pm$ 23.1

**Table B**: Mean and std over 5 averaged scores on DM_Control with distraction setting.
> Re: seeds. My understanding of how this typically works is: 1 algorithm on 1 environment with 1 configuration of h-params is run for $N$ seeds. These $N$ results are averaged together. In my opinion, the real robustness improves by not increasing the number of episodes but by increasing the number of seeds. 

We agree with you a comprehensive comparison is needed for readers to better understand our SCR capabilities. In this comment, we show the results with **std calculated over 5 averaged scores** in the following tables, ***Table A and B***. The std is much lower than the tables in the manuscript.

We did not show the std calculated over $N$ averaged scores in the manuscript, because the averaged scores **decline the effect of the extreme values in final std** and conceal the high variance issue of the challenging tasks. Instead, we show the std calculated over all $N \times num\\_episodes$ scores because it can reflect the **extreme values**.

**A large number of evaluation episodes is necessary** to improve the confidence level of the average scores, due to the high variance across various episodes. The high variance is from the large randomness of the Distracting DM_Control environment. It randomly selects a background video, a body color, and a camera pose at the beginning of each episode. Different episodes will have different distraction effects, which consequently affect the agent's performance.


> I see computational constraint as an issue, but I don't think doing 5 to 10 more seeds is that expensive.

We conduct experiments by running each instance on 0.5 GPU for 12 hours. 0.5 GPU means we run 2 instances in parallel on 1 GPU. For the DM_Control experiment, we have 6 algorithms, 8 tasks, and 2 configurations (default and distraction) for 5 seeds. In total, we need **6 * 8 * 2 * 5 * 0.5 * 12 = 2880 GPU hours**. We understand that more seeds would lead to more robust results. However, due to computational constraints, we don't believe we can train 5 to 10 more seeds during the rebuttal period. 

5 seeds is a common practice in baseline papers, such as SimSR and MICo. We follow the same practice to ensure a fair comparison with the baselines. Increasing the number of seeds can improve confidence, which is beneficial for this paper, and we **will increase the number of seeds in the revised version.** 

The results in ***Table 2 and Figure 4*** in the manuscript and ***Table B*** in this comment show that the performance of SCR is **significantly better** than the baselines in the distraction setting. We believe these result sufficiently prove the effectiveness of SCR.


> Do the aggregate metrics report an IQM? If so, I don't think that is the best reporting especially given that your response mentions that the distraction setting results in high variance. If the setup, by nature, is high variance, then I don't think we want to exclude the extreme values if the goal is to say the algorithm performs well in high-variance settings.

The aggregate metrics in **Figure 4** include median, IQM, mean, and optimality gap. The mean and optimality gap metrics do not exclude the extreme values. 
Your concern regarding the metrics is reasonable. However, the community has **no perfect metric** to measure performance in high-variance environments currently. To mitigate the potential bias, we show the std over all evaluation episodes which can **include the extreme values**.
Thanks to the authors for their response. Two questions:
1. Re: seeds. My understanding of how this typically works is: 1 algorithm on 1 environment with 1 configuration of h-params is run for $N$ seeds. These $N$ results are averaged together. In my opinion, the real robustness improves by not increasing the number of episodes but by increasing the number of seeds. I see computational constraint as an issue, but I don't think doing 5 to 10 more seeds is that expensive.

2. Do the aggregate metrics report an IQM? If so, I don't think that is the best reporting especially given that your response mentions that the distraction setting results in high variance. If the setup, by nature, is high variance, then I don't think we want to exclude the extreme values if the goal is to say the algorithm performs well in high-variance settings.
>3. 

The parameter $n$ is uniformly sampled within the range [1, 100], as described in **lines 338-344** of the manuscript. Our ablation study, detailed in **Section 4.3 / Figure 6**, examines the impact of different $n$. The results indicate that sampling $n$ within the [1, 100] range provides optimal performance and maintains stability across various tasks.

----
[1] Towards Robust Bisimulation Metric Learning. Kemertas et al. NeurIPS 2021.
> Appendix A.2: not finished writing
Apologize for accidentally delete some content in additional related works. Here, we restore the omitted content:

A more recent work introduces quasimetrics learning as a novel RL objective for cost MDPs [1] but it is not for general MDPs.


> Can you talk more about the connection to GCB (Hansen-Estruch et al.)?

GCB utilizes a bisimulation metric to address goal-conditioned RL problems, learning the $L_1$ distance between representations of a current state and a fixed goal state within an episode. In contrast, SCR learns a metric between a current state and a future state that varies, as future states are sampled throughout the episode. This fundamental difference stems from SCR's aim to enhance generalization across general MDPs, while GCB focuses on goal-reaching within goal-conditioned RL.

> According to the SimSR and RAP paper, in some environments it requires more than 1e6 steps to converge. Is 5e5 steps enough for convergence?

We provide learning curves in Appendix Figure 9, which demonstrate that SCR can effectively learn within a 500K step budget. This budget aligns with the standard training parameters set forth in the Distracting Control Suite[2]

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

SCR is designed to be **agnostic** to the deterministic or stochastic nature of the environment dynamics, much like existing methods such as DBC and SimSR. The metric in SCR, akin to a bisimulation metric, is defined as the expectation over next state distributions, making it applicable to stochastic environments. The temporal metric in SCR (discussed in Section  3.3) is not an exact measure but rather is learned to remain within a feasible range, thus accommodating the intrinsic stochasticity of the environment.

-----
[1] Optimal goal-reaching reinforcement learning via quasimetric learning. Wang et al. ICML 2023

[2]The Distracting Control Suite--A Challenging Benchmark for Reinforcement Learning from Pixels. Stone et al. 2021
> To align with the main claim, it should be expected to mainly work on Minigrid/Crafter-like sparse-reward environments.

To validate SCR's effectiveness in **sparse reward environments**, we conducted additional experiments on Minigrid-FourRooms, adapting our base RL algorithm to PPO due to its discrete action space. We compared PPO+SCR with PPO, PPO+DBC, PPO+MICo, and PPO+SimSR over 5M environment steps training, evaluating performance within a 500K environment steps budget. The results, averaged over five seeds, demonstrate SCR's superior performance in these settings:

Method| PPO                | PPO+DBC            | PPO+MICo           | PPO+SimSR          | PPO+SCR            |
------| ------------------ | ------------------ | ------------------ | ------------------ | ------------------ | 
Score| 0.515 $\pm$ 0.029  | 0.515 $\pm$ 0.012  | 0.350 $\pm$ 0.064  | 0.321 $\pm$ 0.046  | 0.546 $\pm$ 0.012  |

These outcomes underscore PPO+SCR's robustness and superior effectiveness in sparse reward environments compared to other methods.

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

As previously noted, the primary motivation behind SCR is to **enhance generalization** capabilities. The ablation study specifically assesses the impact of each SCR component, conducted in a **distraction** setting to underscore their contributions to generalization performance.


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

We appreciate your suggestion and include a comparative analysis below. The table demonstrates SCR's performance relative to RAP and robust DBC across three tasks in distraction setting:

|   | C-SwingUpSparse | Ch-Run | W-Walk
| -| - | - | - |
| SCR| 166.7 $\pm$ 178.3 | 330.4 $\pm$ 152.6 | 530.9 $\pm$ 172.5 |
| RAP | 0.1 $\pm$ 0.6 | 252.4 $\pm$ 113.5 | 205.0 $\pm$ 90.5 |
| robust DBC | 0.1 $\pm$ 0.3 | 20.3 $\pm$ 10.4 | 36.5 $\pm$ 10.1 |

As shown, SCR outperforms RAP and robust DBC in these tasks. 

> For experiments on Meta-World, please also report std in Table 3.

The following table provides the scores along with their standard deviations for the Meta-World experiments:

|SAC |DrQ| DBC| MICo| SimSR |SCR|
|-|-|-|-|-|-|
| 0.495 $\pm$ 0.475| 0.886 $\pm$ 0.125| 0.479 $\pm$ 0.453| 0.495 $\pm$ 0.482|0.258 $\pm$ 0.365| 0.969 $\pm$ 0.03|


> Eq.2 and Eq.5 : is $\phi$ in the target also optimized?

Thank you for highlighting this issue. In our approach, we do not optimize $\phi$ in the target but stop gradient for it. We will make sure to clarify this in the revised version.


> Line 83: Should Meta-World be better to prove the generalizability of your proposed method, rather than distracting DMC?

Meta-World is designed primarily as a benchmark for meta-learning and multi-task learning, whereas the Distracting Control Suite is specifically for the generalization capabilities of RL algorithms. We focus on the distraction setting of DM_Control for our main results because it aligns more closely with common practices in this area of research.

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

Yes, the actions indeed sample from policy $\pi$. Although we omitted mention of $\pi$ in Eq.1, it is implicit. Specifically, $\phi(x_{i+1})$ samples from a marginal distribution shaped by $\sum_{a_{x_i}} \hat{P}(\cdot|\phi(x_i), a_{x_i}) \pi(a_{x_i}|x_i)$. We will make this explicit in the revised manuscript.

> Line 267: it's better to succinctly talk about what IQE is and how you use it for completeness.


IQE is discussed in detail in Appendix A.5.1. We will include a direct reference to this appendix in line 267 for better clarity.
IQE is an asymmetric metric function that calculates the distance between vectors. It's differentiable, allowing gradients from objective Eq.9 to pass through, facilitating updates to the encoder.
Thank you for your insightful comments. We address your concerns as follows:
> 1. Eq.3 is not well-defined. What is the base case (e.g., what if $i=j$ but $i' \neq j'$, thus $i+1>j$ and the recursive step is not well defined?)

Thank you for pointing out this issue. We clarify that $i \neq j$ and $i' \neq j'$, where $j$ represents a later step than $i$, as demonstrated in Section 3.1.1. We will add constraints explicitly stating $i < j$ and $i' < j'$ in Section 3.2 in the revised manuscript.

> 2. 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)$ ?

Eq.6 holds under our assumption that all policies are stochastic. With a value function $V^*$ from an optimal policy $\pi^*$, the n-step Bellman operator at state $x_i$ is defined as $V^*(x_i) = m(x_i, x_j) + \gamma^{j-i}V^*(x_j)$. If there exists another policy $\pi$ where $m_\pi(x_i, x_j) > m(x_i, x_j)$, then $V^*(x_i)$ would be suboptimal instead of optimal, suggesting a need for further updates. However, for deterministic policies, Eq.6 does not hold as $\pi^*$ may not ensure reaching $x_j$ from $x_i$. Thus, we adopt a stochastic policy framework in this study, implemented with a Gaussian policy.


> 3. 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. The overall method is complicated which adds two more temporal embeddings based on state representations.

Proving Eq.8 is challenging due to its general nature. We intuitively propose Eq.8 to confine the value learning of $\hat{m}$. 

In practice, $m$ is intractable and cannot be approximated by regression or recursion. Instead, we constrain $\hat{m}$ to remain within practical bounds (refer to **line 242-243**). Given that Eq. 6 establishes a lower bound for $m$ and our objective in Eq. 7 tends to increase $\hat{m}$'s value. Without proper constraints, $\hat{m}$ could potentially become unbounded. Thus, Eq. 8 is introduced as an upper limit in Section 3.1.1.

Consequently, Eq. 6 and Eq. 8 serve as the lower and upper bounds for $\hat{m}$, respectively. Together with $L_{up}$, Eq. 7 is optimized to ensure $\hat{m}$ remains within a feasible range, fostering a more stable learning process.

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

We would like to clarify that the primary motivation of SCR is to enhance the **generalization** capabilities of RL algorithms. Non-informative reward is one of the challenges in metric-based representation methods. SCR is designed to address generalization enhancement along with non-informative reward by capturing long-term information. 

In our experimental setup, which primarily evaluates generalization, we also include specific tasks in sparse reward settings, i.e. cartpole_sparse and ball_in_cup, to demonstrate SCR's capability in  non-informative rewards. In the default setting (**Table 1**), SCR outperforms other metric-based methods in these tasks, while DBC and MICo fail, and SimSR underperforms in cartpole_sparse. DrQ, which is a data augmentation method, also achieves high performance, further validating SCR’s effectiveness alongside a leading method.

The distraction setting (**Table 2**) introduces additional complexities, where all methods, including SCR, show low scores and high variance in ball_in_cup. Nonetheless, SCR manages better performance in cartpole_sparse compared to baseline methods, which get zero or near zero scores. This indicates SCR's resilience and slightly superior handling of non-informative reward scenarios under more demanding conditions.