Q-Supervised Contrastive Representation: A State Decoupling Framework for Safe Offline Reinforcement Learning

We thank Reviewer 8jhZ for providing valuable suggestions on our research. We address the concerns in your review point by point below.

[W1] Lack of Novelty

We regret to make readers feel that our paper is lack of novelty. As a matter of fact, we believe that our main contribution is the idea of decoupling the global observations into reward-related representation, and cost-related representation for decision making. To the best of our knowledge, we are the first to utilize representation learning in continuous state-based RL tasks, and the first to introduce the concept of state-decoupling for decision-making.

[W2] Lack of Comparison with Model-Based Methods

To the best of our knowledge, there is no model-based safe offline RL algorithms for now. For the comparison with bisimulation, please refer to our response in [Q3]

[Q1] Convergence of Joint Optimization

We agree that the joint optimization of the value functions and representations poses a potential risk of training instability. However, experimental results indicate that, although the inclusion of representation loss results in an increase in critic and value loss, it does not compromise the overall stability of the training. We have added a new section (Appendix F.2, Pages 28-29, lines 1502-1537) with new Figure 9 to present the value function estimation error during the training process. Notably, our proposed neural network architecture, which incorporates an attention-based state encoder, significantly enhances the precision and stability of value function learning compared to the simple MLP used by FISOR. In fact, we had previously recognized this issue intuitively and attempted to address it by employing a soft-start training approach, which involves introducing the representation contrastive term only after the value function converges. However, experimental results indicate that this method is not effective in clustering the representations, leading to suboptimal final performance.

[Q2] Necessity of the Trade-Off Policy

The proposed SDQC framework encompasses three distinct policies: the reward policy $\pi_r$, which prioritizes maximizing rewards while disregarding costs; the trade-off policy $\pi_{to}$, which considers both rewards and costs; and the cost policy $\pi_h$, which aims to escape dangerous regions as swiftly as possible. We emphasize that the trade-off policy $\pi_{to}$ is essential and cannot be omitted.

For any given state $s$ , we denote the set of feasible actions within offline datasets as $\mathcal{A}^s_\beta$. A condition where $V_h^{low} > 0$ indicates that there is no action $a \in \mathcal{A}^s_\beta$ capable of ensuring a completely safe subsequent trajectory. Our reward policy relies solely on reward-related representations and neglects cost-related information. Consequently, if we use only $V_h^{low}$ as the threshold, it may be too late for the agent to transition from $\pi_r$ to $\pi_h$ once it has entered a state where $V_h^{low} > 0$. By introducing $V_h^{up}$, which reflects the existence of an action $a \in \mathcal{A}^s_\beta$ that could potentially lead to an unsafe trajectory in the future, we enable the agent to consider both types of information through the trade-off policy $\pi_{to}$. In fact, the naive implementation of $\pi_{to}$ within our SDQC framework is analogous to the safe policy $\pi_{safe}$ in FISOR. However, due to our incorporation of representation learning and the introduction of $\pi_r$, our SDQC exhibits superior generalization capabilities.

To substantiate this point with experimental evidence, we have conducted ablation studies on the deployment of the three distinct policies in Appendix D (Page 26, lines 1393-1411). Numerical experimental results clearly demonstrate that, in the absence of a trade-off policy, a naive combination of reward and cost policies leads to higher costs. Such outcome can be attributed to the agent's inability to respond promptly to borderline dangers. Additionally, we have included a GIF in the supplementary materials to illustrate the trajectories generated by the collaboration of the three policies, as well as the trajectories produced by each of the three naive policies in isolation.

[Q3] Comparison with Bisimulation

