Paper: SafeDreamer: Safe Reinforcement Learning with World Models (ICLR-2024)

Reviewer comment (Presentation):
A key problem is there is little discussion or elaboration about the choices made in this paper and what the alternatives could be. Unfortunately the related work section is relegated to the appendix, while this does not pose an immediate issue I don't find this section to be particularly comprehensive. Yes, there are references to older approaches from the literature, both model-free and model-based. However, there is no discussion on where this research sits in the broader context of safe RL. Not all safe RL research uses the CMDP framework with simple cumulative constraints. It would be nice for the authors to provide a broader overview of safe RL and the problems associated with deploying RL in the real world, and perhaps discuss and motivate why they only consider simple cumulative constraints and even critique this setup and discuss its limitations.

Author response:
**Re:** Thank you very much for your valuable suggestions. In response to your comments, we have made the following adjustments:
1. We have **moved the related work to the main paper**.
2. We **added some other relevant papers in related work**, including algorithms outside the CMDP framework, such as those addressing constrained problems from the perspective of energy functions, and so on.
3. We have **included additional discussions about SafeRL and its deployment in real-world environments**. Similar to your view, we believe that SafeRL can facilitate the deployment of RL in the real world, and we have further incorporated this discussion into Section 6: Conclusion.

Reviewer comment (Novelty):
Originality. All three algorithms are essentially combinations of existing methods, DreamerV3, PID Lagrangian, CCEM and the Augmented Lagrangian. This is not an issue per se, since this is common for RL research. However, very little discussion is provided for why these methods have been chosen and what the possible alternatives are from the literature and if these may obtain better results or not.

Author response:
**Re:** In the context of model-based RL, there are generally two ways of utilizing a world model. 
1. One is to use the world model for online planning.
2. The other is to use trajectories rolled out by the world model for offline policy optimization, which in this work is referred to as 'background planning'. 

In this work, our main argument is that the world model can facilitate the achievement of zero cost in safe RL, whether through online planning or background planning, as demonstrated by our algorithms OSRP and BSRP-Lag. BSRP-Lag employs an augmented Lagrangian for policy optimization to optimize the action distribution generated by the actor. Conversely, OSRP uses CCEM in conjunction with the world model for online optimization of the action distribution. However, due to the constraints of computational resources and errors in the cost model, online planning is limited in its ability to balance long-term safety and rewards. To address this, we introduced Lagrangian and critics into online planning, resulting in OSRP-Lag. 

Ultimately, the three world model-based algorithms we propose achieve nearly zero cost with both vector and visual inputs, a feat not accomplished by previous work. 

**We argue that the choice of Lagrangian method is not the key to achieving zero cost.** The crucial factor is **how to accurately model safety-relevant states using the world model**. As the environment becomes more complex, the error in each component of the world model also increases. Therefore, it's necessary to adaptively modify to reconcile variances in signal magnitudes and ubiquitous instabilities within all model elements. Additionally, we have incorporated an **ablation study** in Appendix C.5 to examine the impact of various components of the SafeDreamer on its performance. **We believe our work will serve as a reference for the SafeRL community in achieving zero cost.**

Reviewer comment (Evaluation):
Discussion of results. The discussion of results is limited. It is clear all methods achieve similar performance in terms of costs. However, for rewards the three algorithms can perform quite differently across the environments. It would be nice if the authors could address this, and perhaps explain or at least provide some insight into why one algorithm may be working better than another (w.r.t. reward) on a specific task.

Author response:
**Re:** Thank you very much for your valuable suggestions. Here is our analysis of each task according the experimental results (see Figure 4 and Appendix C):
1. **In PointGoal1**, we observed that the cost for OSRP was twice as high as that for OSRP-Lag, which might be attributed to errors in the cost model. OSRP-Lag resulted in a more conservative policy, suggesting that introducing a cost critic through the Lagrangian method can mitigate the errors brought by estimating cost return solely based on the cost model.
2. **In RacecarGoal1**, due to the larger size of the Racecar, it is more prone to touching obstacles. Therefore, OSRP-Lag tends to be more conservative in this environment compared to the other two algorithms.
3. **In PointButton1**, where the agent needs to avoid dynamically moving obstacles, we found that OSRP had higher rewards than BSRP and was close to Unsafe DreamerV3. This is because OSRP can utilize the world model for online planning, allowing real-time prediction of dynamic changes in the environment. BSRP, relying solely on the actor network for action output without online planning, struggles to predict real-time dynamic changes in obstacles.
4. **In PointPush1**, the optimal policy involves the agent learning to wedge its head in the middle gap of the box to push it quickly. We noticed that BSRP could converge to this optimal policy, whereas OSRP and OSRP-Lag had difficulty discovering it during online planning. This may be due to the limited length and number of planning trajectories, making it hard to search for dexterous strategies within a finite time. BSRP's significant advantage in this environment suggests it might be better suited for static environments that require more precise operations.
5. **In PointGoal2**, with more obstacles in the environment, OSRP-Lag, which introduces a cost critic to estimate cost return, tends to be more conservative than OSRP. This leads to a slower increase in reward for OSRP-Lag. Although OSRP's reward increases as quickly as BSRP's, its cost is higher compared to the other two algorithms, indicating that relying solely on the cost model for planning can result in errors.

