State Chrono Representation for Enhancing Generalization in Reinforcement Learning

Thanks for your effort and insightful comments.

Weakness 1.& 2.

We appreciate your concern regarding the number of seeds. Due to resource constraints, we opted for a balance between the number of seeds and computational cost, thus using 5 seeds. To enhance the robustness of our results, each seed was evaluated over 100 episodes, totaling 500(=5x100) episodes across all seeds. The mean and std scores in Table 1 and 2 are calculated over all 500 returns.

The high variance in Table 2 come from the internal randomness of the distraction settings in the DM_Control environment, where backgrounds, body colors, and camera poses vary per episode. The combination of these attributes leads to varying distraction effects, making some episodes more challenging and resulting in lower scores. This variance underscores the challenges of these settings and validates the necessity of our approach for such distraction setting.

To further ensure the reliability of the results, we present aggregate metrics in Figure 4, a common practice to compare different RL algorithms. Figure 4 demonstrates that our SCR significantly outperforms other baselines in the distraction setting.

Regarding the results in the default setting shown in Table 1, our intention is not to claim that SCR outperforms other baselines but to demonstrate that it is comparable to them (lines 206-308). The default setting of DM_Control has been extensively studied, with SOTA methods achieving high scores. The purpose of including results in the default setting is to verify SCR's effectiveness in a standard scenario. Combined with the results in Table 2, this further confirms that SCR enhances generalization rather than merely trading off performance between distraction and default settings. Table 1 shows that the mean scores of SCR are comparable to baseline methods, with variance at a similar level, thereby validating our claims.

Table 3 is also evaluated over 5 seeds, with each seed evaluated over 100 episodes. We update the table to include standard deviations.

Self-predictive representations (SPR) shares similarities with SCR, as both methods learn representations regarding a future state. However, there are several key differences between them. SPR focuses on predicting future states, enforcing the representation of predicted future states to be close to the true future states. In contrast, SCR focuses on predicting a metric/distance between current state and future state. The metrics in SCR (in Section 3.1-3.3) are all related to rewards or values, which is correlated to policy learning. SPR's representation focuses on dynamics learning, without directly considering the reward or value.

Successor Representations (SRs) 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.

We would like to clarify the contributions of Section 3.1. The major technical contributions are as follows:

1. We identify the limitations of existing metric approximations.
2. We propose a novel approximation metric, detailed in Definition 3.3.
3. We provide theoretical support for this new metric in Theorem 3.4 and Lemma 3.5.

These points are original contributions of our paper. Empirical results in the ablation study (Section 4.3 / Figure 5) show that substituting our proposed metric with existing ones like MICo, SimSR, or DBC results in performance degradation, highlighting the effectiveness of our approach in Section 3.1.

From a writing perspective, retaining Section 3.1.1 in the main text is important as it introduces the long-term temporal information issues that motivate the subsequent sections (3.2 and 3.3). Although Eq.1 may appear less novel due to its similar form to existing work, it sets the stage for the necessity of designing a new approximation metric and leads to Eq.2, which is integral to the final objective (Eq.10) of SCR. We will revise Section 3.1 to enhance clarity and explicitly highlight its contributions.

Questions 1.

Representation collapse in bisimulation metrics has been discussed in [1]. This issue can arise in scenarios with sparse rewards, where the reward difference term in Eq.1 is always zero, leading to the collapse of the distance between states $x$ and $y$. Addressing sparse rewards is one of the objectives of our paper.

To mitigate the risk of collapse, we introduce a temporal measurement that extend to multiple states, thereby increasing the likelihood of encountering non-zero rewards. Additionally, the representation function in SCR is continuously updated through the RL objectives, which provide additional gradients and help avoiding the representation collapse.

