Adaptable Safe Policy Learning from Multi-task Data with Constraint Prioritized Decision Transformer

Thank you for your insightful reviews. We appreciate your feedback and hope that our response addresses your concerns.

Q1 More analysis on zero-shot generalization

Currently, for each task type, CoPDT is pre-trained using only 2 similar tasks. We acknowledge that this limited diversity may restrict the model’s ability to achieve stronger zero-shot generalization performance. However, this limitation can be mitigated by pre-training on a broader set of similar tasks.

Therefore, we conducted an additional experiment where we trained policies under a larger number of velocity thresholds ($2 \to 6$) in the AntVelocity and HopperVelocity environments. The goal is to examine whether exposing CoPDT to more similar tasks during training can improve its zero-shot generalization ability. Please refer to our response to Q5 of Reviewer YyH8 for detailed experimental settings and results.

It can be observed that when CoPDT is pre-trained on a larger number of similar tasks—similar to the setup in meta-RL—the model exhibits significant improvements in safety performance on both interpolation (V2) and extrapolation (V7) generalization tasks. The improvement is particularly notable in the Hopper environment, which may be attributed to the large threshold gap between the original two Hopper tasks, making pre-training on only two tasks insufficient. These results support the claim that our constraint-prioritized prompt encoder not only captures task-specific memory but also possesses a certain degree of generalization capability to unseen constraints.

Q2 CPRTG generator with high-variance reward distributions

We acknowledge that this reflects a current limitation of the CPRTG generator under the Gaussian modeling assumption. However, we believe that adopting more powerful generative models with strong multi-modal distribution modeling capabilities—such as diffusion models or flow matching—offers a promising solution. Enhancing the CPRTG generator in this way is an important direction for future research.

Q3 Incrementally learn new constraints

We would like to clarify that CoPDT is inherently capable of incrementally adapting to new constraints and environments.

In the multi-constraint setting, both the prompt encoder and the DT policy in CoPDT adopt a three-layer architecture: the bottom and top layers are environment-specific encoders and decoders, designed to handle different state and action space dimensions across environments, as stated in Sec. 3.1. The middle layer, which forms the core of both the prompt encoder and DT policy, is shared across all tasks and environments to enable effective knowledge transfer.

As a result, when encountering new tasks with previously unseen state and action spaces, we can train new encoders and decoders for those tasks while reusing the shared middle layers of the prompt encoder and DT policy. This design naturally supports incremental adaptation without retraining the entire model from scratch.

We have evaluated this capability in an experimental setup described in App. F.7 (Fig. 12), where we perform few-shot transfer to unseen robots whose state and action spaces were not present during training. The results show that CoPDT with full fine-tuning achieves better safety performance compared to learning from scratch, even in these entirely new environments—demonstrating the effectiveness of multi-task training and knowledge sharing.

Q4 What encoder learns?

Actually, our prompt encoder is capable of effectively extracting and representing information from the cost signal.

Since the cost signal in our prompt encoder is not used to compute gradients, but instead serves as a switch to select between sub-networks, we are unable to provide traditional gradient-based interpretation methods. However, we propose an alternative and more intuitive visualization approach: directly comparing the outputs of the same samples (s,a,r,s’) processed through $p_s$ and $p_u$.

To do this, we collected 5,000 transition samples in PointCircle1 and passed them through both $p_s$ and $p_u$, obtaining two sets of 5,000 embedding vectors. While we are unable to include t-SNE visualizations of these vectors due to rebuttal limitations, we compute the cosine similarity between the mean vectors of the two sets, which yields a value of 0.415. This relatively low similarity indicates that the encodings produced by $p_s$ and $p_u$ differ substantially, confirming that the selection mechanism based on cost information has a significant effect on the encoding.

These findings provide empirical support for the effectiveness of our constraint prioritized prompt encoder, showing that the cost signal meaningfully influences the learned task representation, even without direct gradient supervision.

Q5 Integrate CoPDT with existing safety verification frameworks

Indeed, the constraint representations learned in CoPDT can be naturally integrated with existing safety verification frameworks, such as control barrier functions (CBFs).

