CORMA: Constraint-Routed Decision Mamba for Generalizable Multi-Constraint Offline Safe Reinforcement Learning

Changxu Ni; Feiran You; Yongan Guo; Hongyang Du; Kang Wei; Le Liang; Yiyang Ni; Jun Li

cs.LGOct 25, 2026

CORMA: Constraint-Routed Decision Mamba for Generalizable Multi-Constraint Offline Safe Reinforcement Learning

Authors: Changxu Ni, Feiran You, Yongan Guo, Hongyang Du, Kang Wei, Le Liang, Yiyang Ni, Jun Li

Organizations: Southeast University · The University of Hong Kong · Nanjing University of Posts and Telecommunications

Abstract

Offline safe reinforcement learning (RL) learns safety-aware policies from static datasets while satisfying safety constraints during deployment. Motivated by the strong trajectorymodeling capabilities of Transformers, recent Decision Transformer (DT)-based methods formulate offline safe RL as conditional sequence modeling by conditioning action generation on return-to-go (RTG) and cost-to-go (CTG) tokens, thereby enabling adaptation to varying safety budgets. However, existing DT-based methods are developed for single-constraint tasks. For tasks with multiple constraints, generalizing learned policies to unseen combinations of safety budgets remains challenging due to the conflicting requirements imposed by heterogeneous constraints. Moreover, the quadratic complexity of Transformer architectures significantly increases the computational overhead in multi-constraint tasks, as representing additional constraints requires more CTG tokens. Recently, Mamba has emerged as an efficient sequence-modeling backbone, offering linear computational complexity and strong generalization capabilities. In this paper, we propose CORMA, a constraint-routed decision Mamba framework for generalizable multi-constraint offline safe RL. CORMA decomposes policy inference into constraint-aware RTG estimation and adaptive action generation, where constraint-specific experts model feasible RTG distributions and a router-guided product-of-experts mechanism fuses these distributions to provide safety-aware guidance for a Mamba-based action decoder. A three-stage optimization strategy is designed to improve training stability. Extensive experiments on OSRL benchmarks demonstrate that CORMA achieves superior reward-safety trade-offs, reduces inference latency by 37% compared with DT-based methods, and exhibits stronger zero-shot generalization to unseen safety-budget combinations than offline safe RL baselines.

Figures & tables

Figure 1: Illustration of the challenges. After pre-training on seen safety-budget regions, Transformer-based sequence models struggle under unseen safety-budget combinations, whereas CORMA generalizes more robustly.

Figure 2: Inference efficiency and zero-shot generation analysis of CDT and CDM under multi-constraint settings. (a) Inference latency with increasing constraints. (b) Safety performance under different safety-budget regions.

Figure 3: Overview of CORMA. Stage I pre-trains the Mamba-based action decoder using offline trajectories. Stage II trains constraint-specific RTG experts to model feasible RTG distributions under different constraints. Stage III trains the router to adaptively fuse expert predictions through PoE and generate constraint-aware RTG guidance for safe action execution.

Method	BC-v2		CC-v2		BC-v3		CC-v3		DC-v3		PB-v3		PG-v3		CG-v3		PP-v3		CP-v3
Method	$r\uparrow$	$c\downarrow$	$r\uparrow$	$c\downarrow$	$r\uparrow$	$c\downarrow$	$r\uparrow$	$c\downarrow$	$r\uparrow$	$c\downarrow$	$r\uparrow$	$c\downarrow$	$r\uparrow$	$c\downarrow$	$r\uparrow$	$c\downarrow$	$r\uparrow$	$c\downarrow$	$r\uparrow$	$c\downarrow$
BCQ-Lag	88.3	1.25	75.9	0.99	74.5	1.71	69.7	2.88	92.1	1.92	80.2	1.10	86.3	1.01	81.2	1.18	79.1	1.29	80.3	1.22
BEAR-Lag	90.1	1.02	83.8	1.05	86.3	1.99	73.2	2.42	94.8	2.76	85.9	1.08	89.2	1.17	84.9	1.69	82.7	1.33	83.1	1.29
PDCA	87.7	0.98	81.3	0.97	78.7	1.76	71.6	1.79	89.1	1.57	86.3	1.46	89.3	1.45	87.3	1.67	83.0	1.06	85.0	1.08
POCE	88.8	1.42	81.5	1.08	79.4	1.37	72.4	1.38	86.8	1.33	87.2	1.01	90.1	0.99	88.6	1.01	84.2	1.03	86.1	1.02
OASIS	82.5	0.95	77.8	1.02	73.8	1.18	67.5	1.24	81.6	1.78	84.5	0.99	87.0	0.97	85.5	0.98	80.8	1.01	83.2	0.99