Initially, we attempted to complete the training of representations using bisimulation. However, we encountered challenges during the initial training phase. Specifically, estimating cost functions proved to be difficult due to their non-smooth and sparsely distributed nature. To substantiate this claim, we have added a new section (Appendix F.3, Page 29, lines 1539-1582) where the cost-value estimation accuracy, precision, recall, and corresponding F1 score are reported. The experimental results indicate that the model correctly identifies actual dangerous state-action pairs with a probability of only about 77%, suggesting that the model’s ability to assess dangerous conditions is insufficient to support effective bisimulation training in subsequent stages. Therefore, we propose learning representations based on optimal Q functions to address these issues, owing to the continuity and non-sparsity of the Q functions. From a theoretical standpoint, our optimal Q based representation approach is more advantageous as it introduces a coarser representation.

[Q4] Unclear Evaluation Metrics

We acknowledge Reviewer 8jhZ for the valuable feedback. We assert that both FISOR [1] and our SDQC framework aim to achieve Zero Cost Violations, a goal that has not been addressed by previous baseline algorithms. We believe that in certain domains where safety is critical, a zero-cost setting is essential. For instance, in autonomous driving, it is imperative that a vehicle does not physically collide with any objects. Most prior works have focused on soft constraints, and the implementation of a zero-cost threshold has led to training instability. Consequently, we have adopted the approach used in FISOR [1] by establishing a small threshold. In the revised manuscript, we have revised our description of the evaluation metrics as follows:

Our ultimate objective is to achieve a zero-cost deployment, aligning with the framework established by FISOR [1]. However, most baseline algorithms struggle to operate effectively under a zero-cost threshold. Consequently, in accordance with FISOR [1], we impose a stringent cost limit of 10 for the Safety-Gymnasium environment and 5 for Bullet-Safety-Gym. We employ the metrics of normalized return and normalized cost for evaluation, where a normalized cost below 1 signifies a safe operation.

Furthermore, while the SDQC framework adopts more conservative strategies in certain environments to achieve lower costs, we believe this approach holds practical significance. Additionally, we utilized a unified set of hyperparameters at the domain level, without fine-tuning for specific tasks or environments, which may have the potential to yield improved performance.

Regarding the reward metric, we adhere to the environment settings established by DSRL and report normalized values, as the reward signals vary significantly across different tasks (for instance, the maximum reward for CarGoal2 is approximately 30, whereas for HalfCheetah-Velocity, it is around 2700).

[1] Yinan, Zheng, et al. "Safe Offline Reinforcement Learning with Feasibility-Guided Diffusion Model." International Conference on Learning Representations, 2024.

I appreciate the authors' response and clarifications. However, my concerns and questions are not fully addressed:

Q1: The explanation provided is not convincing. The authors showed results for the value function estimation error during the training process, but this does not directly answer my question. Could the authors provide the actual results for the value functions and Q-functions during training?

Additionally, could the authors explain why introducing the representation contrastive term only after the value function converges does not improve performance? Intuitively, this approach should lead to a better soft similarity measure as per Eq. (5), which could enhance representation learning and ultimately benefit performance, such as lowering costs and increasing rewards.

I also observed in Table 6 that as the Contrast Coefficient ($\delta$) increases, the performance appears to degrade, with both lower rewards and higher costs. Could the authors clarify this? Based on Eqs. (8) and (12), a larger $\delta$ should result in better state representations, which should positively impact performance.

Q4: Could the authors provide results for SDQC and FISOR with a zero cost threshold? The authors highlighted that "their proposed neural network architecture significantly improves the precision and stability of value function learning compared to the simple MLP used by FISOR". Therefore, I would expect SDQC to achieve zero constraint violations and demonstrate better performance than FISOR under such conditions.

Thank you for providing additional clarifications. However, my primary concern about the convergence of joint optimization remains unresolved. In your response to Reviewer FtTa, you stated that "Contrastive Loss Serves as an Auxiliary Loss." Could the authors further clarify this? It seems that every loss term in Eq. (8) and Eq. (12) depends on the learned representations, and the representation learning loss, in turn, depends on $Q$ and $V$. The learning of value functions and representations are fully intertwined and occur simultaneously.

[Q1] Computational Cost During Training/Testing Phase. Ablation studies on the Contrastive-related Hyperparameters