In the case of neural-network-based CBFs, the constraint representation can be directly treated as part of the state and used as an additional input to the neural network, enabling task-adaptive safety filtering.

For non-parametric CBFs, one can construct a set of prototype task embeddings by generating constraint representations from a small number of trajectories for each training task. During testing, a trajectory from the target task can be used to compute its constraint representation, which is then matched to the most similar prototype. The corresponding CBF of the matched prototype task can then be applied to enforce safety in the target task.

Q6 PAC-style bounds on constraint violation probability？

Since CoPDT is trained using the behavior cloning objective from DT, it does not naturally lend itself to PAC-style theoretical guarantees as commonly found in traditional TD-learning frameworks.

To complement this limitation, we provide empirical evidence in App. A by reporting the performance gap between the learned policy and the behavior policy in the offline dataset. These results highlight the importance of using the CPRTG generator to better align the RTG with the offline data distribution, thereby mitigating distributional shift and improving policy performance.

Q7 Adding a "When NOT to use CoPDT” section

Thank you very much for your valuable suggestion. We will incorporate this section in the next version of the paper, explicitly stating:

"Systems that require formal safety guarantees, continuous cost functions, strict real-time constraints, or high-stakes deployment without extensive testing are not suitable for deploying CoPDT."

This clarification will help better delineate the scope and applicability of our method.

Thank you for your insightful reviews. We appreciate your feedback and hope that our response addresses your concerns.

Q1 Novelty?

CoPDT tackles the unique challenges of combining offline safe RL with multi-task sequence modeling, introducing the following key innovations:

• Unified modeling for multi-task offline safe RL: CoPDT is the first to explicitly learn a unified model adaptable to tasks with varying constraints and budgets. This improves offline data efficiency and few-shot transfer to new constraints in limited-data scenarios.
• Resolving reward-safety conflicts in multi-budget settings: In such settings, optimizing for both high reward and strict safety often leads to conflicts that limit safety performance. To address this, CoPDT introduces the CPRTG token, conditioning RTG generation on CTG to prioritize safety while encouraging maximal feasible rewards. Experiments demonstrate that CoPDT achieves significantly better safety performance than CDT and other baselines under multi-budget conditions, validating the effectiveness of this approach. To our knowledge, no prior work has addressed this problem from the perspective of prioritizing CTG during RTG generation, making this a novel solution strategy.
• Constraint prioritized task representation for multi-constraint adaptation: Inspired by hyper-networks, we propose a constraint prioritized prompt encoder explicitly integrates binary cost signals into encoder structure design for more effective constraint representation. As shown in Fig. 3(b) of Sec. 4.3 and the experiments in App. F.6, this encoder outperforms standard MLP-based encoders, especially when prompt quality is low—demonstrating stronger robustness in constraint representation.

In summary, CoPDT introduces key innovations in CPRTG modeling and structured prompt encoding, offering a targeted solution to multi-task offline safe RL rather than merely combining existing techniques.

Q2 Noisy or ambiguous cost signals

In line with most prior work in offline safe RL, we focus on settings with clear, binary cost signals. However, we acknowledge that dealing with noisy or ambiguous cost signals is an important and practical challenge, and is valuable for future works.

To address this, we assume access to a language-based descriptions for the cost-related rules. Under this setting, we believe a promising approach is to apply the language-guided reward generation methods [1,2] to cost modeling, and integrate them with CoPDT. Specifically, cost-related rule descriptions, along with observation and other relevant information, are input into a LLM, which generates a cost function aligned with the described rules. This function can then be verified and refined via human inspection or a few labeled samples, and used to annotate datasets with noisy or ambiguous cost labels.

We conducted preliminary experiments in this direction, discarding the original (potentially noisy) cost labels and using GPT-4o-mini to generate a new, slightly more conservative cost function ( to mitigate the impact of ambiguity or imprecision in the rule-based descriptions) based on the language description. As shown in the table below, CoPDT trained on LLM-labeled data maintains strong safety performance—even surpassing results obtained using the original ground-truth labels in terms of safety. These findings validate the effectiveness of our proposed approach in scenarios with noisy or even absent cost labels.

