BAYESIAN WORLD MODELS
Year: 2022

Abstract:
Improving sample-efficiency and safety are crucial challenges when deploying reinforcement learning in high-stakes real world applications. We propose LAMBDA, a novel model-based approach for policy optimization in safety critical tasks modeled via constrained Markov decision processes. Our approach utilizes Bayesian world models, and harnesses the resulting uncertainty to maximize optimistic upper bounds on the task objective, as well as pessimistic upper bounds on the safety constraints. We demonstrate LAMBDA’s state of the art performance on the Safety-Gym benchmark suite in terms of sample efficiency and constraint violation.

text: 1 INTRODUCTION
A central challenge in deploying reinforcement learning (RL) agents in real-world systems is to avoid unsafe and harmful situations (Amodei et al., 2016). We thus seek agents that can learn efficiently and safely by acting cautiously and ensuring the safety of themselves and their surroundings.
A common paradigm in RL for modeling safety critical environments are constrained Markov decision processes (CMDP) (Altman, 1999). CMDPs augment the common notion of the reward by an additional cost function, e.g., indicating unsafe state-action pairs. By bounding the cost, one obtains bounds on the probability of harmful events. For the tabular case, CMDPs can be solved via Linear Programs (LP) (Altman, 1999). Most prior work for solving non-tabular CMDPs, utilize model-free algorithms (Chow et al., 2015; Achiam et al., 2017; Ray et al., 2019; Chow et al., 2019). Most notably, these methods are typically shown to asymptotically converge to a constraintsatisfying policy. However, similar to model-free approaches in unconstrained MDPs, these approaches typically have a very high sample complexity, i.e., require a large number of – potentially harmful – interactions with the environment.
A promising alternative to improve sample efficiency is to use model-based reinforcement learning (MBRL) approaches
(Deisenroth & Rasmussen, 2011; Chua et al., 2018; Hafner et al., 2019a; Janner et al., 2019). In particular, Bayesian approaches to MBRL quantify uncertainty in the estimated model (Depeweg et al., 2017), which can be used to guide exploration, e.g., via the celebrated optimism in the face of uncertainty principle (Brafman & Tennenholtz, 2003; Auer & Ortner, 2007; Curi et al., 2020a). While extensively explored in the unconstrained setting, Bayesian model-based deep RL approaches for solving general CMDPs remain largely unexplored. In this paper, we close this gap by proposing LAMBDA, a Bayesian approach to model-based policy optimization in CMDPs. LAMBDA learns a safe policy by experiencing unsafe as well as rewarding events through model-generated trajectories instead of real ones. Our main contributions are as follows: ∗Correspondence to: yardas@ethz.ch †Equal contribution.
• We show that it is possible to perform constrained optimization by back-propagating gradients of both task and safety objectives through the world-model. We use the Augmented Lagrangian (Nocedal & Wright, 2006; Li et al., 2021) approach for constrained optimization.
• We harness the probabilistic world model to trade-off between optimism for exploration and pessimism for safety.
• We empirically show that LAMBDA successfully solves the SG6 benchmark tasks in Safety-Gym (Ray et al., 2019) benchmark suite with first-person-view observations as illustrated in Figure 2. Furthermore, we show that LAMBDA outperforms other model-based and model-free methods in this benchmark.

heading: 2 RELATED WORK
text: Interpretations of safety in RL research We first acknowledge that there exist different definitions for safety (Garcı́a et al., 2015). One definition uses reversible Markov decision processes (Moldovan & Abbeel, 2012) that, informally, define safety as reachability between states and use backup policies (Eysenbach et al., 2017) to ensure safety. Another definition adopts robust Markov decision processes (Nilim & Ghaoui, 2005; Tamar et al., 2013; Tessler et al., 2019), which try to maximize performance under the worst-case transition model. Finally, in this work we use nontabular CMDPs (Altman, 1999) that define the safety requirements as a cost function that should be bounded by a predefined threshold. Similarly, Achiam et al. (2017); Dalal et al. (2018); Ray et al. (2019); Chow et al. (2019); Stooke et al. (2020); Bharadhwaj et al. (2021); Turchetta et al. (2021) use non-tabular CMDPs as well, but utilize model-free techniques together with function approximators to solve the constrained policy optimization problem. As we further demonstrate in our experiments, it is possible to achieve better sample efficiency with model-based methods.
Safe model-based RL A successful approach in applying Bayesian modeling to low-dimensional continuous-control problems is to use Gaussian Processes (GP) for model learning. Notably, Berkenkamp et al. (2017) use GPs to construct confidence intervals around Lyapunov functions which are used to optimize a policy such that it is always within a Lyapunov-stable region of attraction. Furthermore, Koller et al. (2018); Wabersich & Zeilinger (2021) use GP models to certify the safety of actions within a model predictive control (MPC) scheme. Likewise, Liu et al. (2021) apply MPC for high-dimensional continuous-control problems together with ensembles of neural networks (NN), the Cross Entropy Method (CEM) (Kroese et al., 2006) and rejection sampling to maximize expected returns of safe action sequences. However, by planning online only for short horizons and not using critics, this method can lead to myopic behaviors, as we later show in our experiments.
Lastly, similarly to this work, Zanger et al. (2021) use NNs and constrained model-based policy optimization. However, they do not use model uncertainty within an optimistic-pessimistic framework but rather to limit the influence of erroneous model predictions on their policy optimization. Even though using accurate model predictions can accelerate policy learning, this approach does not take advantage of the epistemic uncertainty (e.g., through optimism and pessimism) when such accurate predictions are rare.
Exploration, optimism and pessimism Curi et al. (2021) and Derman et al. (2019) also take a Bayesian optimistic-pessimistic perspective to find robust policies. However, these approaches do not use CMDPs and generally do not explicitly address safety. Similarly, Bharadhwaj et al. (2021) apply pessimism for conservative safety critics, and use them for constrained model-free policy optimization. Finally, Efroni et al. (2020) lay a theoretical foundation for the explorationexploitation dilemma and optimism in the setting of tabular CMPDs.