We have detailed the computational costs for both the training and testing phases of SDQC in Appendix E.3 (Page 27, lines 1448-1462). During training, SDQC involves evaluating Q-values within a sampling batch, and the computational cost is significantly influenced by the number of anchors used. Table 5 (Page 26, lines 1359-1367) presents the computing time under different anchor choices. Notably, for the "CarPush2" task, using 8 anchors results in a runtime of 32.7 seconds per epoch (defined as 1000 gradient steps). Without contrastive learning, the runtime decreases to 23.7 seconds--the difference is not substantial. In the testing phase, SDQC requires 11.13 seconds for inference over 1000 RL timesteps, which is slightly higher than FISOR (6.11 seconds), yet significantly lower than TREBI (585.87 seconds).

Furthermore, we have added the ablation studies on the contrastive-related hyperparameters (the term coefficient $\delta$ and the exponential temperature $\nu$). With respect to the temperature $\nu$, employing a very small value tends to destabilize the training process, ultimately resulting in collapse. Conversely, using a larger value produces poorly clustered representations, leading to a marked degradation in performance. Regarding the term coefficient $\delta$, a smaller value results in a slight performance decline. However, a larger coefficient excessively prioritizes the contrastive loss, destabilizing the training of the value function and degrading performance. More detailed results and discussions can be found in Appendix D (Page 26, lines 1369-1390). While fine-tuning these hyperparameters for specific environments and tasks could potentially yield better experimental results on the benchmark, we choose not to do so.

[Q2] Additional Experiments on Bisimulations

Initially, we attempted to complete the training of representations using bisimulation. However, we encountered challenges during the initial training phase. Specifically, estimating cost functions proved to be difficult due to their non-smooth and sparsely distributed nature. To substantiate this claim, we have added a new section (Appendix F.3, Page 29, lines 1539-1582) where the cost-value estimation accuracy, precision, recall, and corresponding F1 score are reported. The experimental results indicate that the model correctly identifies actual dangerous state-action pairs with a probability of only about 77%, suggesting that the model’s ability to assess dangerous conditions is insufficient to support effective bisimulation training in subsequent stages.

[Q3] Extension on Visual-Based Environments

We sincerely thank Reviewer zft8 for the valuable comments. To the best of our knowledge, current safe offline RL algorithms have not yet addressed image-based tasks, and there are no relevant datasets available in the test benchmarks. It is well known that image-based tasks are more complex than state-based tasks and depend more heavily on representation learning. Nevertheless, due to time constraints, we are unable to apply our SDQC in image-based environments at this time. We plan to explore this avenue in our future research.

[1] Castro, Pablo, et al. "Using bisimulation for policy transfer in MDPs." Proceedings of the AAAI conference on artificial intelligence, 2010.

[2] Castro, Pablo. "Scalable methods for computing state similarity in deterministic markov decision processes." Proceedings of the AAAI Conference on Artificial Intelligence, 2020.

[3] Zhang, Amy, et al. "Learning Invariant Representations for Reinforcement Learning without Reconstruction." International Conference on Learning Representations, 2021.

[4] Agarwal, Rishabh, et al. "Contrastive Behavioral Similarity Embeddings for Generalization in Reinforcement Learning." International Conference on Learning Representations, 2021.

[5] Lee, Vint, et al. "DreamSmooth: Improving Model-based Reinforcement Learning via Reward Smoothing." International Conference on Learning Representations, 2024.

We thank Reviewer zft8 for recognizing the importance of our work. Please see our response to your concerns below.

[W1] Why Compared with Bisimulation

As bisimulation is one of the most widely used representation learning methods in the relm of reinforcement learning [1–4], we choose to establish theoretical comparisons with it. In particular, we demonstrate that the optimal Q based representations proposed in our SDQC yield coarser representations than those derived from bisimulation.