In the absence of language-based cost rule descriptions, we advocate exploring robust algorithmic designs and shifting to an online training setup. In offline settings with potentially erroneous cost labels, it is difficult to reliably distinguish trustworthy from corrupted data. In contrast, under an online setting, one can leverage existing safe and robust reinforcement learning frameworks by introducing an adversary that perturbs the cost labels of online samples with a certain probability [3]. This enables learning policies that are both robust to label noise and maintain strong safety—representing a promising direction for future research.

Q3 Sensitivity analysis on $\beta-$decay

Different values of $\beta_{\text{end}}$ do affect reward and cost, but the impact on cost remains within 0.1, indicating minimal influence on policy stability. Instead, it offers CoPDT a useful mechanism for fine-tuning performance during deployment.

We agree that that studying the impact of $\beta$-decay on the risk-reward trade-off is important. To this end, we conduct a sensitivity analysis of $\beta_{\text{end}}$ in App. F.3 (Fig. 7), under the multi-constraint multi-budget setting across all 26 tasks.

As shown in the figure, there is a clear trend: as $\beta_{\text{end}}$ increases from 0.5 to 0.9, both the average reward performance and cost performance of the policy show an increase trend, but the increase of cost remains within 0.1. This indicates that when $\beta_{\text{end}}$ is smaller, the policy emphasizes safety satisfaction more strongly; whereas when $\beta_{\text{end}}$ is larger, the policy places more emphasis on reward maximization.

Since $\beta_{\text{end}}$ is only applied during evaluation, this observation inspired us to explore an Oracle setting, where $\beta_{\text{end}}$ is slightly tuned for each task based on actual test-time performance. The results under this setting are shown in Tab. 5 of App. F.3. We observe that such post-hoc tuning during evaluation can further improve the overall safety performance of the policy.

Q4 Extending CoPDT to online setting

While we focus on the offline safe RL setting in this work, we agree that extending CoPDT to online scenarios is a highly promising direction for future research. In Sec. 4.4 (Fig. 4) and App. F.7 (Fig. 10–12), we have demonstrated that CoPDT can perform efficient adaptation when provided with a small number (e.g., 10) of expert trajectories, indicating its potential for low-data transfer.

However, in fully online settings, directly applying supervised learning-based approaches such as Online Decision Transformer (ODT) [4] may lead to poor learning efficiency, primarily due to the difficulty of exploring trajectories that satisfy safety constraints. This poses unique challenges in safe exploration. To address this, new algorithmic designs will be required—for example, determining how to choose the initial CTG and RTG values during exploration, how to sample them dynamically, and how to update the policy under safety-aware objectives.

We consider this an important and challenging direction, and leave it as a key avenue for future work.

Q5 Zero-shot generalization when learned with more similar tasks

Thanks for the valuable suggestion, pre-training on a larger set of similar tasks can indeed effectively enhance CoPDT’s zero-shot generalization capabilities.

We conducted an additional experiment where we trained policies under a larger number of velocity thresholds (from 2 to 6) in the AntVelocity and HopperVelocity environments. Specifically, in the 2-threshold setting, the thresholds were [2.57,2.62] for Ant and [0.37,0.74] for Hopper. In the 6-threshold setting, we used [2.42,2.47,2.55,2.57,2.62,2.67] for Ant and [0.37,0.42,0.52,0.62,0.67,0.74] for Hopper.

For evaluation, beyond the original zero-shot generalization tasks (AntV2 with threshold 2.52 and HopperV2 with threshold 0.56), we further designed two additional zero-shot generalization tasks: AntV7 (threshold 2.4) and HopperV7 (threshold 0.35). Notably, the V2 tasks correspond to interpolation generalization settings, while the V7 tasks represent extrapolation settings. All results were obtained under a budget of 10, and evaluated using 3 random seeds.

The experimental results are shown in the table below. It can be observed that when CoPDT is pre-trained on a larger number of similar tasks—similar to the setup in meta-RL—the model exhibits significant improvements in safety performance on both interpolation (V2) and extrapolation (V7) generalization tasks. The improvement is particularly notable in the Hopper environment, which may be attributed to the large threshold gap between the original two Hopper tasks, making pre-training on only two tasks insufficient. These results support the claim that our constraint-prioritized prompt encoder not only captures task-specific memory but also possesses a certain degree of generalization capability to unseen constraints.