Table 1: Evaluation results on Bullet-Safety-Gym and Safety-Gymnasium tasks. Entries report normalized reward (

r\uparrow\in[0,100]

) and maximum normalized cost among all constraints (

c\downarrow

), averaged over 5 seeds and 20 episodes. Black bold indicates constraint satisfaction (

c\leq 1

). Black bold highlights the best safety-performance trade-off among all methods.

Figure 4: Fine-tuning convergence under new safety-budget combinations in the DC-v3 tasks. The policies are fine-tuned using trajectories that satisfy the new safety-budget combinations. (a) Normalized accumulative rewards during fine-tuning. (b) Normalized collision costs. (c) Normalized high-speed costs. (d) Normalized low-speed costs.

Figure 5: Zero-shot generalization under unseen safety-budget combinations in the DC-v3 task. (a) Normalized accumulative rewards. (b) Normalized Collision costs. (c) Normalized High-speed costs. (d) Normalized Low-speed costs.

Phase	CDT	CDM	CORDT	CORMA
Backbone (h)	15.82	11.64	15.81	11.62
Experts (h)	–	–	1.42	1.34
Router (h)	–	–	0.71	0.68
Total Train (h)	15.82	11.64	17.94	13.64
Inference (s)	0.008	0.003	0.010	0.005

Table 2: Time comparison during training and deployment. Training times are reported in hours (h), and deployment inference times are reported in seconds (s).

Appendix figures & tables8 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 6: Visualization of the simulation environments and representative tasks (Liu et al., 2023a).

	Hyperparameter	Value
conditional action decoder $\theta^{B}$	Mamba num layers	2
	Mamba embedding dim	128
	Mamba expand factor	2
	Mamba convolution size	4
	update steps	30000
	batch size	1024

Appendix

Table 3: Hyperparameter choices of CORMA.

Task	Method	Safety-budget Combinations
		S1		S2		S3		S4		S5
		Reward	Cost- $N$	Reward	Cost- $N$	Reward	Cost- $N$	Reward	Cost- $N$	Reward	Cost- $N$
BC-v2	BCQ-Lag	88.6	1.06	82.4	1.21	85.7	1.48	80.1	0.94	81.2	1.03
	BEAR-Lag	84.7	1.12	88.2	1.07	91.5	1.01	86.4	1.09	87.1	1.15
	PDCA	88.9	1.15	83.8	0.96	88.5	0.99	81.2	1.05	82.6	0.97
	POCE	89.6	1.02	84.5	0.99	89.2	1.16	82.2	1.01	87.0	1.19

Appendix

Table 4: Detailed evaluation under seen two-safety-budget combinations. Reward is normalized to

0

–

100

, and the safety threshold is normalized to 1. S1–S5 denote threshold combinations appearing in the offline dataset. Bold indicates constraint satisfaction (

c\leq 1

), and underlined bold indicates the highest reward among feasible methods.

Task	Method	safety-budget combinations
		S1		S2		S3		S4		S5
		Reward	Cost- $N$	Reward	Cost- $N$	Reward	Cost- $N$	Reward	Cost- $N$	Reward	Cost- $N$
BC-v3	BCQ-Lag	72.6	1.82	77.8	1.63	75.1	1.48	73.4	1.91	73.6	1.71
	BEAR-Lag	84.7	1.61	88.9	1.86	86.1	1.72	83.8	2.24	88.0	2.02
	PDCA	76.8	1.48	81.5	1.60	80.3	1.49	77.1	1.94	77.8	1.88
	POCE	77.2	1.45	82.6	1.29	80.4	1.18	78.5	1.50	78.3	1.42

Appendix

Table 5: Detailed evaluation under seen three-safety-budget combinations. Reward is normalized to

0

–

100

, and the safety threshold is normalized to 1. S1–S5 denote threshold combinations appearing in the offline dataset. Bold indicates constraint satisfaction (

c\leq 1

), and underlined bold indicates the highest reward among feasible methods.

Task	CDT-MC		CORMA w/o PoE		CORDT		CORMA (Ours)
Task	Reward	Cost- $N$	Reward	Cost- $N$	Reward	Cost- $N$	Reward	Cost- $N$
BC-v2	71.2	0.98	72.8	1.22	83.4	0.98	85.5	0.96
CC-v2	72.1	1.24	78.5	1.02	72.0	1.28	70.6	0.98
DC-v2	77.0	1.11	75.9	1.05	82.4	0.97	84.8	0.93
BC-v3	84.6	2.04	78.6	1.94	81.0	1.46	86.5	0.98
CC-v3	70.7	1.18	73.2	1.26	66.5	1.13	64.2	0.98

Appendix

Table 6: Trajectory stitching ability analysis on multi-constraint tasks. Reward (

r

) is normalized to 0–100, and Cost-N (

c

) denotes the maximum normalized constraint violation. Bold indicates constraint satisfaction (

c\leq 1

), and underlined bold indicates the highest reward among feasible methods.