While implementing our original idea of decoupling global observations into reward- and cost-related representations, bisimulation initially appeared to be a natural choice for training. Nevertheless, we found that bisimulation fails even at the initial model-estimation stage due to the non-smooth nature of the cost function and its sparse value distribution. Similar issues have been reported in sparse reward RL tasks [5]. For further details, please refer to our response in [Q2]. Alternatively, our representation method based on optimal Q values resolves these issues due to the continuity and non-sparsity of the Q functions.

[W2] Additional Computational Costs

We admit that the training of SDQC requires higher computational cost compared with other baseline algorithms, which we have listed as an limitation in Appendix G. However, during the deployment (testing) process, its inference speed is not significantly lower than that of diffusion-based policies. Please refer to our response in [Q1].

[W3] None of Baseline Algorithms Use Representation Learning for State-Based RL

Indeed, this is the key innovation of SDQC compared to existing baseline algorithms. Our experiments demonstrate that representation learning significantly enhances the model's generalization ability. To the best of our knowledge, we are the first to introduce representation learning for state-based tasks within the field of safe offline RL.

We thank Reviewer 6xP2 for providing positive feedback and recognizing the effectiveness of our work. Please see our response to your questions below:

[W1] Lack of Psdudo Code

We apologize for any confusion resulting from the absence of pseudocode. We acknowledge that SDQC necessitates multi-phase training, and a detailed pseudocode can improve clarity. We have included it in our updated PDF documents, available on Page 21, lines 1099-1124.

[W2] Clear Justification

We agree that providing clear justification for each step is essential for an intuitive understanding of the proposed algorithms. To address this issue, we have added a summary section (Appendix C.1, Pages 21-22, lines 1094-1166) on the training and deployment of SDQC to enhance understanding of the process. In Appendix C.1, we explain the purpose of each phase of training and the rationale behind our three distinct policies.

[W3] Scientific Value of the Paper

We sincerely thank Reviewer 6xP2 for recognizing the excellent engineering performance of SDQC. While we acknowledge that SDQC is indeed complex to implement, our algorithm requires minimal hyperparameter tuning and is adaptable to various domains. We believe that our primary contribution lies in the novel idea of decoupling the state into reward-related and cost-related representations. Our experimental results demonstrate that this approach generalizes effectively to unseen states in the dataset. Although SDQC is currently challenging in training complexity, it offers valuable insights for advancing safe reinforcement learning, particularly in enhancing the generalization capabilities of future algorithms.

[W4] Writing Issues

We acknowledge Reviewer 6xP2 for detailed guidance on improving the writing of the paper. We have thoroughly reviewed the manuscript, corrected the typographical errors, and made revisions where necessary.

[Q1&Q2] Temperature-related Parameters

The settings for the temperature-related hyperparameters are detailed in Table 3 (Page 24, lines 1268-1269). For all environments/tasks in our experiments, we consistently set the values $\eta=1.0$ and $\nu=0.1$, which proved to be effective. Additionally, we conducted ablation studies on the contrastive-related hyperparameters, the results of which are discussed in our response to [Q3].

[Q3] Additional Ablation Studies on Contrastive-related Hyperparameters

We added the ablation studies on the contrastive-related hyperparameters (the term coefficient $\delta$ and the exponential temperature $\nu$). With respect to the temperature $\nu$, employing a very small value tends to destabilize the training process, ultimately resulting in collapse. Conversely, using a larger value produces poorly clustered representations, leading to a marked degradation in performance. Regarding the term coefficient $\delta$, a smaller value results in a slight performance decline. However, a larger coefficient excessively prioritizes the contrastive loss, destabilizing the training of the value function and degrading performance. More detailed results and discussions can be found in Appendix D (Page 26, lines 1369-1390). While fine-tuning these hyperparameters for specific environments and tasks could potentially yield better experimental results on the benchmark, we choose not to do so.

We acknowledge Reviewer FtTa for providing the helpful suggestions. Please see our response to your concerns below.

[W1] Unclear contributions

We have revised our introduction to assert that our SDQC is built upon FISOR, with the main distinction being the state-decoupling framework. The modifications can be found on Page 2, lines 72-79, as follows:

