GAS: Enhancing Reward-Cost Balance of Generative Model-assisted Offline Safe RL
Authors: Zifan LIU
Xinran Li
Shibo Chen
Jun Zhang
Venue: ICLR 2026 Poster (2026)
Keywords: offline safe reinforcement learning
constrained decision transformer
improving stitching ability
Primary area: reinforcement learning

Abstract: Offline Safe Reinforcement Learning (OSRL) aims to learn a policy that achieves high performance in sequential decision-making while satisfying safety constraints, using only pre-collected datasets. Recent works, inspired by the strong capabilities of Generative Models (GMs), reformulate decision-making in OSRL as a conditional generative process, where GMs generate desirable actions conditioned on predefined reward and cost return-to-go values. However, GM-assisted methods face two major challenges in constrained settings: (1) they lack the ability to ``stitch'' optimal transitions from suboptimal trajectories within the dataset, and (2) they struggle to balance reward maximization with constraint satisfaction, particularly when tested with imbalanced human-specified reward-cost conditions. To address these issues, we propose Goal-Assisted Stitching (GAS), a novel algorithm designed to enhance stitching capabilities while effectively balancing reward maximization and constraint satisfaction. To enhance the stitching ability, GAS first augments and relabels the dataset at the transition level, enabling the construction of high-quality trajectories from suboptimal ones. GAS also introduces novel goal functions, which estimate the optimal achievable reward and cost goals from the dataset. These goal functions, trained using expectile regression on the relabeled and augmented dataset, allow GAS to accommodate a broader range of reward-cost return pairs and achieve a better tradeoff between reward maximization and constraint satisfaction compared to human-specified values. The estimated goals then guide policy training, ensuring robust performance under constrained settings. Furthermore, to improve training stability and efficiency, we reshape the dataset to achieve a more uniform reward-cost return distribution. Empirical results validate the effectiveness of GAS, demonstrating superior performance in balancing reward maximization and constraint satisfaction compared to existing methods.

Published as a conference paper at ICLR 2026
GAS: ENHANCING REWARD-COST BALANCE OF GEN- ERATIVE
OFFLINE SAFE RL
Zifan Liu', Xinran Li', Shibo C Jun Zhang',
1 The Hong Kong University of Science and Technology, Hong Kong SAR, China
2 South China University of Technology, Guangdong, China
ABSTRACT
Offline Safe Reinforcement Learning (OSRL) aims to learn a policy that achieves high performance in sequential decision decision-making while satisfying safety constraints, using only pre-collected datasets. Recent works, inspired by the strong capabilities of Generative Models (GMs), reformulate decision-maker
in OSRL as a condi- tional generative process, where GMs generate desirable actions conditioned on predefined reward and cost return-to-go values. However, GM-assisted methods face two major challenges in constrained settings: (1) they lack the ability "stitch" optimal transitions from suboptimal trajectories within the dataset, and (2) they struggle to balance reward maximization with constraint satisfaction, partie- ularly when tested with imbalanced human-specified reward-cost conditions. To address these issues, we propose -oss-sss Stitching (GAS), a novel algorithm designed to enhance stitching capabilities while effectively balancing reward max- imization and constraint satisfaction. To enhance the stitching ability, GAS first augments and relabels the dataset at the transition level, enabling the construction of high-quality trajectories from suboptimal ones. GAS also introduces novel goal functions, which estimate the optimal achievable reward and cost goals from the dataset. These goal functions, trained using expectile regression on the relabeled and augmented dataset, allow GAS to accommodate broader range of reward- cost return pairs and achieve a better tradeoff between reward maximization and constraint satisfaction compared human-specified values. The estimated goals then guide policy training, ensuring robust performance under constrained settings. Furthermore, to improve training stability and efficiency, we reshape the dataset achieve a more uniform reward-cost return distribution. Empirical results validate the effectiveness of GAS, demonstrating superior performance in balancing reward maximization and constraint satisfaction compared to existing methods.
INTRODUCTION
Significant progress has been achieved in Reinforcement Learning (RL) that learns policies to maximize rewards through constant interactions with the environment. Limited by the trial-and-error approach, standard RL often fails in scenarios with safety constraints, such as autonomous driving (Fang et al. (002) and investment portfolios (Lee & Moon 2023). It stems from two issues: (1) the exploration process in RL can be inherently risky, as random actions may lead to unsafe outcomes, and (2) the optimized policies often focus solely on maximizing cumulative rewards, neglecting safety constraints. To tackle these issues, Offline Safe RL (OSRL) has emerged as promising paradigm that learns safe policies from pre-collected dataset, eliminating the need for risky online exploration.
While OSRL mitigates the risks associated with online exploration by relying on static datasets, introduces new challenge: the -itttt--Distiitttion (OOD) problem. Specifically, the Bellman backup may extrapolate actions beyond the dataset, leading to unpredictable or unsafe behaviors. Early OSRL methods primarily adapted RL techniques to constrained settings using approaches such as the Lagrange method (Stooke et al. 2020), constraint penalty (Xu et al., 2022), or the DICE-style techniques (Lee et (002)) To mitigate reD-related issues, these methods incorporate strategies like distribution correction (Lee et al., (0022), regularization (Kostrikov et al. (021). and OOD detection (Xu et al. (022). However, they still face challenges effectively addressing OOD problems and adapting to dynamic, real-world constraints.
Both Shiho Chen (shihochen.usteee) gaail.eommmmmmmmmmmmmm and Jun Zhang ceejzhng@uusthil are corresponding authors.

---

Published as a conference paper at ICLR 2026
More recently, Generative Model-assisted (GM) methods have emerged an alternative to traditional OSRL approaches to address these limitations. In OSRL, GM methods reformulate the Constrained Markov Decision Process (CMDP) as conlicondi generating problem, where the generative model is trained to produce trajectories that match predefined reward and cost returns specified as inputs. This formulation provides two key benefits. First, GM methods essentially adopt goal- conditioned behavior cloning scheme, where the model learns to imitate behavior from the dataset conditioned on the desired reward and cost targets. This formulation entirely bypasses the Bellman backup procedure, which is the primary source of the OOD problem in traditional OSRL approaches. Second, GM methods also offer greater flexibility compared to conventional methods, which typically operate under fixed safety constraints. Because the target reward and cost returns are provided as inputs, GM methods can seamlessly adapt to varying objectives at test time without re-training.
Nevertheless, GM methods present their own challenges. First, while GM methods bypass the Bellman backup procedure, they lack trajectory stitching capabilities (Badrinath et al. [023]] Wu et al. 2023 Kim et al., 2024b) the ability to combine transitions from different trajectories to enhance performance. This limitation restricts their ability to fully utilize suboptimal datasets to improve performance. Second, GM methods lack explicit mechanism to balance reward maximization and constraint satisfaction, potentially leading to unsafe or overly conservative policies. In this work, we propose a novel Goal-Assisted Stitching (GAS) method to address these challenges while retaining the advantages of GM methods. The key contributions are summarized as follows.
We propose use goal functions as intermediate values bridge the gap between human-specified (potentially suboptimal) reward-cost targets and the optimal achievable goals of the conditional policy instantiated by GM methods. To achieve this, we introduce three key innovations: 1) Temporal Segmented Return Augmentation and Transition-leve Return Relabeling: To enhance GAS's stitching capabilities, we restructure the offline dataset at the transition level and introduce temporal segmented return augmentation, which extracts richer information by considering reward and cost returns over varying timesteps rather than only at trajectory endpoints. Additionally, we ensure ro- target return-to- to-goes during testing by relaheling reward-cost
bustness to suboptimal human-specing return-to souur-tt--gos at the transition level during training. 2) Goal Functions with Expectile Regressions: We train the novel reward and cost goal functions using expectile regression to estimate the optimal achievable reward and cost goals without relying on Bellman backups. These goal functions guide the policy to stitch transitions effectively, achieving both reward maximization and constraint satisfaction for a wider range of given target reward-cost return pairs. 3) Dataset Reshaping: To improve training stability and efficiency, we address the data imbalance issue by reshaping the dataset create more balanced reward-cost return distribution. Through extensive experiments on 2 benchmarks with 12 scenarios and 8 baselines under various constraint thresholds, GAS shows superior safety ability under tight thresholds and 6% improvement in performance under loose thresholds.
2 RELATED WORK
RL Methods in Offline RL. Offline RL (Fujimoto et al. 2019) aims to find the policy to maximize the cumulative rewards from pre-collected dataset. The primary challenge lies in the OOD problem during the Bellman backup, where the policy may select actions beyond the dataset. This issue arises because offline RL does not allow for environment exploration gather additional data. To this end, most existing works try to constrain the target policy to stay close to the behavior policy with a KL regularization term (Jaques et al. 2020; Peng et al. [019:] Siegel et al. (2020) or Wasserstein distance (Wu et al. (2020) where the behavior policy is estimated by generative model. For example, BCQ (Fujimoto al. 2019) uses a variational autoencoder (VAE) to estimate the behavior policy and restrict the action space during the Bellman backup. Except for explicitly estimating the behavior policy, CQL (Kumar et al. 2020) proposes to learn Q-values conservatively, encouraging those within the dataset to act as lower bound. Alternatively, IQL (Kostrikov et al.] 2022) avoids the OOD problem entirely by employing expectile regression to learn Q-values without querying OOD actions.
GM Methods in Offline RL. More recently, some works formulate offline RL as return-conditioned generation problem and address it using generative models, such as Decision Transformer (DT) (Wul
al. 2023 Chen et al. 2021 Zheng et al. 2022), and Decision Convformer (DC) (Kim et al. 2024a). These methods naturally avoid the OOD problem since the return- condiiooiii generation problem does not need the Bellman backup procedure. However, significant limitation of these methods is their limited stitching ability, which refers to the ability to find better policies beyond the dataset trajectories, particularly when the dataset consists primarily of suboptimal trajectories. This definition
2

---

Published as a conference paper at ICLR 2026
is a little different from the stitching ability in RL since GM methods try avoid the Bellman backup procedure. To improve the stitching ubility of GM methods, several works have been proposed. For instance, QDT Yamagata al. (023a) relabels return-to-go in the dataset with Q-values derived from RL methods. WT (Badrinath et al. 2023) proposes to use a sub-goal as the prompt to guide the DT policy to find shorter path in navigation problems. Building on these ideas, ADT (Ma et al. 2024) proposes a hierarchical framework by replacing the return-to-go with sub-goal learned from IQL. Reinformer (Zhuang et al.) 2024) generates RTG based on the state using expectile regression and achieves strong reward stitching capabilities. In contrast to improving stitching ability in the training procedure, EDT (Wu et al., 2023) boosts the stitching ability of DT at decision time by adaptively adjusting the history length of the attention module.
Offline Safe RI. OSRL, integrates safety constraints into offline learning settings, addressing both safety requirements and limited online interaction with the environment. This emerging field has spawned two main approaches: RL methods and GM methods. For RL methods, CPQ (Xu et al. 2022) employs a conditional VAE model to estimate and penalize the OOD actions. COptiDICE (Lee et al. 2022) utilizes stationary distribution correction to mitigate the distributional shift problem. For GM methods, which are inspired by DT, CDT (Liu et al., 2023b) transforms OSRL into goal-conditioned generative problem and inputs both target reward and cost return-to-go go to the GPT structure. To handle the safety-critiaal cases, FISOR (Zheng al. 2024) proposes feasibility-guided diffusion model to ensure striet satisfaction of constraints.
Although previous works demonstrate great success in enhancing the stitching ability of GM methods in offline RL, they cannot be utilized in OSRL directly due to the fundamentally different objectives introduced by constrained settings. Extending GM methods OSRL is challenging because, while the goal of maximizing rewards remains consistent, the requirements for constraints vary across scenarios. CQDT (Wang& Zhou 2025) tries to improve the CDT by training lots of RL policies under each constrant, which is extremely inefficient. Recently, a concurrent work, named COPDT (Xue et al., 2025) is proposed by considering the reward targets as a generation conditioned on cost targets to prioritize cost targets, focusing on multi-constraints and multi-tasks. Different from these works, we address this critical gap by introducing GAS, a novel framework designed to give up attention structure and purely focus on "transition stitching" only with simple Multi-Layer Perceptron (MLP) structure guided by robust reward and cost goal functions.
3 BACKGROUND
Constrained Markoy Decision Process (CMDP). CMDP (Altman 1998) is a standard framework for safe RL defined by the tuple M = A, P,,, c, 7), where S is the state space, is the action space, P(s'|s s, a) is the transition function, r(s, a) is the reward function, c(s, a) is the cost function, and y is the discount factor. The goal of safe RL to find a policy T(als) that maximizes the cumulative rewards while satisfying the safety constraints. Given the constraint threshold L and the trajectory 7 = {(00, do, ro, co), ..., (ST, UT, rT, (T,er)) er)laa Ja, = [aae(1))) with T being the trajectory length, the optimization problem of safe RL. can be written as:
max V (T),
(I)
s.l. (T)
where TeT = E [E 7're] denotes the reward value function, and (T) =
Er 1 n'aa] denotes the cost value function.
Offfine Safe Reinforcement Learning (OSRL). OSRL maintains the same objectives as safe RL while additionally addressing the OOD problem due to the inability to explore. To mitigate the OOD problem, the target policy is constrained to remain close the behavior policy, as formulated in eq. (2), which augments eq. (I) in the optimization.
D(r, p)) S.
(2)
where μ is the behavior policy used in the pre-collected dataset DD DD.,.)) is an arbitrary dis- tance/divergence function, and is a hyper- parameer..
Constrained Decision Transformer (CDT), CDT (Liu al. 2023b) transfers the OSRL problem into a goal-conditioned generative problem by reformulating the dataset as:
(50, ao, Ro, Co, .... ST, aT, Rr, Cr),
(3)
3

