Adaptable Safe Policy Learning from Multi-task Data with Constraint Prioritized Decision Transformer
Ruiqi Xue, Ziqian Zhang, Lihe Li, Cong Guan, Lei Yuan, Yang Yu
Abstract
Learning safe reinforcement learning (RL) policies from offline multi-task datasets without direct environmental interaction is crucial for efficient and reliable deployment of RL agents. Benefiting from their scalability and strong in-context learning capabilities, recent approaches attempt to utilize Decision Transformer (DT) architectures for offline safe RL, demonstrating promising adaptability across varying safety budgets. However, these methods primarily focus on single-constraint scenarios and struggle with diverse constraint configurations across multiple tasks. Additionally, their reliance on heuristically defined Return-To-Go (RTG) inputs limits flexibility and reduces learning efficiency, particularly in complex multi-task scenarios. To address these limitations, we propose CoPDT, a novel DT-based framework designed to enhance adaptability to diverse constraints (i.e., cost functions) and varying budgets. Specifically, CoPDT introduces a constraint prioritized prompt encoder, which leverages sparse binary cost signals to accurately identify constraints, and a constraint prioritized Return-To-Go (CPRTG) token mechanism, which dynamically generates RTGs based on identified constraints and corresponding safety budgets. Extensive experiments on the OSRL benchmark demonstrate that CoPDT achieves superior efficiency and significantly enhanced safety compliance across diverse multi-task scenarios, surpassing state-of-the-art DT-based methods by satisfying safety constraints in more than twice as many tasks.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Cited by top-tier papers2
- GAS: Enhancing Reward-Cost Balance of Generative Model-assisted Offline Safe RLZifan LIU, Xinran Li, Shibo Chen, Jun ZhangICLR 2026
- Structured Expert Routing with Multi-View Task Priors for Offline Meta-Reinforcement LearningYisen Zhao, Peixi Peng, Xinyu Hu, Cong Li et al.ICML 2026
Builds on25
- LoRA: Low-Rank Adaptation of Large Language ModelsEdward J. Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu et al.ICLR 2022 · 18,833 citations
- Conservative Q-Learning for Offline Reinforcement LearningAviral Kumar, Aurick Zhou, George Tucker, Sergey LevineNeurIPS 2020 · 2,881 citations
- Decision Transformer: Reinforcement Learning via Sequence ModelingLili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee et al.NeurIPS 2021 · 2,557 citations
- Offline Reinforcement Learning with Implicit Q-LearningIlya Kostrikov, Ashvin Nair, Sergey LevineICLR 2022 · 1,402 citations
- MOPO: Model-based Offline Policy OptimizationTianhe Yu, Garrett Thomas, Lantao Yu, Stefano Ermon et al.NeurIPS 2020 · 989 citations
Related papers
- Constrained Decision Transformer for Offline Safe Reinforcement LearningZuxin Liu, Zijian Guo, Yihang Yao, Zhepeng Cen et al.ICML 2023 · 82 citations
- Temporal Logic Specification-Conditioned Decision Transformer for Offline Safe Reinforcement LearningZijian Guo, Weichao Zhou, Wenchao LiICML 2024 · 7 citations
- Constraint-Adaptive Policy Switching for Offline Safe Reinforcement LearningYassine Chemingui, Aryan Deshwal, Honghao Wei, Alan Fern et al.AAAI 2025 · 12 citations
- Decomposed Prompt Decision Transformer for Efficient Unseen Task GeneralizationHongling Zheng, Li Shen, Yong Luo, Tongliang Liu et al.NeurIPS 2024 · 13 citations
- Q-learning Decision Transformer: Leveraging Dynamic Programming for Conditional Sequence Modelling in Offline RLTaku Yamagata, Ahmed Khalil, Raúl Santos-RodríguezICML 2023 · 121 citations
