Title: CONTEXTUAL INVERSE REINFORCEMENT LEARNING

Abstract: We consider the Inverse Reinforcement Learning problem in Contextual Markov Decision Processes. In this setting, the reward, which is unknown to the agent, is a function of a static parameter referred to as the context. There is also an "expert" who knows this mapping and acts according to the optimal policy for each context. The goal of the agent is to learn the expert's mapping by observing demonstrations. We define an optimization problem for finding this mapping and show that when it is linear, the problem is convex. We present and analyze the sample complexity of three algorithms for solving this problem: the mirrored descent algorithm, evolution strategies, and the ellipsoid method. We also extend the first two methods to work with general reward functions, e.g., deep neural networks, but without the theoretical guarantees. Finally, we compare the different techniques empirically in driving simulation and a medical treatment regime.

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 cannot be used and, instead, we focus on directly learning the mapping.

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

(a) COIRL diagram (b) Personalized medicine (Itenov et al., 2018) Figure 1 : The COIRL framework (left): a context vector parametrizes the environment. For each context, the expert uses the true mapping from contexts to rewards, W * , and provides demonstrations.

The agent learns an estimation of this mapping Ŵ and acts optimally with respect to it.

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, and M is a mapping from contexts c ∈ C to MDPs: M(c) = (S, A, P, R c , ξ, γ). 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 f W (c), with parameters W . This notation allows us to present our algorithms for any function approximator f W (c), and in particular, a deep neural network (DNN). For the theoretical analysis, we will further assume a linear setting, where f * (c) = c T W * , f W (c) = c T W and that W * is in some convex set W.

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 γ t R(s t )] = r T µ(π) where µ(π) := E ξ,P,π [ ∞ t=0 γ t φ(s t )] ∈ R k 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 Rc (s) = f W (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.

EQUATION

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 c T W.

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 f W (c) is a linear function, Eq. ( 1) is convex. The following Lemma characterizes Eq. (1) in this case.

Lemma 2. Let L lin (W ) = E c c T W • µ(π c (W )) -µ(π * c ) . We have that: (1) L lin (W ) is a convex function. (2) g(W ) = E c c µ(π c (W )) -µ(π * c ) is a sub gradient of L lin (W ). (3) L lin is a Lipschitz continuous function, with Lipschitz constant L = 2 1-γ 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 ) ∈ R d×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 L Lin (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 cannot 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.

3.1. MIRRORED DESCENT FOR COIRL:

Lemma 2 identifies L Lin (W ) as a convex function and provides a method to compute its subgradients.

A standard method for minimizing a convex function over a convex set is the subgradient projection algorithm (Bertsekas, 1997) :

w t+1 = Proj W {w t -α t g(w t )}, where f (w t ) is a convex function, g(w t

) is a subgradient of f (w t ), and α t the learning rate. W is a convex set, and specifically, we consider the 2 ball (Abbeel & Ng, 2004 ) and the simplex (Syed & Schapire, 2008) 1 . We focus on

Algorithm 1 MDA for COIRL input: a convex set W, T number of itera- tions initialize w 1 ∈ W for t = 1, . . . , T do Observe c, µ(π * c ) Compute πc (W ), µ(π c (W )) Compute g t according to Lemma 2 if PSGD then α t = (1 -γ) 1 2dkt w t+1 = w t -α t g t if w t+1 > 1 then w t+1 = w t+1 / w t+1 2 else if Exponential weights then α t = (1 -γ) log(dk) 2t for i = 1, . . . , dk do w t+1 (i) = w t (i) exp (-α t g t (i)) w t+1 = w t+1 / i w t+1 (i) return 1 t T t=1 w t

Algorithm 2 ES for COIRL input: step sizes {α t } T t=1 s, noise STD σ, number of evaluations m and smoothing parameter ν > 0 initialize:

W ∈ R k for t = 1, . . . , T do Observe c, µ(π * c ) for j = 1, ..., m do u j ∼ N k (0, σ 2 ) Loss j (W ) = Loss W + uj ||uj || ν d Loss(W ) = m j=1 Loss j (W ) uj ||uj || ν If Loss(W -αt b d Loss(W ) ) < Loss(W ) then W = W -αt mσ d Loss(W ) return W

a generalization of the subgradient projection algorithm that is called the mirror descent algorithm (Nemirovsky & Yudin, 1983, MDA) :

w t+1 = arg min w∈W w • ∇ f (w t ) + 1 αt D ψ (w, w t )

, where D ψ (w, w t ) is a Bregman distance2 , associated with a strongly convex function ψ. The following theorem characterizes the convergence rate of MDA. Theorem 1 (Convergence rate of MDA). Let ψ be a σ-strongly convex function on W w.r.t. • , and let D 2 = sup w1,w2∈W D ψ (w 1 , w 2 ). Let f be convex and L-Lipschitz continuous w.r.t. • . Then,

MDA with α t = D L 2σ t satisfies: f 1 T T s=1 x s -f (x * ) ≤ DL 2 σT .

We refer the reader to Beck & Teboulle (2003) and Bubeck (2015) for the proof. Next, we provide two MDA instances (see, for example Beck & Teboulle (2003) for derivation) and analyze them for COIRL.

Projected subgradient descent (PSGD): Let W be an 2 ball with radius 1. Fix || • || 2 , and ψ(w) = 1 2 ||w|| 2 2 . ψ is strongly convex w.r.t. || • || 2 with σ = 1. The associated Bregman divergence is given by D ψ (w 1 , w 2 ) = 0.5||w 1 -w 2 || 2 2 . Thus, mirror descent is equivalent to PSGD. D 2 = max w1,w2∈W D ψ (w 1 , w 2 ) ≤ 1, and according to Lemma 2, L = 2 √ dk 1-γ . Thus, we have that after T iterations L lin

1 T T t=1 w t -L lin (w * ) ≤ O √ dk (1-γ) √ T . Exponential Weights (EW): Let W be the standard dk -1 dimensional simplex. Let ψ(w) = i w(i) log(w(i)). ψ is strongly convex w.r.t. || • || 1 with σ = 1.

We get that the associated Bregman divergence is given by D ψ (w 1 , w 2 ) = i w 1 (i) log( w1(i) w2(i) ), also known as the Kullback-Leibler divergence. In addition, D 2 = max x,y∈W D ψ (w 1 , w 2 ) ≤ log(dk) and according to Lemma 2, L = 2 1-γ . Furthermore, the projection onto the simplex w.r.t. to this distance amounts to a simple renormalization w ← w/||w|| 1 . Thus, we get that MDA is equivalent to the exponential weights algorithm and L lin

1 T T t=1 w t -L lin (w * ) ≤ O √ log(dk) (1-γ) √

T . Practical MDA: One of the "miracles" of MDA is its robustness to noise. If we replace g t with an unbiased estimate gt , such that Eg t = g t and E gt ≤ L, we obtain the same convergence results as in Lemma 2 (Robbins & Monro, 1951 ) (see, for example, (Bubeck, 2015, Theorem 6.1)) . Such an unbiased estimate can be obtained in the following manner: (i) sample a context c t , (ii) compute µ(π * c T t wt ), (iii) observe a single expert demonstration τ E i = {s i 0 , a 0 , s i 1 , a 1 , . . . , }, where a i is chosen by the expert policy π * c T t w * (iv) let μi = t∈[0,...,|τ E i |-1] γ t φ(s i t ) be the accumulated discounted features across the trajectory such that Eμ i = µ(π * c ). The challenge is, that for μi to be an unbiased estimate of µ(π * c T t w * ), τ E i needs to be of infinite length. There are two ways in which we can tackle this issue. We can either (1) execute the expert trajectory online, and terminate it at each time step with probability 1 -γ (as in (Kakade & Langford, 2002) ), or (2) execute a trajectory of length H = 1 1-γ log(1/ H ). The issue with the first approach is that since the trajectory length is unbounded, the estimate μi cannot be shown to concentrate to µ(π * c ) via Hoeffding type inequalities. Nevertheless, it is possible to obtain a concentration inequality using the fact that the length of each trajectory is bounded in high probability (similar to Zahavy et al.) . The second approach can only guarantee that g t -Eg t ≤ H (Syed & Schapire, 2008) . Therefore, using the robustness of MDA to adversarial noise (Zinkevich, 2003) , we get that MDA converges with an additional error of H , i.e., L lin

1 T T t=1 w t -L lin (w * ) ≤ O 1 √ T + H .

While this sampling mechanism comes with the cost of a controlled bias, usually it is more practical, in particular when the trajectories are given as a set demonstrations (offline data).

3.2. EVOLUTION STRATEGIES FOR COIRL:

To minimize Eq. ( 1), we also design a derivative free algorithm (Algorithm 2) that is based on Evolution Strategies (Salimans et al., 2017, ES) . For convex optimization problems, ES is a gradientfree descent method that is based on computing finite differences (Nesterov & Spokoiny, 2017) , whose sample complexity is provided below in Theorem 2. The Theorem is given in terms of the Lipschitz constant, which is upper bounded by 2 √ dk 1-γ (Section 3.1). While this approach has looser upper-bound guarantees compared to MDA (Theorem 1), Nesterov & Spokoiny (2017) observed that in practice, it often outperforms subgradient based methods. Thus, we test this method empirically and compare it with the subgradient method (Section 3.1). ES is also known to perform well in practice, even with nonconvex objectives. Specifically, Salimans et al. (2017) has shown that ES can be used to optimize the parameters of a DNN to solve challenging high dimensional RL tasks like playing Atari.

Theorem 2 (ES Convergence Rate (Nesterov & Spokoiny, 2017) ). Let L lin (W ) be a non-smooth convex function with Lipschitz constant L, such that

||x 0 -x * || ≤ D, step size of α t = D (dk+4) √ T +1L and ν ≤ 2L √ dk then in T = 4(dk+4) 2 D 2 L 2 2 ES finds a solution which is bounded by E U T -1 [L lin (x T )] -L lin (x * ) ≤

, where U T = {u 0 , . . . , u T } denotes the random variables of the algorithm up to time T and xT = arg min t=1,...,T L lin (x t ).

3.3. ELLIPSOID ALGORITHMS FOR COIRL:

Algorithm 3 Ellipsoid algorithm for COIRL

Initialize: Θ 0 ← B ∞ (0, 1) ={x ∈ R d•k : ||x|| ∞ ≤ 1} Θ 1 ← MVEE(Θ 0 ) for t = 1, 2, . . . do Observe c t , let W t be the center of Θ t Play episode using πt = arg max π V π c T t Wt if V * c T t W * -V πt c T t W * > then µ(π * ct ) is revealed Let a t = c t µ(π * ct ) -µ(π t ) Θ t+1 ← MVEE θ ∈ Θ t : θ T a t ≥ W T t a t else Θ t+1 ← Θ t

The final algorithm we consider is an ellipsoid method, introduced to the IRL setting by Amin et al. (2017) . In this section we extend it to the contextual setting, specifically, we focus on finding a linear mapping W and further assume that W = {W : ||W || ∞ ≤ 1}, and that W * ∈ W.

The algorithm maintains an ellipsoidshaped feasibility set that contains W * . At any step, the current estimation W t of W * is defined as the center of the ellipsoid, and the agent acts optimally w.r.t. this estimation. If the agent performs suboptimally, the expert provides a demonstration in the form of the optimal feature expectations for c t , µ(π * ct ). The feature expectations are used to generate a linear constraint (hyperplane) on the ellipsoid that is crossing its center. Under this constraint, we construct a new feasibility set that is half of the previous ellipsoid, and still contains W * . For the algorithm to proceed, we compute a new ellipsoid that is the minimum volume enclosing ellipsoid (MVEE) around this "half-ellipsoid"3 . These updates are guaranteed to gradually reduce the volume of the ellipsoid (a well-known result (Boyd & Barratt, 1991) ) until its center is a mapping which induces -optimal policies. Theorem 3 shows that this algorithm achieves a polynomial upper bound on the number of sub-optimal time-steps. Finally, note that in Algorithm 3 we use an underline notation to denote a "flattening" operator for matrices, and to denote a composition of an outer product and the flattening operator. The proofs in this section are provided in the supplementary material, and are adapted from (Amin et al., 2017) . Theorem 3. In the linear setting where R * c (s) = c T W * φ(s), for an agent acting according to Algorithm 1, the number of rounds in which the agent is not -optimal is O(d 2 k 2 log( dk (1-γ) )). Remark 4. Note that the ellipsoid method presents a new learning framework, where demonstrations are only provided when the agent performs sub-optimally. Thus, the theoretical results in this section cannot be directly compared with those of the descent methods. We further discuss this in the experiments and discussion sections. Remark 5. The ellipsoid method does not require a distribution over contexts -an adversary may choose them. MDA can also be easily extended to the adversarial setting via known regret bounds on online MDA (Hazan, 2016) .

Practical ellipsoid algorithm: In many real-world scenarios, the expert cannot evaluate the value of the agent's policy and cannot provide its policy or feature expectations. To address these issues, we follow Amin et al. (2017) and consider a relaxed approach, in which the expert evaluates each of the individual actions performed by the agent rather than its policy, and provides finite rollouts instead of a policy or feature expectations (see the supplementary material (Algorithm 4) for pseudo code). We define the expert criterion for providing a demonstration to be Q

* c T t W * (s, a) + < V * c T t W * (s)

for each state-action pair (s, a) in the agent's trajectory.

Near-optimal experts: In addition, we relax the optimality requirement of the expert and instead assume that, for each context c t , the expert acts optimally w.r.t. W * t which is close to W * ; the expert also evaluates the agent w.r.t. this mapping. This allows the agent to learn from different experts, and from non-stationary experts whose judgment and performance slightly vary over time. If a sub-optimal action w.r.t. W * t is played at state s, the expert provides a roll-out of H steps from s to the agent. As this roll-out is a sample of the optimal policy w.r.t. W * t , we aggregate n examples to assure that with high probability, the linear constraint that we use in the ellipsoid algorithm does not exclude W * from the feasibility set. Note that these batches may be constructed across different contexts, different experts, and different states from which the demonstrations start. Theorem 4 below upper bounds the number of sub-optimal actions that Algorithm 4 chooses. 4Theorem 4. For an agent acting according to Algorithm 4 , with probability of at least 1 -δ,

for H = 1 1-γ log( 8k (1-γ) ) and n = 512k 2 (1-γ) 2 2 log(4dk(dk + 1) log( 16k √ dk (1-γ) )/δ) , if ∀t : W * t ∈ B ∞ (W * , (1-γ) 8k ) ∩ Θ 0 the number of rounds in which a sub-optimal action is played is O d 2 k 4 (1-γ) 2 2 log dk (1-γ)δ log( dk (1-γ) ) .

4. EXPERIMENTS:

The simulations in this section include two domains: (1) an autonomous driving simulation (Abbeel & Ng, 2004) , that we adapted to the contextual setup and (2) a medical treatment regime, constructed from a data set of expert (clinician) trajectories for treating patients with sepsis5 . In each of these domains we compare the algorithms in two setups: the ellipsoid learning framework and an offline framework. All the results are averaged across 10 random seeds in Section 4.1 and 5 seeds in Section 4.2 (we report the mean and the standard deviation). Due to space considerations we present the simulations in the ellipsoid framework only for the car domain, and the simulations in the offline framework only in the dynamic treatment regime. Complementary simulations can be found in the supplementary material.

4.1. DRIVING SIMULATION -THE ELLIPSOID FRAMEWORK:

In the ellipsoid framework, an expert evaluates the agent policy. If the agent's policy is sub-optimal, the expert provides the agent its feature expectations; otherwise, no demonstration is given. The algorithm performs learning in between demonstrations. This setup enables a proper comparison with the ellipsoid algorithm, which requires the additional expert supervision. We measure performance w.r.t. the following criteria: (1) # demonstrations -the amount of contexts on which each algorithm requested an expert demonstration (y-axis) as a function of time, i.e., the total number of contexts (x-axis). ( 2) Value -the difference in value, between the agent policy and the expert policy w.r.t. the true reward mapping, i.e., c∈Ctest f W

* (c) • µ(π c (W )) -µ(π * c )

, where C test is a holdout (test) set of contexts. The x-axis measures the amount of demonstrations given.

Setup. This domain simulates a three-lane highway with two visible cars -cars A and B (illustration provided in the appendix). The agent, controlling car A, can drive both on the highway and off-road. Car B drives on a fixed lane, at a slower speed than car A. Upon leaving the frame, car B is replaced by a new car, appearing in a random lane at the top of the screen. The feature vector φ(s) is composed of 3 features: (1) a speed feature, (2) a collision feature, which is valued 0 in case of a collision and 0.5 otherwise, and (3) an off-road feature, which is 0.5 if the car is on the road and 0 otherwise.

In this task, the context vector implies different priorities for the agent; should it prefer speed or safety? Is going off-road to avoid collisions a valid option? For example, an ambulance will prioritize speed and may allow going off-road as long as it goes fast and avoids collisions, while a bus will prioritize avoiding both collisions and off-road driving as safety is its primary concern. To demonstrate the effectiveness of our solutions, the mapping f : C → [-1, 1] k is constructed in a way that induces different behaviors for different contexts, making generalization a challenging task. We provide additional details on the domain as well as the hyper parameter selection in the appendix.

Linear: The optimal behavior is defined using a linear mapping W * . In this setting, all three approaches obtain competitive results, in terms of generalization, although the ES is capable of obtaining these results faster, as is seen through the regret and number of required demonstrations. Nonlinear: For the nonlinear task, we consider two reward coefficient vectors r 1 and r 2 , and define the mapping by f * (c) = r 1 if ||c|| ∞ ≥ 0.55, and r 2 otherwise -an illustration is provided in the appendix. In order to learn the nonlinear mapping, we represent f W (c) using a DNN, a multi-layered perceptron, which maps from context to reward vector. DNNs have proven to be capable of extracting meaningful features from complex high-dimensional data, e.g., images -in these scenarios, the linear assumption no longer holds, yet DNNs often overcome such issues. In this setting, the superiority of the descent methods rises; as the linear assumption in the ellipsoid algorithm is not met, it fails to generalize and keeps requiring new demonstrations. We believe these results to be crucial when considering real-life applications, in which the problem is not necessarily linear. Such cases highlight the strength of the descent methods, which, as Fig. 2 shows, are capable of scaling to nonlinear high dimensional mappings.

4.2. DYNAMIC TREATMENT REGIME -THE OFFLINE FRAMEWORK:

In the offline framework, we focus on the ability to learn from previously collected data. A data set of previously collected trajectories is given, such that a single trajectory of finite length is observed for each context and no context is observed more than once. We measure performance w.r.t. the following criteria: (1) Value -as in the ellipsoid framework above, but here the x-axis corresponds to the amount of iterations. Each iteration corresponds to a single subgradient step, where the subgradient is computed from a mini batch of 10 contexts.

(2) Loss -as in Eq. ( 1).

(3) Accuracy % -the percent of actions on which the expert and the agent agree on. All these criteria are evaluated on a holdout set.

Setup. In the dynamic treatment regime, there is a clinician which acts to improve a sick patient's medical condition. The context (static information) represents patient features, which do not change during treatment, such as age and gender. The state summarizes the dynamic measurements of the patient, e.g., blood pressure and EEG readouts. The actions are the forms of intervention a clinician may take, including combinations of various treatments provided in parallel. Dynamic treatment regimes are particularly useful for managing chronic disorders and fit well into the broader paradigm of personalized medicine (Komorowski et al., 2018; Prasad et al., 2017) .

The agent needs to choose the right treatment for a patient that is diagnosed with sepsis. We use the MIMIC-III data set (Johnson et al., 2016) and follow the data processing steps that were taken in Jeter et al. (2019) . As performing off-policy evaluation is not possible using this data set, due to it not satisfying basic requirements (Gottesman et al., 2018; 2019) , we designed a simulator of a CMDP.

The simulator is based on this data set; a complete overview and explanation on how it was created is provided in the appendix. The mapping W * is linear, W * ∈ R 8×42 , which we constructed from the data. In the simulator, the expert acts optimally w.r.t. this W * .

Specifically, when treating a sepsis patient, the clinician has several decisions to make, such as whether or not to provide a patient with vasopressors, drugs which are commonly provided to restore and maintain blood pressure in patients with sepsis. However, what is regarded as healthy blood pressure differs based on the age and weight of the patient (Wesselink et al., 2018) . In our setting, W captures this information -as it maps from contextual (e.g., age) and dynamic information (e.g., blood pressure) to reward.

Results. Fig. 3 presents the ability of the descent methods to generalize to unseen contexts by learning from offline data (without supervision). The data is composed of a set of trajectories, i.e., offline data, that were collected from experts (clinicians treating patients). In each iteration, we sample a minibatch of 10 contexts, i.i.d, from the context distribution. For each context, there is a corresponding expert trajectory of length H = 40. Performance is measured on a holdout set of 300 contexts (that are sampled from the same context distribution) according to Theorem 1. We can see that both ES and PSGD attain near-optimal performance using only previously collected expert trajectories. Looking at Fig. 3a , we can see that all the algorithms manage to minimize the loss to roughly the same error. The small bias is explained by the fact that we use truncated trajectories (as we discussed in the practical MDA paragraph) where in the ellipsoid framework experiments we used feature expectations. We can also see that minimizing the loss leads to policies that attain -optimal value w.r.t. the true reward Fig. 3b . Finally, in Fig. 3 we can see that all the algorithms reach around 70% accuracy with the expert policy. We emphasize here that 100% accuracy should not be expected for two reasons: (i) different policies may have the same feature expectations (hence the same value) but make different decisions (ii) there exists reward for which there is more than one optimal policy. Nevertheless, Fig. 3 suggests that accuracy is correlated with minimizing the COIRL loss (Eq. ( 1)).

Finally, we present results in the nonlinear setting. Here, there is a non-linear function of the context that determines which one of the two reward coefficient vectors is used, i.e., f * (c) = r 1 if age > 0.1, and r 2 otherwise. Where age refers to the normalized age of the patient, which is an element As seen in Fig. 4 , the PSGD algorithm minimizes the loss and achieves a value that is close to that of the expert.

In addition, similarly to Fig. 7 , accuracy and performance do not necessarily correlate one to another.

5. RELATED WORK:

We begin with a short discussion on contextual policies, i.e., a policy that is a function of both the state and the context. While there is empirical evidence that learning such a policy may perform well in practice (e.g., Xu et al. (2018) ; Fu et al. ( 2019)), from a theoretical point of view, there exist hardness results in this setting. Specifically, given an MDP with k + 1 states, there is a reduction from the training problem of the union of k hyperplanes to the policy (see appendix E for proof).

Alternatively, one may consider applying an AL algorithm on a single, large MDP that includes all the states and the contexts. For a concrete example, consider a reduction from the CMDP model to a large MDP where each state is expanded by concatenating the context to it. The new states are s = (s, c), and the new features are φ(s ) = c φ(s). Generally speaking, applying an AL algorithm to this large MDP will give the same scalability and sample complexity as COIRL. However, as the large MDP has |S | = |C||S| states, computing the optimal policy in each iteration of the algorithm will require at least |C| times more time.

To illustrate this problem we conducted a simple grid world experiment on a 3 × 4 grid world MDP (with 12 states) and one-hot features (φ(s i ) = e i ∈ R 12 ). The dynamics are deterministic, and the actions correspond to going up, down, left and right (with cyclic transitions on the borders). The contexts correspond to "preferences" on the grid; mathematically, each context is sampled from a uniform distribution over the simplex. The mapping W * is set to be I 12•12 , and for AL, we let w * to be a flattened version of W * . We compare the performance of PSGD with the projection algorithm of (Abbeel & Ng, 2004) in Fig. 5 . We measure performance by three metrics: run-time, value, and accuracy. Inspecting the results, we can see that AL in the large MDP requires significantly more time to run as the number of contexts grows, while the run time of PSGD (COIRL) is not affected by the number of contexts. We can also see that both methods achieve roughly the same performance: COIRL performs slightly better in terms of value while AL performs slightly better in terms of accuracy. To conclude, AL on a large MDP does not scale to problems with large context spaces. In addition, this construction is only possible when there is a finite number of context and does not provide generalization results. We avoided all of these issues in the COIRL framework.

6. SUMMARY AND DISCUSSION:

In this work, we formulated and studied the COIRL problem. We presented two types of algorithms to solve it: (1) cutting plane methods (ellipsoid) and (2) iterative descent approaches (MDA and ES).

We summarize the theoretical guarantees of the different algorithms in Table 1 .

We can see that the iterative descent approaches have better dependence in dk than the ellipsoid method, i.e., they scale better with the dimensions of the problem. In particular, the EW algorithm has a logarithmic dependence in dk, which makes it computationally comparable to standard IRL/AL (on a single, noncontextual MDP). In addition, the iterative methods extend naturally to the more general scenario where the mapping from contexts to rewards is not linear, and f W is modeled as a DNN. As Sutton (2019) puts it: "The biggest lesson that can be read from 70 years of AI research is that general methods that leverage computation are ultimately the most effective, and by a large margin".

The ellipsoid method has better sample complexity (as a function of ) than the descent methods in the deterministic setting. However, both methods attain the same complexity in the more realistic, stochastic setting. Our empirical findings suggest that the iterative methods always outperform the ellipsoid algorithm. Among these methods, we found the ES method to perform better than the MDA method. Similar findings were reported in (Nesterov & Spokoiny, 2017) for other convex problems.

The iterative methods have another advantage over the ellipsoid method -they can learn from previously collected demonstrations (i.e., offline learning). The ellipsoid framework, on the other hand, requires expert supervision throughout the entire learning process.

Finally, an attractive property of the ellipsoid learning framework is its safety, i.e., an IRL algorithm that is being supervised by an expert will never perform sub-optimally. In each step, either that the agent performs -optimally or that the expert acts on its behalf (provides a demonstration). This property is appealing in mission-critical domains where errors have a high cost; for instance, in health-care, a failure may result in a loss of lives. In the experimental section, we have seen that we can use this learning framework for the iterative methods as well while enjoying improved efficiency.

Scalability Sample Complexity:

Extension to DNNs Similar to Section 4.2, we compare the various methods in the offline framework. We can see that all the algorithms manage to minimize the loss, and achieve near optimal value. We can also see that they achieve high accuracy with the expert policy but not 100%. Similar to Section 4.1, we compare the various methods in the ellipsoid framework . We observe that ES outperforms the ellipsoid method. Additionally, we compare the accuracy, i.e., how often does the policy derived from W match the expert's policy, which is derived from W * . As IRL is not a supervised learning problem, we observe that while there is a correlation between the success in the task and the ability to act similarly to the experts policy -this correlation is not strict, in the sense that the agent is capable of finding near-optimal policies with a relatively high miss rate (accuracy of approximately 80%). For more intuition, see the proof of Lemma 1. In these experiments, we define our mappings in a way that induces different behaviours for different contexts, making generalization a more challenging task. Specifically, for the linear setting we use W * = -1 0.75 0.75 0.5 -1 1 0.75 1 -0.75 , before normalization. For our nonlinear mapping, contexts with ||c|| ∞ > 0.55 are mapped to reward coefficients vector (1, -1, -0.05), otherwise they are mapped to (-0.01, 1, -1), which induce the feature expectations (9.75, 3.655, 5), (5.25, 5, 2.343) respectively. The decision regions for the nonlinear mapping are visualized in Appendix B.1. The contexts are sampled uniformly in the 2-dimensional simplex. We evaluate all algorithms on the same sequences of contexts, and average the results over 20 such sequences. The algorithms were modified to fit the Ellipsoid framework; instead of iterating over the whole data, the algorithms iterate over the given expert feature expectations, one at a time, until convergence, i.e. every timestep a new demonstration is presented. 

Deterministic Stochastic Deterministic Stochastic MDA PSGD O(dk) O 1 2 O 1 2 EW O(log dk) ES O(dk) O(d 2 k 2 ) Ellipsoid O(d 2 k 2 ) O(d 2 k 4 ) O 1 log(1/ )

PSGD::

The algorithm was executed with with the parameters: α 0 = 0.3, α t = 0.9 t α t-1 , and iterated for 40 epochs. An outer decay on the step size was added for faster convergence, the initial α 0 becomes 0.94 • α 0 every time a demonstration is presented. The gradient, g t is normalized to be g t = g t gt ||gt||∞ and the calculated step is taken if:

cW t µ(π t c ) -µ(π * c ) > cW t+1 µ(π t+1 c ) -µ(π * c )

, where πt c denotes the optimal policy for a context c according to W t . ES: The algorithm was executed with the parameters: σ = 10 -3 , m = 250, α = 0.1 with decay rate of 0.95, for 50 iterations which didn't iterate randomly over one the contexts, but rather used the entire training set (all of the observed contexts and expert demonstrations up to the current time-step) for each step. The matrix was normalized according to || • || 2 , and so was the step calculated by the ES algorithm, before it was multiplied by α and applied.

For the nonlinear setting, the model used for the ES method was a fully connected DNN, with layers of sizes 15, 10, 5, 3. The activation function used was the leaky ReLU function, with a parameter of α = 0.1. Note that we can't normalize the parameters here as in the linear case; therefore an L2-normalization layer is added to the output. The same parameters were used as in the linear case, except with 120 iterations over the entire training set. They were originally optimized for this model and setting and worked as-is for the linear environment. As we aim to estimate the gradient, a small σ = 10 -3 was used and performed best. The number of points, m = 250, was selected as fewer points produced noisy results. The step size, decay rate and the number of iterations were 0.1, 0.96, 120 respectively, and were selected in a way that produced fast yet accurate convergence of the loss. Note that here the steps were also normalized before application, and the normalization was applied per layer. For PSGD, a similar network was used. Specifically, it had layer sizes 14, 10, 6, 3, and the same leaky ReLU activation function was used in this network. In parallel to the normalization used for the ES model, here we used batch normalization and gradient clipping. The learning rate was set to 0.1 • 0.98 t and 120 iterations were performed. For this result, as with the ES method, the batch consisted of all available training data.

Offline Framework: In the offline framework we compute the subgradients using expert trajectories of length 40, instead of the feature expectations. In this framework at every iteration we sample a mini batch of 10 contexts (from a finite set) and their corresponding trajectories (sampled from the expert policy and dynamics) then taking one descent step according to them. The generalization is measured over 80 holdout contexts that are referred as the test set, where the W that is used to calculate the feature expectations of the agent is fitted to the EW algorithm requirement to be in the dk -1 simplex. The PSGD and the EW algorithms are configured as the theory specifies, where each descent step is calculated from the whole batch. The ES algorithm is applied with the parameters σ = 10 -3 , m = 500, α = 0.1 with decay rate 0.95, for every iteration.

B.2 DYNAMIC TREATMENT REGIME:

The environment we describe in 4.2 simulates a decision-making process for treating sepsis. Sepsis is a life-threatening severe infection, where the treatment applied to a sepsis patient is crucial for saving his life. To create a sepsis treating simulator, we leverage the MIMIC-III data set (Johnson et al., 2016) . This data set includes data from hospital electronic databases, social security, and archives from critical care information systems, that had been acquired during routine hospital care. We follow the data processing steps that were taken in Jeter et al. (2019) to extract the relevant data in a form of normalized measurements of sepsis patients during their hospital admission and the treatments that were given to each patient. The measurements include dynamic measures, e.g., heart rate, blood pressure, weight, body temperature, blood analysis standard measures (glucose, albumin, platelets count, minerals, etc.), as well as static measures such as age, gender, re-admission (of the patient), and more.

The processed data from Jeter et al. (2019) consists of 5366 trajectories, each representing the sequential treatment provided by a clinician to a patient. At each time-step, the available information for each patient consists of 8 static measurements and 41 dynamic measurements. In addition, each trajectory contains the reported actions performed by the clinician (the number of fluids and vasopressors given to a patient at each time-step and binned to 25 different values), and there is a mortality signal which indicates whether the patient was alive 90 days after his hospital admission.

In order to create a tabular CMDP from the processed data, we separate the static measurements of each patient and keep them as the context. We cluster the dynamic measurements using K-means (MacQueen et al., 1967) . Each cluster is considered a state and the coordinates of the cluster centroids are taken as its features φ(s). We construct the transition kernel between the clusters using the empirical transitions in the data given the state and the performed actions. Two states are added to the MDP and the feature vector is extended by 1 element, corresponding to whether or not the patient died within the 90 days following hospital release. This added feature receives a value of 0 on all non-terminal states, a value of -0.5 for the state representing the patient's death and 0.5 for the one representing survival. In addition, as the data is limited, not all state-action pairs are available. In order to ensure the agent does not attempt to perform such an action for which the outcome is unknown, we add an additional terminal state. At this state, all features are set to -1 to make it clearly distinguishable from all other states in the CMDP.

In our simulator, we used the same structure as the raw data, i.e., we used the same contexts prevalent in the data and the same initial state distribution. Each context is projected onto the simplex and the expert's feature expectations for each context are attained by solving the CMDP. While we focus on a simulator, as it allows us to analyze the performance of the algorithms, our goal is to have a reward structure which is influenced by the data. Hence, we produce W * by running the ellipsoid algorithm on trajectories obtained from the data. As done in the Autonomous Driving Simulation, the algorithms were modified to fit the Ellipsoid framework.

Hyper-parameter selection and adjustments::

Ellipsoid Framework: ES algorithm: the same method as in the autonomous driving is applied, with the parameters σ = 10 -4 , m = 1000, α = 0.25 with decay rate 0.95, for 80 iterations over the entire training set.

Offline Framework: In the offline framework for the linear setting, we compute the subgradients using expert trajectories of length 40, instead of the exact feature expectations. At every iteration, we sample a mini batch of 10 contexts (from a finite set) and their corresponding trajectories (sampled from the expert policy and dynamics) and perform a single descent step. Generalization is measured over a set of 300 holdout contexts, referred to as the test set, where W is fitted to be in the dk -1 simplex. The PSGD and the EW algorithms are configured as specified by the theory, where each descent step is calculated from the entire batch. The ES algorithm is applied with the parameters σ = 10 -4 , m = 1000, α = 0.3 with decay rate 0.95, for every iteration.

For the nonlinear setting, the model used for the ES method was a fully connected DNN, with layers of sizes 24, 42, 42 which include a bias term. The activation function used was the leaky ReLU function, with a parameter of α = 0.1. In this setting we use trajectories of length 80, mini batch of 32 contexts. The PSGD algorithm uses the gradient computed by back-propagation of the loss function value on the DNN, where each descent step is calculated from the entire batch. The algorithm is applied with modified parameters; α = 0.1 with decay rate 0.98, for every iteration.

C PROOFS FOR SECTION 3:

Definition 1 (Bregman distance). Let ψ : W → R be strongly convex and continuously differential in the interior of W. The Bregman distance is defined by D ψ (x, y) = ψ(x) -ψ(y) -(x -y) • ∇ψ(y), where ψ is strongly convex with parameter σ. Definition 2 (Conjugate function). The conjugate of a function ψ(y), denoted by ψ * (y) is max x∈W {x • y -ψ(x)}.

Example: let • be a norm on R n . The associated dual norm, denoted • * , is defined as

z * = sup{z x | x ≤ 1}. The dual norm of • 2 is • 2 , and the dual norm of • 1 is • ∞ .

Before we begin with the proof of Lemma 2, we make the following observation. By definition, πc (W ) is the optimal policy w.r.t. c T W ; i.e., for any policy π we have that

c T W • µ(π c (W )) ≥ c T W • µ(π).

(2)

Proof of Lemma 1. Fix W . For any context c, we have that µ(π c ) is the optimal policy w.r.t. reward f W (c), thus,

f W (c) • µ(π c (W )) -µ(π * c ) ≥ 0.

Therefore we get that Loss(W ) ≥ 0. For W * , we have that µ(π c (W )) = µ(π * c ), thus Loss(W * ) = 0. For the second statement, note that Loss(W ) = 0 implies that ∀c, f W (c)• µ(π c (W ))-µ(π * c ) = 0. This can happen in one of two cases. ( 1) µ(π c (W )) = µ(π * c ), in this case π * c , πc (W ) have the same feature expectations. Therefore, they are equivalent in terms of value. ( 2

) µ(π c (W )) = µ(π * c ), but f W (c) • µ(π c (W )) -µ(π * c ) = 0.

In this case, π * c , πc (W ) have different feature expectations, but still achieve the same value w.r.t. reward f W (c). Since πc (W ) is an optimal policy w.r.t. this reward, so is π * c .

Proof of Lemma 2.

1. We need to show that ∀W 1 , W 2 ∈ W, ∀λ ∈ [0, 1], we have that

L lin (λW 1 + (1 -λ)W 2 ) ≤ λL lin (W 1 ) + (1 -λ)L lin (W 2 ) L lin (λW 1 + (1 -λ)W 2 ) = E c c T (λW 1 + (1 -λ)W 2 ) • µ πc λW 1 + (1 -λ)W 2 -µ(π * c ) = λE c c T W 1 • µ πc λW 1 + (1 -λ)W 2 -µ(π * c ) + (1 -λ)E c c T W 2 • µ πc λW 1 + (1 -λ)W 2 -µ(π * c ) ≤ λE c c T W 1 • µ πc W 1 -µ(π * c ) + (1 -λ)E c c T W 2 • µ πc W 2 -µ(π * c ) = λL lin (W 1 ) + (1 -λ)L lin (W 2 ),

where the inequality follows from Eq. (2).

2. Fix z ∈ W. We have that

L lin (z) = E c c T z • µ(π c (z)) -µ(π * c ) ≥ E c c T z • µ(π c (W )) -µ(π * c ) = L lin (W ) + (z -W ) • E c c µ(π c (W )) -µ(π * c ) ,

where the inequality follows from Eq. (2).

3.:

Recall that a bound on the dual norm of the subgradients implies Lipschitz continuity for convex functions. Thus it is enough to show that ∀W ∈ W, g(W

EQUATION

)

where in Eq. ( 3) we used the fact that ∀π we have that µ(π

) ∞ ≤ 1 1-γ , thus, for any π i , π j , µ(π i ) -µ(π j ) ∞ ≤ 2 1-γ Therefore, L = 2 1-γ w.r.t. • ∞ . Since • 2 ≤ √ dk • ∞ we get that L = 2 √ dk 1-γ w.r.t. • 2 . D PROOFS & PSEUDO CODE FOR SECTION 3.3 D.1 ELLIPSOID ALGORITHM FOR TRAJECTORIES

Algorithm 4 Batch ellipsoid algorithm for COIRL 

Initialize: Θ 0 ← B ∞ (0, 1) = {x ∈ R d•k : ||x|| ∞ ≤ 1} Θ 1 ← MVEE(Θ 0 ) i ← 0, Z ← 0, Z * ← 0 for t =

ξ i = 1 s : x * ,H i = H h=0 γ h φ(s E h ) Let x i be the agent's feature expectations for ξ i : E ξ i ,P,πt [ ∞ h=0 γ h φ(s h )] Denote z i = c t x i , ẑ * ,H i = c t x * ,H i i ← i + 1, Z ← Z + z i , Z * ← Z * + ẑ * ,H i if i = n then Θ t+1 ← MVEE θ ∈ Θ t : θ -W t T • ( Z * n -Z n ) ≥ 0 i ← 0, Z ← 0, Z * ← 0 D.2 MVEE COMPUTATION

This computation is commonly found in optimization lecture notes and textbooks. First, we define an ellipsoid by {x : (x -c)Q -1 (x -c) ≤ 1} for a vector c, the center of the ellipsoid, and an invertible matrix Q. Our first task is computing Θ 1 -the MVEE for the initial feasibility set

Θ 0 = B ∞ (0, 1) = {x ∈ R d•k : ||x|| ∞ ≤ 1}.

The result is of course a sphere around 0:

c 1 = 0, Q 1 = dkI. For the update Θ t+1 ← MVEE θ ∈ Θ t : θ -W t T • a t ≥ 0 , we define ãt = -1 √ a T t Qtat a t

and calculate the new ellipsoid by

c t+1 = c t -1 dk+1 Qã t , Q t+1 = d 2 k 2 d 2 k 2 -1 (Q t -2 dk+1 Q t ãt ãT t Q t ).

D.3 PROOF OF THEOREM 3:

For simpler analysis, we define a "flattening" operator, converting a matrix to a vector:

R d×k → R d•k by W = [w 1,1 , . . . , w 1,k , . . . , w d,1 , . . . , w d,k ].

We also define the operator to be the composition of the flattening operator and the outer product:

u v = [u 1 v 1 , . . . , u 1 v k , . . . , u d v 1 , . . . , u d v k ].

Therefore, the value of policy π for context c is given by V D+1) .