---

Published as a conference paper at ICLR 2026
K
K-10
Threshold
Constraint
Constraint
Constraint 70
Coets Rewards
4.50
Mamarron
400
75
50-
0
20
40
60
80
20
40
60
80
20
40
60
80
Episodes (5e3)
Episodes (5e3)
Episodes (5e3)
Figure 1: Training curve of CDT with different memory length K on task CarCircle.
where R₂ r2 + rr is the cumulative rewards from to the trajectory end in the dataset, named reward return-to-go, and C. = c + ... + cr the cost return-to-go. Then CDT trains a policy x with the GPT structure, as shown in eq. to predict the action given (s, R. C)) at the current timestep and (s, a, R. C) from previous K timesteps as conditions, where is the memory length.
T[aelst, Re, Ct.... Or- K+1 St- K+1, R K+1, Ct-K+1) a.k.a T(delsts RR, Cr K), (4)
During testing, users need to specify the target reward return R and cost return é as the input of the CDT policy, which then generates actions approach these targets R and Õ.
MOTIVATION
In this section, we first analyze two critical limitations of current GM methods for OSRL: (1) insufficient trajectory stitching capabilities and (2) the inability to balance reward maximization and constraint satisfaction. We design our solution to specifically address these challenges.
4.1 INSUFFICIENT TRAJECTORY STITCHING CAPABILITIES
A key strength of traditional RL methods lies in their ability stitch together different suboptimal trajectories, enabling effective generalization across diverse experiences. This stitching ability arises from the fact that RL only considers the current state as the condition and stitches the next optimal transition via Bellman backup. In contrast, GM methods sacrifice this crucial ability by additionally taking previous information conditions to capture the temporal association among transitions. Recent studies (Badrinath et al. 2023 Ma et al., 2024] Xiao et al. 2023 Yamagata et al. 2023b) demonstrate that DTTCIDT's trajectory- level training paradigm is essentially form of goal- conditioned behavior cloning, which predicts the action based on information in previous timesteps via the temporal attention module. Given suboptimal information as the contextual condition, GM methods tend to memorize the suboptimal actions, resulting in suboptimal performance.
However, Kim et (024b) reveals that the attention module designed for natural language pro- cessing often fails to characterize the temporal relations for sequential transitions of MDPs in DT (Chen et al. 202T). Specifically, while MDP policies should primarily focus on immediate previous states, DTTs 8 attention spans up to 20 timesteps backward, potentially diluting the focus relevant temporal information and adversely affecting the performance of offline RI. (Kim al., 2024b). To validate this observation,
conducted empirical experiments comparing CDT variants with different memory lengths (K = 1 and K = 10) for attention modules while keeping other parameters constant. As shown in fig. [ CDT's performance remains largely unchanged across these memory settings under three distinct constraints, suggesting that the attention mechanism fails to effectively leverage temporal information in OSRL.
4.2 INABILITY BALANCE REWARD MAXIMIZATION AND CONSTRAINT SATISFACTION
In GM-assisted OSRL, as opposed to goal-conditioned behavior cloning, the objective should be formulated as:
max (s (arlst, RG, Ct,
s.t. C(st, (a) Rí, Ct, K)) S Cr.
(5)
R2 = R (ro + + rz-1),
C. = Č (co + c--),,
where R..) and CC.) are the non-discounted reward and cost returns under the policy π. DT-ffmily algorithms, due to their goal-conditioning nature (detailed in section 3, struggle to effectively balance
4

---

Published as a conference paper at ICLR 2026
Folicy kameed by CDT Policy leamed by GAS Stitching procedure
A Equation (10)
AWR: Addvantage-Weighe Regression
Transitions
transitions
Abr Equation (12)
ANR = ja 1Arr 0] (d, a.)
Tr
a
5
10 la- M(Arr (22)"
of= 1 < -Atr)
fer)
Loss
201
a
GAS: Goal
GAS: Policy
F2
Enander Time Embedd
Relabel
Dataset
Temporal
Reshapo
Augmentation
Input
Figure 2: Overall view of GAS. Left: Comparison between CDT and GAS. CDT optimizes at a trajectory level, while GAS enables fine- aineigrained trajectory stitching under the guidance of reward maximization and constraint satisfaction. Right: GAS's stitching mechanism, where the goal function learns the optimal reward and cost return- to-n--ttooooooo within the constraint given any target. The policy aims to take actions to achieve optimal goals estimated by the goal function via constrained AWR.
the dual objectives of reward maximization and constraint satisfaction in OSRL. This limitation stems from two key challenges: First, without prior knowledge, it becomes problematic to determine appropriate target reward-cost pairs (ft C) as these objectives often require careful trade-offs. The GPT structure can only search for policies that satisfy given targets (Liu al.) 2023b) without validating their feasibility. Consequently, specifying overly ambitious return targets R alongside stringent constraints Ö may lead to policy degradation when such conditions prove unrealistic. Second, the current CDT architecture simply concatenates target return R and cost Ó as contextual conditions without any mechanism to balance their relative importance. This structural limitation prevents GM methods from properly prioritizing constraint satisfaction over reward maximization erucial requirement in OSRL. This misalignment between architectural design and OSRL objectives frequently causes performance degradation in practice.
5 GOAL-ASSISTED STITCHING
we highlight that extending the temporal attention module to overly long time windows does not enhance memorization and instead hinders the stitching ability for GM methods in OSRL. This inspires us to trade the temporal attention for more focused approach: stitching transitions in the dataset under the guidance of reward maximization and constraint satisfaction in a conditional manner. To achieve this goal, we propose GoslsAssisted Stitching (GAS), a novel algorithm that achieve
In section 4.1
offers more flexible balance between R, and C₁ and stitches high-quality transitions safe and best performance shown in fig. 2 In particular, we first estimate the optimal achievable reward return-to-go satisfying the constraint and corresponding cost return-to-go in the dataset for given pair of RG C, with reward-cost goal function through expectile regression without relying on the Bellman backup procedure. The estimated optimal reward and cost goals are then used to guide the policy optimization through constrained policy optimization paradigm. To further enhance the stability and efficiency of the training procedure, we reshape the dataset in terms of Re, Ct distribution, ensuring more uniform reward-cost-retum distribution to support GAS training. Theoretical analysis and the pseudocode of GAS are presented in section B and section
5.1 OPTIMAL ACHIEVABLE GOALS
Conceptually, the maximum achievable return-to-go V.R and the corresponding cost return-to-go V/ for a given state and the target reward/cost returns R and Ò in the dataset should follow:
max (R, R, C* (s. B4..
R. 1 S C).
(6)
1 (s. R, C) argc, (5:=8,0,, Re C,)
Rg 1(C, S Õ)}.
(7)
However, directly applying eqs. 6 and 7 to estimate the optimal goals raises three issues. First, the standard definitions of R2 and C are the cumulative rewards and costs from the current timestep
5

---