[1] Yu W, Gileadi N, Fu C, et al. Language to rewards for robotic skill synthesis[J]. arXiv preprint arXiv:2306.08647, 2023.

[2] Xie T, Zhao S, Wu C H, et al. Text2reward: Reward shaping with language models for reinforcement learning[J]. arXiv preprint arXiv:2309.11489, 2023.

[3] Li Z, Hu C, Wang Y, et al. Safe reinforcement learning with dual robustness[J]. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2024.

[4] Zheng Q, Zhang A, Grover A. Online decision transformer[C]//International Conference on Machine Learning. PMLR, 2022: 27042-27059.

Thank you for your careful reviews and constructive suggestions. We appreciate your feedback and hope that our response addresses your concerns.

Q1 Clarification of multi-constraint settings

Thank you for your valuable comments. We would like to clarify that our multi-constraint setting does not refer to handling multiple constraints within a single environment (e.g., PointButton1). Instead, as described in Section 4.1, we train a single DT policy using a unified dataset that spans 26 diverse tasks, including PointButton1 through AntVelocity-v1. This universal policy is then evaluated across all these environments and further tested in few-shot transfer scenarios to previously unseen tasks.

Due to the heterogeneity in state and action spaces among the 26 tasks—for example, the state dimension of PointButton is 76, while that of CarCircle is 40—we employ task-specific encoders and decoders. These components help standardize the input and output dimensions for the shared Transformer core, allowing it to effectively operate across diverse environments.

We sincerely apologize for the earlier lack of clarity in our explanation, and we truly appreciate your feedback. We will revise the manuscript to provide a more precise and detailed description of the multi-task setting to avoid any future confusion.

Q2 Definition of safe and unsafe

We thank the reviewer for the opportunity to clarify our problem formulation. Our work is situated within the soft constraint formulation of safe RL, a prominent and widely-adopted paradigm in the field [1, 2]. To be more precise, in this formulation, a trajectory is considered 'safe' if its total accumulated cost does not exceed a predefined budget. Therefore, the core objective of our model is to learn a policy that generates trajectories compliant with these specified cost constraints.

Q3 Zero-cost expert trajectory

CoPDT, is designed to be robust and does not necessitate unsafe examples in the expert data to function. In the specific scenario where an expert trajectory has zero cost, our method gracefully handles it as follows:

• The entire trajectory is effectively considered a single, continuous 'safe patch'.
• This 'safe patch' is then processed exclusively by the safe encoder $p_s$ to generate its representation.
• Since no cost violations exist, the unsafe encoder $p_u$ is naturally bypassed and does not contribute to the final embedding.

Crucially, our formulation in Eqn. (5) is designed to unify these cases, allowing the final representation to be computed seamlessly. We will clarify this point in the revised version to avoid possible misunderstandings.

Q4 When constraint is not directly observable

Typically, we assume that cost labels are directly available in offline datasets. However, in many real-world scenarios, we often encounter situations where only rule-based descriptions of costs are provided, without explicit cost labels or with only ambiguous or noisy cost labels.

Under this setting, we believe a promising approach is to apply the commonly used language-guided reward generation methods [3,4] to cost modeling, and integrate them with CoPDT. Specifically, cost-related rule descriptions, along with observation and other relevant information, can be input into a LLM. The LLM can then generate a cost function that aligns with the provided rule. After cost function generation, its correctness can be verified and refined through human inspection or a small number of labeled samples. Then, this cost function can subsequently be used to annotate the offline dataset without cost labels. We have conducted preliminary experiments along this line, as discussed in our response to Q2 of Reviewer YyH8.

Q5 Meaning of Eqn. (9)?