heading: 3 PRELIMINARIES

heading: 3.1 SAFE MODEL-BASED REINFORCEMENT LEARNING
text: Markov decision processes We consider an episodic Markov decision process with discrete time steps t ∈ {0, . . . , T}. We define the environment’s state as st ∈ Rn and an action taken by the agent, as at ∈ Rm. Each episode starts by sampling from the initial-state distribution s0 ∼ ρ(s0). After observing s0 and at each step thereafter, the agent takes an action by sampling from a policy distribution at ∼ π(·|st). The next state is then sampled from an unknown transition distribution st+1 ∼ p(·|st,at). Given a state, the agent observes a reward, generated by rt ∼ p(·|st,at). We define the performance of a pair of policy π and dynamics p as
J(π, p) = Eat∼π,st+1∼p,s0∼ρ [ T∑ t=0 rt ∣∣s0] . (1)
Constrained Markov decision processes To make the agent adhere to human-defined safety constraints, we adopt the constrained Markov decision process formalism of Altman (1999). In CMPDs, alongside with the reward, the agent observes cost signals cit generated by c i t ∼ p(·|st,at) where i = 1, . . . , C denote the distinct unsafe behaviors we want the agent to avoid. Given cit, we define the constraints as
J i(π, p) = Eat∼π,st+1∼p,s0∼ρ [ T∑ t=0 cit ∣∣s0] ≤ di, ∀i ∈ {1, . . . , C}, (2)
where di are human-defined thresholds. For example, a common cost function is ci(st) = 1st∈Hi , where Hi is the set of harmful states (e.g., all the states in which a robot hits an obstacle). In this case, the constraint (2) can be interpreted as a bound on the probability of visiting the harmful states. Given the constraints, we aim to find a policy that solves, for the true unknown dynamics p?,
max π∈Π
J(π, p?) s.t. J i(π, p?) ≤ di, ∀i ∈ {1, . . . , C}. (3)
Model-based reinforcement learning Model-based reinforcement learning revolves around the repetition of an iterative process with three fundamental steps. First, the agent gathers observed transitions (either by following an initial policy, or an offline dataset) of {st+1, st,at} into a dataset D. Following that, the dataset is used to fit a statistical model p(st+1|st,at, θ) that approximates the true transition distribution p?. Finally, the agent uses the statistical model for planning, either within an online MPC scheme or via offline policy optimization. In this work, we consider the cost and reward functions as unknown and similarly to the transition distribution, we fit statistical models for them as well. Modeling the transition density allows us to cheaply generate synthetic sequences of experience through the factorization p(sτ :τ+H |sτ−1,aτ−1:τ+H−1, θ) = ∏τ+H t=τ p(st+1|st,at, θ) wherebyH is a predefined sequence horizon (see also Appendix I). Therefore, as already shown empirically (Deisenroth & Rasmussen, 2011; Chua et al., 2018; Hafner et al., 2019b), MBRL achieves superior sample efficiency compared to its model-free counterparts. We emphasize that in safetycritical tasks, where human supervision during training might be required, sample efficiency is particularly important since it can reduce the need for human supervision.
Learning a world model Aiming to tighten the gap between RL and real-world problems, we relax the typical full observability assumption and consider problems where the agent receives an observation ot ∼ p(·|st) instead of st at each time step. To infer the transition density from observations, we base our world model on the Recurrent State Space Model (RSSM) introduced in Hafner et al. (2019a). To solve this inference task, the RSSM approximates the posterior distribution sτ :τ+H ∼ qφ(·|oτ :τ+H ,aτ−1:τ+H−1) via an inference network parameterized by φ. We utilize the inference network to filter the latent state as new observations arrive, and use the inferred latent state as the policy’s input. Furthermore, the RSSM models the predictive distribution p(sτ :τ+H |sτ−1,aτ−1:τ+H−1, θ) as a differentiable function. This property allows us to perform the constrained policy optimization by backpropagating gradients through the model. We highlight that the only requirement for the world model is that it models p(sτ :τ+H |sτ−1,aτ−1:τ+H−1, θ) as a differentiable function. Hence, it is possible to use other architectures for the world model, as long as p(sτ :τ+H |sτ−1,aτ−1:τ+H−1, θ) is differentiable (e.g., see GP-models in Curi et al. (2020b)).

