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

We thank Reviewer y1rM for providing the positive feedback and recognizing the importance of our work. Please see our response to your questions below.

(Weakness/Question1) More Generalization Experiments on "Button"

We would like to clarify why we did not conduct generalization experiments on the Button task. As shown in Table 1 (page 6), all baseline algorithms exhibit extremely high costs on Button-related tasks. While our SDQC framework achieves safe results, the corresponding reward remains exceptionally low.

In the Button task, the agent is required to reach the orange button (marked with a green cylinder) without colliding with obstacles, stepping into traps, or activating incorrect buttons. At difficulty level 2, however, the environment includes 17 distractors, six of which are dynamic and move around the agent. (For a visualization, please refer to the results [here].)

The optimal policy learned by the agent trained with SDQC involves moving alongside the dynamic obstacles to avoid collisions, but this behavior largely ignores the target button. In contrast, agents trained with other baseline algorithms tend to disregard the moving obstacles entirely and recklessly navigate toward the target button. These completely contradictory behaviors make comparisons between algorithms meaningless. Conversely, in other tasks, such as Goal and Push, agents trained with different algorithms exhibit similar behaviors, enabling fair generalization tests in these environments. (For reference, additional generalization experiments on the "PointGoal" and "PointPush" tasks are provided in Figure 9 (page 25) of our main paper.)

We will include a more detailed discussion of this phenomenon in the final version of the paper.

(Question2) Visualization results of difficulty-level.

Thank you for your insightful comments. We agree that clear visualizations can better illustrate the generalization performance of our SDQC.

In the Goal-related tasks, the agent is required to reach the goal (marked with a green cylinder) while avoiding traps and collisions with obstacles. A visualization of this task is provided [here]. At difficulty level 1 (row 1 in the GIF), the environment contains 8 traps and 1 obstacle. At difficulty level 2 (row 2), the complexity increases, with 10 traps and 10 obstacles.

In the Push-related tasks, the agent is required to push a yellow box to the goal (marked with a green cylinder) while avoiding traps and collisions with solid pillars. A visualization of this task is available [here]. At difficulty level 1 (row 1 in the GIF), the environment consists of 2 traps and 1 pillar. At difficulty level 2 (row 2), the environment becomes more challenging, with 4 traps and 4 pillars.

We thank Reviewer JBF8 for the valuable suggestions. We would like to address your concerns point by point below.

(Method1) Potential instability for Eq. 5

In Eq. 5, the exponential coefficient $\Gamma(s,\tilde{s})$ is a detached soft measure introduced in [1] that helps stabilize training. Regarding the potential instability caused by coupling value function updates with representation learning, a detailed analysis is provided in Appendix F.2 (pages 25–26). Experimental results demonstrate that the contrastive term has no adverse effect on value estimation, while our proposed ATN-based representation network significantly reduces estimation error.

(Method2) Why limited anchors lead to better results

We apologize for the unclear description. We use the full dataset for training, not just a subset. For each gradient step, we sample a batch of data and perform contrastive learning within the batch. Since our total loss includes multiple components, using more anchors could cause the neural network to overly focus on learning the representation, which might result in suboptimal performance of the value function.

(Method3) Upper-bound of cost value function

Please note that $Q_{h}$ is updated based on the TD error computed relative to $V_{low}$ rather than $V_{up}$. In contrast, $V_{up}$ is updated toward $Q_h$ via expectile regression parameterized by $\tau$. According to Lemma. 1 in [2], it is straightforward that $\lim_{\tau\rightarrow 1}V_{up}(z(s))=\max_{a\in\mathcal{A}_\beta^s}Q_h(z(s),a)$.

(Theory1) Why introduce a comparison with "bisimulation"

We conduct a theoretical comparison with bisimulation, which is one of the most widely used representation learning methods in RL [1, 3]. While implementing our idea of decoupling states for decision-making, bisimulation initially seemed like a natural choice. However, this approach proved unsuccessful. It fails during the initial model-estimation stage due to the non-smooth nature of the cost function and its sparse value distribution. Similar challenges have been observed in sparse reward RL tasks [4]. In contrast, our SDQC method avoids these issues by using smooth Q-values and provides stronger theoretical advantages. We will include this failed attempt in the final version of our paper.

(Theory2) Contrastive representations do not have good properties for decision-making.

In fact, contrastive-based representations are widely used in RL [1]. Similar to previous studies, our analysis is built upon the MDP framework, while our simulations address high-dimensional, infinite state and action spaces. We do not claim that contrastive learning achieves a rigorous Q-irrelevance representation; indeed, such a result is unattainable for infinite problems approximated by neural networks. Rather, we cluster the Q-indicated representations in high-dimensional space, analogous to previous works that cluster representations with their associated indicators.

