# Approximate Model-Based Shielding
for Safe Reinforcement Learning

Alexander W. Goodall Corresponding Author. Email<EMAIL_ADDRESS>Francesco Belardinelli  Imperial College London

###### Abstract

Reinforcement learning (RL) has shown great potential for solving complex
tasks in a variety of domains. However, applying RL to safety-critical systems
in the real-world is not easy as many algorithms are sample-inefficient and
maximising the standard RL objective comes with no guarantees on worst-case
performance. In this paper we propose approximate model-based shielding
(AMBS), a principled look-ahead shielding algorithm for verifying the
performance of learned RL policies w.r.t. a set of given safety constraints.
Our algorithm differs from other shielding approaches in that it does not
require prior knowledge of the safety-relevant dynamics of the system. We
provide a strong theoretical justification for AMBS and demonstrate superior
performance to other safety-aware approaches on a set of Atari games with
state-dependent safety-labels.

##  1 Introduction

Due to the inefficiencies of deep reinforcement learning (RL) and lack of
guarantees, safe RL [5] has emerged as an increasingly active area of
research. Ensuring the safety of RL agents is crucial for their widespread
adoption in safety-critical applications where the cost of failure is high.
From formal verification, _Shielding_ [3] has been developed as an approach
for safe RL, that comes with strong safety guarantees on the agents
performance. However, classical shielding approaches make quite restrictive
assumptions which limit their capabilities in real-world and high-dimensional
tasks. To this end, we propose approximate model-based shielding (AMBS) to
address these limitations and obtain a more general and widely applicable
algorithm.

The constrained Markov decision process (CMDP) [4] is a popular way of framing
the safe RL paradigm. In addition to maximising reward, the agent must satisfy
a set of safety-constraints encoded as a cost function that penalises unsafe
actions or states. In effect, this formulation constrains the set of feasible
policies, which results in a tricky non-smooth optimisation problem. In high-
dimensional settings, several model-free methods have been proposed based on
trust-region methods [1, 45, 30] or Lagrangian relaxations of the constrained
optimisation problem [37, 11]. Many of these methods rely on the assumption of
convergence to obtain any guarantees.

Furthermore, model-based approaches to safe RL have gained increasing
traction, in part due to recent significant developments in model-based RL
(MBRL) [18, 21, 22] and the superior sample complexity of model-based
approaches [19, 27]. In addition to learning a reward maximising policy, MBRL
learns an approximate dynamics model of the system. Approximating the dynamics
using Gaussian process (GP) regression [42] or ensembles of neural networks
are both principled ways to quantify uncertainty, and develop risk- and
safety-aware policy optimisation algorithms [41, 32, 6] and model predictive
control (MPC) schemes [31].

Shielding for RL was introduced in [3] as a correct by construction reactive
system that forces hard constraints on the learned policy, preventing it from
entering unsafe states determined by a given temporal logic formula. Classical
approaches to shielding [3] require a priori access to a safety-relevant
abstraction of the environment, where the dynamics are known or at least
conservatively estimated. This can be quite a restrictive assumption in many
real-world scenarios, where such an abstraction is not known or too complex to
represent. However, these strong assumptions do come with the benefit of
strong guarantees (i.e., _correctness_ and _minimal-interference_).

In this paper we operate in a less restrictive domain and only assume that
there exists some expert labelling of the states; we claim that this is a more
realistic paradigm which has been studied in previous work [24, 16]. As a
result, our approach is more broadly applicable to a variety of environments,
including high-dimensional environments which have been rarely studied in the
shielding literature [34]. We do unfortunately loose the strict formal
guarantees obtained by classical shielding, although we develop tight
probabilistic bounds and a strong theoretical basis for our method.

Our approach, AMBS, is inspired by _latent shielding_ [24] an approximate
shielding algorithm that verifies policies in the latent space of a world
model [19, 21, 22] and is closely related to previous work on _approximate
shielding of Atari agents_ [16]. The core idea of our approach is that we
remove the requirement for a suitable safety-relevant abstraction of the
environment by learning an approximate dynamics model of the environment. With
the learned dynamics model we can simulate possible future trajectories under
the learned policy and estimate the probability of committing a safety-
violation in the near future. In contrast to previous approaches, we provide a
stronger theoretical justification for utilising world models as a suitable
dynamics model. And while we leverage DreamerV3 [22] as our stand-in dynamics
model, we propose a more general purpose and model-agnostic framework for
approximate shielding.

#### Contributions.

Our main contributions are as follows: (1) we formalise the general AMBS
framework which captures previous work on latent shielding [24] and shielding
Atari agents [16]. (2) We formalise notions such as _bounded safety_ [15] in
Probabilistic Computation-tree Logic (PCTL), a principled specification
language for reasoning about the temporal properties of discrete stochastic
systems. (3) We provide PAC-style probabilistic bounds on the probability of
accurately estimating a constraint violation with both the ‘true’ dynamics
model and a learned approximation. (4) We more rigorously develop the theory
from [16] and theoretically justify the use of world models as the stand-in
dynamics model. (5) We provide a much richer set of results on a small set of
Atari games with state-dependent safety-labels and we demonstrate that AMBS
significantly reduces the cumulative safety-violations during training
compared to other safety-aware approaches. And in some cases AMBS also vastly
improves the learned policy w.r.t. episode return. All technical proofs are
provided in full in the Appendix, along with extended results (learning
curves) and implementation details, including hyperparameters and access to
code. These materials are also provided online at
https://github.com/sacktock/AMBS.

##  2 Preliminaries

In this section we introduce the relevant background material and notation
required to understand our approach and the main theoretical results of the
paper. We start by formalising the problem setup and the specification
language used to define bounded safety. We then introduce prior look-ahead
shielding approaches and world models.

###  2.1 Problem Setup

In our experiments we evaluate our approach on Atari games provided by the
Arcade Learning Environment (ALE) [8]. Atari games are partially observable,
as the agent is only given raw pixel data and does not have access to the
underlying memory buffer of the emulator. This means one observation can map
to many underlying states by a probability distribution. Therefore, we model
the problem as a partially observable markov decision process (POMDP) [35].

To capture state-dependent safety-labels, we extend the POMDP tuple with a set
of atomic propositions and a labelling function; a common formulation used in
[7]. Formally, we define a POMDP as a 10-tuple
$\mathcal{M}=(S,A,p,\iota_{init},R,\gamma,\Omega,O,AP,L)$, where, $S$ is a
finite set of states, $A$ is a finite set of actions, $p:S\times A\times
S\rightarrow[0,1]$ is the transition function, where $p(s^{\prime}\mid s,a)$
is the probability of transitioning to state $s^{\prime}$ by taking action $a$
in state $s$, $\iota_{init}:S\rightarrow[0,1]$ is the initial state
distribution, where $\iota_{init}(s)$ is the probability of starting in state
$s$, $R:S\times A\rightarrow\mathbb{R}$ is the reward function, where $R(s,a)$
is the immediate reward received for taking action $a$ in state $s$,
$\gamma\in(0,1]$ is the discount factor, $\Omega$ is a finite set of
observations, $O:S\times\Omega\rightarrow[0,1]$ is the observation function,
where $O(o\mid s)$ is the probability of observing $o$ in state $s$, $AP$ is a
finite set of atomic propositions (or atoms), $L:S\rightarrow 2^{AP}$ is the
labeling function, where $L(s)$ is the set of atoms that hold in state $s$.

We note here that MDPs are a special case of POMDPs without partial
observability. For example, we can suppose that in MDPs, $\Omega=S$ and the
observation function $O:S\times\Omega\rightarrow[0,1]$ collapses to the
identity relation.

In addition, we are given a propositional safety-formula $\Psi$ that encodes
the safety-constraints of the environment. For example, a simple $\Psi$ might
look like,

$\Psi=\neg\textbf{collision}\land(\textbf{red-light}\Rightarrow\textbf{stop})$

which says “don’t have a collision and stop when there is a red light”, for
$AP=\\{\textbf{collision},\textbf{red-light},\textbf{stop}\\}$.

At each timestep $t$, the agent receives an observation $o_{t}$, reward
$r_{t}$ and a set $L(s_{t})$ of labels. The agent can then determine that
$s_{t}$ satisfies the safety-formula $\Psi$ by applying the following
relation,

$\begin{array}[]{@{}r@{{}\mathrel{}}c@{\mathrel{}{}}l@{}}s\models
a&\text{iff}&a\in L(s)\\\ s\models\neg\Psi&\text{iff}&s\not\models\Psi\\\
s\models\Psi_{1}\land\Psi_{2}&\text{iff}&s\models\Psi_{1}\text{ and
}s\models\Psi_{2}\end{array}$

where $a\in AP$ is an atomic proposition, negation ($\neg$) and conjunction
($\land$) are the familiar logical operators from propositional logic. The
goal is to find a policy $\pi$ that maximises reward, that is
$\pi^{*}=\arg\max_{\pi}\mathbb{E}[\sum_{t=0}^{\infty}\gamma^{t}R(s_{t},\pi(s_{t}))]$,
while minimising the cumulative number of violations of the safety-formula
$\Psi$ during training and deployment.

###  2.2 Probabilistic Computation-tree Logic

Probabilistic Computation Tree Logic (PCTL) is a branching-time temporal logic
based on CTL that allows for probabilistic quantification along path formula
[7]. PCTL is a convenient way of stating soft reachability and safety
properties of discrete stochastic systems which makes it a commonly used
specification language for probabilistic model checkers. A well-formed PCTL
formula can be constructed with the following grammar,

$\displaystyle\Phi::=$ $\displaystyle\mid
a\mid\neg\Phi\mid\Phi\land\Phi\mid\mathbb{P}_{J}(\phi)$ $\displaystyle\phi::=$
$\displaystyle X\Phi\mid\Phi U\Phi\mid\Phi U^{\leq n}\Phi$

where $J\subset[0,1]$, $J\neq\varnothing$ is a non-empty subset of the unit
interval, and next $X$, until $U$, and bounded until $U^{\leq n}$ are the
temporal operators from CTL [7]. We make the distinction here between state
formula $\Phi$ and path formula $\phi$ which are interpreted over states and
paths respectively.

For state formula $\Phi$, we write $s\models\Phi$ to indicate that the state
$s\in S$ satisfies the state formula $\Phi$, where the satisfaction relation
is defined in a similar way as before (Section 2.1), but also includes
probabilistic quantification, see [7] for details. On the other hand, path
formula $\phi$ are interpreted over sequences of states or _traces_. In the
following section we provide the satisfaction relation for path formula $\phi$
for the fragment of PCTL that we use. We also note that the common temporal
operators eventually $\lozenge$ and always $\square$, and their bounded
counterparts $\lozenge^{\leq n}$ and $\square^{\leq n}$ can be derived from
the grammar above in a straightforward way, see [7] for details.

###  2.3 Bounded Safety