Reviewer comment (Writing):
Some references are wrong in section 5.2 and these need fixing.

Author response:
**Re:** Thank you very much for your careful reading. We have updated the correct citations for MBPPO-Lag and DreamerV3 in Section 5.2 and have highlighted them in orange.

Reviewer comment (Experiments):
While the paper compares SafeDreamer with several existing SafeRL methods, it lacks a comprehensive comparison with a broader range of related work in the field. A more extensive comparison could provide a clearer context for SafeDreamer's contributions and limitations. For example, the authors noticed the failure of achieving zero cost of Lagrangian methods and CPO and concluded the reason behind it is the lack of long-horizon planning. Actually there are several literatures [1, 2] that have discussed this issue from the perspective of cost function design. In general, more informative cost functions could solve this issue of not achieving zero costs of Lagrangian methods and CPO.

[1] Ma, H., Liu, C., Li, S. E., Zheng, S., & Chen, J. (2022, May). Joint synthesis of safety certificate and safe control policy using constrained reinforcement learning. In Learning for Dynamics and Control Conference (pp. 97-109). PMLR.

[2] He, T., Zhao, W., & Liu, C. (2023). Autocost: Evolving intrinsic cost for zero-violation reinforcement learning. arXiv preprint arXiv:2301.10339.

Author response:
Thank you very much for your valuable suggestions. Following your advice, **we have expanded the Related Work section, and moved the Related Work from Appendix A to the main paper**. We have carefully readed the two works [1][2] and found that their focus on **Cost Function design aligns with our view** that existing methods such as CPO and Lagrangian approaches cannot achieve zero cost. We have now included these works in the Related Work section and have discussed them there.

**[1]** Ma, H., Liu, C., Li, S. E., Zheng, S., & Chen, J. (2022, May). Joint synthesis of safety certificate and safe control policy using constrained reinforcement learning. In Learning for Dynamics and Control Conference (pp. 97-109). PMLR.

**[2]** He, T., Zhao, W., & Liu, C. (2023). Autocost: Evolving intrinsic cost for zero-violation reinforcement learning. arXiv preprint arXiv:2301.10339.

Reviewer comment (Evaluation):
Why choose different cost thresholds for different environments in Figure 6? I think the threshold of value 0 makes the most sense to me. Also, I noticed that OSRP and OSRP-Lag seem to be too conservative when the cost threshold is around 14 in SafetyPointButton1, which might reveal the imbalance between performance and safety satisfaction.

Author response:
Firstly, we apologize for the misunderstanding and thank you for your careful reading. **We did not choose different cost thresholds for different environments.** In Figure 6, the cost thresholds for all environments are set to 2.0. Regarding your observation that **'OSRP and OSRP-Lag seem to be too conservative when the cost threshold is around 14 in SafetyPointButton1'**, we believe you might have mistaken the dashed line of the baseline algorithm MPC:sim for the cost threshold.

The reason we did not set the cost threshold to zero is **due to potential errors in the trained cost model and cost value model**. Therefore, even in a completely safe state, the estimated cost return by these models might not be zero but a very small positive value. This implies that any trajectory in the planning process may not meet the constraints based on these estimations, making online planning overly conservative and unable to converge.

Reviewer comment (Experiments):
Also, If you were to deploy safeDreamer in the real world, what control frequency can be achieved at maximum? The simulation of series of neural networks rollout might be feasible in simulation (the simulation could pause to wait for the next action), but might be infeasible in real-world practice.

Author response:
In our experiments, we utilized an NVIDIA GeForce RTX 3090 Ti GPU for inferencing the SafeDreamer model, with an average inference time of approximately 0.05 seconds. Theoretically, this suggests that we could achieve a control frequency of 20Hz. However, the actual control frequency is influenced not only by the GPU's inference speed but also by other factors in the real-world deployment environment, such as network latency, sensor data processing time, etc.

Reviewer comment (Novelty):
SafeDreamer (OSRP, OSRP-Lag, and BSRP-Lag) seems like a straightforward combination of prior works (Dreamerv3, CCEM, Augmented Lagrangian, and PID Lagrangian). For example, if I understand correctly, OSRP (Sec 3.2) is exactly CCEM but using the TD($\lambda$) objectives from DreamerV3 (for the cost and rewards estimates). It is not clear what are nuances that make this combination not as straightforward as one would expect. No theory nor empirical analysis is given to gain better insights into the behavior of SafeDreamer when all these components are combined.

