Q-learning Penalized Transformer for Safe Offline Reinforcement Learning

Shengchao Hu; Peng Wang; Jifeng Hu; Qiyang Zhou; Anning Hu; Li Shen; Ya Zhang; Dacheng Tao

cs.LGSep 28, 2026

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

Abstract

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 $\uparrow$	cost $\downarrow$	reward $\uparrow$	cost $\downarrow$	reward $\uparrow$	cost $\downarrow$	reward $\uparrow$	cost $\downarrow$	reward $\uparrow$	cost $\downarrow$	reward $\uparrow$	cost $\downarrow$	reward $\uparrow$	cost $\downarrow$	reward $\uparrow$	cost $\downarrow$	reward $\uparrow$	cost $\downarrow$
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

\uparrow

symbol denotes that the higher reward, the better. The

\downarrow

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 $Q^{r}$	Train with $Q^{c}$	Inf. with ensemble	harddense Reward	harddense Cost	HopperVelocity Reward	HopperVelocity Cost
1					$0.37\pm 0.19$	$1.00\pm 0.08$	$0.04\pm 0.02$	$1.49\pm 0.12$
2	✓				$0.40\pm 0.05$	$0.94\pm 0.04$	$0.54\pm 0.03$	$0.65\pm 0.03$
3	✓	✓			$0.48\pm 0.08$	$0.94\pm 0.05$	$0.85\pm 0.05$	$0.99\pm 0.04$
4	✓		✓		$0.43\pm 0.06$	$0.14\pm 0.10$	$0.15\pm 0.02$	$0.22\pm 0.01$
5	✓	✓	✓		$0.49\pm 0.04$	$0.84\pm 0.02$	$0.69\pm 0.04$	$0.51\pm 0.02$
6	✓			✓	$0.46\pm 0.05$	$0.91\pm 0.06$	$0.56\pm 0.02$	$0.60\pm 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

Q^{r}

” and “Train with

Q^{c}

” 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

\eta_{1}

and

\eta_{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
SafetyGymnasium	Mujoco	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\pm 0.04$	$0.50\pm 0.02$	$0.52\pm 0.03$	$0.50\pm 0.02$	$0.48\pm 0.04$
Cost	$0.84\pm 0.02$	$0.83\pm 0.03$	$0.90\pm 0.05$	$0.81\pm 0.03$	$0.80\pm 0.03$

Appendix

Table 5 : Impact of the number of candidate target reward and cost pairs in the harddense task in the MetaDrive setting.

Explore similar work

Safe Flow Q-Learning: Offline Safe Reinforcement Learning with Reachability-Based Flow Policies
Mar 16, 2026Mumuksh Tayal, Manan Tayal, Ravi PrakashFlow Policies Model-Based Reinforcement Learning
Indian Institute of Science, India
Safe-Support Q-Learning: Learning without Unsafe Exploration
Apr 28, 2026Yeeun Lim, Narim Jeong, Donghwan LeeQ-Learning Safety Constraints
HD Korea Shipbuilding & Offshore Engineering · Department of Electrical Engineering, KAIST, Daejeon, South Korea
Lyapunov-Guided Self-Alignment: Test-Time Adaptation for Offline Safe Reinforcement Learning
Apr 29, 2026Seungyub Han, Hyungjin Kim, Jungwoo LeeModel-Based Reinforcement Learning Offline Reinforcement Learning
Seoul National University

Mar 16, 2026, cs.LGJ/K move · Enter open · S save

Safe Flow Q-Learning: Offline Safe Reinforcement Learning with Reachability-Based Flow Policies

Mumuksh Tayal, Manan Tayal, Ravi Prakash

Indian Institute of Science, India

Flow Policies Model-Based Reinforcement Learning Q-Learning

Open paper AskPDF arXiv

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.

Abstract

Explore similar work

Safe Flow Q-Learning: Offline Safe Reinforcement Learning with Reachability-Based Flow Policies

Safe-Support Q-Learning: Learning without Unsafe Exploration

Lyapunov-Guided Self-Alignment: Test-Time Adaptation for Offline Safe Reinforcement Learning