Now we are equipped to formalise the notion of _bounded safety_ introduced in
[15]. Consider a fixed (stochastic) policy $\pi:O\times A\to[0,1]$ and a POMDP
$\mathcal{M}=(S,A,p,\iota_{init},R,\gamma,\Omega,O,AP,L)$. Together $\pi$ and
$\mathcal{M}$ define a transition system $\mathcal{T}:S\times S\to[0,1]$ on
the set of states $S$, where $\sum_{s^{\prime}\in S}\mathcal{T}(s^{\prime}\mid
s)=1$. A finite trace of the transition system $\mathcal{T}$ with length $n$
($n$ transitions) is a sequence of states $s_{0}\to s_{1}\to\ldots\to s_{n}$
denoted by $\tau$, where each state transition is induced by the transition
probabilities of $\mathcal{T}$. Let $\tau[i]$ denote the $i^{\textrm{th}}$
state of $\tau$. A trace $\tau$ is said to satisfy bounded safety if all of
its states satisfy the given propositional safety-formula $\Psi$, that is,

$\displaystyle\tau\models\square^{\leq n}\Psi$ iff $\displaystyle\forall
i\;0\leq i\leq n,\tau[i]\models\Psi$

for some bounded look-ahead parameter $n$. In words, the path formula
$\phi=\square^{\leq n}\Psi$ says “always satisfy $\Psi$ in the next $n$
steps”. Now in PCTL we say that a given state $s\in S$ satisfies
$\Delta$-bounded safety or formally $s\models\mathbb{P}_{\geq
1-\Delta}(\square^{\leq n}\Psi)$ iff,

$\mu_{s}(\\{\tau\mid\tau[0]=s,\forall i\;0\leq i\leq
n,\tau[i]\models\Psi\\})\in[1-\Delta,1]$ (1)

where $\mu_{s}$ is a well-defined probability measure (induced by the
transition system $\mathcal{T}$) on the set of traces starting from $s$ and
with finite length $n$, see [7] for additional details on the semantics of
PCTL.

Throughout the rest of the paper we will denote $\mu_{s\models\phi}$ as
shorthand for the measure $\mu_{s}(\\{\tau\mid\tau[0]=s,\forall i\;0\leq i\leq
n,\tau[i]\models\Psi\\})$, where $\phi=\square^{\leq n}\Psi$ is the path
formula that corresponds to bounded safety. By grounding bounded safety in
PCTL we obtain a principled way to trade-off safety and exploration with the
$\Delta$ parameter. Since strictly enforcing bounded safety can lead to overly
conservative behaviour during training, which is not always desirable if we
are to make progress toward the optimal policy.

###  2.4 Look-ahead Shielding

The general principle of look-ahead shielding is as follows, from the current
state we compute the set of possible future paths and check if they incur any
unsafe states [15]. If the agent proposes an action that leads the agent down
a path with an unavoidable unsafe state, then the look-ahead shielding
procedure should override the agent’s action with an action leading down a
safe path instead, see Fig. 1.

Figure 1: With a Manhattan distance look-ahead of 2 (blue squares) the robot
can avoid the left-most vat of acid, but is unable to determine that the
conveyor belt leads to an unavoidable unsafe state. Without a further look-
ahead horizon, the robot is doomed to fall in the right-most vat of acid
during exploration. _This image was created with the assistance of DALL·E 2._

Bounded Prescience Shielding (BPS) [15] is an approach to shielding Atari
agents that assumes access to a black-box simulator of the environment, which
can be queried for look-ahead shielding. In the worst case BPS must enumerate
all paths of length $H$ from the current state, to find a safe path for the
agent to follow. This of course comes with a substantial computational
overhead, even with a relatively small horizon of $H=5$.

Our approach most closely resembles _latent shielding_ , [24] which rolls out
learned policies in the latent space of a _world model_ [18, 19] that
approximates the ‘true’ dynamics of the environment for look-ahead shielding.
By simulating futures in a low-level latent space rather than with a high-
fidelity simulator, latent shielding can be used with a much larger look-ahead
horizon ($H=15$) and can be deployed during training. As a result we can hope
to detect unavoidable unsafe states that require a larger look-ahead horizon
to avoid.

Unfortunately _latent shielding_ [24] is not the be all and end all and is
still fairly short-sighted in comparison to our approach. Rolling out dynamics
models for long horizons is impractical as we will quickly run into
accumulating errors that arise from the approximation error of the learned
model. Instead we use _safety critics_ for bootstrapping the end of simulated
trajectories, to obtain further look-ahead capabilities without running into
accumulating errors.

###  2.5 World Models

_World models_ were first introduced in the titular paper [18]. Built on the
recurrent state space model (RSSM) [20], world models are trained to learn a
compact latent representation of the environment state and dynamics. Once the
dynamics are learnt the world model can be used to ‘imagine’ possible future
trajectories.

We leverage DreamerV3 [22] as our stand-in dynamics model for look-ahead
shielding and policy optimisation. DreamerV3 consists of the following
components: (1) the image encoder $z_{t}\sim q_{\theta}(z_{t}\mid
o_{t},h_{t})$ that learns a posterior latent representation given the
observation $o_{t}$ and recurrent state $h_{t}$. (2) The recurrent model
$h_{t}=f_{\theta}(h_{t-1},z_{t-1},a_{t-1})$, which computes the next
deterministic latents given the past state $s_{t-1}=(h_{t-1},z_{t-1})$ and
action $a_{t-1}$. (3) The transition predictor $\hat{z}_{t}\sim
p_{\theta}(\hat{z}_{t}\mid h_{t})$, which is used as the prior distribution
over the stochastic latents (without $o_{t}$). (4) The image decoder
$\hat{o}_{t}\sim p_{\theta}(\hat{o}_{t}\mid h_{t},z_{t})$ which is used to
provide high quality image gradients by minimising reconstruction loss. (5)
The prediction heads $\hat{r}_{t}\sim p_{\theta}(\hat{r}_{t}\mid h_{t},z_{t})$
and $\hat{\gamma}_{t}\sim p_{\theta}(\hat{\gamma}_{t}\mid h_{t},z_{t})$
trained to predict reward signals and episode termination respectively.

In DreamerV3 the latents represent categorical variables and the reward and
image prediction heads both parameterise a _twohot symlog_ distributions [22].
All components are implemented as neural networks and optimised with straight
through gradients. In addition, KL-balancing [21] and free-bits [19] are used
to regularise the prior and posterior representations and prevent a degenerate
latent space representation.

Policy optimisation is performed entirely on experience ‘imagined’ in the
latent space of the world model. We refer to $\pi^{\text{task}}$ as the task
policy trained in the world model to maximise expected discounted accumulated
reward, that is, optimise the following,

$\max\mathbb{E}_{\pi^{\text{task}},p_{\theta}}\left[\sum_{t=1}^{\infty}\hat{\gamma}^{t-1}\hat{r}_{t}\right]$
(2)

In addition to the task policy $\pi^{\text{task}}$, a task critic denoted
$v^{\text{task}}$, is used to construct TD-$\lambda$ targets [39] which are
used to help guide the task policy $\pi^{\text{task}}$ in a TD-$\lambda$
actor-critic style algorithm. We refer the reader to [22] for more precise
details.

##  3 Approximate Model-Based Shielding

In this section we present the AMBS framework, our novel method for look-ahead
shielding RL policies with a learned dynamics model. We start by giving an
overview of the approach, followed by a strong theoretical basis for AMBS.
Specifically, we derive PAC-style probabilistic bounds on estimating the
probability of a constraint violation under both the ‘true’ transition system
and a learned approximation. We show these hold under reasonable assumptions
and we demonstrate when these assumptions hold in specific settings.

###  3.1 Overview

AMBS is split into two main phases: the _learning phase_ , which includes (i)
dynamics learning, (ii) task policy optimisation with the standard RL
objective, and (iii) safe policy learning and safety critic optimisation, and
the _environment interaction phase_ , which consists of collecting experience
from the environment with the shielded policy. The new experience is then used
to improve the dynamics model.

To obtain a shielded policy, recall that we are interested in verifying PCTL
formulas of the form $\Phi=\mathbb{P}_{\geq 1-\Delta}(\square^{\leq n}\Psi)$
($\Delta$-bounded safety) on the transition system $\mathcal{T}:S\times
S\to[0,1]$ induced by fixing the task policy $\pi^{\textrm{task}}$ for a given
POMDP $\mathcal{M}$, where $\Psi$ is the propositional safety-formula. Quite
simply, the shielded policy picks actions according to $\pi^{\textrm{task}}$
provided $\Delta$-bounded safety is satisfied, otherwise the shielded policy
picks actions with a backup policy that should prevent the agent from
committing the possible safety-violation.

To check $\Delta$-bounded safety, we can simulate the transition system
$\mathcal{T}$ by rolling-out the learned world model $p_{\theta}$ with actions
sampled from the task policy $\pi^{\textrm{task}}$; effectively obtaining an
approximate transition system $\widehat{\mathcal{T}}:S\times S\to[0,1]$. Even
with access to the ‘true’ transition system $\mathcal{T}$, exact PCTL model
checking is $O(\textrm{poly}(size(\mathcal{T}))\cdot n\cdot|\Phi|)$ which is
too big for high-dimensional settings. Rather, we rely on Monte-Carlo sampling
from the approximate transition system $\widehat{\mathcal{T}}$ to obtain an
estimate $\tilde{\mu}_{s\models\phi}$ of the measure $\mu_{s\models\phi}$,
which specifies the probability of satisfying $\phi=\square^{\leq n}\Psi$
(bounded safety) from $s$, under the ‘true’ transition system $\mathcal{T}$.

###  3.2 Fully Observable Setting

We start by considering the fully observable setting, that is we are concerned
with checking PCTL formula of the form $\Phi=\mathbb{P}_{\geq
1-\Delta}(\square^{\leq n}\Psi)$ ($\Delta$-bounded safety) on the transition
system $\mathcal{T}:S\times S\to[0,1]$ induced by fixing some given policy
$\pi$ in an MDP $\mathcal{M}=(S,A,p,\iota_{init},R,\gamma,AP,L)$. We state the
following result,

######  Theorem .

Let $\epsilon>0$, $\delta>0$, $s\in S$ be given. With access to the ‘true’
transition system $\mathcal{T}$, with probability $1-\delta$ we can obtain an
$\epsilon$-approximate estimate of the measure $\mu_{s\models\phi}$, by
sampling $m$ traces $\tau\sim\mathcal{T}$, provided that,

$m\geq\frac{1}{2\epsilon^{2}}\log\left(\frac{2}{\delta}\right)$ (3)

###### Proof.

The proof is an application of Hoeffding’s inequality [26]. See Appendix B for
details. ∎

Suppose that we only have access to an approximate MDP, where the transition
function $p$ is estimated with a learned approximation $\hat{p}$, that is,
$\widehat{\mathcal{M}}=(S,A,\hat{p},\iota_{init},R,\gamma,AP,L)$. As earlier,
$\pi$ and $\widehat{\mathcal{M}}$ define an approximate transition system
$\widehat{\mathcal{T}}$. We show that we can obtain an estimate
$\tilde{\mu}_{s\models\phi}$ for $\mu_{s\models\phi}$ by sampling traces from
$\widehat{\mathcal{T}}$.