Author response:
Re: Recently, many works have been integrating traditional MPC (Model Predictive Control)-Based methods with TD (Temporal Difference) learning. Examples like TDMPC[1] and LOOP[2] represent a combination of MPC with TD learning, known as MPC with a terminal value function.
**However, directly transferring these methods to the SafeRL domain presents challenges.**
1. Due to the need to meet safety constraints, it's crucial for the world model to accurately predict changes in safety-related states, such as the relative position of obstacles to the agent. The prediction of obstacle positions can be particularly problematic; a single incorrect prediction can lead to a cascade of errors in subsequent position predictions. Therefore, using methods like CCEM to calculate the total cost of a trajectory might only yield accurate cost estimates for the initial steps. In approaches like TDMPC or LOOP, which use a value model to estimate future costs at the terminal step of planning, the accumulated error in the final step's state can be significant, leading to large inaccuracies in cost value estimation. Instead, we use TD-lambda, where each step's cost estimation depends not only on the cost model but also on the predicted cost value.
2. Traditional CCEM requires a ground truth simulator to perform safe planning on ground truth states. However, in visual tasks, accurately predicting every pixel's change is challenging and resource-intensive. Our method involves planning on the latent hidden state, which avoids the need for precise pixel predictions. Experimental evidence shows that this approach can achieve nearly zero-cost performance, indicating that the latent hidden state implicitly captures almost all safety-related visual information. We believe that conducting safe planning on the latent hidden state could offer new insight for the SafeRL community.
3. We have added an ablation study of SafeDreamer's design and configuration choices in Appendix C.5.


**[1]** Hansen, N., Wang, X., & Su, H. (2022). Temporal difference learning for model predictive control.

**[2]** Sikchi, H., Zhou, W., & Held, D. (2022, January). Learning off-policy with online planning. In Conference on Robot Learning (pp. 1622-1633). PMLR.

Reviewer comment (Novelty):
While the shown empirical evaluations of SafeDreamer reporting great reward and cost performances in SafetyGym tasks are decent contributions by themselves, I think they are not sufficiently significant for publication. Some theoretical/empirical analysis of SafeDreamer and evaluations in other domains (other than SafetyGym) would improve the contributions.

Author response:
Re: Thank you for your valuable suggestions. To further validate our method, we conducted additional experiments beyond Safety-Gym, incorporating tests in the Safe Drive scenario, focused on safe autonomous driving tasks. Specifically, we conducted supplementary tests of the SafeDreamer in three tasks on MetaDrive, Car-Racing, and FormulaOne. The results demonstrate that SafeDreamer significantly outperforms the baseline algorithms. We have updated the descriptions of these environments and the corresponding experimental data in the revised version of our manuscript. For more details, please refer to Appendix C.1 - C.3. Some video demos can be seen in https://sites.google.com/view/safedreamer#:~:text=SafeDreamer-,MetaDrive,-SafeDreamer

Reviewer comment (Evaluation):
The paper doesn't state how many runs the results are averaged over.

Author response:
Re: Thank you very much for your detailed reading. All our experiments were conducted over 3 runs. We have added a description of this in Section 6.3, highlighted in orange. Additionally, we have uploaded all the Meta Data from the experiments in this work for the community to observe and use. The data can be found at https://sites.google.com/view/safedreamer#:~:text=Other-,Resources,-Data

Reviewer comment (Experiments):
The paper doesn't state all the main hyperparameters used for the baselines (e.g all model architectures and planning hyperparameters for LAMBDA/Safe SLAC). It is unclear if the model-based baselines were given approximately the same number of planning samples as SafeDreamer, and if the model-free baselines were given enough additional model parameters and gradient steps to compensate for the model size and runtime of SafeDreamer. In general, it is unclear if the hyperparameters for each experiment were chosen to be as fair as possible to the baselines.

Author response:
Re: Thank you for your valuable feedback. During our experiments, we placed a high emphasis on the reproducibility and fairness. 
In response to concerns about reproducibility,
1. **We have adopted your suggestion and included detailed information about the baseline hyperparameters in Appendix B (see Tables 2, 3, 4, 5, 6) of revised paper.** This includes model and planning architectures, as well as hyperparameters and gradient steps for model-free algorithms. 
2. Regarding fairness, when comparing SafeDreamer with other model-based methods, we ensured consistency in model parameter size and update frequency to guarantee fairness and comparability in all aspects other than the algorithms themselves. 
3. In comparisons with model-free methods, considering the significant differences between these two types of methods, we indeed increased the model parameters and gradient steps to pursue fairness. However, this approach does not guarantee a completely fair comparison between model-based and model-free methods due to the significant differences in their settings. Nevertheless, our comparison with model-free methods primarily aims to demonstrate the superiority of the world model. Such comparisons have also been reflected in the literature [1] and [2].

**[1]** As, Y., Usmanova, I., Curi, S., & Krause, A. (2022). Constrained policy optimization via bayesian world models.

**[2]** Jayant, A. K., & Bhatnagar, S. (2022). Model-based safe deep reinforcement learning via a constrained proximal policy optimization algorithm. Advances in Neural Information Processing Systems, 35, 24432-24445.

Reviewer comment (Reproducibility):
The lack of information on number of runs and hyperparameters makes it hard to evaluate the significance of the empirical results.