Thank you for pointing this out. We would like to clarify that Eqn. (9) is not used to train $q_\phi$—rather, it is trained using Eqn. (8). Eqn. (9) defines the target we aim to sample from $q_\phi$. In this formulation, $G_t$ serves as an expert indicator. Under our setup, among trajectories that satisfy the CTG budget constraint, those with higher returns are considered more "expert”. Therefore, $p(\hat{R}_t \mid G_t, \hat{C}_t, s_t)$ represents the desired sampling distribution over $\hat{R}_t$ conditioned on $\hat{C}_t$ and $s_t$, such that the sampled $\hat{R}_t$ is more likely to correspond to trajectories aligned with the expert indicator $G_t$. Eqn. (9) then applies Bayes’ rule to decompose this distribution into two components: The first term reflects the likelihood under $q\_\phi$, i.e., the probability is higher when the sample is closer to the mean of $q\_\phi$; The second term, $p(G_t \mid \hat{C}_t, \hat{R}_t)$, denotes the likelihood that a trajectory with the given CTG and RTG is considered an expert. When CTG is fixed, trajectories with higher RTG are more likely to satisfy $G_t$ (importantly, $G_t$ is not used to directly label trajectories—it is an interpretable latent variable). The takeaway from Eqn. (9) is that, during sampling from $q\_\phi$, the objective is not only to sample high-probability RTGs but also to prefer relatively larger RTGs that are more likely to satisfy the expert indicator. We will improve the clarity in the next version.

Q6 Comparison with FISOR?

While FISOR is not inherently budget-aware, we include it as a strong baseline due to its SOTA performance in safety. FISOR is designed for hard constraint settings and does not accept a budget as input, we conducted 12 separate runs for each task and grouped them into four categories, using $[10, 20, 40, 80]$ to normalize the cost accordingly. From the safety perspective, this setup gives FISOR an advantage during evaluation, as the normalization scale of cost is relaxed for it. Despite this, CoPDT achieves comparable safety performance to FISOR, demonstrating that CoPDT maintains strong safety properties. We acknowledge that this comparison may not be perfectly aligned in terms of problem assumptions, but we believe it still provides useful insights into safety performance across settings.

Q7 New baselines

Thank you very much for your valuable suggestion. We have additionally provided experimental results comparing CoPDT with OASIS [5] and LSPC [6] in the Single-constraint Multi-budget setting, and the conclusions remain consistent with our previous findings. Due to time constraints, these experiments were conducted on the first 16 tasks. Since the official OASIS code does not include hyperparameter settings for the SafetyGymnasium tasks we tested, we followed a similar practice to that used in their BulletGym experiments and adopted the hand-set (RTG, CTG) pairs provided in the OSRL codebase as generation-condition hyperparameters for OASIS.

The results are shown in the table below. OASIS, which leverages diffusion models for data generation, achieved promising results in the BulletGym environment. However, due to time constraints, we were unable to perform extensive hyperparameter tuning for OASIS in the SafetyGymnasium environments, which may have limited its performance and failed to fully showcase its potential. On the other hand, LSPC—a hard-constraint baseline similar to FISOR—demonstrates strong safety performance and achieves constraint satisfaction in most environments. Nevertheless, CoPDT still outperforms LSPC in terms of overall safety, underscoring the effectiveness of the CPRTG mechanism in balancing reward and safety under the multi-budget setting, in line with our prior conclusions. We will include the corresponding experimental results in the next version.

[1] Garcıa J, Fernández F. A comprehensive survey on safe reinforcement learning[J]. Journal of Machine Learning Research, 2015, 16(1): 1437-1480.

[2] Gu S, Yang L, Du Y, et al. A review of safe reinforcement learning: Methods, theories and applications[J]. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2024.

[3] Yu W, Gileadi N, Fu C, et al. Language to rewards for robotic skill synthesis[J]. arXiv preprint arXiv:2306.08647, 2023.

[4] Xie T, Zhao S, Wu C H, et al. Text2reward: Reward shaping with language models for reinforcement learning[J]. arXiv preprint arXiv:2309.11489, 2023.

[5] Yao Y, Cen Z, Ding W, et al. Oasis: Conditional distribution shaping for offline safe reinforcement learning[J]. Advances in Neural Information Processing Systems, 2024, 37: 78451-78478.