EQUATION

Proof of Theorem 3. We prove the theorem by showing that the volume of the ellipsoids Θ t for t = 1, 2, ... is bounded from below. In conjunction with Lemma 3, which claims there is a minimal rate of decay in the ellipsoid volume, this shows that the number of times the ellipsoid is updated is polynomially bounded. We begin by showing that W * always remains in the ellipsoid. We note that in rounds where 

V π * c T t W * -V πt c T t W * > , we have W * T c t µ(π * ct ) -µ(π t ) > .

EQUATION

This shows that W * is never disqualified when updating Θ t . Since W * ∈ Θ 0 this implies that ∀t : W * ∈ Θ t . Now we show that not only W * remains in the ellipsoid, but also a small ball surrounding it. If θ is disqualified by the algorithm:

(θ -W t ) T • c t µ(π * ct ) -µ(π t ) < 0 .

Multiplying this inequality by -1 and adding it to (4) yields:

(W * -θ) T • c t µ(π * ct )-µ(π t ) > .

We apply Hölder inequality to LHS:

< LHS ≤ ||W * -θ|| ∞ • || c t µ(π * ct ) -µ(π t ) || 1 ≤ 2k 1-γ ||W * -θ|| ∞ . Therefore for any disqualified θ: ||W * -θ|| ∞ > (1-γ) 2k , thus B ∞ (W * , (1-γ) 2k ) is never disqualified. This implies that ∀t : vol(Θ t ) ≥ vol(Θ 0 ∩ B ∞ (W * , (1-γ) 2k )) ≥ vol(B ∞ (W * , (1-γ) 4k )).

