Transformer policy improves safety and rewards in offline reinforcement learning
Q-learning Penalized Transformer for Safe Offline Reinforcement Learning
Machine Learning
Summary
Safe offline reinforcement learning means teaching a computer to make good decisions by learning from past information without taking unsafe actions. The authors address the tricky challenge of balancing safety, following past behaviors, and getting high rewards. They introduce a new method that uses a Transformer model guided by learned value functions that penalize unsafe choices, helping the model pick better and safer actions both while learning and when used later. Their approach performs better than previous methods on many tasks and can quickly adjust to new safety limits without extra training.
What this means in practice
- •For autonomous vehicle engineers: Train driving policies from offline data that strictly avoid unsafe actions while maximizing performance across varied driving conditions.
- •For robotics software developers: Build control policies for robots that safely operate using previously collected data, adapting instantly to new safety requirements without retraining.
Authors
Shengchao Hu, Peng Wang, Jifeng Hu, Qiyang Zhou, Anning Hu, Li Shen, Ya Zhang, Dacheng Tao
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.