Author response:
Re: We highly appreciate your valuable feedback on the experimental setup, including the addition of a hyperparameter table for the baseline, detailed experimental details, and more. We have updated these in the revised version of our manuscript. 

**The fairness and comparability of our experiments are of utmost importance to us.**

Reviewer comment (Evaluation):
The performance of the model-free baselines (e.g in Fig 1) is not consistent with their performance in prior works (e.g in Ray 2019). For example, in almost all of the shown plots, the reward returns of TRPO-Lag and PPO-Lag behave weirdly (they increase early in training and then quickly decrease seemingly converging towards zero). One may think that it is because of the vision-based observations, but Fig 7 shows equally poor performance for the baselines. This puts into question the validity of all of the empirical results.

Author response:
Re: **There seems to be a misunderstanding regarding the Ray 2019 [1] paper.** It proposed a benchmark environment for SafeRL and provided five basic algorithms, such as CPO, PPO, PPO-Lagrangian, TRPO, and TRPO-Lagrangian. All experiments in that paper were conducted under the constraint of Cost Limit = 25.0, but our Fig 1 was conducted under a setting of Cost Limit = 2.0.
Here is the description from the original Safety Gym paper: **Hyperparameters: In all constrained cases, we set d = 25 for the expected cost limit.**
We attribute this to the fact that model-free algorithms have a low data utilization rate. In situations with a low cost-limit, this leads to the Lagrangian multiplier continuously penalizing the policy, resulting in overly conservative strategies – that is, situations with low costs and also low rewards. This is also one of our motivations for using a World Model for planning.

**[1]** Ray, A., Achiam, J., & Amodei, D. (2019). Benchmarking safe exploration in deep reinforcement learning.

Reviewer comment (Experiments):
It is not clear how exactly the baselines from Ray 2019 were modified for the vision-based partial observations (Ray 2019 uses state-vector observations). According to the caption of Fig 4, it seems like the *exact* experimental procedure of Ray 2019 was used. This possibly implies that no CNN (to handle the image observations) and RNN (to handle the partial observations) were used for the model-free baselines. Meanwhile, SafeDreamer uses recurrence in its model and possibly also uses a CNN (whether it uses a CNN or not was not stated in the paper, so this is only guesswork). This is a severe concern that could actually explain why almost all of the model-free baselines have extremely poor performance inconsistent with prior works.

Author response:
We did not add a vision module (such as CNN or RNN) to the Model Free algorithms. A major advantage of SafeDreamer is its ability to support both vector and visual inputs. In our comparisons with Model Free algorithms, we only used vector inputs. We tried adding specific vision modules to Model Free algorithms, but the performance was not satisfactory. Extending existing Model Free Safe RL algorithms to support visual input scenarios presents significant challenges. Additionally, we have added detailed descriptions of the hyperparameters and network structures for both Model Free and Model Based algorithms. For more details, please refer to Appendix B (see Tables 2, 3, 4, 5, 6).
Regarding the statement, **'This is a severe concern that could actually explain why almost all of the model-free baselines have extremely poor performance inconsistent with prior works,'** prior works [1][2][3][4] used a setting of **cost limit = 25.0**. Under this more relaxed safety constraint, they were able to achieve more satisfactory performance. However, **our paper sets the cost limit = 2.0, which is very close to zero cost,** posing a significant challenge to existing algorithms. In situations with a low cost limit, this leads to the Lagrangian multiplier continuously penalizing the policy, resulting in overly conservative strategies – that is, situations with low costs and also low rewards, which is the reason for the Poor Performance noted in our paper.

Once again, we reiterate: The experiments in our paper are based on factual experimentation and are fair and comparable. We strictly respect the scientific norms of research paper submission!

[1] Ray, A., Achiam, J., & Amodei, D. (2019). Benchmarking safe exploration in deep reinforcement learning.

[2] Zhang, Y., Vuong, Q., & Ross, K. (2020). First order constrained optimization in policy space. Advances in Neural Information Processing Systems, 33, 15338-15349.

[3] Yang, L., Ji, J., Dai, J., Zhang, L., Zhou, B., Li, P., ... & Pan, G. (2022). Constrained update projection approach to safe policy optimization. Advances in Neural Information Processing Systems, 35, 9111-9124.

[4] As, Y., Usmanova, I., Curi, S., & Krause, A. (2022). Constrained policy optimization via bayesian world models. Eleventh International Conference on Learning Representations.

Reviewer comment (Experiments):
The model-free baselines are extremely old relative to the model-based baselines. There are a number of relevant recent model-free baselines like Saute RL and ROSARL that also aim to achieve near-zero cost.

Author response:
Re: Your suggestion is valuable, and within the limited time available for our rebuttal, we have added the SauteRL you recommended, including PPO-Saute and TRPO-Saute, with their cost limits set to 2.0. Detailed hyperparameters are provided in Appendix B. However, due to time constraints and limited computing resources, we were unable to add ROSARL as a baseline, which would require certain modifications to ROSARL, including:
1. Changing the Safety Gym setting to reset the episode upon collision with an obstacle, along with modifications to the reward settings. 
The original text states: We modify the environment so that any collision with a hazard results in episode termination with a reward of −1, thereby making the problem much harder.
2. Fixing the position of obstacles so that their location does not update with each episode reset. 
The original text states: The goal’s location is randomly reset when the agent reaches it, while the locations of the obstacles remain unchanged.