######  Theorem .

Let $\epsilon>0$, $\delta>0$ be given. Suppose that for all $s\in S$, the
total variation (TV) distance between $\mathcal{T}(s^{\prime}\mid s)$111Unless
$s^{\prime}$ is fixed we define $\mathcal{T}(s^{\prime}\mid s)$ as the
distribution of next states $s^{\prime}$ given $s$ rather than as an explicit
probability. and $\widehat{\mathcal{T}}(s^{\prime}\mid s)$ is bounded by some
$\alpha\leq\epsilon/n$. That is,

$D_{\text{TV}}\left(\mathcal{T}(s^{\prime}\mid
s),\widehat{\mathcal{T}}(s^{\prime}\mid s)\right)\leq\alpha\;\forall s\in S$
(4)

Now fix an $s\in S$, with probability $1-\delta$ we can obtain an
$\epsilon$-approximate estimate of the measure $\mu_{s\models\phi}$, by
sampling $m$ traces $\tau\sim\widehat{\mathcal{T}}$, provided that,

$m\geq\frac{2}{\epsilon^{2}}\log\left(\frac{2}{\delta}\right)$ (5)

###### Proof.

The proof follows from Theorem 3 and the _simulation lemma_ [28] adapted to
our purposes, see Appendix B. ∎

The assumption we make in Theorem 5 on the $TV$ distance being bounded for all
states $s\in S$ (Eq. 4) is quite strong. In reality we may only require that
the $TV$ distance between $\mathcal{T}(s^{\prime}\mid s)$ and
$\widehat{\mathcal{T}}(s^{\prime}\mid s)$ is bounded for the safety-relevant
subset of the state space.

#### Tabular Case.

It may also be interesting to consider under what conditions the $TV$ distance
between $\mathcal{T}(s^{\prime}\mid s)$ and
$\widehat{\mathcal{T}}(s^{\prime}\mid s)$ are bounded for a given state. In
the tabular case we can obtain an approximate dynamics model
$\hat{p}(s^{\prime}\mid s,a)$ for the ‘true’ dynamics $p(s^{\prime}\mid s,a)$
by using the maximum likelihood principle for categorical variables. Quite
simply let,

$\hat{p}(s^{\prime}\mid s,a)=\frac{c(s^{\prime},s,a)}{v(s,a)}$ (6)

where $c(s^{\prime},s,a)$ is the visit count of $s^{\prime}$, $s$, $a$ or
equivalently the number of times $(s,a)$ has lead $s^{\prime}$ and $v(s,a)$ is
the visit count of $(s,a)$ or similarly the number of times $a$ has been
picked from $s$.

######  Theorem .

Let $\alpha>0$, $\delta>0$, $s\in S$ be given. With probability $1-\delta$ the
total variation (TV) distance between $\mathcal{T}(s^{\prime}\mid s)$ and
$\widehat{\mathcal{T}}(s^{\prime}\mid s)$ is upper bounded by $\alpha$,
provided that all actions $a\in A$ with non-negligable probability
$\eta\geq\alpha/(|A||S|)$ (under $\pi$) have been picked from $s$ at least $m$
times, where

$m\geq\frac{|S|^{2}}{\alpha^{2}}\log\left(\frac{2|A||S|}{\delta}\right)$ (7)

###### Proof.

The proof is an application of Hoeffding’s inequality [26] to categorical
distributions, see Appendix B. ∎

The strong dependence on $|S|$ here can make this result quite unwieldy.
However, for low-rank or low-dimensional MDPs this dependence on $|S|$ can be
suitably replaced. For example in _gridworld_ environments, where there are
only 4 possible next states, $|S|$ can be replaced with 4. Additionally, for
deterministic policies we can ignore the dependence on $|A|$.

###  3.3 Partially Observable Setting

Sample complexity bounds like the one we obtained for the tabular MDP case are
a lot harder to obtain for general POMDPs. To say anything meaningful about
learning in POMDPs, various assumptions can be made about the structure of the
environment [29]. We will reason about the underlying POMDP $\mathcal{M}$ in
terms of _belief states_. Belief states infer a distribution over possible
underlying states given the history of past observations and actions. Formally
they are defined as a filtering distribution $p(s_{t}\mid o_{t\leq t},a_{\leq
t})$. Importantly, conditioning on the entire history of past observations and
actions exposes more information about the possible underlying states.

However, maintaining an explicit distribution over possible states is
difficult, particularly when we have no idea what the space of possible states
is. Instead it is common to learn a latent representation $b_{t}=f(o_{t\leq
t},a_{\leq t})$ of the history of past observations and actions that captures
the important statistics of the filtering distribution, so that $p(s_{t}\mid
o_{t\leq t},a_{\leq t})\approx p(s_{t}\mid b_{t})$. We will denote $b_{t}$ the
belief state, although it is actually a compact latent state representation of
$o_{t\leq t},a_{\leq t}$, rather than a distribution over possible states.

######  Theorem .

Let $b_{t}$ be a latent representation (belief state) such that $p(s_{t}\mid
o_{t\leq t},a_{\leq t})=p(s_{t}\mid b_{t})$. Let the fixed policy
$\pi(\cdot\mid b_{t})$ be a general probability distribution conditional on
belief states $b_{t}$. Let $f$ be a generic $f$-divergence measure (TV or
similar). Then the following holds:

$D_{f}(\mathcal{T}(s^{\prime}\mid b),\widehat{\mathcal{T}}(s^{\prime}\mid
b))\leq D_{f}(\mathcal{T}(b^{\prime}\mid
b),\widehat{\mathcal{T}}(b^{\prime}\mid b))$ (8)

where $\mathcal{T}$ and $\widehat{\mathcal{T}}$ are the ‘true’ and approximate
transition system respectively, defined now over both states $s$ and belief
states $b$.

###### Proof.

The proof is a direct application of the data-processing inequality [2] and we
note similar bounds have been derived in previous work [14]. See Appendix B
for details. ∎

What Theorem 3.3 says, is that by minimising the divergence between the next
belief state $b^{\prime}$ in the ‘true’ transition system $\mathcal{T}$ and
the approximate transition system $\widehat{\mathcal{T}}$ (given the current
belief $b$), we minimise an upper bound on the divergence between the next
underlying state $s^{\prime}$ in $\mathcal{T}$ and $\widehat{\mathcal{T}}$.
This objective is precisely encoded in the world model loss function of
DreamerV3 [22]. And so by using DreamerV3 (or similar architectures) we aim to
minimise the $TV$ distance between $\mathcal{T}$ and $\widehat{\mathcal{T}}$
in the hope that we may use the results stated earlier in this section.

##  4 AMBS with DreamerV3

In this section we present our practical implementation of AMBS with DreamerV3
[22]. We start by detailing the important components of the algorithm, before
describing the shielding procedure in detail and justifying some of the
important algorithmic decisions. An outline for the full procedure is
presented in Appendix A.

###  4.1 Algorithm Components

#### RSSM with Costs.

We augment DreamerV3 [22] with a cost predictor head $\hat{c}_{t}\sim
p_{\theta}(\hat{c}_{t}\mid h_{t},z_{t})$ used to predict state dependent
costs. Similar to the reward head, the cost predictor is also implemented as a
neural network that parameterises a _twohot symlog_ distribution [22]. Targets
for the cost predictor are constructed as follows,

$c_{t}=\begin{cases}0,&\text{if $s_{t}\models\Psi$}\\\
C,&\text{otherwise}\end{cases}$ (9)

where $s_{t}\models\Psi$ can be determined using the labels $L(s_{t})$
provided at each timestep (see Section 2.1), and $C>0$ is a hyperparameter
that determines the cost of violating the propositional safety-formula $\Psi$.
We claim that using a cost function in this way rather than a binary
classifier, allows the agent to distribute its uncertainty about a constraint
violation over consecutive states, since the predicted cost at a latent state
$(h_{t},z_{t})$ may lie anywhere on the real number line.

#### Safe Policy.

The safe policy denoted $\pi^{\text{safe}}$ is used as the backup policy [23]
if we detect that a safety-violation is likely to occur in the near future
when following the task policy $\pi^{\text{task}}$. The goal of the safe
policy is to prevent the agent from committing the safety-violation that was
‘likely’ to happen in the near future.

Unlike classical shielding methods that synthesise a backup policy ahead of
time [3], we must learn a suitable backup policy as we have no knowledge of
the safety-relevant dynamics of the system a priori. Although, we note that if
such a backup policy is available ahead of time, either because it is easy to
determine (i.e. breaking in a car) or significant engineering effort has gone
into creating a backup policy, then this can be easily integrated into our
framework. However, in general we train a separate policy $\pi^{\text{safe}}$
and critic $v^{\text{safe}}$ using the same TD-$\lambda$ actor-critic style
algorithm used for the task policy [22]. However, we minimise expected
discounted costs and ignore any reward signals entirely and so specifically we
optimise the following objective,

$\min\mathbb{E}_{\pi^{\text{safe}},p_{\theta}}\left[\sum_{t=1}^{\infty}\hat{\gamma}^{t-1}\hat{c}_{t}\right]$
(10)

In theory, by minimising expected costs we should learn a good backup policy,
since by construction we only incur a cost $>0$ when $\Psi$ is not satisfied.

#### Safety Critics.

Safety critics are used to estimate the expected discounted costs under the
state distribution of the task policy. In essence, they give us an idea of how
safe a given state $s\in S$ is if we were to follow the task policy
$\pi^{\text{task}}$ from that state $s$. As a result we can use them to
bootstrap the end of ‘imagined’ trajectories for further look-ahead
capabilities, without having to roll-out the dynamics model any further. We
jointly train two safety critics $v^{C}_{1}$ and $v^{C}_{2}$ with a TD3-style
algorithm [13] to prevent overestimation, which is important for reducing
overly conservative behaviour. We also train a separate safety discount
predictor head $\hat{\gamma}^{\text{safe}}_{t}\sim
p_{\theta}(\hat{\gamma}^{\text{safe}}_{t}\mid h_{t},z_{t})$ to help with
overestimation. Specifically, $v^{C}_{1}$ and $v^{C}_{2}$ estimate the
following quantity,

$V^{C}(s)=\mathbb{E}_{\pi^{\text{task}},p_{\theta}}\left[\sum_{t=1}^{\infty}(\hat{\gamma}^{\text{safe}}_{t})^{t-1}\hat{c}_{t}\>\Big{|}\>s_{0}=s\right]$
(11)

The safety discount head is a binary classifier that predicts safety-
violations, we use it to implicitly transform the MDP into one where violating
states are terminal, this ensures that the safety critics are always upper
bounded by $C$ and any costs incurred after a safety-violation are not
counted. We note that this is an important property to maintain for checking
bounded safety.

###  4.2 Shielding Procedure

During environment interaction we deploy our approximate shielding procedure
to try and prevent the task policy $\pi^{\text{task}}$ from violating the
safety-formula $\Psi$. Actions that lead to violations in the learned world
model $p_{\theta}$ do not count since they are not actually committed in the
‘real world’.

