C ONSTRAINED P OLICY O PTIMIZATION VIA B AYESIAN W ORLD M ODELS
A BSTRACT
Improving sample-efﬁciency 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 efﬁciency and constraint violation.
I NTRODUCTION
[Figure]
A M
[Figure]
D A
[Figure]
C
[Figure]
C PP O [Figure] 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 efﬁciently 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 efﬁciency 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:
[Figure]
• 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 ﬁrst-person-view observations as illustrated in Figure 2.
Furthermore, we show that LAMBDA outperforms other model-based and model-free methods in this benchmark.
R ELATED WORK
Interpretations of safety in RL research We ﬁrst acknowledge that there exist different deﬁnitions for safety (Garc ´ ıa et al., 2015).
One deﬁnition uses reversible Markov decision processes (Moldovan & Abbeel, 2012) that, informally, deﬁne safety as reachability between states and use backup policies (Eysenbach et al., 2017) to ensure safety.
Another deﬁnition 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 deﬁne the safety requirements as a cost function that should be bounded by a predeﬁned threshold.
Similarly, Achiam et al.
; Dalal et al.
Ray et al.
(2019); Chow et al.
; Stooke et al.
; Bharadhwaj et al.
; 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 efﬁciency 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 conﬁdence 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.
; 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 inﬂuence 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 ﬁnd 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.
P RELIMINARIES 3. S AFE MODEL - BASED REINFORCEMENT LEARNING
Markov decision processes We consider an episodic Markov decision process with discrete time steps t ∈ { 0 , . .
, T } .
We deﬁne the environment’s state as s t ∈ R n and an action taken by the agent, as a t ∈ R m .
Each episode starts by sampling from the initial-state distribution s 0 ∼ ρ ( s 0 ) .
After observing s 0 and at each step thereafter, the agent takes an action by sampling from a policy distribution a t ∼ π ( ·| s t ) .
The next state is then sampled from an unknown transition distribution s t +1 ∼ p ( ·| s t , a t ) .
Given a state, the agent observes a reward, generated by r t ∼ p ( ·| s t , a t ) .
We deﬁne the performance of a pair of policy π and dynamics p as
[Equation]
Constrained Markov decision processes To make the agent adhere to human-deﬁned safety constraints, we adopt the constrained Markov decision process formalism of Altman (1999).
In CMPDs, alongside with the reward, the agent observes cost signals c it generated by c it ∼ p ( ·| s t , a t ) where i = 1 , . . .
, C denote the distinct unsafe behaviors we want the agent to avoid.
Given c it , we deﬁne the constraints as
[Equation]
where d i are human-deﬁned thresholds.
For example, a common cost function is c
i ( s t )
= 1 s t ∈H i , where H i 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 ﬁnd a policy that solves, for the true unknown dynamics p (cid:63) ,
[Equation]
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 ofﬂine dataset) of { s t +1 , s t , a t } into a dataset
D .
Following that, the dataset is used to ﬁt a statistical model p ( s t +1 | s t , a t , θ ) that approximates the true transition distribution p (cid:63) .
Finally, the agent uses the statistical model for planning, either within an online MPC scheme or via ofﬂine policy optimization.
In this work, we consider the cost and reward functions as unknown and similarly to the transition distribution, we ﬁt 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 , θ ) = (cid:81) τ
+ H t = τ p ( s t +1 | s t , a t , θ ) whereby H is a predeﬁned 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 efﬁciency compared to its model-free counterparts.
We emphasize that in safetycritical tasks, where human supervision during training might be required, sample efﬁciency 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 o t ∼ p ( ·| s t ) instead of s t at each time step.
To infer the transition density from ob- servations, 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 ﬁlter 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)).
3. A B AYESIAN VIEW ON MODEL - BASED REINFORCEMENT LEARNING
The Bayesian predictive distribution A common challenge in MBRL is that policy optimization can “overﬁt” 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 quantiﬁcation 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
[Equation]
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 ) .
Brieﬂy, 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 ﬂexible and admits other variants of approximate Bayesian inference, such as variational techniques (Graves, 2011; Kingma & Welling, 2014).
L AGRANGIAN M ODEL - BASED A GENT (LAMBDA)
We ﬁrst 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 ﬁnding the optimistic and pessimistic models with which the algorithm solves the CMDP before collecting more data.
4. S OLVING CMPD S WITH MODEL - BASED CONSTRAINED POLICY OPTIMIZATION
The Augmented Lagrangian We ﬁrst consider the optimization problem in Equation (3).
To ﬁnd 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
[Equation]
−∞ 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 ( π ) ≤ d i ∀ i ∈
, C } and so the maximum over λ
i is satisﬁed 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:
[Equation]
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:
[Equation]
otherwise .
[Equation]
We take gradient steps of the following unconstrained objective:
[Equation]
where
[Equation]
Task and safety critics For the task and safety critics, we use the reward value function v ( s t ) and the cost value function of each constraint v i ( s t ) ∀ i ∈ { 1 , . . .
, C } .
We model v ( s t ) and v i ( s t ) 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
[Equation]
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, V iλ 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 π ξ ( a t | s t ) =
N ( a t ; NN µξ ( s t ) , NN σξ ( s t )) .
We sample a sequence s τ : τ + H ∼ p ( s τ : τ + H | s τ − 1 , a τ − 1: τ + H − 1 , θ ) , utilizing π ξ to generate actions on the ﬂy (see Appendix I).
To estimate ˜ J ( π ; λ k , µ k ) , we compute V λ and V iλ for each state in the sampled sequence.
This approximation leads to the following loss function for policy learning
[Equation]
We approximate J i ( π ξ ) in Ψ i ( π ξ ; λ ik , µ k ) (recall Equation (8)) as
[Equation]
Moreover, the ﬁrst 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).
[Figure]
4. A DOPTING OPTIMISM AND PESSIMISM TO EXPLORE IN CMDP S
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 conﬁdence reinforcement learning (UCRL)
(Auer et al., 2009; Curi et al., 2020a), we deﬁne 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 (cid:63) is within the support of P , and that by sampling θ ∼ p ( θ |D )
the conditional predictive density satisﬁes p ( s t +1 | s t , a t , θ ) =
p θ ∈ P .
These assumptions lead us to the following constrained problem
[Equation]
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 conﬁdence bounds estimation via posterior sampling
Require: N, p ( θ |D ) , p ( s τ : τ + H | s τ − 1 , a τ − 1: τ
+ H − 1 , θ ) , s τ − 1 , π ( a t , | s t ) .
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.
Append V ← V ∪ (cid:80) τ
+ H t
=
τ V λ ( s t ) .
6: end for 7: return max V .
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 sample N models from the posterior and use them to generate novel sequences with horizon H from every state in the replay buffer sampled sequences.
1: Initialize D by following a random policy or from an ofﬂine dataset.
2: while not converged do 3: for u = 1 to U update steps do 4:
Sample uniformly at random B sequences { ( a τ (cid:48) − 1: τ (cid:48)
+ L − 1 , o τ (cid:48) : τ (cid:48)
+ L , r τ (cid:48) :
τ (cid:48)
+ L ) } ∼ D . 5:
Update model parameters θ and φ .
# E.g., see Hafner et al. (2019a) for the RSSM.
Infer s τ (cid:48) : τ (cid:48)
+ L ∼ q φ ( ·| o τ : τ + L , a τ (cid:48) − 1: τ (cid:48)
+ L − 1 ) .
Compute (cid:80) τ
+ H t
=
τ V λ ( s t ) , (cid:80) τ
+ H t =
τ V iλ ( s t ) via Algorithm 1.
Use each state in s τ (cid:48) : τ (cid:48)
+ L as an initial state for sequence generation.
Update ψ and ψ i via Equation (9) with (cid:80) τ
+ H t
=
τ V λ ( s t ) and (cid:80) τ
+ H t =
τ V iλ ( s t ) .
Update ξ according to Equation (10) with (cid:80) τ
+ H t
=
τ V λ ( s t ) and (cid:80) τ
+ H t =
τ V iλ ( s t ) .
Update λ i via Equations (6) and (11).
11: end for 12: for t = 1 to T do 13:
Infer s t ∼ q φ ( ·| o t , a t − 1 , s t − 1 ) .
Sample a t ∼ π ξ ( ·| s t ) .
Take action a t , observe r t , c i t , o t +1 received from the environment.
16: end for 17:
Update dataset D ← D ∪ (cid:8) o 1: T , a 1: T , r 1: T , c i 1: T (cid:9) .
18: end while
E XPERIMENTS
We conduct our experiments with the SG6 benchmark as described by Ray et al.
, aiming to answer the following questions:
• How does our model-based approach compare to model-free variants in terms of performance, sample efﬁciency and constraint violation?
• What is the effect of replacing our proposed policy optimization method with an online planning method?
More speciﬁcally, how does LAMBDA’s policy compare to CEM-MPC of Liu et al.
• How does our proposed optimism-pessimism formulation compare to greedy exploitation in terms of performance, sample efﬁciency 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 .
5. SG6 B ENCHMARK
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 J c ( π ) ≡ J 1 ( π ) .
We measure performance with the following metrics, as proposed in Ray et al.
• Average undiscounted episodic return for E episodes: ˆ J ( π ) = 1 E (cid:80)
Ei =1 (cid:80) T ep t =
r t • Average undiscounted episodic cost return for E episodes: ˆ J c ( π ) = 1 E (cid:80)
Ei =1 (cid:80) T ep t =
c t • Normalized sum of costs during training, namely the cost regret : for a given number of total interaction steps T , we deﬁne ρ c ( π ) =
(cid:80) Tt =
c t 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 T ep = 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 PointGoal PointGoal
CarGoal
PointButton PointPush DoggoGoal
4 ¯ J ( π )
SG6 (average) 0 .0 .0 .0 .
6 ¯ J c
( π ) 0 .0 .1 .1 .¯
ρ c ( π )
LAMBDA
CEM-MPC
PPO-Lagrangian
TRPO-Lagrangian
CPO 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.
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 satisﬁes 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 signiﬁcantly outperform the other approaches.
In Appendix B we examine further LAMBDA’s sample efﬁciency compared to the model-free baseline algorithms.
Additionally, in the PointGoal2 task, which is denser and harder to solve in terms of safety, LAMBDA signiﬁcantly 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.
A BLATION STUDY
Unsafe LAMBDA As our ﬁrst 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, PointGoal 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 ﬁeld of view
[Figure]
t
[Figure]
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 ﬁnd 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.
Speciﬁcally, 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 signiﬁcant 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 speciﬁcally, 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 ﬁnd safer and more performant policies than its greedy version.
C ONCLUSIONS
We introduce LAMBDA, a Bayesian model-based policy optimization algorithm that conforms to human-speciﬁed 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 efﬁciency 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 speciﬁcations (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 satisﬁes 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 satisﬁed during learning and not only at the end of the training process.
A CKNOWLEDGMENTS AND D ISCLOSURE OF F UNDING
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.
A L EARNING CURVES FOR THE SG6 BENCHMARK
[Figure]
PointButton
[Figure]
PointGoal
[Figure]
CPO (proprio
[Figure]
B S AMPLE EFFICIENCY
We record LAMBDA’s performance ˆ J ( π ) at the end of training and ﬁnd the average number of steps (across seeds) required by the baseline model-free methods to match this performance.
We demonstrate LAMBDA’s sample efﬁciency 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 efﬁcient.
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
[Figure]
Required amount of steps to reach [Figure] s
[Figure]
C H YPERPARAMETERS
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 speciﬁc (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., speciﬁc 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 classiﬁer.
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 classiﬁcation loss.
Posterior over model parameters in supervised learning, SWAG is typically used under the assumption of a ﬁxed 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.
The policy outputs actions such that each element in a t is within the range [ − 1 , 1] .
However, since we model π ξ ( a t | s t ) as N ( a t ; NN µξ ( s t ) , NN σξ ( s t )) , we perform squashing of the Gaussian distribution into the range [ − 1 , 1] , as proposed in Haarnoja et al.
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 σξ ( s t ) 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).
[Table]
[Figure]
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 signiﬁcantly harder in terms of partial observability, as the goal is occluded by the box while the robot pushes the box.
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
[Figure]
A v e r ag e r e w a r d r e t u r n
[Figure]
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.
[Figure]
t r e g r e t
LAMBDA
Unsafe LAMBDA
[Figure]
[Table]
[Figure] [Table] [Figure] [Table]
[Figure]
[Table] [Figure] [Table] [Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
PointButton1PointButton
[Figure]
G O PTIMISM AND PESSIMISM COMPARED TO GREEDY EXPLOITATION
[Figure]
PointButton1PointButton
[Figure]
H C OMPARING ALGORITHMS IN THE SG6 BENCHMARK
To get a summary of how different algorithms behave across all of the SG6 tasks, we follow the proposed protocol in Ray et al.
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 PPO c and ˆ ρ PPO c .
We then normalize the recorded metrics for each environment as follows:
[Equation]
[Equation]
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.
I B ACKPROPAGATING GRADIENTS THROUGH A SEQUENCE
p θ ( s τ : τ + H | s τ − 1 , a τ − 1: τ
+ H − 1 , θ )
Require: π ξ ( a t | s t ) , p θ ( s t +1 | s t , a t ) , s τ − 1 1: for t = τ − 1 to τ + H do 2: a t ∼ π ξ ( ·| stop gradient ( s t ))
# Stop gradient from s t when conditioning the policy on it.
3: s t +1 ∼ p ( ·| s t , a t , θ ) 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 s t − 1 .
We do so to avoid any recurrent connections between an action a t − 1 and the preceding states to s t such that eventually, backpropagation to actions occurs only from their dependant succeeding values.
We further illustrate this in Figure 13.
[Figure]
[Figure] [Figure] [Figure]
[Figure]
J S CORES FOR SG
The .json format ﬁles, summarizing the scores for the experiments of this work are available at https://github.com/yardenas/la-mbda .
[Table]
We present the results with the tuple ( ¯ J ( π ) , ¯ J c ( π ) , ¯ ρ c ( π )) of the normalized metrics.
[Table]