We hope you can see our commitment and take every one of your suggestions as valuable input. If you could consider increasing our score, we promise to fully incorporate the two baselines you suggested in the ready version.

Reviewer comment (Presentation):
There is no related works section in the main paper. The authors say "Owing to space constraints, the related work is provided in Appendix A.", but I think moving the related works to the appendix is going a bit too far. The related works section is important to contextualise the contributions of this work in relevant literature properly.

Author response:
Re: Thank you for your valuable suggestion. **We have moved the related work from Appendix A to the main paper.** Additionally, **we have included a detailed discussion of these three papers[1][2][3]**, especially in terms of achieving near-zero cost returns. In addition, we have expanded the range of our discussion in the related work section, including algorithms outside the CMDP framework, such as those addressing constrained problems from the perspective of energy functions.

[1] Sootla et al, Saute RL: Almost Surely Safe Reinforcement Learning Using State Augmentation, 2022

[2] Tasse et al, ROSARL: Reward-Only Safe Reinforcement Learning, 2023

[3] Kwanda et al, Safe Trajectory Sampling in Model-based RL, 2023

Reviewer comment (Novelty):
There are missing relevant related works. Given that the aim of this paper is to achieve near-zero cost returns, the following works are relevant: Saute RL (model free) [1], ROSARL (model free) [2], STS-MBRL (model based) [3].

Author response:
Re: Thank you for your valuable suggestion. **We have moved the related work from Appendix A to the main paper.** Additionally, **we have included a detailed discussion of these three papers[1][2][3]**, especially in terms of achieving near-zero cost returns. In addition, we have expanded the range of our discussion in the related work section, including algorithms outside the CMDP framework, such as those addressing constrained problems from the perspective of energy functions.

[1] Sootla et al, Saute RL: Almost Surely Safe Reinforcement Learning Using State Augmentation, 2022

[2] Tasse et al, ROSARL: Reward-Only Safe Reinforcement Learning, 2023

[3] Kwanda et al, Safe Trajectory Sampling in Model-based RL, 2023

Reviewer comment (Presentation):
Figure 2 is hard to understand where it is currently placed. It could be moved a few pages below to where the methods are being described.

Author response:
Re: In response to your suggestions in the Minor section, we have taken each one carefully and made the corresponding adjustments, summarized as follows:

1. We have removed Figure 1, as you correctly pointed out, the discussion about it was out of place in the introduction.
2. We have moved Figure 2 to Section 4 where the methods are described.
3. We have revised 'world models' to 'world model'.
4. We changed 'return and cost return' to 'reward and cost returns' and 'value and cost value' to 'reward and cost values', as these terms are indeed more common and easier to understand.
5. Regarding the grammar errors you mentioned, we have corrected them and have carefully reviewed the rest of the document for any other grammar errors.

Reviewer comment (Writing):
It may be less confusing to say "reward and cost returns" and "reward and cost values", instead of "return and cost return" and "value and cost value" (e.g. In Fig 3, but this holds throughout the paper).

Author response:
Re: In response to your suggestions in the Minor section, we have taken each one carefully and made the corresponding adjustments, summarized as follows:

1. We have removed Figure 1, as you correctly pointed out, the discussion about it was out of place in the introduction.
2. We have moved Figure 2 to Section 4 where the methods are described.
3. We have revised 'world models' to 'world model'.
4. We changed 'return and cost return' to 'reward and cost returns' and 'value and cost value' to 'reward and cost values', as these terms are indeed more common and easier to understand.
5. Regarding the grammar errors you mentioned, we have corrected them and have carefully reviewed the rest of the document for any other grammar errors.

Reviewer comment (Writing):
If there are not multiple world models being used in SafeDreamer, remove the "s" in world models where unnecessary (e.g Fig 3 caption).

Author response:
Re: In response to your suggestions in the Minor section, we have taken each one carefully and made the corresponding adjustments, summarized as follows:

1. We have removed Figure 1, as you correctly pointed out, the discussion about it was out of place in the introduction.
2. We have moved Figure 2 to Section 4 where the methods are described.
3. We have revised 'world models' to 'world model'.
4. We changed 'return and cost return' to 'reward and cost returns' and 'value and cost value' to 'reward and cost values', as these terms are indeed more common and easier to understand.
5. Regarding the grammar errors you mentioned, we have corrected them and have carefully reviewed the rest of the document for any other grammar errors.

Reviewer comment (Writing):
There are a couple of grammar errors throughout the paper that make it hard to parse. Examples include:
- Page 1 "proffer"
- Fig 3 caption "using policy"
- Page 4 "This notation offers our model components:"
- Page 7 "How does SafeDreamer against safe model-based algorithms?"