(Theory3) Why conditional entropy can indicate "generalization ability"

Due to space limitations, please refer to our response to Reviewer zA7P on Explanations of "conditional entropy".

(Experiment1) "Fair comparison" with baselines.

The required "fair comparison based on same computational cost" is hard to execute. First, each baseline algorithm has unique optimal parameter settings suggested by their authors. Moreover, some baselines, such as CDT with GPT-2 backbone and TREBI with U-Net structure, are even more computational expensive than our SDQC.

According to reviewer's requirement, we implement FISOR with ATN-based structure and separate diffusion policies [here].

(Experiment2) Temperature ablations for diffusion policies

Please find ablations [here]. The default parameters follow FISOR and are kept consistent across all environments.

(Question) Consistent hyperparameters for all environments

We would like to emphasize that our choice of hyperparameters is domain-specific. (We use the same hyperparameters across 12 tasks in safety-gym domain.) Please refer to our response to Reviewer zA7P on "Complex hyperparameter settings". We argue that no RL algorithm can employ a unified set of hyperparameters across all domains and environments, including most commonly used SAC and PPO.

[1] Agarwal, Rishabh, et al. Contrastive behavioral similarity embeddings for generalization in reinforcement learning. ICLR (2021).

[2] Kostrikov, Ilya, et al. Offline Reinforcement Learning with Implicit Q-Learning. ICLR (2021).

[3] Castro, Pablo. Scalable methods for computing state similarity in deterministic markov decision processes. AAAI (2020).

[4] Lee, Vint, et al. DreamSmooth: Improving Model-based Reinforcement Learning via Reward Smoothing. ICLR (2024).

We acknowledge Reviewer zA7P for appreciating the significance and novelty of our work. Below, we respond to the your questions point by point.

(References) Missed discussing on relevant references

The suggested papers are indeed highly relevant to our work. They represent extensions of bisimulation-series representation methods. We will incorporate a discussion of these papers in Section 3.4 and the Related Works section.

(Weakness1) Complex hyperparameter settings

We aknowledge SDQC involves multi-step training and many hyperparameters. However, as shown in Table 3, we have set most of the hyperparameters to be consistent across all environments and domains. 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, we provide detailed guidelines for selecting them (in Appendix C.3, lines 1097-1133).

(Weakness2/Question2) Computational costs

The computational costs for SDQC are indeed slight higher than other baseline algorithms. We discussed this point in Appendix E.3, lines 1288-1302. To provide a clear comparison, we implement all algorithms with pytorch on a single machine equipped with one GPU (NVIDIA RTX 4090) and one CPU (AMD Ryzen 7950X), and report the results on CarPush2 task. During the training phase, the total training time (hours) are presented as follows.

During the testing (inference) phase, we measure the time consuption (seconds) over 1000 RL timesteps as follows.

(Weakness3/Question4) Explanations on "conditional entropy"

We apologize for any conceptual confusion caused by the missing definition. Here, we clarify that we denote $H(s|z)$ as conditional entropy, which is used measure information content of random variables, rather than cross-entropy or KL divergence, which are used to measure the distance between prob-distributions. Conditional entropy is a concept in information theory, and its formal definition was introduced by Shannon[1]:

$H(s|z)=-\sum_{s,z} p(s,z) \log p(s|z)$

where $p$ represents the probability function. According to [2], coarser representation in RL generally introduce better generalization performance due to reduced state space. Thus, when treating the representation as a random variable $z$, we initially aim to minimize the information content of $z$, i.e., the entropy $H(z)$. Note that the following relationship always holds:

$H(s|z) = H(s,z) - H(z) = H(s) - H(z)$

where the second equality is valid since $z$ is a function of $s$. For any given offline dataset, $H(s)$ is fixed. Therefore, minimizing $H(z)$ is equivalent to maximizing $H(s|z)$.

We will add above detailed explanation in the final version of our paper.

(Typos)

Thanks for pointing out these! We have carefully checked entire paper to correct all typos.

(Question1) What is "reverse expectile regression"

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 [3], we refer to this as "reverse expectile regression".

(Question3) What is Safe Agents and Unsafe Agents in Table 1

As introduced in Section 4.1 (lines 320–325, column 1, and lines 299–303, column 2), our ultimate objective is to achieve zero cost during testing, in alignment with the framework established by FISOR [3]. However, most baseline algorithms struggle to perform effectively under a zero-cost threshold. Therefore, following FISOR [3], we impose a strict cost limit of 10 for the Safety-Gymnasium environment and 5 for the Bullet-Safety-Gym environment. Agents with an average cost below this threshold during testing are classified as safe agents, while those exceeding the threshold are classified as unsafe agents.