Algorithm 1 details the approximate shielding procedure. It takes as input an
approximation error $\epsilon$, desired safety level $\Delta$ (from the PCTL
formula), the current latent state $\hat{s}=(z,h)$, an action $a$ proposed by
the task policy, the ‘imagination’ horizon $H$, the look-ahead shielding
horizon $T$, and the number of samples $m$, along with the RSSM components and
the learned policies. The shielding procedure works by sampling $m$ traces
from the approximate transition system $\widehat{\mathcal{T}}$ (obtained by
sampling actions with $\pi^{\text{task}}$ in $p_{\theta}$), checking whether
each trace satisfies bounded safety and returning the proportion of satisfying
traces. If the estimate obtained $\tilde{\mu}_{s\models\phi}$ is in the
interval $[1-\Delta+\epsilon,1]$ then the proposed action $a$ is verified
safe, otherwise we pick an action with the safe policy $\pi^{\text{safe}}$
instead.

Algorithm 1 Approximate Model-Based Shielding (AMBS)

Input: approximation error $\epsilon$, desired safety level $\Delta$, current
state $\hat{s}=(z,h)$, proposed action $a$, ‘imagination’ horizon $H$, look-
ahead shielding horizon $T$, number of samples $m$, RSSM $p_{\theta}$, safe
policy $\pi^{\text{safe}}$ and task policy $\pi^{\text{task}}$.
Output: shielded action $a^{\prime}$.

for $i=1,...,m$ do

Let $\hat{s}_{0}=(h,z)$

// Sample trace

From $\hat{s}_{0}$ play $a_{0}$ and sample trace $\tau=\hat{s}_{1:H}$ with
$\pi^{\text{task}}$ and $p_{\theta}$.

Using $\tau=\hat{s}_{1:H}$ compute sequences $\hat{c}_{1:H}$ and
$\hat{\gamma}_{1:H}$ with $p_{\theta}$.

// Check trace is satisfying

$X_{i}=\mathbbm{1}\left[\text{cost}(\tau)<\gamma^{T-1}\cdot C\right]$ using
Eq. 12 or Eq. 13.

end for

Let $\tilde{\mu}_{s\models\phi}=\frac{1}{m}\sum^{m}_{i=1}X_{i}$

If $\tilde{\mu}_{s\models\phi}\in[1-\Delta+\epsilon,1]$ return $a^{\prime}=a$
else return $a^{\prime}\sim\pi^{\text{safe}}$

######  Proposition .

Suppose we have an estimate
$\tilde{\mu}_{s\models\phi}\in[\mu_{s\models\phi}-\epsilon,\mu_{s\models\phi}+\epsilon]$,
if $\tilde{\mu}_{s\models\phi}\in[1-\Delta+\epsilon,1]$ then it must be the
case that $\mu_{s\models\phi}\in[1-\Delta,1]$ and therefore
$s\models\mathbb{P}_{\geq 1-\Delta}(\square^{\leq n}\Psi)$.

###### Proof.

Suppose
$\tilde{\mu}_{s\models\phi}\in[\mu_{s\models\phi}-\epsilon,\mu_{s\models\phi}+\epsilon]$,
$\tilde{\mu}_{s\models\phi}\in[1-\Delta+\epsilon,1]$ and
$\mu_{s\models\phi}\not\in[1-\Delta,1]$. Then
$\tilde{\mu}_{s\models\phi}-\mu_{s\models\phi}>\epsilon$ which contradicts
$\tilde{\mu}_{s\models\phi}\in[\mu_{s\models\phi}-\epsilon,\mu_{s\models\phi}+\epsilon]$.
This implies that indeed $\mu_{s\models\phi}\in[1-\Delta,1]$ and therefore
$s\models\mathbb{P}_{\geq 1-\Delta}(\square^{\leq n}\Psi)$ by Eq. 1. ∎

Proposition 4.2 justifies checking that the condition
$\tilde{\mu}_{s\models\phi}\in[1-\Delta+\epsilon,1]$ holds in Algorithm 1.
Since if $\tilde{\mu}_{s\models\phi}$ is an $\epsilon$-estimate of
$\mu_{s\models\phi}$ then $\tilde{\mu}_{s\models\phi}\in[1-\Delta+\epsilon,1]$
necessarily implies that the current state $s$ satisfies $\Delta$-bounded
safety. We present one more proposition that is required to understand how we
determine (or verify) that a trace satisfies bounded safety.

######  Proposition .

Suppose we have a trace $\tau$ with cost given by,

$\text{cost}(\tau)=\sum^{H}_{t=1}(\hat{\gamma}_{t})^{t-1}\hat{c}_{t}$ (12)

Under the ‘true’ transition system $\mathcal{T}$ if
$\textrm{cost}(\tau)<\gamma^{H-1}\cdot C$ then necessarily
$\tau\models\square^{\leq H}\Psi$.

###### Proof.

The proof is a straightforward argument. By construction $c_{t}=C$ if and only
if $\tau[t]\not\models\Psi$, therefore $\textrm{cost}(\tau)<\gamma^{H-1}\cdot
C$ implies that $\forall t\;1\leq t\leq H$ we have $c_{t}=0$, which implies
that $\forall t\;1\leq t\leq H$ we have $\tau[t]\models\Psi$. ∎

Furthermore, with safety critics we can compute the cost of a trace in a
similar way:

$\textrm{cost}(\tau)=\sum^{H-1}_{t=1}(\hat{\gamma}_{t})^{t-1}\hat{c}_{t}+\min\left\\{v^{C}_{1}(h_{H},\hat{z}_{t}),v^{C}_{2}(h_{H},\hat{z}_{t})\right\\}$
(13)

Provided that the safety critics accurately capture the expected discounted
costs from beyond the ‘imagination’ horizon $H$ then we can check bounded
safety with a larger horizon $T>H$ and derive something similar to Proposition
4.2. In this situation it may be clear now why we want the safety critics to
be upper bounded by $C$. Since a state $s\in S$ that satisfies bounded safety
could have a value $V^{C}(s)>\gamma^{T-1}\cdot C$ if two or more violations
occurred beyond the horizon $T$; further in the future than what we care
about.

##  5 Experimental Evaluation

In this section we present the results of running AMBS equipped with DreamerV3
[22] on a set of Atari games with state-dependent safety-labels [15]. First,
we detail the corresponding safety-formula for each of the games. We then
cover with the training details, followed by the results and a brief
discussion.

#### Atari Games with Safety-labels.

We evaluate each agent on a set of five Atari games provided in the ALE [8].
Following [33] we use sticky actions (action repeated with probability 0.25),
frame skip ($k=4$), and provide no life information unless it is explicitly
encoded in the safety-formula. We leverage the state-dependent safety-labels
provided in prior work [15] to construct the safety-formula, Table 1 outlines
each of the Atari environments used in our experiments. For a more precise
description of each of the environments we refer the reader to [33]. We opt
for this setting, as opposed to more standard safe RL benchmarks like
SafetyGym [37], as it has been used in previous work [15, 16] and it provides
a richer set of safety-constraints and corresponding behaviours that much be
learnt.

Table 1: Atari environments and their corresponding safety-formula Environment |  safety-formula $\Psi$
---|---
Assault | $\neg\textbf{hit}\land\neg\textbf{overheat}$
DoubleDunk | $\neg\textbf{out-of-bounds}\land\neg\textbf{shoot-bf-clear}$
Enduro | $\neg\textbf{crash-car}$
KungFuMaster | $\neg\textbf{loose-life}\land\neg\textbf{energy-loss}$
Seaquest | $(\textbf{surface}\Rightarrow\textbf{diver})\land\neg\textbf{hit}\land\neg\textbf{out-of-oxygen}$

#### Training Details

Each agent is trained on a single Nvidia Tesla A30 (24GB RAM) GPU and a
24-core/48 thread Intel Xeon CPU with 256GB RAM. Due to constraints on compute
resources, each agent is trained for 10M environment interactions (40M frames)
as opposed to the typical 50M [21, 22] and our experiments are only run on one
seed.

To test robustness, all hyperparameters are fixed over all runs for all
agents. The hyperparameters of DreamerV3 for the Atari benchmark are detailed
in [19], most notably the ‘imagination’ horizon $H=15$. The AMBS
hyperparameters are also fixed in all experiments as follows, bounded safety
level $\Delta=0.1$, approximation error $\epsilon=0.09$, number of samples
$m=512$, and look-ahead shielding horizon $T=30$. Additional implementation
details and hyperparameter settings can be found in Appendix C.

#### Results

We evaluate the effectiveness of AMBS equipped with DreamerV3 [22] by
comparing it with the following algorithms, vanilla DreamerV3 without any
shielding procedure or access to the cost function, a safety-aware DreamerV3
that implements an augmented Lagrangian penalty framework (LAG) [43, 37],
since this method is based on DreamerV3 it should act as a more meaningful
baseline compared to similar model-free counterparts like PPO- and TRPO-
Lagrangian [37] that are not optimised for Atari.

In addition, we provide results for two state-of-the-art (single GPU) model-
free algorithms, Rainbow [25] and Implicit Quantile-Network (IQN) [12]. While
these model-free algorithms don’t really provide a meaningful baseline, we
include them (to the same end as [15]) to demonstrate that state-of-the-art
Atari agents frequently violate quite straightforward safety properties that
we would expect them to satisfy. We also note that both BPS [15] and classical
shielding [3] are too unwieldy to be used during training in this setting, and
so we cannot obtain a meaningful comparison between these approaches and AMBS.

We compare the five algorithms by recording the cumulative number of
violations and the best episode score obtained during a single run of 10M
environment interactions (40M frames). Table 2 presents the results. We also
provide learning curves for each of the agents in Appendix D.

Table 2: Best episode scores and total violations for the five Atari games. Blank ($-$) denotes that the agent failed to converge. |  | DreamerV3 | DreamerV3 (AMBS) | DreamerV3 (LAG) | IQN | Rainbow
---|---|---|---|---|---|---
Assault | Best Score $\uparrow$ | 14738 | 44467 | 19832 | 9959 | 9632
# Violations $\downarrow$ | 18745 | 12638 | 16802 | 24462 | 24019
DoubleDunk | Best Score $\uparrow$ | 24 | 24 | 24 | 24 | -
# Violations $\downarrow$ | 877499 | 66248 | 359018 | 188363 | -
Enduro | Best Score $\uparrow$ | 2369 | 2367 | 2365 | 2375 | 2383
# Violations $\downarrow$ | 167933 | 132147 | 174217 | 129012 | 108000
KungFuMaster | Best Score $\uparrow$ | 97000 | 117200 | 97200 | 51600 | 59500
# Violations $\downarrow$ | 427476 | 10936 | 567559 | 284909 | 612762
Seaquest | Best Score $\uparrow$ | 4860 | 145550 | 1940 | 34150 | 1900
# Violations $\downarrow$ | 73641 | 40147 | 64679 | 53516 | 67101

#### Discussion.