Author response:
Re: In response to your suggestions in the Minor section, we have taken each one carefully and made the corresponding adjustments, summarized as follows:

1. We have removed Figure 1, as you correctly pointed out, the discussion about it was out of place in the introduction.
2. We have moved Figure 2 to Section 4 where the methods are described.
3. We have revised 'world models' to 'world model'.
4. We changed 'return and cost return' to 'reward and cost returns' and 'value and cost value' to 'reward and cost values', as these terms are indeed more common and easier to understand.
5. Regarding the grammar errors you mentioned, we have corrected them and have carefully reviewed the rest of the document for any other grammar errors.

Reviewer comment (Evaluation):
"On the other hand, SafeDreamer was still able to converge rapidly to zero cost, even with a cost limit set at 25" (regarding Figure 5). This suggests that the proposed approach could be overly conservative, compared to prior works that attempt to converge to the given cost threshold. This is not necessarily a good or bad thing but requires further investigation.

Author response:
Re: From your comments, we can sense your deep familiarity with SafeRL and Safety Gym[2]. **We want to clarify that we did not intentionally make unfair adjustments or subjectively bias the Model-free baseline.** We opted not to use the baseline provided by Safety Gym[2], Safety-Start-Agent[1], but instead used OmniSafe[3], for the following reasons:

1. Safety-Start-Agent is currently based on TensorFlow 1.X, which posed some compatibility issues on our 4090 servers, although it is still runnable. 
2. Safety-Start-Agent currently includes only three SafeRL algorithms (CPO, PPO-Lag, TRPO-Lag) and PPO and TRPO, whereas OmniSafe encompasses a wider range of algorithms and is based on Pytorch, **making it more user-friendly and easier to replicate**.

**Regarding your viewpoint, we strongly agree.** Yesterday, we reran the three SafeRL algorithms from Safety-Start-Agent, and the results were as you described: “they all still maximize rewards (at different rates) even with a near-zero cost threshold (but with very poor cost performance).” **We believe this phenomenon can be explained.** Considering that the goal of SafeRL is to maximize rewards while meeting constraints, both high rewards without meeting constraints and low rewards with constraints are actually signs of the algorithm’s failure. **We attribute this to the control of the algorithm's hyperparameters.**

Reviewer comment (Experiments):
The concern about apple-to-apple comparisons remains. The model-free baselines should also use CNNs and RNNs, or the model-based ones plus the proposed approach should use vector observations with no RNN. That said, the model-based baselines are at least apple-to-apple with the proposed approach (which is also model-based).

Author response:
Re: Firstly, we would like to **explain the rationale behind our experiment setting**. **In Figure 5, we compared SafeDreamer with a Model-free baseline under the same vector input.** **We also drew inspiration from Figure 2 in the main paper of Safe SLAC[1]**, where Model-Based and Model-free algorithms are compared under different inputs. However, Safe SLAC only represented the final performance of Model-free algorithms with dashed lines. Our approach differs in two aspects:

1. We replaced the dashed lines indicating convergence performance with learning curves. We did this because we believe that the learning process of Model-free algorithms can also provide valuable information to the readers, helping them understand the differences between vision-based and vector-based inputs.
2. Due to an oversight, we omitted the 'sensor' suffix in our baseline descriptions. We believe that including this suffix will clarify the input differences between the two. Consequently, we have made this adjustment in the latest revised version of our manuscript.

> **We did not intend to place Model-based and Model-free algorithms in an unfair comparative environment.**

Moreover, **we agree with your viewpoint and have accordingly adjusted Figure 15**. In the legend, we have added suffixes to the Model-free algorithms, for example, **CPO (sensor)**. **Following the approach of Safe SLAC[1], we have removed the learning curves of the Model-free algorithms, only depicting their final convergence with dashed lines, to maintain consistency with Safe SLAC.**

**[1]** Hogewind, Y., Simao, T. D., Kachman, T., & Jansen, N. (2022). Safe reinforcement learning from pixels using a stochastic latent representation. The Eleventh International Conference on Learning Representations (ICLR 2023).

Reviewer comment (Experiments):
Given what the authors have reported, it seems like there is a very significant performance difference (in the model-free baselines) between the safety-starter-agents codebase and the omnisafe codebase (the latter reportedly performs worse than the former, but the authors use the latter). This is potentially because of a difference in hyperparameters between the two codebases. This realisation was beyond the control of the authors though, since omnisafe is a relatively popular baseline in the community.

Author response:
Re: From your comments, we can sense your deep familiarity with SafeRL and Safety Gym[2]. **We want to clarify that we did not intentionally make unfair adjustments or subjectively bias the Model-free baseline.** We opted not to use the baseline provided by Safety Gym[2], Safety-Start-Agent[1], but instead used OmniSafe[3], for the following reasons:

1. Safety-Start-Agent is currently based on TensorFlow 1.X, which posed some compatibility issues on our 4090 servers, although it is still runnable. 
2. Safety-Start-Agent currently includes only three SafeRL algorithms (CPO, PPO-Lag, TRPO-Lag) and PPO and TRPO, whereas OmniSafe encompasses a wider range of algorithms and is based on Pytorch, **making it more user-friendly and easier to replicate**.

