OASIS: Conditional Distribution Shaping for Offline Safe Reinforcement Learning

We gratefully thank the reviewer for recognizing the novelty, comprehensive experiment validation, and theoretical analysis contribution of our work. We provide our response to the comments below:

W2: Assumptions 1 and 2 seem kind of idealized to directly bound the distribution and policy discrepancy, since and cannot be directly calculated or estimated. Admittedly, analyzing the distribution and policy discrepancy based on diffusion-generated data may be difficult and beyond the scope of this work. Maybe more discussion and explanation could help justify these assumptions.

Assumption 1 and 2 bound the function approximation error of diffusion model and inverse dynamics model, which depends on the implementation (e.g., network architecture, learning rate, etc.). As you mentioned, we also believe it may be out of the scope to directly analyze this type of error. Meanwhile, those assumptions bound the expectation of error instead of maximal error and they are also adopted by previous work [1,2]. Therefore, we argue that the assumptions are not overly idealized but necessary to derive theoretical analysis.

W3: Directly excluding mismatched data from training datasets could be an another intuitive approach for SDM problem. So it would be better to discuss the performance if tempting and conservative data were directly excluded from the full datasets, utilizing the remaining datasets (i.e., Full dataset - Hybrid dataset) for training. A comparison between these two approaches (i.e., adding matched data vs. reducing mismatched data) would be valuable.

According to the reviewer's suggestion, we conduct the experiments with these datasets in the Ball-Circle task with threshold 40. The used hybrid, hybrid-R(rewarding only), hybrid-S(safe-only), and hybrid-Com(complementary) datasets are visualized in Figure R-4 presented in the supplementary PDF file. We can observe that when using the safe-only dataset, the learned policy is conservative with low cost and low reward. When using the rewarding-only dataset, the learned policy is tempting with high reward and high cost. When training with the complement of the hybrid dataset, we can get a better safe performance compared to using the hybrid dataset. However, since it still contains a large amount of imperfect demonstration with low reward, the reward performance is not satisfactory compared to our OASIS.

Q1: What is the size of the generated dataset? I would like to confirm whether you directly train the policy on or a mixture of and the original datasets?

The number of transition pairs of the original tempting dataset and the generated datasets used in Table 1 is shown below. Each value presents the number of transition pairs.

For OASIS, we directly train RL agents on the generated datasets.

Q2: How do you calculate the reward condition based on the given cost?

We select the conditions with similar study in Figure 7 to find a good condition that has a high reward while sticking to the cost limit.

Q3: Why are the thresholds in Figure 5 different for different types of datasets? Did you use the same threshold when constructing different types of datasets? This is a concern because one tempting dataset constructed under the thresholds 20 may be considered tempting for one threshold 20 but not for another threshold 60.

We change the thresholds in Figure 5 to show that our method can perform well under varying threshold conditions. We used different thresholds to construct different types of datasets to make them tempting/conservative/hybrid. The visualization of the dataset and corresponding threshold is also available in Appendix C.2.

According to the reviewer's suggestion, we also conducted experiments with the same thresholds on different datasets. The datasets and the experiment results are presented in Figures R-2 and R-3 in the provided PDF file. We can observe that when setting threshold=40, in all tested datasets including tempting, conservative, and hybrid, our method OASIS exhibits the best performance, achieving the highest reward among the safe agents.

Q4: In Figure 6, how did you obtain the performance of baseline methods (e.g., CDT) under a specific $\alpha$? Did you use the data of size generated by OASIS for these baselines? It would be helpful to provide more details about the data efficiency experiments.

We apologize for any confusion. In Figure 6, $\alpha$ represents the size of the RL agent training dataset. For OASIS, it denotes the size of the generated data. A subsequent BCQ-Lag agent is trained on this generated dataset to obtain the safe RL policy. For the baseline methods, we create the training dataset by randomly sampling $\alpha\%$ of trajectories from the original dataset for RL training. In this experiment, we aim to demonstrate that, for offline safe RL, the agent can learn a good policy in a data-efficient manner if the dataset has minimal safe dataset mismatch (SDM) issues. OASIS offers a solution to shape the dataset distribution, which can reduce the required dataset size for RL training while maintaining good performance. We have added a detailed explanation and analysis of this experiment in the revision.