As we can see in Table 2, our approach (AMBS) compared with the other
DreamerV3 [22] based agents dramatically reduces the total number of safety-
violations during training in all five of the Atari games. In terms of reward,
AMBS also achieves comparable or better best episode scores in all five of the
Atari games. This is likely because the objective of maximising rewards and
minimising constraints are suitably aligned, which we note might not always be
the case. LAG fails to reliably reduce the total number of safety-violations
in all the Atari games, which could suggest this method is more sensitive to
hyperparameter tuning. Clearly the extra machinery introduced by AMBS helps us
to obtain a dynamic algorithm that can effectively switch between reward
maximising and constraint satisfying policies, which on this set of Atari
games demonstrates a significant improvement in performance compared to other
safety-aware and state-of-the-art algorithms. Interestingly, the model-free
algorithms do better across the board on _Enduro_ , it is not immediately
clear why this is the case and this result may need further investigation.
Perhaps it is the case that DreamerV3 is not quite as suited to Enduro as IQN
or Rainbow, although we note that AMBS improves on DreamerV3 w.r.t. safety-
violations while maintaining comparable performance.

##  6 Related Work

#### Model-based RL.

Model-based RL as a paradigm for learning complex policies has become
increasingly popular in recent years due to its superior sample efficiency
[19, 27]. Dyna – “an integrated architecture for learning, planning and
reacting” [38] proposed an architecture that learns a dynamics model of the
environment in addition to a reward maximising policy. In theory, by utilising
the dynamics model for planning and/or policy optimisation we should learn
better policies with fewer experience. Model-based policy optimisation (MBPO)
[27] is a sample-efficient neural architecture for optimising policies in
learned dynamics model. More recently, more sophisticated neural architectures
such as Dreamer [19] and DreamerV2 [21] have been proposed that have
demonstrated state-of-the-art performance on the DeepMind Visual Control Suite
[40] and the Atari benchmark [8, 33] respectively. Our work is built on
DreamerV3 [22] which has demonstrated superior performance in a number of
domains (with fixed hyperparameters) including the Atari benchmark [8, 33] and
MineRL [17], which is regarded as a notoriously difficult exploration
challenge in the RL literature.

#### Shielding.

Shield synthesis for RL was first introduced in [3] as a method for preventing
learned policies from entering an unsafe region of the state space defined by
a given temporal logic formula. In contrast to other safety-aware approaches
based on soft penalties [1, 41], shielding enforces hard constraints by
directly overriding unsafe actions with a verified backup policy. The reactive
shield, which includes the backup policy, is computed ahead of time to ensure
no training violations (_correctness_). To compute the reactive shield a
safety automaton is constructed from a safety-relevant abstraction of the game
with known dynamics and a safety game [9] is then solved.

_Look-ahead shielding_ [15, 24, 44, 16] is a more dynamic approach to
shielding that aims to alleviate some of the restrictive requirements that
come with classical shielding. The key benefits of look-ahead shielding
include, (1) better computational efficiency as only the reachable subset of
the state space needs to be checked, (2) can be applied online in real-time
from the current state. Although we note that look-ahead shielding can be done
ahead of time [44] similar to classical shielding. Inspired by _latent
shielding_ [24] our approach can very much be categorised as an online
(approximate) look-ahead shielding approach.

##  7 Conclusions

In this paper we introduce and describe approximate model-based shielding
(AMBS), a general purpose framework for shielding RL policies by simulating
and verifying possible futures with a learned dynamics model. In contrast with
the majority of work on shielding [3], we apply our approach in a much less
restrictive setting, where we only require access to an expert labelling of
the states, rather than having the safety-relevant dynamics be known. In
addition, we obtain further look-ahead capabilities in comparison to previous
work, such as _latent shielding_ [24] and BPS [15], by utilising safety
critics that are trained to predict expected costs.

Compared with other safety-aware approaches, we empirically show that
DreamerV3 [22] augmented with AMBS demonstrates superior performance in terms
of episode return and cumulative safety-violations, on a set of Atari games
with state-dependent safety-labels. We also develop a rigorous set of
theoretical results that underpin AMBS, which includes strong probabilistic
guarantees in the tabular setting. We stress the importance of these results,
as it allows operators to meaningfully choose hyperparameter settings based on
their specific requirements w.r.t. safety-guarantees.

All the components of our algorithm must be learnt from scratch, which
unfortunately means that there are few guarantees during early stage training.
Although in principle we could jump start the learning process with offline or
expert data, this would be an interesting paradigm to explore. Important
future work also includes, empirically verifying our theoretical results for
the tabular setting and applying AMBS to more standard safe RL benchmarks,
like SafetyGym [37], that are not frequently used in the shielding literature.

This work was supported by UK Research and Innovation [grant number
EP/S023356/1], in the UKRI Centre for Doctoral Training in Safe and Trusted
Artificial Intelligence (www.safeandtrustedai.org).

## References

  * [1] Joshua Achiam, David Held, Aviv Tamar, and Pieter Abbeel, ‘Constrained policy optimization’, in International conference on machine learning, pp. 22–31. PMLR, (2017).
  * [2] Syed Mumtaz Ali and Samuel D Silvey, ‘A general class of coefficients of divergence of one distribution from another’, Journal of the Royal Statistical Society: Series B (Methodological), 28(1), 131–142, (1966).
  * [3] Mohammed Alshiekh, Roderick Bloem, Rüdiger Ehlers, Bettina Könighofer, Scott Niekum, and Ufuk Topcu, ‘Safe reinforcement learning via shielding’, in Proceedings of the AAAI Conference on Artificial Intelligence, volume 32, (2018).
  * [4] Eitan Altman, Constrained Markov decision processes: stochastic modeling, Routledge, 1999.
  * [5] Dario Amodei, Chris Olah, Jacob Steinhardt, Paul Christiano, John Schulman, and Dan Mané, ‘Concrete problems in ai safety’, arXiv preprint arXiv:1606.06565, (2016).
  * [6] Yarden As, Ilnura Usmanova, Sebastian Curi, and Andreas Krause, ‘Constrained policy optimization via bayesian world models’, arXiv preprint arXiv:2201.09802, (2022).
  * [7] Christel Baier and Joost-Pieter Katoen, Principles of model checking, MIT press, 2008.
  * [8] M. G. Bellemare, Y. Naddaf, J. Veness, and M. Bowling, ‘The arcade learning environment: An evaluation platform for general agents’, Journal of Artificial Intelligence Research, 47, 253–279, (jun 2013).
  * [9] Roderick Bloem, Bettina Könighofer, Robert Könighofer, and Chao Wang, ‘Shield synthesis: Runtime enforcement for reactive systems’, in International conference on tools and algorithms for the construction and analysis of systems, pp. 533–548. Springer, (2015).
  * [10] Pablo Samuel Castro, Subhodeep Moitra, Carles Gelada, Saurabh Kumar, and Marc G. Bellemare, ‘Dopamine: A Research Framework for Deep Reinforcement Learning’, (2018).
  * [11] Yinlam Chow, Ofir Nachum, Aleksandra Faust, Edgar Duenez-Guzman, and Mohammad Ghavamzadeh, ‘Lyapunov-based safe policy optimization for continuous control’, arXiv preprint arXiv:1901.10031, (2019).
  * [12] Will Dabney, Georg Ostrovski, David Silver, and Rémi Munos, ‘Implicit quantile networks for distributional reinforcement learning’, in International conference on machine learning, pp. 1096–1105. PMLR, (2018).
  * [13] Scott Fujimoto, Herke Hoof, and David Meger, ‘Addressing function approximation error in actor-critic methods’, in International conference on machine learning, pp. 1587–1596. PMLR, (2018).
  * [14] Tanmay Gangwani, Joel Lehman, Qiang Liu, and Jian Peng, ‘Learning belief representations for imitation learning in pomdps’, in Uncertainty in Artificial Intelligence, pp. 1061–1071. PMLR, (2020).
  * [15] M Giacobbe, Mohammadhosein Hasanbeig, Daniel Kroening, and Hjalmar Wijk, ‘Shielding atari games with bounded prescience’, in Proceedings of the International Joint Conference on Autonomous Agents and Multiagent Systems, AAMAS, (2021).
  * [16] Alexander W Goodall and Francesco Belardinelli, ‘Approximate shielding of atari agents for safe exploration’, arXiv preprint arXiv:2304.11104, (2023).
  * [17] William H Guss, Cayden Codel, Katja Hofmann, Brandon Houghton, Noboru Kuno, Stephanie Milani, Sharada Mohanty, Diego Perez Liebana, Ruslan Salakhutdinov, Nicholay Topin, et al., ‘Neurips 2019 competition: the minerl competition on sample efficient reinforcement learning using human priors’, arXiv preprint arXiv:1904.10079, (2019).
  * [18] David Ha and Jürgen Schmidhuber, ‘Recurrent world models facilitate policy evolution’, in Advances in Neural Information Processing Systems, eds., S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett, volume 31. Curran Associates, Inc., (2018).
  * [19] Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi, ‘Dream to control: Learning behaviors by latent imagination’, in International Conference on Learning Representations, (2020).
  * [20] 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. 2555–2565. PMLR, (2019).
  * [21] Danijar Hafner, Timothy P Lillicrap, Mohammad Norouzi, and Jimmy Ba, ‘Mastering atari with discrete world models’, in International Conference on Learning Representations, (2021).
  * [22] Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap, ‘Mastering diverse domains through world models’, arXiv preprint arXiv:2301.04104, (2023).
  * [23] Alexander Hans, Daniel Schneegaß, Anton Maximilian Schäfer, and Steffen Udluft, ‘Safe exploration for reinforcement learning.’, in ESANN, pp. 143–148. Citeseer, (2008).
  * [24] P He, B Gonzalez Leon, and F Belardinelli, ‘Do androids dream of electric fences? safety-aware reinforcement learning with latent shielding’. CEUR Workshop Proceedings.
  * [25] Matteo Hessel, Joseph Modayil, Hado Van Hasselt, Tom Schaul, Georg Ostrovski, Will Dabney, Dan Horgan, Bilal Piot, Mohammad Azar, and David Silver, ‘Rainbow: Combining improvements in deep reinforcement learning’, in Proceedings of the AAAI conference on artificial intelligence, volume 32, (2018).
  * [26] Wassily Hoeffding, ‘Probability inequalities for sums of bounded random variables’, The collected works of Wassily Hoeffding, 409–426, (1994).
  * [27] Michael Janner, Justin Fu, Marvin Zhang, and Sergey Levine, ‘When to trust your model: Model-based policy optimization’, Advances in neural information processing systems, 32, (2019).
  * [28] Michael Kearns and Satinder Singh, ‘Near-optimal reinforcement learning in polynomial time’, Machine learning, 49, 209–232, (2002).
  * [29] Jonathan N Lee, Alekh Agarwal, Christoph Dann, and Tong Zhang, ‘Learning in pomdps is sample-efficient with hindsight observability’, arXiv preprint arXiv:2301.13857, (2023).
  * [30] Yongshuai Liu, Jiaxin Ding, and Xin Liu, ‘Ipo: Interior-point policy optimization under constraints’, in Proceedings of the AAAI conference on artificial intelligence, volume 34, pp. 4940–4947, (2020).
  * [31] Zuxin Liu, Hongyi Zhou, Baiming Chen, Sicheng Zhong, Martial Hebert, and Ding Zhao, ‘Constrained model-based reinforcement learning with robust cross-entropy method’, arXiv preprint arXiv:2010.07968, (2020).
  * [32] Yuping Luo and Tengyu Ma, ‘Learning barrier certificates: Towards safe reinforcement learning with zero training-time violations’, Advances in Neural Information Processing Systems, 34, 25621–25632, (2021).
  * [33] Marlos C. Machado, Marc G. Bellemare, Erik Talvitie, Joel Veness, Matthew J. Hausknecht, and Michael Bowling, ‘Revisiting the arcade learning environment: Evaluation protocols and open problems for general agents’, Journal of Artificial Intelligence Research, 61, 523–562, (2018).
  * [34] Haritz Odriozola-Olalde, Maider Zamalloa, and Nestor Arana-Arexolaleiba, ‘Shielded reinforcement learning: A review of reactive methods for safe learning’, in 2023 IEEE/SICE International Symposium on System Integration (SII), pp. 1–8. IEEE, (2023).
  * [35] Martin L Puterman, ‘Markov decision processes’, Handbooks in operations research and management science, 2, 331–434, (1990).
  * [36] Aravind Rajeswaran, Igor Mordatch, and Vikash Kumar, ‘A game theoretic framework for model based reinforcement learning’, in International conference on machine learning, pp. 7953–7963. PMLR, (2020).
  * [37] Alex Ray, Joshua Achiam, and Dario Amodei, ‘Benchmarking safe exploration in deep reinforcement learning’, arXiv preprint arXiv:1910.01708, 7(1), 2, (2019).
  * [38] Richard S Sutton, ‘Dyna, an integrated architecture for learning, planning, and reacting’, ACM Sigart Bulletin, 2(4), 160–163, (1991).
  * [39] Richard S Sutton and Andrew G Barto, Reinforcement learning: An introduction, MIT press, 2018.
  * [40] Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, et al., ‘Deepmind control suite’, arXiv preprint arXiv:1801.00690, (2018).
  * [41] Garrett Thomas, Yuping Luo, and Tengyu Ma, ‘Safe reinforcement learning by imagining the near future’, Advances in Neural Information Processing Systems, 34, 13859–13869, (2021).
  * [42] Christopher KI Williams and Carl Edward Rasmussen, Gaussian processes for machine learning, volume 2, MIT press Cambridge, MA, 2006.
  * [43] Jorge Nocedal Stephen J Wright.  Numerical optimization, 2006.
  * [44] Wenli Xiao, Yiwei Lyu, and John Dolan, ‘Model-based dynamic shielding for safe and efficient multi-agent reinforcement learning’, arXiv preprint arXiv:2304.06281, (2023).
  * [45] Tsung-Yen Yang, Justinian Rosca, Karthik Narasimhan, and Peter J Ramadge, ‘Projection-based constrained policy optimization’, in International Conference on Learning Representations.