Lastly, and importantly, we are committed to open-sourcing all code upon acceptance.

[1] Shannon, C. E. A mathematical theory of communication. The Bell system technical journal (1948).

[2] Liu, Guoqing, et al. Return-based contrastive representation learning for reinforcement learning. ICLR (2021).

[3] Zheng, Yinan, et al. Safe Offline Reinforcement Learning with Feasibility-Guided Diffusion Model. ICLR (2024).

We sincerely acknowledge Reviewer cp6m for providing the insightful feedback. Please see our response to your concerns below.

(Claim1) Motivation: How are safe offline datasets collected?

In offline RL, datasets are typically collected from past human experiences. For example, thousands of car accidents occur every day. These unsafe driving records, combined with safe ones, can be used to build a safe offline RL dataset for training autonomous driving systems. Compared to allowing UGVs to explore the real world through an online trial-and-error manner, it is both safer and more efficient.

(Claim2) How is it possible to decouple the representations according to $Q^{\*}$?

When focusing solely on the reward $r$, the optimal value $Q_r^{*}$ is inherently independent of cost information $c$, as defined by its formulation. For instance, as shown in Figure 1 (page 1), all states in the figure share the same $Q_r^{\*}$ across all actions. Leveraging this property, we cluster the mid-layer representations of states with similar $Q_r^{\*}$ values across actions within the support, ensuring these representations remain independent of cost-related information. Likewise, the same principle applies when considering only the cost. For further intuition, please refer to the GIF in the supplementary material, which visualizes policies derived from the decoupled states.

(Claim4) What does "generalization" refer to?

We apologize for any confusion caused by the lack of explanation regarding this terminology. In machine learning, "generalization" usually refers to a model's ability to perform well on new, unseen data not included in its training set.

In real-world decision-making, state spaces are often high-dimensional and infinite, making it impossible to cover all possible states during training phase. This is especially true in offline settings, where the agent is trained on a fixed dataset without exploration [1]. The online testing performance of an offline-trained agent directly reflects the algorithm's generalization ability [1]. During deployment, the agent may encounter unfamiliar states and take suboptimal or unsafe actions. As shown in Table 1, our SDQC algorithm achieves SOTA performance in both rewards and costs.

Moreover, our approach extends generalization to more complex scenarios, where the testing environment (containing more obstacles) differs significantly from the one in which the offline data was collected, as displayed in Figure 3 and Figure 9. (For online RL, generalization is usually evaluated in a similar way [2].)

(Question1) Applying SDQC to online settings.

We agree that SDQC has great potential for extension to online settings. However, the in-sample learning method (implicit-Q learning) used in our work is specifically designed for offline RL. Our primary contribution is the introduction of the concept of "decoupling the states for decision-making," which can be seamlessly integrated with many existing online RL methods that use representations like bisimulation. Unfortunately, due to time constraints, we were unable to verify this point during the rebuttal period.

(Question2) Possibilities of abstracting representations from pre-trained value functions.

If we understand correctly, the reviewer is asking whether it is possible to first train the value functions through in-sample learning and then extract the representations based on the converged value functions. We argue that this approach would make the training process more complex and prone to failure.

Specifically, given any global observation $s$, we should firstly obtain its cost-related representation $z_h$ then make safety assessment on representations $z_h$ according to value function $V_h(z_h)$. This process requires learning two mappings: $\mathcal{S} \mapsto \mathcal{Z}_h \mapsto \mathcal{V}_h$. However, if we were to train the value functions first and subsequently extract the representations based on these pre-trained value functions, the task would become more complex. This is because an additional mapping would be introduced: $\mathcal{S} \mapsto \mathcal{V}_h \mapsto \mathcal{Z}_h \mapsto \mathcal{V}_h$. In this case, the mapping between the vector space $\mathcal{Z}_h$ and the scalar space $\mathcal{V}_h$ would be bijective, which is extremely challenging for neural networks to learn. Furthermore, this approach would be susceptible to additional errors due to the approximation nature of neural networks.

In contrast, our approach utilizes the neural network structure presented in Figure 5 (Appendix C.2, page 19). The representations are explicitly set as mid-layer outputs, and only the necessary mappings are learned.

[1] Levine, Sergey, et al. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv:2005.01643 (2020).

[2] Agarwal, Rishabh, et al. Contrastive behavioral similarity embeddings for generalization in reinforcement learning. ICLR (2021).