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
r↑
c↓
r↑
c↓
r↑
c↓
r↑
c↓
r↑
c↓
r↑
c↓
r↑
c↓
r↑
c↓
r↑
c↓
r↑
c↓
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↑∈[0,100] ) and maximum normalized cost among all constraints ( c↓ ), averaged over 5 seeds and 20 episodes. Black bold indicates constraint satisfaction ( c≤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 θ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≤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≤1 ), and underlined bold indicates the highest reward among feasible methods.
Task
CDT-MC
CORMA w/o PoE
CORDT
CORMA (Ours)
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≤1 ), and underlined bold indicates the highest reward among feasible methods.
Task
CDT-MC
CORMA w/o PoE
CORDT
CORMA
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≤1 ), and underlined bold indicates the highest reward among feasible methods.
K
Transformer
Mamba
r↑
c↓
r↑
c↓
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≤1 ), and underlined bold indicates the highest reward among feasible settings.
Method
DC-v3
PB-v3
CG-v3
CP-v3
Average
r↑
c↓
r↑
c↓
r↑
c↓
r↑
c↓
r↑
c↓
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.
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.
Shengchao Hu, Peng Wang, Jifeng Hu +5
Shanghai Jiao Tong University, China · Shenzhen Campus of Sun Yat-sen University, China · Jilin University, China +1
Sequential decision making using Markov Decision Process underpins many realworld applications. Both model-based and model free methods have achieved strong results in these settings. However, real-world tasks must balance reward maximization with safety constraints, often conflicting objectives, that can lead to unstable min/max, adversarial optimization. A promising alternative is safety reachability analysis, which precomputes a forward-invariant safe state, action set, ensuring that an agent starting inside this set remains safe indefinitely. Yet, most reachability based methods address only hard safety constraints, and little work extends reachability to cumulative cost constraints. To address this, first, we define a safetyconditioned reachability set that decouples reward maximization from cumulative safety cost constraints. Second, we show how this set enforces safety constraints without unstable min/max or Lagrangian optimization, yielding a novel offline safe RL algorithm that learns a safe policy from a fixed dataset without environment interaction. Finally, experiments on standard offline safe RL benchmarks, and a real world maritime navigation task demonstrate that our method matches or outperforms state of the art baselines while maintaining safety.
Janaka Chathuranga Brahmanage, Akshat Kumar
School of Computing and Information Systems, Singapore Management University.
Multi-task offline safe reinforcement learning (RL) promises to learn a shared optimal safe policy from offline data across multiple tasks. This paradigm provides an effective means for the widespread application of RL in multi-task scenarios with high risk and interaction costs. However, the triple challenges of multi-tasking, safety constraints, and out-of-distribution (OOD) actions pose a significant hurdle for existing methods to ensure safety while maximizing reward returns. In this work, we propose a Conditional Diffusion model with Contextual Prompts (CDCP) to address these challenges. Concretely, we first rethink the requirements and challenges in current multi-task decision-making and control scenarios and establish the objectives of multi-task offline safe RL. Subsequently, we transform the multi-task constrained optimization problem into a conditional generation problem using the diffusion model. Based on this, we design a classifier-free guided cost-constraint strategy to provide flexible cost constraints and eliminate extrapolation errors from OOD actions via supervised learning. Additionally, we introduce a novel contextual prompting method to enhance multi-task representation accuracy and adaptability to unseen tasks. A gradient loss synchronization strategy is also introduced to eliminate gradient interference, improving training stability. Finally, extensive experiments demonstrate that the CDCP algorithm exhibits higher performance and safety in multi-task scenarios than the current state-of-the-art baseline methods. It meets different cost constraints without further training, providing a more flexible cost-constraint solution for the multi-task safe RL.
Jiayi Guan, Tianle Zhang, Li Shen +8
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