##  Appendix A Algorithms

###  A.1 DreamerV3 with AMBS

Algorithm 2 DreamerV3 [22] with AMBS

Initialise: replay buffer $\mathcal{D}$ with $S$ random episodes and RSSM
parameters $\theta$ randomly.

while not converged do

// World model learning

Sample $B$ sequences $\\{\langle
o_{t},a_{t},r_{t},c_{t},\gamma^{\text{safe}}_{t},o_{t+1}\rangle^{k+H}_{t=k}\\}\sim\mathcal{D}$.

Update RSSM parameters $\theta$ with representation learning [22].

// Task policy optimisation

‘Imagine’ trajectory $\sigma$ from every $o_{t}$ in every seq. with
$\pi^{\text{task}}$.

Train $\pi^{\text{task}}$ with TD-$\lambda$ actor-critic to optim. Eq. 2.

// Safety critic optimisation

Train $v^{C}_{1}$ and $v^{C}_{2}$ with maximum likelihood to estim. Eq. 11.

// Safe policy optimisation

‘Imagine’ trajectory $\sigma$ from every $o_{t}$ in every seq. with
$\pi^{\text{safe}}$.

Train $\pi^{\text{safe}}$ with TD-$\lambda$ actor-critic to optim. Eq. 10.

// Environment interaction

for $k=1,...,K$ do

Observe $o_{t}$ from environment and compute $\hat{s}_{t}=(z_{t},h_{t})$.

Sample action $a\sim\pi^{\text{task}}$ with the task policy.

Shield the proposed action $a^{\prime}=\text{shield}(a)$ and play
$a^{\prime}$.

Observe $r_{t},o_{t+1}$ and $L(s_{t})$ and construct $c_{t}$ and
$\gamma^{\text{safe}}_{t}$.

Append $\langle
o_{t},a_{t},r_{t},c_{t},\gamma^{\text{safe}}_{t},o_{t+1}\rangle$ to
$\mathcal{D}$.

end for

end while

##  Appendix B Proofs

We provide proofs for the theorems introduced in Section 3.

###  B.1 Proof of Theorem 3

######  Theorem 3 Restated.

Let $\epsilon>0$, $\delta>0$, $s\in S$ be given. With access to the ‘true’
transition system $\mathcal{T}$, with probability $1-\delta$ we can obtain an
$\epsilon$-approximate estimate of the measure $\mu_{s\models\phi}$, by
sampling $m$ traces $\tau\sim\mathcal{T}$, provided that,

$m\geq\frac{1}{2\epsilon^{2}}\log\left(\frac{2}{\delta}\right)$

###### Proof.

We estimate $\mu_{s\models\phi}$ by sampling $m$ traces
$\langle\tau_{j}\rangle^{m}_{j=1}$ from $\mathcal{T}$. Let $X_{1},...,X_{m}$
be indicator r.v.s such that,

$X_{j}=\begin{cases}1&\text{if $\tau_{j}\models\square^{\leq n}\Psi$,}\\\
0&\text{otherwise}\end{cases}$

Let,

$\tilde{\mu}_{s\models\phi}=\frac{1}{m}\sum^{m}_{j=1}X_{j},\;\text{where}\;\>\mathbb{E}_{\mathcal{T}}[\tilde{\mu}_{s\models\phi}]=\mu_{s\models\phi}$

Then by Hoeffding’s inequality,

$\mathbb{P}\left[|\tilde{\mu}_{s\models\phi}-\mu_{s\models\phi}|\geq\epsilon\right]\leq
2\exp\left(-2m\epsilon^{2}\right)$

Bounding the RHS from above with $\delta$ and rearranging completes the proof.
∎

The only caveat is arguing that $\tau_{j}\models\square^{\leq n}\Psi$ is
easily checkable. Indeed this is the case (for polynmial $n$) because in the
fully observable setting we have access to the state and so we can check that
$\forall i\;\tau[i]\models\Psi$.

###  B.2 Proof of Theorem 5

We split the proof into two parts for Theorem 5, first we present the _error
amplification_ lemma [36], followed by the full proof of Theorem 5.

######  Lemma  (error amplification).

Let $\mathcal{T}(s^{\prime}\mid s)$ and $\widehat{\mathcal{T}}(s^{\prime}\mid
s)$ be two transition systems with the same initial state distribution
$\iota_{init}$. Let $\mathcal{T}^{t}(s)$ and $\widehat{\mathcal{T}}^{t}(s)$ be
the marginal state distribution at time $t$ for the transitions systems
$\mathcal{T}$ and $\widehat{\mathcal{T}}$ respectively. That is,

$\displaystyle\mathcal{T}^{t}(s)=\mathbb{P}_{\tau\sim\mathcal{T}}[\tau[t]=s]$
$\displaystyle\widehat{\mathcal{T}}^{t}(s)=\mathbb{P}_{\tau\sim\widehat{\mathcal{T}}}[\tau[t]=s]$

Suppose that,

$D_{\text{TV}}\left(\mathcal{T}(s^{\prime}\mid
s),\widehat{\mathcal{T}}(s^{\prime}\mid s)\right)\leq\alpha\;\forall s\in S$

then the marginal distributions are bounded as follows,

$D_{\text{TV}}\left(\mathcal{T}^{t},\widehat{\mathcal{T}}^{t}\right)\leq\alpha
t\;\forall t$

###### Proof.

First let us fix some $s\in S$. Then,

$\displaystyle\big{|}\mathcal{T}^{t}(s)-\widehat{\mathcal{T}}^{t}(s)\big{|}$
$\displaystyle=\Big{|}\sum_{\bar{s}\in
S}\mathcal{T}(s\mid\bar{s})\mathcal{T}^{t-1}(\bar{s})-\sum_{\bar{s}\in
S}\widehat{\mathcal{T}}(s\mid\bar{s})\widehat{\mathcal{T}}^{t-1}(\bar{s})\Big{|}$
$\displaystyle\leq\sum_{\bar{s}\in
S}\big{|}\mathcal{T}(s\mid\bar{s})\mathcal{T}^{t-1}(\bar{s})-\widehat{\mathcal{T}}(s\mid\bar{s})\widehat{\mathcal{T}}^{t-1}(\bar{s})\big{|}$
$\displaystyle\leq\sum_{\bar{s}\in
S}\big{|}\mathcal{T}^{t-1}(\bar{s})\left(\mathcal{T}(s\mid\bar{s})-\widehat{\mathcal{T}}(s\mid\bar{s})\right)\big{|}$
$\displaystyle\qquad+\big{|}\widehat{\mathcal{T}}(s\mid\bar{s})\left(\mathcal{T}^{t-1}(\bar{s})-\widehat{\mathcal{T}}^{t-1}(\bar{s})\right)\big{|}$

Using the above inequality we get the following,

$\displaystyle
2D_{\text{TV}}\left(\mathcal{T}^{t},\widehat{\mathcal{T}}^{t}\right)$
$\displaystyle=\sum_{s\in
S}\big{|}\mathcal{T}^{t}(s)-\widehat{\mathcal{T}}^{t}(s)\big{|}$
$\displaystyle\leq\sum_{\bar{s}\in S}\mathcal{T}^{t-1}(\bar{s})\sum_{s\in
S}\big{|}\mathcal{T}(s\mid\bar{s})-\widehat{\mathcal{T}}(s\>|\>\bar{s})\big{|}$
$\displaystyle+\sum_{\bar{s}\in
S}\big{|}\mathcal{T}^{t-1}(\bar{s})-\widehat{\mathcal{T}}^{t-1}(\bar{s})\big{|}$
$\displaystyle\leq
2\alpha+2D_{\text{TV}}\left(\mathcal{T}^{t-1},\widehat{\mathcal{T}}^{t-1}\right)$
$\displaystyle\leq 2\alpha t$

The final inequality holds by applying the the recursion obtained on $t$ until
$t=0$ where $\mathcal{T}$ and $\widehat{\mathcal{T}}$ start from the same
initial state distribution $\iota_{init}$. ∎

######  Theorem 5 Restated.

