Wasserstein Robust Reinforcement Learning

1 INTRODUCTION . Reinforcement learning ( RL ) has become a standard tool for solving decision-making problems with feedback , and though significant progress has been made , algorithms often over-fit to training environments and fail to generalise across even slight variations of transition dynamics ( Packer et al. , 2018 ; Zhao et al. , 2019 ) . Robustness to changes in transition dynamics is a crucial component for adaptive and safe RL in real-world environments . Motivated by real-world applications , recent literature has focused on the above problems , proposing a plethora of algorithms for robust decisionmaking ( Morimoto & Doya , 2005 ; Pinto et al. , 2017 ; Tessler et al. , 2019 ) . Most of these techniques borrow from game theory to analyse , typically in a discrete state and actions spaces , worst-case deviations of agents ’ policies and/or environments , see Sargent & Hansen ( 2001 ) ; Nilim & El Ghaoui ( 2005 ) ; Iyengar ( 2005 ) ; Namkoong & Duchi ( 2016 ) and references therein . These methods have also been extended to linear function approximators ( Chow et al. , 2015 ) , and deep neural networks ( Peng et al. , 2017 ) showing ( modest ) improvements in performance gain across a variety of disturbances , e.g. , action uncertainties , or dynamical model variations . In this paper , we propose a generic framework for robust reinforcement learning that can cope with both discrete and continuous state and actions spaces . Our algorithm , termed Wasserstein Robust Reinforcement Learning ( WR2L ) , aims to find the best policy , where any given policy is judged by the worst-case dynamics amongst all candidate dynamics in a certain set . This set is essentially the average Wasserstein ball around a reference dynamics P0 . The constraints makes the problem well-defined , as searching over arbitrary dynamics can only result in non-performing system . The measure of performance is the standard RL objective , the expected return . Both the policy and the dynamics are parameterised ; the policy parameters θk may be the weights of a deep neural network , and the dynamics parameters φj the parameters of a simulator or differential equation solver . The algorithm performs estimated descent steps in φ space and - after ( almost ) convergence - performs an update of policy parameters , i.e. , in θ space . Since φj may be high-dimensional , we adapt a zero ’ th order sampling method based extending Salimans et al . ( 2017 ) to make estimations of gradients , and in order to define the constraint set which φj is bounded by , we generalise the technique to estimate Hessians ( Proposition 2 ) . We emphasise that although access to a simulator with parameterisable dynamics is required , the actual reference dynamics P0 need not be known explicitly nor learnt by our algorithm . Put another way , we are in the “ RL setting ” , not the “ MDP setting ” where the transition probability matrix is known a priori . The difference is made obvious , for example , in the fact that we can not perform dynamic programming , and the determination of a particular probability transition can only be estimated from sampling , not retrieved explicitly . Hence , our algorithm is not model-based in the traditional sense of learning a model to perform planning . We believe our contribution is useful and novel for two main reasons . Firstly , our framing of the robust learning problem is in terms of dynamics uncertainty sets defined by Wasserstein distance . Whilst we are not the first to introduce the Wasserstein distance into the context of MDPs ( see , e.g. , Yang ( 2017 ) or Lecarpentier & Rachelson ( 2019 ) ) , we believe our formulation is amongst the first suitable for application to the demanding application-space we desire , that being , high-dimensional , continuous state and action spaces . Secondly , we believe our solution approach is both novel and effective ( as evidenced by experiments below , see Section 5 ) , and does not place a great demand on model or domain knowledge , merely access to a simulator or differentiable equation solver that allows for the parameterisation of dynamics . Furthermore , it is not computationally demanding , in particular , because it does not attempt to build a model of the dynamics , and operations involving matrices are efficiently executable using the Jacobian-vector product facility of automatic differentiation engines . 2 BACKGROUND . A Markov decision process ( MDP ) 1 is denoted byM = 〈S , A , P , R , γ〉 , where S ⊆ Rd denotes the state space , A ⊆ Rn the action space , P : S × A × S → [ 0 , 1 ] is a state transition probability describing the system ’ s dynamics , R : S × A → R is the reward function measuring the agent ’ s performance , and γ ∈ [ 0 , 1 ) specifies the degree to which rewards are discounted over time . At each time step t , the agent is in state st ∈ S and must choose an action at ∈ A , transitioning it to a new state st+1 ∼ P ( st+1|st , at ) , and yielding a rewardR ( st , at ) . A policy π : S×A→ [ 0 , 1 ] is defined as a probability distribution over state-action pairs , where π ( at|st ) represents the density of selecting action at in state st . Upon subsequent interactions with the environment , the agent collects a trajectory τ of state-action pairs . The goal is to determine an optimal policy π ? by solving : π ? = argmax π Eτ∼pπ ( τ ) [ RTotal ( τ ) ] , ( 1 ) where pπ ( τ ) denotes the trajectory density function , andRTotal ( τ ) the return , that is , the total accumulated reward : pπ ( τ ) = µ0 ( s0 ) π ( a0|s0 ) T−1∏ t=1 P ( st+1|st , at ) π ( at|st ) and RTotal ( τ ) = T−1∑ t=0 γtR ( st , at ) , ( 2 ) with µ0 ( · ) denoting the initial state distribution . We make use of the Wasserstein distance to quantify variations from a reference transition density P0 ( · ) . The latter being a probability distribution , one may consider other divergences , such as Kullback-Leibler ( KL ) or total variation ( TV ) . Our main intuition for choosing Wasserstein distance is explained below , but we note that it has a number of desirable properties : Firstly , it is symmetric ( Wp ( µ , ν ) = Wp ( ν , µ ) , a property K-L lacks ) . Secondly , it is well-defined for measures with different supports ( which K-L also lacks ) . Indeed , the Wasserstein distance is flexible in the forms of the measures that can be compared - discrete , continuous or a mixture . Finally , it takes into account the underlying geometry of the space the distributions are defined on , which can encode valuable information . It is defined as follows : Let X be a metric space with metric d ( · , · ) . Let C ( X ) be the space of continuous functions on X and letM ( X ) be the set of probability measures on X . Let µ , ν ∈M ( X ) . Let K ( µ , ν ) be the set of couplings between µ , ν : K ( µ , ν ) : = { κ ∈M ( X × X ) ; ∀ ( A , B ) ⊂ X × X , κ ( A×X ) = µ ( A ) , κ ( X ×B ) = ν ( B ) } ( 3 ) That is , the set of joint distributions κ ∈ M ( X × X ) whose marginals agree with µ and ν respectively . Given a metric ( serving as a cost function ) d ( · , · ) for X , the p ’ th Wasserstein distance Wp ( µ , ν ) for p ≥ 1 between µ and ν is defined as : Wp ( µ , ν ) : = ( min κ∈K ( µ , ν ) ∫ X×Y d ( x , y ) pdκ ( x , y ) ) 1/p ( 4 ) ( in this paper , and mostly for computational convenience , we use p = 2 , though other values of p are applicable ) . 1Please note that we present reinforcement learning with continuous states and actions . This allows us to easily draw similarities to optimal control as detailed later . Extending these notions to discrete settings is relatively straight-forward . 3 WASSERSTEIN ROBUST REINFORCEMENT LEARNING . The desirable properties of Wasserstein distance aside , our main intuition for choosing it is described thus : Per the definition , constraining the possible dynamics to be within an -Wasserstein ball of a reference dynamics P0 ( · ) means constraining it in a certain way . Wasserstein distance has the form mass× distance . If this quantity is constrained to be less than a constant , then if the mass is large , the distance is small , and if the distance is large , the mass is small . Intuitively , when modelling the dynamics of a system , it may be reasonable to concede that there could be a systemic error - or bias - in the model , but that bias should not be too large . It is also reasonable to suppose that occasionally , the behaviour of the system may be wildly different to the model , but that this should be a low-probability event . If the model is frequently wrong by a large amount , then there is no use in it . In a sense , the Wasserstein ball formalises these assumptions . 3.1 PROBLEM DEFINITION : ROBUST OBJECTIVES AND CONSTRAINTS . Due to the continuous nature of the state and action spaces considered in this work , we resort to deep neural networks to parameterise policies , which we write as πθ ( at|st ) , where θ ∈ Rd1 is a set of tunable hyper-parameters to optimise . For instance , these policies can correspond to multilayer perceptrons for MuJoCo environments , or to convolutional neural networks in case of highdimensional states depicted as images . Exact policy details are ultimately application dependent and , consequently , provided in the relevant experiment sections . In principle , one can similarly parameterise dynamics models using deep networks ( e.g. , LSTMtype models ) to provide one or more action-conditioned future state predictions . Though appealing , going down this path led us to agents that discover worst-case transition models which minimise training data but lack any valid physical meaning . For instance , original experiments we conducted on CartPole ended up involving transitions that alter angles without any change in angular velocities . More importantly , these effects became more apparent in high-dimensional settings where the number of potential minimisers increases significantly . It is worth noting that we are not the first to realise such an artifact when attempting to model physics-based dynamics using deep networks . Authors in ( Lutter et al. , 2019 ) remedy these problems by introducing Lagrangian mechanics to deep networks , while others ( Koller et al. , 2018 ; Chen et al. , 2018 ) argue the need to model dynamics given by differential equation structures directly . Though incorporating physics-based priors to deep networks is an important and challenging task that holds the promise of scaling model-based reinforcement learning for efficient solvers , in this paper we rather study an alternative direction focusing on perturbing differential equation solvers and/or simulators with respect to the dynamic specification parameters φ ∈ Rd2 . Not only would such a consideration reduce the dimension of parameter spaces representing transition models , but would also guarantee valid dynamics due to the nature of the simulator . Though tackling some of the above problems , such a direction arrives with a new set of challenges related to computing gradients and Hessians of black-box solvers . In Section 4 , we develop an efficient and scalable zero-order method for valid and accurate model updates . Unconstrained Loss Function : Having equipped agents with the capability of representing policies and perturbing transition models , we are now ready to present an unconstrained version of WR2L ’ s loss function . Borrowing from robust optimal control , we define robust reinforcement learning as an algorithm that learns best-case policies under worst-case transitions : max θ [ min φ Eτ∼pφθ ( τ ) [ Rtotal ( τ ) ] ] , ( 5 ) where pφθ ( τ ) is a trajectory density function parameterised by both policies and transition models , i.e. , θ and φ , respectively : pφθ ( τ ) = µ0 ( s0 ) π ( a0|s0 ) T−1∏ t=1 Pφ ( st+1|st , at ) ︸ ︷︷ ︸ specs vector and diff . solver πθ ( at|st ) ︸ ︷︷ ︸ deep network . At this stage , it should be clear that our formulation , though inspired from robust optimal control , is , truthfully , more generic as it allows for parameterised classes of transition models without incorporating additional restrictions on the structure or the scope by which variations are executed2 . Constraints & Complete Problem Definition : Clearly , the problem in Equation 5 is ill-defined due to the arbitrary class of parameterised transitions . To ensure well-behaved optimisation objectives , we next introduce constraints to bound search spaces and ensure convergence to feasible transition models . For a valid constraint set , our method assumes access to samples from a reference dynamics model P0 ( ·|s , a ) , and bounds learnt transitions in an -Wasserstein ball around P0 ( ·|s , a ) , i.e. , the set defined as : W ( Pφ ( · ) , P0 ( · ) ) = { Pφ ( · ) : W22 ( Pφ ( ·|s , a ) , P0 ( ·|s , a ) ) ≤ , ∀ ( s , a ) ∈ S ×A } , ( 6 ) where ∈ R+ is a hyperparameter used to specify the “ degree of robustness ” in a similar spirit to maximum norm bounds in robust optimal control . It is worth noting , that though we have access to samples from a reference simulator , our setting is by no means restricted to model-based reinforcement learning in an MDP setting . That is , our algorithm operates successfully given only traces from P0 accompanied with its specification parameters , e.g. , pole lengths , torso masses , etc . – a more flexible framework that does not require full model learners . Though defining a better behaved optimisation objective , the set in Equation 6 introduces infinite number of constraints when considering continuous state and/or actions spaces . To remedy this problem , we target a relaxed version that considers a constraint of average Wasserstein distance bounded by a hyperparameter : Ŵ ( average ) ( Pφ ( · ) , P0 ( · ) ) = { Pφ ( · ) : ∫ ( s , a ) P ( s , a ) W22 ( Pφ ( ·|s , a ) , P0 ( ·|s , a ) ) d ( s , a ) ≤ } ( 7 ) = { Pφ ( · ) : E ( s , a ) [ W22 ( Pφ ( ·|s , a ) , P0 ( ·|s , a ) ) ] ≤ } The sampling ( s , a ) in the expectation is done as follows : We sample trajectories using reference dynamics P0 and a policy π that chooses actions uniformly at random ( uar ) . Then ( s , a ) pairs are sampled uar from those collected trajectories . For a given pair ( s , a ) , W22 ( Pφ ( ·|s , a ) , P0 ( ·|s , a ) ) is approximated through the empirical distribution : we use the state that followed ( s , a ) in the collected trajectories as a data point . Estimating Wasserstein distance using empirical data is standard , see , e.g. , Peyré et al . ( 2019 ) . One approach which worked well in our experiments , was to assume that the dynamics are given by deterministic functions plus Gaussian noise with diagonal convariance matrices . This makes estimation easier in high dimensions since sampling in each dimension is independent of others , and the total samples needed is a constant factor of the number of dimensions . Gaussian distributions also have closed-form expressions for Wasserstein distance , given in terms of mean and covariance . We thus arrive at WR2L ’ s optimisation problem allowing for best policies under worst-case yet bounded transition models : Wasserstein Robust Reinforcement Learning Objective : max θ [ min φ Eτ∼pφθ ( τ ) [ Rtotal ( τ ) ] ] s.t . E ( s , a ) [ W22 ( Pφ ( ·|s , a ) , P0 ( ·|s , a ) ) ] ≤ ( 8 )

Summary: The paper studies the problem of training policies which maximize returns under worst case environmental perturbations. They propose to solve this max min problem using iterative optimization with the max step performed using standard RL. The min step relies on sample-based estimates of gradients and hessians with respect to environment parameters. The problem of robust control is an important one, and the paper presents interesting results, however it lacks a proper comparison to baselines.