Q-Supervised Contrastive Representation: A State Decoupling Framework for Safe Offline Reinforcement Learning
Authors: Zhihe Yang
Yunjian Xu
Yang Zhang
Venue: ICML 2025 poster (2025)
Keywords: Safe Reinforcement Learning
Offline Reinforcement Learning
Representation Learning
Contrastive Learning
Self-Supervised Learning
Primary area: reinforcement_learning->batchoffline
TL;DR: To address the OOD issue during testing for safe offline RL, we propose the first framework that decouple the global observations into reward- and cost-related representations through Q-supervised contrastive learning for decision-making.

Abstract: Safe offline reinforcement learning (RL), which aims to learn the safety-guaranteed policy without risky online interaction with environments, has attracted growing recent attention for safety-critical scenarios. However, existing approaches encounter out-of-distribution problems during the testing phase, which can result in potentially unsafe outcomes. This issue arises due to the infinite possible combinations of reward-related and cost-related states. In this work, we propose State Decoupling with Q-supervised Contrastive representation (SDQC), a novel framework that decouples the global observations into reward- and cost-related representations for decision-making, thereby improving the generalization capability for unfamiliar global observations. Compared with the classical representation learning methods, which typically require model-based estimation (e.g., bisimulation), we theoretically prove that our Q-supervised method generates a coarser representation while preserving the optimal policy, resulting in improved generalization performance. Experiments on DSRL benchmark problems provide compelling evidence that SDQC surpasses other baseline algorithms, especially for its exceptional ability to achieve almost zero violations in more than half of the tasks, while the state-of-the-art algorithm can only achieve the same level of success in a quarter of the tasks. Further, we demonstrate that SDQC possesses superior generalization ability when confronted with unseen environments.