Let $\epsilon>0$, $\delta>0$ be given. Suppose that for all $s\in S$, the
total variation (TV) distance between $\mathcal{T}(s^{\prime}\mid s)$ and
$\widehat{\mathcal{T}}(s^{\prime}\mid s)$ is bounded by some
$\alpha\leq\epsilon/n$. That is,

$D_{\text{TV}}\left(\mathcal{T}(s^{\prime}\mid
s),\widehat{\mathcal{T}}(s^{\prime}\mid s)\right)\leq\alpha\;\forall s\in S$

Now fix an $s\in S$, with probability $1-\delta$ we can obtain an
$\epsilon$-approximate estimate of the measure $\mu_{s\models\phi}$, by
sampling $m$ traces $\tau\sim\widehat{\mathcal{T}}$, provided that,

$m\geq\frac{2}{\epsilon^{2}}\log\left(\frac{2}{\delta}\right)$

###### Proof.

Recall that,

$\mu_{s\models\phi}=\mu_{s}(\\{\tau\mid\tau[0]=s,\text{ for all }0\leq i\leq
n,\tau[i]\models\Psi\\})$

where $\tau\sim\mathcal{T}$. Equivalently we can write,

$\mu_{s\models\phi}=\mathbb{P}_{\tau\sim\mathcal{T}}[\tau\models\square^{\leq
n}\Psi]$

Similarly, let $\hat{\mu}_{s\models\phi}$ be defined as the _true_ probability
under $\widehat{\mathcal{T}}$,

$\hat{\mu}_{s\models\phi}=\mathbb{P}_{\tau\sim\widehat{\mathcal{T}}}[\tau\models\square^{\leq
n}\Psi]$

Let the following denote the average state distribution for $\mathcal{T}$ and
$\widehat{\mathcal{T}}$ respectively,

$\displaystyle\rho_{\mathcal{T}}(s)=\frac{1}{n}\sum^{n}_{i=1}\mathbb{P}_{\tau\sim\mathcal{T}}(\tau[i]=s)$
$\displaystyle\rho_{\widehat{\mathcal{T}}}(s)=\frac{1}{n}\sum^{n}_{i=1}\mathbb{P}_{\tau\sim\widehat{\mathcal{T}}}(\tau[i]=s)$

By following the simulation lemma [28], we get the following,

$\displaystyle\big{|}\mu_{s\models\phi}-\hat{\mu}_{s\models\phi}\big{|}$
$\displaystyle=\big{|}\mathbb{P}_{\tau\sim\mathcal{T}}[\tau\models\square^{\leq
n}\Psi]-\mathbb{P}_{\tau\sim\widehat{\mathcal{T}}}[\tau\models\square^{\leq
n}\Psi]\big{|}$ $\displaystyle\leq 1\cdot
D_{TV}(\rho_{\mathcal{T}},\rho_{\widehat{\mathcal{T}}})$
$\displaystyle=\frac{1}{2}\sum_{s\in
S}\big{|}\rho_{\mathcal{T}}(s)-\rho_{\widehat{\mathcal{T}}}(s)\big{|}$
$\displaystyle=\frac{1}{2n}\sum_{s\in
S}\Big{|}\sum^{n}_{i=1}\mathbb{P}_{\tau\sim\mathcal{T}}(\tau[i]=s)-\mathbb{P}_{\tau\sim\widehat{\mathcal{T}}}(\tau[i]=s)\Big{|}$
$\displaystyle\leq\frac{1}{2n}\sum_{s\in
S}\sum^{n}_{i=1}\Big{|}\mathbb{P}_{\tau\sim\mathcal{T}}(\tau[i]=s)-\mathbb{P}_{\tau\sim\widehat{\mathcal{T}}}(\tau[i]=s)\Big{|}$
$\displaystyle\leq\frac{1}{2n}\sum^{n}_{i=1}\alpha n\qquad\text{(Using Lemma
\ref{lemma:marginal})}$ $\displaystyle=\frac{\alpha n}{2}$

Now we have that
$\big{|}\mu_{s\models\phi}-\hat{\mu}_{s\models\phi}\big{|}\leq(\alpha
n)/2\leq\epsilon/2$. It remains to obtain an $\epsilon/2$-approximation of
$\hat{\mu}_{\models\phi}$. Using the exact same reasoning as in the proof of
Theorem 3, we estimate $\hat{\mu}_{\models\phi}$ by sampling $m$ traces
$\langle\tau_{j}\rangle^{m}_{j=1}$ from $\widehat{\mathcal{T}}$. Then
provided,

$m\geq\frac{2}{\epsilon^{2}}\log\left(\frac{2}{\delta}\right)$

with probability $1-\delta$ we obtain an $\epsilon/2$-approximation of
$\hat{\mu}_{\models\phi}$ and by extension an $\epsilon$-approximation of
$\mu_{\models\phi}$. ∎

###  B.3 Proof of Theorem 7

######  Theorem 7 Restated.

Let $\alpha>0$, $\delta>0$, $s\in S$ be given. With probability $1-\delta$ the
total variation (TV) distance between $\mathcal{T}(s^{\prime}\mid s)$ and
$\widehat{\mathcal{T}}(s^{\prime}\mid s)$ is upper bounded by $\alpha$,
provided that all actions $a\in A$ with non-negligable probability
$\eta\geq\alpha/(|A||S|)$ (under $\pi$) have been picked from $s$ at least $m$
times, where

$m\geq\frac{|S|^{2}}{\alpha^{2}}\log\left(\frac{2|S||A|}{\delta}\right)$

###### Proof.

Recall that $p(s^{\prime}\mid s,a)$ denotes the probability of transitioning
to $s^{\prime}$ from $s$ when action $a$ is played. Lets first fix
$s^{\prime}$ and just consider approximating one of these probabilities. Let,

  * •

$p=p(s^{\prime}\mid s,a)$ and

  * •

$m$ be the number of times $a$ is played from $s$.

  * •

Let $X_{1},...,X_{m}$ be the indicator r.v.s such that,

$X_{i}=\begin{cases}1&\text{if $s^{\prime}$ given $(s,a)$}\\\
0&\text{otherwise}\end{cases}$

  * •

$\hat{p}=\frac{1}{m}\sum^{m}_{i=1}X_{i}$, where $\mathbb{E}[\hat{p}]=p$.

By Hoeffding’s Inequality we have,

$\mathbb{P}\left[|\hat{p}-p|\geq\frac{\alpha}{|S|}\right]\leq
2\cdot\exp\left(-2m\frac{\alpha^{2}}{|S|^{2}}\right)$

Bounding the RHS from above by $\delta/(|S||A|)$ gives us,

$m\geq\frac{|S|^{2}}{\alpha^{2}}\log\left(\frac{2|S||A|}{\delta}\right)$

And so with probability $1-\delta/(|S||A|)$ we have an
$\alpha/|S|$-approximation for $p$ provided $m$ satisfies the above bound.
Taking a union bound over all $s^{\prime}\in S$ and all $a\in A$ we have with
probability at least $1-\delta$ an $\alpha/|S|$-approximation
$p(s^{\prime}\mid s,a)$ for all $s^{\prime}\in S$ and all actions $a\in A$
with non-negligible probability $\eta\geq\frac{\alpha}{|A|}$ (under $\pi$). It
remains to show that the TV distance between $\mathcal{T}(s^{\prime}\mid s)$
and $\widehat{\mathcal{T}}(s^{\prime}\mid s)$ is upper bounded by $\alpha$,

$\displaystyle 2D_{TV}\left(\mathcal{T}(s^{\prime}\mid
s),\widehat{\mathcal{T}}(s^{\prime}\mid s)\right)$
$\displaystyle=\sum_{s^{\prime}\in S}\big{|}\mathcal{T}(s^{\prime}\mid
s)-\widehat{\mathcal{T}}(s^{\prime}\mid s)\big{|}$
$\displaystyle=\sum_{s^{\prime}\in S}\sum_{a\in A}\big{|}p(s^{\prime}\mid
s,a)\pi(a\mid s)-\hat{p}(s^{\prime}\mid s,a)\pi(a\mid s)\big{|}$
$\displaystyle=\sum_{s^{\prime}\in S}\Bigg{(}\sum_{a\in A:\pi(a\mid
s)\geq\eta}\big{|}p(s^{\prime}\mid s,a)\pi(a\mid s)-\hat{p}(s^{\prime}\mid
s,a)\pi(a\mid s)\big{|}$ $\displaystyle\qquad+\sum_{a\in A:\pi(a\mid
s)<\eta}\big{|}p(s^{\prime}\mid s,a)\pi(a\mid s)-\hat{p}(s^{\prime}\mid
s,a)\pi(a\mid s)\big{|}\Bigg{)}$ $\displaystyle=\sum_{s^{\prime}\in
S}\Bigg{(}\sum_{a\in A:\pi(a\mid s)\geq\eta}\pi(a\mid
s)\big{|}p(s^{\prime}\mid s,a)-\hat{p}(s^{\prime}\mid s,a)\big{|}$
$\displaystyle\qquad+\sum_{a\in A:\pi(a\mid s)<\eta}\pi(a\mid
s)\big{|}p(s^{\prime}\mid s,a)-\hat{p}(s^{\prime}\mid s,a)\big{|}\Bigg{)}$
$\displaystyle\leq\sum_{s^{\prime}\in S}\Bigg{(}\sum_{a\in A:\pi(a\mid
s)\geq\eta}\pi(a\mid s)\frac{\alpha}{|S|}+\sum_{a\in A:\pi(a\mid
s)<\eta}\pi(a\mid s)\Bigg{)}$ $\displaystyle\leq\sum_{s^{\prime}\in
S}\left(\frac{\alpha}{|S|}+|A|\cdot\eta\right)$
$\displaystyle\leq\sum_{s^{\prime}\in
S}\left(\frac{\alpha}{|S|}+\frac{\alpha}{|S|}\right)$ $\displaystyle\leq
2\alpha$

∎

###  B.4 Proof of Theorem 3.3

######  Theorem 3.3 Restated.

Let $b_{t}$ be a latent representation (belief state) such that $p(s_{t}\mid
o_{t\leq t},a_{\leq t})=p(s_{t}\mid b_{t})$. Let the fixed policy
$\pi(\cdot\mid b_{t})$ be a general probability distribution conditional on
belief states $b_{t}$. Let $f$ be a generic $f$-divergence measure (TV or
similar). Then the following holds:

$D_{f}(\mathcal{T}(s^{\prime}\mid b),\widehat{\mathcal{T}}(s^{\prime}\mid
b))\leq D_{f}(\mathcal{T}(b^{\prime}\mid
b),\widehat{\mathcal{T}}(b^{\prime}\mid b))$

where $\mathcal{T}$ and $\widehat{\mathcal{T}}$ are the ‘true’ and approximate
transition system respectively, defined now over both states $s$ and belief
states $b$.

###### Proof.

First for clarity, we define the transitions systems $\mathcal{T}$ and
$\widehat{\mathcal{T}}$ over states $s$ and belief states $b$. First we have
as before