Attributable to the successful application of Hamilton-Jacobi (HJ) reachability analysis in Safe RL, which introduces a safety analysis method iterated through Q-learning with convergence guarantees, our approach makes safety assessments on the cost-related representations and make decisions based on the assessment results. Our SDQC, developed based on FISOR, distinguishes itself from FISOR and other classical methods which rely on global observations for decision-making (as depicted in the left subplot of Figure 2), by being the first to utilize decoupled representations for decision-making in safe RL tasks (see the right subplot of Figure 2).

[W2] Lack of Pseudo Code

We apologize for any confusion resulting from the absence of pseudocode. We acknowledge that SDQC necessitates multi-phase training, and a detailed pseudocode can improve clarity. We have included it in our revised manuscript, on Page 21, lines 1099-1124. Additionally, we have added a summarization section (Appendix C.1) on the training and deployment of SDQC, providing an intuitive understanding of the process (Pages 21-22, lines 1094-1166).

[W3] Convergence of Q for SDQC

We agree that the joint optimization of the value functions and representations poses a potential risk of training instability. However, the experimental results indicate that, although the inclusion of representation loss results in an increase in critic and value loss, it does not compromise the overall stability of the training. We have added a new section (Appendix F.2, Pages 28-29, lines 1502-1537) with the new Figure 9 to present the value function estimation error during the training process. Notably, our proposed neural network architecture, which incorporates an attention-based state encoder, significantly enhances the precision and stability of value function learning compared to the simple MLP used by FISOR.

[W4] Usage of Expectile Regression

We would like to emphasize that the expectile regression employed in IQL [1] is intended to fit the optimal value function, specifically to ensure that $V(s) \rightarrow \max_{a\in\mathcal{A_\beta}}Q(s,a)$. It is used for approximating the upper bound rather than merely for naive parameter tuning to enhance the accuracy of the value estimate. Here we provide a copied sentence from Section 4.2 (Page 4) of IQL [1]:

We can use expectile regression to predict an upper expectile of the TD targets that approximates the maximum $r(s, a) + \gamma Q(s',a')$ over actions $a'$ constrained to the dataset actions.

It can also be referred to Theorem 3 in IQL [1]: $\lim_{\tau\rightarrow1}V_{\tau}(s)=\max_{a\in\mathcal{A}, s.t. \pi_{\beta}(a|s)>0} Q^*(s,a)$ The reason for not setting $\tau\rightarrow1.0$ is that the overestimation issue cannot be corrected under such circumstances. Experimental results indicate that $\tau = 0.9$ is the optimal choice.

Note that the optimal value function under the general Bellman operator is given by $V_r(s)=\max_a Q_r^*(s,a)$, and is $V_h(s)=\min_a Q_h^*(s,a)$ under the cost-related safe Bellman operator. The expectile regression $L^{\tau}(u)=\vert \tau - \mathbb{I}(u<0)\vert u^2$, with $\tau \in (0.5, 1)$ can be used to approximate the maximum value. Conversely, the reverse version $L^{rev}_{\tau}(u)=\vert \tau - \mathbb{I}(u>0)\vert u^2$ with $\tau \in (0.5, 1)$ can be utilized to approximate the minimum value. The difference lies in the direction of the indicator function $\mathbb{I}(u)$. Following the setting in FISOR [2], we refer to this as "reverse expectile regression".

[W5] Generalization Experiments

Thanks for pointing this out. We note that "generalization" in the field of machine learning typically refers to a model's ability to perform well on new, unseen data that was not part of its training set. It does NOT strictly pertain to the model's ability to handle entirely new tasks or scenes. To the best of our knowledge, most studies on representation learning in reinforcement learning use "generalization" to describe a model's ability to handle the same task but with a slightly different observation space (e.g., shifted or distorted images) [3, 4]. In our experiments, the number of obstacles in the training sets differs significantly from those in the testing sets, thereby the radar observations for the agent will be totally different. We do not agree that "generalization" is overstated for our settings. Nonetheless, in accordance with the reviewer's suggestions, we have emphasized this point in our revised introduction (Page 2, lines 103-104) as follows:

