SafeDreamer: Safe Reinforcement Learning with World Models
Authors: Weidong Huang
Jiaming Ji
Chunhe Xia
Borong Zhang
Yaodong Yang
Venue: ICLR 2024 poster (2024)
Keywords: Safe Reinforcement Learning
SafeRL
World Model
Primary area: reinforcement learning
TL;DR: A novel algorithm for Model-based Safe Reinforcement Learning.

Abstract: The deployment of Reinforcement Learning (RL) in real-world applications is constrained by its failure to satisfy safety criteria.
Existing Safe Reinforcement Learning (SafeRL) methods, which rely on cost functions to enforce safety, often fail to achieve zero-cost performance in complex scenarios, especially vision-only tasks. These limitations are primarily due to model inaccuracies and inadequate sample efficiency. The integration of the world model has proven effective in mitigating these shortcomings. In this work, we introduce SafeDreamer, a novel algorithm incorporating Lagrangian-based methods into world model planning processes within the superior Dreamer framework. Our method achieves nearly zero-cost performance on various tasks, spanning low-dimensional and vision-only input, within the Safety-Gymnasium benchmark, showcasing its efficacy in balancing performance and safety in RL tasks. 
Further details can be found in the code repository: https://github.com/PKU-Alignment/SafeDreamer.

Published as a conference paper at ICLR 2024
SAFEDREAMER: SAFE REINFORCEMENT LEARNING WITH WORLD MODEL
Weidong Huang +1.1 Jiaming Ji ni3
Chunhe Xia
Borong Zhang 1
Yaodong Yang 1.8
IInstitute of Artificial Intelligence, Peking University "Scooo of Cyber Science and Technology, Beihang University yaodong. yang@pku. edu.cn
ABSTRACT
The deployment of Reinforcement Learning (RL) in real-world applications is constrained by its failure to satisfy safety criteria. Existing Safe Reinforcement Learning (SafeRL. methods, which rely on cost functions to enforce safety, often fail to achieve zero-cost performance complex scenarios, especially vision-only tasks. These limitations are primarily due to model inaccuracies and inadequate sample efficiency. The integration of the world model has proven effective in mitigating these shortcomings. In this work, introduce SafeDreamer, a novel algorithm incorporating Lagrangian-basse processes within the superior Dreamer framework. Our method achieves nearly zero-cost performance on various tasks, spanning low-dimensional and vision-only benchmark, showcasing its efficacy in balanc-
methods into world model planning
input, within the ing performance and safety in RL tasks. Further details can be found in the code repository: https: ggittb..thuuuugggggggggggg..w.tt
INTRODUCTION
challenge in the real-world deployment of RL agents to prevent unsafe situations (Feng et al., 2023 Ji al. 2023a). SafeRL proposes practical solution by defining constrained Markov decision process (CMDP) (Altman 1999) and integrating an additional cost function to quantify potential hazardous behaviors. In this process, agents aim to maximize rewards while maintaining costs below predefined constraint thresholds. Several remarkable algorithms have been developed on
this foundation (Achiam et al. 2017 Yang al. 2022 Dai al. 2023 Guan al. 2024
However, with the cost threshold nearing zero, existing Lagrangian- based methods often fail to meet constraints meet them but cannot complete the tasks. Conversely, by employing internal dynamics model, an agent can effectively plan action trajectories that secure high rewards and nearly zero costs (González al.) 2015). This underscores the significance of dynamics models and planning in such contexts. However, obtaining a ground-trute dynamics model for planning in complex scenarios like vision-only autonomous driving is impractical (Li et al. 2022). Additionally,
the high expense of Long-horizon planning forces optimization over a finite horizon, yielding locally optimal and unsafe solutions. To fill such gaps, we aim to answer the following question:
How can we develop safety-aware world model to balance long-term rewards and costs of agent?
Autonomous Intelligence Architecture LeCun 2022) introduces a world model and incorporates costs that reflect the agent's level of discomfort. Subsequently, Hogewind et al. (2022); As et al. (2022); Jayant & Bhatnagar (002)) offer solutions to ensure cost and reward balance via the world model. However, these methods fail to realize nearly zero-cost performance across several environments due to inaccuracies in modeling the agent's 's safety in current or future states. In this work, we introduce SafeDreamer (see Figure I1),, which integrates safety planning and the Lagrangian method within world model to balance errors between cost models and critics. A detailed comparison with various algorithms can be found in Table In summary, our contributions are:
We present the online safety-reward planning algorithm (OSRP) (as shown in Figure Tp) and substantiate the feasibility of using online planning within the world model satisfy constraints in
"Equal Contribution. CCCoresppnddigg Author, Work done when Weidong Huang visited Peking University.

---

Published as a conference paper at ICLR 2024
(DBRP)
SenDhwemer (ISRMiaa))
seticn
aotion
actio
aeller
Porospt
Parcept
Percepe
Por
(4)
(b)
(C)
(d)
Figure I: The Architecture of Dreamrr (a) illustrates all components of Sreaerr,mrr, which distinguishes costs as safety indicators from rewards and balances them using the Lagrangian method and safe planner. The OSRP (b) and OSRP-Lag (c) variants execute online safety-reward planning (OSRP) within the world model for action generation, especially OSRP-Lag integrates online planning with the Lagrangian approach to balance long-term rewards and costs. The BSRP-Lag variant of SafeDreamer (d) employs background safety-reward planning (BSRP) via the Lagrangian method within the world model to update safe actor.
vision-only tasks. In particular, we employ the Constrained Cross-Entropy Method (Wen & Topcu 2018) in the planning process.
We integrate Lagrangian methods with the safety-rewan online and background planning within the world model to balance long-ten term reward and cost. This gives rise to two algorithms, OSRP-Lag and BSRP-Lag, as depicted in Figure [ and Td.
SafeDre handles both low-dimensional and visual input tasks, achieving nearly zero-cost per- formance in the Safety-Gyynnsínn benchmark. We further illustrate that SafeDreamer outperforms competing model-basee and model-free methods within this benchmark.
2 RELATED WORK
Safe Reinforcement Learning SafeRL aims to manage optimization objectives under constraints (Altman 1999) Achiam et al. (017)) introduced a general-purpose policy search algorithm focusing on iterative adherence to constraints, albeit with high computational costs due to second-orde methods. |Yang et al. (022)) deliver superior results by mitigating approximation errors inherent in Taylor approximations and inverting high-dimensional Fisher information matrices. (Sooll et al. (2022b) introduces Sauté MDPs, which climinate safety constraints by incorporating them into the state-space and reshaping the objective. This approach satisfies the Bellman equation and advances us toward solving tasks with constraints that are almost surely met. (Tase et al. (203) investigated optimizing policies based the expected upper bound of rewards in unsafe states. However, due to the inherent lack of robustness in the policies, these SafeRL. methods struggle to cope with the randomness of the environment and fail achieve zero violations, even when the cost limit is set to zero. He et al. (2023) proposed AutoCost that automatically identifies well-defined cost functions crucial for achieving zero violation performance. From the perspective of energy functions, MM et al. (2022a) simultancously synthesizes energy energy-function- based safety certificates and learns safe control policies, without relying on prior knowledge of control laws or perfect safety certificates.
Despite these advances, our experiments reveal that the convergence of safe model-free algorithms still requires substantial interaction with the environment, and model-based methods can effectively boost sample efficiency. Recently, SafeRL has been widely used in the safety alignment of large
language models (e.g., Safety Layer (Ji et al. 2024a), PPO-Lagrangian (Dai et al. 2024)). )
Safe Model-based RL Model-based RL approaches (Polydoros Nalpantidis 2017) serve as promising alternatives for modeling environment dynamics, categorizing into online plan- ning-utilizing the world model for action selection (Hafner et al. [019bb Hansen et al. (022). and background planning-leverren methods like MPC exemplify online planning by ensuring safety action generation (Camacho et al.)
the world model for policy updates (Hafner et al. (019a). Notably,
2007: Koller al. 2018 Wabersich & Zeilinger 2021 Liu al.) 2020 Zwane et al. 2023). though this can lead to sight-sighted decisions due to the limited scope of planning and absence of critics. Recent progress seeks to embed terminal value functions into online model planning, fostering
2

---

Published as a conference paper at ICLR 2024
Table 1: Comparison to prior works. We compare key components of SafeDreamer prior model-based and model-free algorithms, Predominantly, contemporary algorithms are tailored for low-dimensional input, while only a subset supports vision input. Our approach accommodates both input modalities. Moreover, Lagrangian-baser and planning based constitute two predominant strands of SafeRL algorithms, which we consolidate into a unified framework.
Method
Vision Input
Low con-dimensional Input
Lagrangian- based
Online Planning
Background Planning
CPO Achinm al 2017
PRO-LES Ray uE
2019
TRPO-Lag
2019
MBPPO-L
Davant Restnager 2022
LAMBDA As 2024
X
X
SafeLOOP Sikchhi
2023
X
X
DreamerV3s
Faater of 4125
X
SafeDreamer
consideration of long-ter term rewards (Sikchi et al., 2022 Moerland al. 2023 but not addressing long-term safety. On the other hand, background planning methods like those in Jayant & Bhatnagar (2022); Thomas al. (202T) employ ensemble Gaussian models and safety value functions to update policy with PPO (Schulman et al.] 2017) and SAC (Haarnoja et al. 2018). respectively. However, challenges remain adapting these methods tasks that require processing vision input. In this regard, LAMBDA (As et al. 2022) extends the DreamerV1 I (Hafner et al. 2019a) with principles from the Lagrangian and Bayesian methods. However, the constraints intrinsic to DreamerV limit its efficacy due to disregarding necessary adaptive modifications to reconcile variances in signal magnitudes and ubiquitous instabilities within all model elements (Hafner et al., 2023). This original framework's 's misalignment with online planning engenders suboptimal results and a deficiency in low- dimensional input tasks, thereby greatly reducing the benefits of the world model. Meanwhile, Safe SLAC (Hogewind et al., (002)) integrates the Lagrangian mechanism into SLAC (Lee al. 20020), achieving comparable performance to LAMBDA on vision-only tasks. Yet. it does not maximize the potential of the world model to augment safety, overlooking online background planning.
3 PRELIMINARIES
Constrained Markov Decision Process (CMDP) SafeRL is often formulated CMDP M AS,A, R, C, C,,4, 7) (Altman 1999). The state and action spaces are S and A, respectively. Transition probability IP(s'|s, a) refers to transitioning from reward acquired upon transitioning from to s' through action a. The cost function set C [(C,, H² encompasses cost functions C, : x A + IR and cost thresholds b;, 1,... m. The initial state distribution and the discount factor are denoted by (a-) : S - j0, 1] and y E (0, 1). respectively. A stationary parameterized policy my represents the action-taking probability To(a in state All stationary policies are represented as TT (mo 0 E IR?), where 0 is the learnable network parameter. We define the infinite-horize reward function and cost function as J (TTT) and o (T²). respectively, as follows:
s' under action a. R(s'(s, (s, a) stands for the
Do
x
j (TT) Eag
Σ
Σ x *C((se+1 at)
The goal of CMDP is to achieve the optimal policy:
T, arg max (T(T))
(I)
moElle
where TTc Nl n Jf < bi) denotes the feasible set of policies.
3

---

Published as a conference paper at ICLR 2024
Reward/Coss 12.C,
Reward and Cost Balance
Vafue/ Cost Valse
(441))
max min JØ ToElle Ap20
Ap (Jg b)
Signal Decoders
D(5s)
and Critics
V(s))
Sequence Model
Return and Cost Return Estimation
Observation
Encoerr
Latent Rollout
Optimized
Objective
Model
Optimization
Inlontpuz
wartable
Figure 2: Safety-reward planning process. The agent acquires an observation and employs the encoder distill it into a latent state $1. Subsequently, the agent generates action trajectories via policy and executes them within the world model, predicting latent rollouts of the model state and reward, cost, reward value, and cost value with each latent state. We employ TD(A) (Hafner et al. 2020 to estimate reward and cost return for each trajectory that used to update the policy.
Safe Model- basel RL Problem
We formulate a Safe Model-based RL problem as follows:
max J (TO) Ealll
{ S b, where
(2)
0
J
14 PP((1s),at)) no].
E
Σ t=
(()))1 81 Pol-ssi,y), ~ ~T
(4)
In the above, Pa(- az) is d-parameterized world model, assume the initial state so is sampled from the true initial state distribution μ and 81+1 P- ax), VI. > O. We would use the world model & to roll out imaginary trajectories and then estimate their reward and cost returns required for policy optimization algorithms. Without loss of generality, we will restrict our discussion to the case of one constraint with cost function C' and a cost threshold
4 METHODS
In this section, we introduce SafeDreamer, a framework for integrating safety-reward planning of the world model with the Lagrangian methods balance rewards and costs. The world model is trained through a replay buffer of past experiences as the agent interacts with the environment. Meanwhile, we elucidate the notation for our safe model-based agent, assuming access the learned world model, which generates latent rollouts for online or background planning, depicted in Figure 2 The design and training objectives of these models are described Section
4.1 MODEL COMPONENTS
SafeDreamer includes the world model and actor-critic models. At each time step, the world model receives an observation o, and an action at. The observation condensed into discrete representation žς. Then, along with the action, are used by the sequence model to predict the next representation 44+1- We represent the model state s, = {he, 21} by concatenating ht and 2;, where h, is a recurrent state. Decoders employ to predict observations, rewards, and costs. Mcanwhile, s, serves as the input of the actor-critic models to predict reward value or,, cost value U and action at Our model components are follows:
4