[6] Koirala P, Jiang Z, Sarkar S, et al. Latent Safety-Constrained Policy Approach for Safe Offline Reinforcement Learning[C]//The Thirteenth International Conference on Learning Representations, 2025.

Thank you for your kind and thoughtful reviews. We appreciate your feedback and hope that our response addresses your concerns.

Q1 Choosing $\beta_t$ when the underlying distribution is not Gaussian?

In our current implementation, we model $q_\phi$ as a Gaussian distribution, as it is generally sufficient to handle most practical scenarios. However, we acknowledge that this assumption may not hold in more complex scenarios. In such cases, we believe that leveraging the strong multi-modal distribution modeling capabilities of generative models—such as diffusion models or flow matching—provides a more suitable alternative for modeling $q_\phi$. Under this setup, the choice of $\beta_t$ can remain consistent with the Gaussian case. Specifically, to sample $\hat{R}_t$, we first draw $X$ samples from a standard normal distribution. These samples are then transformed using a diffusion model or flow matching to generate $X$ candidate values. Finally, we select the $\beta_t$-quantile among these candidates as the final $\hat{R}_t$.

Q2 Intuitive explanations about modeling RTG on CTG?

Intuitively: To illustrate the motivation behind our design, consider a simple example. In our framework, the CTG represents the goal of safety satisfaction, while the RTG corresponds to reward maximization. The DT policy acts as a goal-conditioned policy, where both CTG and RTG jointly specify the intended behavior.

Suppose an agent operates along a 1D axis and starts at the origin. The CTG defines a safety region—e.g., $X \in [-10, 10]$—within which the agent must remain. Meanwhile, the reward function incentivizes movement in the positive X direction, so the RTG effectively specifies a target position. If the RTG is set to 20, the agent will attempt to reach $X = 20$, which lies outside the safe region, resulting in a conflict between the safety and reward objectives. In this case, the agent may violate safety constraints in pursuit of the reward.

To avoid such conflicts, our CPRTG mechanism retrieves feasible (CTG, RTG) pairs from the offline dataset. For instance, for a given CTG, the dataset might contain trajecotries with RTG values of 5, 8, 9, and 9.8. Since these trajectories already satisfy the safety constraints, their associated RTGs are guaranteed to be compatible with the CTG. Among these candidates, we select the one with the highest return (e.g., 9.8), thereby maximizing reward while ensuring safety.

Formally: More detailly, modeling RTG using $q_\phi$ conditioned on CTG can be interpreted from two complementary perspectives: (1) Resolving the conflict between RTG and CTG in multi-budget safe RL; (2) Reducing extrapolation error from an offline RL perspective.

For the first point, we have visualized the conflict between RTG (representing reward maximization) and CTG (representing adherence to safety constraints) in Fig. 2(b) and 2(c) of Sec. 4.3. The plots show that, for a DT model, when CTG is held fixed and RTG is increased, both the reward and cost performance of the resulting policy increase accordingly. In contrast, when RTG is fixed and CTG is varied, the performance remains largely unchanged. This asymmetry indicates that the DT model tends to prioritize RTG over CTG, thereby potentially neglecting safety requirements. To address this issue in safe RL—where satisfying safety constraints is often a primary concern—we propose using $q_\phi$ to automatically generate RTG values conditioned on a given CTG. These RTG values are learned from offline data and are generally more conservative than externally specified RTG targets, which helps mitigate the conflict between reward and safety objectives.

From the offline RL perspective, DT learns policies through behavior cloning via supervised learning. Hence, its performance heavily relies on how well the input distribution aligns with that of the offline dataset. If CTG is regarded as an externally specified safety target and RTG is treated as part of the input state, minimizing the extrapolation error requires aligning the policy input distribution with that of the offline data. To achieve this, we model RTG using $q_\phi$ learned from the offline dataset. This interpretation is further elaborated in App. A.

Q3 The notation system does not highlight the multi-task nature

Thank you for the constructive suggestion. We will revise our notation system accordingly in the next version. Specifically, we will explicitly incorporate the task embedding into the formulations in Sec. 3.2 and 3.3 to better highlight the multi-task nature of our approach.