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 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.