Q-learning Penalized Transformer for Safe Offline Reinforcement Learning
Authors: Shengchao Hu, Peng Wang, Jifeng Hu, Qiyang Zhou, Anning Hu, Li Shen, Ya Zhang, Dacheng Tao
Organizations: Shanghai Jiao Tong University, China · Shenzhen Campus of Sun Yat-sen University, China · Jilin University, China · Nanyang Technological University, Singapore
This paper addresses the problem of safe offline reinforcement learning, which involves training a policy to satisfy safety constraints using an offline dataset. This problem is inherently challenging as it requires balancing three highly interconnected and competing objectives: satisfying safety constraints, maximizing rewards, and adhering to the behavior regularization imposed by the offline dataset. To tackle this trilogy challenge, we propose Q-learning Penalized Transformer policy (QPT), a \emph{training--inference consistent} framework that bridges conditional sequence modeling with constraint-aware value estimation. QPT trains a Transformer policy that generates actions conditioned on trajectory context and target return/cost, retaining strong behavior regularization. To inject explicit safety semantics during learning, we augment sequence-model training with a Q-shaped penalty using learned reward and cost Q-functions to favor high return under low constraint violation. At inference, the same Q-functions enforce the cost threshold and choose the highest-reward feasible action, closing the loop between training and deployment. We provide a principled analysis under stylized near-deterministic CMDPs, characterizing how Q-penalized conditional generation improve safety and performance. Empirically, QPT consistently outperforms strong safe offline RL baselines across 38 tasks on the DSRL benchmark, and exhibits robust zero-shot adaptation to different constraint thresholds.
Figures & tables
QPT (Ours)
BC-Safe
CDT
BCQ-Lag
CPQ
COptiDICE
FISOR
OASIS
CAPS
Task
reward ↑
cost ↓
reward ↑
cost ↓
reward ↑
cost ↓
reward ↑
cost ↓
reward ↑
cost ↓
reward ↑
cost ↓
reward ↑
cost ↓
reward ↑
cost ↓
reward ↑
cost ↓
PointButton1
0.13
0.81
0.10
0.63
0.62
7.17
0.24
1.73
0.69
3.2
0.13
1.35
0.03
0.81
0.26
1.33
0.02
0.30
PointButton2
-0.01
0.88
0.04
0.58
0.31
5.15
0.4
2.66
0.58
4.3
0.15
1.51
0.02
0.69
0.31
1.89
0.01
0.92
PointCircle1
0.58
0.93
0.45
0.67
0.57
0.75
0.17
1.04
0.43
0.29
0.78
15.64
0.60
12.8
0.32
0.00
0.25
2.10
PointCircle2
0.62
0.92
0.49
0.44
0.61
1.39
0.53
8.35
0.28
0.77
0.78
25.94
0.70
11.79
0.36
0.20
0.38
0.40
PointGoal1
0.68
0.65
0.42
0.70
0.70
1.54
0.59
1.30
0.68
0.76
0.35
1.75
0.54
2.73
0.66
4.35
0.22
1.05
Table 1 : Complete evaluation results of the normalized reward and cost. The cost threshold is 1. The ↑ symbol denotes that the higher reward, the better. The ↓ symbol denotes that the lower normalized cost (up to threshold 1, corresponding to original cost limit 10), the better. Each value is averaged over 20 evaluation episodes and 3 random seeds. Bold : Safe agents whose normalized cost is smaller than 1. Gray : Unsafe agents with normalized costs exceeding 1. Blue : Safe agent with the highest reward.
Exp
Data aug.
Train with Qr
Train with Qc
Inf. with ensemble
harddense Reward
harddense Cost
HopperVelocity Reward
HopperVelocity Cost
1
0.37±0.19
1.00±0.08
0.04±0.02
1.49±0.12
2
✓
0.40±0.05
0.94±0.04
0.54±0.03
0.65±0.03
3
✓
✓
0.48±0.08
0.94±0.05
0.85±0.05
0.99±0.04
4
✓
✓
0.43±0.06
0.14±0.10
0.15±0.02
0.22±0.01
5
✓
✓
✓
0.49±0.04
0.84±0.02
0.69±0.04
0.51±0.02
6
✓
✓
0.46±0.05
0.91±0.06
0.56±0.02
0.60±0.04
Table 2 : Impact of different components. Average scores and standard deviations are reported over three random seeds for the harddense task in the MetaDrive setting and the HopperVelocity setting. “Train with Qr ” and “Train with Qc ” indicate whether the corresponding penalization in Equation 8 is applied. “Data aug.” refers to the use of data augmentation, while “Inf. with ensemble" denotes ensemble applied at inference time.
Figure 1 : Results of the impact of hyper-parameters η1 and η2 . Each column is a task with one hyper-parameter stable. The x-axis is the training steps. The first row shows the evaluated normalized reward, and the second row shows the evaluated normalized cost. All plots are averaged among 3 random seeds and 20 trajectories for each seed. The solid line is the mean value, and the light shade represents the area within one standard deviation.
Figure 2 : Results of zero-shot adaption to different cost returns. Each column is a task. The x-axis is the target cost return. The first row shows the evaluated normalized reward, and the second row shows the evaluated normalized cost under different target costs. All plots are averaged among 3 random seeds and 20 trajectories for each seed. The solid line is the mean value, and the light shade represents the area within one standard deviation.
Appendix figures & tables5 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 3 : Illustrative example of the data-augmentation procedure ( Liu et al., 2023b ) .
Figure 4 : Visualization of the simulation environments and representative tasks ( Liu et al., 2023a ) .
Benchmarks
Backends
Environments
Agents
Difficulty Levels
Total Tasks
Dataset Trajectories
SafetyGymnasium
Mujoco
Goal, Button, Push, Circle
Point, Car
2
16
40310
Velocity
Ant, HalfCheetah, Hopper, Swimmer, Walker2d
1
5
11399
BulletSafetyGym
PyBullet
Run, Circle
Ball, Car, Drone, Ant
1
8
14498
MetaDrive
Panda3D
Driving
Vehicle
3
9
9000
Appendix
Table 3 : Overview of the safe RL benchmarks and tasks for dataset collection ( Liu et al., 2023a ) .
Parameter
All tasks
Parameter
All tasks
Number of layers
3
Number of attention heads
8
Embedding dimension
128
Batch size
2048
Context length K
10
Learning rate
0.0001
Droupout
0.1
Adam betas
(0.9, 0.999)
Grad norm clip
0.25
Cost threshold
10
Training steps (BulletGym, SafetyGym)
100000
Training steps (MetaDrive)
200000
Appendix
Table 4 : Hyperparameters for QPT
1
10
30
50
100
Reward
0.49±0.04
0.50±0.02
0.52±0.03
0.50±0.02
0.48±0.04
Cost
0.84±0.02
0.83±0.03
0.90±0.05
0.81±0.03
0.80±0.03
Appendix
Table 5 : Impact of the number of candidate target reward and cost pairs in the harddense task in the MetaDrive setting.
Offline safe reinforcement learning (RL) seeks reward-maximizing policies from static datasets under strict safety constraints. Existing methods often rely on soft expected-cost objectives or iterative generative inference, which can be insufficient for safety-critical real-time control. We propose Safe Flow Q-Learning (SafeFQL), which extends FQL to safe offline RL by combining a Hamilton--Jacobi reachability-inspired safety value function with an efficient one-step flow policy. SafeFQL learns the safety value via a self-consistency Bellman recursion, trains a flow policy by behavioral cloning, and distills it into a one-step actor for reward-maximizing safe action selection without rejection sampling at deployment. Empirically, SafeFQL trades modestly higher offline training cost for substantially lower inference latency than diffusion-style safe generative baselines, which is advantageous for real-time safety-critical deployment. Across boat navigation, and Safety Gymnasium MuJoCo tasks, SafeFQL matches or exceeds prior offline safe RL performance while substantially reducing constraint violations.
Ensuring safety during reinforcement learning (RL) training is critical in real-world applications where unsafe exploration can lead to devastating outcomes. While most safe RL methods mitigate risk through constraints or penalization, they still allow exploration of unsafe states during training. In this work, we adopt a stricter safety requirement that eliminates unsafe state visitation during training. To achieve this goal, we propose a Q-learning-based safe RL framework that leverages a behavior policy supported on a safe set. Under the assumption that the induced trajectories remain within the safe set, this policy enables sufficient exploration within the safe region without requiring near-optimality. We adopt a two-stage framework in which the Q-function and policy are trained separately. Specifically, we introduce a KL-regularized Bellman target that constrains the Q-function to remain close to the behavior policy. We then derive the policy induced from the trained Q-values and propose a parametric policy extraction method to approximate the optimal policy. Our approach provides a unified framework that can be adapted to different action spaces and types of behavior policies. Experimental results demonstrate that the proposed method achieves stable learning and well-calibrated value estimates and yields safer behavior with comparable or better performance than existing baselines.
Yeeun Lim, Narim Jeong, Donghwan Lee
HD Korea Shipbuilding & Offshore Engineering · Department of Electrical Engineering, KAIST, Daejeon, South Korea
Offline reinforcement learning (RL) agents often fail when deployed, as the gap between training datasets and real environments leads to unsafe behavior. To address this, we present SAS (Self-Alignment for Safety), a transformer-based framework that enables test-time adaptation in offline safe RL without retraining. In SAS, the main mechanism is self-alignment: at test time, the pretrained agent generates several imagined trajectories and selects those satisfying the Lyapunov condition. These feasible segments are then recycled as in-context prompts, allowing the agent to realign its behavior toward safety while avoiding parameter updates. In effect, SAS turns Lyapunov-guided imagination into control-invariant prompts, and its transformer architecture admits a hierarchical RL interpretation where prompting functions as Bayesian inference over latent skills. Across Safety Gymnasium and MuJoCo benchmarks, SAS consistently reduces cost and failure while maintaining or improving return.