Contextual Inverse Reinforcement Learning

1 INTRODUCTION . We study sequential decision-making in a Contextual Markov Decision Process ( CMDP , Hallak et al . ( 2015 ) ) , where the reward , while unknown to the agent , depends on a static parameter referred to as the context . For a concrete example , consider the dynamic treatment regime ( Chakraborty & Murphy , 2014 ) . Here , there is a patient and a clinician which acts to improve the patient ’ s health . The context is composed of static information of the patient ( such as age and weight ) ; the state is composed of the patient ’ s dynamic measurements ( such as heart rate and blood pressure ) ; and the clinician ’ s actions are a set of intervention categories ( e.g. , infusion ) . The reward is different for each patient ( context ) , and there is a mapping from the context to the reward . Recent trends in personalized medicine motivate this model – instead of treating the ” average patient ” , patients are separated into different groups for which the medical decisions are tailored ( Fig . 1b ) . For example , in Wesselink et al . ( 2018 ) , the authors study organ injury , which may occur when a specific measurement ( mean arterial pressure ) decreases below a certain threshold . They found that this threshold varies across different patient groups ( context ) . In other examples , clinicians set treatment goals for the patients , i.e. , they take actions to make the patient measurements reach some pre-determined values . For instance , in acute respiratory distress syndrome ( ARDS ) , clinicians argue that these treatment goals should depend on the static patient information ( the context ) ( Berngard et al. , 2016 ) . There are serious issues when trying to manually define a reward signal in real-world tasks . When treating patients with sepsis , for example , the only available signal is the mortality of the patient at the end of the treatment ( Komorowski et al. , 2018 ) . This signal is sparse , and it is unclear how to manually tweak the reward to maximize the patient ’ s health condition ( Leike et al. , 2017 ; Raghu et al. , 2017 ; Lee et al. , 2019 ) . To address these issues , we propose the Contextual Inverse Reinforcement Learning ( COIRL ) framework . Similarly to Inverse Reinforcement Learning ( Ng & Russell , 2000 , IRL ) , we focus on trying to infer the mapping from contexts to rewards by observing experts . The main challenge in our problem is that for each context there is a different reward , hence , a different optimal policy for each context . Therefore , Apprenticeship Learning algorithms ( Abbeel & Ng , 2004 ; Syed & Schapire , 2008 ) that try to mimic the expert can not be used and , instead , we focus on directly learning the mapping . In particular , our main contributions are : 1 . We formulate COIRL with a linear mapping as a convex optimization problem . 2 . We propose and analyze the sample complexity of three algorithms for COIRL : the mirrored descent alg . ( MDA ) , evolution strategies ( ES ) , and the ellipsoid method . 3 . For nonlinear mappings , we implement a deep learning version for MDA and ES ( without theoretical guarantees ) . 4 . We compare these methods empirically on two frameworks : an autonomous driving simulator ( Abbeel & Ng , 2004 ) and a dynamic treatment regime ( Komorowski et al. , 2018 ) . 2 PRELIMINARIES . Contextual MDPs : A Markov Decision Process ( Puterman , 1994 , MDP ) is defined by the tuple ( S , A , P , ξ , R , γ ) where S is a finite state space , A a finite action space , P : S × S ×A→ [ 0 , 1 ] the transition kernel , ξ the initial state distribution , R : S → R the reward function and γ ∈ [ 0 , 1 ) is the discount factor . A Contextual MDP ( Hallak et al. , 2015 , CMDP ) is an extension of an MDP , and is defined by ( C , S , A , M , γ ) where C is the context space , andM is a mapping from contexts c ∈ C to MDPs : M ( c ) = ( S , A , P , Rc , ξ , γ ) . In addition , each state is associated with a feature vector φ : S → [ 0 , 1 ] k. Note that P and ξ are not context dependent . We consider a setting in which the reward for context c is a linear combination of the state features : R∗c ( s ) = f ∗ ( c ) Tφ ( s ) . The goal is to approximate f∗ ( c ) using a function fW ( c ) , with parameters W . This notation allows us to present our algorithms for any function approximator fW ( c ) , and in particular , a deep neural network ( DNN ) . For the theoretical analysis , we will further assume a linear setting , where f∗ ( c ) = cTW ∗ , fW ( c ) = cTW and that W ∗ is in some convex setW . We assume that c ∈ C = ∆d−1 , the standard d− 1 dimensional simplex . This assumption makes the contexts bounded ( which we use in our proofs ) , and it also allows a straight-forward expansion to a model in which the transitions are also a linear mapping of the context ( Modi et al. , 2018 ) . One way of viewing this model is that each row in the mapping W ∗ is a base rewards coefficient vector , and the reward for a specific context is a convex combination of these base rewards . We consider deterministic policies π : S → A which dictate the agent ’ s behaviour at each state . The value of a policy π for reward coefficients vector r is : V πr = Eξ , P , π [ ∑∞ t=0 γ tR ( st ) ] = r Tµ ( π ) where µ ( π ) : = Eξ , P , π [ ∑∞ t=0 γ tφ ( st ) ] ∈ Rk is called the feature expectations of π . For the optimal policy with respect to ( w.r.t . ) a reward coefficients vector r , we denote the value by V ∗r . For any context c , π∗c denotes the optimal policy w.r.t . reward R ∗ c ( s ) = f ∗ ( c ) Tφ ( s ) and π̂c ( W ) denotes the optimal policy w.r.t . reward R̂c ( s ) = fW ( c ) Tφ ( s ) . Inverse Reinforcement Learning in CMDPs : In standard IRL , the goal is to learn a reward which best explains the behavior of an observed expert . The model describing this scenario is the MDP\R - an MDP without a reward function ( also commonly called a controlled Markov chain ) . Similarly , we denote a CMDP without a mapping of context to reward by CMDP\M . The goal in Contextual IRL is to approximate the mapping f∗ ( c ) by observing an expert . The expert knows f∗ ( c ) , and for each context c , can provide a demonstration from π∗c . Contextual dynamics : Learning a transition kernel and an initial state distribution that is parametrized by the context is an orthogonal problem to COIRL . Therefore , we focus only on a contextual reward which simplifies our analysis . Existing methods , such as in Modi et al . ( 2018 ) , can be used to learn the mappings for the transition kernel and initial distribution in a contextual model . In conjunction with the simulation lemma ( Kearns & Singh , 2002 ) , these methods can extend our results to the more general CMDP setting . 3 OPTIMIZATION METHODS FOR COIRL . In this section , we propose and analyze optimization algorithms for minimizing the following loss function ; Lemma 1 below justifies its use for COIRL . Loss ( W ) = Ec max π [ fW ( c ) · ( µ ( π ) − µ ( π∗c ) ) ] = Ec [ fW ( c ) · ( µ ( π̂c ( W ) ) − µ ( π∗c ) ) ] . ( 1 ) Lemma 1 . Loss ( W ) satisfies the following properties : ( 1 ) ∀W , Loss ( W ) ≥ 0 , and Loss ( W ∗ ) = 0 . ( 2 ) If Loss ( W ) = 0 then ∀c ∈ C , the expert policy π∗c is the optimal policy w.r.t . reward cTW . To evaluate the loss , the optimal policy π̂c ( W ) and its features expectations µ ( π̂c ( W ) ) must be computed for all contexts . For a specific context , finding π̂c ( W ) can be solved with standard RL methods such as Value or Policy Iteration . Computing µ ( π̂c ( W ) ) is equivalent to policy evaluation ( solving linear equations ) . The challenge is that Eq . ( 1 ) is is not differentiable in W . We tackle this problem using two methods for computing descent directions that do not involve differentiation : ( i ) computing subgradients and ( ii ) randomly perturbing the loss function . In addition , as the loss is defined in expectation over the contexts , computing it requires to calculate the optimal policy for all contexts . We deal with this issue at the end of Section 3.1 . In the special case that fW ( c ) is a linear function , Eq . ( 1 ) is convex . The following Lemma characterizes Eq . ( 1 ) in this case . Lemma 2 . Let Llin ( W ) = Ec [ cTW · ( µ ( π̂c ( W ) ) −µ ( π∗c ) ) ] . We have that : ( 1 ) Llin ( W ) is a convex function . ( 2 ) g ( W ) = Ec [ c ( µ ( π̂c ( W ) ) − µ ( π∗c ) ) ] is a sub gradient of Llin ( W ) . ( 3 ) Llin is a Lipschitz continuous function , with Lipschitz constant L = 21−γ w.r.t . ‖·‖∞ and L = 2 √ dk 1−γ w.r.t . ‖·‖2 . A technical proof ( by definition ) is provided in the supplementary material . Note that g ( W ) ∈ Rd×k ; we will sometimes refer to it as a matrix and sometimes as a flattened vector , no confusion will arise . Remark 1 . The Lipschitz of LLin ( W ) is related to the simulation lemma ( Kearns & Singh , 2002 ) ; a small change in the reward results in a small change in the optimal value . Remark 2 . As g ( W ) is a subgradient of Loss ( W ) , it can be used to back-propagate DNNs . Clearly , we can not guarantee convexity ( hence no theoretical guarantees ) , but we can design Loss ( W ) to be Lipschitz continuous in W using the methods presented in Cisse et al . ( 2017 ) ; Arjovsky et al . ( 2017 ) . Remark 3 . The subgradient g ( W ) is given in expectation over contexts , and in expectation over trajectories ( feature expectations ) . We will later see how to replace it with an unbiased estimate , which can be computed by observing a single expert trajectory for a single context .

Summary: This work focuses on the problem of 'contextual' inverse reinforcement learning, where the reward is a function of the current state of the MDP, and a set of context features, which remain constant within each episode.  The primary contribution of this work is the formulation of inverse reinforcement learning (for restricted spaces of context-dependent reward functions) as a convex optimization problem.  Based on this formulation, the paper describes several IRL algorithms based on approaches to solving convex and non-convex optimization problems, including variations of mirror descent, and evolution strategies (in principle allowing for the optimization of reward functions with arbitrary parametric representations).  The algorithms presented in this work all assume that computing an optimal policy for a specific reward function is a relatively inexpensive subroutine, which limits their applicability to domains where such planning is straightforward.  Experimental results are presented for a simple highway driving domain, as well as a simulated patient treatment domain constructed from real-world clinical data.