Published as a conference paper at ICLR 2026
Suhoptimal trajectory from may optimul from
{ Benefit
Transitions
Good transitions
Benefit
Rr
R
R
R
R
Rg
St Rn Pr. Fe.
Sr Rn
Sy
dr
Pio Kr
Wr
Figure 3: Necessity of temporal segmented return augmentation.
the trajectory end. However, good transitions often occur within shorter segments. This motivates our introduction of Re.r and Car in section 5.2 representing cumulative values from t to T. Second, in the training stage, the target reward/cost returns are derived from cumulative values along the trajectory, which may not consistently align with the predefined targets in the testing stage. To address this problem, we propose a transition-level return relabeling in the training stage in section 5.3 Third, rare "lucky" transitions with high rewards and low costs can lead to value overestimation for maximum achievable goals due to transition function stochasticity. To address this problem, we design a goal function with expectile regression estimate the expectation the upper quantile of the returntto-o distribution in section 5.4
shorter window
5.2 TEMPORAL SEGMENTED RETURN AUGMENTATION
Processing the dataset in smaller temporal segments provides GM training with more abundant information to enhance cross-trajectoon stitching. Specifically, we leverage such insights and augment the return information in the dataset according to:
(st, G1, R2, Ct) + {(st. at, RA., Car) f..... {(s av, ReCT, Ce = =t +
where REr = + + TT and Ce. = - +... cr. When sampling (sp, dy Rg Cr), GAS can seek better R&T Ry C.T S Cy for the same state from other transitions under augmentation as long as r-f= T - t. As illustrated by 3 this data augmentation scheme provides two benefits: (1) It substantially expands the training data by including transitions of varying temporal lengths; (2) It enables more flexible transition stitching across different timesteps by providing diverse time intervals.
5.3 TRANSITION-LEVES RETURN RELABELING
First observed in Liu et al. 2023b the misalignment between human-specified reward-cost targets the testing stage and the training inputs can lead to degraded performance for GM methods due to their nature of behavior cloning. Inspired by the trajectory-Ievel labeling in Liu al. (2023b), propose a more fine-grained, transition-level return relabeling mechanism fitting in with transition- level stitching. For sampled transitions D {(s, a,, Rer, Car, t' = + T I')], we relabel Ber and Cur in eq. ( and take relabeled values as input of goal functions (st, de, Bar, Cer,t').
REr = U((1 - 5) Rar, ; (1 + 8) Rer),
(8)
ČAr UCCrr, C
where U (a, b) is a uniform distribution between a and 6, 8 E (0, 1) is a hyper-parameter and Omax is the maximum value of cost returns. In this way, goal functions can be trained under more imbalanced and comprehensive reward-cost targets. Notably, our proposed GAS does not directly update the policy guided by the relabeled values. Instead, we utilize them assist training through intermediate optimal goals and update the policy based on these goals during optimization. As a result, GAS retains the robustness of the policy without affecting reward maximization and constraint satisfaction.
5.4
GOAL FUNCTIONS WITH EXPECTILE REGRESSIONS
Since naively taking the maximum operator in the dataset can be prone to rare "lucky" samples, we adopt distributional perspective and optimize goal functions that focus more high return-to-go samples and less on low return-to-ge samples. To this end, we employ expectile regression for iteratively updating the estimated goal functions.

---

Published as a conference paper at ICLR 2026
The reward goal function should output the largest reward-to-ge that satisfies the constraint. To formalize this, we first define the advantage function as:
A 1 (V < C.r) REr - (s REr, Cer,t'=t+TTDD
where 1VE C.r) an indicator function of constraint satisfaction.
In this way, transitions that violate the constraint or have low return Rer down-weighted during expectile regression. Then the loss function with expectile regression can be defined as
Ln = Eplla ( 0)l )A..))).
(10)
With this loss function, the reward goal function (V converges to the expectile of the largest reward return-to-ge go that satisfies the constraint controlled by a.
The optimization objective for the cost goal function differs from that of the reward goal function. While the reward goal function aims to find the largest reward return-to-go in the dataset via expectile regression, the cost goal function seeks to estimate the cost value associated with the optimal reward goal. To address this, we modify the loss function of the cost goal function eqs. [1]) and (12).
assigning higher weights to transitions with higher reward goals under the constraint.
A = Cur - VEr (8,, RET, C.r, =t+T
(11)
T.c = E 1 0)l (4..)))
(12)
is worth noting that the reward and cost goal functions are not trained separately but derived under the same optimization framework with a theoretical guarantee.
To ensure that the policy comprehends optimal targets under relabeled targets returns, incorporate the optimal reward and cost goals obtained from the goal functions as inputs to the policy. Then we optimize the policy utilizing a constrained version of Advantage Weight Regression (AWR), as shown in:
Goal-guided Policy Optimization:
L = EDLI(VE < C.r) la 1(A < 0)] (I(alse, Rer, Car, V VGr t') -a,)%. (13)
5.5 DATASET RESHAPING
Existing GM methods take reward and cost returns as input, while the reward-cost distribution in the dataset is often highly imbalanced, problem referred to data imbalance (Kang et al. 2021
Bagui Li 2021: Yang et al., 2021 Ren et al. 2022). The issue seriously affects the ability of
RL-based methods constraint satisfaction Yao et al., 2024), as well as E-Massisted methods.
We define transitions with low costs and low rewards "conservative transitions".
Osiginal Distrilation
80
Reshuped Distibuttio
transitions with high costs and high rewards as "aggressive transitions" and transitions with low costs but high rewards as "ideal transitions". As shown in fig. most tran- sitions in the dataset are concentrated in re- gions where both cost and reward returns are extremely low. When training with uni- form sampling, GM methods tend to learn predominantly from these conservative transitions, while under-representing ideal transitions that exhibit higher reward returns with lower cost returns. To mitigate this issue, we propose to re- shape the dataset distribution during the training phase. In particular, first estimate the reward- return distribution conditioned on cost returns from the offline dataset and then select all transitions that fall within the top 9% reward returns for each cost return, thereby creating new dataset Dr = {(8, a, R. C) ~ 'D) (R(CC > - 4). where pe indicates cumulative distribution function. Throughout the training procedure, D will be sampled with probability e and the original dataset D will be sampled with probability 1 - e. As shown in fig. the reshaped dataset distribution is more balanced compared to the original one.
480-
480-
860
360
Rowed
240
Rowar
240
120
120
20
40
50
30
20
40
60
80
Cost Retuns
Cosi Retums
Figure 4: Original and reshaped dataset distribution.
EXPERIMENT
In this section, we aim to evaluate our proposed GAS and empirically answer three questions: 1) Can GAS achieve both safe and better performance with improved stitching ability? 2) Can GAS
7

---