Finally, let M T be the number of rounds by T in which

V π * c T t W * - V πt c T t W * > . Using Lemma 3 we get that: M T 2(dk+1) ≤ log vol(Θ 1 ) -log vol(Θ T +1 ) ≤ log vol MVEE(B ∞ (0, 1)) -log vol(B ∞ (0, (1-γ) 4k )) ≤ log vol MVEE(B 2 (0, √ dk)) - log vol(B 2 (0, (1-γ) 4k )) ≤ log ( 4k √ dk (1-γ) ) dk ≤ dk log 4k √ dk (1-γ) . Therefore M T ≤ 2dk(dk + 1) log 4k √ dk (1-γ) = O(d 2 k 2 log( dk (1-γ) )) .

D.4 PROOF OF THEOREM 4:

Lemma 4 (Azuma's inequality). For a martingale

{S i } n i=0 , if |S i -S i-1 | ≤ b a.s. for i = 1, ..., n: P |S n -S 0 | > b 2n log( 2 δ ) < δ

Proof of Theorem 4. We first note that we may assume that for any t:

||W * -W t || ∞ ≤ 2. If W t ∈ Θ 0 , we update the ellipsoid by Θ t ← MVEE θ ∈ Θ t : θ -W t T • e j ≶ 0

where e j is the indicator vector of coordinate j in which W t exceeds 1, and the inequality direction depends on the sign of (W t ) j . If W t ∈ Θ 0 still, this process can be repeated for a finite number of steps until W t ∈ Θ 0 , as the volume of the ellipsoid is bounded from below and each update reduces the volume (Lemma 3). Now we have W

t ∈ Θ 0 , implying ||W * -W t || ∞ ≤ 2.

As no points of Θ 0 are removed this way, this does not affect the correctness of the proof. Similarly, we may assume ||W * t -W t || ∞ ≤ 2 as W * t ∈ Θ 0 . We denote W t which remains constant for each update in the batch by W . We define t(i) the timesteps corresponding to the demonstrations in the batch for i = 1, ..., n. We define z * ,H i to be the expected value of ẑ * ,H i , and z * i to be the outer product of c t(i) and the feature expectations of the expert policy for W * t(i) , c t(i) , ξ t(i) . We also denote W * t(i) by W * i . We bound the following term from below, as in Theorem 3:

EQUATION

(1): Since the sub-optimality criterion implies a difference in value of at least for the initial distribution which assigns 1 to the state where the agent errs, we may use identical arguments to the previous proof. Therefore, the term is bounded from below by .

( ) is never disqualified, and the number of updates is bounded by 2dk(dk + 1) log( 16k√ dk

H = (1 -(1 -γ)) H ≤ (1 -(1 -γ)) 1 1-γ (1-γ)H = (1 -(1 -γ))

(1-γ) ), and multiplied by n this yields the upper bound on the number of rounds in which a sub-optimal action is chosen. By union-bound, the required bound for term (4) holds in all updates with probability of at least 1 -δ.

E CONTEXTUAL POLICIES:

Consider the following problem. Given a complete specification of an MDP, and a hypothesis class H, for each state s assign a hypothesis h a : C → A such that the return is maximized.

The following theorem shows that it is NP-complete to find such a policy. We will use the class of linear separators. We define the following contextual MDP training problem. We are given an MDP with only transitions and a sample of m contexts and their reward (namely for each context c i we specify the rewards for each state and action). Theorem 5. There is a reduction from training problem of the union of k hyperplanes to the contextual MDP training problem of k + 1 states.

Proof. Consider the following MDP, which has two parameters r 0 and r 1 which will define the rewards for each context. The MDP composed from a line of k internal states i, and one sink state. Each internal state i ∈ [1, k -1] has two actions: action 1 leads to the sink state with reward r 1 and the action 0 lead to the next internal state i + 1 with reward 0. In internal state k action 1 leads to the sink state with reward r 1 and action 0 leads to the sink state with reward r 0 . In the sink state there is a singe action that stays in the sink state and has reward 0.

The context is C = R d . The hypotheses class includes all hyperplanes, each hyperplane is characterize by a weight vector w ∈ R d , and if w c ≥ 0 then the action is 1 and otherwise it is 0.

Given a sample of size m to the training problem (c i , y i ), we generate rewards for the contextual MDP training problem by having the context c i and specifying the parameters r 0 and r 1 . Specifically, we will set r 0 = 1 -y i and r 1 = y i .

Given k hyperplanes w 1 , . . . , w k , and m example, assume that they make e errors. By using those k hyperplanes in the k internal nodes for each context c i we have a reward 1 iff the union of the k hyperplanes classify it correctly. This implies that we have reward m -e. This is since each example that is labeled correctly will get a reward of 1 and each incorrect example will get a reward of 0.

Given an assignment of k hyperplanes to the internal nodes, which gets a return of m -e, we output as a hypothesis the union of the k hyperplanes. Again, the number of errors on the sample is exactly e errors.

:

Scaling of the reward by a constant does not affect the resulting policy, thus, these sets are not restricting.

We refer the reader to Appendix C for definitions of the Bregman distance, the dual norm, etc.

This procedure follows a sequence of linear algebra operations which we explain in the appendix.

MDA also works with near optimal experts due to the robustness of MDA. The analysis of this case is identical to the analysis of biased trajectories, as we discuss in the end of Section 3.1.

The data, code and implementation of our algorithms can be found in github.com/CIRLMDP/CIRL.

Figures and Tables:

Figure 2: Experimental results in the autonomous driving simulation with a linear mapping (a & b) and a nonlinear mapping (c & d)

Figure 3: Experimental results in the dynamic treatment regime with a linear mapping

Figure 4: Experimental results in the dynamic treatment regime with a non-linear mapping

Figure 5: Comparison between COIRL and AL on a large MDP

driving simulation . . . . . . . . . . . . . . . . . . . . . . . . . . . 14 A.2 Dynamic treatment regime . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 14 B Experimental details 15 B.1 Autonomous driving simulation . . . . . . . . . . . . . . . . . . . . . . . . . . . 15 B.2 Dynamic treatment regime . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 16 C Proofs for Section 3 18 D Proofs & pseudo code for Section 3.3 20 Ellipsoid Algorithm for trajectories . . . . . . . . . . . . . . . . . . . . . . . . . . 20 D.2 MVEE computation . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 20 D.3 Proof of Theorem 3 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 20 D.4 Proof of Theorem 4 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 21

Figure 6: Offline framework, autonomous driving simulation, linear mapping

Figure 7: Ellipsoid framework, dynamic treatment regime, linear mapping

Figure 9: Visualization of nonlinear decision boundaries

γ) ) ≤ e -log( 8k (1-γ) ) =(1-γ)  8k . (4): The partial sumsN i=1 (W * -W ) T • (z * ,H i -ẑ * ,H i ) for N = 0, ..., n form a martingale sequence. Note that ||z * ,H i || 1 ≤ k 1-γ , and ||ẑ * ,H i || 1 ≤ k 1-γ . Also, we have that ||W * -W t || ∞ ≤ 2, thus, we can apply Azuma's inequality (Lemma 4) with b = 4k(1-γ) and with our chosen n this yields:n i=1 (W * -W ) T • (z n ) > 4and as in Theorem 3 this shows B ∞ (W * , (1-γ) 8k