---

[1] Holden Lee, et al. Convergence for score-based generative modeling with polynomial complexity.

[2] Sitan Chen, et al. Sampling is as easy as learning the score: theory for diffusion models with minimal data assumptions.

(continued)

Q6: In Fig. 6, are all models using data generated by the same diffusion model? OASIS handles the diffusion model training and data generation with BCQ-Lag as the actual Offline RL policy learning. This makes me a little confused. How are the curves related in Fig. 6?

We apologize for any confusion. In Figure 6, $\alpha$ represents the size of the RL agent training dataset. For OASIS, it denotes the size of the generated data. A subsequent BCQ-Lag agent is trained on this generated dataset to obtain the safe RL policy. For the baseline methods, we create the training dataset by randomly sampling $\alpha\%$ of trajectories from the original dataset for RL training. In this experiment, we aim to demonstrate that, for offline safe RL, the agent can learn a good policy in a data-efficient manner if the dataset has minimal safe dataset mismatch (SDM) issues. OASIS offers a solution to shape the dataset distribution, which can reduce the required dataset size for RL training while maintaining good performance. We have added a detailed explanation and analysis of this experiment in the revision.

Q7: Why is a hybrid dataset an interesting setting vs. a full dataset?

The hybrid dataset is more realistic in some real-world tasks. For example, in autonomous driving, we define the cost as the distance to the nearest surrounding obstacle and the reward as the time to arrive at the destination. In data collection, some conservative drivers achieve low cost but medium rewards. Some aggressive drivers achieve high reward but high cost. The combination of these two results in a hybrid dataset.

Q8: typos.

We fixed these typos in our revised version and carefully checked the manuscript.

We appreciate the reviewer for the insightful suggestions and the praise of our theoretical results and experimental performance. We provide our response to the reviewer's comments and questions as follows:

W1: Data generation and training can be a slow process.

We agree that generating new data could cause extra computational cost to the training process. However, since we focus on the offline RL setting, this extra cost only occurs before online inference, which we believe is still acceptable.

W2: Selecting hyperparameters such as the number of denoising steps might vary results significantly. For example, Table 2 has greatly varying policy costs for different values of $K$.

For offline safe RL tasks, agents sometimes achieve similar rewards while obtaining different cost values when the cost limit is small and the cost signal is binary. The difference in cost performance in Table 2, where the threshold is set at $\kappa=20$, may come from this. Nevertheless, our method consistently outperforms the baselines in terms of safety constraint satisfaction and reward maximization.

We also tested our method with moderate thresholds ($\kappa=30, 40$) and show the results below. We observe that the cost performance of our method is more consistent under this setting. (All the values are normalized.)

BallCircle ($\kappa=30$):

BallCircle ($\kappa=40$):

BallRun ($\kappa=30$):

BallRun ($\kappa=40$):

W3: While results are mostly consistent, it is hard to say when the proposed method prefers to act more conservatively with lower cost. A study on a hyperparameter change (apart from cost/return targets) to control this balance would be helpful.

In our method, the balance between reward and cost of generated data for the same model is only controlled by the input condition, i.e., the target reward/cost.

To further study the influence of other hyperparameters on the learned generative model, we provide more ablation experiments on the sequence length $L$. We use the Ball-Circle task with a threshold of $\epsilon=30$ and show results on the testing dataset in the following table. We find that our method is not sensitive to $L$ within the tested range.

Q1: How are the initial states decided to generate the data from the diffusion model?

They are randomly sampled from the training dataset trajectories.

Q2: How long is the inference time for the OASIS model? Is that included in the training time?

The generation time in all experiments is less than 1 minute with an A6000 GPU. Taking Ball-Circle as an example, we train the diffusion model of OASIS with a sequence length of $L=32$. During generation, we randomly sample 2000 states from the training dataset. After one-time generation, we obtain a dataset with $2000 \times (32-1) = 62000$ transitions. This process only takes 11 seconds. The generation time is not included in the training time presented in L999, App C.3.

Q3: How do we decide the “target” cost given the cost threshold to get the best performance? Is the study done in Fig. 7 required for each setting or is there a good heuristic to set these targets? How expensive is this study (Fig. 7)? On a related note, how were these targets set for Table 1 on OASIS and the compared baselines?