**Regarding your viewpoint, we strongly agree.** Yesterday, we reran the three SafeRL algorithms from Safety-Start-Agent, and the results were as you described: “they all still maximize rewards (at different rates) even with a near-zero cost threshold (but with very poor cost performance).” **We believe this phenomenon can be explained.** Considering that the goal of SafeRL is to maximize rewards while meeting constraints, both high rewards without meeting constraints and low rewards with constraints are actually signs of the algorithm’s failure. **We attribute this to the control of the algorithm's hyperparameters.**

Reviewer comment (Novelty):
The paper proposes OSPR, OSRP-Lag and BSRP-Lag but the differences between them is not well motivated and presented.

Author response:
In the realm of model-based RL, there are generally two ways to utilize a world model
1. one involves using the world model for online planning. This approach allows the use of the world model for inference during testing, thereby offering enhanced adaptability to dynamically changing environments.
3. the other involves using the world model’s rollouts for offline policy optimization, referred to as background planning in our work. This method does not require the use of a world model for online planning during testing, thus reducing computational overhead.

Online planning denotes the real-time optimization of action distributions by the agent through the world model during interaction with the environment. In contrast, background planning refers to training an actor offline using the world model, which is then employed to output action distributions when interacting with the real environment.
In this work, we aim to articulate the viewpoint that the world model can facilitate the implementation of SafeRL at nearly zero cost, whether through online planning or background planning. This is substantiated by the performance of our algorithms: OSRP and BSRP-Lag. 
1. BSRP-Lag employs augmented lagrangian methods for policy optimization to enhance the action distributions produced by the actor. 
2. OSRP uses the constrained cross-entropy method (CCEM) in conjunction with the world model for online optimization of action distributions. 

Given the computational resource constraints that limit the online planning horizon, it's challenging to consider long-term trade-offs between rewards and costs. Therefore, we introduce the Lagrangian method in online planning to balance these long-term costs and rewards, as demonstrated in OSRP-Lag.

In our experiments, we observed that OSRP performs better in complex, dynamically changing environments, while BSRP excels in static environments requiring dexterous action. For instance, in PointButton1, where the agent must avoid dynamically moving obstacles, OSRP achieved higher rewards compared to BSRP, nearing the DreamerV3 (only considering the reward). This is attributable to OSRP's ability to use the world model for online planning, enabling real-time prediction of environmental dynamics. In contrast, BSRP, relying solely on action outputs from the actor network without online planning, struggles to predict the real-time dynamics of obstacles. Furthermore, in PointPush1, the optimal policy involves the agent learning to quickly push the box by wedging its head in the middle gap of the box. BSRP converged to this optimal policy, whereas OSRP and OSRP-Lag found it challenging to discover this approach through online planning, likely due to the limited length and number of planning trajectories, hindering the discovery of dexterous strategies within a finite time. The significant advantage of BSRP in this scenario suggests that it may be better suited for static environments requiring dexterous operations.

Ultimately, the three world model-based algorithms achieve nearly zero-cost with both vector and visual inputs, a feat not previously accomplished in prior works.

Reviewer comment (Presentation):
In Section 3.2, it seems the sampled trajectories contain two parts, one set of the trajs are deduced from a Normal action distribution and another set is from current policy. This part is confusing and not well presented. The notation is not aligned between SafeDreamer paragraph and the Algorithm 1. I would recommend rephrase the method part and mark the Algorithm line number in the text.

Author response:
Thank you very much for the valuable suggestion. Indeed, the action trajectory is generated by both the **normal action distribution** and **the actor network**. We have employed techniques similar to LOOP [1] and TDMPC [2]. By integrating a small number of actions generated by the actor network into the action trajectory produced by the Gaussian distribution, we can guide the planning toward a more rapid convergence. We have rephrased the method section in the main paper and marked the algorithm line numbers for clarity.

**[1]** Hansen, N., Wang, X., & Su, H. (2022). Temporal difference learning for model predictive control.

**[2]** Sikchi, H., Zhou, W., & Held, D. (2022, January). Learning off-policy with online planning. In Conference on Robot Learning (pp. 1622-1633). PMLR.

Reviewer comment (Presentation):
Mixing up 3 algorithms in 1 Algorithm diagram (like Algo 1 and Algo 2) is a terrible writing practice. I would highly recommend revising this. It's not increasing the density of information but instead increasing the confusion.

Author response:
Thank you very much for your valuable suggestions. Following your advice, we have made corresponding adjustments to our paper. We extracted the common operations of the three algorithms in Algorithm 1 and Algorithm 2, describing them in the main paper using the line numbers of the algorithm. This approach enhances the readability of the algorithm. For more detailed about these modifications, please refer to Section 4.2 and Section 4.3.