Comparison between various approaches.

Table content: <table/>

1, 2, 3, ... do c t is revealed, Let W t be the center of Θ t Play episode using πt = arg max π V π

Table content: <table><tr><td>c T t Wt</td><td/></tr><tr><td>Θ t+1 ← Θ t</td><td/></tr><tr><td>if a sub-optimal action a is played at state s then</td><td/></tr><tr><td>Expert provides H-step trajectory (s E 0 = s, s E 1 , ..., s E H ). Let x * ,H i</td><td>be the H-step sample of</td></tr><tr><td>the expert's feature expectations for</td><td/></tr></table>

): By assumption ||W * -W * i || ∞ ≤ (1-γ) 8k thus since ||(z * i -z i )|| 1 ≤ 2k1-γ by Hölder's inequality the term is bounded by 4 .(3):We have ||x * ,H i -x * i || 1 ≤ kγ H 1-γ from definitions, thus ||z * ,H i -z * i || 1 ≤ kγ H 1-γ since c ∈ ∆ d-1. As mentioned previously we may assume ||W * -W t || ∞ ≤ 2, therefore by Hölder's inequality the term is bounded by 4 due to our choice of H: γ

Table content: <table/>

Peer review (iclr2020):
The authors consider the problem of inverse reinforcement learning for CMDPs in which the reward function is a function of the context. They propose different algorithms for learning the function and evaluate their algorithms on a driving simulator and a sepsis treatment problem (based on real data from the MIMIC corpus). I think this is a good paper, studying an interesting problem and proposing useful solutions, so it should be accepted. The paper could be improved by putting more focus on the presentation of the studied problem and its variants (also the ones the authors mention their approach is easy to generalize to) and reducing a little the focus on the algorithm. In total they propose three algorithms and give convergence guarantees for all of them. Of course this analysis is important but I found it somewhat distracting from the main flow of the paper. Maybe these guarantees could be moved into a separate section. A few more points: * Please comment on why you see estimating contextual transition dynamics as an orthogonal problem. * In the last paragraph of the conclusion you talk about a safety critical application. For which applications do you think this is practical? I would assume that constantly reviewing of an AI systems* actions is very impractical. (But I agree on these aspects being important.) * One of the main limitations of this work seems to be that the CMDPM has to be known. Please comment on how one could expand the analysis/applications/experiments to extend to the case where the CMDP is not known.