Task	CDT-MC		CORMA w/o PoE		CORDT		CORMA
Task	Reward	Cost- $N$	Reward	Cost- $N$	Reward	Cost- $N$	Reward	Cost- $N$
BC-v3	78.0	2.15	68.5	1.80	72.4	1.45	80.2	0.98
CC-v3	86.7	1.65	77.2	1.41	75.5	1.28	72.0	0.97
DC-v3	90.5	2.14	87.0	2.33	85.2	1.65	85.3	0.99
PB-v3	78.2	1.99	76.4	1.54	75.8	0.99	75.1	0.98
CB-v3	85.6	1.59	84.8	1.39	84.2	1.08	84.0	0.98

Appendix

Table 7: Zero-shot generalization to unseen safety-budget combinations. Reward is normalized to 0–100, and Cost-

N

denotes the maximum normalized constraint violation. Bold indicates constraint satisfaction (

c\leq 1

), and underlined bold indicates the highest reward among feasible methods.

$K$	Transformer		Mamba
$K$	$r\uparrow$	$c\downarrow$	$r\uparrow$	$c\downarrow$
5	89.8	2.44	88.5	2.21
10	86.1	2.12	82.4	1.96
15	85.9	1.96	83.6	1.68
20	80.7	1.61	77.6	1.42
25	79.4	1.58	75.0	1.24

Appendix

Table 8: Effect of sequence length

K

on Stage-1 action decoder training. Results are averaged over three representative three-constraint tasks: DC-v3, PB-v3, and CG-v3. Reward is normalized to 0–100, and Cost-

N

denotes the maximum normalized constraint violation. Bold indicates constraint satisfaction (

c\leq 1

), and underlined bold indicates the highest reward among feasible settings.

Method	DC-v3		PB-v3		CG-v3		CP-v3		Average
Method	$r\uparrow$	$c\downarrow$	$r\uparrow$	$c\downarrow$	$r\uparrow$	$c\downarrow$	$r\uparrow$	$c\downarrow$	$r\uparrow$	$c\downarrow$
CORMA w/o 3-stage Training	91.3	2.86	92.1	2.24	88.5	2.65	89.7	2.18	90.4	2.48
CORMA w/ Mean-Regression Router	84.7	1.31	88.6	1.18	85.1	1.35	85.9	1.20	86.1	1.26
CORMA w/o Balanced Router Sampling	85.8	1.17	88.7	1.09	89.6	1.21	86.5	1.11	87.7	1.15
CORMA w/o Expert Freezing	83.1	1.45	87.1	1.28	83.2	1.48	82.5	1.30	84.0	1.38
CORMA (Ours)	88.0	0.95	90.2	0.94	91.1	0.90	88.3	0.95	89.4	0.94

Appendix

Table 9: Ablation study on the three-stage optimization strategy of CORMA across four multi-constraint tasks. Reward (

r

) is normalized to 0–100, and Cost-

N

(

c

) denotes the maximum normalized constraint violation.

Safe RL

Safe reinforcement learning studies how agents can learn and act while respecting safety constraints and limiting harmful outcomes.

See all 217 papers

Q-learning Penalized Transformer for Safe Offline Reinforcement Learning
Sep 28, 2026Shengchao Hu, Peng Wang, Jifeng Hu +5Reinforcement Learning Transformer-Based RL
Shanghai Jiao Tong University, China · Shenzhen Campus of Sun Yat-sen University, China · Jilin University, China +1
Beyond Hard Constraints: Budget-Conditioned Reachability For Safe Offline Reinforcement Learning
Mar 8, 2026Janaka Chathuranga Brahmanage, Akshat KumarSafe RL Constrained RL
School of Computing and Information Systems, Singapore Management University.
CDCP: Conditional Diffusion Model with Contextual Prompts for Multi-task Offline Safe Reinforcement Learning
Jul 4, 2026Jiayi Guan, Tianle Zhang, Li Shen +8Reinforcement Learning Constrained RL
Tongji University, Shanghai, China · JD Explore Academy, Beijing, China · School of Cyber Science and Technology, Shenzhen Campus of Sun Yat-sen University, Shenzhen, China +2

Sep 28, 2026J/K move · Enter open · S save

Q-learning Penalized Transformer for Safe Offline Reinforcement Learning

Shengchao Hu, Peng Wang, Jifeng Hu, Qiyang Zhou+4

Shanghai Jiao Tong University, China · Shenzhen Campus of Sun Yat-sen University, China · Jilin University, China+1

Reinforcement Learning Transformer-Based RL Constrained RL Offline RL

Open paper AskPDF arXiv

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 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.

Abstract

Safe RL

Explore similar work

Q-learning Penalized Transformer for Safe Offline Reinforcement Learning

Beyond Hard Constraints: Budget-Conditioned Reachability For Safe Offline Reinforcement Learning

CDCP: Conditional Diffusion Model with Contextual Prompts for Multi-task Offline Safe Reinforcement Learning