Published as a conference paper at ICLR 2026
Table I: Normalized evaluation results. The normalized cost threshold is set to 1. Values shown as "mean+su" represent the mean and standard deviation. Each value represents the average performance over 10 evaluation episodes with 5 seeds and thresholds. Bold/gray/blur indicate safe/unsafc/safe and best-performing results. (L) indicates that higher (lower) values are better.
Methods
CPQ
COgoDDiE
WSAC
VOUE
CDT
FISOR
CAPS
DCAC
GAS
C1
C1
Tight Constaiir Threddole Average reults thrnstilld win 19%, 30% and 30% of the maximiem
AsiRun 162 suM uud8ss 0.55aa
8,20 a.50 0.33 aa 0.72aa
8.3 8:29
Id0.am
R664ss 1U96sun 0.14
881
6.72.as 0.70
Ballkan m:30.a4
1.S
4.54p0
Lildess 7:10sas 108 10 8:3
1.00.. [ilh-su Libiia 8:7law
365
6,33
0.Mose 0M6umm
047xan
0.67ass
4.3amm
10luo
a,7.a 1.1
addi..s 6.22.x0
AncCede 842 1U00 0.23wam
2.44
0.26 68tas 0.17au
0.3 0.54cca
1.biox 6.13-ax 8.00 1.47 -m 0.3 0,22
0.75
6.52.am 0.9%
Baliiini
866 42 166 0.70wa
353sox
a. a.30-am 0.76aam 1.10sam 0.22ccm
123smm R.28 8,30-am am RU3 wu 0.47 0.77
0.35
6.71-am 0.54
CarCade
a. 1,7
1.0 1.3tarm
1M0 HMd... 1166 117haa
6.70iaan
005,re 2.0
1.2gpu
0.00gres
8.9,.n
0.7k.r
Aserge 0.39.ai 1.42se 0.59pmm 2.25
655 1:36am 0.57aai 3. 3.10sam 0.67 E.l2sm 6.3 863 0.S0 1.09 0,55psa 131gs
6.80iim 0.67
NounCinke] 8:60-a
1,80
5.76axx
3.00aun
0.5aaai 4 4.33 0.20 8:6-am
0.33aa 1:23-a. 1.02-sm 1.37ass 0.7 L.5kxm
0.60-am 0.38
RoinCicle
4.5
392
0.55.a
0.8 5.3
1.l8on
0.5
1.42-a
0.7
0.10
3.8
6:75
CaCimkk
2.5 0.70pm 3.22gs
6,5 1,39 0.57 4 722 0/22go 2.56
8,57.MM 8:40 (.56.s
(182gam 0.57gst
2.46g
0,50 0,74
Loooe Costtiia Tareshold Avergge soulls oz durdsolis with 20%, SIS, and 90% of the nasimiin costa for esch talk.
0.53. 6.6..
0.30.so
0M8 0.00...
BallRin
1.30amm 1.5..
0.8k.om
1.10im
6.4
MM0... 11.7
0:76.
0.%%
Carkn
0.5 0.73-am 0.96.cm 0.15vou
0.54. 6.3
0.95-mm 200 0.99 uc 0,87-xx 6.52
0.01 (.99 0.28 0.86.om
0.19.
6,50 0.00
DecedRus
8.2-am am 1259-4a 0.0iia 0,7 100
1:32
0.9242 1 1/2 0.70 9.b-.s 8,37-ax 8.22 -m 1.S3-a 0.22 0.0boa 0.11.
6.50-am 0.92
AncCecl
a. I.I.... (1:5 0.50.
66 6.3,
8.K 0.32 BA a.tt.am 8.1..a 1M00 ILT0 1102 0.19 0,40;
6.77.
Buliiaii 0.37 1192-42
0.5
6.2-au
0.52.sm 8:32-m
8.28-sx
0:00-am 1192-.m:
0.2..m
6.2
162. ILl1:am go 0,27ga
0,8%
661:aw
0.8.a
(.51gan
0,79
8.49:a
1.01 (68 (194 0,21ga0
9.57.m
0.87.sm 0.80.am
Avenge a. 1Mili 11M0u 0,50um
0.74.o0
0.66umm
0.0
8.76ue B.Muss B.0... 1L77..o 1175pm WAlu 0,65.
8.MG.s 0.87
Tonniinl::
IL76..w
PoinCirded (66
0.91
1.23
3.08ua
0.53am
0.7 0,15.s
11.70sx
0.06uom 0.89.a
0.30.
Carlidel 1.54.n 1.0.es 0.50
1.28aam 0d
0.29amm 1 1:5aa 0.7 9,57 875 8.88 1150 11300amm 0.22ion L.bkaon 0.52.m
0.80
CarCiek2 0. 1.0rew 0.79 0,79 4.7 6.75:am 0.53aam 1 09aan 0.73 0.55aan 8,57.am 1.9.. (L72.sa 1154icm 0.40poo 1:21p 6.50. 0.82 am
preserve the zero-shot adaptation ability to different constraint thresholds? 3) Is GAS robust to imbalanced and human- -peean-pe target reward-cost return-to-goess Accordingly, we design the following experimental setup to evaluate GAS.
Tasks: We evaluate GAS on the widely used Bullet-Safety-Gym (Gronauer 2022) and Safety- Gymnasium (Ji et al., 2023) benchmarks. For all tasks, we use the dataset provided in DSRL. (Liu
ct al.
2023a) as our offline dataset, following the D4RL (Fu ct al. 2020) benchmark format.
Baselines: We compare our proposed GAS with the following bascline methods: (1) Constraint penalized method: CPQ (Xu et al. 2022); (2) Distribution correction estimation: COptiDICE (Le
et al. 2022): (3) Variational optimization with conservative estimation: VOCE (Guan al. 2023 (4)
Weighted safe actor-critic: WSAC (Wei et al. 2024); (5) Generative model assisted algorithms: CDT
(Liu et al. 2023b) and FISOR (Zheng al. 2024 RL methods with zero-shot ability: CAPS
(Chemingui et al. 2025) and CCAC (Guo et al. 2025).
Metries: We evaluate performance using normalized reward and cost returns. The normalized reward return is defined by Raomm R__, where R2 is the reward return achieved by policy π and and Rmma Runs'
is the maximum reward return in the dataset. The normalized cost return is defined by Cnom = C. where C2 is the cost return achieved by policy π and L is the selected threshold. To provide better interpretation to decouple the two objectives in OSRL compared to traditional fixed thresholds, we use percentage- based thresholds calibrated to each task's cost range. Specifically, we define tight constraints as 10%. 20%, and 30% of the maximum cost to emphasize constraint satisfaction, and loose constraints as 70%. 80%, and 90% of the maximum cost to focus on reward maximization.
L
6.1 CAN GAS ACHIEVE BOTH SAFE AND BETTER PERFORMANCE WITH IMPROVED STITCHING ABILITY?
The experiment results of different baselines under various tasks on tight and loose constraints are presented in table In tight constraint settings, only GAS achieves the best and safe performance in all tasks. Traditional RL methods, such as CPQ, COptiDICE, WSAC, and VOCE, suffer from severe constraint violations both on average and for each task. Even among GM methods, CDT, while outperforming RL-assisted methods, still fails to ensure constraint satisfaction in multiple scenarios where GAS succeeds. FISOR maintains safety in most tasks but at a substantial cost to performance, with rewards significantly lower than GAS. In such tight constraint settings, the superiority of GAS over CDT on constraint satisfaction comes from the stitching ability, where GAS can stitch safe transitions among different timesteps and trajectories together. In loose constraint settings, although most baselines exhibit safe performance, GAS achieves the best performance on reward maximization. This performance gap between GAS and CDT highlights GAS's advanced reward maximization
8

---

Published as a conference paper at ICLR 2026
Meal
BallKun
DroneCiroke
DroneRa
1200
900
500
Rowards
1000
800
500
800
600
400
500
300
40
70-
16
24
32
10
48
56
64
72
10
20
30
40
50 Threshold
60
70
80
90
14
12
56
70
84
98
112
126
Threshold
Threshold
Figure 5: Evaluation results zero-short adaptation. The x-axis indicates different selected thresh- olds and the y-axis indicates corresponding performance on cumulative rewards and costs. "Ideal" line indicates the case when the cumulative costs are equal to the constraint thresholds.
capabilities, leveraging its innovative transition stitching approach to combine high-reward segments while maintaining robust safety guarantees.
6.2 CAN GAS PRESERVE THE ZERO-SHOT ADAPTATION ABILITY TO DIFFERENT CONSTRAINT THRESHOLDS?
critical advantage of GM methods is their zero-shot adaptation ability in handling different constraint thresholds without retraining. To validate this ability in GAS, we compare our method with CDT and test thresholds from 10% 90% of the maximum costs, increased by 10% each time. The results shown in fig. Compared with CDT, the cumulative costs of GAS are consistently below the constraint thresholds, whereas CDT performs unsafely some cases, especially when the constraint thresholds are smaller than 30% of the maximum costs. In addition, the constraint threshold increases, GAS flexibly adjusts behavior and achieves progressively higher rewards while ensuring safety, but CDT tends to be overly conservative.
6.3 GAS ROBUST TO IMBALANCED AND HUMAN HUMAN-SPECIFIED SPECIFIED TARGET REWARD-COST RETURN-TO-GOBS?
To demonstrate GAS's robustness against imbalanced target reward-cost return-to-goss to-gs,,,, we compare GAS under different return relabeling levels (denoted as "GAS-8" and "GAS w/o Relabel")) in settings with the constraint threshold 20% the maximum costs and varying target reward return- to-goes. As shown in fig. 6 GAS consistently achieves safe performance as the target reward return-to-go increases. In contrast, GAS w/o Relabel can only achieve safe performance for a narrow band of reward targets. As 8 increases, GAS becomes more robust against imbalanced targets. This highlights one of the GAS's s key innovations: the transition- return relabeling method in enhancing robustness.
GAS wio Rolabel
GAS-0.05
GAS-0.10
GAS-0.20
GAS-0.30
GAS-0.40
Threshald
BallRan
DroneClucle
DroncRun
Rewards
500
450
400
300
360 400 440 490
340
360
380
400
50
40
50
25
360 380 400
420 440 460 480 Tärget Reward Retam-tt-go
550 580
340 360
350 400 440 460 480 Targer Return-to-go
Tirget Reward Retum-to-go
Figure 6: Evaluation results robustness of imbalanced target reward-cost returntto-goes. The x-axis indicates different target reward return-to to--ootoooooooooo and the y-axis indicates corresponding performance on cumulative rewards and costs. "Threshold" line indicates constraint thresholds.

---

Published as a conference paper at ICLR 2026
GAS w/o Stitching
GAS-0.7
GAS-0.9
CDT
Ideal
GAS-0.6
GAS-0.8
GAS-0.99
GAS w/o TSRA
BallRum
DroneCirele
DroneRum
1200
600
900
1000
Rewards
500
800
780
400
660
400
600
300
72
90
126
56
70
98
Cods
40
70
30
14
10
16
24
32
40
48
56
64
72
10
20
30
40
50
112 126
Constraint Threshold
Constraint Threshold
Constraint Threshold
Figure 7: Ablation studies on the stitching ability (expectile regression) and temporal segmented return augmentation.
6.4 ABLATION STUDY ON STITCHING ABILITY AND TSRA
To verify the theoretical improvement for GAS's component in section we compare GAS under different expectile level a from 0.5 to 0.99 (where = 0.5 is denoted as "GAS w/o Stitching" and others are denoted as GAS-a), and a variant without Temporal Segmented Return Augmentation (denoted as "GAS w/o TSRA") in fig. As shown in fig. 7 without expectile regression, GAS w/o Stitching degrades to conditional behavior cloning (a 0.5), which has similar performance and problems to CDT. This result is also consistent with section 4.1 where both GAS w/o Stitching and CDT suffer from insufficient stitching ability. a is extremely high, GAS-0.99 fails to keep safety in DroneRun tasks. This result is consistent with claims of the problem of "lucky" biases in section 5.1 As a increase from 0.6 to 0.9, GAS well-performed and demonstrates robustness on a. The performance of GAS w/o TSRA is almost completely inferior GAS. This result validates the theoretical analysis result, where the optimality without augmentation is just a lower bound of that with augmentation. Furthermore, GAS w/o TSRA suffers from constraint violation for tight constraint thresholds compared to GAS. This result indicates that the flexible sub-trajer sub-trajectories augmented by TSRA provide additional benefits when original trajectories fail to satisfy the tight constraint thresholds.
7 CONCLUSION
Aiming to address the limitations of existing GM-assisted OSRL methods, we propose a novel algo- rithm named GAS that concentrates enhaneing the stitching ability for better balance of reward maximization and constraint satisfaction. GAS focuses on achieving better reward maximization and constraint satisfaction training specialized reward and cost goal functions via expectile regression. These goal functions estimate the optimal achievable reward and cost returns and are used to guide the policy in effectively stitching transitions under relabeled reward-cost return-to-goes Experiment results demonstrate GAS's 's superiority on reward maximization, constraint satisfaction, zero-shot adaptation, and robustness on imbalanced target reward-cost return-to-goes. A potential weakness of GAS is the trade-off between memory capability and stitching ability, the current attention module struggles to fully capture true temporal dependencies in CMDPs. Besides, due to the limitation of the existing goal-conditioned framework, GAS's stitching ability focuses on improving performance by better selecting high-quality sub-segments in the safe dataset, rather than constructing new safe trajectories by combining safe and unsafe data in RL. Future work could address this limitation by designing more advanced memory mechanisms and frameworks into GAS to further enhance the robustness and performance of GM-assisted policies.
10

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGEMENT
This work was supported by the Hong Kong Research Grants Council under the Areas Ex- cellence scheme grant AoE/E-601/22-R and NSFC/RGC Collaborative Research Scheme grant CRS_ 2KUS603/2/2
REPRODUCIBILITY STATEMENT
This paper introduces a new algorithm, named GAS, clear theoretical explanation of GAS is illustrated in sections A and B
The clear pseudo code is shown in algorithm [ Detailed experiment settings, including hyperparameters, are shown in section The experimental code available at com/ htmm/ifffer.. wtwp:tooooooo..f.eeeerooorrre
https://g http://gg.ttt..ocooo...
REFERENCES
Eitan Altman. Constrained markov decision processes with total cost criteria: Lagrangian approach and dual linear program. Mathematical methods of operations research, 48:387-417, 1998.
Anirudhan Badrinath, Yannis Flet-Berliac, Allen Nie, and Emma Brunskill. Waypoint transformer: Reinforcement learning via supervised learning with intermediate targets. In A. Oh, T. Nau- mann, A. Globerson, K. Saenko, M. Hardt, and S. Levine (eds.), Advances in Neural Information Processing Systems, volume 36, pp. 78006-78027. Curran Associates, Inc., 2023.
Sikha Bagui and Kunqi Li. Resampling imbalanced data for network intrusion detection datasets.
Journal of Big Data, 8(1):6, 2021.
Yassine Chemingui, Aryan Deshwal, Honghao Wei, Alan Fern, and Jana Doppa. constraint-ddapiii policy switching for offline safe reinforcement learning. In Proceedings of the AAAI Conference Artificial Intelligence, volume 39, pp. 15722-15730, 2025.
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems, 34:15084-15097 2021.
Will Dabney, Mark Rowland, Mare Bellemare, and Rémi Munos. Distributional reinforcement learning with quantile regression. In Proceedings the AAAI conference on artificial intelligence. volume 32, 2018
Xing Fang, Qichao Zhang, Yinfeng Gao, and Donghin Zhao. Offline reinforcement learning for autonomous driving with real world driving data. 2022 IEEE 25th International Conference on Intelligent Transportation Systems (ITSC), pp. 3417-3422. IEEF, 2022.
Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning. arXiv preprint arXiv:2004.07219, 2020.
Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings Machine Learning Research, 2052 2062. PMLR, 09-15 Jun 2019. URL bttps://proceediing
mlr
ppress/v7/ Fu http://////uimmoolll... html
Sven Gronauer. Bullet-safety-gym A framework for constrained reinforcement learning. Technical report, TUM Department of Electrical and Computer Engineering, Jan 2022.
Jiayi Guan, Guang Chen, Jiaming Ji, Long Yang, ao zhou, Zhijun Li, and changjun jiang. Voce: Variational optimization with conservative estimation for offline safe reinforcement learning. In
A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine (eds.), Advances Neural Information Processing Systems, volume 36, pp. -3758-33780.. Curran Associates, Inc., 2023.
Zijian Guo, Weichao Zhou, Shengao Wang, and Wenchao Li. constraint-oodditoon actor- critic for offline safe reinforcement learning. In The Thirteenth International Conference on Learning Representations, 2025. h htttpa://ppeneviiiiiie

---

Published as a conference paper at ICLR 2026
Philippe
Ilya Kostrikov, Michael Janner, Jakub Grudzien Kuba, and Sergey Levine. Idql: Implicit q- q-learning as an actor-critic method with diffusion policies. arXiv preprint arXiv:2304.10573. 2023.
Natasha Jaques, Asma Ghandeharioun, Judy Hanwen Shen, Craig Ferguson, Agata Lapedriza, Noah Jones, Shixiang Gu, and Rosalind Picard. Way off-policy batch deep reinforcement learn- ing of human preferences in dialog. 2020. URL https:://pperrevisws
EJ15ERVEV!!
Jiaming Ji, Borong Zhang, Jiayi Zhou, Xuchai Pan, Weidong Huang, Ruiyang Sun, Yiran Geng, Yifan Zhong, Josef Dai, and Yaodong Yang. Safety gymnasium: A unified safe reinforcement learning benchmark. Advances Neural Information Processing Systems,
2023.
Haeyong Kang, Thang Vu, and Chang Yoo. Learning imbalanced datasets with maximum margin loss. In 2021 IEEE International Conference on Image Processing (ICIP), pp. 1269-1273, IEEE, 2021.
Jeonghye Kim, Suyoung Lee, Woojun Kim, and Youngchul Sung. Decision convformer: Local filtering in metaformer is sufficient for decision making. In International Conference on Learning Representations, 2024a.
Jeonghye Kim, Suyoung Lee, Woojun Kim, and Youngchul Sung. Decision convformer: Local filter- ing in metaformer is sufficient for decision making. In The Twelfth International Conference on Learning Representations, 2024b. URL
net/ forummrmm d=
af2c8RaK18
Roger Koenker and Kevin F Hallock. Quantile regression. Journal of economic perspectives, 15(4): 153-156, 2001.
lly Kostrikov, Rob Fergus, Jonathan Tompson, and Ofir Nachum. Offline reinforcement learning with fisher divergence critic regularization. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp.
PMLR, 18 24 Jul 2021. URL http:tts////ppceediccee html
mlr. wtr.spres/vv333 wttess///1////strkkkv///
Ilya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit q-learning. In International Conference on Learning Representations, 2022. URL https://
openreview n forum? wooorrmmii:nnnm... 68 92JWFF
Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative gelarnning for offline reinforeement learning. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.). Advances Neural Information Processing Systems, volume 33, pp. 1 119911991 Curran Asso- ciates, Inc., 2020. URL http:tttss//ppppceeeeeegg/ge neur ips cc caper fi fllesppppe 2020 htt.f/lle/ddaa/ (rl.2222222222222222211116553530
82
Jongmin Lee, Cosmin Paduraru, Daniel Mankowitz. Nicolas Heess, Doina Precup, Kee-Eung Kim, and Arthur Guez. COptiDICE: Offline constrained reinforcement learning via stationary distribution correction estimation. In International Conference on Learning Representations, 2022.
URL httphhtttss///peeeeeeee http://wwww.wweeeeeeeeeeeeee
wwt//forrrr d==FLA55mBees
Namyeong Lee and Jun Moon. Offline reinforcement learning for automated stock trading. IEEE Access, 11:1 11:112577-11288 2023.
Zuxin Liu, Zijian Guo, Haohong Lin, Yihang Yao, Jiacheng Zhu, Zhepeng Cen, Hanjiang Hu, Wenhao Yu, Tingnan Zhang, Jie Tan, et al. Datasets and benchmarks for offline safe reinforcement learning. arXiv preprint arXiv:2306.0933 093003 2023a.
Zuxin Liu, Zijian Guo, Yihang Yao, Zhepeng Cen, Wenhao Yu, Tingnan Zhang, and Ding Zhao. Con- strained decision transformer for offline safe reinforcement learning. In International Conference on Machine Learning, pp. 21611-21630. PMLR, 2023b.
Yi Ma, Jianye Hao, Hebin Liang, and Chenjun Xiao. Rethinking decision transformer via hierarchical reinforcement learning. In Proceedings of the 41st International Conference on Machine Learning, pp. 33730-33745, 2024.
12

---

Published as a conference paper at ICLR 2026
Xue Bin Peng, Aviral Kumar, Grace Zhang, and Sergey Levine. weighte regression: Simple and scalable poflcc reinforcement learning. arXiv preprint arXiv:1910.00177, 2019.
Jiawei Ren, Mingyuan Zhang. Cunjun Yu, and Ziwei Liu. Balanced mse for imbalanced visual regres- sion. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 1922,77335, 2022.
Noah Siegel, Jost Tobias Springenberg, Felix Berkenkamp. Abbas Abdolmalcki, Michael Ne- unert, Thomas Lampe, Roland Hafner, Nicolas Heess, and Martin Riedmiller. Keep doing what worked: Behavior modelling priors for offline reinforcement learning. International Conference on Learning Representations, 2020. URL.tttps://pperevviie hwtp:/ee//ffruuta
id=rke od=rke7geeto. geHttH
Adam Stooke, Joshua Achiam, and Pieter Abbeel. Responsive safety in reinforcement learning by pid lagrangian methods. In International Conference on Machine Learning, pp. 9133-9143, PMLR, 2020.
Ruhan Wang and Dongruo Zhou, Safe decision transformer with learning-based constraints. In 7th Annual Learning for Dynamics)& Control Conference, pp. 245-258. PMLR, 2025.
Honghao Wei, Xiyue Peng, Arnob Ghosh, and Xin Liu. Adversarially trained weighted actor-critic for safe offline reinforcement learning. Advances in Neural Information Processing Systems, 37: 52806- 52835, 2024.
Yifan Wu, George Tucker, and Ofir Nachum. Behavior regularized offline reinforcement learning, 2020.
net/forumm 9
Yueh-Hua Wu, Xiaolong Wang, and Masashi Hamaya. Elastic decision transformer. Advances in neural information processing systems, 36:18532-18550, 2023.
Chenjun Xiao, Han Wang, Yangchen Pan, Adam White, and Martha White. The in-sample soft- max for offline reinforeement learning. In The Eleventh International Conference on Learning Representations, 2023. URL https: http:///ppeeeiieieeeeeeeeeeeeeee wtt.:e//ooouuddddaaauurrrroaaaaa
Haoran Xu, Xianyuan Zhan, and Xiangyu Zhu. Constraints penalized q-learning for safe offline reinforcement learning. Proceedings of the AAAI Conference on Artificial Intelligence, 36(8): 8753 8760, Jun. 2022. doi: 10.1609/aaai, 33336688825555. URL
aaai. org
Undex.ppp/aie AAA T aerriiclee wewiiileevviii.c ew 20855
Ruiqi Xue, Ziqian Zhang, Lihe Li, Cong Guan, Lei Yuan, and Yang Yu. Adaptable safe policy learning from mulli-task data with constraint prioritized decision transformer. The Thirty-ninth Annual Conference on Neural Information Processing Systems, 2025. URL htttps://operreeee net/fooum?? Him
Taku Yamagata, Ahmed Khalil, and Raul Rodriiue. Q-Jearning decision transformer:
Leveraging dynamic programming for conditional sequence modelling in offline rl. In International Conference on Machine Learning, pp. 38989-39007. PMLR, 2023a.
Taku Yamagata, Ahmed Khalil, and Raul Santos- Rodrigue. Q-learning decision transformer:
Leveraging dynamic programming for conditional sequence modelling in offline RL. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.). Proceedings of the 40th International Conference Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. -9089. PMLR, 23-29 Jul 2023b.
Yuzhe Yang, Kaiwen Zha, Yingcong Chen, Hao Wang, and Dina Katabi. Delving into deep im- balanced regression. In International conference machine learning, pp. 11842 11851. PMLR, 2021.
Yihang Yao, Zhepeng Cen, Wenhao Ding. Haohong Lin, Shiqi Liu, Tingnan Zhang, Wenhao Yu, and Ding Zhao. Oasis: Conditional distribution shaping for offline safe reinforcement learning. Advances in Neural Information Processing Systems, 37:78451-78478, 2024.
Qinqing Zheng, Amy Zhang, and Aditya Grover. Online decision transformer. In international conference machine learning, pp. 27042-27059. PMLR, 2022.
13

---

Published as a conference paper at ICLR 2026
Yinan Zheng, Jianxiong Li, Dongjie Yu, Yujie Yang, Shengbo Eben Li, Xianyuan Zhan, and Jingjing Liu. Safe offline reinforcement learning with feasibility-guiie diffusion model. arXiv preprint arXiv:2401.10700, 2024.
Zifeng Zhuang, Dengyun Peng, Jinxin Liu, Ziqi Zhang, and Donglin Wang. Reinformer: Max- return sequence modeling for offline rl. In International Conference on Machine Learning, pp. 62722. PMLR, 2024.
14

---

Published as a conference paper at ICLR 2026
A EXPECTILE REGRESSION
Expectile regression (Koenker & Hallock 2001) is statistical modeling technique that generalizes traditional mean regression to obtain the weighted means across different parts of the distribution. Given of samples {a:li ... N} under distribution, the expectile regression aims to achieve eq. (14 where a is the expectile level and u 2 - 0².
min Blaa 1 (u < 0) u²] + min Bl 1(xx~ x²) < <0) -x- xx-x=))),
(14)
Property (Kostrikov et al. 2022);
As a increases from 0.5 to 1, _a increases from the mean value to the largest value of a distribution:
lim F" max(x;li 1, 1...,.N
(15)
For example, if 0.5, is the same MSE since |0.5 1(.) 0.5; if 0.9, it gives weight of 0.9 to samples æ 2 z² but only gives weight of 0.1 to samples < ~
This property makes expectile regression widely used in the estimation of the optimal value function in reinforcement learning (Kostrikov et al. 2021 Dabney et al. (018)) since the optimal value function naturally defined as the largest reward return in task.
Property 2 MlansnnEEstuule et al. 2023):
Let us represent the expectile regression f') = la 1(u 0)1 u² and f'(u) df(w) du
for easy
and clear writing. We have:
f'(u) (u] u
(16)
where f(u) is convex and f0( O.
B THEORETICAL ANALYSIS OF GAS
B. ALGORITHM DERIVATION
Fowllowing the eq.
15 property eq. and eq. 110.
directly have:
max lim (1 Ro., Čar,t' =t+ T r) (f Ber. C.r,''
(Ra Car) (17)
This indicates that the reward goal function finally converges to the largest reward return-to-go within the constraint.
Then consider La in eq. 10p as f(A
When the reward goal function converges, we have
ÖLn
A
u(als) (A
R
[S'(AR)) )
((aas))
-A
(18)
T(als) R
where π (a|s)) = p(als))
(A Naculu A
is the target policy.
With the target policy, we can define the loss function of the cost goal function as L = Ex((Crr VC j]. When converges, have:
15

---

Published as a conference paper at ICLR 2026
2Ex(C21 Virl
2 J. T(als) (OuT ViT))
(19)
2//Nooale [ p((ass) - la- -1(A < 0)l (Cur Vi.T-)
Thus, we can obtain the loss function of the cost goal function as eq.
(A(A
During the policy extraction, the target policy supposed to be (a|s) = p(als)
Nacate 14
However, to further ensure safety, we add the constraint to the policy extraction:
<(a|s) = 1/Neoale. 1( S C.r) la 1 < 0) plals)
(20)
Thus, the loss function of the target policy eq. [3
B.2
ABLATION ANALYSIS
This ablation analysis aims to study the contribution of each separate component of GAS theoretically,
B.2.1 STITCHING ABILITY (EXPECTiLE REGRESSION)
Based on property section A easy to prove that:
lim (s Rar, Cer,t'= lim (et, Rar, Cer,t') Rea (RerlVer Car) a=0.5
a-1
(21)
Thus the optimality guarantee is lost without expectile regression.
B.2.2 TEMPORAL SEGMENTED RETURN AUGMENTATION
We can partition the dataset after augmentation as D = D U D where D is the original dataset and D is the augmented part.
lim EDIVE (8,, Re., C.r, = a+1
max (.e.,,,.ur. Cur, (RarlVer Car)
max max Rer VVC < Crr), max(RarlVal Car)l (22) (D")
(D"
(Rer/ < C.r)
max [s Rett C)r-e)-D")
This result indicates that the optimality guarantee in the original dataset is just lower bound of that after augmentation. With augmentation, the optimality guarantee is always larger than that in the original dataset, which further improves the stitching ability.
B.2.3 TRANSITION-LEV LEVEL RETURN RELABELING
Since the purpose of this method is to make GAS more robust, this paper will discuss it from the perspective of the input and output. For simplicity, the influence the reward goal function is analyzed as example, considering that it has no difference from that on the cost goal function and the policy.
Without Relabeling, both inputs and loss functions utilize (R, C)) pairs following the dataset distribu- tion. This indicates that only with the appropriate R and C, the goal functions and policy can achieve the optimal value:
max lim (.,, Ren, Cur, ((s)a)) C.r.C.r (eerVVe Car)
(23)
However, during the test stage, the targets R and Ô cannot be accurately known and need to be specified by users. Thus it will suffers from inaccurate R and signals and cannot determine which target should be prioritized when R and C conflict with each other.
16

---

Published as a conference paper at ICLR 2026
With Relabeling, the targets relabeled to consider more robust input distribution while loss functions still utilize true (R, C*) pairs to update the outputs with expectile regression. This indicates that even with inaccurate R. and C, the goal functions still obtain the optimal value, and the policy can provide the corresponding action:
lim Vel (st, Ber, Crr,t)= (er Cer. a-1
(BerlVer < Cer)
(24)
C ALGORITHM DETAILS
We present the full algorithm of our method in algorithm T
Algorithm 1 Goal-Assisted Stitching (GAS)
V.. 1: Network: Initialize two goal functions V and policy (a s, R, C, VR, ,' -t)
R2
2: for iteration= 0,... do
3: Sample transitions D {(s, a, REr, Car)) ~ 'D with c cand 2 with c probability.
4: Get augmented return and cost return:
5:
Aer (( (1 Rrrr (1 + Rar).
6:
Č.r UCrrr C!MA)).
7:
Get advantage function:
8:
A 1V < Cer) RET V Ra., Car, t' -)).
9:
AC Car Vi Ber, Car, -
10:
Update reward goal function:
LR = E 1A 0)] (41
12:
Update cost goal function:
L.c = 1- (A <0] (4(
14:
Policy Extraction:
15:
LE EVL < C.r) laa 1(4 < 0)1 (T(aas,, Rar, Cur, ETs VGG, t') -a,)²]
D RETURN RELABELING TECHNIQUE DETAILS
The misalignment problem between the training and test phases is first proposed in Liu et al. (202bb) which indicates that users may select target reward-cost pairs different from the training inputs. This misalignment problem has become serious when users select imbalanced target reward-cost pairs, such as extremely large target rewards and extremely small target costs. To address this issue, CDT (Liu et al. 2023b) proposes trajectory-level return relabel method CDT only learns policy within the trajectory. This method, although improves the robustness of CDT, degrade the ability to maximize rewards and satisfy constraints, the relabeled returns provide wrong information about the trajectory. Inspired by the trajectory-Ieww return labeling in Liu et al. (2023b), we propose more fine- ane-grrined, transition-leve return relabeling mechanism fitting in with transition-level stitching. Different from CDT, our proposed GAS does not directly update the policy guided by only the relabeled values. Instead, utilize them assist training through intermediate optimal goals and update the policy based on these goals in the loss function during optimization. As a result, GAS retains the robustness of the policy without affecting reward maximization and constraint satisfaction.
EXPERIMENT DETAILS
E.1 BENCHMARK AND TASKS
Bullet-safety-ee gym Gronauer 2022) and gaffet-gaamnsium (Ji al. 2023) are utilized for experi- ments. We consider 8 cases in Bullet- seweee--sfftyyyyye and cases in gaaetay-mmmmiuum involving two tasks (Run and Circle) and multiple types of robots (Ant, Ball, Car, Drone, and Point). The tasks AntCircle, PointCirllel, PointCircle2, CarCirclel, and CarCircle2 are considered as complex tasks with episode length T 500 and maximum cost mmax > 200, while other tasks regarded as
17

---

Published as a conference paper at ICLR 2026
Table 2: Benchmark details. The cost range indicates the maximum cumulative among the
offline trajectories. Offline trajectories indicate the number of trajectories in the offline dataset.
Benchmarks
Task
State Space
Action Space
Cost Range
Episode Length (T)
Offiine Trajectorics
AntRun
33
150
200
1816
BallRun
80
CarRun
40
200
651
DroneRun
17
4
140
200
1990
Bullet-safety--yym
AntCircle
200
500
5728
BallCircle
80
CarCirele
100
300
1450
DroneCircle
18
100
300
1923
PointCirclel
28
200
500
1098
safery-gymnasiun
PoinCCircle2
300
CarCirclel
250
500
1271
CarCircle2
40
2
400
500
940
simpler tasks with episode length less than 301. Details of the parameter settings for each tasks are shown in table which is part of DSRL (Liu et 2023a). As for the cost definition and reward definition, we follow the same standard with DSRL.
E.2 HYPERPARAMETER
Table 3: Hyperparameters of GAS.
Parameter
Value
Parameter
Value
Number of layers
7
Hidden size
128
Emhedding size Learning rate
64
Batch Size
2048
0.0001
Adam betas
(0.9,0.999)
Grad norm clip Expectile level a
0.25
Weight decay
0.0001
0.8
Reward relahel level 8
0.1
Dataset reshape threshold 4%
10%
Sample probability e 0.5
The table 3 shows the detailed hyperparameter for GAS used in section and table shows the target reward and cost return- to-go -o-g pairs used in the testing stage. Notably, GAS is not sensitive to the target reward-cost return-to-go pairs as we demonstrate section 6.3
Table 4: Target reward and cost return-to-go pairs for GAS in testing stage.
Benchmark
Task
Cost range
10%
20%
30%
70%
80%
90%
AntRun
150
(690, 15)
(690. 30)
(700. 45)
(750, 105)
(800, 120)
(820, 135)
BallRun
80
(420, 8)
(420. 16)
(500, 24)
(900, 56)
(1000, 64)
(1200. 72)
CarRun
40
(572, 4)
(572,
(572, 12)
(572, 28)
(572, 32)
(572, 36)
Bullet-safety-gynn
DroneRun
140
(400, 14)
(420, 28)
(45, 42)
(600, 98)
(620, 112)
(640, 126)
AntCircle
200
(160, 20)
(200, 40)
(240, 60)
(320, 140)
(350, 160)
(400, 180)
BallCirele
80
(500, 8)
(600, 16)
(690, 24)
(800, 56)
(810, 64)
(820, 72)
CarCirele
100
(370, 10)
(390,
(410, 30)
(480,
(480, 80)
(480, 90)
DroneCircle
100
(600, 10)
(650, 20)
(700, 25)
(830, 70)
(850, 80)
(870. 90)
PointCirelel
200
(43.20)
(45.40)
(46,60)
(52.140)
(54.160)
(58,180)
Safety gymeysymmmsiumm
PointCirele2
300
(36,30)
(41.60)
(42,90)
(45,210)
(47.240)
(48,270)
CarCircle1
250
(4.25)
(8,50)
(10.75)
(13.175)
(15,200)
(18.225)
CarCirele2
400
(8,40)
(14,80)
(15,120)
(22,280)
(23,320)
(27,660)
E.3
MORE EXPERIMENTS ON IMPROVED STITCHING ABILITY
Detailed results on each threshold. We provide more detailed results on each threshold for all
baselines under Bullet-safety-e gymm for reference in fig.
Additional results on velucity tasks. As shown in table we compare GAS with other baselines on
5 additional velocity tasks in Safety-gamnasium.
18

---

Published as a conference paper at ICLR 2026
34
Frando
Figure 8: Normalized evaluation results on each threshold for all baselines under sallet-safety-aee gym.
Table 5: Additional normalized evaluation results on Velocity tasks.
Methods
CDT
CAPS
CCAC
GAS
Tighr Constraint Threshold.
Hopper Velocity
0.07
0.06
0.90)=
0.99
0.25
1.94
0.78
0.80
0.97 nat
0.28 1.44
SwimmerVelles
0.67
1.06
0.46
2.26:109
z.20
Walker2dVelacity
a.78
0.12
1.78
0.76
0,23
0.92:
0.77 aat
0.75 4:35
AntVolocit
0.99
0.52
0,97
0.62
0,00
0.00
0.98
0,98
Loose Constraint Thresbold.
HopperVelocity
0.70
0.27:
0.24
0.42
0.02
0.04:aux
0.93
0.94
0.99.am
0.13
1.98
0.67
0.81 :aue
0.95
1.000
0.71
SwimmerVelaciity Velocity
0.57
0.37
0.35
0.60
0.19 0.19
0.24
0.64
0.67..41
Walker2Veelccity
0.78:
0.01
0.74:au
0.00:aue
0.00
0.88-134
AntVolocity
0.98
0.13
0.99
0.43
0.00 co1e
0.00
0.98
0.97:
E.4 MORE EXPERIMENTS ON THE ZERO-SHOP ADAPTATION ABILITY
The fig. demonstrates the zero-shot adaptation ability comparison of GAS and CDT on more complex tasks. Even on more complex tasks, GAS can preserve the reward maximization and constraint satisfaction ability for all cost ranges. CDT, in contrast, suffers from constraint violation in tight constraint settings or is overly conservative as the threshold increases.
E.5 ABLATION STUDY DATASET RESHAPING METHOD
To assess the importance of the dataset reshaping on the training stability and efficiency, we compare GAS under two hyperparameters, including dataset reshape threshold 9% (GAS-9%) in fig, I0 and sample probability (GAS-e) in fig. 11 on DroneCircle task under different constraint thresholds. In these two figures, GAS w/o DR represents GAS without dataset reshaping, which is the same as the case when 4% = 100% or e = 0. In fig. 10 GAS is robust to the dataset reshape threshold when 9% < 40% and is sensitive when 4% is extremely large. In fig. I1 the sample probability mainly influences the learning efficiency the early training stage without affecting the final results. These results demonstrate that GAS is robust under different 4% and of dataset reshaping.
GAS
CDT
Hieal
22
24
47-
18
43
35
180
370
125
340
100
125
250
75
40
160 120 140 150 180 Conatnin: Throbld
60 90
30 60
90
120 150 180 210 340 370
25
50 75 100
135 200 225
40
90
120
200 240 100 320 350 Cunsiraiaa Threshold
160
Comstrin Throdold
Conatrinal Thrshill
Figure Evaluation results on zero-shot adaptation with complex tasks. The x-axis indicates different selected thresholds and the y-axis indicates corresponding performance on cumulative rewards and costs. "Ideal" line indicates the case when the cumulative costs are equal to the constraint thresholds.
19

---

Published as a conference paper at ICLR 2026
GAS- 5AS-100%
GAS 20%
GAS 40%
GAS 80%
GAS w/o DR
Threshold
Threshold
Threshold
Threshold
Threshold 70
Threshold 80
Threshold 90
Rewards
600
200
Costs
120
80
40
20
40
60
80
20
40
60
80
20
40
60
80
20
40
60
80
20
40
60
80
20
40
60
80
Iterations (5e3)
Iserations (5c3)
Iterations (5e3)
Inerations (5c3)
Iterations (5c3)
Iterations (5c3)
Figure 10: Hyperparameter analysis on dataset reshape threshold 4% on task DroneCircle.
GAS w/o DR
GAS-25%
GAS-50%
GAS-75% 75%
Threshold
Threshold
Threshold
Threshold
Threshold
Threshold
Threshold
Rewards
600
400
160
Costs
120
80
40
Iterations (5c3)
Iterations (5c3)
Iterations (Sc3)
Iterations (5c3)
Iterations (5c3)
Iterations (Sc3)
Figure 11:
analysis on sample probability e on task DroneCircle.
E.6 INFRASTRUCTURE AND TIME COST
We run all the experiments and all the baselines with:
1. GPT version: NVIDIA GeForce RTX 3080.
2.
CPU version: Intel(R) Xeon(R) Platinum 8375C CPU @ 2.90GHz.
Table 6: Training iterations and time for GAS and baselines on task CarCircle.
Method
CPQ
COpúDICE
WSAC
VOCE
CDT
FISOR
GAS
Iteration times
4c5
4c5
3e4
4e4
4c5
4e5
4e5
Time
2h19min
2h43min
3h10min
3h38min
11h32min
1h44min
8h20min
We also test the training time of GAS and baselines on task CarCircle, as shown in table 6 For most baselines, we follow the same iteration times. As for WSAC and VOCE, we follow the traiiin standard the original paper and make sure their convergence since training 4e5 times much too long for them. Notably, both GAS and CDT only need to be trained once for all thresholds, even though the training time much longer, while other baselines need to re-trained for each threshold. Besides, GAS is faster than CDT as GAS only utilizes the MI. architecture, but CDT utilize the Transformer architecture.
F
THE USE OF LARGE LANGUAGE MODELS (LLMS)
In this paper, LLMs are utilized to polish writing (e.g., grammar, spelling, word choice),
20


Reviews:
[{"review_id": "CBYSpRM45y", "reviewer": "Reviewer_gTZL", "summary": "Goal-Assisted Stitching (GAS) is a generative-model–assisted approach to offline safe reinforcement learning (OSRL) that improves trajectory stitching and the reward–cost balance using Temporal Segmented Return Augmentation (TSRA) plus transition-level relabeling, expectile-trained goal functions, and dataset reshaping; those goals guide a constrained Advantage-Weighted Regression (AWR) policy, yielding safer and higher-reward policies across two benchmarks and multiple constraint levels.", "questions": "- Please check weaknesses.\n\n- Could you train GAS solely on an offline dataset of unsafe trajectories? If it succeeds, this may indicate an ability to stitch together feasible segments from otherwise failed trajectories.\n\n- In the tight constraint setting (Table 1), an interesting observation is that the rewards achieved by GAS are too similar to that of CDT, yet the costs are lower and safely below the constraint, while CDT's costs violate the constraint often.  How can GAS maintain the same reward scale as CDT while successfully reducing costs to a safe level, given that these two objectives are conflicting?\n\n- FISOR is designed to ensure the strict satisfaction of constraints. They focus on identifying the largest feasible region where zero violation is enforced.  How is FISOR able to adapt and show varying performance across the different loose and tight constraint limits presented in Table 1? \n\n- How would GAS perform on additional tasks from SafetyGym or MetaDrive?\n\n- What is the segment length used in TSRA?\n\n- The main results in Table 1 are reported as averages across multiple tight and loose constraint thresholds. Please provide the separate, unaveraged normalized reward and cost returns for GAS and the baseline methods at each individual constraint level.", "limitations": "", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- The paper pinpoints two gaps in GM-assisted offline safe RL: weak stitching and inadequate reward–cost balancing.\n- Using goal functions to guide a constrained Advantage-Weighted Regression policy, aligning actions with feasible reward–cost targets is interesting.\n- GAS adapts to different test-time constraint levels without retraining.", "weaknesses": "- In RL, \"stitching\" traditionally refers to combining different trajectories via bootstrapping (the Bellman backup). The paper itself identifies this as a \"crucial ability\" that Generative Model (GM) methods lack. However, the paper's entire premise is to avoid the Bellman backup, which it calls the \"primary source of the OOD problem\". It then calls its own mechanism, supervised goal estimation using expectile regression, \"stitching\". This is a misleading contradiction. The proposed method does not \"stitch\" in the established sense.\n\n- In Fisor, the tight budget is 10 for safety gym and 5 for bullet gym. Based on the cost range table 2, only for CarRun the tight budget of 10% is below 5, while for antrun its 15, 3 times the “tight” budget of fisor.\n\n- The main results table (Table 1) shows that CDT is an unsafe policy in the average tight budget setting on the tasks chosen for the zero-shot ablation (BallRun, DroneRun, DroneCircle, etc.). Comparing GAS against a baseline that already fails is not informative.\n\n- The paper includes comparisons against recent algorithms like CAPS and CCAC in Table 6, but the main results in Table 1 instead feature baselines like CPQ, COptiDICE, and VOCE . Given that CAPS and CCAC are more recent and appear to be stronger competitors than some of those in the main table, they should have been included in the primary comparison.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "ZKoaorPCZ6", "reviewer": "Reviewer_RPXE", "summary": "The paper introduces a method called Goal-Assisted Stitching (GAS) to improve how reinforcement learning agents make safe and effective decisions using only offline data. It builds on recent generative model approaches but points out that these methods often fail to combine good parts of different trajectories and cannot properly balance rewards and safety limits. GAS addresses these issues by estimating achievable reward and cost goals directly from the dataset, instead of depending on user-specified targets. It does this through a statistical technique called expectile regression and adds two simple but effective steps: segmenting trajectories into shorter pieces to create more training samples, and reshaping the dataset so that the model sees a more even mix of safe and risky examples. Tested on common safe RL benchmarks, GAS shows more reliable safety under strict constraints and higher returns when the safety limits are relaxed, performing better than previous approaches in both stability and adaptability.", "questions": "- How sensitive is the performance of GAS to the choice of the expectile level \\alpha?\n- What are the training time and GPU resource requirements of GAS compared with CDT?\n- How would GAS perform on tasks with sparse or delayed rewards?\n- Since the reward and cost goal functions are trained independently, could the authors analyze their mutual consistency and verify that the resulting goal pairs correspond to feasible points on the reward–cost frontier?\n- Could the authors include additional baselines, as mentioned in the weakness section, to further verify the performance of GAS?\n\n## Reference\n\n**[1]** Zhuang, Zifeng, et al. \"Reinformer: Max-return sequence modeling for offline rl.\" arXiv preprint arXiv:2405.08740 (2024).\n\n**[2]** Chemingui, Yassine, et al. \"Constraint-adaptive policy switching for offline safe reinforcement learning.\" Proceedings of the AAAI Conference on Artificial Intelligence. Vol. 39. No. 15. 2025.\n\n**[3]** Cen, Zhepeng, et al. \"Learning from sparse offline datasets via conservative density estimation.\" arXiv preprint arXiv:2401.08819 (2024).\n\n**[4]** Rengarajan, Desik, et al. \"Reinforcement learning with sparse rewards using guidance from offline demonstration.\" arXiv preprint arXiv:2202.04628 (2022).\n\n**[5]** Yamagata, Taku, Ahmed Khalil, and Raul Santos-Rodriguez. \"Q-learning decision transformer: Leveraging dynamic programming for conditional sequence modelling in offline rl.\" International Conference on Machine Learning. PMLR, 2023.\n\n**[6]** Wang, Ruhan, and Dongruo Zhou. \"Safe Decision Transformer with Learning-based Constraints.\" 7th Annual Learning for Dynamics\\& Control Conference. PMLR, 2025.\n\n**[7]** Guan, Jiayi, et al. \"Voce: Variational optimization with conservative estimation for offline safe reinforcement learning.\" Advances in Neural Information Processing Systems 36 (2023): 33758-33780.\n\n**[8]** Wei, Honghao, et al. \"Adversarially trained weighted actor-critic for safe offline reinforcement learning.\" Advances in Neural Information Processing Systems 37 (2024): 52806-52835.", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- This work tackles a practical and important problem in Offline Safe Reinforcement Learning (OSRL) with Generative Models (GMs), focusing on two key challenges: balancing reward and cost, and improving the model’s ability to stitch useful transitions from different trajectories.\n- The paper presents solid theoretical support for the proposed GAS algorithm, especially through its use of expectile regression to estimate optimal reward and cost goals without relying on Bellman backups. This design effectively avoids common offline RL problems like value overestimation and out-of-distribution action selection. The clear theoretical foundation not only strengthens the methodological soundness of the approach but also enhances confidence in the reported empirical results, marking it as a notable strength of the work.\n- The paper shows strong empirical rigor through comprehensive experiments on multiple benchmarks and baselines under various constraint settings and ablation studies, convincingly demonstrating the robustness and effectiveness of the proposed GAS method.\n- The paper is clearly written and well organized, moving smoothly from the motivation to the method and then to the experiments, making it easy for readers to follow the main ideas and technical details.", "weaknesses": "- The reward and cost goal functions are learned from offline data using expectile regression, which makes them susceptible to the biases of the dataset. If the data are unbalanced or lack high-quality transitions, the estimated goals can become either too optimistic or too conservative, leading the policy in the wrong direction. Because these functions guide the policy’s optimization, even small estimation errors can distort the balance between reward and cost or result in unsafe actions.\n- The use of expectile regression introduces a hyperparameter \\alpha that controls the aggressiveness of the upper-tail estimation. The paper does not analyze the sensitivity of \\alpha or provide adaptive tuning [1].\n- Although the paper includes some information on infrastructure and training time, it does not clearly analyze computational efficiency. The method adds several components—segmented augmentation, dual goal functions, and dataset reshaping—but lacks a direct runtime or scalability comparison with CDT, making it difficult to assess the true efficiency impact [2].\n- While the proposed GAS framework demonstrates strong performance on dense-reward benchmarks such as SafetyGymnasium and Bullet-Safety-Gym, it lacks evaluation in sparse-reward settings [3-5].\n- Although the paper builds on recent progress in generative model–assisted offline safe RL, earlier works [6] have already explored similar ideas such as goal-conditioned augmentation and return relabeling. It would strengthen the paper to more clearly distinguish GAS from these methods, include [6] as a baseline for direct comparison, and incorporate additional recent strong baselines to ensure a fair and comprehensive evaluation [7-8].\n- The reward and cost goal functions are trained separately, without accounting for their inherent correlation in constrained settings. This independent training can produce inconsistent or even infeasible goal pairs, but the paper does not discuss or analyze this potential issue.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "dxoEoP4VFE", "reviewer": "Reviewer_oGCt", "summary": "This paper starts from two challenges in offline safe RL: 1) most existing methods lack capability to stitch good sub-trajectories from different sub-optimal trajectories, and 2) they struggle to balance reward and cost. The authors propose Goal-Assisted Stitching (GAS), which includes return augmentation and relabeling, expectile regression, dataset distribution reshaping, to overcome those issues. The authors run experiments on safe RL benchmarks and the results show that GAS show superior performance to baselines.", "questions": "- In fig 2, the stitched trajectory (denoted by red arrow) is confusing. For example, in the stitch of GAS showed by the second arrow (from the 2nd block of $\\tau_j$ to the 1st of $\\tau_k$), since the states in $\\tau_j,\\tau_k$ will be different, how do you make sure the 1st block of $\\tau_k$ can be stitched after the 2nd block of $\\tau_j$? Otherwise it's not even a valid trajectory. \n- In sec.5.2, GAS seeks better $R,C$ based on other sub-trajectory / segmentation as long as their lengths are the same. However, since the starting states $s_t, s_{t'}$ are different, their optimal $R$ and $C$ will also be different. Why do you think the $R,C$ of $s_{t'}$ can be a good reference for $s_t$?\n- The dataset reshaping seems to be also an important component. Do you have any ablation on that? E.g., the performance comparison with v.s. without reshaping.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Two motivations of paper, i.e., lack of stitching and failure to balance reward and cost, are clearly presented and important to safe offline RL community.\n- GAS exhibits consistent advantage over most baselines on DSRL in terms of reward improvement and cost constraint satisfaction, in both situations with tight or loose constraint.\n- The experiment is comprehensive. The authors also provide ablations to explain the motivation or show the effectiveness of relabel module.", "weaknesses": "- The \"stitching\" ability of the GAS is a little over-claimed. The authors claim the GAS has the capability to stitch the sub-trajectories. However, GAS seems to learn new return-to-go and cost-to-go targets instead of stitching the data. See more in the question section. I understand that the true stitching can be hard but I believe there is a large mismatch between \"stitching\" and GAS's implementation.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "xlqMPBaepM", "reviewer": "Reviewer_aKUM", "summary": "This paper proposes GAS, a method designed to balance the reward and cost conditions in offline safe reinforcement learning. GAS achieves a favorable trade-off between safety and reward by combining techniques such as return augmentation, return relabeling, goal generation, and dataset reshaping.", "questions": "No other questions.", "limitations": "", "rating": 4, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The paper adopts a reasonable approach to tackling offline safe reinforcement learning by focusing on balancing reward and cost.\n- The use of goal generation to address the reward–cost trade-off is well-motivated.\n- The proposed method demonstrates strong reward performance in tasks with relatively low safety requirements.", "weaknesses": "- Methodology\n    - The proposed method combines return augmentation, return relabeling, goal generation, and dataset reshaping, making it overly complex and heavily reliant on manual engineering design, with limited novelty. In fact, similar ideas have been explored in prior work.\n        - Reinformer (ICML 2024 [1]) generates RTG based on state using expectile regression, achieving strong reward stitching capabilities.\n        - CoPDT (NeurIPS 2025 [2]) applies a similar concept in offline safe RL, automatically generating RTG distributions conditioned on state and CTG, and sampling from them using expectile regression. It demonstrates strong cost prioritization and achieves safe performance in Safety Gymnasium Navigation tasks under more stringent thresholds ([10, 20, 40, 80]).\n        \n        Compared to these works, GAS requires extensive dataset relabeling and augmentation, while methods like Reinformer and CoPDT only rely on learning the expectile distribution, making them simpler and more flexible.\n        \n    - The model’s input design—comprising four types of tokens (reward goal, cost goal, RTG, and CTG)—appears unnecessarily complicated. Since the reward goal is already generated, why are RTG and CTG still needed as additional inputs? Although the paper states that reward goal and cost goal are designed to avoid conflicts, what happens when RTG and CTG themselves conflict? How does the model ensure that reward goal and cost goal remain the primary optimization targets, rather than being overshadowed by RTG and CTG? Furthermore, if RTG and CTG do conflict, how is CTG’s priority guaranteed?\n- Experiments\n    - The experiments are mainly conducted on BulletSafetyGym and the Circle tasks from Safety Gymnasium, but they omit more standard and safety-critical environments, such as other Navigation and Velocity tasks from Safety Gymnasium.\n    - The safety requirements in the experimental settings are too lenient.\n        - In BulletSafetyGym, satisfying safety constraints is already relatively easy.\n        - In Safety Gymnasium, most tasks have constraint thresholds above 20 (for a 10% constraint ratio), which are too loose to evaluate the model’s ability under strict safety conditions.\n        - In the “loose constraint” settings, the thresholds remain excessively high, making the evaluation of safety performance insufficient.\n        \n        Overall, the experiments mainly verify that GAS achieves high reward under low safety requirements, but its applicability to high-safety-demand scenarios remains questionable.\n        \n    - The method involves numerous hyperparameters, including expectile level $\\alpha$, reward relabel level $\\delta$, dataset reshape threshold q% and sample probability $\\epsilon$, among others. The authors should conduct a sensitivity analysis to evaluate how these hyperparameters affect performance.\n\n[1] Reinformer: Max-Return Sequence Modeling for offline RL\n\n[2] Adaptable Safe Policy Learning from Multi-task Data with Constraint Prioritized Decision Transformer", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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:

Decision: Accept (Poster)