---

Published as a conference paper at ICLR 2024
Observation encoder:
21 ~ Es (zt
h4 0z)
Sequence model:
hes &t = Sq (h4-1, --1 ac-1)
Observation decoder:
Ōr O& (ō:
8t)
Actor:
as (at 8t)
Reward decoder:
ftt ~ Rs (ft
St)
Reward critic:
Ür Vv (Ur, 3t)
Cost decoder: G Co (Ce
8t)
Cost critie:
De V (e: 8t)
4.2 ONLINE SAFETY-REWARD PLANNING (OSRP) VIA WORLD MODEL
SafeDreamer (OSRP) We introduce online safety-reward planning (OSRP) through the world model, depicted in Algorithm [T] The on- line planning procedure is conducted at every decision time 4, generating state-action trajecto- ries from the current state 8 within the world model. Each trajectory is evaluated by learned reward and cost models, along with their crit- ics, and the optimal safe action trajectory is se- lected for execution in the environment. Specif- ically, we adopt the Constrained Cross-Entropy Method (CCEM) (Wen Topcu 2018) for plan- ning. To commence this process, we initial- ize (n² (f."))tt+ H. with nr, g E RRAI, which represent sequence action mean and stan- dard deviation spanning over the length of plan- ning horizon H. Following this, we indepen- dently sample N,,, state trajectories using the current action distribution N eration within the world model. Addi- tionally, we sample No, state trajectories using reward-driven actor, similar to Hansen al. (2022); Sikchi et al. (2022), accelerating the convergence planning. Afterward, we esti- mate the reward return JR and cost return J of each trajectory, as detailed in lines 9 14 of Algorithm 1 The estimation of JR is obtained using the TD(A) value of reward, denoted as Rd. which balances the bias and variance of the critics through bootstrapping and Monte Carlo The total cost over H steps, denoted as C. is computed Σ. H+ ((C) We use TD(A) value of cost CA to estimate
Algorithm 1: Online Safety-Reward Planning.
Input: current model state St. planning horizon H, num sample/policy/safe trajectories . Ney y No N.. Lagrangian multiplier Ap. cost limit b, 4t" for
for j f to / do
Init. empty safe/candiddtte/liite actions set As, Ac, Ae
Sample NEw + No traj. {&i, as s4+1)= 14+ using N( (o"-1)31), no within Sy with 8y as the initial state
for all New + No. trajectories do
Init. trajectory cost JC. (T)) 0- for t H to I do Compute R (st), (st) via
Equation (6)
J TJ C.H + Co(st)
12
end
13
J _ + R (s))
o ) at it-
S /H)H
14
if J < b then A. + As U {am44)]
16
17
end
end
19
Select sorting key Q candidate action set As Select the top-k action trajectories with highest Q values among A, as elite actions A, p. MEAN (A.), STD (A.)
29
21
end
Return: Nic 11)
value estimation Hafner et al.
2023
using the cost model Cs: J
JC and define J (JC. /H) alternative estimation for avoiding errors of the cost critic, where L signifies the episode length. Concurrently, trajectory satisfies the cost threshold b as J < b. We AA to denote the number of trajectories that meet the cost threshold and represent the desired number of safe trajectories N.. Upon this setup, one of the following two conditions will hold:
set the criterion for evaluating whether
1. If [AAl No: We employ -J as the sorting key Ω. The complete set of sampled action N New-N NTo serves as the candidate action set Ac.
trajectories (0u.tH h
2. If [A, 2 N.: We adopt Jh as the sorting key Ω for SafeDreamer (OSRP), and JR Ap for SafeDreamer (OSRP-Lag), respectively, with further discussion in Section 4.3 The safe action trajectory set A, is selected the candidate actions set Ac.
We select the top-k action trajectories, those with the highest Ω values, from the candidate action set A, to be the elite actions Ag. Subsequently, we obtain new parameters ue and o at iteration j:
_ Σ Al. 0d
V E (AL wi)2 The planning process is concluded after reaching
5

---

Published as a conference paper at ICLR 2024
predetermined number of iterations J. An action trajectory is sampled from the final distribution Ni (I)) Subsequently, the first action from this trajectory is executed in the real environment.
4.3 LAGRANGIAN METHOD WITH THE WORLD MODEL PLANNING
The Lagrangian method stands as a general solution for SafeRL. and the commonly used ones are Augmented Lagrangian (Dai et al. 2023) and PID Lagrangian (Stooke al. 2020). However, it has been observed that employing Lagrangian approaches for model-free SafeRI. results in suboptimal performance under low-cost threshold settings, attributable inaceuracies in the critic's estimation. This work integrates the Lagrangian method with online and background planning to balance errors between cost models and crities. By adopting the relaxation technique of the Lagrangian method (Nocedal & Wright 2006), the Equation (2) is transformed into an unconstrained safe model-based optimization problem, where >p the Lagrangian multiplier:
max min J (xo) Xp (Jg (mT) b)
(5)
ToEll, Ap20
SafeDreamer (OSRP-Lag) We introduced the PID Lagrangian method (Stooke et al. 2020 into our online safety-reward planning. yielding SafeDreamer (OSRP-Lag), as shown in Figure It. In the online planning process, the sorting key Ω is determined by J A when |A,1 2 N., where J approximated using the TD(A) value of cost, denoted as CA computed with Vo. and Cg:
(s,) ( +7 ((1 X)Vo, (SL+1) + XC (st-)) Vp. (8t)
if t< H if 1=H
(6)
The intuition of OSRP-Lag is to optimize a conservative policy under comparatively safe conditions, considering long-term risks. The process for updating the Lagrangian multiplier remains consistent with Stooke al. (2020) and depends on the episode cost encountered during the interaction between agent and environment. Refer to Algorithm 3 for additional details.
SafeDreamer (BSRP-Lag) Due to the time-consuming property of online planning, in order to meet the real-time requirements of certain scenarios, we extend our algorithm to support background planning. We use the Lagrangian method in background safety-reward planning (BSRP) for the safe actor training, which is referred to as SafeDreamer (BSRP-Lag) as shown Figure Dd. During the actor training, we produce imagined latent rollouts of length T = 15 within the world model in the background. We begin with observations from the replay buffer, sampling actions from the actor, and observations from the world model. The world model also predict rewards and costs, from which we compute TD(A) value (sx)) and (sx). These values are then utilized in stochastic backpropagation (Hafner et al., 2023) update the safe actor, a process
denominate background sewary planning. The training loss in this process guides the safe actor to maximize expected reward return and entropy while minimizing cost return, utilizing the Augmented Lagrangian method (Simo & Laursen
1992
As et al. 2022)::
(()) =
Σ sg (s)) + nh (To (a, | st) - V (sg(C Ap. 1=1
where sg(-) represents the stop- stop-gradient operator, / = max max + 1.0), 1.0) represents non- decreasing term that corresponds to the current gradient step k. v > Define A (C* b). The update rules for the penalty term in the training loss and the Lagrangian multiplier are follows:
x A +
A2
N
A
if +
V (sg(C> 21. 1²) ,2411_
(1#)²
(8)
otherwise.
2
PRACTICAL IMPLEMENTATION
Leveraging the framework above, we develop SafeDreamer, a safe model-based RL algorithm that extends the architecture of DreamerV3 (Hafner et al. 2023).

---

Published as a conference paper at ICLR 2024
World Model Implementation The world model, trained via variational auto-encoding (Kingmal & Welling 2013; Yang et al. 2024), transforms observations or into latent representations 2;. These representations used in reconstructing observations, rewards, and costs, enabling the evaluation of reward and safety action trajectories during planning. The representations 2 are regularized towards prior by regularization loss, ensuring the representations are predictable. We utilize the predicted distribution over & from the sequence model as this prior. Mcanwhile, we utilize representation 2; as a posterior, and the future prediction loss trains the sequence model leverage historical information from times before t to construct a prior approximating the posterior at time t. These models are trained by optimizing the oog-liielihod and KL divergence:
T
(()) =
KL b || sg($t)) + KI. (sg(2)) || &x]
regularization loss
Future prediction loss
B. In Og (oz 8,) B, In Ro (re st) B In Co (ce (1)
observation loss
roward loss
cost loss
Here, sg(-) is the stop-gradient operator. The world model based on the Recurrent State Space Model (RSSM) (Hafner et al. 2019b), utilizing GRU (Cho et al. (0014) for sequence modeling.
The SafeDreamer algorithm. Algorithm illustrates how the introduced components interrelate to build safe model-based agent. We sample batch of B sequences of length L from replay buffer to train the world model during each update. In SafeDreamer (OSRP), the reward-driven actor is trained by minimizing Equation (IO). aiming maximize both the reward and entropy. The reward-driven actor serves to guide the planning process, which avoids excessively conservative policy. In OSRP-Lag, the cost return of each episode utilized to update the Lagrangian multiplier Ap via Algorithm 3 Specifically, SafeDreamer (BSRP-Lag) update the safe actor and Lagrangian multiplier A via Equation (7) and Equation (8). respectively. Subsequently. the safe actor is employed to generate actions during exploration, which speeds up decision -aaoing and avoids the requirement for online planning. See Appendix A the design of crities and additional details.
EXPERIMENTAL RESULTS
We use different robotic agents in Safety-Gymnasien GGmnnsiimm Jji et al. 2023b). The goal in the environments is to navigate robots to predetermined locations while avording collisions with other objects. We evaluate algorithms using five fundamental environments (refer to Appendix C.4). Performance is
assessed using metrics from Ray al. (2019):
Average undiscounted reward return over E episodes: j(T) = 1 Σe Tt.
E Ei
Average undiscounted cost return over E episodes: j((z)= = Σ Ete
Average cost throughout the entire training phase, namely the cost rate: Given a total of T interaction steps, we define the cost rate Pe(TT) 2i
We calculate (T) and J.(TT) by averaging episode costs and rewards over E 10 episodes, each of length Tep = 1000, without network updates. Unlike other metrics, (T( is calculated using costs incurred during training, not evaluation.
Baseline algorithms. The baselines include: (1). PPO-Lag, TRPO-Lag (Ray et al. 2019 (Model- free): Lagrangian versions of PPO and TRPO. (2). CPO (Achiam et al. 2017) (Model free): A policy search algorithm with near-constraint satisfaction guarantees. (3). Sauté PPO, TRPO (Sootla et al.)
2022a (Model-free): Eliminates safety constraints into the state-space and reshape the objective function. (4). Safe SLAC (Hogewind al. (022) (Model-based):: Combines SLAC (Lee al.
2020) with the Lagrangian methods. (5). LAMBDA (As et al., 2022) (Model-based)) Implemented
DreamerVI. 1, combines Bayesian and Lagrangian methods. (6). MPC:sim (Wen Topeu 2018
Liu et al. 2020) (Model-based): Employs CCEM for MPC with ground-truth simulator, termed
https: thubittubboooo cttp.oooompkk-jo/p http:////J.Allllgnnnggaaaaaet www.nnt//aaetyyggmmeeemm sat.:aaffeyyyyayaaa

---

Published as a conference paper at ICLR 2024
Median
IQM
Mean
LAMBDA
Safe SLAC
Ours (OSkP-Lag)
Ours (OSRP)
(BSRP-Lag)
Normalized reward performance convergence
Median
IQM
Mean
LAMBDA
Safe SLAC
Ours (OSkP-Lag)
Ours (OSRP)
(BSRP-Lag)
Normalized cost performance convergence
Figure 3: Experimental results from the five vision tasks for the model-based methods. The results recorded after the agent completes the 2M training steps. We normalize the metrics following Ray et al. (2019) and utilize the rliable library (Agarwal et al. 202T) to calculate the median, inter-quartile mean (IQM), and mean estimates for normalized reward and cost returns.
MPC:sim. (7). MBPPO-Lag (Jayant & Bhatnagar 2022) (Model-based): Trains the policy through
ensemble Gaussian models and Lagrangian method. (8). DreamerV3 (Hafner al. 2023) (Model-
based): Integrates practical techniques and excels across domains. For further hyperparameter configurations, refer to Appendix B
The results of experiments. In Safety- Gymnasium tasks, the agent begins in safe region at episode reset without encountering obstacles. trivial feasible solution apparent to humans to keep the agent stationary, preserving its position. However, with this simplistic policy, realizing cost with the prior algorithms either demands substantial updates or remains elusive some tasks. As shown Figure SafeDreamer (BSRP-Lag) attains rewards similar to model-based RL methods like LAMBDA, with substantial 94.3% reduction in costs. The training curves of experiments are shown in Figure 4 and Figure 5 and each experiment conducted across runs. SafeDreamer achieves nearly zero-cost performance and surpasses existing model-free and
algorithms.
Dual objective realization: balancing enhanced reward with minimized cost. Figure SafeDreamer uniquely attains minimal costs while achieving higher rewards in the five visual-only safety tasks. In contrast, model-based algorithms such as Safe SLAC attain a cost limit beyond which further reductions are untenable due to the inaccuracies of the world model. On the other hand, in environments with denser or more dynamic obstacles, such as SafetyyPointuutoon 1, MPC struggles to ensure safety due to the absence of cost critic within limited online planning horizon. From the beginning of training, our algorithms demonstrate safety behavior, ensuring extensive safe exploration. Specifically, in the SafetyPointGoall and SafetyPointPushl environments, SafeDreamer matches the performance of DreamerV3 in reward while preserving nearly zero cost.
depicted in
Mastering in visual and low-dimensional tasks. We conducted evaluations within two low- dimensional input environments, namely SafetyPointGoall (vector) and SafetyCarGoall (vector) (refer to Figure [5). Although model-free algorithms can decrease costs, they struggle to achieve higher rewards. This challenge stems from their reliance on learning policy purely through trial- and-error, devoid of world model assistance, which hampers optimal solution discovery with limited data samples. Mcanwhile, the reward of MBPPO-Lag ceases to increase once the cost decreases to relatively low level. SafeDreamer surpasses them regarding both rewards and costs.
Discussion and ablation study. In the SafetyPointButtonl, which requires the agent navigate around dynamically moving obstacles, OSRP outperformed BSRP-Lag in terms of rewards. This may be attributed to its of the world model for online planning, enabling it to predict dynamic
8

---

Published as a conference paper at ICLR 2024
Sadety/oocnGoull
CurGalll
SafeetyPinnGeal2
SafetyPuintPush
Rewud
13
2.
0.0
0.5
1.0
35
20
155
240
Average Cost
100
15
05
1.0
15
1.3
16
0.
2.10
0.52
Rune
2.00
0.03
2.03
010
0.05
0.00
0.03-
0.00
0.09
0.0
085
1.0
2.0
0.5
1.11
1.5
3:1
0.5
1.0
1.5
3.0
0.00
3.0
0.5
1-0
155
5.0
0.00
0.0
0.1
1.3
1.5
20
Enn Steps (X20*)
Bay Steps (XIDP))
Eno Skepe (X100))
Iaw Steps (KCUP))
Eav Steps (X109)
MPC
DrcamerV33
Safo SLACC
FAMIIDA
Oher (OSRP)
Dane (05R0-
Ours (USRR-aaa))))
Figure 4: Comparing SafeDreamer to model-based baselines across five image-based safety tasks.
environmental changes in real time. In the SafetyPointPushl, the optimal policy involves the agent learning to wedge its head the middle gap the box to push quickly. We noticed that BSRP-Lag could converge this optimal policy, whereas OSRP and OSRP-Lag had difficulty discovering 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 finite time. BSRP's significant advantage in this environment suggests might be better suited for static environments that require more precise operations. In the SafetyPointGoal2, with more obstacles in the environment, OSRP-Lag, which introduces a cost critic to estimate cost return, tends to be more conservative than OSRP. Through the ablation experiments, we found that factors like the reconstruction loss of the world model and the size of its latent hidden states greatly impact SafeDreamer's 's ability to reduce cost. This indicates that when the world model accurately predicts safety-related states, is more likely to contribute to safety improvements, thereby achieving zero cost. For more experimental results, refer to Appendix
7 CONCLUSION
Averape Return
Averge Covl Reiarn
Cost Rate
So0ffhettoo aoooooalll (wers)
In this work, we tackle the issue of zero-cost performance within SafeRL., to find an optimal policy while satisfying safety constraints. We introduce SafeDreamer, safe model-based RI. algorithm that utilizes safety-reward planning of world models and the Lagrangian methods to balance rewards and costs. To knowledge, SafeDreamer is the first algorithm to achieve nearly zero-cost in final performance, utiliz- ing vision-only input in the Giynnasimmmm benchmark. However, SafeDreamer trains each task independently. incurring substantial costs with each task. Future research should leverage offline data from multiple tasks to pre-train the world model, examining its ability to facilitate the safety alignment of new tasks. Considering the various constraints that robots must adhere to in the real world, utilizing a world model for effective environmental modeling and deploying it in robots to accomplish tasks safely is potential direction (Chen al., 2022 Zhang et al. (003).. Another future work that can be explored is how to enhance the safety of large language models from the perspective of LLMs as World Models, through constrained optimization (Ji et al. 2024b Dai et al. 2024).
304
305
300
313
0.30
0.34
0.55
Sxpe
Exe Shape (C10P)
P0O
Figure 5: Results in low-dimensional input tasks.

---

Published as a conference paper at ICLR 2024
8 ACKNOWLEDGEMENT
This work is sponsored by the National Natural Science Foundation of China (62376013, 62272024) and the Beijing Municipal Science & Technology Commission (Z231 10000743305)
REPRODUCIBILITY STATEMENT
To advance open-source science, release 80+ model checkpoints along with the code for training and evaluating SafeDreamer agents. These resources are accessible at: https //github. Becthuu..oom/Ba. www.com///// www..//KO-ALLLimmttttteeeeeeeeee
All SafeDreamer agents are trained on one Nvidia 3090Ti GPU each and experiments are conducted utilizing the Safety-Gymnasium MetaDrive and Gymnasium benchmark We implemented the baseline algorithms based on the following code repository: LAMBDA Safe SLAC MBPPO-Lag" and Dreamerv respectively. We further included additional baseline algorithms: CPO, TRPO-Lag. PPO-Lag, Sauté PPO and Sauté TRPO, sourced from the Omnisafe (Ji et al., 2023c For additional hyperparameter configurations and experiments, please refer to Appendix and Appendix
REFERENCES
Joshua Achiam, David Held, Aviv Tamar, and Pieter Abbeel. Constrained policy optimization. In International conference machine learning, pp. 22-31. PMLR, 2017.
Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron C Courville, and Marc Bellemare. Deep reinforcement learning at the edge of the statistical precipice. Advances in neural information processing systems,
2021.
Eitan Altman. Constrained Markov decision processes: stochastic modeling. Routledge, 1999.
Yarden As, Ilnura Usmanova, Sebastian Curi, and Andreas Krause. Constrained policy optimization via bayesian world models. arXiv preprint arXiv:2201.0980 009002, 2022.
Marc G Bellemare, Will Dabney, and Rémi Munos. A distributional perspective on reinforcement learning. International conference on machine learning, pp. 449 458. PMLR, 2017.
Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojeiech Zaremba. Openai gym. arXiv preprint arXiv:1606. 01540, 2016.
Eduardo F Camacho, Carlos Bordons, Eduardo F Camacho, and Carlos Bordons. Model predictive controllers. Springer, 2007.
Yuanpei Chen, Tianhao Wu, Shengjie Wang, Xidong Feng, Jiechuan Jiang, Zongqing Lu, Stephen McAleer, Hao Dong, Song-Chun Zhu, and Yaodong Yang. Towards human-level bimanual dexterous manipulation with reinforcement learning. Advances in Neural Information Processing Systems, 15:5150-5166. 2022.
Kyunghyun Cho, Bart Van Merriënboer, Caglar Gulcehre, Dzmitry Bahdanau, Fethi Bougares, Holger Schwenk, and Yoshua Bengio. Learning phrase representations using rnn encoder-decoder for statistical machine translation. arXiv preprint urXiv: /406. 1078, 2014.
Josef Dai, Xuehai Pan, Ruiyang Sun, Jiaming Ji, Xinbo Xu, Mickel Liu, Yizhou Wang, and Yaodong Yang. Safe rihf: Safe reinforcement learning from human feedback. In The Twelfth International Conference on Learning Representations, 2024.
https//gitnbbbcoom/kklllggggggg safe
10

---

Published as a conference paper at ICLR 2024
Juntao Dai, Jiaming Ji, Long Yang, Qian Zheng, and Gang Pan. Augmented proximal policy optimization for safe reinforcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 37, pp. 7288-7295, 2023.
Shuo Feng, Haowei Sun, Xintao Yan, Haojie Zhu, Zhengxia Zou, Shengyin Shen, and Henry X Liu. Dense reinforcement learning for safery validation of autonomous vehicles. Nature, 615(7953): 620-627, 2023.
Farama Foundation. A standard api for single-agent reinforcement learning environments, with popular reference environments and related utilities (formerly gym). httpp:://githbb.co
Taraana-Poumaatioon/yymaaa uum 2022.
David González, Joshué Pérez, Vicente Milanés, and Fawzi Nashashibi. A review of motion planning techniques for automated vehicles. IEEE Transactions on intelligent transportation systems, 17(4): 1135-11455 2015.
Jiayi Guan, Guang Chen, Jiaming Ji, Long Yang. Zhijun Li, al. Voce: Variational optimization with conservative estimation for offline safe reinforcement learning. Advances in Neural Information Processing Systems, 36, 2024.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In International conference on machine learning, pp. 1861-1870. PMLR, 2018.
Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. arXiv preprint arXiv:1912.01603, 2019a.
Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In International conference on machine learning, pp. 3555-2565 PMLR, 2019b.
Danijar Hafner, Timothy Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. arXiv preprint arXiv:2010.02/93, 2020.
Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models. arXiv preprint arXiv:2301.044 044144, 2023.
Nicklas Hansen, Xiaolong Wang, and Hao Su, Temporal difference learning for model predictive control. arXiv preprint arXiv:2203.0495: 2022.
Tairan He, Weiye Zhao, and Changliu Liu. Autocost: Evolving intrinsic cost for zero-violation reinforcement learning. Proceedings of the AAAI Conference on Artificial Intelligence. 37(12): 14847-14855, Jun. 2023, doi: 10..1609/aaiivv111111 w37.1/..2..7111...7...2 URL httpp:ss:///os index. Atdex.php/AAAA/// http://Awwarrttcceeeeeeee
org/
Yannick Hogewind, Thiago Simao, Tal Kachman, and Nils Jansen. Safe reinforcement learning from pixels using stochastic latent representation. arXiv preprint wxXiv:22/0.0180)) 2022.
Ashish K Jayant and Shalabh Bhatnagar. Model-based safe deep reinforcement learning via constrained proximal policy optimization algorithm. Advances in Neural Information Processing Systems, 35:24432-2444 2022.
Jaming Ji, Tianyi Qiu, Boyuan Chen, Borong Zhang, Hantao Lou, Kaile Wang, Yawen Duan, Zoonhaa He, Jiayi Zhou, Zhaowei Zhang, et al. Ai alignment: A comprehensive survey, arXiv preprint arXiv:2310. 19852, 2023a.
Jiaming Ji, Borong Zhang. Jiayi Zhou, Xuehai Pan, Weidong Huang, Ruiyang Sun, Yiran Geng, Yifan Zhong, Josef Dai, and Yaodong Yang. Safety gymnasium: unified safe reinforcement learning benchmark. Advances in Neural Information Processing Systems, 36, 2023b.
Jiaming Ji, Jiayi Zhou, Borong Zhang, Juntao Dai, Xuehai Pan, Ruiyang Sun, Weidong Huang, Yiran Geng, Mickel Liu, and Yaodong Yang. Omnisafe: An infrastructure for accelerating safe reinforcement learning research. arXiv preprint M.Xiv:2305509304 2023c.

---

Published as a conference paper at ICLR 2024
Jiaming Ji, Boyuan Chen, Hantao Lou, Donghai Hong, Borong Zhang, Xuehai Pan, Juntao Dai, and Yaodong Yang. Aligner: Achieving efficient alignment through weak-to-strong correction. arXiv preprint ar 402.:2240..02416, 2024a.
Jiaming Ji, Mickel Liu, Josef Dai, Xuehai Pan, Chi Zhang, Bian, Boyuan Chen, Ruiyang Sun,
Yizhou Wang, and Yaodong Yang. Beavertails: Towards improved safety alignment of Ilm via human-preete cemannpreeerenee dataset. Advances in Neural Information Processing Systems, 36, 2024b.
Diederik P Kingma and Max Welling. arXiv:1312.6114, 2013.
Auto-encoding variational bayes.
arXiv preprint
Torsten Koller, Felix Berkenkamp, Matteo Turchetta, and Andreas Krause. Learning- based model predictive control for safe exploration. In 2018 IEEE Conference on Decision and Control (CDC), pp. 6059-6066. IEEE, 2018.
Yann LeCun. A path towards autonomous machine intelligence version 0.9. 2, 2022-06-27. Open Review, 62, 2022.
Alex X Lee, Anusha Nagabandi, Pieter Abbeel, and Sergey Levine. Stochastic latent actor-critic: Deep reinforcement learning with a latent variable model. Advances in Neural Information Processing Systems, 33:741-752, 2020.
Quanyi Li, Zhenghao Peng, Lan Feng, Qihang Zhang, Zhenghai Xue, and Bolei Zhou. Metadrive: Composing diverse driving scenarios for generalizable reinforcement learning. IEEE transactions on pattern analysis and machine intelligence, 45(3):3461-3444 2022.
Zuxin Liu, Hongyi Zhou, Baiming Chen, Sicheng Zhong, Martial Hebert, and Ding Zhao. Con- strained model-based reinforcement learning with robust cross-entropy method. arXiv preprint arXiv:2010.07968, 2020.
Haitong Ma. Changliu Liu, Shengbo Eben Li, Sifa Zheng, and Jianyu Chen. Joint synthesis of safety certificate and safe control policy using constrained reinforcement learning. In Learning for Dynamics and Control Conference, pp. 97-109. PMLR, 2022a.
Yecheng Jason Ma, Andrew Shen, Osbert Bastani, and Jayaraman Dinesh. Conservative and adaptive penalty for model-based safe reinforcement learning. In Proceedings of the AAAI conference on artificial intelligence, volume 36, pp- 5404-5412, 2022b.
Thomas M Moerland, Joost Broekens, Aske Plaat, Catholijn M Jonker, et al. Model-based rein- forcement learning: A survey. Foundations and Trends@ in Machine Learning, 16(1):1-118, 2023.
Jorge Nocedal and Stephen J. Wright. Numerical Optimization. Springer, New York, NY, USA, second edition, 2006.
Athanasios S Polydoros and Lazaros Nalpantidis. Survey of model-based reinforcement learning: Applications on robotics. Journal of Intelligent Robotic Systems, 86(2):153-173, 2017.
Alex Ray, Joshua Achiam, and Dario Amodei. Benchmarking safe exploration in deep reinforcement learning. arXiv preprint 1911v:19,, 7:1, 2019.
Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, et al. Mastering atari, go, chess and shogi by planning with a learned model. Nature, 588(7839):604-699 2020.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347 2017.
Harshit Sikchi, Wenxuan Zhou, and David Held. Learning off off-policy with online planning. In Conference on Robot Learning, pp. 1622-1633. PMLR, 2022.
J Ci Simo and TA1143885 Laursen. An augmented lagrangian treatment of contact problems involving friction. Computers & Structures, 42(1))97-111 1992.
12

---

Published as a conference paper at ICLR 2024
Aivar Sootla, Alexander RiversRivers, Taher Jafferjee, Ziyan Wang, David H Mguni, Jun Wang, and Haitham Ammar. Sauté rl: Almost surely safe reinforcement learning using state augmentation. In International Conference on Machine Learning, pp.
PMLR, 2022a.
Aivar Sootla, Alexander Cowen-Rivers, Taher Jafferjee, Ziyan Wang, David H Mguni, Jun Wang, and Haitham Ammar. Sauté rl: Almost surely safe reinforcement learning using state augmentation. In International Conference on Machine Learning, pp. 20443. PMLR, 2022b.
Adam Stooke, Joshua Achiam, and Pieter Abbeel. Responsive safety in reinforcement learning by pid lagrangian methods. In International Conference on Machine Learning, pp. 9133-9143, PMLR, 2020.
Geraud Nangue Tasse, Tamlin Love, Mark Nemecek, Steven James, and Benjamin Rosman. Rosarl: Reward-only safe reinforcement learning. arXiv preprint arXiv:23006. 000000030000035 2023.
Garrett Thomas, Yuping Luo, and Tengyu Ma. Safe reinforcement learning by imagining the near future. Advances in Neural Information Processing Systems, 34:13899-13899 2021.
Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-based control. In 2012 IEEE/RSJ international conference on intelligent robots and systems, pp. 5026-5033. IEEE, 2012.
Kim P. Wabersich and Melanie N. Zeilinger. A predictive safety filter for learning baased control of constrained nonlinear dynamical systems, 2021.
J Beau W Webber. A bi-symmetric log transformation for wide-range data. Measurement Science and Technology. 24(2):027001, 2012.
Min Wen and Ufuk Topcu. Constrained cross-entropy method for safe reinforcement learning. Advances in Neural Information Processing Systems, 31, 2018.
Long Yang. Jiaming Ji, Juntao Dai, Linrui Zhang, Binbin Zhou, Pengfei Li, Yaodong Yang, and Gang Pan. Constrained update projection approach to safe policy optimization. Advances in Neural Information Processing Systems, 35:9111-995 2022.
Zhiqin Yang. Yonggang Zhang. Yu Zheng, Xinmei Tian, Hao Peng, Tongliang Liu, and Bo Han. Fedfed: Feature distillation against data heterogeneity in federated learning. Advances in Neural Information Processing Systems, 36, 2024.
Yiming Zhang, Quan Vuong, and Keith Ross. First order constrained optimization in policy space. Advances in Neural Information Processing Systems, 01211333-153499 2020.
Zeyu Zhang, Muzhi Han, Baoxiong Jia, Ziyuan Jiao, Yixin Zhu, Song-Chun Zhu, and Hangxin Liu. Learning a causal transition model for object cutting. In 2023 JEEE/RSJ International Conference on Intelligent Robots and Systems (IROS), pp. 1996-2003. IEBE, 2023.
winelukwanda Zwane, Denis Hadjivelichkov, Yicheng Luo, Yasemin Bekiroglu, Dimitrios Kanoulas, and Mare Peter Deisenroth. Safe trajectory sampling in model- based reinforcement learning. In 2023 IEEE 19th International Conference on Automation Science and Engineering (CASE), pp- 1-6. IEEE, 2023.
13

---

Published as a conference paper at ICLR 2024
A DETAILS OF SAFEDREAMER TRAINING PROCESS
A.1 THE TRAINING PROCESS SAPEDREAMER
Algorithm SafeDreamer.
Input: batch length T, batch size B, episode length L, initial Lagrangian multiplier Ap Initialize the world model parameters #, actor-critic parameters 0, V Vo. Initialize dataset D using random policy;
while not converged do
Sample B trajectories for, de, 01+1, T1+1; 441h1+47 D
Update the world model via minimize Equation 9
Condense 0t:++7 into Sta++ via the world model;
Generate latent rollouts using the actor within the world model with 86447 as the initial state; Update the reward and cost critics via minimize Equation (1I
Update the actor and Lagrangian multiplier Xp;
10
Te < for = 1
11
L do
12
Condense o, into latent state via 8, So(or, $1 a,-1) Sample ae from the safe actor or Online Safety-Rewan Planner; Execute action at observe 0++1, T1+1, C++1 returned from the environment;
14
Update dataset D < D U for, ae, 04+1, (t+1.C(41)
16
Tc - Ic + C++1;
17
end
18 end
A.2 THE PROCESS OF PID LAGRANGIAN
Algorithm 3: PID Lagrangian (Stooke et al.
2020
Input: Proportional coefficient Ky. integral coefficient K,, differential coefficient K₂;
Initialize previous integral item: _ +- 0;
Initialize previous episode cost: Jo 0;
while iteration continues do Receive cost Jh:
P - JJ - d; D + max (0, J&C- - -
Ih + mar(0,, +
Ap maz(0, KpP + K₁I + KzD); Return Lagrangian multiplier Api
10 end
PID Lagrangian with Planning The safety planner might fail to ensure zero-cost when the re- quired hazard detection horizon surpasses the planning horizon, in accordance with the constraint (J..H (H)L < during planning. To overcome this limitation, we introduced the PID Lagrangian method (Stooke et al.) 2020) into our planning scheme for enhancing the safety of exploration, yielding SafeDreamer (OSRP-Lag).
Cost Threshold Settings Our empirical studies show that for all SafeDreamer variations, cost threshold of 2.0 leads to consistent convergence towards nearly zero cost in vision tasks. The reason we did not set the cost threshold to zero due to potential errors in the trained cost model and cost value model. Therefore, even in 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 these estimations. making online planning overly conservative and unable to converge. contrast, SafeDreamer (BSRP-Lag) manages to operate with zero-cost threshold some tasks but requires more data samples to achieve convergence. For the
14

---

Published as a conference paper at ICLR 2024
low-dimensional input tasks, we evaluate the cost via observation reconstruction and do not directly use the cost model, so we set the cost threshold at zero for our method.
Reward-driven Actor A reward-driven actor is trained to guide the planning process, enhancing reward convergence. The actor loss aims to balance the maximization of expected reward and entropy.
The gradient of the reward term is estimated using stochastic backpropagation (Hafner et al. 2023). whereas that of the entropy term is calculated analytically:
(()) = Σ sg (s:)) + nH (To (ae st)]
(10)
Reward and Cost Critics The sparsity and heterogeneous distribution costs within the environment complicate the direct regression of these values, thereby presenting substantial challenges for the
learning efficiency of the cost critic. Leveraging the methodology from (Hafner et al., 2023), we
employ trio of techniques for critic training: discrete regression, twohot encoded targets (Bellemare it al. 2017 Schrittwieser et al.) 2020), and symlog smoothing (Hafner et al. 2023 Webber 2012);
Lreward Critiic (V,,)
sg (twohot (symlog (R*)))² In Poi. (c)
(11)
1=1
where the critic network output the Twohot softmax probability p& (b; (1) and
twohot(z); ={ [b++1 1 - x / 1bx+1 bul lb 21 bb+1 bkl
if n if + else
[B|
(by
B 1020...
We define the bi-symmetric logarithm function (Webber 2012 symlog(x) sign In([ + 1) and its inverse function, as symexp(x) sign(z)) (exp r 1). The cost critic loss, represented Leest.critee ( is calculated similarly, using TD(A) value of cost CA, Additionally, extending the principle of Onehot encoding to continuous values, Twohot encoding allows us to reconstruct
target values after encoding:vy (8t) symexp (pv (. [x)=² B). For further details, please refer to (Hafner et al.
2023
15

---

Published as a conference paper at ICLR 2024
B HYPERPARAMETERS
The SafeDreamer experiments were executed utilizing Python3 and Jax 0.3.25, facilitated by CUDA 11.7, on an Ubuntu 20.04.2 LTS system (GNU/Linux 5.8.0-59-generic x86 64) equipped with 40 Intel(R) Xcon(R) Silver 4210R CPU cores (operating at 240GHz), 251GB of RAM, and an array of GeForce RTX 3090Ti GPUs.
Table 2: Hyperparameters for SafeDre When addressing other safety tasks, we suggest tuning the initial Lagrangian multiplier, proportional coefficient, integral coefficient, and difference coefficient at various scales. Name
Symbol
Value
Description
World Model
Number of laten
Classes per laten
48
Batch size Batch length
B
64
T
16
Learning rate
10-4
Coefficient of kl divergence loss Coefficient of decoder in loss
Clg- Bor Pr. Bc
0.1, 0.5
1.0, 1.0, 1.0
Planner
Planning horizon Number of samples Mixture coefficient
H
15
Nay M
500
0.05/0.0
No. M * Nay different vec- tor/visual input
Number of iterations Initial variance
J
dg
1.0
PID Lagrangian
Proportional coefficient Integral coefficient
K
0.1
K;
0.01
Differential coefficient
Ka
Initial Lagrangian multiplier Lagrangian upper bound
X
0.0
0.1
Maximum of
Augmented Lagrangian
Penalty term Initial Penalty multiplier Initial Lagrangian multiplier
5-9
4
1 1
2
0.01
Actor Critic
Sequence generation horizon
155
Discount horizon
0.997
Reward lambda
X
0.95
Cost lambda
A
0.95
Learning rate
3 10-
General
Number of other MLP layers Number of other MI.P layer units Train ratio Action repeat
512
MPC:sim. We apply CCEM to MPC using ground-truth simulator, denoted as MPC:sim, omitting the of value function, thus establishing a non-parametric baseline. We restrict the number of sample trajectories to 150, with planning horizon of 15 and 6 iterations, due to the computational demands of simulator-bass based planning.
16

---

Published as a conference paper at ICLR 2024
Table 3: Hyperparameters for Safe SLAC. We set the hyperparameters for Safe SLAC following Hogewind et al. (2022). We maintain the original hyperparameters unchanged, with the exception of the action repeat, which we adjust from its initial value of 2 to This adjustment is driven by empirical observation that employing more rapid update rate in visual environments yields advantages for both LAMBDA and our approach. Furthermore, we increase the length of sampled sequences from 10 15 and reduce the cost limit from 25.0 to 2.0 to ensure a fair comparison with our method.
Name
Value
Length of sequences sampled from replay buffer Discount factor
15
0.99
Cost discount factor
0.995
Replay buffer size
2 * 10
Latent model update batch size Actor-critic update batch size
32
64
Latent model learning rate Actor-critic learning rate
x 10
2 x 10
Safety Lagrange multiplier learning rate Action repeat Cost limit
2e
2.0
Initial value for a Initial value for À Warmup environment steps Warmup latent model training steps
10
2 10-2
60 * 103
30 10
Gradient clipping norm
40
Target network update exponential factor
* 10-3
Table 4: Hyperparameters for LAMBDA.
We evaluate our approach against the official implementa- tion presented by As et al. (2022). We employ a sequence generation horizon of 15, in line with our algorithm's 's configuration. Additionally, we have modified the cost limit from 25.0 2.0 to assess its convergence under a lower cost limit setting.
Name
Value
Sequence generation horizon
15
Sequence length
50
Learning rate
le-4
Burn-in steps
500
Period steps
200
Models
20
Decay
0.8
Cyclic LR factor
5.0
Posterior samples
5
Safety critic learning rate
2c-4
Initial penalty
5e-9
Initial Lagrangian
le-6
Penalty power factor
le-5
Safety discount factor
0.995
Update steps
100
Critic learning rate
8e-5
Policy learning rate
8e-5
Action repeat
4
Discount factor
0.99
TD(X) factor
0.95
Cost limit
2.0
Batch size
32
17

---

Published as a conference paper at ICLR 2024
Table 5: Hyperparameters for MBPPO-Lag. We utilize the official implementation provided by
Jayant Bhatnagar (2022) and include a listing of its hyperparameters for the sake of completeness.
Name
Value
Number of models ensemble
8
Hidden layers in single ensemble NN
4
Hidden layers in single ensemble NN
200, 200, 200]
Hidden layers in Actor NN
2
Hidden layers in Actor NN Hidden layers in Critic NN Hidden layers in Critic NN Gradient descent algorithm Actor learning rate
[64, 64]
2
[64, 64]
Adam
3e-4
Critic learning rate
le-3
Lagrange multiplier learning rate
5e-2
Initial Lagrange multiplier value
1
Cost limit
2.0
Activation function
tanh
Discount factor
0,99
GAE parameter
0.95
Validation dataset/Train dataset
10%/90%
PR threshold
66%
used in (29)
0.02
Table 6: Hyperparameters for model-free algorithms, We list the most important hyperparameters for the model-free baselines. This enumeration encompasses the salient hyperparameters essential for the model-free baselines. It is imperative to note that our approach closely adheres to the implementation presented by et al. (003c)), with certain enhancements. Specifically, we have elevated the hidden layer sizes of the MLP from 64 to 512 and augmented the number of hidden layers from 2 to 4, ensuring a fair comparison with model-based algorithms. our experiments, we found that cost discount factor has substantial impact the efficacy of Sauté RL methods, and setting it to 0.999 resulted in better performance in the Bafety-Gyymasiingg Moreover, it is noteworthy that our chosen cost limit is set at 2.0, opposed to the conventional value of 25.0, which is the standard practice in Ji et al. (2023b): Ray et al. (2019). This discrepancy highlights a challenge inherent in model-free algorithms, especially in situations with tight budget constraints.
Name
CPO
PPO-Tag
TRPO-Tag
Santé TRPO
Snuté PRO
Batch size
128
64
128
128
64
Target KI.
0.01
0.02
0.01
0.01
0.02
Max gradicat norm
40.0
40.0
40.0
40.0
Critic norm coefficient
0.001
0.001
0.001
0.001
0.001
Discount
0.99
0.99
Cost discount factor
0.99
0.99
0.99
0.999
0.999
Lambda for gae
Lambda for cost gae
0.95
0.95
0.95
0.95
0.95
Advantage estimation method
gae
gae
Unsafe reward
None
Nome
Nome
0.2
0.2
Actor-critic
MLP Sidden sizes
[512,
312
5122 512, 512, 512]
[512, 512, 512
[512, 512, 512
[122, 512,
Activation
tunh
tah
lanh
tanb
uah
Leamiing rate
0.001
0.0003 20
0.001
0.001
0.0003
Cost limit
2.0
2.0
2.0
2.0
Lagrange
Trital I Lagrange Multiphier value
None
0.001
0101
None
None
Lagrango learning
None
0.035
None
Lagrange optimizer
None
Adam
Adam
None
None
18

---

Published as a conference paper at ICLR 2024
C EXPERIMENTS
C.I EXPERIMENTS ON METADRIVE
Nundrr Mova the
(a) Safe
(b) Unsafe
Figure 6: The MetaDrive benchmark. The objective for the car is navigate successfully predetermined destination. During this process, the vehicle incurs a cost penalty in instances of collision with obstacles other vehicles, well when deviating from the designated roadway.
2022) stands comprehensive, effective simulation environment for au- MetaDrive (Li et al.
tonomous vehicle research. It features designed environments optimized for developing safe policies. The observation in MetaDrive combines vector input and first-person view input. The cost function in this setting is defined as follows:
C(s,o) (1,
if collides of the road otherwise
Table 7: Results Metadrive.
Arrival Rate Episode Cost Return
Ours (BSRP-Lag)
1.0 + 0.0
2.2 0.1
PPO-Reward Shaping 0.72 + 0.1
14.1 6.2
We maintained consistent environment by conducting both training and testing the identical roadway. To augment the complexity, vehicle positions were randomized at the beginning of each episode reset. As illustrated in Table 7 SafeDreamer successfully achieved 100% rate in 10 evaluation tests after engaging in 4M steps of environmental interaction. Notably, in comparison to IPO-Reward Shaping, SafeDreamer demonstrated enhanced efficiency in cost reduction. We observed that SafeDreamer has adopted a safe policy. Specifically, pauses to monitor traffic flow when encountering dense vehicle presence ahead, resuming forward progression only when the vehicular density diminishes.
19

---

Published as a conference paper at ICLR 2024
(a) Safe
(b) Unsafe
Figure 7: The Car-Racing task, introduced by Brrockman al. (016)), involves 2D image-based racing simulation with tracks that vary in each episode. The observation comprises a 64x64x3 pixel top-down view of the car. The action space is continuous and three-dimensional, encompassing steering, acceleration, and braking actions.
C.2 EXPERIMENTS ON CAR-RACIA CAR-RACING
We utilize the Car- Rac-Rining environment from OpenAI Gym's gymnas uum envs box2d inter- face, as established by Brockman et al. (2016). Our adaptation involves augmenting this environment with cost function analogous the one described in Ma et al. (2022b). A cost of 1 per step is assigned for any instance of wheel skid due to excessive force:
{1, if any wheel of the car skids
C(s,a)
0,
otherwise
Average Reward Return
Average Cost Return
Cost Rate
1000
800
0.60
750
600
0.45
500
400
0.30
250
200
0.15
0.00
0.0
0.5
2.0
0.0
0.5
1.0
1.5
2.0
1.5
2.0
Env Steps (x105)
Env Steps (x105)
Env Steps (x10*)
DreamerV3
Ours (BSRP-Lag)
Figure 8: Training curve on Car-Racing.
Table 8: Results on Car-Racing.
Average Reward Return Average Cost Return
Ours (BSRP-Lag)
763.1 + 1.34
0.04 + 0.01 580. 1 : 2.35
DreamerV3
775.1 + 0.56
SafeDreamer demonstrates rapid convergence, achieving zero-cost within 0.2M steps. Compared to DreamerV3, SafeDreamer significantly reduces vehicle slippage, thereby enhancing driving robustness.
As depicted in Figure 8 and Table g
20

---

Published as a conference paper at ICLR 2024
C.3 FORMULAONE
(a) Environment map
(b) First-person perspective
Figure 9: The FormulaOne benchmark.
We utilized the visual environment FormulaOne from Ji et al. (003b)) where the agent receives 64x64x3 image input, as illustrated in Figure 9 The environment' 's complexity and visual aspects significantly increase the demands the algorithm. We employed Level 1 of FormulaOne, which requires the agent to reach the goal while avoiding barriers and racetrack fences. Each episode begins with the agent randomly placed at one of checkpoints.
As illustrated in Figure 10 and Table SafeDreamer attains nearly zero-cost after 1.5M training steps and demonstrates efficient obstacle avoidance in task completion. This indicates that SafeDreamer can achieve convergence and accomplish nearly zero-cost even in complex visual environments.
Average Reward Return
Average Cost Return
Cost Rate
0.06
30
60
0.04
40
0.03
0.01
0-
0.00
0.0
1.2
Env Steps (x10%)
Env Steps (x106)
Env Steps (x (x105)
DreamerV3
Ours (BSRP-Lag)
Figure 10: Results FormulaOne.
Table 9: Results on FormulaOne.
Average Reward Return Average Cost Return
DreamerV3
32,2 1.22
56.36 + 2.32
Ours (BSRP-Lag)
25.3 0.51
0.43 + 0.08
21

---

Published as a conference paper at ICLR 2024
C.4 EXPERIMENTS ON SAFETY-GYMNABE
Safety-Gymmasiin um is an environment library specifically designed for SafeRL. This library builds the foundational Gymna s API (Brockman et al. 2016) Foundation 2022). utilizing the high-performance MuJoCo engine (Todorov et al. 2012). We conducted experiments in five different environments, namely, SafetyPointGoll. SafetyPoimtGoale SafetyPoimtPushl and SafetyCarGoall, as illustrated in Figure [T] Following Jji ct al. (2023b), we adjusted the arrangement of the cameras to provide more information to planning- based algorithms in processing visual inputs compared to Ray et al.] (2019), as shown in Figure 12] Following Liu et al. (2020);
SufetyPointButton..
Jayant & Bhatnagar (2022): Sikehi et al. (2022), we modified the state representation to facilitate model learning better. All experiments utilize identical settings.
PaittoPointiiall
(b): SafetyPoiniGeal2
(c): PafttyPoiniPush!
(d): SaferyPoinnButionn
CarCoolll
Figure 11: The tasks in
Global view
Forward view
Rear view
Figure 12: Different views in the
C.4.1 AGENT SPECIFICATION
We consider three robots: Point, Car and Racecar (as shown in Figure 13 To potentially enhance learning with neural networks, we maintain all actions as continuous and linearly scaled to the range of [-1, +1]. Detailed descriptions of the robots are provided as follows:
(a) Point
(b) Racecar
(c) Car
Figure 13: Robots in the Safety- Gymnasiuu.
22

---

Published as a conference paper at ICLR 2024
The Point robot, functioning in 2D plane, is controlled by two distinct actuators: one governing rotation and another controlling linear movement. This separation of controls significantly cases its navigation. A small square, positioned in the robot's front, assists in visually determining its orientation and crucially supports Point effectively manipulating boxes encountered during tasks.
Point.
Car. The Car robot apparatus operating within a direeddimnsiooal space. is equipped with two independently powered wheels positioned in parallel, complemented by rear wheel that rolls freely, This configuration necessitates a coordinated manipulation of the dual propulsion systems to achieve both navigational steering and linear movement in the forward and reverse directions. Although this robot exhibits characteristics akin to those of a rudimentary Point robot, it introduces additional complexities due its design.
The Racecar robot exhibits realistic car dynamics, operating in three dimensions, and controlled by a velocity and a position servo. The former adjusts the rear wheel speed to the target, and the latter fine-tunes the front wheel steering angle. The dynamics model is informed by the widely recognized MIT Racecar project. To achieve the designated goal, it must appropriately coordinate steering angle and speed, mirroring human car operation.
Racecar.
C.4.2 TASK REPRESENTATION
Tasks within fameay-Gymmasium distinct and confined single environment each, shown in Figure 14
(a) Goal
(b) Button
(c) Push
Figure 14: Tasks in the Symmas-Gmmmmsiumm.
The task requires robot to navigate towards multiple target positions. Upon each successful arrival, the robot's goal position is randomly reset, retaining the global configuration. Autaining target location, signified by entering the goal circle, provides sparse reward. Additionally, dense reward encourages the robot's progression through proximity to the target.
Goal:
The task requires robot manipulate box towards several target locations. Like the goal task, a new random target location is generated after each successful completion. The sparse reward granted when the box enters the designated goal circle. The dense reward comprises two parts: one for narrowing the agent-box distance, and another for advancing the box towards the final target.
Push:
The task requires the activation of numerous target buttons distributed across the environment. The agent navigates and interacts with the currently highlighted button, the goal button. Upon pressing the correct button, new goal button is highlighted, while maintaining the overall environment. The sparse reward is issued upon successfully activating the current goal button, with the dense reward component encouraging progression toward the highlighted target button.
Button:
C.4.3 CONSTRAINT SPECIFICATION
These are used to symbolize substantial cylindrical obstacles within the environment, typically incurring costs upon contact.
Pillars:
Hazards:
These are integrated to depict risky areas within the environment that induce costs when agent navigates into them.
23

---

Published as a conference paper at ICLR 2024
Vases: Exclusively incorporated for Goal tasks, vases denote static and delicate objects within the environment. Contact or displacement of these objects yields costs for the agent.
Gremlins: Specifically employed for Button tasks, gremlins signify dynamic objects within the environment that can engage with the agent. Contact with these objects yields costs for the agent.
C.4.4 EVALUATION METRICS
In our experiments, we employed a specific definition of finite horizon undiscounted return and cumulative cost. Furthermore, we unified all safety requirements into single constraint (Ray et al. 2019). The safety assessment of the algorithm was conducted based on three key metrics: average episodic return, average episodic cost, and the cost rate. These metrics served as the fundamental basis for ranking the agents, and their utilization as comparison criteria has garnered widespread recognition within the SafeRI. community (Achiam al. 2017 Zhang al. 2020 As et al. 2022).
Any agent that fails to satisfy the safety constraints is considered inferior to agents that meet these requirements or limitations. other words, meeting the constraint is prerequisite for considering agent superior.
When comparing two agents, A and B, assuming both agents satisfy the safety constraints and have undergone the same number of interactions with the environment, agent A is deemed superior to agent B if it consistently outperforms agent B in terms of return. Simply put, if agent consistently achieves higher rewards over time, it considered superior to agent B.
In scenarios where both agents satisfy the safety constraint and report similar rewards, their relative superiority is determined by comparing the convergence speed of the average episodic cost. This metric signifies the rate at which the policy can transition from an initially unsafe policy feasible set. The importance of this metrie in safe RI. research cannot be overstated.
C.4.5 SAFEDREAMER TRAINING RESULTS WITHIM 8M STEPS
SaferyPointGoall
CaretyCa
SafetyPoiiniButttnn
Safety PaintPushi
PoinsyPoitt
Average Revard
12-
Avergee
0.98
0.24
Cost. Batt
0.52
0.52
0.02
0.33
0.00
0.00
Eas Stape (x10f)
Enr Stzpa 0c10%)
Baw Stppppt00) (X10))
Dmw Steps (X11P)
Pae Swppse(011000 (XTOF)
Ouex (OSRT)
Ouns (rr.(000PP5aap)
Figure 15: The SafeDreamer training results within 8M environment steps.
24

---

Published as a conference paper at ICLR 2024
C.4.6 VIDEO PREDICTION
Using the past 25 frames context, our world model predicts the next 45 steps in Symmay-Gmm based solely on the given action sequence, without intermediate image access.
Model
fonward Predicion
Safety PointGostt
Trae
Sufety
True
Safety Nointluttint
Modde
selety
PoirtButoo2
Thue
Sefetyy
PointPush1
Moddl
Thue
Safetyy
NointPus12
T=0
T-25
Figure 16: The video predictions in the tasks of the Point agent. In SafetyPointGoal. the model leverages observed goals forecast subsequent ones in future frames. In PafetyyointGoal2, the oncoming rightward navigational movement of the robot to avoid an obstacle is predicted by the model. the SafetyPPintButton1,,, the model predicts the robot's direction toward the green goal. For PaifttyPointuuttoo2,, the model anticipates the robot's 's trajectory. bypassing the yellow sphere on its left. In the SafetyPointPushl, the model foresees the robot's intention to utilize head to mobilize the box. Finally, in Paftty,oin the model discerns the emergence of hitherto unseen crates in future frames, indicating the model's 's prediction ability of environmental transition dynamics.
25

---

Published as a conference paper at ICLR 2024
Model
Fowward Przdittion
BacwcarGeatt
Moddl
Safeey RacocarGoa
1
Model
2
Model
Safety
Safety
RacourrPaht ahd
Figure 17:
The video predictions of the racecar agent. In SafetyRacecargoall,, the world model anticipates the adjustment of agent direction towards a circular obstacle. Similarly, within the SafetyRacecargoal2,, the model predicts the Racecar's incremental deviation from vase. In Safe- IyRacccaBButton,, the world model predicts the Racecar's 's nuanced navigation to avoid right-side obstacle. In SafetyRacecarButtonn the model predicts the Racecar's 's incremental distance toward circular obstacle. In SafetyRacecarPuus and RafetyRacecarPushh tasks, the model predicts the emergence of the box and predicts the Racecar's 's direction towards box, respectively.
26

---

Published as a conference paper at ICLR 2024
C.4.7
SAFEDREAMER VS DREAMERV3
Average Roward Return
30
Average Cost Reeurn
20
40
SafetyPointGoool
20
0.0
0.5
1.0
1.5
2.0
0.0
0.5
1.0
1.5
2.0
Average Cost Return
20
200
SaielyPaintGoul
10-
100
Averaage
0.0
0..4
0.8
1.2
1.6
2.0
0.0
0.4
0.8
1.2
1.6
2.0
Average Reward Return
80
Average Cost Return
20
100
SalitPPinnBtttoon
10-
50-
0.0
0.4
0.8
1.2
1.6
2.0
0.0
0.4
0.8
1.2
1.6
2.0
Average Cost Return
150
20
100
SaketyPolntButton22
20
Average
50
0.0
0.5
1.0
1.
2.0
0.0
0.5
1.0
1.5
200
Average Cost Return
60
PufehyPointuush
10
Average
20
0
0.0
0.5
1.0
1.5
2.0
0.0
2.0
Average Cost Return
Average Reward
50
PufeyyPiitPush2
25
0.0
0.5
1.0
1.5
2.0
Env Steps (x10%)
Env Steps (x100)
SofeDreamer (BSRP-Lag)
DreamerV3
Figure 18: The comparison of SafeDreamer (BSRP-Lag) with Dreamer V3 in the task of Point agent.
27

---

Published as a conference paper at ICLR 2024
Average Rewarrd Return
Average Cost Return
CafetyCarCoalt
20
20
0.0
0.5
1.0
1.5
2.0
0.0
0.5
1.0
1.5
2.0
Average Rewarr Retur
Average Cost Retur
2001
20
Gule
100
10-
0.0
0.55
1.0
1.5
2.0
0.0
0.5
1.0
1.5
2.0
Average Reward Return
Average Cost Return
200
20
SafetyCarButtoni
10-
0.0
0.5
1.0
0.5
1.0
1.5
2.0
Average Reward Return
Average Cost Return
200
20
SafatyCaBButton2
100
0.0
0.5
1.0
1.5
0.5
1.0
1.5
2.0
Averagn Reward Return
15
Average Cost Return
75
10
30
SafetyCarPush1
25
Average Reward Return
Average Cost Return
SarehyCarPush2
0.5
1.5
Eny Steps (x10%)
Env Steps (x10%)
SafeDearrer (BgRP-Lag)
Dramerk3
Figure 19; The comparison of SafeDreamer (BSRP-Lag) with DreamerV3 in the task of Car.
28

---

Published as a conference paper at ICLR 2024
Average Reward Return
30
Average Return
100
.00
SufetyRacecarGaall
50
1.-
0.0
0.5
1.5
2.0
0.0
0.5
1.0
2.0
Average Reward Return
Average Return
600
20
100
SafebfRacccarGaal
200
0-
0
Average Reward Return
.20
150
BafetyRaccearrB.toon
Average Cost
10-
50
0.0
0.5
1.5
2.0
0.5
1.0
2.0
Average Reward Return
Average Return
.0-
150
SafetyRaccccrBB Bu rton2
10-
100
0.0
0.5
1.0
1.5
2.0
0.0
0.5
1.0
1.5
2.0
Average Reward Return
20
Average Return
Safety Pushh
10
40
20
0-
0.0
0.5
1.0
1.5
2.0
0.0
0.5
1.0
1.5
2.0
100
RafeyyRacoca
Average
50
Average
0.0
2.0
0.5
Eny Steps (x10%)
Env Steps (x10%)
Sakoroamo
Figure 20: The comparison of Saeeere (BSRP-Lag) with DreamerV3 in the task of Racecar.
29

---

Published as a conference paper at ICLR 2024
C.5
ABLATION STUDIES
Average Reward Return
Average Cost Return
Cost Rate
30
0.08
0.06
20
0.04
0.00
2
Steps (x110))
Steps (x110)
Steps (x:10)
Cost Limit=1,0
Cost Limit=0,1
Cost Limit=0.01
Figure Ablation studies on the cost limit SafetyPointGoall. We run BSRP-Lag on SafetyPoint- Goall for 8M steps, and we observe that lower cost limit leads to the cost approaching closer to zero. Additionally, we notice that a lower cost limit tends to result decrease in the reward.
Average Reward Return
Average Cost Return
Cost Rate
30
0.34
0.50
10
0.30
0
Enn (x10%)
(x10")
Env (x10%)
Cost Limit=1.0
Cost Limit-0.1 0..1
Cost Limit-0.01
Figure 22: Ablation studies on the cost limit in SafetyPointGoall. We run SafeDreamer (BSRP-Lag) on SafetyPointGoal2 for 4M steps and observe that lower cost limit enables the cost to approach closer zero. This experimental outcomes is consistent with our finding on SafetyPointGoall.
Average Reward
Average Cost Retum
Cost Rate
75-
0.24
50-
01
25
0.08
0.00
0.23
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1.00
0.00
0.25
0..00
0.75
1.00
Env Steps (X110)
Env Steps (x10%)
Bav Steps (x10*))
16 x
32
x
Figure 23: Ablation studies on the selection of the world model's state in Safety Our analysis indicates that in the world model, larger stochastic hidden state 2, corresponds more substantial decrease in cost. This observation highlights the impact of the perception module on the predictive capabilities of the cost model, particularly in complex visual environments.
30

---

Published as a conference paper at ICLR 2024
Average Reward Return
Average Cost Return
Cost Rate
100
0.10
0.12
10
0.08
0.04
-10
0.00
0.75
1.00
0.50
0.75
Env (x10")
Steps (x10")
Steps (X10%)
OSRP-Lag without recomstruction loss
Ours (OSRP-Lag)
Figure 24: Ablation studies on the design and configuration of the SafeDreamer's world model in PointyPoint We observed that upon removing the reconstruction loss, the convergence speed of the reward slowed down, and the cost failed to decrease. This suggests a correlation between the predictive ability of the cost model and the reconstruction capability of the observation. More precise reconstruction of the observation may lead to increased accuracy in the cost model's predictions.
Average Reward Return
Average Cost Return
Cost Rate
0.10
0.08
0.06
10
0.00
0.25
Eav Steps (x10")
Eny Steps (x10*)
Env Steps (x110)
Cost Weight-1.0
Cost Weight-10.0
Cost Weight-30.0
Cost Weight-50.0
Figure 25: Ablation studies on the weight of the cost model loss in SafetyPointGoall. We run SafeDreamer (BSRP-Lag) for 1M steps SafetyPointGoall. We find that applying different weights to the unsafe interactions in the cost model's loss has varying effects on the cost's convergence. higher weight might aid in the cost's reduction. We hypothesize that this effect is due to the unbalanced distribution of cost in the environment. Different weights can mitigate this imbalance, thereby accelerating the convergence of the cost model.
Average Reward Return
Average Cost Return
Cost Rate
6.0
0.10
0.00
3.0-
0.06
0.0
0.5
0.5
0.00
0.5
1.0
1.5
Env Steps (X10*)
Frw Steps (x10
Env Steps (X10*))
1,0
Cost Weight:10.0
Weighr=30.0
Cost Weighr:50.0
Figure 26: Ablation studies on the weight of the cost model loss in SafetyPointGoal2. We run SafeDreamer (BSRP-Lag) for 2M steps SafetyPointGooll. The experimental results were similar to those on SafetyPointGoall, but we suggest fine-tuning this hyperparameter based on the cost distribution in different environments.
31


Reviews:
[{"review_id": "XMyKdyaB12", "reviewer": "Reviewer_fQYU", "summary": "Safe RL method like Lagrangian-based methods hard to achieve zero-cost performance. In this paper, the authors use model-based method, introduce the Dreamer as the world model, and conduct the constrained planning in the latent rollout. The experiment shows in Safety Gym environment the proposed method achieves better reward and lower cost.", "questions": "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?\n\n\nIn test-time, does the BSRP-Lag method still run planning for each real step? Or the world model is only used during training?\n\nHave 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?\n\n\nHave you ran experiments in other safe RL benchmark such as safe autonomous driving environments?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 1, "contribution": 3, "strengths": "The method is straightforward and is promising. Using Lagrangian-based method is proven to be effective but you will need to set a cost threshold and in test-time the cost will be around the threshold when the agent is fully trained. But in this work, by using the world model and trained cost estimator the agent can achieve much higher safety performance due to the internal planning in the latent space.", "weaknesses": "The paper proposes OSPR, OSRP-Lag and BSRP-Lag but the differences between them is not well motivated and presented.\n\n\nIn 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.\n\n\nMixing 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "LZNp0dGccQ", "reviewer": "Reviewer_96af", "summary": "This paper focuses on the problem of safe exploration in in vision-only tasks, where an agent needs to learn to maximise rewards while minimising the number of safety violations during learning. They formalise the problem as a constrained Markov decision process (CMDP), and propose SafeDreamer to address it. SafeDreamer combines prior work in model-based RL (DreamerV3) and constrained cross-entropy methods (CCEM) to achieve near zero-cost performance in continuous robot tasks with state-vector observations and pixel-image partial observations. They conduct various experiments demonstrating that their approach outperforms prior model-based and model-free works on various tasks in the SafetyGym domain.", "questions": "It would be great if the authors could address the major weaknesses I outlined above. I am happy to increase my score if they are properly addressed, as I may have misunderstood pieces of paper. \n\nOn the points requiring additional space to address (e.g. related works), I recommend the authors prioritise which experiments they show in the main paper to save on space. There is a lot of space that can be freed up in the main paper. Another example is Algorithm 1 (and maybe even Algorithm 2 if needed) which can be moved to the appendix.\n\n**POST REBUTTAL________________________________________________________________________________**\n\nThank you to the authors for their very detailed response, and for going the extra mile to run additional experiments to address my concerns regarding the model-free baselines. That's really appreciated. The author's responses in general have greatly helped in clarifying most of my concerns. I also believe that the revised paper is significantly clearer, more transparent, and stronger than before. Hence, I have increased my score from 5 to 6 to reflect my current belief that the strengths of this paper outweigh its weaknesses. I have also increased my confidence level accordingly. Here are a couple of final points I have:\n\n**Contributions**: \"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.\n\n**Baselines**: My concern about apple-to-apple comparisons remains. The model-free baselines should also use CNNs and RNNs, or the model-based ones + 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).\n\n**Emperical significance**: 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 of course beyond the control of the authors though, since omnisafe is a relatively popular baseline in the community.", "limitations": "", "rating": 6, "confidence": 5, "soundness": 2, "presentation": 4, "contribution": 3, "strengths": "- The paper is mostly well-written and investigates an important problem. The proposed approach of combining DreamerV3 with CCEM and Lagrangian SafeRL methods is novel, interesting, and very well-motivated.\n\n- I especially liked the impressive empirical results showing that SafeDreamer is able to achieve near-zero cost returns while still having good reward returns.\n\n- The paper shows numerous empirical results comparing SafeDreamer with many relevant model-based and model-free baselines, in a number of tasks in the SafetyGym domain. Impressively, SafeDreamer is shown to outperform all these baselines in terms of cost returns, while still achieving decent reward returns.", "weaknesses": "**MAJOR**\n\n1- **Contributions**:\n  - 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.\n\n  -  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.  \n\n2- **Emperical significance**:\n  - The paper doesn't state how many runs the results are averaged over. \n  - 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 where chosen to be as fair as possible to the baselines.\n  - These make it hard to evaluate the significance of the empirical results.\n\n3- **Baselines**:\n  - 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. I believe this puts into question the validity of all of the empirical results.\n\n  - 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 also not stated in the paper, so I can only guess). 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.\n\n  - 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. \n\n4- **Related works**:\n  - 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. \n\n  - There are missing relevant related works. Given that the aim of this paper is to achieve near-zero cost returns, I think the following works are relevant: Saute RL (model free) [1], ROSARL (model free) [2], STS-MBRL (model based) [3]\n\n\n[1] Sootla et al, Saute RL: Almost Surely Safe Reinforcement Learning Using State Augmentation, 2022\n\n[2] Tasse et al, ROSARL: Reward-Only Safe Reinforcement Learning, 2023\n\n[3] Kwanda et al, Safe Trajectory Sampling in Model-based RL, 2023\n\n\n**MINOR**\n\n- Figure 1 with its discussion is out of place in the introduction. I think it can be removed.\n- Figure 2 is hard to understand where it is. Maybe it can be moved a few pages below to where the methods are being described.\n- 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.\n- Are there multiple world models being used in SafeDreamer (I assume not)? If not, remove the \"s\" in world models where unnecessary (e.g Fig 3 caption).\n- There are a couple of grammar errors throughout the paper that make it hard to parse. I have listed some examples here:\n  - Page 1 \"proffer\"\n  - Fig 3 caption \"using policy\"\n  - Page 4 \"This notation offers our model components:\"\n  - Page 7 \"How does SafeDreamer against safe model-based algorithms?\"", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "FtZCcn3pAz", "reviewer": "Reviewer_7enF", "summary": "The paper introduces SafeDreamer, an algorithm addressing the challenge of safe reinforcement learning (SafeRL). SafeDreamer integrates Lagrangian-based methods into world model planning processes within the Dreamer framework. The approach achieves nearly zero-cost performance on various tasks, including vision-only scenarios, within the Safety-Gymnasium benchmark. SafeDreamer's contributions include online safety-reward planning, integration of Lagrangian methods, and balancing long-term rewards and costs within world models.", "questions": "1. I might have missed some details. But why the blues in Figure 5 is not complete? The robustness of converged performance is worth reporting. Also, I am curious to see whether OSRP-Lag is able to converge to zero costs.\n\n2. 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.\n\n3. Considering the computational demands of SafeDreamer, how scalable is the approach to more complex tasks or larger-scale environments? Additionally, what measures have been taken to enhance training efficiency, especially for tasks with high-dimensional sensory input? \n\n4. 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.\n\nI would be happy to raise my scores if my concerns can be solved.", "limitations": "", "rating": 8, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "1. This is a well-written paper and I enjoyed reading it. Though some claims need further assessment, It is easy to follow in general. The paper provides a detailed explanation of the SafeDreamer algorithm, including its components and integration methods. The inclusion of algorithms, figures, and experimental results demonstrates a high level of technical rigor and quality in the research.\n\n2. SafeDreamer effectively balances long-term rewards and costs, a crucial aspect of SafeRL. The ability to achieve high rewards while minimizing costs, especially in vision-only tasks, highlights the algorithm's practical utility and robustness.", "weaknesses": "1. 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.\n\n[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.\n\n[2] He, T., Zhao, W., & Liu, C. (2023). Autocost: Evolving intrinsic cost for zero-violation reinforcement learning. arXiv preprint arXiv:2301.10339.\n\n2. SafeDreamer is a combination of Dreamer and Lagrangian methods. The technique contribution is incremental in that sense. But considering the experiment part is solid and sound, that should be fine. This paper provides the safeRL community with another perspective on solving the zero costs issues where safeRL achieving near zero costs is difficult when it comes to sparse cost functions, and that's where planning-based methods could make a difference.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ARYc5xsWGg", "reviewer": "Reviewer_PSGc", "summary": "The authors propose SafeDreamer - a safe model-based reinforcement learning (RL) algorithm built on the DreamerV3 code base. The author's propose three variants of SafeDreamer, one that implements the constrained cross entropy method (CCEM), a simple model-predictive control (MPC) algorithm that filters high quality trajectories and updates the action selection parameters with each iteration to converge to a set of safe trajectories. The other two variants implement Lagrangian methods. OSRP-Lag implements PID Lagrangian for constrained online planning in a similar fashion as CCEM. The final variant implements the augmented Lagrangian, which directly modifies the policy gradient to include cost TD-lambda estimates, weighted by a Lagrange multiplier. This final version called BSRP-Lag does not use online planning at each step in comparison to the two other variants of SafeDreamer.\n\nThe authors evaluate their algorithms on five separate environments from Safety-Gymnasium - a newer version of Safety Gym which is a well known benchmark. The authors evaluate against several suitable (model-free and model-based) algorithms and demonstrate good performance.", "questions": "Is there room for improvement of your algorithm? For example, are there other online planning algorithms you may consider and would you expect these to perform better or worse?\n\nSimilarly, for the policy optimisation would you consider an alternative to the Augmented Lagrangian?\n\nWhat would you expect the outcome to be if you incorporated Augmented Lagrangian policy optimisation with an online (constrained) planning algorithm, either PID Lagrangian or CCEM? Would you expect this to improve the performance of SafeDreamer or would it become overly conservative?\n\nWould you consider incorporating the Bayesian method from (As et al. 2022)? And how would you expect this to affect peformance?\n\nIn your implementation, is there a reason for the cost predictor head being a real valued predictor rather than a binary classifier? And what is the effect of scaling the costs by 10.0 (line 208 agent.py), is this necessary to learn a good cost predictor?", "limitations": "", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "(1) The problem is well motivated (if quite briefly).  \n\n(2) The contributions are novel and the strengths of all three algorithms are demonstrated soundly on a \"popular\" set of environments. \n\n(3) Codebase provided is useful and the supplementary material is very thorough.", "weaknesses": "While the quality and significance of the results are without a doubt the strongest point of this paper. The paper unfortunately has several weaknesses. \n\n(1) A key problem is there is little dicussion 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 realworld. And perhaps discuss and motivate why they only consider simple cumulative constraints and even critique this setup and discuss its limitations. \n\n(2) 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. \n\n(3) Discussion of results. The discusison of results is limited, it is clear all methods achieve the 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 adress 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. \n\n*Some references are wrong in section 5.2 these need fixing*\n\nFor these reasons I am (weakly) rejecting the paper. However, since the results are actually quite good, if the authors were to adress some of my concerns I would consider changing my opinion.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)