heading: 3.2 A BAYESIAN VIEW ON MODEL-BASED REINFORCEMENT LEARNING
text: The Bayesian predictive distribution A common challenge in MBRL is that policy optimization can “overfit” by exploiting inaccuracies of the estimated model. A natural remedy is to quantify uncertainty in the model, to identify in which situations the model is more likely to be wrong. This is especially true for safety-critical tasks in which model inaccuracies can mislead the agent into taking unsafe actions. A natural approach to uncertainty quantification is to take a Bayesian perspective, where we adopt a prior on the model parameters, and perform (approximate) Bayesian inference to obtain the posterior. Given such a posterior distribution over the model parameters p(θ|D), we marginalize over θ to get a Bayesian predictive distribution. Therefore, by maintaining such posterior, p(sτ :τ+H |sτ−1,aτ−1:τ+H−1) is determined by the law of total probability p(sτ :τ+H |sτ−1,aτ−1:τ+H−1) = Eθ∼p(θ|D) [p(sτ :τ+H |sτ−1,aτ−1:τ+H−1, θ)] . (4) Such Bayesian reasoning forms the basis of celebrated MBRL algorithms such as PETS (Chua et al., 2018) and PILCO (Deisenroth & Rasmussen, 2011), among others.
Maintaining a posterior over model parameters Since exact Bayesian inference is typically infeasible, we rely on approximations. In contrast to the popular approach of bootstrapped ensembles (Lakshminarayanan et al., 2017) widely used in RL research, we use the Stochastic Weight Averaging-Gaussian (SWAG) approximation (Maddox et al., 2019) of p(θ|D). Briefly, SWAG constructs a posterior distribution over model parameters by performing moment-matching over iterates of stochastic gradient decent (SGD) (Robbins, 2007). The main motivation behind this choice is the lower computational and memory footprint of SWAG compared to bootstrapped ensembles. This consideration is crucial when working with parameter-rich world models (such as the RSSM). Note that our approach is flexible and admits other variants of approximate Bayesian inference, such as variational techniques (Graves, 2011; Kingma & Welling, 2014).

heading: 4 LAGRANGIAN MODEL-BASED AGENT (LAMBDA)
text: We first present our proposed approach for planning, i.e., given a world model representing a CMDP, how to use it to solve the CMPD. Following that, we present our proposed method for finding the optimistic and pessimistic models with which the algorithm solves the CMDP before collecting more data.