We first set the “target” cost as the normalized cost threshold, then adjust the conditions according to the results in Figure 7 to find a good condition that has a high reward while adhering to the cost limit. It is a one-time inference for one set of conditions, which only takes about 15 seconds in total. Meanwhile, the performance of each condition is evaluated by the learned reward and cost and does not require online data. For CDT, which requires an additional reward condition, the reward conditions are adopted from the source code released by the authors.

Q4: Is it right to say the primary reason for the success of the proposed method over the diffusion baselines (like FISOR), is a more realistic conditional generative model that yields optimal generated data satisfying our cost constraints?

In the offline RL setting, the distribution of the dataset matters a lot. If the data quality is low, i.e., the dataset is imbalanced and mostly contains unsafe or low-rewarding demonstrations, our OASIS method, which shapes the dataset distribution towards the target distribution, is more effective than baselines (e.g., FISOR) directly modeling the policy by conditional generative model.

Q5: Are the learned labeling models (inverse dynamics, reward, and cost) only used for labeling the data generated from the diffusion model?

Yes, they are only used to label the data generated from the diffusion model.

(continued)

We would like to express our gratitude to the reviewer for the valuable feedback. We are glad to know that the reviewer recognizes the clear logic, sufficient theoritical analysis, and experiments proving the effectiveness. We provide our response to the questions and concerns below.

W1: Some baselines are not introduced in related work, such as CVAE, FISOR, etc. This may cause difficulties in understanding.

We thank the reviewer for pointing this out. We have added the following sentences to the related work section. FISOR identifies the largest feasible region where agents can operate safely, and optimizes for high rewards within this region and minimizes risks outside it with a diffusion policy. CVAE is an extension of the standard VAE that incorporates additional conditional information, such as class labels or attributes, into data generation. This enables more controlled and specific output generation.

W2: Typos.

Thanks for your careful review. we fixed these typos in our revised version.

Q1: Compared with CVAE, OASIS shows that the conditional diffusion model has a better ability to generate according to the condition information, as shown in Figure 7 c, but this example is a bit simple. Can the author show more comparisons of the two similar to Figure 7 c? For example, add the OASIS generation results to Figure 8 c?

We present additional comparisons in Figure R-1 of the supplementary PDF file. This figure illustrates the generation results of OASIS and CVAE under two conditions: medium-cost-high-reward and low-cost-medium-reward. While CVAE can reconstruct the trajectories, it fails to integrate conditions into the generation process. In contrast, OASIS successfully controls the generated trajectories, avoiding the restricted area when conditioned on low-cost-medium-reward.

Q2: When showing the effectiveness of the newly generated dataset, in addition to showing the distribution of the generated state like Figure c, it is also necessary to verify whether the annotations of the inverse dynamics model and reward & cost models are accurate. Can the author supplement the accuracy of these three models?

We provided the scaled MSE loss values of inverse dynamics, reward and cost models in the following tables. We provide the inital value of the loss when the training begins for a comparison. The error of these models is small.

Inverse dynamics

reward model

cost model

Q3: By comparing the results of OASIS and CDT, I came to a conclusion: in offline safe RL, both methods conditioned on cost and reward, generating data is more effective than generating policy. I wonder if the author agrees with this conclusion? This conclusion seems to be uncommon in the field of offline RL.

In the offline RL setting, the distribution of datasets matters a lot. If the data quality is low, i.e. the dataset is imbalanced, and most contain unsafe or low-rewarding demonstrations, our OASIS, which shapes the dataset distribution towards the target distribution, is more effective than directly generating safe policy by conditional sequential modeling such as CDT. The importance of distribution shaping for imbalanced dataset in offline RL has also been discussed in some related works [1, 2], where the authors proposed sampling strategies to solve this issue.

We acknowledge that the performance comparison between general policy generation and data generation in offline safe RL remains an open question. This represents a fascinating area for future research.

---

[1] Hong, Zhang-Wei, et al. "Beyond uniform sampling: Offline reinforcement learning with imbalanced datasets." NeurIPS 2023.

[2] Hong, Zhang-Wei, et al. "Harnessing mixed offline reinforcement learning datasets via trajectory weighting." ICLR 2024.