# GAS: Enhancing Reward-Cost Balance of Generative Model-assisted Offline Safe RL

arXiv: https://arxiv.org/abs/2602.05323 · HTML: https://mlsanity.com/paper/2602.05323/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Safe offline RL, one of 13 directions in Offline RL and imitation learning from data.
- Builds on 32 of its 41 references that are in the map.
- First posted on arXiv 5 Feb 2026.

## Placed in

- Offline RL and imitation learning from data / [Safe offline RL](https://mlsanity.com/offlinerl/safe-offline-rl/)

## Builds on (most cited first, 30 of 32)

- [Conservative Q-Learning for Offline Reinforcement Learning](https://mlsanity.com/paper/2006.04779/) (2020)
- [Decision Transformer: Reinforcement Learning via Sequence Modeling](https://mlsanity.com/paper/2106.01345/) (2021)
- [D4RL: Datasets for Deep Data-Driven Reinforcement Learning](https://mlsanity.com/paper/2004.07219/) (2020)
- [Offline Reinforcement Learning with Implicit Q-Learning](https://mlsanity.com/paper/2110.06169/) (2021)
- [Off-Policy Deep Reinforcement Learning without Exploration](https://mlsanity.com/paper/1812.02900/) (2018)
- [Advantage-Weighted Regression: Simple and Scalable Off-Policy Reinforcement Learning](https://mlsanity.com/paper/1910.00177/) (2019)
- [Behavior Regularized Offline Reinforcement Learning](https://mlsanity.com/paper/1911.11361/) (2019)
- [Distributional Reinforcement Learning with Quantile Regression](https://mlsanity.com/paper/1710.10044/) (2017)
- [IDQL: Implicit Q-Learning as an Actor-Critic Method with Diffusion Policies](https://mlsanity.com/paper/2304.10573/) (2023)
- [Responsive Safety in Reinforcement Learning by PID Lagrangian Methods](https://mlsanity.com/paper/2007.03964/) (2020)
- [Offline Reinforcement Learning with Fisher Divergence Critic Regularization](https://mlsanity.com/paper/2103.08050/) (2021)
- [Keep Doing What Worked: Behavioral Modelling Priors for Offline Reinforcement Learning](https://mlsanity.com/paper/2002.08396/) (2020)
- [Online Decision Transformer](https://mlsanity.com/paper/2202.05607/) (2022)
- [Delving into Deep Imbalanced Regression](https://mlsanity.com/paper/2102.09554/) (2021)
- [Safety-Gymnasium: A Unified Safe Reinforcement Learning Benchmark](https://mlsanity.com/paper/2310.12567/) (2023)
- [Q-learning Decision Transformer: Leveraging Dynamic Programming for Conditional Sequence Modelling in Offline RL](https://mlsanity.com/paper/2209.03993/) (2022)
- [Constraints Penalized Q-learning for Safe Offline Reinforcement Learning](https://mlsanity.com/paper/2107.09003/) (2021)
- [Balanced MSE for Imbalanced Visual Regression](https://mlsanity.com/paper/2203.16427/) (2022)
- [COptiDICE: Offline Constrained Reinforcement Learning via Stationary Distribution Correction Estimation](https://mlsanity.com/paper/2204.08957/) (2022)
- [Constrained Decision Transformer for Offline Safe Reinforcement Learning](https://mlsanity.com/paper/2302.07351/) (2023)
- [Elastic Decision Transformer](https://mlsanity.com/paper/2307.02484/) (2023)
- [Safe Offline Reinforcement Learning with Feasibility-Guided Diffusion Model](https://mlsanity.com/paper/2401.10700/) (2024)
- [Datasets and Benchmarks for Offline Safe Reinforcement Learning](https://mlsanity.com/paper/2306.09303/) (2023)
- [The In-Sample Softmax for Offline Reinforcement Learning](https://mlsanity.com/paper/2302.14372/) (2023)
- [Decision ConvFormer: Local Filtering in MetaFormer is Sufficient for Decision Making](https://mlsanity.com/paper/2310.03022/) (2023)
- [Reinformer: Max-Return Sequence Modeling for Offline RL](https://mlsanity.com/paper/2405.08740/) (2024)
- [Waypoint Transformer: Reinforcement Learning via Supervised Learning with Intermediate Targets](https://mlsanity.com/paper/2306.14069/) (2023)
- [Rethinking Decision Transformer via Hierarchical Reinforcement Learning](https://mlsanity.com/paper/2311.00267/) (2023)
- [Constraint-Adaptive Policy Switching for Offline Safe Reinforcement Learning](https://mlsanity.com/paper/2412.18946/) (2024)
- [OASIS: Conditional Distribution Shaping for Offline Safe Reinforcement Learning](https://mlsanity.com/paper/2407.14653/) (2024)

## Abstract (arXiv)

Offline Safe Reinforcement Learning (OSRL) aims to learn a policy to achieve high performance in sequential decision-making while satisfying 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 values. However, GM-assisted methods face two major challenges in OSRL: (1) lacking the ability to "stitch" optimal transitions from suboptimal trajectories within the dataset, and (2) struggling to balance reward targets with cost targets, particularly when they are conflict. 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.