$\displaystyle\mathcal{T}(s^{\prime}\mid
b)=\mathbb{P}_{\pi,p}[s_{t}=s^{\prime}\mid b_{t-1}=b]$
$\displaystyle\widehat{\mathcal{T}}(s^{\prime}\mid
b)=\mathbb{P}_{\pi,\hat{p}}[s_{t}=s^{\prime}\mid b_{t-1}=b]$

Additionally, let,

$\displaystyle\mathcal{T}(b^{\prime}\mid
b)=\mathbb{P}_{\pi,p}[b_{t}=b^{\prime}\mid b_{t-1}=b]$
$\displaystyle\widehat{\mathcal{T}}(b^{\prime}\mid
b)=\mathbb{P}_{\pi,\hat{p}}[b_{t}=b^{\prime}\mid b_{t-1}=b]$

From these definitions, note that we can immediately define conditional and
joint probabilities (i.e. $\mathcal{T}(s^{\prime},b^{\prime}\mid b)$,
$\mathcal{T}(s^{\prime}\mid b)$, $\mathcal{T}(b^{\prime}\mid s^{\prime},b)$
and similarly for $\widehat{\mathcal{T}}$) using the standard laws of
probability.

Now we are ready to apply the data-processing inequality [2] for f-divergences
as follows,

$\displaystyle D_{f}(\mathcal{T}(b^{\prime}\mid
b),\widehat{\mathcal{T}}(b^{\prime}\mid b))$
$\displaystyle=\mathbb{E}_{b^{\prime}\sim\widehat{\mathcal{T}}}\left[f\left(\frac{\mathcal{T}(b^{\prime}\mid
b)}{\widehat{\mathcal{T}}(b^{\prime}\mid b)}\right)\right]$
$\displaystyle=\mathbb{E}_{s^{\prime},b^{\prime}\sim\widehat{\mathcal{T}}}\left[f\left(\frac{\mathcal{T}(s^{\prime},b^{\prime}\mid
b)}{\widehat{\mathcal{T}}(s^{\prime},b^{\prime}\mid b)}\right)\right]$
$\displaystyle=\mathbb{E}_{s^{\prime}\sim\widehat{\mathcal{T}}}\left[\mathbb{E}_{b^{\prime}\sim\widehat{\mathcal{T}}}f\left(\frac{\mathcal{T}(s^{\prime},b^{\prime}\mid
b)}{\widehat{\mathcal{T}}(s^{\prime},b^{\prime}\mid b)}\right)\right]$
$\displaystyle\geq\mathbb{E}_{s^{\prime}\sim\widehat{\mathcal{T}}}\left[f\left(\mathbb{E}_{b^{\prime}\sim\widehat{\mathcal{T}}}\frac{\mathcal{T}(s^{\prime},b^{\prime}\mid
b)}{\widehat{\mathcal{T}}(s^{\prime},b^{\prime}\mid
b)}\right)\right]\qquad\text{(Jensen's)}$
$\displaystyle=\mathbb{E}_{s^{\prime}\sim\widehat{\mathcal{T}}}\left[f\left(\mathbb{E}_{b^{\prime}\sim\widehat{\mathcal{T}}}\frac{\mathcal{T}(s^{\prime},b^{\prime}\mid
b)\widehat{\mathcal{T}}(b^{\prime}\mid
s^{\prime},b)}{\widehat{\mathcal{T}}(s^{\prime},b^{\prime}\mid
b)\mathcal{T}(b^{\prime}\mid s^{\prime},b)}\right)\right]$
$\displaystyle=\mathbb{E}_{s^{\prime}\sim\widehat{\mathcal{T}}}\left[f\left(\mathbb{E}_{b^{\prime}\sim\widehat{\mathcal{T}}}\frac{\mathcal{T}(s^{\prime}\mid
b)}{\widehat{\mathcal{T}}(s^{\prime}\mid b)}\right)\right]$
$\displaystyle=\mathbb{E}_{s^{\prime}\sim\widehat{\mathcal{T}}}\left[f\left(\frac{\mathcal{T}(s^{\prime}\mid
b)}{\widehat{\mathcal{T}}(s^{\prime}\mid b)}\right)\right]$
$\displaystyle=D_{f}(\mathcal{T}(s^{\prime}\mid
b),\widehat{\mathcal{T}}(s^{\prime}\mid b))$

∎

##  Appendix C Code and Hyperparameters

Here we specify the most important hyperparameters for each of the agents in
our experiments. For more precise implementation details we refer the reader
to https://github.com/sacktock/AMBS.

#### World Model.

As discussed we leverage DreamerV3 [22]. For all architectural details and
hyperparameter choices please refer to [22].

#### Predictor Heads

The cost function and safety discount predictor heads are implemented as
neural network architectures similar to those used for the reward predictor
and termination predictor of DreamerV3 [22]. Specifically the cost function
predictor head is implemented in exactly the same was as the reward predictor
head except it is used to predict cost signals instead of reward signals.
Similarly, the safety discount predictor is a binary classifier implemented in
the exact same was as the termination predictor, except one predicts safety-
violations and the other predicts episode termination.

#### Safe Policy

As discussed, the safe policy is trained with the same TD-$\lambda$ style
actor-critic algorithm that is used for the standard reward maximising (task)
policy of DreamerV3 [22]. The actor and critic are implemented as neural
networks with the same architecture that is used for the task policy and the
hyperparameters are consistent between the safe policy and the task policy.
Other details are outlined in Table 3.

#### Safety Critics

The safety critics are trained with a TD3-style algorithm [13]. The two
critics themselves (and their target networks) are implemented as neural
networks with the same architectures used for the critics that help train the
task policy and safe policy. The hyperparameters also are mostly the same and
any changes are outlined in Table 4.

###  C.1 DreamerV3 Hyperparameters

Table 3: DreamerV3 hyperparameters [22]. Other methods built on DreamerV3 such as AMBS and LAG use this set of hyperparameters as well, unless otherwise specified. Name | Symbol | Value
---|---|---
General
Replay capacity | - | $10^{6}$
Batch size | $B$ | 16
Batch Length | - | 64
Num. Envs | - | 8
Train ratio | - | 64
MLP Layers | - | 5
MLP Units | - | 512
Activation | - | LayerNorm + SiLU
World Model
Num. latents | - | 32
Classes per latent | - | 32
Num. Layers | - | 5
Num. Units | - | 1024
Recon. loss scale | $\beta_{\text{pred}}$ | 1.0
Dynamics loss scale | $\beta_{\text{dyn}}$ | 0.5
Represen. loss scale | $\beta_{\text{rep}}$ | 0.1
Learning rate | - | $10^{-4}$
Adam epsilon | $\epsilon_{\text{adam}}$ | $10^{-8}$
Gradient clipping | - | 1000
Actor Critic
Imagination horizon | H | 15
Discount factor | $\gamma$ | 0.997
TD lambda | $\lambda$ | 0.95
Critic EMA decay | - | 0.98
Critic EMA regulariser | - | 1
Return norm. scale | $S_{\text{reward}}$ | $\text{Per}(R,95)-\text{Per}(R,5)$
Return norm. limit | $L_{\text{reward}}$ | 1
Return norm. decay | - | 0.99
Actor entropy scale | $\eta_{\text{actor}}$ | $3\cdot 10^{-4}$
Learning rate | - | $3\cdot 10^{-5}$
Adam epsilon | $\epsilon_{\text{adam}}$ | $10^{-5}$
Gradient clipping | - | 100

###  C.2 AMBS Hyperparameters

Table 4: AMBS hyperparameters. We note that $m>512$ is sufficient for $\Delta=0.1$, $\epsilon=0.09$, $\delta=0.01$ using a bound similar to Eq. 5 that gives a bound on overestimating $\mu_{s\models\phi}$. Name | Symbol | Value
---|---|---
Shielding
Safety level | $\Delta$ | 0.1
Approx. error | $\epsilon$ | 0.09
Num. samples | $m$ | 512
Failure probability | $\delta$ | 0.01
Look-ahead horizon | $T$ | 30
Cost Value | $C$ | 10
Safe Policy
See ‘Actor Critic’ table 3
…
Safety Critic
Type | - | TD3-style [13]
Slow update freq. | - | 1
Slow update fraction | - | 0.02
EMA decay | - | 0.98
EMA regulariser | - | 1
Cost norm. scale | $S_{\text{cost}}$ | $\text{Per}(R,95)-\text{Per}(R,5)$
Cost norm. limit | $L_{\text{cost}}$ | 1
Cost norm. decay | - | 0.99
Learning rate | - | $3\cdot 10^{-5}$
Adam epsilon | $\epsilon_{\text{adam}}$ | $10^{-5}$
Gradient clipping | - | 100

###  C.3 LAG Hyperparameters

Table 5: LAG hyperparameters [37]. We use the default hyperparameters provided in [6] for the safety gym benchmark [37].  Name | Symbol | Value
---|---|---
Penalty Multiplier | $\mu_{k}$ | $5\cdot 10^{-9}$
Lagrange Multiplier | $\lambda^{k}$ | $10^{-6}$
Penalty Power | $\sigma$ | $10^{-5}$
Safety Horizon | $T$ | 30
Cost Value | $C$ | 10
Cost Threshold | $d=C\cdot\gamma^{T}$ | $\approx 9$

###  C.4 IQN Hyperparameters

Table 6: IQN hyperparameters [12] adapted for a fairer comparison as in [10]. Name | Symbol | Value
---|---|---
Replay capacity | - | $10^{6}$
Batch size | $B$ | 32
IQN kappa | $\kappa$ | 1.0
Num. $\tau$ samples | - | 64
Num. $\tau^{\prime}$ samples | - | 64
Num. quantile samples | - | 32
Discount factor | $\gamma$ | $0.99$
Update horizon | $n_{\text{hor}}$ | 3
Update freq. | - | $4$
Target update freq. | - | 8000
Epsilon greedy min | $\epsilon_{\text{min}}$ | 0.01
Epsilon decay period | - | 250000
Optimiser | - | Adam
Learning rate | - | $5\cdot 10^{-5}$
Adam epsilon | $\epsilon_{\text{adam}}$ | $3.125\cdot 10^{-4}$

###  C.5 Rainbow Hyperparameters

Table 7: Rainbow hyperparameters as recommended in [25]. Name | Symbol | Value
---|---|---
Replay capacity | - | $10^{6}$
Batch size | $B$ | 32
Discount factor | $\gamma$ | $0.99$
Update horizon | $n_{\text{hor}}$ | 3
Update freq. | - | $4$
Target update freq. | - | 8000
Epsilon greedy min | $\epsilon_{\text{min}}$ | 0.01
Epsilon decay period | - | 250000
Optimiser | - | Adam
Learning rate | - | $6.25\cdot 10^{-5}$
Adam epsilon | $\epsilon_{\text{adam}}$ | $1.5\cdot 10^{-4}$

##  Appendix D Learning Curves

Figure 2: Learning curves for each of the five agents on a small set of Atari
games. Each line represents one run over 10M environment interactions (40M
frames). The left plots represent the episode return and the right plots
represent the cumulative safety-violations during training.