Constrained Policy Optimization via Bayesian World Models

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 . 2 RELATED WORK . 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 . 3 PRELIMINARIES . 3.1 SAFE MODEL-BASED REINFORCEMENT LEARNING . 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 ) ) .

Summary: 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).