Q-Supervised Contrastive Representation: A State Decoupling Framework for Safe Offline Reinforcement Learning
Zhihe Yang'
Yunjian Xu
Yang Zhang
Abstract
Training environment
Safe offline reinforcement learning (RL., which aims to learn the geeettggaa policy with- out risky online interaction with environments, has attracted growing recent attention for safety- critical scenarios. However, existing approaches encounter out-of of--iiitriuution problems during the testing phase, which result in potentially safe outcomes. This issue arises due to the infinite possible combinations of reward-related and cost- related states. In this work, we propose State Decoupling with Q-supervised Contrastive rep- resentation (SDQC), novel framework that de- couples the global observations into reward- and cost-related representations for decision-makine thereby improving the generalization capability for unfamiliar global observations. Compared with the classical representation learning methods, which typically require model-based estimation (e.g., bisimulation). we theoretically prove that our Q-supervised method generates coarser rep- resentation while preserving the optimal policy, resulting in improved generalization performance. Experiments on DSRL benchmark provide com- pelling evidence that SDQC surpasses other base- line algorithms, especially for its exceptional abil- ity to achieve almost zero violations in more than half of the tasks. Further, we demonstrate that SDQC possesses superior generalization ability when confronted with unseen environments.
Testing environment
Agent (L.GV)
Agent's observation
Target
Obstacle
Optimal trajectory
Figure I. OOD issue for offline trained UGV in the testing phase.
problems under uncertainty (Mnih et al., 2015; Silver et al., 2017: Schrittwieser al., 2020). Nevertheless, safety con- remain significant obstacle the extensive adoption of RL in safety-criticll domains (Garcıa & Fernández, 2015; Gu et al., 2022; Xu et al., 2022b; Li, 2023), such as indus- trial management, and robot control. these contexts, the potential for catastrophic outcomes necessitates a signifi- cant emphasis on preventing unsafe actions (Andersen et al., 2020; Brunke al., 2022). As promising method that received growing attention, safe RL provides safety guaran- tees by formulating the problem constrained Markov decision process (CMDP) (Altman, 1998; 2021).
Over the past few years, a multitude of safe RL algorithms have been introduced (Achiam et al., 2017: Tessler et al., 2018; Zhao al., 2021: Sootla et al., 2022; Yu al., 2022). Regrettably, most existing methods address safety concerns within online settings, relying the high-fidelity simulators or agent- eecentenvvrooment interactions during the training pro- cess, which introduces additional risks of safety violations (Liu et al., 2023a). Safe offline RL, on the other hand, pro- vides a promising solution that learns the safety-guarantedd policy fully offline manner. Its training requires no risky interaction with the environment and relies only on the pre- collected offline dataset. However, empirical observations indicate that most existing safe offline RI. algorithms fail to thoroughly ensure pre-defined safety constraints during test- ing (Liu et al., 2023a; Zheng et al., 2024). Such occurrences tend to be more pronounced in environments characterized by higher observation dimensions.
1. Introduction
Reinforcement learning (RL) has been proven to pow- erful tool for solving digh-dimensional eeciisionmmaingg
Deparrment of Mechanical and Automation Engineering, The Chinese University of Hong Kong, Hong Kong SAR, China. "Department of Electrical and Electronic Engineering, Hong Kong Polytechnic University, Hong Kong SAR, China. Correspondence to: Yunjian Xu <yjxu@ mac www..cuk..cu..kk..
Proceedings of the ndnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Q-Supervised Contrastive Representation: A State Decoupling Framework for Safe Offline Reinforcement Learning
Classical approaches for safe decision-making
Our state-decoupling framework for safe decision-making
Safe Policy z7 (s,)
a2
Env p(
Reward Policy [a(a((s))
St
R-dee-
s
Assessment Switch
St
sc+1
Tradeoff Poliey Tooo (ar sf)
a2
Env a,,
St
Conl
Reward Policy My, (s,
Safe
C-dee
st
Safety Assessment
Con2
Safety
Assessment Switch
a2
Eny p(52,115;. a,)
Cost Policy IT. [s,()
Con3
Unsafe
Safe Policy Te (e( (s+)
St 8c41
Bonunlectttdd decosple vin Rich geueral poneeieeeeeeeeeeee. Codec.
R-dac-
Cont: Absolute safe (V Vi""
St St+1
Con2: Bordertine safe (vi. 0 VH*)) Con3: Unsafe
Costerdatel decssple
(0 < v s VT*)
Figure 2. Overview diagram of classical approaches (left) and our proposed stale-decoupling framework (right) for safe decision-making.
In offline RL, it is imperative that the states visited during testing have been included in (or at least not far away from) the training dataset to ensure robust performance (Fujimoto et al., 2019; Wang et al., 2022). However, Safe offline RL problems have various combinations of reward-related and cost-related states. For instance, as illustrated in Figure I, a UGV (unmanned ground vehicle) needs to navigate around traps to reach its final destination. During testing, if the relative positions of traps and the target haven't 't occurred in the training dataset, the agent may struggle to make in- formed decisions based on such unknown observations. It is reasonable to suspect that the primary reason for the sub- par performance of safe offline RL during tests lies in the -itt---diisributiinn (OOD) issue.
Nevertheless, effective differentiation between reward- and cost-related information from global observations poses formidable challenge, especially when certain dimensions of the observations contain intertwined information. For instance, some information, like speed and accelcration of UGV, should be included in both reward- and cost cost-related representations. On the other hand, information from the enviromment-detectiin sensors, which include positions of destinations and obstacles, should be distinctly decoupled. Manual separation proves impractical in most cases.
Towards this end, our Q-supervised contrastive representa- tion decouples the global observations through clustering representations that demonstrate similar learned- across the actions in support The representations solely cap- ture either reward cost information, independent of an- other factor, as determined by the training of Q Unlike model-based representations learning (e.g., bisimulation), our SDQC circumvents model estimation, thus mitigating the challenges posed by severe estimation errors in scenar- ios with sparse rewards or costs. Moreover, we demonstrate that our representations can be trained concurrently with the Q*-Jearning process by incorporating an additional loss term within the framework of implicit Q-learning.
To tackle this problem and improve the generalization of safe offline RL, we propose State Decoupling with Q- supervised Contrastive representation (SDQC), novel rep- resentation learning method that decouples the global ob- servations into reward- and cost-related representations. At- tributable to the successful application of Hamilton-Jacohe (HJ) reachability analysis in Safe RL, which introduces safety analysis method iterated through Q-learning with convergence guarantees, our approach makes safety assess- ments on the cost-related representations and make deci- sions based on the assessment results. Unlike classical methods (as depicted in the left subplot of Figure 2), which rely on global observations for decision -aaiing our SDQC the first to utilize decoupled representations for decision- making in safe RL tasks (see the right subplot of Figure 2). It employs seward-rlatedd representations to make decisions when the assessment confirms absolute safety, switches to cost-related representations when the assessment deems the situation unsafe, and integrates both when the assessment indicates borderline safety.
Further, we provide theoretical evidence that our method produces a coarser representation compared to bisimulation, while still preserving the optimal policy. This is supported by our argument that SDQC leads to a higher information entropy of the global observations when conditioned on the representations. This attribute grants SDQC superior gen- cralization capabilities, bolstering its efficacy in handling OOD observations during the testing phase.
For simplicity, denote Q" as generic notaiton to repre- sent the optimal Q-value functions for both reward (Q;) and cost (Q%) in this paper. Similarly, Q refers to both Q, and Q+-
2

---

Q-Supervised Contrastive Representation: A State Decoupling Framework for Safe Offline Reinforcement Learning
The experimental results showcase that our SDQC outper- forms other safe offline RI. algorithms in the DSRL bench- mark, especially in its exceptional ability to achieve zero violations in the majority of tasks. Further, in generaliza- tion tests where agents are evaluated in environments that differ from the training ones, 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 slight decay in reward.
(Stooke et al., 2020). Furthermore, the joint optimization of V,", V.*, and ι leads to significant instability, minor approximation errors can bootstrap across them and prop- agate, thereby undermining the ability to provide robust safety guarantees (Kumar et al., 2019; Zheng et al., 2024).
Hamilton-Jacohh reachability. As a prospective method to perform safety assessment rooted in control theory, HJ reachability has been proven to be applicable in Safe RL tasks for both online settings (Chen et al., 202 Yu et al., 2022) and offline settings (Zheng et al., 2024). In addi- tion to the tuple formulation .M in CMDP, we introduce a constraint violation function h : S + I. which is pos- itive if the state constraint is violated and negative other- wise. For given state s, the safe value function V == maxten (h(s)) | So s,0; [T-ss,), Vi E {0, {0.....t)) +t} represents the worst constraint violations among all possible trajectories induced by policy т. The corresponding safe Q-function is given by Q% (s. a) == max,en(((s.) 80 = s, Gg = a, a, ~ (s)) Vi E fl..... f)). The optimal safe value function, defined as Vi == min, Vi (s), stands for the smallest violation one can obtain through adjusting the policy æ. Similarly, the corresponding optimal safe Q- function can be expressed as (s,aa) a) == ming Q% (7s a). With the discount factor defined as 7, Fisac et al. (2019) introduce the following safety Bellman operators:
2. Preliminaries
Safe Offline RL. Safe RL tasks are generally modeled as CMDP in the form of M = AS P. r, 7, where S is the state space, A is the action space, P is the model dynamics, r: A + IR represents the reward function, c: S x A + [0, Cmaz] represents the cost function, 7 E [o. 1) is the discount factor, and do E (S) is the distribution of initial state 80 (the set of the probability distribution over S is denoted as (S)). P(s (ds[|s,00 a): S S + ASS) represents the transition function from state s to s' when taking action a. The satte-aciion-rewardcooo transitions over trajectory recorded as 5 (st, at, F2, Ct)20 The goal of Safe RL is to learn a policy S + (AA) that maximizes the expec- tation of the cumulated discounted reward while restricting the expected cumulative costs below a predefined cost limit K, which can be denoted by max, EE iE (rr(s,, a;)l, s.t. E 10 cc(8,, a+)] S K.
Bi a) == (1 - ())(s) + max)))((s), (s)), Vi.,(s') min Qu(s', ,a),
(2)
In offline settings, the training is performed on statisti- cal dataset denoted D3- == (s, a, sr; c). This offline dataset comprises both safe and unsafe trajectories and is acquired from a behavior policy Ty. During training, most existing safe offline RL. algorithms utilize the temporal dif- ference (TD) method to learn the reward state-value func- tion Vy*(s)), (st), which models the expected cumulative reward Ey y'r(s,, d:)], as well as the cost state-value func- tion V. which models the expected cumulative cost E Σ t
which is a contraction mapping satisfying lim,-+1 Qhi + Qū, lim Vi., - V%- A direct safety inference can be made after training converges. V,, (s) S 0 implies the exis- tence of policies that guarantee adherence to the hard con- straints throughout the trajectory. Conversely, Vi(()
0 indicates that the destiny towards unsafe states regardless of the chosen policy. In the offline settings, Zheng et al. (2024) pioneered the application of HJ reachability analysis for safety assessment. They present that the decision decision-making making for Safe RL with hard constraints can be decoupled as:
00 7'c(s;, a:)]. The primal training objective of
safe offline RI can be expressed as follows:
max E, ~ Dr VV (st)
x
Safe : max, BaNn (Ai(s,a) (lig,a)-Tyggss [lg s.L. ((ugnnanngg m(als)da 1; D( Se
s.t. E,,,
k; ((r(as)) <T
where (p) is the divergence term that prevents the distributional shift in offline training.
Unsafe : :max, Bann -A, Iyy (yy
s.t. Sa (als)dd - 1; D(TTTT)) Se
A commonly employed approach for solving Eq. (1) in- volves reformulating the training objective using the La- grangian dual form as minx20 max, ED, (V7 (sx) - X(V - k)], s.t. ((T(TT) < Ex, where the learnable Lagrange multiplier À is iteratively updated to enforce the constraint. However, the Lagrangian approach can be sensi- tive to the learning rate and initialization of the multiplier
(3)
where a, == (;(s (s) and (s,a) = Q (8, a) (s). Eq. (3) theoretically ensures zero con- straint violations. However, challenges arise from estima- tion errors and OOD problems during the testing phase. As a result, the empirical results demonstrate the inability of their algorithm (FISOR) to achieve absolute safety guarantees.
3

---

Q-Supervised Contrastive Representation: A State Decoupling Framework for Safe Offline Reinforcement Learning
3. State Decoupling with Q-supervised Contrastive Representations
icy (cf. Appendix B.2,C.1 for details). This allows to generate in-support actions for any given states in the offline dataset, denoted as A As result, we have the approxima- tion d(s, 3) = su [Q* - Q (z)(5), a) for calculating the soft similarity measure.
In our state-decoupling framework, we aim to decouple the state into two separate representations: one related to rewards, denoted as Sy, and the other related to costs, denoted as Sc- We slightly abuse the notation (0(s) (or simplified as z) to represent the neural network embedded representations of either s, or Se in this section.
In practice, we employ random sampling approach to se- lect a subset of states, denoted as S", from the offline state set. Within the subset, we further randomly choose a set of anchor states, denoted as {s, E S' | I}, where I represents the index set of the selected anchor states. For each anchor state 8;, we use its nearest neighbor in S' based on the similarity measure I to define the posi- tive pairs &.,, where 5; arg min (xx:) [(st, s). The remaining states in S' are considered as negative sam- ples. Attention- or multiple- perepppt based neural networks are utilized to encode the state as normal- ized vector on the unit hypersphere, i.c., [s]] = 1 (cf. Appendix C.2 for detailed network selection and structure design). Finally,
3.1. Motivation
Manually abstracting representations of reward-related or cost-related aspects directly from original state observations can be challenging due to the entanglement of information within certain observation dimensions. It is observed that the optimal Q-value, whether associated with reward or cost, exclusively encompasses the information it was trained with, independent of another factor. For instance, concerning the states the agent depicted in Figure I, the optimal Q-values related to reward are the same across all actions regardless of the cost-related observations. These states should be embedded as the same reward-related representation. To achieve coarser abstraction while maintaining the optimal Q-value unchanged, we design the objective for both reward- and cost-related representations as follows
have the following contrastive loss, which encourages the embedding of states with similar Q" values across all actions to have similar representations:
POS,
Lo
log POS; + NEG.
EE
POS: = [(s., 3,) explzi 2c/v)
max, (s)) s.t. ((==Q(=(ss), - ((zo(s). s EB, Va E A,
NEG; E Esjezn (1 - -(=, ex))
(4)
(5)
where is temperature parameter. 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 small changes in the net- work can result in variations in the Q-values. Therefore, it is necessary to integrate the training process of the representa- tion with the Q-learning process. Such coupling ensures that both the representation and Q-values are jointly updated, accounting for the interdependencies between them.
where ([---) represents conditional entropy (Shannon, 1948), eg is an arbitrary small number and E is the optimal general/safety Bellman operator. We define d(s1, 82) == Sup.EAIl' IQ" (20 (21)))),)) a), Q" (2o a)] as the distance mea- sure between pair of states 81, 82 E S. One can always find an arbitrarily small number ea such that the objective in Eq. (4) can be achieved through embedding the states C(s')= = (5 E S | d(), < {2) with the same repre- sentation for any s' E S', where S' is a smallest subset of S such that for any S s/ E S',, we have d(s{,,%)) 2 2cd and Urrss C(s')= =. S. Contrastive learning, which aims to bring artificially defined similar instances closer and push other instances further apart in the representation space (Oord et al., 2018; Bachman al., 2019; Chen al., 2020), provides a promising solution for our embedding task.
3.3. Practical Implementation
Building upon in-sample learning methods (Kostrikov et al., 2021; Xu et al., 2023; Garg et al., 2023; Zheng et al., 2024), our approach follows a two-step process. In the initial phase, we undertake the learning process for the value functions and representations associated with cost and reward sepa- rately. Following that, we extract the policy based on the acquired value functions and representations.
3.2. Q-supervised Contrastive Representation
Inspired by Agarwal et al. (2021), we adopt a soft simi- larity measure, denoted as [(ss $) exp(-d(s 8)/n), to quantify the distance between two states (with 1) represent- ing the temperature factor). Notably, directly calculating the distance measure involves querying out-of-distribution (OOD) actions in the offline setting. To address this issue, we pre-train a generative model to capture the behavior pol-
Reward-related Representation. We use implicit Q- learning (IQL) (Kostrikov et al., 2021) (cf. Appendix B. for details) to approximate the reward-related optimal value functions Q and Vi- within the support of data distribution through upper expectile regression:
4

---

Q-Supervised Contrastive Representation: A State Decoupling Framework for Safe Offline Reinforcement Learning
where var represents the variables involved in the expecta- tion, with t ~ (1, ~ (0 T), and (s,a) D3. The noised action d; = aga+ ore satisfies the forward transition distribution N (arla,, o,I) in the diffusion models, and at, of are noised schedules. The weights in Eq. (13) are
Lv,= L ()D (),,a.""")r) (rw. (za. a))"] (7)
(s)))] (6)
where L (u) = |T I[(u < 0)(u², + E (0.5, 1). The reward-relate representations are learned with the addition of a contrastive loss term (Eq. (5)), with weighting factor of 6. Consequently, the overall loss for the reward-related value functions and representations is formulated as:
wr (2o, wr(20, (5), a) exp((x(Q)(E0, (s), V=(20, we (2,, (6). = expl -a (Qn (20, (8). - Vi, (3))) Who (20 (5), zR, (5), = explesa
Va (s))) (2 (o))))
(14)
Lyeward = Lv, + Lq. + +Lo..
(8)
where 4r 4/ and 4to are temperatures that control the behav- ior regularization strength.
Cost-related Representation. Similar to Zheng et al. (2024), we employ the safety Bellman operator, as denoted in Eq. (2), and utilize lower expectile regression to learn the cost-related optimal value functions Q7 and Vi::
After obtaining Sv- Sv, and Cu... the three approximated optimal policies can be sampled through the reverse diffu- sion chain starting from random Gaussian noise (Ho et al., 2020; Song et al., 2020) (cf. Appendix B.3 for details). During the testing phase, we first perform safety assess- ments on the cost-related representations. If the assessment verifies absolute safety (V/lw S V,, I 0), we employ the policy T.. If the assessment indicates borderline safety (V2ow < 0 < Vi"",, we utilize the policy Nio- In the case of an unsafe condition (0 < V/low < Vi we rely on the policy Te. See the right subplot of Figure 2 for details.
(Qa(s,, (5), 1/l (z0, (o))) (9)
-)) ((1)))))) maxth, Val (20,, (5')
(10)
- -Q (s). {)
where (w) = - 1- I(u 0)/u², - E (0.5, 1). Addition- ally, we learn an upper-bound cost-related value function h to model the maximum Q7, across all actions in support:
3.4. Comparison with Bisimulation
Bisimulation has been established as a useful tool for abstracting state representations (Definition 3.1), where states with identical transition and reward/cost functions are grouped together (Givan et al., 2003; Castro & Pre- cup, 2010; Castro, 2020; Castro al., 2021). However, employing bisimulation typically entails an additional step of training model-based estimator to learn the state tran- sition and reward/cost functions. Notably, the estimation of reward/cost functions becomes particularly challenging when the values are sparsely distributed (Lee et al., 2024). In contrast to such a madel- representation approach, learning the representations based on Q* (Definition 3.2) eliminates the necessity for estimating the exact model dy- namics (Givan al., 2003; Li et al., 2006). With θ de- noting a generic surjective mapping from the ground-truth state space to representation space, we have the following definitions:
((o==Be,))D (z0, (0), -1, (s)))] (11)
By incorporating an additional contrastive loss term (Eq. (5)) with a weighting factor of 6, we express the overall loss for the cost-related value functions and representations as:
Lenst Lyyow + Lyy + +Qn 5.a.,
(12)
Policy Extraction. As illustrated in the right subplot of Figure 2, we divide the global policy into three components: the reward policy Ty, which solely depends on the reward- related representation; the cost policy Th which solely relies on the cost-related representations; and the tradeoff policy Tto which depends on both. We independently train the three policies using weighted regressed diffusion models, an approach pioneered by Zheng et al. (2024). They present that the optimal policy satisfies n (a]=) x (a|=)) w0( a), and the optimal policy can be obtained through weighted training of diffusion models. The weighted loss function for the three policies can be expressed as follows:
Definition 3.1. A bisimulation representation Opisim is such that for any action a and any represented state 2, Obisim ($1) (bssi((s2) implies r($1, a) = r(s2, a)
(or c(s1,a) = c(s2,a)) and Exe (z) P(s'(s1,
Ever S (=) P(d||82,0)) a).
Trr : Lx,, -Bar (20- (s), a). Gur 20, (x), (]]
Th : LTT, Evar |wn(2o, (z0, (s),
Cah (ae, 20, (8),
Tio : Lai. = Ever (Wro [2a (z_ (5): z0, (s), a).
Definition 3.2. A -rreeeeeven representation Og- is such that, θq (s1) = θq (s2) implies Q"(s1,a) = Q* 2"(82, a) for any action a.
115 Sen (ac, 20, (s), z0,, (s), O]]
(13)
5

---

Q-Supervised Contrastive Representation: A State Decoupling Framework for Safe Offline Reinforcement Learning
Table 1. Normalized DSRL benchmark results. The evaluation results averaged over 3 random seeds (20 episodes for each seed).
Gray; Unsafe agents. Bold: Safe agents whose normalized cost is smaller than 1. Red: Safe agents with the highest reward. Blue: Safe agents with the lowest cost.
BCQ-Lay
CPQ
COPiDDIE
CDT
TREBI
FISOR
SDQC(ourS)
Task
reward"
cost,
reward
cost
reward?
cost.
reward
cost.
reward]
cust]
reward
cust
reward"
cost.
"PainGGoall
0.71
4.29
0.56
0.93
0.40
5.53
0.21
1.59
0.36
2.79
0.68
4.19
0.35
0.36
"PointGoal2
3.81
0.41
5.03
0.43
2.78
0.22
1.19
0.28
3.86
0.21
0.09
"PuintPushl
0.32
3.08
0.14
1.35
0.13
3.80
0.27
2.81
0.31
2.02
0.27
1.38
0.12
0.00
"PuintPush2
0.21
1.86
0.16
2.36
0.02
2,90
0.18
1.69
0.13
3.85
0.24
2.41
0.19
0.28
"PointButton1
0.21
4.45
0.61
11.80
0.08
4.29
0.48
10.88
0.12
4.55
0.04
0.97
0.08
0.46
PPointButoon2
0.38
8.04
0.35
12.09
0.17
6.12
0.42
9.97
0.02
2.18
0.08
4.49
0.06
0.57
CarGoall
0.44
2.76
0.33
4.93
0.43
2.81
0.60
3.15
0.41
1.16
0,49
0.83
0.38
0.01
CarGoal2
0.34
4.72
0.10
6.31
0.19
2.83
0.45
6.05
0.13
1.16
0.06
0.33
0.23
0.00
CarPush
0.23
1.33
0.08
0.77
0.21
1.28
0.27
2.12
0.26
1.03
0.28
0.28
0.30
0.00
CarPush2
0.10
2.78
-0.03
10.00
0.10
4.55
0.16
4.60
0.12
2.65
0.14
0.89
0.31
0.04
CarBattonl
0.13
6,68
0.22
40,06
-0.16
4.63
0.17
7.05
0.07
3.75
-0,02
0.26
0.03
0.32
CarBulion2
0.04
4.43
0.08
19.03
-0.77
3.40
0.23
12.87
-0.03
0.97
0.01
0.58
0.02
0.42
Ant
18.54
-1.01
1.00
10.29
0.91
0.31
0.00
0.89
0.00
HallCheetaldel
1.04
57.06
0.08
2.56
0.43
0.00
0.97
0,55
0.87
0.23
0.89
0.00
0.81
0.00
SwimmerVel
0.29
4.10
0.31
11.58
0.58
23.64
0.67
1.47
0,42
1.31
-0.04
0.00
-0.04
0.00
SafetyGym
0.39
9.17
0.16
9.05
0.26
5,66
0.46
4.93
0.25
2.10
0.28
1.25
0.26
0.17
Average
AniRun
0.65
3.30
0.00
0.00
0.62
3.64
0.70
1.88
0.63
5.43
0.45
0.03
0.31
0.00
BallRun
0.43
6.25
0.85
13.67
0.55
11.32
0.32
0.45
0.29
4.24
0.18
0.00
0.20
0.00
CarRun
0.84
2.51
1.06
10.49
0.92
0.00
0.99
1.10
0.97
1.01
0.73
0.14
0.56
0.00
DroneRun
17.98
0.02
7.95
0.72
13:77
0.30
0.59
1.41
0.30
0.36
0.56
AntCircle
0.67
19.13
0.00
0.00
0.18
13.41
0.48
7.44
0.37
2.50
0.20
0.04
0.38
0.00
BallCircle
0.67
8.50
0.40
4.37
0.70
9.06
0.68
2,10
0.63
1.89
0.34
0.00
0.42
0.00
CarCirele
8.84
0.49
4.48
0.44
7.73
0.71
2.19
0.49
0.73
DroneCircle
0.95
18.56
-0.27
1.29
0.24
2.19
0.55
1.29
0.54
2.36
0.48
0.00
0.36
0.07
BulletGym
10.63
0.32
0.55
7.64
0.63
2.09
0.56
2.44
0.39
0.08
Average
For any state representation O,, 02, we say θ is finer than O2, denoted Q 20,, if and only if for any states (s2) implies 02(s1) = O₂ (s2).
We compare the proposed SDQC with several state-of- the-art baseline safe offfine RL algorithms the DSRL benchmark(Liu et al., 2023a), which provides extensive datasets and environment wrappers for safe offline RL per- formance evaluation. Evaluation results" are presented in Table 1. The baseline algorithms include i) BCQ-Lag: A PID-Lagrangian-base based method (Stooke et al., 2020) that con- siders cost threshold based on Batch Constrained Q-learning (BCQ) (Fujimoto al., 2019), ii) CPQ (Xu et al., 2022a): A constrained Q-updating method that incorporates penalties for OOD actions and unsafe actions, iii) COptiDICE (Lee et al., 2022): A DICE (distribution correction estimation) based Lagrangian method that builds upon OptiDICE (Lee et al., 2021), iv) CDT (Liu et al., 2023b): A future cost in- ference method based on Decision Transformer (DT) (Chen et al., 2021b), v) TREBI (Lin et al., 2023): A real-time cost budget inference method on the basis of Diffuser (Janner et al., 2022), vi) FISOR (Zheng et al., 2024): An HJ reach- ability guided method with diffusion policies that firstly considers the hard constraints in safe offline RL. problems.
$1, $2 E S, (s1) Givan et al. (2003) established the relationship between bisi- umulation and the Q- -irrelevance representations (Ohisim > @Q.) ) for finite-horizon MDPs with respect to the general Bellman operator. In the following theorem, we extend this relationship to
MDPs and incorporate the safety Bellman operator, as described below.
Theorem 3.3. For any MDP, the optimal Q induced by ei- ther the general Bellman operator or the safety Bellman op- erator satisfy Obisim ž θq. The optimal policies derived from both bisimulation representation and Q" -irrelevance representation are also optimal in the ground MDP.
Theorem 3.3 shows that neither of the representations alters the optimal policy, while the bisimulation representation is finer than the Q-based representation, which implies that
0 H (s (s)) S H(s (s)).
(15)
Our ultimate objective to achieve zero-cost during test, aligning with the framework established by FISOR (Zheng et al., 2024). However, most baseline algorithms struggle to operate effectively under a zero-cost threshold. Conse- quently, following FISOR (Zheng et al., 2024), we impose a
Refer to Appendix A for detailed proof. Since our objec- tive is to maximize the conditional entropy H (s)) our Q-supervised contrastive learning method theoretically sur- passes bisimulation in terms of generalization.
*The baseline algorithm evaluation results are sourced from FISOR (Zheng et al., 2024), except for the evaluation of the Point agent on Gymmasiummm (marked with *), which is conducted independently as it is not in the source.
4. Experiment
4.1. Evaluation on DSRL benchmark
6

---

Q-Supervised Contrastive Representation: A State Decoupling Framework for Safe Offline Reinforcement Learning
(a)
Reward
Cost of the policy S strrindd S tested
(b)
Reward
Cost of the policy C-trained C-tester Cost of the policy C-trained S-tested
Reward
Cost of the policy S-trained C-tested
Reward
sufety-cymmassuu---arooo
Safety-Gymnnsium-caruuu
sefettycyymassuum-cardaalll
Safety-Gymmssium-Caruuhhh
0.8
Normalizdd Reward
Normalized Reward
0.6
0.6
0.6
0.8
0.4
0.4
0.4
0.4
0.2
0.0
0.0
0.0
0.0
-0.2
1,000
copuaitor #soM sooc
-0.2
EISOn sooc
-0.2
coppolitor AsoR spac
-0.2
u20 copuokeer roso
sooc
5.0
5.0
5.0
5.0
Normalized Cost
4.0
4.0
Normalized Cost
4.0
4.0
3.0
3.0
3.0
3.0
2.0
2.0
2.0
1.0
1.0
1.0
1.0
800,00000 contibi
Eol
eISOR
sogc
ComtuM Eo1
nIS08
5000
800 800.04000 L
Coppotic toa
FSDR
5p0c
L P0
COPDLAE ED1
PISOR
500c
(a) The agent is trained the dataset from a simple environment (S -rtrinnddd, and its performance is evaluated in both the simple environment (S-tested) and the complicated one (C-tested). Conversely, (b) the agent trained on the dataset from a complicated environment (C-trained), and its performance is assessed in both the original environment (C-tested) and the simple one (S-tested). The evaluation results are obtained from 3 random seeds, with 20 tests on each seed. Outlier data points are omitted for clarity.
Figure 3. The generalization tests the agent "Car" in Bafety-Gymnassin
stringent cost limit of 10 for the ment and 5 for Sullet-Safety-Gyni We employ the metrics of normalized return and normalized cost for evaluation, where a normalized cost below I signifies a safe operation.
environ-
while avoiding hazardous areas and obstacles. The difficulty level varies between tasks, with the simple tasks (CarGoall, CarPushI) I) having fewer hazards and obstacles than the chal- lenging tasks (CarGoal2, CarPush2).
The former five baseline algorithms exhibit either significant constraint violations or suboptimal returns when subjected to stringent safety requirements, partially due to the fact that they consider only soft constraints. Despite incorporating hard constraints, FISOR still encounters high costs in tasks with high complexity. As discussed in Section 1, this issue can be attributed to estimation errors and OOD problems during the testing phase. In contrast to FISOR, our proposed SDQC conducts safety assessments the cost-related rep- resentation abstracted from the original observations and makes decisions accordingly. The utilization of decoupled representations in SDQC substantially improves the accu- racy of state safety assessment and enhances the general- ization capability of the policy, thereby providing a higher level of safety assurance. The experimental results clearly demonstrate that SDQC outperforms FISOR in terms of higher rewards and lower costs. Remarkably, SDQC even achieves zero violations in the majority of tasks.
It is reasonable to be concerned about the performance of an agent when it is tested in environments that differ from the ones it was trained on, especially if the testing environ- ment is more complex or comprehensive. The experimental results provide evidence that our proposed SDQC algorithm is the only algorithm that ensures no increase in cost under such circumstances. In fact, SDQC achieves almost zero violations in the majority of tests, with only a slight decay in reward performance. In contrast, other algorithms exhibit sharp increase in cost and/or a significant decrease in reward. Generalization in ensuring safety is crucial in safety-critical scenarios like autonomous driving. It is impractical for the agent to traverse every possible radar observation that may arise in real-worke scenarios during training. Our proposed SDQC offers potent and promising solution for addressing these complex safety-critical scenarios.
4.3. Ablation Study
To validate the efficacy of our proposed Q-supervised con- trastive learning approach in acquiring meaningful repre- sentations and enhancing performance, as discussed in Sec- tion 3.2, we conducted ablation studies on the "Safety- Gymm CarGoal2" task (cf. Figure 4). In the absence of contrastive loss during the critic and representation training phase, the agent experiences considerably lower rewards and higher costs compared to the agent trained with contrastive loss.
4.2. Generalization Tests
To showcase the superior generalization capabilities of our proposed SDQC compared to other safe offline RL algo- rithms, we perform generalization tests on the "CarGoal" and "CarPush" tasks (in in Figure In these tasks, the "Car" agent is tasked with reaching the goal point or pushing the box the goal point
as illustrated

---

Q-Supervised Contrastive Representation: A State Decoupling Framework for Safe Offline Reinforcement Learning
(a)
Reward
Cost of [+)
(b)
Original States
Reward Representations()
Reward Representations(x)
Rewardd
Cost (K)
100
Normalized Reward
0.3
80
80
0.1
0.0
6.0
300
500
Original States
Cost Representations( *
Cost Representations/x x]
Now
0.0
Normalized Cost
4.0
-0.2
-0.4
2.0
-0.6
-0.8
100 200 300 400 500
-1.0
Gradient steps (x10P)
Figure 4. Ablation studies on the Q-supervised contrastive loss in CarGoal2. (a) The actor-training process evaluations of SDQC with (marked by *) and without (marked by x) contrastive loss. The curves are averaged over 3 random seeds and smoothed with a window size of 3. (b) t-SNE visualization of the distribution of the original state, the reward-related and the cost-related representations with and without contrastive loss across 10 different safe trajectories, where the policies are from the agent trained with contrastive loss. The original states (first column) are colored according to the critic trained with contrastive loss (i.e., the same as the second column).
The t-SNE visualization results (in Figure 4b) reveal that the Q-supervised contrastive loss effectively clusters repre- sentations with similar values. This aligns with our original intention, which is cluster states with similar Q-values for any actions in the representation space. The clusterings facilitate the learning of the conditional diffusion model (i.e., the actor) by promoting the generation of similar out- put policies for similar representations. Furthermore, the inclusion of the Q-supervised contrastive loss enables a more reliable evaluation of the states' safety. In the depicted 10 trajectories, despite the cumulative cost being zero, the agent trained without contrastive loss erroneously identifies majority of the experienced states as unsafe (V/ow 2 0). Conversely, the agent trained with contrastive loss provides a more accurate assessment, demonstrating the effective- ness the proposed approach. For more ablation studies on the impact of anchor number ([Z)) and neural network structures, please refer Appendix C.
mistic/pessimistic planning problem (Wachi et al., 2018; Kalagarla et al., 2021). ii) Combining the safe RL problem with the field of safe control. A prevalent method entails representing a safety certificate through a learned safety assessment function, such as the Control Barrier Function (CBF) (Ma et al., 2021a; Luo & Ma, 2021) or Hamilton- Jacobbi (HJ) reachability (Yu et al., 2022; Fisac et al., 2019; Chen et al., 2021a). An agent can switch between optimal and safe policies based on safety assessment results (Chen et al., 2021a; Thananjeyan al., 2021), thereby theoretically ensuring hard constraints with state-wise zero violations. Recent research endeavors have embraced the integration of safety constraint problems with existing reliable offline reinforcement learning algorithms (Xu et al., 2022a; Lee et al., 2022; Liu et al., 2023b; Lin et al., 2023). However, most existing methods only provide soft constraints without any guarantees of zero violations. FISOR (Zheng al., 2024) is the first safe offline RI algorithm that tackles hard constraints issues, while the limited offline data still makes it difficult to guarantee safety during tests thoroughly. As complementary algorithm to FISOR, our SDQC decouples the global observations for safe decision-making substan- tially improves the accuracy of state safety assessment, and enhances the generalization capability of the policy, thereby providing a higher level of safety assurance.
5. Related works
Safe RL. In online settings, safe RL problems are gen- erally tackled with two mainstream approaches (Xu et al., 2022b). i) Formulating the problem as a CMDP and solv- ing it from an optimization perspective. Solution tech- niques include updating the policy constrained in trust region (Achiam et al., 2017; Liu et al., 2022), reformulating the problem into its Lagrangian dual form (Tessler et al., 2018; Chow et al., 2018; Ma et al., 2021b; Duan et al., 2022), and addressing the constraints by framing an opti-
Representation Learning. Representation learning in RL involves compressing the large observation space into a smaller latent vector that captures relevant aspects of the environment (Watter et al., 2015; Finn et al., 2016; Gelada et al., 2019), often applied in image-based tasks (Kostrikov
8

---

Q-Supervised Contrastive Representation: A State Decoupling Framework for Safe Offline Reinforcement Learning
al., 2020; Yarats et al., 2021; Cetin et al., 2022). Con- trastive learning has been widely acknowledged as po- tent technique for unsupervised representation learning (Liu et al., 2021; Zhu et al., 2022), primarily achieved by aug- menting data through introducing noise to the original image (Laskin et al., 2020; Agarwal et al., 2021). In state-based tasks, this approach is not directly applicable as the noise may distort the underlying information. Unlike previous works that conduct contrastive learning among the gener- ated samples, we employ contrastive learning within the dataset itself in Q-supervised manner. The most relevant works to ours are from Bellemare et al. (2019) and Le Lan et al. (2021), who learn representations via Bellman value functions. To the best of our knowledge, we are pioneers in utilizing representation learning in state-based Safe RL tasks. We are the first to introduce the concept of decou- pling states into reward- and cost-related representations specifically for decision-making purposes.
and reliability of RL applications in real-world scenarios, such as autonomous driving, robot control, and industrial management, where safe decision-making holds exceptional significance.
We would like to emphasize that our SDQC does not have any negative societal impacts. Our proposed approach is benign, free from potential malicious or unintended uses, and raises no concerns regarding fairness or privacy.
References
Achiam, J., Held, D., Tamar, A., and Abbeel, P. Constrained policy optimization. In International Conference on Ma- chine Learning, pp. 22-31. PMLR, 2017.
Agarwal, R., Machado, M. C., Castro, P. S., and Bellemare, M. G. Contrastive behavioral similarity embeddings for generalization in reinforcement learning. arXiv preprint arXiv:2101.05265 2021.
6. Conclusion
Altman, E. Constrained Markov decision processes with total cost criteria: Lagrangian approach and dual linear program. Mathematical Methods of Operations Research, 48:387-417, 1998.
In this work, we propose the first framework of state decou- pling for safe decision-making to tackle the OOD problem of offline safe RI. during the testing phase. We propose Q-supervised contrastive learning method to learn the repre- sentations without relying on additional system model esti- mation such as bisimulation. Theoretical analysis demon- strates that our Q-supervised approach generates coarser representations while preserving the optimal policy, leading to enhanced generalization performance. Experiments on DSRL benchmarks showcase that SDQC surpasses other baseline algorithms, especially for its exceptional ability to achieve almost zero violations in more than half of tasks. Further, SDQC possesses superior generalization ability when confronted with unseen environments.
Altman, E. Constrained Markov decision processes. Rout- ledge, 2021.
Andersen, P.-A., Goodwin, M., and Granmo, O.-C. To- wards safe reinforcementllemming in industrial grid- warehousing. Information Sciences, 537:467-484, 2020.
Bachman, P., Hjelm, R. D., and Buchwalter, W. Learn- ing representations by maximizing mutual information across views. Advances in Neural Information Processing Systems, 32, 2019.
Bellemare, M.. Dabney, W., Dadashi, R., Ali Taiga, A., Castro, P. S., Le Roux, N., Schuurmans, D., Lattimore, T., and Lyle, C. A geometric perspective optimal representations for reinforcement learning. Advances in Neural Information Processing Systems, 32, 2019.
Acknowledgments
This work was supported in part by the General Research Fund (GRF) project 14200720 of the Hong Kong University Grants Committee and the National Natural Science Foun- dation of China (NSFC) Project 62073273. The authors would like to thank the anonymous reviewers for valuable discussion.
Brunke, L., Greeff, M., Hall, A. W., Yuan, Z., Zhou, S., Panerati, J., and Schoellig, A. P. Safe learning in robotics:
From learning- ment learning. Annual Review of Control, Robotics, and Autonomous Systems, 5:411 14144444, 2022.
control to safe reinforce-
Impact Statement
Castro, and Precup, D. Using bisimulation for policy transfer in MDPs. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 24, pp. 1065-1070, 2010.
The impact of our work lies in its potential to develop trust- worthy decision-making systems and advance the field of safe offline RI. Focusing on the hard constraint problems, our SDQC achieves almost zero violations in the DSRL benchmark tasks and possesses great generalization capa- bility to unfamiliar global observations. Its superior per- formance and adaptability significantly enhance the safety
Castro, P. S. Scalable methods for computing state similarity in deterministic Markov decision processes. In Proceed- ings of the AAAI Conference on Artificial Intelligence, volume 34, pp. 10069-10076, 2020.

---

Q-Supervised Contrastive Representation: A State Decoupling Framework for Safe Offline Reinforcement Learning
Castro, P. S., Kastner, T., Panangaden, P., and Rowland, M. MICo: Improved representations via sampling-base state similarity for Markov decision processes. Advances in Neural Information Processing Systems, 34:301 13- 30126, 2021.
Gelada, C., Kumar, S., Buckman, J., Nachum, O., and Belle-
mare, M. G. DeepMDD: Learning continuous latent space models for representation learning, In International Con- ference on Machine Learning. pp. 2170-2179. PMLR, 2019.
Cetin, E., Ball, P. J., Roberts, S., and Celiktutan, O. Stabi- lizing off-policc deep reinforcement learning from pixels. arXiv preprint arXiv:2207.00 000966 2022.
Givan, R., Dean, T., and Greig, M. Equivalence notions and model minimization in Markov decision processes. Artificial Intelligence,
2003.
Chen, B., Francis, J., Oh, J., Nyberg, E., and Herbert, S. L.. Safe autonomous racing via approximate reachability on ego-vision, arXiv preprint arXiv:21 2021a.
Gronauer, S. Bullet-safety- gym: A framework for con- strained reinforcement learning. 2022.
Gu, S., Yang. L., Du, Y., Chen, G., Walter, F., Wang. J., Yang, Y., and Knoll, A. A review of safe reinforce- ment learning: Methods, theory and applications. arXiv preprint arXiv:2205. 10330, 2022.
Chen, L., Lu, K., Rajeswaran, A., Lee, K., Grover, A., Laskin, M., Abbeel, P., Srinivas, A., and Mordatch, I. De- cision transformer: Reinforcement learning via sequence modeling. Advances in Neural Information Processing Systems, 14:15084-15099 2021b.
Ho, J., Jain, A., and Abbeel, Denoising diffusion proba- bilistic models. Advances in Neural Information Process- ing Systems, 33:6840-653 6851, 2020.
Chen, T., Kornblith, S., Norouzi, M., and Hinton, G. A simple framework for contrastive learning of visual rep- resentations. In International Conference on Machine Learning, pp. 1597-1607. PMLR, 2020.
Janner, M., Du, Y.. Tenenbaum, J. B., and Levine, S. Plan- ning with diffusion for flexible behavior synthesis. arXiv preprint arXiv:2205.099 29991, 2022.
Chow, Y., Ghavamzadeh, M., Janson, L., and Pavone, M. Risk-constrained reinforeement learning with percentile risk criteria. Journal of Machine Learning Research, (167):1-51, 2018.
Kalagarla, K. C.,, Jain, R., and Nuzzo, P. A sample-efficient algorithm for episodic knite-horizon MDP with con- straints. In Proceedings of the AAAI Conference on Artifi- cial Intelligence, volume 35, pp. 8030-8037, 2021.
Duan, J., Liu, Z., Li, E., Sun, Q., Jia, Z., and Cheng, B. Adaptive dynamic programming for nonaffine nonlinear optimal control problem with state constraints. Neuro- computing, 484:122 128-141, 2022.
Kostrikov, I., Yarats, D., and Fergus, R. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. arXiv preprint arXiv: 2002200044.644 2020.
Finn, C., Tan, X. Y., Duan, Y., Darrell, T., Levine, S., and Abbeel, P. Deep spatial autoencoders for visuomotor learning. In 2016 IEEE International Conference on Robotics and Automation (ICRA), pp. 512-519. IEEE, 2016.
Kostrikov, I,, Nair, A., and Levine, S. Offline reinforce- ment learning with implicit Q-learning. arXiv preprint arXiv:2 10,0616,,, 2021.
Kumar, A., Fu, J., Soh, M., Tucker, G.. and Levine, S. Stabilizing off- aff-policy Q-learning via bootstrapping error reduction. Advances in Neural Information Processing Systems, 32, 2019.
Fisac, J. F., Lugovoy, N. F., Rubies-Royo, V., Ghosh, S., and Tomlin, C. Bridging Hamilton- safety analysis and reinforcement learning. 2019 International Confer- ence on Robotics and Automation (ICRA), pp. 8550-8556. IEEE, 2019.
Laskin, M., Srinivas, A., and Abbeel, CURL: Contrastive unsupervised representations for reinforcement learning. In International Conference on Machine Learning, pp. 5639-5650. PMLR, 2020.
Fujimoto, S., Meger, D., and Precup, D. Off Off-policy deep reinforcement learning without exploration. In Interna- tional Conference on Machine Learning, pp. 2052-200 2062. PMLR, 2019.
Le Lan, C., Bellemare, M. G., and Castro, P. S. Metries and continuity in reinforcement learning- In Proceed- ings of the AAAI Conference on Artificial Intelligence, volume 35, pp. 8261-8269, 2021.
Garcia, J. and Fernández, F. A comprehensive survey on safe reinforcement learning. Journal of Machine Learning Research, 163)::4477-,48, 2015.
Lee, J., Jeon, W., Lee, B., Pineau, J., and Kim, K.-E. Op- tiDICE: Offline policy optimization via stationary distri- bution correction estimation. In International Conference Machine Learning, pp. 6 6120-6130. PMLR, 2021.
Garg, D., Hejna, J., Geist, M., and Ermon, S. Extreme Q-learning: Maxent RL without entropy. arXiv preprint arXiv:2301.02.22 2023.
10

---

Suppervised Contrastive Representation: A State Decoupling Framework for Safe Offline Reinforcement Learning
Lee, J., Paduraru, C., Mankowitz, D. J., Heess, N., Precup, D., Kim, K.-F., and Guez, A. COptiDICE: Offline con- strained reinforcement learning via stationary distribution correction estimation. arXiv preprint arXiv:2204.0895. 2022.
Luo, Y. and Ma, T. Learning barrier certificates: Towards safe reinforcement learning with zero training-time vi- olations. Advances in Neural Information Processing Systems, 14:25621-256333 2021.
Ma. H., Chen, J., Eben, S., Lin, Z., Guan, Y., Ren, Y., and Zheng, S. Model based constrained reinforcement learning using generalized control barrier function. In 2021 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS), pp. 4552 4559, IEEE, 2021a.
Lee, V., Abbeel, P., and Lee, Y. DreamSmooth: Improving model-basee reinforcement learning via reward smooth- ing. In The Twelfth International Conference on Learning Representations, 2024.
Li, L., Walsh, T. J., and Littman, M. L. Towards unified theory of state abstraction for MDPs. AI&M, 1(2):, 2006.
Ma, H., Guan, Y., Li, S. E., Zhang, X., Zheng, S., and Chen, J. Feasible actor-critic: Constrained reinforcement learning for ensuring statewise safety, arXiv preprint arXiv:2105.10682. 2021b.
Li, S. E. Reinforcement learning for sequential decision and optimal control. Springer, 2023.
Mnih, V., Kavukcuoglu, K., Silver, D., Rusu, A. A., Veness, J., Bellemare, M. G., Graves, A., Riedmiller, M., Fidje- land, A. K., Ostrovski, G., et al. Human-level control through deep reinforcement learning. Nature, 518(7540): 529-533, 2015.
Lin, Q. Tang, B., Wu, Z., Yu, C., Mao, S., Xie, Q., Wang, X., and Wang, D. Safe offline reinforcement learning with real-time budget constraints. In International Conference on Machine Learning, pp. 21127-21152. PMLR, 2023.
Oord, A. v. d., Li, Y., and Vinyals, O. Representation learn- ing with contrastive predictive coding. arXiv preprint arXiv: 1807.03778.. 2018
Liu, G., Zhang, C., Zhao, I.., Qin, T., Zhu, J., Li, J., Yu, N., and Liu, T.-Y. Return- based contrastive representa- tion learning for reinforcement learning. arXiv preprint arXiv:2102. 2021.
Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., Antiga, L., al. Pytorch: An imperative style, high-performance deep learning library. Advances in Neural Information Processing Systems, 32, 2019.
Liu, Z., Cen, Z.,, Isenbaev, V., Liu, W., Wu, S., Li, B., and Zhao, D. Constrained variational policy optimization for safe reinforcement learning. In International Conference on Machine Learning. pp.
PMLR, 2022.
Ray, A., Achiam, J., and Amodei, D. Benchmarking safe ex- ploration in deep reinforcement learning. arXiv preprint arXiv: 7(1):2, 2019.
Liu, Z., Guo, Z., Lin, H., Yao, Y., Zhu, J., Cen, Z., Hu, H., Yu, W., Zhang, T., Tan, J., et al. Datasets and benchmarks for offline safe reinforeement learning. arXiv preprint arXiv:2306. 09303, 2023a.
Schrittwiessr, J., Antonoglou, I., Hubert, T., Simonyan, K., Sifre, L., Schmitt, S., Guez, A., Lockhart, E., Hassabis, D., Graepel, T., et al. Mastering Atari, Go, chess and shogi by planning with a learned model. Nature, 588 (7839):604 609, 2020.
Liu, Z., Guo, Z., Yao, Y., Cen, Z., Yu, W., Zhang, T., and Zhao, D. Constrained decision transformer for offline safe reinforcement learning. In International Conference on Machine Learning, pp. 2-611-2 PMLR, 2023b.
Shannon, C. E. A mathematical theory of communication.
The Bell system technical journal, 27:379 423, 1948.
Lu, C., Zhou, Y., Bao, F., Chen, J., Li, C., and Zhu, J. DPM-solver: A fast ode solver for diffusion probabilistic model sampling in around 10 steps. Advances in Neurul Information Processing Systems, 35:5775-5787 2022a.
Silver, D., Schrittwieser, J., Simonyan, K., Antonoglou, I. Huang, A., Guez, A., Hubert, T., Baker, L., Lai, M., Bolton, A., et al. Mastering the game of Go without human knowledge. Nature, 550(7666):354-33555 2017.
Lu, C., Zhou, Y., Bao, F., Chen, J., Li, C., and Zhu, J. DPM- solver++: Fast solver for guided sampling of diffusion probabilistic models. arXiv preprint arXiv:2211.01095, 2022b.
J., Weiss, E., Maheswaranathan. N., and Ganguli, S. Deep unsupervised learning using nonequi- librium thermodynamics. In International Conference on Machine Learning, pp. 2256-2265. PMLR, 2015.
Lu, C., Chen, H., Chen, J., Su, H., Li, C., and Zhu, J. Con- trastive energy prediction for exact energy-guided diffu- sion sampling in offline reinforcement learning. In Inter- national Conference on Machine Learning, pp. 22825- 22855, PMLR, 2023.
Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Er- mon, S., and Poole, B. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:201 1.1456, 2020.

---

Q-Supervised Contrastive Representation: A State Decoupling Framework for Safe Offline Reinforcement Learning
Sootla, A., Cowen-Rivers, A. I., Jafferjee, T., Wang, Z., Mguni, D. H., Wang, J., and Ammar, H. Sauté RL: Almost surely safe reinforcement learning using state augmentation. In International Conference on Machine Learning, pp. 20423-20443 PMLR, 2022.
intrinsic vulnerabilities: Robustness, safety, and general- izability. arXiv preprini arXiv:2209.08023, 2022b.
Yarats, D., Zhang, A., Kostrikov, L,, Amos, B., Pineau, J., and Fergus, R. Improving sample efficiency in model- free reinforcement learning from images. In Proceedings of the AAAI Conference Artificial Intelligence, vol- ume 35, pp. 10674-10681, 2021.
Stooke, A., Achiam, J., and Abbeel, P. Responsive safety in reinforcement learning by pid lagrangian methods. In International Conference on Machine Learning, pp. 9133- 9143. PMLR, 2020.
Yu, D., Ma, H., Li, S., and Chen, Reachability constrained reinforcement learning. In International Conference on Machine Learning, pp. 25636-25655. PMLR, 2022.
Subramanian, J.. Sinha, A., Seraj, R., and Mahajan, A. Ap- proximate information state for approximate planning and reinforcement learning in partially observed systems. Journal of Machine Learning Research, 23(12):1-83, 2022.
Zhao, W., He, T., and Liu, C. Model-free safe control for zero-violation reinforcement learning. In 5th Annual Conference on Robot Learning, 2021.
Tessler, C., Mankowitz, D. J., and Mannor, S. Re- ward constrained policy optimization. arXiv preprint 1rXvv:180..11074 2018.
Zheng, Y., Li, J., Yu, D., Yang. Y., Li, S. E., Zhan, X., and Liu, J. Safe offline reinforcement learning with feasibility-guide diffusion model. arXiv preprint arXiv:2401. 10700, 2024.
Thananjeyan, B., Balakrishna, A., Nair, S., Luo, M., Srini- vasan, K., Hwang, M., Gonzalez, J. E., Ibarz, J., Finn, C., and Goldberg, K. Recovery RL: Safe reinforcement learning with learned recovery zones. IEEE Robotics and Automation Letters, 6(3):4915 4922, 2021.
Zhu, J., Xia, Y., Wu, L., Deng, J., Zhou, W., Qin, T. Liu, T--Y,, and Li, H. Masked contrastive representation learn- ing for reinforcement learning. IEEE Transactions on Pattern Analysis and Machine Intelligence, 45(3):3421- 3433, 2022.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., and Polosukhin, I. Atten- tion is all you need. Advances in Neural Information Processing Systems, 30, 2017.
Wachi, A., Sui, Y., Yue, Y., and Ono, M. Safe exploration and optimization of constrained MDPs using gaussian processes. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 32, 2018.
Wang, 7., Hunt, J., and Zhou, M. Diffusion policies as expressive policy class for offline reinforcement learning. arXiv preprint 2r0i 060093,..6.......0 2022.
Watter, M.. Springenberg, J., Boedecker, J., and Riedmiller, M. Embed control: A locally linear latent dynamies model for control from raw images. Advances in Neural Information Processing Systems, 28, 2015.
Xu, H., Zhan, X., and Zhu, X. Constraints penalized Q- learning for safe offline reinforcement learning. Pro- ceedings of the AAAI Conference Artificial Intelli- gence, volume 36, pp. 8753-8766 2022a.
Xu, H.. Jiang L., Li, J., Yang, Z., Wang, Z. Chan, V. W. K., and Zhan, X. Offline RL with no OOD actions: In-sample learning via implicit value regularization. arXiv preprint arXiv:2303. 158/0, 2023.
Xu, M., Liu, Z.. Huang, P., Ding. W. Cen, Z., Li, B., and Zhao, D. Trustworthy reinforcement learning against
12

---

Q-Supervised Contrastive Representation: A State Decoupling Framework for Safe Offline Reinforcement Learning
A. Theoretical Interpretations
The first comparison between the bisimulation representation and --iirrelevance representation for finite horizon MDPs was given by Givan al. (2003). The systematic state abstraction theory for MDPs was summarized in Li al. (2006). The expansion of the theory in Partially Observable MDPs (POMDPs) are introduced recently (Subramanian et al., 2022). It is worth noting that their formulation does not incorporate the safety Bellman operator, and comprehensive proof for infinite-horizon MDPs is not provided. We now provide a complete proof for Theorem 3.3 as follows.
Definition A.1. For any state representation 1, O2 we say θ is finer than 2₂ (or 22 is coarser than θ1), denoted as 1 = 2g if and only if for any states 81, 82 E S, (81) = (82) implies (x))) = (x2))
To clarify, let 21 and 22 represent the representations O (s) and (s))) for any s E S, respectively. It is always possible to find a function f: 2 - Z2 that is surjective. The equality holds (θ1 θ2) if and only if the surjective function is also injective (i.e., bijective).
Theorem 3.1. For any MDP, the optimal Q-functions induced by either the general Bellman operator or the safety Bellman operator satisfy Onisim č Og.. The optimal policies derived from both bisimulation representation and Q" Qr-irrelevance representation are also optimal in the ground MDP, i.e., " (Obssim (s)) = nT" (@ (s)) = n =s(ss for any state 8 E S.
Proof. We start by considering a finite-horizon MDP with a maximum timestep T. For any timestep 1. E {1, 2.., T), we
denote Q7 as the optimal-Q function at timestep t. Then, for Vs E S and Va E A, we have:
(s, = r(s, + E P(s') 81, (maxxEd Q 17 aa)).
(16)
For timestep T' + 1, we define (114 a) = for Vs E S and Va E A, which implies that (T, (8, = r(s, a). Now, for any 81, 82 E S that bisimilar (i.e., Obisim (81) Oissim (onsim((2)) (82)). we have (r, other words, for any z' E Zbisim, maxa'e (T:((r) (s', a') is the same for all s' E O (E').
(s,
(81, a) Q Q, In
Considering any 81, 82 E S that are bisimilar, and for any action a E A, we perform backward induction on timestep / from - 1 to following the proof sketch of Theorem 5 in Givan et al. (2003):
(s
= a)) + +7 Eves P(d' 81, (1,a))maxeee (rsttaryy6' d')l
- a)()) +7 E Eve(Urege 1(9)) P(s 81, s1, alaaxom Qf (5
(17)
2 0(11 +7 EvezE
P(s'l 81, mmaxwea U 1ir (8'
- 0)) +7 EvezErro-1e P(s') s s2,a)(max, O 0 ($ a')l
= C (82,
Equalities (a) and (b) hold due to the surjective relationship between and Equality (c) holds due to the definition of bisimulation and the fact that, for Vz' E Znisim, max 17 (s', a²) is the same for all s' E O (z'), as established by the induction hypothesis.
We denote Q; = Q" as the optimal Q function for the infinite-horizon MDP. The uniqueness of Q is guaranteed by
the fixed-point property of Bellman operator. For any S, E A, and timestep t E 21,2,... T - 1}. The optimal Q value gap between Q and Q can be expressed as:
(s, C: (s,
r(s,a) + Σ P(s' 1,
max (7 a')] a'EA
e'ES
r(s,a) - Σ
P(s'ls, ahl max O (2 (8,
a'EA
s'ES
13

---

Q-Supervised Contrastive Representation: A State Decoupling Framework for Safe Offline Reinforcement Learning
Σ P(o''s,0)11 max (f(),0) - max s'ES
a'EA (114) [')]
Q
d'Ed
Σ P(s'|s, ma (
max C 1ny (s', a'EA
d'EA
eEES
$
Σ P( a) max d'Ed
s'ES
$7 7 max a' E'
Q a') Q (s)
(18)
where inequality (a) holds due to the triangle inequality property, and inequalities (b) and (c) follow from the properties of the maximum function. Assuming that O and Qi, are bounded for any E and a E A, we conclude that for discount factor E (0, 1). the difference (Q7 a) - Q (T)) (s. converges to 0 as T - co. Applying backward induction as introduced in Eq. (17), deduce that, for any 81, 82 E S that are bisimilar (i.e., Obisim (81) Obisim (82)).
(51, a) ($2, a). This completes the proof that Obisim Qq. for general Bellman operators.
For the safety Bellman operator, we have an analogous definition that for any timestep t E {1,2, T), Thaery is the optimal-Q function at timestep t. Then, for Vi E S and Va E A. we have:
(haTT) = (1 7))ss) + Eves P(s' s. s.a))max(( mina'e (1,1s a').
(19)
We denote 17, (s. a) h(s) for all S and A, which implies that Pi (s.a) h(s) 0 and Qi, (8,0) a) = (1 h))ss) otherwise. Given any state 8 E S, for all E Zgisim, max (()) mina'en 2 (5', d')} the same for all s' m (z'). For any 81, 82 E S that are bisimilar, and for any action A, we apply backward induction from timestep T to I such that:
(
(1)-)(((1)-- +7 E ExE((NEE0-(1)) P(s' s1, (a.a)))ax(((s))) mina Qi
( a')))
(1 11--11 -))((11)+ +7 Evez E Eiec-eoees) P(s') 61,0) (h(s1), (mxx))))))((())) min Qh ,8)) = (1 -))((8)) +7 E Eveg- P(s'
(20)
minare Q Qi ,
Similarly, we denote Q7 ))) as the optimal Q function with safety Bellman operator for the infinite-horizon MDP. The uniqueness of Q is also guaranteed by the fixed-point property of safety Bellman operator. For any E S, a E A, and timestep & E 21,2,.... T 1). The optimal Q value gap between Q7 and 1h,rrr can be expressed as:
(7, a)- - (s,
(1 (1-)))()) + Σ
P(s i s, max( min (s','')) s'ES
d'EA
(1 (1-7)h(ss-
Σ s'ES
min Q 1 (6') d'EA
Σ
P(s'|s, a)ly max) min Qh(s', max(h(s), a'EA
min Qiatrer) (')) d'A
=
EE
< P(s' s, a) max(h(s), min (5 a²)) - max(h(s), min Q (s', d'EAA
d'A
SS
14

---

Q-Supervised Contrastive Representation: A State Decoupling Framework for Safe Offline Reinforcement Learning
Σ 8'ES
a(s')s,a) max (Qh(s','') (s', d'EA
$
max d'EAA''S
(2)(s','') - (s))))))))
(21)
similar conclusion can be drawn that for discount factor 7 E (0, 1), the difference 1Q% - ((2)) (s, a)] approaches 0 as T + sc, provided that O and Qú,T(T) are bounded for any E S and a A. With backward induction as introduced in Eq. (20), conclude that Oiisim -q- for safety Bellman operators.
Lemma A.2. For any MDP, given the optimal Q-functions induced by either the general Bellman operator or the safety Bellman operator, the optimal policy for Qq- remains optimal in the ground MDP.
Proof. For any state s S, we observe that Q" ("((q (s), (s, for any A, as per Definition 3.2. It is evident that for given state E S:
a. = arg m ( arg max (s.
(22)
a = arg min Qī(0g; (s), = arg minge.A (%,
(23)
Therefore, conclude that the optimal policy is preserved for Q" -irrelevant representations.
Lemma A.3. For any MDP, given the optimal Q-functions induced by either the general Bellman operator or the safety Bellman operator, and any representation O₁ that is finer than Qq i.e. θ ż Qqq. it holds that Q* (s), as) for any E S and E A. The optimal policy for O1 is also optimal in the ground MDP.
Proof. We denote the optimal value function for representation θ1 OQ- and ground state as Qé Uzg. and Q respectively. is evident that 06, (+), a) Geg. (s)- (7 a) for any & S, E A one of the solutions for Qe, due to the subjective relationship between O O2(s) and O (s). To show that the optimal value function for representation θ is unique, suppose that there exist two optimal value functions (,, a) and Qi2(z, a) for any representation z E Ze, and action a E A. The gap between them can be expressed as
(Q:(z,a) = 10; (z, (z)
= =,,, P(e'|2, a)> (maxa (2',''a'') max. (22(2' Symaxy', a Aq: (z',a')
(24)
where the inequality directly arises from the reasoning outlined in Eq. (18). With the discount factor y E (0, 1), Aq: tends zero for any finite value of Qi1 a) and (7) This implies that the fixed point Q' for the general Bellman operator is always unique. For the safety Bellman operators, we also have
Dez (z,a)
= = (2, a) ())(=aa)
Σ. P(=12z, a) (max(h(2)) mina Qin('',')) max(f((2), min. (ha
(25)
S 7 maxx Aoz (z',
where the inequality is a straightforward result of the proof sketch given in Eq. (21). It can be concluded that (*)) (",s a) holds for any O QQ., both for the general Bellman operator and the safety Bellman operator. Therefore, the optimal policy for θ is also optimal in the ground MDP, following the proof sketch provided in Lemma A.2.
Combining Eqs. 17 18 and 20 21, we conclude that Obesmm > θQ- holds for both the general Bellman operator and the safety Bellman operator. Combining Lemma A.2 and Lemma A.3, we conclude that both Opisim and Oo. preserve optimality for ground MDP. The proof of Theorem 3.1 is complete.
15

---

Q-Supervised Contrastive Representation: A State Decoupling Framework for Safe Offline Reinforcement Learning
Theorem 3.1 shows that our Q'-irelavancee representation leads to smaller representation space than bisimulation representa- tion, while preserving the optimal policy. A smaller representation space typically implies higher generalization capabilities and higher sampling efficiency during the policy learning process (Liu et al., 2021).
Proposition A.4. For any MDP, 1 - 2₂ indicates H (s (s)) < (s)02(s)).
Proof. We denote 21 and 22 as representations of 1 (s) and (s)) respectively. We have
21)
H(s = Σ p(s, 21 log p(2))
8,21
>> p(s, 21
ΣΣ
((s))
8,21
21
Σ logg(s(((s) + Σ log
21
_ Σ p(s) + Σ p(z2) Σ ((21122)logp(3)
3 p(s) + Σ p(2)) log Σ p(21 ((21122)p((1)
82
21
1 p(s)) + Σ p(z2) log
[nte (200)())))))
22
+ Σ p(22) p(22)
22
H (s|22).
(26)
where pp(z l2s [2)) if 22) ب O, and = otherwise. Equality (a) holds as is function of s. Note
that Σziz p(2 (22) = for V22 2, in equality (6). Inequality (c) is consequence of Jensen's 's inequality. Inequality (d) holds since for 2/2 E 22, the conditional probability p(21|22) does not exceed I for all 21 E {a E 21lp(z(22) # 0}. Equality (e) holds due to the surjective relationship between 21 and ø2.
Based on Theorem 3.1 and Proposition A.4, we conclude that:
0 S ((s Opisim (s)) S H(s(Og-
(27)
Given our primary objective of maximizing the conditional entropy H(s)=s (s)), the proposed Q-supervised contrastive learning method theoretically exhibits superior generalization capabilities compared to bisimulation.
B. Methodology Clarifications
B.1. Implicit Q-Learning
Implicit Q-Learning (IQL) is the pioneering in-sample offline RL algorithm proposed by Kostrikov et al. (2021). It decouples the estimation of optimal Q-values from policy optimization, enabling implicit policy learning through the value function. Unlike standard Q-lcarning, which explicitly derives a policy by maximizing Q-values, IQL avoids direct maximization, reducing susceptibility to instability issues such as overestimation or divergence. The core technique in IQL is upper expectile regression. Given random variable X with an unknown distribution, the (0, 1) expectile can be estimated by solving:
arg min Ex m.
(æ [L m.)), )} where L (u) = |T Ilu < 0) u
(28)
Specifically, as r - 1, the solution to Eq. (28) approximates the upper bound of the random variable X. Extending this to conditional distributions, the optimal value function can be approximated by minimizing:
Lv (s,a)-,, (Q(s,)) - V(s))),
(29)
and the optimal function can be updated accordingly with the TD loss
Lq "e D [(r >V(s') - ((s,
(30)
16

---

Q-Supervised Contrastive Representation: A State Decoupling Framework for Safe Offline Reinforcement Learning
This process yields in-support optimal value and Q functions without training optimal policies. In this paper, extend IQL. with safety Bellman Operator and estimate both upper bound and lower bound of the value functions for safety assessments.
B.2. Diffusion Behavior Cloner
The diffusion model initially introduced as an iterative denoising framework for image generation in the domain of computer vision Dickstinc al., 2015; Ho et al., 2020). More recently, it has been adapted for decision decision-making in state-based tasks, due to its superior performance in capturing action distributions within dataset. As introduced in Section 3.2, SDQC requires a behavior cloner to capture and reproduce in-support actions for each state within the offline datasets. Following the approach of Lu et al. (2022a:b), we employ score-based diffusion and the DPM-Solver. The training loss for the behavior cloner Thehau is expressed as follows:
Lasehev = -.r.
Serenav (at, 8, t)ll]-
(31)
After training converges, we use second-order DPM-Solver (Lu et al., 2022a) to form Thehae and sample |Agg] actions for each state offline datasets Dg. These actions will be utilized the subsequent joint optimization of Q-functions and representations.
B.3. Diffusion Policy
As described in Eqs. 13 and 14, we train three distinet diffusion policies using weighted regression (Zheng et al., 2024). In line with most existing ifffusion-basee policies (Wang et al., 2022; Garg al., 2023; Lu al., 2023), our three policies can be formulated follows:
Th (al 20, (s)) Pv- (ao:rl 20, (s)) Mar 0, I) I) Po., (at-1104, 20, (s))
ITA (a) z0,, (s)) Par (0o.T z0,, (5) NNar; 0, 1) I] Per, (ax- 110 20t (0)
(32)
Teo (a) 20, (s), 20, (s)) = Prioo (a0:T 20, (s), 20, (s))
-aaa 0, I) IT Perro (at-1 ( (5), zan (5))
where the reverse transitions are modeled as Gaussian process:
Par (a.-1 ac, 2a, (s)) (Vaa (a,, 20, (s),t),((t))
Pon (at-1lae, 20, (s)) = (at-1i (ae, 20, (s) E(t))
(33)
Porto (ac-1 Iace =0, (8), 28,, (s)) Na 1i parro (ac, =0, (8), 20,, (s), E()).
Given variance schedule defined by 3r Cit, we proceed to define:
ā II az
B.
1 - ac- 3,-
(34)
1 a₁
i=1
The mean of the Gaussian process is then given by:
Par (ar, 20, 20, (s),
| (u+ B. Co, (ar, 20, (8), ())
Va,
vI-ai
Per, (ot, 20, (8), = 1 (at B.
So, (at, 20,, (s). 1))
(35)
Vac
VI-a.
Per (ui, z0, (8), 20,, ()), Vat
(at VI-aa So. (at, 20, (8), 20, (8), ())
and the covariance matrix is expressed as Ett = I.I. During the testing phase, actions can be sampled from the reverse diffusion chain for each diffusion policy. To account for safety considerations, we sample multiple actions and select the one with the lowest Qn value as the final action to be executed (cf. Appendix E.2 for details).
17

---

Q-Supervised Contrastive Representation: A State Decoupling Framework for Safe Offline Reinforcement Learning
C. Implementation Details
C.1. SDQC Summarization
To provide an intuitive understanding the mechanism of SDQC, we present a brief summary of its training and deployment scheme in this subsection.
Algorithm SDQC Training
Phase 1: Behavior cloner training and behavior actions generation
Require: Initial network Sereamee datasets Dn = {s, a, srr,, 4] for each iteration {s, a} M2 D do
2: Update Gusehan using Eg. (31)
3: end for
4: for each state D do
5: Given s, generate multiple behavior actions Ab with Gisshav using DoM-Srlver (Lu et al., 2022a)
6: end for
7: return Updated datasets D3 = {s, 43, 9, h,
Phase 2: Joint optimization for value functions and representations
Require: Initial ewaardrelltted network V,, Q ze,. cost-related network Viu,, V/lo Qu, 20, and datasets D, ,r,,,,
8: for each iteration As, a, rr]'² D3 do
9: Update Vr, Qr z0, jointly using Eq. (8)
10: end for
11: for each iteration (s, a, 8', 81, ,/ ~ D3 do
12: Update V/2 Vi Qni 20,, jointly using Eq. (12)
13: end for
14: return Reward/Cost-related value function and representation networks Vr, Qr, 20,, VilP, Vilo Qn. 20,
Phase 3: Three policies extraction
Require: Initial policy network Sua. SWi.. tixed pre-trained network Vr, Qr, z0,, Vi Qu, 20,, and datasets D3 =s,
15: for each iteration {s, a} Ma ~ DD do
16: Calculate regression weight for three policies with Vr, Zr, 2B,, V/rt Qn, 20, using Eq. (14)
17:
Update Ce, Soa. Seve using Eq. (13)
18:
end for
return Three distinct policies Iy, Nh, Tto
As introduced in Section 3, SDQC requires three-phase training process (see Algorithm 1). The primary objective of the first stage (lines 1-7) is to generate a set of behavior actions As for each state E Dg. The set is then utilized in the second stage to measure the similarity between states, represented by d(xx,8) 8) Rs SuPaE (=))=)()), a) Q" (=a(8), a)}- The second phase (lines 8- 8-14), known the joint optimization of value functions and representations, employs expectile regression (Kostrikov et al., 2021) to learn in-support optimal value functions while simultaneously using contrastive learning to cluster similar states (measured by d(s, 6)) within the representation space. The final phase (lines 15-19) introduces three distinct policies, which are trained using weighted regression diffusion models as introduced by Zheng et al. (2024). It is important state space. During
to note that these policies are conditioned on the representation space rather than the ground deployment, the selection of the specific policy to be adopted is guided by the evaluation of states using the cost-related value function. For further details, please refer the subsequent paragraph.
Upon completion of training, SDQC is deployed with three distinct policies nr The Tio derived from the third training phase, along with two representation networks 20, and three cost-related value functions TiuP Vllo Qn from the second training phase (cf. Algorithm 2). Given any state s, we initially conduct safety assessments based on cost-related representations. The condition V (=n, (s)) > 0 (line 2) suggests the existence of in-support actions that might lead to unsafe outcomes, necessitating a joint consideration of safety and reward. A positive Vydow(2e (s)) (line 3) indicates that no action can ensure safety in future trajectories; the agent's primary objective therefore exit the unsafe region by deploying policy The regardless of the reward considerations. Conversely, Vllow <0 Vrtt (line 5) reflects borderline safe condition, requiring the agent consider both reward and cost, thereby deploying policy moo- On the other hand, Vi (ze, (s)) < 0 (line 8) confirms absolute safety, obviating the need for the agent consider cost-related information. An illustration diagram is presented in the right subplot of Figure 2. Note that single action sampled by diffusion model is not trustworthy enough, thereby sampling batch actions and conduct the one with the lowest Qn ( (zo, (s), a) leads to safer outcomes (line 11). Please refer to Appendix E.2 for further details.
18

---

Q-Supervised Contrastive Representation; A State Decoupling Framework for Safe Offline Reinforcement Learning
Algorithm 2 SDQC Deployment
Require: Three policies ny, Rh, Too representation network 20,; , and cost-related falue-funciions V V. Qh
1: Given any
2: if V/ (s)) > then
(There exist in- in-support actions that may lead to unsafe outcomes)
3:
if Vi (ze, (s)) > 0 then
4:
sample actions: [a;]cand: jeand ~ TT (-1120, (=))
5: else
6:
sample actions: {a. caa]aaud. Tro(-120 (s), =0, (s))
7:
end
8:
else
(The state is absolute safe)
9:
sample actions: {a, yeand. n, (-120- (s))
10:
end if
11:
return Final action arg (aa jound. Qn (20, a))
C.2. SDQC Network Structure
As described Section 3.3, the representations in our proposed SDQC framework are trained concurrently with the optimal Q value learning process. The neural network structure illustrated in Figure 5 is utilized for training both the reward- and cost-related representations as well as the value functions. The global observation s is encoded into the representation 2, and the value functions (both V and Q) are computed based on this representation with separate (MLP) neural networks.
MLP
L2 Norm
State Encoder (MLP or ATN)
S
MLP
In certain safe RL benchmark problems, it is observed that the majority of dimensions in the global observa- tion share similar physical meanings. For instance, in domain, a significant number of dimensions in the global observations correspond to lidar measurements, which provide information about the distances between the agent and the destination obstacles in specific directions. This reminds us of the self-attention mechanism (Vaswani et al., 2017), which is known for its superior ability to capture relationships among input information that share similar representations in comparison to traditional MLP architectures.
Q
a
the
Figure Neural network structure for training value functions and representations of SDQC.
Nevertheless, attention mechanisms typically rely on vector multiplication to compute attention weights, which presents a challenge when dealing with global observations where each dimension contains scalar information. Towards this end, we propose transform each scalar observation dimension into vector representation using fixed Gaussian Fourier Encoder (Ho et al., 2020). Subsequently, attention is applied the encoded vector representations. The output of the attention module is then flattened and passed through an MLP to obtain the final representation. Please refer to Figure 6 for a detailed illustration of the network structure.
To demonstrate the efficacy of our attention-based state encoder (ATN) in effectively capturing information from global observations and identifying the relevance of specific dimensions to reward/cost, we present the attention patterns (i.e., softmax(( KT) / /d of the reward- and cost-related state encoders in the task "PointGoal2" (cf. Figure 7). In "PointGoal2", the global observations consist of 60 dimensions. Among them, the first 12 dimensions represent the self- status of the agent, the subsequent 16 dimensions contain reward-related information, and the last 32 dimensions contain cost-related information. Ideally, reward-related attention pattern should assign higher attention weights to the first 28 dimensions while ignoring the last 32 dimensions. On the other hand, cost-related attention pattern should focus on the first 12 and last 32 dimensions while disregarding the middle 16 dimensions. The observed attention patterns during our experiments align with the relevance of specific dimensions to reward and cost. For the ablation study on the state encoder network structure, please refer Appendix D.
19

---

Q-Supervised Contrastive Representation: A State Decoupling Framework for Safe Offline Reinforcement Learning
Matrix Operation Block
Transformer Encoder
MLP
Multi-head
Add Norm
Feedforward
Add Norm
Attention
Z
Feature Embedding
Flatten
Figure Neural network structure of the antention-hased state encoder.
Reward Representations
Cost Representations
12
44
60
28
44
60
12
Dimension Observation Info
1-12
Self Status
28
13-28
Goal Lidar
28-44
Hazard Lidar
44-60
Vases Lidar
44
44
60
60
Figure 7. Attention pattern of the reward- and cost-related state encoder, and the meaning of each observation dimension in the task "PointGoal22" from Gymnasium. Darker colors represent higher values. The "Goal Lidar" dimensions observe the position of the destination, indicating reward information. The "Hazard Lidar" and "Vases Lidar" dimensions observe the position of obstacles, indicating cost information. pattern averaged over 3000 observations randomly chosen from DSRL datasets.
C.3. SDQC Hyperparameters
As discussed in Section 3.2, to calculate the soft similarity measure for con- trastive learning, our SDQC framework requires pre-training a generative model capture the behavior policy of the offline datasets TT. For this purpose, employ diffusion probabilistic models (DPM) (Ho et al., 2020; Song et al., 2020) and utilize the DPM-Solver, fast high-fidelity ODE solver proposed by Lu et al. (2022a;b), to generate the behavior actions of each state s in the offline datasets, denoted as Ag. We utilize the default network configurations outlined in (Lu et al., 2022a; 2023). Specific hyperparameter settings can be found in Table
Table 2. Hyperparameters of the DPM- solver for generating behavior actions.
Hyperparameters
Value
Learning rate Batch size
3e-4
4096
Training steps
5e5
Diffusion timesteps
15
SDQC can simultaneously train the reward and value functions and their respective representations using Eqs. 8 and 12. The network structures described in detail in Appendix C.2, and generic hyperparameters can be found in Table 3. Regarding the updating of the safety Bellman operator, we follow the settings in FISOR (Zheng et al., 2024). The constraint violation function defined as h(s) = when the cost function satisfies c(s) = 0, and h(s) = 25 when c(s) >
Generated action numbers for each state [.4']
Considering the significant variation in physical meanings among the observation dimensions of different tasks, we employ different state encoder structures accordingly. For tasks that have observation dimensions with diverse physical meanings,
20

---

Q-Supervised Contrastive Representation: A State Decoupling Framework for Safe Offline Reinforcement Learning
Table Generic hyperparameters of SDQC in value functions and representations training phase.
Module
Hypee-paramettrs
Value
Optimizer
Adam
Learning rate
3e-4
Batch size
512
General
Training steps
5c5
Soft measure temperature factor ")
1.0
Contrastive temperature factor
0.1
Contrastive term cocf 8
1.0
Number of hidden layers (Q & V)
2
Number of neurons in hidden layer (Q & V)
256
Activation function (Q & V)
Mish
Critic
Expectile T
0.9
Discount factor
0.99
Target eritic soft update
0.005
Number of hidden layers
MLP
Number of neurons in hidden layer
256
Activation function
Mish
State Encoder
Number of head
2
ATN
Embed dimension for each head
64
Dropout rate
0.1
we utilize the MLP structure. Conversely, for tasks where most observation dimensions have consistent physical meanings, we utilize the attention-based state encoder (ATN).
is observed that the performance of SDQC with an ATN-based state encoder improves when trained with a larger contrastive loss coefficient (ō in Eqs. 8 and 12) and higher number of anchor points ((Z) in Eq. (5)). On the other hand, the SDQC performs better with smaller values of 5 and fewer anchor points if the MLP based encoder is used. Besides, the global observation dimensions vary across different tasks. For the based state encoder, we select the encoded state dimension (i.e., the dimensionality of z) to be approximately half of the global observations. On the other hand, for MLP, we choose the encoded state dimension to be roughly twice the size of the global observations. Optimal hyperparameters achieving the best performance on different tasks are presented in Table 4.
Table Hyperparameters of SDQC for different tasks. "All" denotes all different tasks for the same agent, while "Vel" refers to the velocity task.
State Encoder
Encoded
Contrastive
Anchor Number
Domain
Agent
Task
State Dim
Loss Coef
Point
ATN
32
1.0
Car
All
Safety
HalfCheetah
Vel
Gymnasium
Ant
Vel
MLP
32
0.5
4
Swimmer
Vel
Ball
All
16
Bullet
Car
All
32
MLP
0.1
4
Safety
Drone
All
64
Ant
All
64
For the final training phase, which involves policy extraction using weighted regressed diffusion models as described in Eq. (13), we follow the network structure design and generie diffusion parameter selection described by Zheng et al. (Zheng al., 2024). We train three separate policies (x, Ty and Tio) with learning rate of 0.0003, a batch size of 1024, and the total number of training steps is set 500,000. The temperature parameters that control the strength of behavior
21

---

Q-Supervised Contrastive Representation: A State Decoupling Framework for Safe Offline Reinforcement Learning
regularization (in Eq. (14)) are chosen by Lto 3.0 and 'n 5.0.
D. Additional Ablation Studies
Ablation studies on network structure. To demonstrate the effectiveness of our proposed iiiention-baseed state encoder (cf. Appendix C.2), we conduct ablation studies on the neural network architecture, as depicted in Figure 8. By substituting the attention-based state encoder with an MLP-based counterpart, we observe a deterioration the performance of the SDQC, terms of diminished rewards and increased costs. The t-SNE visualization results depicted Figure 8b demonstrate that while the MLPP-based state encoder does cluster representations with similar values in the high-dimensional space, the clustering effect is not robust as that achieved by the attention- basedd approach. Consequently, this leads to an overestimation of the cost value, resulting in inaccurate assessments of the safety condition.
(a)
Reward
of Cost of (MLP)
(b)
Original States
Reward Represencations(ATN) Reward Representations(MLPP
V
Normalized Reward
0.3
100
0.1
0.0
20
6.0
400
Original States
Cost Representations(ATN)
Cost Representations(MLP
v vlow
0.0
Normalized Cost
4.0
-0.
-0.4
2.0
-0.6
-0.8
0.0
100
200
300
400 500
-1.0
Gradient steps (x10P)
Figure 8. Ablation studies on the network structure in CarGoal2. (a) The actor-training- proccesss evaluations of SDQC with attention-based (ATN) and MLP-based (MLP) state encoder. (b) t-SNE visualization the distribution of the original state, the reward- and cost-related representations with ATN/MLP state encoders across 10 different safe trajectories.
Table Ablation studies the choice of anchor number.
CarGoal2
CarPush2
Anchor Number
Runtime (slepoch)
Runtime
Reward
Cost
Reward
Cost
(s/epoch)
16
0.22
0.00
31.4
0.28
0.06
36.2
8
0.23
0.00
28.8
0.31
0.04
32.7
0.20
0.05
26.0
0.25
0.06
29.9
0.15
0.13
23.6
0.10
0.18
27.6
0.05
0.86
19.4
0.21
2.15
23.7
Ablation studies on anchor number choice. learning method the anchor number, [Z] in Eq. (5). This parameter determines the number of representation pairs be clustered in the high-dimensional space during each gradient step. The ablation study results summarized in Table 5, Our experimentation reveals that anchor number result in subpar clustering outcomes, consequently impairing the performance of SDQC. Conversely, overly large anchor numbers lead to an excessive influence of the contrastive loss term
essential hyperparameter in our proposed Q-supervised contrastive
22

---

Q-Supervised Contrastive Representation: A State Decoupling Framework for Safe Offline Reinforcement Learning
in the overall loss function, increasing the computational costs. To strike balance and attain optimal performance, we let [ 8 for our attention-based state encoders.
Ablation studies on contrastive-rehated hy hyperparameters. In our SDQC framework, one of the critical components is the contrastive representation loss, as described in Equations 8 to 12. This involves selecting appropriate values for the term coefficient 5 and the exponential temperature v. As shown in Tables 3 and 4, we vary ó across different domains but maintain a consistent V = 0.1 across all environments. The effects of these parameter choices are detailed in Table 6. With respect to the temperature v, employing a very small value (0.01) tends to destabilize the training process, ultimately resulting in collapse. Conversely, using a larger value (1.0) produces poorly clustered representations, leading to marked degradation in performance. Regarding the term coefficient a smaller value results in slight performance decline. However, larger coefficient excessively prioritizes the contrastive loss, destabilizing the training of the value function and significantly degrading performance. 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.
Table 6, Ablation studies on contrastive-related hyperparameters.
Contrast Contrast
Contrast Contrast Coef.
Env
Reward Cost
Env
Reward
Cost
Coef. (8)
Temp. (v)
Temp.
0.01
NaN
NaN
0.01
NaN
NaN
0.1
0.31
0.04
1
0.1
0.29
0.09
PointGoal2
1
0.22
0.16
CarPush2
1
1
0.20
0.48
0.1
0.1
0.31
0.17
0.1
0.1
0.30
0.10
10
0.1
0.23
0.48
10
0.1
0.22
0.53
Ablation studies on the deployment of three distinet policies As introduced in Section 3.3 and Appendix C.I, SDQC coordinates three distinct policies, reward policy Ty, trade-off policy Tto. and cost policy The to ensure excellent safety performance. To verify the necessity of each policy, we conduct ablation studies examining their individual deployments, with results presented Table 7. Notably, naive reward policy मr focuses solely on maximizing rewards while ignoring costs, a naive cost policy Th prioritizes minimizing costs but disregards rewards, and a naive trade-off policy Tto takes both into account but fails excel in either maximizing rewards minimizing costs. The best performance consistently results from the collaboration of all three policies. When the trade-off policy nto is omitted (combining т, and TT). the agent incurs higher costs as it cannot respond promptly to borderline dangers. Combining the trade-off policy too and cost policy ga does not increase costs, but results in decline in reward accumulation. While combining the reward policy गr, and trade-off' policy to achieves comparable performance to using all three policies, it results in slightly higher costs due to the agent's ss reduced ability to quickly escape dangerous situations.
Table 7. Ablation studies on the deployment of three distinct policies.
Policy Num.
One
Two
Three
Env Name
Naive
Naïve teo reward cost
Naive Th reward
Tr and reward
and Thh
Tto and Th
Tr, Tto
and Wh
reward
cost
cost
cost
reward
cost
reward
cost
reward
cost
PointGoall
0.69
4.92
0.27
0.69
0.01
0.00
0.32
0.42
0.34
0.51
0.22
0.18
0.35
0.36
PointGoal2
0.75
13.28
0.28
0.18
-0.11
0.00
0.23
0.24
0.20
0.14
0.23
0.12
0.29
0.09
CarPushl
0.38
1.32
0.26
0.00
0.05
0.00
0.27
0.00
0.27
0.12
0.21
0.00
0.30
0.00
CarPush2
0.42
4.34
0.27
0.01
0.02
0.00
0.31
0.01
0.28
0.16
0.18
0.03
0.31
0.04
E. Experimental Details
E.1. Task Description
-ymnns-Gmmmmsiumm (Ray et al., 2019). A collection of environments based on the Mujoco physics simulator. In the obstacle avoidance series environments, there are two agents (Point and Car) and three main tasks (Goal, Button, and Push), each with two levels of difficulty (1 and 2). Agents aim to reach the goal while avoiding any contact with obstacles. The environments are named using the following convention: {Agent) (agent)[Taask] (Task) ) (Difficulty). ). In the velocity environments, there are three agents: Ant, HalfCheetah, and Swimmer. The primary objective of these agents is to maximize
23

---

Q-Supervised Contrastive Representation: A State Decoupling Framework for Safe Offline Reinforcement Learning
their rewards while adhering the imposed velocity constraints. The environments are named in the convention of {Agent) Velocity.
Bullet-Safety-Gyy (Gronauer, 2022). A suite of environments built upon the PyBullet physics simulator. These environments are similar to Lafety--Gymnasium but feature a broader range of agents (including Ball, Car, Drone, and Ant). The tasks are relatively straightforward, with only two options available (Circle and Run). The environments are named through { {Agent] ( (Task).
E.2. Experiment Settings
We train the baseline algorithms using the recommended hyperparameters specific to each task, for BCQ-Lag (Fujimoto et al., 2019; Stooke et al., 2020), CPQ (Xu et al., 2022a), COptiDICE (Lee et al., 2022), CDT (Liu et al., 2023b), TREBI (Lin et al., 2023), and FISOR (Zheng et al., 2024). To ensure a fair comparison, we train the baseline algorithms with three different random seeds and save the final output policy for safety evaluation. For each output policy, we conduct evaluations over 20 episodes to obtain reliable performance measures.
As for the training process of SDQC, follow the neural network structure design and hyperparameter settings in Appendix C. Analogously, we select three different random seeds for training and perform evaluation over 20 episodes for each seed. To improve safety performance, we follow Zheng et al. (2024) to sample 16 candidate actions for each RL timestep, regardless of the safety assessment results and policy usage. The safest action is then selected based on the lowest (z, (s), a) value and executed as the final action.
E.3. Computational Costs
We implement SDQC using PyTorch (Paszke et al., 2019) and conduct experiments on a single machine equipped with one GPU (NVIDIA RTX 4090, 24GB) and one CPU (AMD Ryzen 9 7950X). The training process comprises three phases, The first phase, known as the diffusion behavior cloner, demands approximately hour for each task. For the second phase, which involves training representations and crities, the duration varies depending on the chosen network architecture. Attention-bast based architectures typically require over 4 hours, whereas MLP-based architectures typically demand around I hour. Finally, in the last training phase, the diffusion actor, convergence typically occurs in about I hour (without online testing).
Besides, we assess the inference time consumption of all baseline algorithms across 1000 RL timesteps on the CarPush2 task, averaging the results over 10 trials. Although SDQC is relatively slower compared to other non-autoregressive policies, it remains within an acceptable range.
Table 8. Inference time (seconds) comparison over 1000 RL timesteps.
BCQ-Lag
CPQ
COptiDICE
CDT TREBI
FISOR
SDQC
1.51
1.85
1.86
3.45
585.87
6.11
11.13
F. Additional Experimental Results
F.1. Addtional Generalization tests
In addition to the generalization tests (on the agent "Car") presented in Section 4.2, we perform generalization tests on the "PointGoal" and "PointPush" tasks (in Bafety-Gyynnsiumm as illustrated in Figure 9, Similarly, the "Point" agent is challenged with tasks that involve reaching goal point or pushing a box to a goal point in hazardous areas with obstacles, with the difficulty level varying between simple (Po1ntGoall, PointPush1) and challenging (PointGoal22, PointPush2).
Experimental observations reveal that, in comparison to the "Car" agent, the "Point" agent demonstrates a higher degree of inertia during its motion within the environment. Specificallyy the "Point" agent lacks the ability to instantancously halt or promptly alter its direction, thereby rendering the maintenance of safety more challenging in equivalent tasks when compared to the "Car" agent. This significantly undermines the generalization capability of most algorithms the "Point" agent. For instance, the state-of-the-art (SOTA) safe offline RI algorithm FISOR performs poorly on the "Point" agent, exhibiting high costs across multiple cnvironments. In contrast, our SDQC algorithm still achieves nearly zero violations in the majority of environments.
24

---

Q-Supervised Contrastive Representation: A State Decoupling Framework for Safe Offline Reinforcement Learning
(a)
Reward
Cost the policy S-trained S-tested
(b)
Reward
Cost the policy C-trained Citested
Reward
Cost of the policy S-trained C-tested
Reward
Cost of the policy C-trained S-tested
Safety-Gymmasium Ponaeyyyy
Po
PointPush
Safety-Gymnasium Pointtaaal
Safety-Gymnazium Gymnssiuu-Pointtushh PointPush
0.3
0.8
0.8
0.8
Normalized aizez Reward
Normalized izzd Reward
0.6
0.6
0.6
0.6
0.4
0.4
0.4
-0.2
coptiDiCEo FISOR spac
-0.2
860, UMC00
CoT
FOOR
spoc
-0.2
8C0,
coptiIIEET
FISOR spcc
-0.2
800, L20co0
coptoll
FSOR spoc.
8.0
8.0
8.0
6.0
Normalized Cost
Normalized Cost
6.0
6.0
6.0
6.0-
4.0
4.0
copcorcest
F50R
sooc
coptitic 600
PSDR
spoc
8C0,100
EP0
coppocic
500c
EP
Coptiric
LD7
FSDR
5000
Figure
The generalization tests on the agent "Point" in Safety-Gymnann Gymnasium.
Value Functions During Training
SDQC without representation loss
SDQC with representation loss
FISOR
Cost-Value
Cost-Critic(Q)
Reward-Value
Reward-Critic(())
15
15
100
100
CarPush2 Value
10
0
100
200
300
400
500
0
100
200
300
400
500
0
100
200
300
400
500
0
100
200
300
400
500
MSE Loss
0.20
0.5
0.15
0.4
10-
0.10
0.05
0.00
200
400
200
400
200
400
500
200
400
500
Value
8-
8
20
20
BallCircle
100
200
300
400
500
100
200
300
400
500
100
200
300
400
500
0
100
200
300
400
500
0.8
MSE Loss
0.6
15
0.75
0.4
0.50
5
0.25
0.00
0.0
100
200
300
400
500
0
100
200
300
400
500
100
200
300
400
500
0
100
200
300
400
500
Gradient steps (x10²)
Gradient steps (x10²)
Gradient steps (x10²)
Gradient steps (x10²)
Figure 10. Comparison of value function loss during the training process among SDQC with and without representation loss, and FISOR, on the 'CarPush2" and BaalCCircle" tasks.
25

---

Q-Supervised Contrastive Representation: A State Decoupling Framework for Safe Offline Reinforcement Learning
F.2. Impact of Representation Loss Value Estimations
major concern for SDQC is that the joint optimization process of the value functions and representations (refer to Eqs. and 12) may lead instability in value estimation, subsequently affecting final performance. We analyze the performance of SDQC, both with and without representation loss, and compare it to FISOR by examining their respective Critic (Q) loss and Value loss patterns in relation to reward/cost metrics. This comparative analysis is conducted throughout the training process using two tasks: 'CarPush2' and 'BallCirele', with results illustrated in Figure 10. The experimental results indicate that while the inclusion of representation loss does lead to an increase in critic and value loss, it does not compromise the overall stability of the training. Furthermore, our proposed neural network architecture (used for "CarPush2"), with the incorporation of an attention-based state encoder, markedly improves the precision and stability value function learning compared to the simple MLP utilized by FISOR.
F33 Learning Curves Comparison between SDQC and FISOR
Reward (SDQC)
Reward (FISOR)
Cost (SDQC)
Cost (FISOR)
Safety-Gymmaslun-Carrr
5afetyyGymmasiuum-arrro
Nomalized Reward
3.0
Normllze Reward
0.25
3.0
0.4
2.5-
0.20
2.5
0.3
Normal
2.0
0.15
Rormal zed
2.0
0.2
1.5
1.5
0.05
0.1
1.0
0.00
1.0
0.0
-0.05
100
200
300 400
500
100
200
300
400
500
-0.10
100
200
300
400
500
100
200
300
400
500
Gradient steps (x10%)
Gradient (x10P)
Gradient (x109)
Gradient steps (x10*)
Carpu
Safety-Gymnasee Gymnasium Syymnssimm-CrPush22 CarPush2
0.35
4.0
0.35
4.0
0.30
3.5
0.30
3.5
Normelieed Reward
Normelized Reward
0.25
Nornalteed Cost
3.0
3.0
2.5
Normaltzdd
2.5
0.15
2.0
0.10
0.05
0.05
0.00
-0.05
0.05
100
200
300
0.10
400
500
0.0
Gradient steps [x102)
Gradient steps (x10²)
Gradient steps (x10?)
Gradient steps (x10?)
Snlety.Gymmasluum-CaaBtttooo
0.10
Nomelzed Reward
Normelized Rewerd
203 Cost
0.00
Nomalized Cost
-0.05
M
-0.55
2.5
-0.10
2.0
-0.10
2.0
-0.15
1.5
-0.15
1.3
-0.20
1.0
0.20
1.0
-0.25
0.5
0.5
-0.30
100
200
300
0.0
100
200
300
100
500
-0.30
100
200
300
400
500
100
200
300
400
500
Gredient steps (x10")
Gradient steps (x10')
Gradient steps (x10?)
Gradient steps (x110?)
Figure Training curves of SDQC and FISOR on the "Car" agent with tasks "Goal," "Push," " and "Button" the fyfety-Gymmaiium domain.
As complementary to the SOTA safe offline RL algorithm FISOR (Zheng et al., 2024), our SDQC employs the same implicit Q-learning method (Kostrikov al., 2021) learn optimal value functions and utilizes the safety Bellman operator in (Fisac al., 2019) for safety assessment. Additionally, we adopt their approach for policy extraction through training weighted regressed diffusion model. However, it should be noted that our decision decision-making process is based on decoupled representations rather than global observations. Furthermore, our policies are completely decoupled, and different policies are selected based on varying safety assessment results. To provide futher comparisons between SDQC and FISOR, we plot the training curves of both algorithms on the DSRL benchmark (Liu et al., 2023a) in Figure 11-14. The experimental results indicate that SDQC exhibits higher level of safety assurance during training and achieves higher rewards the majority of tasks.
26

---

Q-Supervised Contrastive Representation: A State Decoupling Framework for Safe Offline Reinforcement Learning
Reward (SDQC)
Reward (FISOR)
Cost (SDQC)
Cost (FISOR)
Safety-Gyymnasum--Ponneeoo
0.8
4.0
0.7
0.30
Normaized Reward
0.6
Normalized Cas
Nomal Roward
0.25
Nomalized Cast
0.5
0.20
2.5
0.4
2.0
2.0
0.3
1.5
1.5
0.2
1.0
0.05
0.00
1.0
0.1
0.5
-0.05
0.5-
0.0
0.0
100
200
300
400
500
100
200
300
400
500
100
200
300
400
500
100
200
300
400
500
Gradient steps (x10%)
Grodient steps (x10%)
Gradlent steps (x10*)
Gradient steps (x101)
Safety- Cyynnaalum-Ponnteeeeee
0.30
3.5
Normllized Reward
0.25
Hormalized Cast
3.0
Aomallzed Roward
0,25
3.0
0.20
2,5
0.20
2.5
0.15
Homalized
0.15
2.0
2.0
1.5
0.05
1.5
0.00
-0.05
-0.05
100
200
300
400
500
100
200
300
400
500
-0.10
100
200
300
400
500
100
200
300
400
500
Gradient
Gradient
(x10º)
Gradient (x10f)
Gradi
(x107)
Safety-Gymnesium
0.15
4.0
0.15
4.0
3.5
3.5
Normalzeed Reward
0.10
mellieddd Cost
3.0
Mormaizzdd Reward
0.10
3.0-
0.05
2.5
0.05
2.5
2.0
0.00
0.00
-0.05
-0.05
-0.0
100
200
300
400
0.30
300
400
500
0.0
staps (X10P)
Gradient steps (x10P)
Gradient steps (x10*)
ent steps (x110²)
Figure Training curves of SDQC and FISOR on the "Point" agent with tasks "Goal," "Push," and "Button" in the domain.
Reward (SDQC)
Reward (FISOR)
Cost (SDQC)
Cost (FISOR)
SafetyyGymmasium-HHllCCCe HaaChhheea
Sifety-Gymnasium-Ant
2,00
2.00
Normalized Reward
Normalize Reward
1.50
1.50
0.6
0.6
1.25
0.4-
0.4-
0.2
0.50
0.2
0.50
0.25
0.25
0.0
0.0
0.00
0.00
100
200
300
400
500
100
200
300
400
500
100
200
300
400
500
100
200
300
400
500
Gradient steps (x102)
Gradient steps [x10")
Gradient stops (x10")
Gradient steps (x10?)
0.30
Normalized Roward
0.05
N
Let
1.50
1.75
0.00
1.25
1.00
0.75
0.50
-0.15
0.25
-0.20
Sáo
400
500
Gradient steps (x10?)
Gradient steps (x10?)
Figure 13. Training curves SDQC and FISOR on the agents "HalfChectah", "Ant", and "Swimmer" with velocity constraints task in the Sefety-Gymmasium domain.
27

---

Q-Supervised Contrastive Representation: A State Decoupling Framework for Safe Offline Reinforcement Learning
Reward (SDQC)
Reward (FISOR)
Cost (SDQC)
Cost (FISOR)
Safety-GymnasiiumBBan Ballirrr
Safety-Gymmnsiuu-Ban Balllnu
0.5
2.00
0.30
2:00
1.75
Nomallze Reverd
CoR
1.50
Normalized Reward
M
1.50
Nomalizv
1.25
Horedllsed
0.2
1.00
0.05
0.1
0.50
0.00
0.0
0.25
-0.05
0.00
-0.10
0.00
100
200
300
400
500
Gradient steps ix
Gradient steps (x109)
Gradient steps (x13')
Gradient steps (x10%)
Sefety-Gymnasium-CrCCii
Safety-Gymnasium-Garrf
2.00
0.5
1.75
0.8
1.75
Normaized Reward
0.4
Hermalized Coa
1.50
Normalized Reward
0.6
1.50
1.25
1.25
1.00
0.4
Normalized
1.00
0.2
0.75
0.75
0.50
0.2
0.50
0.0
0.25
0.25
-0.1
100
200
300
400
500
0.00
100
200
300
400
500
100
200
300
400
500
0.00
100
200
300
400
500
Gradient
Gradient (x10%)
Gradiant
(x102)
Gradient
(x10²)
Safety.Gyymasium-AAntiie
Safety- Gymnaslum-Anter
0.5
2.00
0.5
2.00
Nomalzed Roward
0.4
1.50
1.75
Normalized Reward
0.4
1.75
1.50
1.25
0.3
Normalized
1.25
1.00
1.00
0.75
0.1
0.0
0.0
0.00
-0.1
300
400
500
Gradient steps (x10%)
Gradient steps (x10%)
steps (x16²)
steps [x10P)
2.00
0.5
2,00
1.75
Normaized Reward
0.4
Nemalzed Codl
Nomalize Reward
1.50
0.3
1.25
1.25
Nomalized
0.2
1.00
0.2
0.75
0.75
0.1
0.1
0.50
0.50
0.25
0.0
0.25
-0.1
100
200
300
400
500
0.00
100
200
300
400
500
100
200
300
400
500
0.00
100
200
Gradient (X30%)
Gradient (x10%)
Gradient steps (x10")
Gradient steps (x10')
Figure 14. Training curves of SDQC and FISOR on the agents "Ball", "Car", "Ant", and "Drone" with tasks "Circle" and "Run" in the Bullet-Safety domain.
G. Limitations and Future Works
One limitation our current study arises from the substantial computational demands associated with training the SDQC model. This particularly notable due to the necessity executing three distinct training phases and the utilization of complex network architectures certain scenarios. Despite this challenge, the remarkable cost-effectiveness and robustness to seed variance exhibited by our model mitigate these weaknesses. Looking ahead, future research endeavors will prioritize the optimization of the training pipeline and the simplification of network structures to enhance training efficiency while maintaining performance standards.
28


Reviews:
[{"review_id": "HBFXgY6mFK", "reviewer": "Reviewer_cp6m", "summary": "This paper introduces State Decoupling with Q-supervised Contrastive representation (SDQC), a framework that decouples the global observations of an RL agent into reward- and cost-related representations to deal with OOD data and improve generalisation.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely acknowledge Reviewer cp6m for providing the insightful feedback. Please see our response to your concerns below.\n## (Claim1) Motivation: How are safe offline datasets collected?\nIn 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.\n\n## (Claim2) How is it possible to decouple the representations according to $Q^{\\*}$?\n\nWhen 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**.\nLikewise, the same principle applies when considering only the cost.\nFor further intuition, please refer to the GIF in the supplementary material, which visualizes policies derived from the decoupled states.\n\n## (Claim4) What does \"generalization\" refer to?\n\nWe 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.\n\nIn real-world decision-making, state spaces are often high-dimensional and infinite, making it impossible to cover all possible states during training phase.\nThis is especially true in offline settings, where the agent is trained on a fixed dataset without exploration [1].\nThe 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.\n\nMoreover, 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. \n(For online RL, generalization is usually evaluated in a similar way [2].)\n\n## (Question1) Applying SDQC to online settings.\n\nWe 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.\nOur 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.\n\n## (Question2) Possibilities of abstracting representations from pre-trained value functions.\n\nIf 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.\nWe argue that this approach would make the training process more complex and prone to failure.\n\nSpecifically, 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$.\nHowever, 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.\n\nIn 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.\n\n[1] Levine, Sergey, et al. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv:2005.01643 (2020).\n\n[2] Agarwal, Rishabh, et al. Contrastive behavioral similarity embeddings for generalization in reinforcement learning. ICLR (2021).", "claims_and_evidence": "1) The authors claim that \"Safe offline RL [...] provides a promising solution that learns the safety-guaranteed policy in a fully offline manner. Its training requires no risky interaction with the environment and relies only on the precollected offline dataset.\"\n\nI am not familiar with offline safe RL, but it looks like the problem has only been moved, since how is the precollected offline dataset of execution going to be collected?\n\n2) \"The representations solely capture either reward or cost information, independent of another factor, as determined by the training of Q*.\"\n\nHow is it possible to decouple this information determined by the training of Q*?\n\n3) The authors claim superior experimental performance, which looks backed by empirical evidence.\n\n4) On p. 5 the authors say that \"our Q-supervised contrastive learning method theoretically surpasses bisimulation in terms of generalization.\"\nBut what it is meant by generalisation at this point? So far generalization has only been discussed informally in the introduction.", "methods_and_evaluation_criteria": "The empirical evaluation is rather thorough. It considers 5 other relevant methods in the benchmark DSRL.\n\nAs states above, I am not sure that offline learning is the best solution for safety in RL, as the agent has to interact with an unsafe environment to collect experiences.", "theoretical_claims": "Proofs are provided in the appendix. I did not check all details, but the theoretical claims look correct.", "experimental_designs_or_analyses": "The experiments look fine in terms of baselines and benchmarks for comparison.", "relation_to_broader_scientific_literature": "The authors compare frequently to [Zheng et al., 2024], as one of the first examples of application of Hamilton-Jacobi reachability for safety analysis. However, the empirical results demonstrate the inability of their algorithm (FISOR) to achieve absolute safety guarantees.\nIt looks like the main innovation wrt [Zheng et al., 2024] is the decoupled representation, which leads to a superior experimental performance.", "essential_references_not_discussed": "It doesn't look like.", "other_strengths_and_weaknesses": "1) The authors stress the relevance of OOD generalization for their contribution, but then generalization is only discussed in Sec. 4.2, rather marginally.\nI wonder whether a stronger case for the paper can be made by simply considering the improvement in performance.\n\n2) The authors manage to fit safe (offline) learning in a clear and concise manner in a single column, which is a feat in itself.", "other_comments_or_suggestions": "1) The example in the introduction does not really add much to the discussion, as I'd assume that most people are familiar with issues with OOD data and generalisation.", "questions_for_authors": "1) To what extent your method can be applied to online learning?\n\n2) The goal of the authors is to achieve a decoupled representation of rewards and costs. However, can't this goal be achieved by defining corresponding Q and value function for costs and rewards? These could then be used directly in their framework.\n\nAlso, please answers points (1), (2), and (4) in the Claims section.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "LpRjEXe07u", "reviewer": "Reviewer_zA7P", "summary": "This paper introduces a framework \"State Decoupling with Q-Supervised Contrastive Representation\" (SDQC) to improve safe offline reinforcement learning. They do this by using two distinct state representations (the \"decoupling\"), one being for rewards and one for costs, and learn these using contrastive objectives. They then use these representations to extract three policies: a \"reward policy\", \"cost policy\", and \"tradeoff policy\", which they train using diffusion models. In practice, they use a safety assessment function on the two representations of a given observation, which dictates which policy is used at that timestep. They evaluate their method against a number of baselines across the DSRL benchmark, as well as further experiments to compare generalization capabilities of different algorithms.\n\n\n--- \n\nNot sure where I'm supposed to acknowledge the rebuttal so I guess I'll do it here. \n\nI thank the authors for a strong rebuttal, it cleared up the uncertainties I had and I will adjust my score to an accept. One thing I will note is that I'm not a fan of the term ``reverse expectile regression'' in its use as the authors described, but if it's used in prior work I can't fault them for it.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We acknowledge Reviewer zA7P for appreciating the significance and novelty of our work. Below, we respond to the your questions point by point.\n\n## (References) Missed discussing on relevant references\n\nThe 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.\n\n## (Weakness1) Complex hyperparameter settings\n\nWe 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).\n\n## (Weakness2/Question2) Computational costs\n\nThe computational costs for SDQC are indeed slight higher than other baseline algorithms. We discussed this point in Appendix E.3, lines 1288-1302.\nTo 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. \nDuring the training phase, the total training time (hours) are presented as follows.\n\n|BCQ-Lag|CPQ|COptiDICE|CDT|TREBI|FISOR|SDQC(ours)|\n|-|-|-|-|-|-|-|\n|0.46|0.45|0.41|3.12|16.62|1.77|5.43|\n\nDuring the testing (inference) phase, we measure the time consuption (seconds) over 1000 RL timesteps as follows.\n\n|BCQ-Lag|CPQ|COptiDICE|CDT|TREBI|FISOR|SDQC(ours)|\n|-|-|-|-|-|-|-|\n|1.51|1.85|1.86|3.45|585.87|6.11|11.13|\n\n## (Weakness3/Question4) Explanations on \"conditional entropy\"\n\nWe 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.\nConditional entropy is a concept in information theory, and its formal definition was introduced by Shannon[1]:\n\n$H(s|z)=-\\sum_{s,z} p(s,z) \\log p(s|z)$\n\nwhere $p$ represents the probability function.\nAccording 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:\n\n$H(s|z) = H(s,z) - H(z) = H(s) - H(z)$ \n\nwhere 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)$**.\n\nWe will add above detailed explanation in the final version of our paper.\n\n## (Typos)\n\nThanks for pointing out these! We have carefully checked entire paper to correct all typos.\n\n## (Question1) What is \"reverse expectile regression\"\n\nThe 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\".\n\n## (Question3) What is Safe Agents and Unsafe Agents in Table 1\n\nAs 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**.\n\n**Lastly, and importantly, we are committed to open-sourcing all code upon acceptance.**\n\n[1] Shannon, C. E. A mathematical theory of communication. The Bell system technical journal (1948).\n\n[2] Liu, Guoqing, et al. Return-based contrastive representation learning for reinforcement learning. ICLR (2021).\n\n[3] Zheng, Yinan, et al. Safe Offline Reinforcement Learning with Feasibility-Guided Diffusion Model. ICLR (2024).", "claims_and_evidence": "I believe that the claims made are mostly supported by clear and convincing evidence (I have a couple questions on some theoretical aspects below, but I believe these should be squashable in the rebuttal).", "methods_and_evaluation_criteria": "I believe the proposed methods and evaluation criteria make sense.", "theoretical_claims": "I checked all proofs and theoretical claims made. I had some questions which I highlighted in the sections below.", "experimental_designs_or_analyses": "Their experiments on the DSRL benchmark appears sound to me (although no code was provided).", "relation_to_broader_scientific_literature": "This paper is at the intersection of safe RL, offline RL, and representation learning in RL. It borrows some techniques from previous work in safe offline RL (such as the use of expectile losses, using regressed diffusion models to train the policies, and using a diffusion behaviour cloner. They primarily build on this work by improving representation learning in a novel way.", "essential_references_not_discussed": "The authors state that their method is a useful alternative to bisimulation as it doesn't require model-based methods in order to estimate. However there is a line of work on estimating bisimulation-type distances from samples in a model-free setting, for example beginning with [1] from NeurIPS 2021 and a number of works which followed.\n\n[1] Pablo Samuel Castro, Tyler Kastner, Prakash Panangaden, and Mark Rowland. MICo: Learning improved representations via sampling-based state similarity for Markov decision processes. In Advances in Neural Information Processing Systems, 2021.", "other_strengths_and_weaknesses": "**Strengths**\n- The method introduced appears novel and non-incremental (at least to my knowledge), and obtains strong results against baselines on the environments considered.\n- I find Section C of the appendix to be really great to help the reader understand and build intuition for the algorithm (pseudocode + architecture diagrams).\n- The GIF in the supplementary material is super helpful to build intuition and understand the 3 policies working together (this could've been advertised to the reviewers better! I came across this accidentally).\n\n\n**Weaknesses**\n- The proposed algorithm seems quite complex with a number of moving parts, which may make it brittle with respect to hyperparameter choices. (I appreciate that the authors demonstrate that this isn't the case with respect to the diffusion hyperparameters at least).\n- Similar to the above, I can imagine the number of moving parts may make it run potentially slow (multiple diffusion models, nearest neighbour search, contrastive objective, etc).\n- Some of the theory presented in Section 3.4. feels a bit hard to follow for me and I do not entirely see the motivation (I don't \"type-check\" the variant of cross-entropy they use, and I also don't see how this relates to generalization). I expand on this in my comments and questions in the sections below.", "other_comments_or_suggestions": "- I think the cross-entropy $H(s| z_\\theta(s))$ should be defined prior to its use in Equation (4). This definition confuses me, as cross entropy is normally defined with respect to two probability distributions, but here your arguments are a state and the state's representation.\n- Following the above: the way $H(s| z_1)$ is used in the proof of Proposition A.4. seems to be a variation of the KL divergence rather than cross entropy. \n- It is stated a number of times that a coarser representation has better generalization properties. Perhaps briefly arguing why this might be the case, or citing references that go into detail on this fact, can be helpful.\n\nMinor typos/nits:\n- Lemma A.3: \"finner\" -> \"finer\".\n- For quotation marks, use ``<>'' instead of ''<>'' (otherwise the opening one is reversed).\n- A representation being finer than another is defined but a representation being coarser than another is used without definition.", "questions_for_authors": "- What is \"reverse expectile regression\" in **Cost-related Representation** of Section 3? It seems like the standard expectile loss?\n- Following my second weakness: how does the wall-clock time of your method compare to alternative methods?\n- In Table 1, what is a \"Safe agent\" vs a \"Unsafe agent\"?\n- Following my comments in the previous section, what definition of cross-entropy are you using in the terms $H(s|z)$?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "tUo2VxQRUJ", "reviewer": "Reviewer_JBF8", "summary": "This paper proposes a new safe offline RL algorithm called SDQC, the key idea of the algorithm is to decouple observations into reward and cost-related representations, and use them in a value & policy learning framework similar to the one proposed in FISOR. The key intuition is that the values related to reward and cost may depend on different state representations, and decoupling them can help enhance generalization and safety. The proposed method achieves impressive performance in the DSRL benchmark. However, the algorithm is also quite heavy, and needs to learn two distinct state representations plus three different diffusion policy models.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank Reviewer JBF8 for the valuable suggestions. We would like to address your concerns point by point below.\n## (Method1) Potential instability for Eq. 5\nIn Eq. 5, the exponential coefficient $\\Gamma(s,\\tilde{s})$ is a detached soft measure introduced in [1] that helps stabilize training.\nRegarding the potential instability caused by coupling value function updates with representation learning, a detailed analysis is provided in Appendix F.2 (pages 25–26).\nExperimental results demonstrate that the contrastive term has no adverse effect on value estimation, while our proposed ATN-based representation network significantly reduces estimation error.\n## (Method2) Why limited anchors lead to better results\nWe 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.\nSince 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.\n## (Method3) Upper-bound of cost value function\nPlease 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)$.\n## (Theory1) Why introduce a comparison with \"bisimulation\"\nWe 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.\nHowever, 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].\nIn 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.\n## (Theory2) Contrastive representations do not have good properties for decision-making.\nIn 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. \nWe 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.\n## (Theory3) Why conditional entropy can indicate \"generalization ability\"\nDue to space limitations, please refer to our response to Reviewer zA7P on **Explanations of \"conditional entropy\"**.\n## (Experiment1) \"Fair comparison\" with baselines.\nThe required \"fair comparison based on same computational cost\" is hard to execute. \nFirst, each baseline algorithm has unique optimal parameter settings suggested by their authors.\nMoreover, some baselines, such as CDT with GPT-2 backbone and TREBI with U-Net structure, are even more computational expensive than our SDQC.\n\nAccording to reviewer's requirement, we implement FISOR with ATN-based structure and separate diffusion policies [**[here]**](https://github.com/wrk8/SDQC-Rebuttal/blob/main/to_reviewer_JBF8/compare_FISOR_SDQC.png).\n## (Experiment2) Temperature ablations for diffusion policies\nPlease find ablations [**[here]**](https://github.com/wrk8/SDQC-Rebuttal/blob/main/to_reviewer_JBF8/Ablation_diffusion_temp.png). The default parameters follow FISOR and are kept consistent across all environments.\n## (Question) Consistent hyperparameters for all environments\nWe would like to emphasize that our choice of hyperparameters is domain-specific. \n(We use the same hyperparameters across 12 tasks in safety-gym domain.)\nPlease refer to our response to Reviewer zA7P on \"Complex hyperparameter settings\".\nWe argue that **no RL algorithm can employ a unified set of hyperparameters across all domains and environments**, including most commonly used SAC and PPO.\n\n[1] Agarwal, Rishabh, et al. Contrastive behavioral similarity embeddings for generalization in reinforcement learning. ICLR (2021).\n\n[2] Kostrikov, Ilya, et al. Offline Reinforcement Learning with Implicit Q-Learning. ICLR (2021).\n\n[3] Castro, Pablo. Scalable methods for computing state similarity in deterministic markov decision processes. AAAI (2020).\n\n[4] Lee, Vint, et al. DreamSmooth: Improving Model-based Reinforcement Learning via Reward Smoothing. ICLR (2024).", "claims_and_evidence": "The paper claims that it can achieve strong safety and generalization performances, which are well-supported by the DSRL benchmark results and additional generalization tests.", "methods_and_evaluation_criteria": "**Regarding Method:**\n- In Equation 5, the contrastive loss includes two exponential terms, $exp(s, \\tilde{s})$ and $exp(z, \\tilde{z})$, which might lead to instability during training. Since training is coupled with the value function update, this instability could be further amplified, especially in applications using image inputs.\n- The authors use only a subset of the dataset for training, whereas contrastive learning typically requires a large amount of training data. It is unclear why using a limited number of anchor states would lead to better performance. Clarification on this point would be helpful.\n- The authors introduce an upper-bound cost value function but lack sufficient theoretical analysis to explain its significance. Since this function is used to identify safe conditions, further clarification is needed.\n\n**Regarding Evaluation:**\n- Using the DSRL benchmark and additional generalization tests is reasonable.", "theoretical_claims": "I do not think the theoretical analysis in Section 3.4 has much relevance to the proposed algorithm. Theorem 3.3 only discusses the properties of bisimulation representations and $Q^*$-irrelevance representation. The proposed contrastive representation has nothing to do with bisimulation representation, and there is also no formal proof or discussion given to show the proposed representation is a $Q^*$-irrelevance representation. Hence Theorem 3.3 is completely irrelevant. The authors then use Theorem 3.3 and the entropy maximization design (Eq.(4)) to claim the proposed method has better generalization (has larger conditioned entropy), which in my opinion, is seriously flawed. First, whether Theorem 3.3 holds for your representation is questionable; second, why does larger conditioned entropy justify better generalization?\n\nI suggest removing Section 3.4 and the \"theoretical analysis\" entirely from the paper, as it is irrelevant and misleading, and cannot provide convincing arguments to demonstrate the superiority of the proposed method. If the authors really want to keep the theoretical analysis in the paper, at least the authors should:\n- Rigorously prove the proposed contrastive representation is a $Q^*$-irrelevance representation, which I personally believe is impossible, as contrastive representations typically don't have many good properties for decision-making problems.\n- Use a more rigorous generalization metric for discussion, instead of using the magnitude of conditioned entropy.\n- Lastly, even if the current logic of the theoretical analysis holds (if both the above issues can be proved), the authors can only show their method is better than bisimulation representations, which is a rather weak conclusion and does not offer much real-world relevance.", "experimental_designs_or_analyses": "The evaluation is generally reasonable, but there are also some aspects that can be further improved:\n- SDQC utilizes **two sets of representations, three diffusion policies, and transformer-based (ATN) encoder** during training, making it very complex and computationally expensive. The authors should compare SDQC with baselines with similar scales of parameters, training time, and inference time. Additionally, given the similarity between FISOR and SDQC, it would be valuable to evaluate FISOR's performance using the same transformer encoder and separated diffusion policy model. Furthermore, testing FISOR with representation learning could help isolate and highlight the contribution of state-decoupled representation learning.\n- Adding more ablation and discussion on temperature hyperparameters $\\iota_r, \\iota_h, \\iota_{to}$ in Eq. (14).", "relation_to_broader_scientific_literature": "Safe offline RL methods are relevant to industrial control, robotics, and autonomous driving.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "**Strengths:**\n- SDQC focuses on the hard constraint setting and achieves state-of-the-art performance.\n\n**Weaknesses:**\n- See my above comments in the Methods, Theoretical Claims, and Experimental Designs sections.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "In Table 4, different parameter settings are used for different tasks. How does the performance compare when using the same parameters across all tasks?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "nYMJJ0Qydt", "reviewer": "Reviewer_y1rM", "summary": "This paper introduces SDQC, a method designed to address the OOD problem in safe offline RL. By decoupling states into two independent representations, SDQC improves decision-making, safety, and generalization when facing unseen states. The method employs Q-supervised contrastive learning to guide the learning of state representations based on Q-values and selects policies based on safety evaluations. Experimental results show that SDQC outperforms baseline algorithms, achieving zero violations in most tasks. The paper also provides theoretical evidence that SDQC generates coarser representations while still preserving the optimal policy.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank Reviewer y1rM for providing the positive feedback and recognizing the importance of our work. Please see our response to your questions below.\n\n## (Weakness/Question1) More Generalization Experiments on \"Button\"\n\nWe 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. \n\nIn 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]**](https://github.com/wrk8/SDQC-Rebuttal/blob/main/to_reviewer_y1rm/CarButton.gif).)\n\nThe 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.\nConversely, 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.)\n\nWe will include a more detailed discussion of this phenomenon in the final version of the paper.\n\n## (Question2) Visualization results of difficulty-level.\n\nThank you for your insightful comments. We agree that clear visualizations can better illustrate the generalization performance of our SDQC.\n\nIn 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]**](https://github.com/wrk8/SDQC-Rebuttal/blob/main/to_reviewer_y1rm/PointGoal.gif). 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.\n\nIn 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]**](https://github.com/wrk8/SDQC-Rebuttal/blob/main/to_reviewer_y1rm/CarPush.gif). 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.", "claims_and_evidence": "The claims in the paper are supported by clear and convincing evidence. This theoretical analysis supports SDQC’s ability to better handle OOD scenarios and perform well in unseen environments, which is validated through experimental results. Moreover, the theoretical analysis shows that SDQC can effectively separate these two aspects, allowing precise decisions under different safety conditions.", "methods_and_evaluation_criteria": "The proposed method and evaluation standards are highly relevant to current problems and applications, particularly in the field of safe offline RL. The SDQC method addresses the challenge of OOD problems, which is crucial in safety-critical applications like autonomous driving. The use of the DSRL benchmark dataset allows comprehensive evaluation across various tasks, demonstrating SDQC's ability to achieve zero violations.", "theoretical_claims": "The theoretical analysis is clear and the proof is rigorous. Theorem 3.1 highlights the relationship between bisimulation and Q*-irrelevance representations in MDPs, showing that the latter is a coarser but equally effective representation for deriving the optimal policy.", "experimental_designs_or_analyses": "The comparative experiments with other baselines on the DSRL benchmark are sufficient, showing safer performance in all tasks. \nThe generalization experiments involved Goal and Push tasks of different difficulties, showing that SDQC achieves safer behavior and good rewards.\nThe ablation experiment is very thorough, involving not only the ablation of Q-supervised contrastive learning but also the ablation of contrastive learning parameters. The experiment fully demonstrates the impact of the contrastive learning module on the performance of the SDQC algorithm.", "relation_to_broader_scientific_literature": "This article innovatively combines contrastive learning with offline safe RL, which greatly improves the performance of existing algorithms and is of great significance to the development of offline safe RL algorithms.", "essential_references_not_discussed": "No.", "other_strengths_and_weaknesses": "Strengths:\n1. For the first time, Q-supervised contrastive learning is integrated into offline safe RL.\n2. On the DSRL benchmark, comparative experiments on all tasks fully demonstrate the excellent security performance of the SDQC algorithm.\n3. Ablation experiments on contrastive learning are sufficient and complete.\n\nWeaknesses:\n1. The generalization experiment is relatively simple, only dealing with Goal and Push tasks, without Button tasks, and does not vividly demonstrate the changes in task difficulty.", "other_comments_or_suggestions": "No.", "questions_for_authors": "1. Can provide other generalization experiments that are not Goal and Push tasks?\n2. Can visualize scenarios of different difficulties of these tasks to illustrate the generalization performance?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This work proposes a method for safe offline RL that decomposes states into reward- and cost-specific representations. They theoretically contrast their representation learning approach with bisimulation-metrics-based approaches, and prove that they still maintain optimality. The empirical evaluations are rather thorough, and appendix C is quite commendable for its efficacy in explaining the method.

While there are some concerns about the complexity of the method, the ablations provided suggest it is relatively robust to changes to them.

Most reviewers are supportive of acceptance, as am I.
Decision: Accept (poster)