Further, in generalization tests where agents are evaluated in environments with a different number of obstacles than those in the training dataset, all baseline algorithms show a substantial increase in cost and/or a significant decline in reward. In contrast, SDQC stands out as the only approach that guarantees no increase in cost while experiencing only a slight decay in reward.

[W6] Incomplete Ablation Studies

Thank you for your suggestions. We have added ablation studies on the deployment of the three distinct policies in Appendix D (Page 26, lines 1393-1411). The experimental results demonstrate that none of these policies can be omitted, as the best performance consistently arises from their collaboration. It is important to note that if we rely solely on safe/unsafe conditions and make decisions based on global observations rather than representations, our SDQC method effectively reduces to FISOR[2].

We have also added the ablation studies on the contrastive-related hyperparameters (the term coefficient $\delta$ and the exponential temperature $\nu$). With respect to the temperature $\nu$, employing a very small value tends to destabilize the training process, ultimately resulting in collapse. Conversely, using a larger value produces poorly clustered representations, leading to a marked degradation in performance. Regarding the term coefficient $\delta$, a smaller value results in a slight performance decline. However, a larger coefficient excessively prioritizes the contrastive loss, destabilizing the training of the value function and degrading performance. More detailed results and discussions can be found in Appendix D (Page 26, lines 1369-1390)

[W7] Lack of Visualization

We agree that effective visualizations can provide an intuitive understanding of our proposed SDQC. To this end, we have included a GIF in the supplementary materials. It showcases the trajectory induced by the collaboration of the three policies, as well as the trajectories induced by each of the three naive policies individually. These results align with our discussion in response to [W6]. Due to size limits of supplementary materials, we are only presenting the visualization for the "PointGoal2" task. Additional results will be included in the final version of our paper.

[W8] Over-tuning on Hyperparameters

We would like to emphasize that the excellent performance of SDQC DOES NOT rely on over-tuning the hyperparameters. As shown in Table 3, we have set most of the hyperparameters to be consistent across all environments and domains. Furthermore, we have nearly unified all hyperparameters within the same environment, even when dealing with different tasks and agents, as illustrated in Table 4. For instance, in the Safety Gymnasium domain, we use the same hyperparameters for 12 tasks (PointGoal1, PointGoal2, PointPush1, PointPush2, PointButton1, PointButton2, CarGoal1, CarGoal2, CarPush1, CarPush2, CarButton1, CarButton2). Users only need to select the network structure and adjust the encoded state dimension according to the global observation space, which is largely determined by the physical characteristics of each domain. In fact, the guidelines for selecting hyperparameters are provided in our paper (Appendix C.3, lines 1289-1295). We believe that such domain-level hyperparameter unification is acceptable. While we acknowledge that fine-tuning hyperparameters for specific tasks may yield better results, we have chosen not to pursue it.

[Q1] Convergence of Joint Optimization

Please refer to our response in [W3].

[Q2] Usage of Expectile Regression

Please refer to our response in [W4].

[Q3] Terminology Generalization

Please refer to our response in [W5].

[Q4] State-Decoupling Framework Ablation

Please refer to our response in [W6].

[Q5] Content in Table 4

In Table 4, we use "All" to represent all tasks/environments for the same agent. We apologize for any confusion this may have caused and we have included further details in the corresponding caption.

[Q6] Typos in Paper

We apologize for the typos we made. We have carefully checked entire paper to correct all typos.

[1] Kostrikov, Ilya, et al. "Offline Reinforcement Learning with Implicit Q-Learning." International Conference on Learning Representations, 2022.

[2] Yinan, Zheng, et al. "Safe Offline Reinforcement Learning with Feasibility-Guided Diffusion Model." International Conference on Learning Representations, 2024.

[3] Agarwal, Rishabh, et al. "Contrastive Behavioral Similarity Embeddings for Generalization in Reinforcement Learning." International Conference on Learning Representations, 2021.