heading: 4.1 SOLVING CMPDS WITH MODEL-BASED CONSTRAINED POLICY OPTIMIZATION
text: The Augmented Lagrangian We first consider the optimization problem in Equation (3). To find a policy that solves Equation (3), we take advantage of the Augmented Lagrangian with proximal relaxation method as described by Nocedal & Wright (2006). First, we observe that
max π∈Π min λ≥0
[ J(π)−
C∑ i=1 λi ( J i(π)− di
)] = max
π∈Π { J(π) if π is feasible −∞ otherwise (5)
is an equivalent form of Equation (3). Hereby λi are the Lagrange multipliers, each corresponding to a safety constraint measured by J i(π). In particular, if π is feasible, we have J i(π) ≤ di ∀i ∈
{1, . . . , C} and so the maximum over λi is satisfied if λi = 0 ∀i ∈ {1, . . . , C}. Conversely, if π is infeasible, at least one λi can be arbitrarily large to solve the problem in Equation (5). Particularly, Equation (5) is non-smooth when π transitions between the feasibility and infeasibility sets. Thus, in practice, we use the following relaxation:
max π∈Π min λ≥0
[ J(π)−
C∑ i=1 λi ( J i(π)− di ) + 1 µk C∑ i=1 ( λi − λik
)2] , (∗)
where µk is a non-decreasing penalty term corresponding to gradient step k. Note how the last term in (∗) encourages λi to stay proximal to the previous estimate λik and as a consequence making (∗) a smooth approximation of the left hand side term in Equation (5). Differentiating (∗) with respect to λi and substituting back leads to the following update rule for the Lagrange multipliers:
∀i ∈ {1, . . . , C} : λik+1 = { λik + µk(J
i(π)− di) if λik + µk(J i(π)− di) ≥ 0 0 otherwise . (6)
We take gradient steps of the following unconstrained objective:
J̃(π;λk, µk) = J(π)− C∑ i=1 Ψ i(π;λik, µk), (7)
where
Ψ i(π;λik, µk) =
{ λik(J i(π)− di) + µk2 ( J i(π)− di )2 if λik + µk(J
i(π)− di) ≥ 0 − (λ i k) 2
2µk otherwise.
(8)
Task and safety critics For the task and safety critics, we use the reward value function v(st) and the cost value function of each constraint vi(st) ∀i ∈ {1, . . . , C}. We model v(st) and vi(st) as neural networks with parameters ψ and ψi respectively. Note that we omit here the dependency of the critics in π to reduce the notational clutter. As in Hafner et al. (2019b), we use TD(λ) (Sutton & Barto, 2018) to trade-off the bias and variance of the critics with bootstrapping and Monte-Carlo value estimation. To learn the task and safety critics, we minimize the following loss function
Lvπψ (ψ) = Eπ,p(sτ:τ+H |sτ−1,aτ−1:τ+H−1,θ) [ 1
2H τ+H∑ t=τ ( vπψ(st)− Vλ(st)
)2] (9)
which uses the model to generate samples from p(sτ :τ+H |sτ−1,aτ−1:τ+H−1, θ). We denote vπψ as the neural network that approximates its corresponding value function and Vλ as the TD(λ) value as presented in Hafner et al. (2019b). Similarly, Viλ is the TD(λ) value of the i
th constraint ∀i ∈ {1, . . . , C}. Note that, while we show here only the loss for the task critic, the loss function for the safety critics is equivalent to Equation (9).
Policy optimization We model the policy as a Gaussian distribution via a neural network with parameters ξ such that πξ(at|st) = N (at; NNµξ (st),NNσξ (st)). We sample a sequence sτ :τ+H ∼ p(sτ :τ+H |sτ−1,aτ−1:τ+H−1, θ), utilizing πξ to generate actions on the fly (see Appendix I). To estimate J̃(π;λk, µk), we compute Vλ and Viλ for each state in the sampled sequence. This approximation leads to the following loss function for policy learning
Lπξ(ξ) = Eπξ,p(sτ:τ+H |sτ−1,aτ−1:τ+H−1,θ) [ 1
H τ+H∑ t=τ
−Vλ(st) ] +
C∑ i=1 Ψ i(πξ;λ i k, µk). (10)
We approximate J i(πξ) in Ψ i(πξ;λik, µk) (recall Equation (8)) as J i(πξ) ≈ Eπξ,p(sτ:τ+H |sτ−1,aτ−1:τ+H−1,θ) [ 1
H τ+H∑ t=τ Viλ(st)
] . (11)
Moreover, the first expression in the policy’s loss (10) approximates J(πξ). We backpropagate through p(sτ :τ+H |sτ−1,aτ−1:τ+H−1, θ) using path-wise gradient estimators (Mohamed et al., 2020).

heading: 4.2 ADOPTING OPTIMISM AND PESSIMISM TO EXPLORE IN CMDPS
text: Optimistic and pessimistic models As already noted by Curi et al. (2020a), greedily maximizing the expected returns by averaging over posterior samples in Equation (4) is not necessarily the best strategy for exploration. In particular, this greedy maximization does not deliberately leverage the epistemic uncertainty to guide exploration. Therefore, driven by the concepts of optimism in the face of uncertainty and upper confidence reinforcement learning (UCRL) (Auer et al., 2009; Curi et al., 2020a), we define a set of statistically plausible transition distributions denoted by P and let pθ ∈ P be a particular transition density in this set. Crucially, we assume that the true model p? is within the support of P , and that by sampling θ ∼ p(θ|D) the conditional predictive density satisfies p(st+1|st,at, θ) = pθ ∈ P . These assumptions lead us to the following constrained problem
max π∈Π max pθ∈P J(π, pθ) s.t. max pθi∈P J i(π, pθi) ≤ di, ∀i ∈ {1, . . . , C}. (12) The main intuition here is that jointly maximizing J(π, pθ) with respect to π and pθ can lead the agent to try behaviors with potentially high reward due to optimism. On the other hand, the pessimism that arises through the inner maximization term is crucial to enforce the safety constraints. Being only optimistic can easily lead to dangerous behaviors, while being pessimistic can lead the agent to not explore enough. Consequently, we conjecture that optimizing for an optimistic task objective J(π, pθ) combined with pessimistic safety constraints J i(π, pθi) allows the agent to explore better task-solving behaviors while being robust to model uncertainties.
Estimating the upper bounds We propose an algorithm that estimates the objective and constraints in Equation (12) through sampling, as illustrated in Figure 3. We demonstrate our method in Algorithm 1. Importantly, we present Algorithm 1 only with the notation of the task critic. However, we use it to approximate the model’s upper bound for each critic independently, i.e., for the costs as well as the reward value function critics.
Algorithm 1 Upper confidence bounds estimation via posterior sampling Require: N, p(θ|D), p(sτ :τ+H |sτ−1,aτ−1:τ+H−1, θ), sτ−1, π(at, |st).
1: Initialize V = {} # Set of objective estimates, under different posterior samples. 2: for j = 1 to N do 3: θ ∼ p(θ|D). # Posterior sampling (e.g., via SWAG). 4: sτ :τ+H ∼ p(sτ :τ+H |sτ−1,aτ−1:τ+H−1, θ) . # Sequence sampling, see Appendix I. 5: Append V ← V ∪∑τ+Ht=τ Vλ(st). 6: end for 7: return maxV .
The LAMBDA algorithm Algorithm 2 describes how all the previously shown components interact with each other to form a model-based policy optimization algorithm. For each update step, we sample a batch of B sequences with length L from a replay buffer to train the model. Then, we sampleN models from the posterior and use them to generate novel sequences with horizonH from every state in the replay buffer sampled sequences.
Algorithm 2 LAMBDA 1: Initialize D by following a random policy or from an offline dataset. 2: while not converged do 3: for u = 1 to U update steps do 4: Sample B sequences {(aτ ′−1:τ ′+L−1,oτ ′:τ ′+L, rτ ′:τ ′+L, ciτ ′:τ ′+L)} ∼ D uniformly. 5: Update model parameters θ and φ. # E.g., see Hafner et al. (2019a) for the RSSM. 6: Infer sτ ′:τ ′+L ∼ qφ(·|oτ :τ+L,aτ ′−1:τ ′+L−1). 7: Compute ∑τ+H t=τ Vλ(st), ∑τ+H t=τ V i λ(st) via Algorithm 1. Use each state in sτ ′:τ ′+L as an
initial state for sequence generation. 8: Update ψ and ψi via Equation (9) with ∑τ+H t=τ Vλ(st) and ∑τ+H t=τ V i λ(st).
9: Update ξ according to Equation (10) with ∑τ+H t=τ Vλ(st) and ∑τ+H t=τ V i λ(st).
10: Update λi via Equations (6) and (11). 11: end for 12: for t = 1 to T do 13: Infer st ∼ qφ(·|ot,at−1, st−1). 14: Sample at ∼ πξ(·|st). 15: Take action at, observe rt, cit,ot+1 received from the environment. 16: end for 17: Update dataset D ← D ∪ { o1:T ,a1:T , r1:T , c i 1:T } . 18: end while

heading: 5 EXPERIMENTS
text: We conduct our experiments with the SG6 benchmark as described by Ray et al. (2019), aiming to answer the following questions:
• How does our model-based approach compare to model-free variants in terms of performance, sample efficiency and constraint violation?
• What is the effect of replacing our proposed policy optimization method with an online planning method? More specifically, how does LAMBDA’s policy compare to CEM-MPC of Liu et al. (2021)?
• How does our proposed optimism-pessimism formulation compare to greedy exploitation in terms of performance, sample efficiency and constraint violation?
We provide an open-source code for our experiments, including videos of the trained agents at https://github.com/yardenas/la-mbda.

heading: 5.1 SG6 BENCHMARK
text: Experimental setup In all of our experiments with LAMBDA, the agent observes 64×64 pixels RGB images, taken from the robot’s point-of-view, as shown in Figure 2. Also, since there is only one safety constraint, we let Jc(π) ≡ J1(π). We measure performance with the following metrics, as proposed in Ray et al. (2019):
• Average undiscounted episodic return for E episodes: Ĵ(π) = 1E ∑E i=1 ∑Tep t=0 rt
• Average undiscounted episodic cost return for E episodes: Ĵc(π) = 1E ∑E i=1 ∑Tep t=0 ct
• Normalized sum of costs during training, namely the cost regret: for a given number of total interaction steps T , we define ρc(π) = ∑T t=0 ct T as the cost regret.
We compute Ĵ(π) and Ĵc(π) by averaging the sum of costs and rewards across E = 10 evaluation episodes of length Tep = 1000, without updating the agent’s networks and discarding the interactions made during evaluation. In contrast to the other metrics, to compute ρc(π), we sum the costs accumulated during training and not evaluation episodes. The results for all methods are recorded once the agent reached 1M environment steps for PointGoal1, PointGoal2, CarGoal1 and 2M steps
for PointButton1, PointPush2, DoggoGoal1 environments respectively. To reproduce the scores of TRPO-Lagrangian, PPO-Lagrangian and CPO, we follow the experimental protocol of Ray et al. (2019). We give further details on our experiment with CEM-MPC at the ablation study section.
Practical aspects on runtime The simulator’s integration step times are 0.002 and 0.004 seconds for the Point and Car robots, making each learning task run roughly 30 and 60 minutes in real life respectively. The simulator’s integration step time of the Doggo robot is 0.01 seconds thereby training this task in real life would take about 167 minutes. These results show that in principle, the data acquisition time in real life can be very short. However in practice, the main bottleneck is the gradient step computation which takes roughly 0.5 seconds on a single unit of Nvidia GeForceRTX2080Ti GPU. In total, it takes about 18 hours to train an agent for 1M interaction steps, assuming we take 100 gradient steps per episode. In addition, for the hyperparameters in Appendix C, we get a total of 12M simulated interactions used for policy and value function learning per episode.
Results The results of our experiments are summarized in Figure 4. As shown in Figure 4, LAMBDA is the only agent that satisfies the safety constraints in all of the SG6 tasks. Furthermore, thanks to its model-based policy optimization method, LAMBDA requires only 2M steps to successfully solve the DoggoGoal1 task and significantly outperform the other approaches. In Appendix B we examine further LAMBDA’s sample efficiency compared to the model-free baseline algorithms. Additionally, in the PointGoal2 task, which is denser and harder to solve in terms of safety, LAMBDA significantly improves over the baseline algorithms in all metrics. Moreover, in Figure 1 we compare LAMBDA’s ability to trade-off between the average performance and safety metrics across all of the SG6 tasks. One main shortcoming of LAMBDA is visible in the PointPush1 environment where the algorithm fails to learn to push to box to the goal area. We attribute this failure to the more strict partial observability of this task and further analyse it in Appendix D.

heading: 5.2 ABLATION STUDY
text: Unsafe LAMBDA As our first ablation, we remove the second term in Equation (10) such that the policy’s loss only comprises the task objective (i.e., with only optimistic exploration). We make the following observations: (1) Both LAMBDA and unsafe LAMBDA solve the majority of the SG6 tasks, depending on their level of partial observability; (2) in the PointButton1, PointGoal1 and CarGoal1, LAMBDA achieves the same performance of unsafe LAMBDA, while satisfying the safety constraints; (3) LAMBDA is able to reach similar performance to unsafe LAMBDA in the PointGoal2 task which is strictly harder than the other tasks in terms of safety, as shown in Figure 5. We note that partial observability is present in all of the SG6 tasks due to the restricted field of view
of the robot; i.e., the goal and obstacles are not always visible to the agent. In the DoggoGoal1 task, where the agent does not observe any information about the joint angles, LAMBDA is still capable of learning complex walking locomotion. However, in the PointPush1 task, LAMBDA struggles to find a task-solving policy, due to the harder partial observability of this task (see Appendix D). We show the learning curves for this experiment in Appendix E.
Ablating policy optimization Next, we replace our actor-critic procedure and instead of performing policy optimization, we implement the proposed MPC method of Liu et al. (2021). Specifically, for their policy, Liu et al. (2021) suggest to use the CEM to maximize rewards over action sequences while rejecting the unsafe ones. By utilizing the same world model, we are able to directly compare the planning performance for both methods. As shown in Figure 11, LAMBDA performs considerably better than CEM-MPC. We attribute the significant performance difference to the fact that the CEM-MPC approach does not use any critics for its policy, making its policy short-sighted and limits its ability to address the credit assignment problem.
Optimism and pessimism compared to greedy exploitation Finally, we compare our upper bounds estimation procedure with the more common approach of greedily maximizing the expected performance. More specifically, instead of employing Algorithm 1, we use Monte-Carlo sampling to estimate Equation (4) by generating trajectories with N = 5 sampled models and taking the average trajectory over the samples. By doing so, we get an estimate of the mean trajectory over sampled models together with the intrinsic stochasticity of the environment. We report our experiment results in Appendix G and note that LAMBDA is able to find safer and more performant policies than its greedy version.

heading: 6 CONCLUSIONS
text: We introduce LAMBDA, a Bayesian model-based policy optimization algorithm that conforms to human-specified safety constraints. LAMBDA uses its Bayesian world model to generate trajectories and estimate an optimistic bound for the task objective and pessimistic bounds for the constraints. For policy search, LAMBDA uses the Augmented Lagrangian method to solve the constrained optimization problem, based on the optimistic and pessimistic bounds. In our experiments we show that LAMBDA outperforms the baseline algorithms in the Safety-Gym benchmark suite in terms of sample efficiency as well as safety and task-solving metrics. LAMBDA learns its policy directly from observations in an end-to-end fashion and without prior knowledge. However, we believe that introducing prior knowledge with respect to the safety specifications (e.g., the mapping between a state and its cost) can improve LAMBDA’s performance. By integrating this prior knowledge, LAMBDA can potentially learn a policy that satisfies constraints only from its model-generated experience and without ever violating the constraints in the real world. This leads to a notable open question on how to integrate this prior knowledge within the world model such that the safety constraints are satisfied during learning and not only at the end of the training process.

heading: 7 ACKNOWLEDGMENTS AND DISCLOSURE OF FUNDING
text: This project has received funding from the European Research Council (ERC) under the European Union’s Horizon 2020 research and innovation programme grant agreement No 815943, the Swiss National Science Foundation under NCCR Automation, grant agreement 51NF40 180545 and the Swiss National Science Foundation, under grant SNSF 200021 172781.

heading: A LEARNING CURVES FOR THE SG6 BENCHMARK

heading: B SAMPLE EFFICIENCY
text: We record LAMBDA’s performance Ĵ(π) at the end of training and find the average number of steps (across seeds) required by the baseline model-free methods to match this performance. We demonstrate LAMBDA’s sample efficiency with the ratio of the amount of steps required by the baseline methods and amount of steps at the end of LAMBDA’s training. As shown in Figure 7, in the majority of tasks, LAMBDA is substantially more sample efficient. In the PointPush1 task LAMBDA is outperformed by the baseline algorithms, however we assign this to the partial observability of this task, as further analyzed in Appendix D. It is important to note that by taking LAMBDA’s performance at the end of training, we take a conservative approach, as convergence can occur with many less steps, as shown in Figure 6.
Required amount of steps to reach LAMBDA's performance

heading: C HYPERPARAMETERS
text: In this section we specify the most prominent hyperparameters for LAMBDA, however we encourage the reader to visit https://github.com/yardenas/la-mbda as it holds many specific (but important) implementation details. Table 1 summarizes the hyperparameters used in our algorithm.
World model As mentioned before, for our world model we use the RSSM proposed in Hafner et al. (2019a). Therefore, most of the architectural design choices (e.g., specific details of the convolutional layers) are based on the work there.
Cost function We exploit our prior knowledge of the cost function’s structure in the Safety-Gym tasks. Since we know a-priori that the cost function is implemented as an indicator function, our neural network approximation for it is modeled as a binary classifier. As training progresses, the dataset becomes less balanced as the agent observes less unsafe states. To deal with this issue, we give higher weight to the unsafe interactions in the binray classification loss.
Posterior over model parameters in supervised learning, SWAG is typically used under the assumption of a fixed dataset D, such that weight averaging takes place only during the last few training epochs. To adapt SWAG to the dataset’s distributional shifts that occur during training in RL, we use an exponentially decaying running average. This allows us to maintain a posterior even within early stages of learning, as opposed to the original design of SWAG that normally uses the last few iterates of SGD to construct the posterior. Furthermore, we use a cyclic learning rate (Izmailov et al., 2019; Maddox et al., 2019) to help SWAG span over different regions of the weight space.
Policy The policy outputs actions such that each element in at is within the range [−1, 1]. However, since we model πξ(at|st) as N (at; NNµξ (st),NNσξ (st)), we perform squashing of the Gaussian distribution into the range [−1, 1], as proposed in Haarnoja et al. (2019) and Hafner et al. (2019b) by transforming it through a tangent hyperbolic bijector. We scale the actions in this way so that the model, which takes the actions as an input is more easily optimized. Furthermore, to improve numerical stability, the standard deviation term NNσξ (st) is passed through a softplus function.
Value functions To ensure learning stability for the critics, we maintain a shadow instance for each value function which is used in the bootstrapping in Equation (9). We clone the shadow instance such that it lags U update steps behind its corresponding value function. Furthermore, we stop gradient on Vλ when computing the loss function in Equation (9).
General Parameter update for all of the neural networks is done with ADAM optimizer (Kingma & Ba, 2014). We use ELU (Clevert et al., 2016) as activation function for all of the networks except for the convolutional layers in the world model in which we use ReLU (Nair & Hinton, 2010).
SWAG
Burn-in steps 500 Steps before weight averaging starts. Period steps 200 Use weights every ‘Period steps’ to update weights running average. Models 20 Averaging buffer length. Decay 0.8 Exponential averaging decay factor. Cyclic LR factor 5.0 End cyclic lr period with the base LR times this factor. Posterior samples N 5
Safety
Safety critic learning rate 2e-4 Initial penalty µ0 5e-9 Initial Lagrangian λ10 1e-6 Penalty power factor 1e-5 Multiply µk by this factor at each gradient step. Safety discount factor 0.995
General
Update steps U 100 Critic learning rate 8e-5 Policy learning rate 8e-5 Action repeat 2 Repeat same action for this amount of steps. Discount factor 0.99 TD(λ) factor λ 0.95 Sequence generation horizon H 15

heading: D POINTPUSH1 ENVIRONMENT WITH A TRANSPARENT BOX
text: In all of our experiments with the PointPush1 task, we provide the agent image observations in which the box is observed as a solid non-transparant object. As previously shown, we note that in this setting, LAMBDA and unsafe LAMBDA fail to learn the task. We maintain that this failure arises from the fact that this task is significantly harder in terms of partial observability, as the goal is occluded by the box while the robot pushes the box.1 Furthermore, in Ray et al. (2019), the authors eliminate issues of partial observability by using what they term as “pseudo-LiDAR” which is not susceptible to occlusions as it can see through objects. To verify our hypothesis, we change the transparency of the box such that the goal is visible through it, as shown in Figure 8 and compare LAMBDA’s performance with the previously tested PointPush1 task. We present the learning curves of the experiments in Figure 9. As shown, LAMBDA is able to safely solve the PointPush1 task if
the goal is visible through the box. We conclude that the PointPush1 task makes an interesting test case for future research on partially observable environments.
1Please see https://github.com/yardenas/la-mbda for a video illustration.

heading: E UNSAFE LAMBDA
text: 0.0 0.5 1.0 1.5 2.0 ×106
0
5
10
A ve
ra g e
re w
a rd
re tu
rn
PointButton1
0.0 0.5 1.0 1.5 2.0 ×106
0
100
200
300
A ve
ra ge
co st
re tu
rn
0.0 0.5 1.0 1.5 2.0 ×106
0.1
0.2
C os
t re
g re
t
0.0 0.5 1.0 1.5 2.0 ×106
0
20
40
60
A ve
ra ge
re w
ar d
re tu
rn
DoggoGoal1
0.0 0.5 1.0 1.5 2.0 ×106
0
50
100
150
A ve
ra g e
co st
re tu rn 0.0 0.5 1.0 1.5 2.0
×106
0.04
0.06
0.08
C o st
re g re
t
0.0 0.2 0.4 0.6 0.8 1.0 ×106
0
5
10
15
20
A ve
ra g e
re w
a rd
re tu
rn
PointGoal1
0.0 0.2 0.4 0.6 0.8 1.0 ×106
0
50
100
A ve
ra g e
co st
re tu
rn
0.0 0.2 0.4 0.6 0.8 1.0 ×106
0.02
0.04
0.06
0.08
C o st
re g re
t
0.0 0.5 1.0 1.5 2.0 ×106
−5.0
−2.5
0.0
2.5
A ve
ra ge
re w
a rd
re tu
rn
PointPush1
0.0 0.5 1.0 1.5 2.0 ×106
0
50
100
A ve
ra ge
co st
re tu
rn
0.0 0.5 1.0 1.5 2.0 ×106
0.02
0.04
0.06 0.08 C os t re g re t
0.0 0.2 0.4 0.6 0.8 1.0 ×106
0
5
10
15
20
A ve
ra g e
re w
ar d
re tu
rn
CarGoal1
0.0 0.2 0.4 0.6 0.8 1.0 ×106
50
100
A ve
ra g e
co st
re tu
rn
0.0 0.2 0.4 0.6 0.8 1.0 ×106
0.025
0.050
0.075
0.100
0.125
C os
t re
gr et
0.0 0.2 0.4 0.6 0.8 1.0
Training steps ×106
0
5
10
15
20
A ve
ra ge
re w
ar d
re tu
rn
PointGoal2
0.0 0.2 0.4 0.6 0.8 1.0
Training steps ×106
0
100
200
A ve
ra ge
co st
re tu
rn
0.0 0.2 0.4 0.6 0.8 1.0
Training steps ×106
0.1
0.2
0.3
C os
t re
gr et
LAMBDA Unsafe LAMBDA
Figure 10: Benchmark results of LAMBDA and its unsafe implementation. In the majority of the tasks, LAMBDA is able to find policies that perform similarly to the unsafe version while satisfying the constraints. Interestingly, apart from the DoggoGoal1 and PointGoal2 tasks LAMBDA’s policies are able to achieve similar returns while learning to satisfy the constraints.

heading: F COMPARISON WITH CEM-MPC

heading: G OPTIMISM AND PESSIMISM COMPARED TO GREEDY EXPLOITATION

heading: H COMPARING ALGORITHMS IN THE SG6 BENCHMARK
text: To get a summary of how different algorithms behave across all of the SG6 tasks, we follow the proposed protocol in Ray et al. (2019). That is, we obtain “characteristic metrics” for each of the environments by taking the metrics recorded at the end of training of an unconstrained PPO agent (Schulman et al., 2017). We denote these metrics as ĴPPO, ĴPPOc and ρ̂ PPO c . We then normalize the recorded metrics for each environment as follows:
J̄(π) = Ĵ(π)
ĴPPO
J̄c(π) = max(0, Ĵc(π)− d)
max(10−6, ĴPPOc − d)
ρ̄c(π) = ρc(π)
ρPPOc .
(13)
By performing this normalization with respect to the performance of PPO, we can take an average of each metric across all of the SG6 environments.
To produce Figure 1, we scale the normalized metrics to [0, 1], such that for each metric, the best performing algorithm attains a score of 1.0 and the worst performing algorithm attains a score of 0.

heading: I BACKPROPAGATING GRADIENTS THROUGH A SEQUENCE
text: Algorithm 3 Sampling from the predictive density pθ(sτ :τ+H |sτ−1,aτ−1:τ+H−1, θ) Require: πξ(at|st), pθ(st+1|st,at), sτ−1
1: for t = τ − 1 to τ +H do 2: at ∼ πξ(·|stop gradient(st)) # Stop gradient from st when conditioning the policy on it. 3: st+1 ∼ p(·|st,at, θ) 4: end for 5: return sτ :τ+H
We use the reparametrization trick (Kingma & Welling, 2014) to compute gradients through sampling procedures as both πξ and pθ are modeled as normal distributions. Backpropagating gradients through the model can be easily implemented with modern automatic differentiation tools.
Importantly, we stop the gradient computation in Algorithm 3 when conditioning the policy on st−1. We do so to avoid any recurrent connections between an action at−1 and the preceding states to st such that eventually, backpropagation to actions occurs only from their dependant succeeding values. We further illustrate this in Figure 13.

heading: J SCORES FOR SG6
text: The .json format files, summarizing the scores for the experiments of this work are available at https://github.com/yardenas/la-mbda.

Review summaries:
The manuscript introduces LAMBDA a Bayesian model-based policy optimization algorithm that adheres to supplied safety constraints. The approach relies on the model to simulate trajectories and therefore improve efficiency of learning and effectiveness of safety. Experiments on SG6 compare the proposed algorithm to previous approaches and ablations of LAMDA. Based on experimental results, authors conclude that LAMDA is more efficient and effective for learning and safety.  

This paper is about solving CMDPs in an uncertainty-informed model-based fashion, without going through the usual LP route. The blueprints for this paper is the following:   1. Take a standard CMDP formulation 2. Express the objective and constraints in a UCRL (Auer, 2009) fashion (joint maximization over the set of possible dynamics), 3. Represent and learn the dynamics as a  Recurrent State Space Model from Hafner et al. (2019) 4. Estimate the objective and constraints themselves  (which involve a max operator) by sampling, and where the Bayesian posterior over the "world model" parameters is computed by SWAG (Maddox et al. 2019), where you average out the iterates generated over a trajectory of stochastic gradient descent  5. Using 3 and 4, compute a step of Augmented Lagrangian (a variant of) to solve the CMDP in 2.  The authors provide motivation as to why the UCRL aspect is important, and how the underlying optimism principle gives rise to persimism in the constraints, which is important from a "safety" perspective. This gives rise to an interesting tension (perhaps more "balance") between optimism in maximizing the objective while remaining cautions in abiding to the constraints.  

The authors present LAMBDA, a novel algorithm for learning in CMDPs. The approach seeks to improve on the sample efficiency of previous algorithms by learning a probabilistic model of the environment. This model is then used to train a policy to maximise the return and respect the constraints by backpropagating through imaginary trajectories. The uncertainty of the model can further be used to steer the agent towards for optimistic or pessimistic behavior. The method is benchmarked against one online planning method and previous model-free approaches on the SG6 set of tasks of the Safety Gym environment. The results show that the proposed method compares favourably to competing algorithms, particularly in terms of constraint satisfaction and sample efficiency.

This paper proposed a method called LAMBDA for optimizing an agent's policy in constrained Markov decision processes. LAMBDA is a model-based approach leveraging Bayesian world models, which deals with reward optimistically and safety pessimistically. The authors demonstrate the effectiveness of SG6 tasks in Safety-Gym in terms of sample efficiency and satisfaction of safety constraint(s).