Reviewer comment (Writing):
IIUC the SafeDreamer part Sec 3.2 is doing a trajectory-level cost / return estimation without the interaction with the transition model (either a real RL environment or a latent world model). And in OSRP-Lag, we allow "online planning"? What's the difference between OSPR-Lag and OSPR? What is the online means here?

Author response:
**Online** refers to the process of real-time interaction between an agent and its environment. **Online planning** is when the agent uses the world model in real-time during this interaction to **'dream' the outcomes of executing different action trajectories, selecting the optimal action**, and immediately implementing it in the environment. Evaluating trajectory-level cost or return in online planning requires the use of the world model's reward and cost models, as well as its transition model. As such, online planning in the description of Algorithm is an interaction with the world model, where only the first action of the returned action trajectory is executed in the real environment.
**The key difference** between OSRP-Lag and OSRP lies in their planning method. 
1. OSRP uses reward return as the key to sort these trajectories, when OSRP identifies a substantial number of trajectories that satisfy safety constraints in the online planning within the world model.
2. OSRP-Lag sorts action trajectories using 'reward return - $\lambda$ cost return' as the key, where $\lambda$ is the Lagrangian multiplier, when OSRP-Lag identifies a substantial number of trajectories that satisfy safety constraints in the online planning within the world model. This indicates that OSRP-Lag optimizes a conservative policy under relatively safe conditions, taking into account long-term risks, while OSRP focuses exclusively on the reward.

Reviewer comment (Reproducibility):
In test-time, does the BSRP-Lag method still run planning for each real step? Or the world model is only used during training?

Author response:
In testing, BSRP-Lag does not execute planning for each real step. At every step, it simply requires the actor to output actions. The world model is utilized solely during the training process.

Reviewer comment (Experiments):
Have you experimented on the setting where the training environment is different from the test environment? I suspect the generalizability of the world model will limit the performance in the unseen test environments. And if so to what expend we can believe the BSRP-Lag agent? Will it be more conservative in unseen environment? Or vice versa?

Author response:
The advantage of the World Model lies in its ability to accurately model the current environment. However, when faced with an entirely new environment, the World Model inevitably struggles to predict future states accurately. This is a limitation inherent at the algorithmic level. We conducted tests under conditions where the distribution of obstacles in training and testing environments differed:

1. **Experiment 1 - PointGoal1 to PointGoal2**: We deployed the policy trained in PointGoal1 directly in PointGoal2. PointGoal2 contained more vases, and if the agent collided with a vase, especially a moving one, it incurred a cost. In contrast, PointGoal1 had only one vase, and colliding with it did not result in a cost; costs were only incurred upon hitting hazards. We observed that after deploying the policy in PointGoal2, the agent avoided hazards, a skill it had acquired in PointGoal1. This demonstrated a certain robustness of the trained policy, even when the test environment changed (with more vases). However, the agent failed to reduce collisions with vases, as it was only trained to avoid hazards, not vases.
2. **Experiment 2 - PointGoal2 to PointGoal1**: We deployed the policy trained in PointGoal2 directly in PointGoal1. The policy performed remarkably robustly, and surprisingly, it was even lower costs than the policy directly trained in PointGoal1. We speculate that this may be because PointGoal2 was a more complex and challenging task, so a policy trained for it might perform better when deployed in simpler tasks.

In summary, these experiments highlight the robustness of World Model in varying environments, despite its limitations in predicting states in unseen settings. The training in more complex environments seems to enhance the policy's performance in simpler ones, indicating a beneficial transfer of learning.


- PointGoal1 to PointGoal2:

Training Env| Testing Env|Cost of Hazards Contact | Cost of Vases Contact |Cost of Vases Velocity | Cost | Reward
| -------- | -------- | -------- | -------- | -------- | -------- | -------- |
|PointGoal1|PointGoal2|   1.4    | 16.8     | 94.1  | 95.8 | 17.89 


- PointGoal2 to PointGoal1:

Training Env| Testing Env| Cost of Hazards Contact |Reward
| -------- | -------- |  -------- |  -------- |
|PointGoal2|PointGoal1|   0.0    | 19.0    | 
PointGoal1|PointGoal2| 0.58 | 21.24    | 

More details can be referred to: https://sites.google.com/view/safedreamer#:~:text=Assessment%20in%20Unseen-,Testing,-Environments

Reviewer comment (Experiments):
Have you ran experiments in other safe RL benchmark such as safe autonomous driving environments?

Author response:
Thank you for your valuable suggestions. To further validate our method, we conducted additional experiments beyond Safety-Gym, incorporating tests in the Safe Drive scenario, focused on safe autonomous driving environments. Specifically, we conducted supplementary tests of the SafeDreamer algorithm in three tasks on MetaDrive, Car-Racing and FormulaOne. The results demonstrate that SafeDreamer significantly outperforms the baseline algorithms. We have updated the descriptions of these environments and the corresponding experimental data in the revised version of our manuscript. For more details, please refer to Appendix C.1 - C.3. Some video demos can be seen in https://sites.google.com/view/safedreamer#:~:text=SafeDreamer-,MetaDrive,-SafeDreamer