[4] Kirk, Robert, et al. "A survey of zero-shot generalization in deep reinforcement learning." Journal of Artificial Intelligence Research 76 (2023): 201-264.

Question 4: Clarifying the network structure is also important.

We utilized the whole Section (Appendix C.2, Pages 22-23, lines 1168-1239) to introduce the network structure. We even mentioned it twice in the main text of our paper (Page 5, lines 231 and Page 9, line 460).

Additional t-SNE Visualization on the Representations

We would like to kindly remind the reviewer FtTa that, the deadline for submitting the final version of the revised paper is set for Nov.27, not Dec.3. During the extended period, only forum activities are permitted. We are unable to complete such a huge amount of work (refactoring the codebase of SafetyGymnasium) in such a short time. We estimate this will take at least 1-2 weeks. We have committed to including this in the final version of our paper and hope for your understanding regarding the time constraints (only one day left).

[1] Agarwal, Rishabh, et al. "Contrastive Behavioral Similarity Embeddings for Generalization in Reinforcement Learning." International Conference on Learning Representations, 2021.

[2] Khosla, Prannay, et al. "Supervised contrastive learning." Advances in neural information processing systems, 2020.

[3] Yinan, Zheng, et al. "Safe Offline Reinforcement Learning with Feasibility-Guided Diffusion Model." International Conference on Learning Representations, 2024.

[4] Kostrikov, Ilya, et al. "Offline Reinforcement Learning with Implicit Q-Learning." International Conference on Learning Representations, 2022.

We thank Reviewer FtTa for providing further feedback. We would like to clarify on the

Differences between PSEs [1] and SDQC (ours)

Firstly, we would like to emphasize that contrastive learning serves only as a tool employed by both PSEs [1] and SDQC for learning representations. (In fact, it is a common technique originally derived from computer vision [2]). However, the methodologies underlying PSEs [1] and SDQC are fundamentally different. PSEs leverage behavioral similarity to supervise representation learning, whereas our SDQC utilizes optimal Q values for this purpose.

Question 1: In [1], it states that contrast learning is an auxiliary loss that helps the RL agent learn the representation. But this work first explains how to distinguish between reward-related and cost-related states using different representations through contrastive learning, and then trains the RL agents based on the representations. It seems that the use of contrastive learning to learn the representation and RL learning policies are separate while it does not.

In PSEs [1], contrastive learning serve as an auxiliary loss, assisting the RL agent in learning representations, and is utilized for end-to-end policy learning. This is not the case in our proposed SDQC. As you just mentioned, in SDQC, contrastive learning is employed solely for representation learning and is not applied for RL policy learning. Furthermore, the training processes for representations and polices are completely separate. This distinction is clearly articulated in Section 3.3:

In the initial phase, we undertake the learning process for the value functions and representations associated with cost and reward separately. Following that, we extract the policy based on the acquired value functions and representations.

Additionally, you can refer to the pseudocode provided on Page 21, lines 1099-1124, and our corresponding summarization on the same page, lines 1125-1136.

Question 2: I suggest that you should stress that contrast learning is used as an auxiliary loss for distinguishing between reward-related and cost-related states, while the end-to-end RL plays a main role in decision-making by optimizing $Q_r$ or $Q_h$. The representation is not only learned from contrast learning but also from RL training objectives. I believe better structure can make it easier for readers to understand and less confusing. So you should consider my suggestion and try to adjust your writing.

We would like to clarify that PSEs [1] and SDQC are fundamentally distinct concepts. Importantly, our approach does NOT integrate the policy learning process into the representation learning. We respectfully request that Reviewer FtTa re-evaluate the structure of our writing on methodology as outlined below:

