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

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.

💡 Research Summary

The paper addresses two fundamental shortcomings of recent generative‑model‑based approaches for Offline Safe Reinforcement Learning (OSRL). First, existing methods cannot “stitch” together optimal transitions from sub‑optimal trajectories, limiting their ability to exploit the full information contained in offline datasets. Second, they rely on manually specified reward‑cost targets, which often lead to infeasible or overly conservative policies because the trade‑off between performance and safety is not automatically calibrated.

Goal‑Assisted Stitching (GAS) is introduced to solve both problems. GAS begins by augmenting and relabeling the dataset at the transition level. For each transition it computes temporally segmented returns – i.e., reward‑to‑go and cost‑to‑go values over multiple horizons – and creates new transition tuples that combine pieces from different trajectories. This “temporal segmented return augmentation” yields a richer set of state‑action‑return triples, enabling the construction of high‑quality synthetic trajectories that were not present in the original data.

Next, GAS learns two goal functions, one for reward (G_R) and one for cost (G_C), using expectile regression on the augmented dataset. Expectile regression estimates conditional expectiles (e.g., τ = 0.9 for reward, τ = 0.1 for cost), thereby providing data‑driven estimates of the best achievable reward and the safest achievable cost. These goal functions replace human‑specified (R̂, Ĉ) pairs: during policy inference the model queries G_R and G_C to obtain feasible target returns that are guaranteed to lie within the support of the offline data. Consequently, the policy can automatically balance reward maximization against constraint satisfaction without manual tuning.

To further improve training stability, GAS reshapes the dataset to obtain a more uniform distribution over the reward‑cost return space. By re‑sampling or re‑weighting under‑represented regions, the expectile regressors receive balanced supervision, and the downstream policy avoids bias toward densely populated return regions.

The authors evaluate GAS on two benchmark environments (e.g., CarCircle and a drone navigation task) across twelve distinct safety scenarios and compare against eight strong baselines, including Constrained Decision Transformer (CDT), Q‑Conditioned DT (QDT), and recent stitching‑focused methods. Results show that under tight cost thresholds (L ≤ 0.1) GAS improves safety compliance by roughly 15 % relative to the best prior GM‑based method, while under looser thresholds (L ≥ 0.5) it achieves a 6 % increase in cumulative reward. Importantly, the automatically inferred goal pairs lead to a 20 % better reward‑cost trade‑off than manually set targets, demonstrating the practical advantage of the goal‑function mechanism.

Additional ablation studies confirm that increasing the attention memory length in DT‑style models does not enhance stitching ability; performance remains flat when K is varied from 1 to 10, supporting the authors’ claim that temporal attention is ineffective for OSRL. In contrast, GAS’s transition‑level stitching, guided by the learned goal functions, yields consistent gains.

In summary, GAS contributes three novel components: (1) transition‑level temporal segmented return augmentation and relabeling for enhanced stitching, (2) expectile‑regression‑based reward and cost goal functions that automatically calibrate feasible targets, and (3) dataset reshaping for balanced return distributions. Together these innovations enable generative‑model‑based OSRL to achieve superior performance while respecting safety constraints, opening a path toward safe, offline learning in real‑world domains such as autonomous driving, robotics, and finance.