We store entire episodes in the replay buffer in a contiguous manner, with each step stored as $(s, a, s')$, similar to other one-step off-policy methods. The only additional memory burden is the storage of indices for the start and end of each episode. These indices are integers, therefore, the memory overhead is negligible.

When sampling a batch of training data from the replay buffer, we first uniformly sample a batch of steps $i$. We then sample steps $j$ based on the indices of the start and end of the episode, ensuring that steps $i$ and $j$ are within the same episode. Finally, we swap $i$ and $j$ if $j > i$ to ensure that $j$ is a future step relative to $i$. This method keeps the memory overhead minimal while effectively tracking future states.

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

Table B: Mean and std over 5 averaged scores on DM_Control with distraction setting.

Thank you for your insightful reviews. We address your concerns as follows:

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.

We appreciate the reviewer's observations. We would like to clarify the core aim of this paper is to investigate the generalization capabilities of SCR. The major contributions of this paper are demonstrated in the results for the distraction setting (Table 2). In the default setting (Table 1), our claim is not that SCR outperforms baseline methods significantly but rather that it performs comparably (see lines 306-308). These experiments (Table 1 and 2) serve to demonstrate that SCR facilitates improved generalization rather than only offering a trade-off between default and distraction settings.

The SOTA methods achieve high scores in the default setting of DM_Control (Table 1), this study does not aim to surpass these in the default configuration. Table 1 illustrates that SCR maintains competitive mean scores and similar variance level to baseline methods, with the exception of the Reacher-Easy task. This Reacher-Easy task involves stretching a robotic arm to reach a randomly placed target in 2D plane at the start of each episode. The random cropping data augmentation technique (DrQ) proves effective in Table 1, while data augmentation currently absent in SCR. Nonetheless, SCR could incorporate such data augmentation to enhance performance in tasks like Reacher-Easy.

The results in Table 2 and Figure 4 robustly show that SCR significantly outperforms other baselines in the distraction setting. This setting, characterized by random selection of background videos, body colors, and camera poses at the beginning of each episode, presents a more complex challenge. These results underscore SCR's efficiency in generalizing across varied distractions. Our ablation study in Section 4.3 further delineates the contributions of individual SCR components to this performance.

1. What factors contribute to SCR's occasional underperformance in the default setting?

As detailed earlier, the lack of random cropping data augmentation in SCR compared to DrQ particularly affects its performance in the Reacher-Easy task. SCR focuses on measuring the distance between states, which is beneficial for generalization over distractions. Integrating data augmentation into SCR could potentially improve its performance in similar tasks.

2. Is the distraction setting providing SCR with an advantage, and would the findings be consistent in an alternative setting demanding generalization abilities?

The distraction setting from [1], used as a benchmark in DM_Control to assess the generalization abilities of RL algorithms, is not uniquely providing SCR with an advantage. Unlike implementations by others, e.g. DBC and SimSR, which modify only the background video, our approach also modify body color and camera pose. Such combined modifications increase the complexity of the setting, leading to more challenging environments. This rigorous setup ensures that our findings regarding SCR's generalization are consistent and reliable across various settings.

---

[1] The Distracting Control Suite--A Challenging Benchmark for Reinforcement Learning from Pixels. Stone et al. 2021

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

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.

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:

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:

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:

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.

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

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

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.

Thank you for your thorough and insightful feedback. We address your concerns as follows:

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.

The Reacher-Easy task presents a challenge because it requires the agent to locate a randomly placed target on a 2D plane, a task that heavily relies on understanding spatial relationships within the environment. DrQ leverages random crop data augmentation to effectively randomize the target’s absolute position, thereby help the agent in capturing the relative positioning, which is important for task success. SCR lacks of such augmentation and have less efficiency in this scenario. However, DrQ and SCR are different kinds of methods, data augmentation vs metric-based representation, integrating data augmentation into SCR would potentially utilizes DrQ's advantages and improves performance in Reacher-Easy.

Additionally, in the distraction setting of Reacher-Easy, background variations further distract perception of target object, leading to high variance and lower mean scores as shown in Table 2. While DBC appears to perform better by learning to stretch the robot arm towards a location that commonly spawns target objects, this strategy does not help to learn a robust policy for exact target reaching. We will to extend the limitations section to discuss these scenarios comprehensively in the revised manuscript.

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.

We acknowledge the necessity to discuss SPR and PlayVirtual. SCR focuses on predicting a metric or distance between the current and future states which is linked to reward, value function and policy learning. This is in contrast to SPR, which focuses on future states predictions to true future states, and thus centers more on dynamic learning without direct reward integration. PlayVirtual extends SPR’s approach by incorporating a backward dynamics model and consistency loss, but like SPR, remains focused on dynamics model. The revised manuscript will include a comparison to clarify these differences with SPR and PlayVirtual.

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.

We agree that data augmentation is effective for enhancing generalization in RL. Currently, SCR does not incorporate data augmentation techniques because it focuses on metric-based representation learning. However the robust performance of SCR in the distraction setting of DM_Control demonstrates its generalization ability, even comparing to data augmentation methods DrQ. Given the distinct mechanisms of SCR and DrQ, we believe integrating data augmentation with SCR is a promising way to further improve its effectiveness.

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?

SCR utilizes future information prediction/regularization by estimating a metric between current and future states rather than predicting exact future state dynamics, such as SPR, PlayVirtual and model-based methods. SCR is parameter-efficient and related closely with the RL’s $n$-step Bellman operator that can introduce higher value variance with increased prediction steps. However, SCR mitigates the higher variance issue by ensuring the predicted metrics in a feasible range rather than exact approximation.

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.

we will extend the limitations section in the revised manuscript to discuss scenarios where SCR may underperform, specifically emphasizing tasks like Reacher-Easy.

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.

Thanks for your insightful suggestion. We will improve the related works on how established practices such as SPR, PlayVirtual, and other model-based methods contrasting to SCR in integrating future information.

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