• We first introduce the motivation behind SDQC, which involves abstracting coarser reward- and cost-related representations, in Section 3.1.
• Then we present our unique approach to abstracting the representations based on optimal Q values, alongside the corresponding methodology (contrastive learning), in Section 3.2.
• Following this, we provide a detailed implementation of our SDQC in Section 3.3: initially learning the representations using in-sample Q-learning methods, and then extracting the policies based on these pre-learned representations. We would like to kindly remind Reviewer FtTa that learning optimal Q values without learning policies is the fundamental feature of in-sample learning, as introduced by IQL [4].

If reviewer FtTa has any further suggestions, we are pleased to follow your instruction to polish our writing.

Question 3: Besides, after reading [1], I found it more like a work applying [1] to safe offline RL areas. I would like to ask for more novelty beyond combining [1] and FISOR. (such as what problem you found in safe offline RL and why using representation can perfectly solve it)

We reclaim that, PSEs and SDQC are fundamentally different. The sole similarity lies in our use of supervised contrastive learning [2] to derive representations. PSEs learn representations based on policy similarities and function as an online RL algorithm tailored for image-based RL tasks. In contrast, our SDQC approach learns representations based on optimal Q values and operates as an offline RL algorithm designed for state-based RL tasks. To the best of our knowledge, we are the first to apply representation learning to continuous state-based RL tasks and the first to introduce the concept of state-decoupling in decision-making. Indeed, we believe that our main contribution is the innovative concept of decoupling the global observations into reward-related and cost-related representations for decision-making. We hope that Reviewer FtTa can reconsider the novelty of our proposed method.

We acknowledge reviewer FtTa for providing further suggestions, and we are willing to follow your instruction to further polish our paper.

Similar Network Structure

Your understanding is correct: SDQC learns representations based on optimal Q-values instead of RL policies (as in PSEs). SDQC and PSEs indeed share the common characteristic of learning representations with contrastive loss derived from reinforcement learning features. However, we would like to emphasize that PSEs is designed for image-based RL tasks in online settings and require a data augmentation process and further projection layers. In contrast, SDQC is designed for state-based RL tasks in offline settings and is implemented alongside in-sample learning methods. While the network structures of PSEs and SDQC may appear similar, we would like to clarify that, for representation learning in reinforcement learning, such a network structure is a necessary choice. This is because representations must originate from the ground-truth observations, and value functions or policies must build upon these representations. This sequentially connected neural network design is a common choice but not unique to PSEs or SDQC.

Contrastive Loss Serves as an Auxiliary Loss

Your understanding is accurate, and we have followed your guidance to revise our manuscript from the following three perspectives. (The detailed changes are marked in blue in the updated manuscript.)

• We presented our representation learning methods predicated on optimal Q in Section 3.2, detailing the corresponding methodology in Eq. 5. Furthermore, we added an explanation of how the contrastive loss serves as an auxiliary loss on Page 5, lines 239-245, as follows:

It is important to note that Eq. 5 requires precise calculation of optimal Q-values for all states across all actions, i.e., the constraints in Eq. 4 are satisfied. However, the Q-values are derived from the representation network, and even a small change in the network can result in variations in the Q-values. Therefore, it is necessary to integrate the training process of the representation network with the Q-learning process. This integration can be achieved by incorporating the contrastive loss as an auxiliary objective during the learning of the optimal value functions. Please refer to Section 3.3 for further details.

• Additionally, we have emphasized this point in the detailed methodology section. On Page 5, lines 259-260, we have revised the introduction to reward-related representation learning as follows:

The reward-related representations are acquired using the auxiliary contrastive loss term described in Eq. 5, with a weighting factor denoted by $\delta$. Consequently, the overall loss for the reward-related value functions and representations is formulated as follows:

$L_{reward} = L_{V_r} + L_{Q_r} + \delta L_{\theta_r}$

• In Page 6, lines 274-275, we made the similar modifications as follows:

By incorporating the auxiliary contrastive loss term (Eq. 5) with a weighting factor of $\delta$, we express the overall loss for the cost-related value functions and representations as:

$L_{cost} = L_{V_h^{\rm low}} + L_{V_h^{\rm up}} + L_{Q_h} + \delta L_{\theta_h}.$