Year: 2021

Abstract:
We consider the problem of using expert data with unobserved confounders for imitation and reinforcement learning. We begin by defining the problem of learning from confounded expert data in a contextual MDP setup. We analyze the limitations of learning from such data with and without external reward, and propose an adjustment of standard imitation learning algorithms to fit this setup. We then discuss the problem of distribution shift between the expert data and the online environment when the data is only partially observable. We prove possibility and impossibility results for imitation learning under arbitrary distribution shift of the missing covariates. When additional external reward is provided, we propose a sampling procedure that addresses the unknown shift and prove convergence to an optimal solution. Finally, we validate our claims empirically on challenging assistive healthcare and recommender system simulation tasks.

heading: 1 INTRODUCTION
text: Reinforcement Learning (RL) is increasingly used across many fields to create agents that learn via interaction and reward feedback. In many such cases, we rely on experts to perform certain tasks, integrating their knowledge to improve learning efficiency and overall performance. Imitation Learning (IL, Hussein et al. (2017)) is concerned with learning via expert demonstrations without access to a reward function. Similarly, RL settings often use expert data to boost performance, eliminating the need to learn from scratch. In this work we consider the IL and RL paradigms in the presence of partially observable expert data.
While expert demonstration data is useful, in many realistic settings such data may be prone to hidden confounding (Gottesman et al., 2019), i.e., there may be features used by the expert which are not observed by the learning agent. This can occur due to, e.g., privacy constraints, continually changing features in ongoing production pipelines, or when not all information available to the human expert was recorded. As we show in our work, covariate shift of unobserved factors between the expert data and the real world may lead to significant negative impact on performance, frequently rendering the data useless for imitation (see Figure 1 and Theorem 2).
In this paper we define the tasks of imitation and reinforcement learning using expert data with unobserved confounders and possible covariate shift. We focus on a contextual MDP setting (Hallak et al., 2015), where a context is sampled at every episode from some distribution, affecting both the reward and the transition between states. We assume that the agent has access to additional expert data, generated by an optimal policy, for which the sampled context is missing, yet is observed in the online environment.
We begin by analyzing the imitation-learning problem, (i.e., without access to reward) in Section 3. Under no covariate shift in the unobserved context, we characterize a sufficient and necessary set of optimal policies. In contrast, we prove that in the presence of a covariate shift, if the true reward depends on the context, then the imitation-learning problem is non-identifiable and prone to catastrophic errors (see Section 3.2 and Theorem 2). We further analyze the RL setting (i.e., with access to reward and confounded expert data) in Section 4. Figure 1 depicts a possible failure case of using confounded expert data with unknown covariate shift in a dressing task. Unlike the imitation setting, we show that optimality can still be achieved in the RL setting while using confounded expert data with arbitrary covariate shift. We use a corrective data sampling procedure and prove convergence to an optimal policy.
Our contributions are as follows. (1) We introduce IL and RL with hidden confounding and prove fundamental characteristics w.r.t. covariate shift and the feasibility of imitation. (2) In the RL setting, under arbitrary covariate shift, we provide a novel algorithm with convergence guarantees which uses a corrective sampling technique to account for the unknown context distribution in the expert data. (3) Finally, we conduct extensive experiments on recommender system (Ie et al., 2019) and assistive-healthcare (Erickson et al., 2020) environments, demonstrating our theoretical results, and suggesting that confounded expert data can be used in a controlled manner to improve the efficiency and performance of RL agents.
2 PRELIMINARIES
Online Environment. We consider a contextual MDP (Hallak et al., 2015) defined by the tuple M = (S,X ,A, P, r, ρo, ν, γ), where S is the state space, X is the context space, A is the action space, P : S × S × A ×X 7→ [0, 1] is the context dependent transition kernel, r : S × A ×X 7→ [0, 1] is the context dependent reward function, and γ ∈ (0, 1) is the discount factor. We assume an initial distribution over contexts ρo : X 7→ [0, 1] and an initial state distribution ν : S × X 7→ [0, 1]. The environment initializes at some context x ∼ ρo(·), and state s0 ∼ ν(·|x). At time t, the environment is at state st ∈ S and an agent selects an action at ∈ A. The agent receives a reward rt = r(st, at, x) and the environment then transitions to state st+1 ∼ P (·|st, at, x). We define a Markovian stationary policy π as a mapping π : S × X × A 7→ [0, 1], such that π(·|s, x) is the action sampling probability. We define the value of a policy π by vM(π) = Eπ[ (1− γ) ∑∞ t=0 γ
tr(st, at, x) | x ∼ ρo, s0 ∼ ν(· | x) ], where
Eπ denotes the expectation induced by the policy π. We denote by Π the set of all Markovian policies and Πdet the set of deterministic Markovian policies. We define the optimal value and policy by v∗M = maxπ∈Π vM(π), and π ∗ M ∈ arg maxπ∈Π vM(π), respectively. Whenever appropriate, we simplify notation and write v∗, π∗. We use Π∗M to denote the set of optimal policies inM, i.e., Π∗M = arg maxπ∈Π vM(π). We also define the set of catastrophic policies Π † M as the set
Π†M = arg min π∈Π vM(π). (1)
We later use this set to show impossibility of imitation under arbitrary covariate shift and a contextindependent transition function.
Expert Data with Unobserved Confounders. We assume additional access to a confounded dataset consisting of expert trajectories D∗ = { (si0, a i 0, s i 1, a i 1, . . . , s i H , a i H) }n i=1
, where aij ∼ π∗ ∈ Π∗M. The trajectories in the dataset were sampled i.i.d. from the marginalized expert distribution (under possible context covariate shift) P ∗(s0, a0, s1, a1, . . . , sH) = ∑ x ρe(x)ν(s0|x) ∏H−1 t=0 P (st+1|st, at, x)π∗(at|st, x), where ρe is some distribution over contexts. Importantly, ρe does not necessarily equal ρo – the distribution of contexts in the online environment. Notice that it is assumed that π∗ that generated the data had access to the context xi (i.e., π∗ is context-dependent), though it is missing in the data.
In this work, we consider two settings:
1. Confounded Imitation Learning (Section 3): The agent has access to confounded expert data (with context distribution ρe) as well as real environment (S,X ,A, P, ρo, ν, γ), without access to reward.
2. Reinforcement Learning with Confounded Expert Data (Section 4): The agent has access to confounded expert data (with context distribution ρe) as well as real environment M = (S,X ,A, P, r, ρo, ν, γ), with access to reward.
In both settings we aim to find a context-dependent policy which maximizes the cumulative reward. The confounding factor here is w.r.t. the unobserved context and distribution ρe in the offline data.
Marginalized State Action Distribution. We denote the state-action frequency of a policy π ∈ Π given context x ∈ X by dπ(s, a|x) = (1 − γ) ∑∞ t=0 γ
tPπ(st = s, at = a|x, s0 ∼ ν(·|x)), where Pπ denotes the probability measure induced by π. Similarly, given a distribution over contexts, we define the marginalized state-action frequency of a policy π under the corresponding context distribution by
dπρo(s, a) = Ex∼ρo [d π(s, a | x)] (online environment), dπ ∗ ρe (s, a) = Ex∼ρe [ dπ ∗ (s, a | x) ] (offline expert data).
A Causal Perspective. Our work sits at an intersection between the fields of RL and Causal Inference (CI). We believe it is essential to bridge the gap between these two fields, and include an interpretation of our model using CI terminology in Appendix D, where we equivalently define our objective as an intervention over the unknown distribution ρe in a specified Structural Causal Model, as depicted in Figure 2.

heading: 3 IMITATION LEARNING WITH UNOBSERVED CONFOUNDERS
text: In this section, we analyze the problem of confounded imitation learning, namely, learning from expert trajectories with hidden confounders and without reward. Similar to previous work, we consider the task of imitation learning from expert data in the setting where the agent is allowed to interact with the environment (Ho & Ermon, 2016; Fu et al., 2017; Kostrikov et al., 2019; Brantley et al., 2019). In the first part of this section we assume no covariate shift between the online environment and the data is present, i.e., ρe = ρo. We lift this assumption in the second part, where we focus on the covariate shift of the hidden confounders.
3.1 NO COVARIATE SHIFT: ρo = ρe
We first consider the scenario in which no covariate shift is present between the offline data and the online environment, i.e., ρo = ρe. We begin by defining the marginalized ambiguity set, a central component of our work.
Definition 1. For π ∈ Π, we define the set of deterministic policies that match the marginalized state-action frequency of π by Υπ = { π′ ∈ Πdet : dπ ′ ρo(s, a) = d π ρe(s, a) ∀ s ∈ S, a ∈ A } .
Recall that, in general, π∗ ∈ Π∗M may depend on the context x ∈ X . Therefore, the set Υπ∗ corresponds to all deterministic policies that cannot be distinguished from π∗ based on the observed
expert data. The following theorem shows that for any policy π∗ ∈ Π∗M and any policy π0 ∈ Υπ∗ , one could design a reward function r0, for which π0 is optimal, while the set Υπ∗ is indiscernible from Υπ0 , i.e., Υπ∗ = Υπ0 (see Appendix F for proof). In other words, Υπ∗ is the smallest set of candidate optimal policies. Theorem 1. [Sufficiency of Υπ∗] Assume ρe ≡ ρo. Let π∗ ∈ Π∗M and let π0 ∈ Υπ∗ . Then, Υπ∗ = Υπ0 . Moreover, if π0 6= π∗, then there exists r0 such that π0 ∈ Π∗M0 but π
∗ /∈ Π∗M0 , where M0 = (S,A,X , P, r0, ρo, ν, γ).
The above theorem shows that any policy in Υπ∗ is a candidate optimal policy, yet without knowing the context the expert used. Such ambiguity can result in selection of a suboptimal or even catastrophic policy. We provide a practical algorithm in Appendix B which calculates the ambiguity set Υπ∗ , and returns an average policy, with computational guarantees. In the next subsection we analyze a more challenging scenario, for which ρo 6= ρe. In this case, Υπ∗ may not be sufficient for the imitation problem.
3.2 WITH COVARIATE SHIFT: ρo 6= ρe
Next, we assume covariate shift exists between the online environment and the expert data, i.e., ρo 6= ρe. Particularly, without further assumptions on the extent of covariate shift, we show two extremes of the problem. In Theorem 2 we prove that whenever the transitions are independent of the context, the data cannot in general be used for imitation. In contrast, in Theorem 3 we prove that, whenever the reward is independent of the context, the imitation problem can be efficiently solved.
Clearly, if Supp(ρo) 6⊆ Supp(ρe)1 then there exists x ∈ Supp(ρo) for which π∗ is not identifiable from the expert data2. We therefore assume throughout that Supp(ρo) ⊆ Supp(ρe). We begin by defining the set of non-identifiable policies as those that cannot be distinguished from their respective state-action frequencies without information on ρe.
Definition 2. We say that {πi}ki=1 are non-identifiable policies if there exist {ρi} k i=1 such that dπiρi (s, a) = d πj ρj (s, a) for all i 6= j.
Next, focusing on catastrophic policies (recall Equation (1)), we define catastrophic expert policies as those which could be either optimal or catastrophic under ρo for different reward functions.
Definition 3. We say that {πi}ki=1 are catastrophic expert policies if there exist {ri} k i=1 such that for all i, πi ∈ Π∗Mi , and ∃j ∈ [k], j 6= i such that πi ∈ Π † Mj , whereMj = (S,X ,A, P, rj , ρo, ν, γ).
Using the fact that both ρe and r are unknown, the following theorem shows that whenever P (s′|s, a, x) is independent of x, one could find two policies which are non-identifiable, catastrophic expert policies (see Appendix F for proof). In other words, in the case of context-independent transitions, without further information on ρe or r the expert data is useless for imitation. Furthermore, attempting to imitate the policy using the expert data could result in a catastrophic policy. Theorem 2. [Catastrophic Imitation] Assume |X | ≥ |A|, and P (s′|s, a, x) = P (s′|s, a, x′) for all x, x′ ∈ X . Then ∃πe,1, πe,2 s.t. {πe,1, πe,2} are non-identifiable, catastrophic expert policies.
1For a distribution P we denote by Supp(P) the support of P. 2We define non-identifiability in Definition 2. We use a similar notion of identifiability as in Pearl (2009b)
While Theorem 2 shows the impossibility of imitation for context-free transitions, whenever the reward is independent of the context, the imitation problem becomes feasible. In fact, as we show in the following theorem, for context-free rewards, any policy in Υπ∗ is an optimal policy.
Theorem 3. [Sufficiency of Context-Free Reward] Assume Supp(ρo) ⊆ Supp(ρe) and r(s, a, x) = r(s, a, x′) for all x, x′ ∈ X . Then Υπ∗ ⊆ Π∗M.
Theorems 2 and 3 suggest that the hardness of the imitation problem under covariate shift lies on a wide spectrum (as depicted in Figure 3). While dependence of the transition P (s′|s, a, x) on x provides us with information to identify x in the expert data, the dependence of the reward r(s, a, x) on x increases the degree of confounding in the imitation problem. Both of these results are concerned with arbitrary confounding.
Bounded Confounding: A Sensitivity Perspective. A common approach in causal inference is to bound the bias of unobserved confounding through sensitivity analysis (Hsu & Small, 2013; Namkoong et al., 2020; Kallus & Zhou, 2021). In our setting, this confounding bias occurs due to a covariate shift of the unobserved covariates. As we’ve shown in Theorem 2, though these covariates are observed in the online environment, their shifted and unobserved distribution in the offline data can render catastrophic results. Therefore, we consider the odds-ratio bounds of the sensitivity in distribution between the online environment and the expert data, as stated formally below.
Assumption 1 (Bounded Sensitivity). We assume that Supp(ρe) ⊆ Supp(ρo) and that there exists some Γ ≥ 1 such that for all x ∈ Supp(ρe), Γ−1 ≤ ρo(x)(1−ρe(x))ρe(x)(1−ρo(x)) ≤ Γ.
Next, we define the notion of δ-ambiguity, a generalization of the ambiguity set in Definition 1.
Definition 4 (δ-Ambiguity Set). For a policy π ∈ Π, we define the set of all deterministic policies that are δ-close to π by Υδπ = { π′ ∈ Πdet : ∣∣∣dπ′ρo(s, a)− dπρe(s, a)∣∣∣ < δ, s ∈ S, a ∈ A}. Similar to Definition 1, the δ-ambiguity set considers all deterministic policies with a marginalized state-action frequency of distance at most δ from π. The following results shows that ΥΓ−1π∗ is a sufficient set of candidate optimal policies, as long as Assumption 2 holds for some Γ ≥ 1.
Theorem 4. [Sufficiency of ΥΓ−1π∗ ] Let Assumption 2 hold for some Γ ≥ 1. Then π∗ ∈ ΥΓ−1π∗ .
For the interested reader, we further analyze the case of bounded confounding in Appendix C. We also demonstrate the effect of bounded confounding in Section 5. In the following section, we show that, while arbitrary confounding may result in catastrophic results for the imitation learning problem, when coupled with reward, one can still make use of the expert data.

heading: 4 USING EXPERT DATA WITH UNOBSERVED CONFOUNDERS FOR RL
text: In the previous section we showed sufficient conditions under which imitation is possible, with and without covariate shift. When covariate shift is present, but unknown, the imitation learning problem may be hard, or even impossible (see Theorem 2, catastrophic imitation). We ask, had we had access to the reward function, would the expert data be useful under arbitrary covariate shift? In this section we show that expert data with unobserved confounders can be used to converge to an expert policy, even when arbitary covariate shift is present. In our experiments (Section 5) we empirically show that using our method can also improve overall performance.
We view the confounded expert data as side information to the RL problem. Specifically, we assume access to the true reward signal in the online environment and wish to leverage the offline expert data to aid the agent in converging to an optimal policy. To do this, we define an optimization problem that maximizes the cumulative reward, while minimizing an f -divergence (e.g., KL-divergence, TVdistance, χ2-divergence) of state-action frequencies in Υπ∗ ,
max π∈Π
Ex∼ρo,s,a∼dπ(s,a|x)[r(s, a, x)]− λDf (d π ρo(s, a)||d
π∗ ρe (s, a)). (P1)
Algorithm 1 RL using Expert Data with Unobserved Confounders (Follow the Leader)
1: input: Expert data with missing context D∗, λ > 0, policy optimization algorithm ALG-RL 2: init: Policy π0 3: for k = 1, . . . do 4: ρs ← argminρDKL(dπ k−1 ρo (s, a)||d π∗ ρ (s, a))
5: gk ← 1 k
( gk−1 + E
s,a∼dπk−1ρo
[ 1
dπ ∗ ρs (s,a) ]) 6: πk ← ALG-RL(r(s, a, x)− λgk(s, a)) 7: end for
Algorithm 2 RL using Expert Data with Unobserved Confounders (Online Gradient Descent)
1: input: Expert data with missing context, λ,B,N > 0, policy optimization alg. ALG-RL 2: init: Policy π0, bonus reward network gθ 3: for k = 1, . . . do 4: ρs ← argminρDf (d πk−1 ρo (s, a)||dπ ∗ ρ (s, a)) 5: for e = 1, . . . N do 6: Sample batch {si, ai}Bi=1 ∼ d πk−1 ρo (s, a) 7: Sample batch {sei , aei}Bi=1 ∼ d π∗ ρs (s, a) 8: Update gθ according to ∇θL(θ) = 1 B ∑B i=1∇θ[f ∗(gθ(s e i , a e i ))− gθ(si, ai)]
9: end for 10: πk ← ALG-RL(r(s, a, x)− λgθ(s, a)) 11: end for
Here, λ > 0 and Df is the f -divergence, where f is a convex function f : (0,∞) 7→ R. The solution to Problem (P1) is an optimal policy π∗ ∈ Π∗M as long as ρo ≡ ρe. Rewriting Df using its variational form (see Appendix A for background on the variational form of f -divergences), we get the following equivalent optimization problem, motivated by Nachum et al. (2019):
max π∈Π min g:S×A7→R Ex∼ρo,s,a∼dπ(s,a|x)[r(s, a, x) + λg(s, a)]− λEs,a∼dπ∗ρe (s,a)[f ∗(g(s, a))], (P1b)
where f∗ is the convex conjugate of f , i.e., f∗(y) = supx xy − f(y). Unfortunately, when covariate shift exists (i.e., ρo 6= ρe), Problems (P1) and (P1b) are not ensured to converge to an optimal policy (Theorem 2). Instead, we propose to reformulate Problem (P1b) using a distribution ρs which minimizes the f -divergence, as follows,
max π∈Π min g:S×A7→R ρs∈B(X )
Ex∼ρo,s,a∼dπ(s,a|x)[r(s, a, x) + λg(s, a)]− λEs,a∼dπ∗ρs (s,a)[f ∗(g(s, a))]. (P2)
Here, B(X ) denotes the set of probability measures on the Borel sets of X , and dπ ∗ ρs (s, a) = Ex∼ρs [ dπ ∗ (s, a | x) ] . Indeed, whenever Supp(ρo) ⊆ Supp(ρe), we have that (π, ρs) = (π ∗, ρo) is a solution to Problem (P2). That is, unlike Problems (P1) and (P1b), Problem (P2) can achieve an optimal solution to the RL problem which still uses the expert data.
Corrective Trajectory Sampling (CTS). Solving Problem (P2) involves an expectation over an unknown distribution, dπ ∗
ρs (s, a). Fortunately, d π∗
ρs (s, a) can be equivalently written as an expectation over trajectories inD∗, rather than expectation over unobserved contexts, as shown by the following proposition (see Appendix F for proof): Proposition 1. [Trajectory Sampling Equivalence] Let ρ∗s which minimizes Problem (P2) for some π ∈ Π, g : S × A 7→ R, and assume Supp(ρo) ⊆ Supp(ρe). Then, there exists pn ∈ ∆n such that dπ ∗
ρ∗s (s, a) = lim n→∞ Ei∼pn
[ (1− γ) ∑∞ t=0 γ t1 { (sit, a i t) = (s, a) }] .
Proposition 1 allows us to estimate the inner minimization problem over ρs in Problem (P2) using trajectory samples. Particularly, we uniformly sample k distributions pn1 , . . . p n k , where p n j ∈ ∆n, and then estimate
min ρs
Df (d π ρo ||d
π∗ ρs ) ≈ minj∈{1,...,k}
{ Df ( dπρo(s, a) ∣∣∣∣∣∣ Ei∼pnj [ ∞∑ t=0 γt1 { (sit, a i t) = (s, a) }])} , (2)
which can be estimated by using the variational form ofDf (see Appendix A). We call this procedure Corrective Trajectory Sampling (CTS), as it uses complete trajectory samples to account for the unknown context distribution ρe.
Solving Problem (P2). Algorithm 1 provides an iterative procedure for solving the optimization problem in Problem (P2). It uses alternative updates of a cost player (line 5) and policy player (line 6). In line 5 the gradient of DKL w.r.t. dπ is taken using a Follow the Leader (FTL) cost player
to estimate the next bonus iterate. Finally, in line 6, an efficient, approximate policy optimization algorithm ALG-RL is executed using an augmented reward. The following theorem, provides convergence guarantees for Algorithm 1 (see Appendix F for proof based on Zahavy et al. (2021)).
Theorem 5. Let ALG-RL be an approximate best response player that solves the RL problem in iteration k to accuracy k = 1√k . Then, Algorithm 1 will converge to an -optimal solution to
Problem (P2) in O ( 1 4 ) samples.
Notice that, while Theorem 5 shows Algorithm 1 converges to an optimal policy, it does not determine whether the expert data improves overall learning efficiency. We leave this theoretical question for future work. Nevertheless, in the following section we conduct extensive experiments to show that such data can indeed improve overall performance on various tasks.
A drawback of Algorithm 1 is that it needs to estimate the state-action frequencies. A practical implementation of Algorithm 1 using online gradient descent (OGD) is provided in Algorithm 2 – the algorithm does not require approximate estimates of the state-action frequencies, but rather, only the ability to sample from them. Similar to Algorithm 1, we use CTS (see Equation (2)) to estimate ρs in line 4 according to some f -divergence. Here, samples are drawn from the current policy as well as samples from D∗ (with CTS). We write Df in its variational form, and use a neural network representation for gθ. We then use the aforementioned samples to minimize the f -divergence using OGD. Finally, the policy is updated using ALG-RL and an augmented reward.

heading: 5 EXPERIMENTS
text: We tested our proposed approach for using expert data with hidden confounding in recommendersystem and assistive-healthcare environments. For all our experiments we used χ2-divergence as our choice of f -divergence, as we found it to work best. Comparison to other divergences is provided in Figure 5 (left). We used PPO (Schulman et al., 2017) implemented in RLlib (Liang et al., 2018) for both the imitation as well as RL settings. We include specific choice of hyperparameters and an exhaustive overview of further implementation details in Appendix E.
Assistive Healthcare. Consider the challenge of providing physical assistance to disabled persons. A recently proposed set of tasks for assistive-healthcare, simulating autonomous robots as versatile caregivers (Erickson et al., 2020). Each task has a unique goal, affected by both the physical world as well as the patient specific preferences and disabilities.
We tested our algorithm on four tasks: feeding, dressing, bathing, and drinking. In these, we used the following features to define user context: gender, mass, radius, height, patient impairment, and patient preferences. The patient’s mass, radius, and height distributions were dependent on gender. The patient’s impairment was given by either limited movement, weakness, or tremor (with sporadic movement). Finally, the patient’s preferences were affected by the velocity and pressure of touch forces applied by the robot. For the context distribution ρo we used the default values as provided by the original environment. To enforce a distribution shift in the expert data, we shifted each distribution randomly with an additive factor β · d̃x, where β ∈ [0, 1], and d̃x was a random distribution chosen from a set of shifting distributions (see Appendix E). Here β corresponds to the covariate shift strength. The expert data was generated by a fixed policy trained using a dense reward function. A sparse reward signal was used for executing our experiments with the confounded expert data. For further details, we refer the reader to Appendix E.
Figure 4 depicts results for executing Algorithm 2 on four assistive-gym environments with various covariate shift strengths. As evident in most of the enviornments, covariate shift strongly affected overall performance. Particularly in the feeding, drinking, and dressing environments, the success of reaching the goal (i.e., spoon to mouth, cup to mouth, and sleeve to hand) was highly affected by the degree of covariate shift. This is due to the changing distribution of size, movement, and preferences of the patient, and thus of the goal. Nevertheless, in all environments, using the expert data (with and without CTS) was found to help induce better policies than executing the same RL algorithm without expert data. This suggests that expert data can assist in improving overall RL performance, yet correcting for covariate shift may significantly improve it in these domains.
Recommender Systems. In practical recommender systems, sequential interaction with users presents a great challenge for optimizing user long-term engagement and overall satisfaction (Ie et al., 2019). Leveraging expert data collected using, e.g., surveys to users, may greatly benefit future solutions. Because features are repeatedly added to these systems, full information in the data is rarely available. Here, we use the recently proposed RecSim Interest Evolution environment (Ie et al., 2019), simulating sequential user interaction with a slate-based recommender system. The environment consists of a document model for sampling documents, a user model for defining a distribution over user context features, and a user choice model, which defines the intent of the user based on observable document features and the user’s sampled context (e.g., personality, satisfaction, interests, demographics, and other behavioral features such as session length or visit frequency).
We used a slate of 10 documents and a user context of dimension 20. To test the severity of the implications of Theorem 2 in the confounded imitation setting, we used a user-model sampled from a Beta-distribution. Specifically, for the expert data the user context features x = (x0, . . . , x19) were sampled from a Beta-distribution, where xi ∼ Beta(αi, 4), and αi = 1.5 + 8.519 i. In contrast, the online environment features were sampled from a shifted Beta-distribution with αi = (1− β) ( 1.5 + 8.519 i ) + β ( 10− 8.519 i ) , where β ∈ [0, 1] defined the shift strength. While the
original environment used a uniform distribution to generate user contexts, the Beta-distribution let us analyze severe forms of covariate shifts, testing the limits of our results in Sections 3 and 4.
Figure 5(a) depicts the effect of increased covariate shift on imitation in the RecSim environment with a dataset of 100 expert trajectories (generated by an optimal policy that had access to the full context). Without covariate shift (β = 0) an optimal score is achieved, and as β increases, performance decreases. Particularly, as the mirrored distribution is reached (β = 1), a catastrophic policy is reached. While the imitator “believes” to have reached an optimal policy, it has in fact reached a catastrophic one, as shown by the orange plot. Conversely, Figure 5(b) depicts the benefit of using confounded expert data in the RL setting (with an online reward signal). Though strong confounding is present, the agent is capable of leveraging the data to improve overall learning performance.

heading: 6 RELATED WORK
text: Imitation Learning. The imitation learning problem has been extensively studied in both the fully offline (Pomerleau, 1989; Bratko et al., 1995) as well as online setting (Ho & Ermon, 2016; Fu et al., 2018; Kim & Park, 2018; Brantley et al., 2019). Specific to our work are GAIL (Ho & Ermon, 2016), AIRL (Fu et al., 2017), and DICE (Kostrikov et al., 2019), which use distribution matching methods. Our work generalizes these settings to imitation with hidden confounders.
Reinforcement Learning with Expert Data. Much work has revolved on leveraging offline data for RL. Recently, offline RL (Levine et al., 2020) has shown great improvement over regular offline imitation techniques (Kumar et al., 2020; Kostrikov et al., 2021; Tennenholtz et al., 2021a; Fujimoto & Gu, 2021). In the online RL setting, the combination of offline data to improve RL efficiency has shown great success (Nair et al., 2020). KL-regularized techniques (Peng et al., 2019; Siegel et al., 2019) as well as DICE-based algorithms (Nachum et al., 2019) have also shown efficient utilization of offline data. Our work generalizes the latter to the confounded setting.
Intersection of Causal Inference and Imitation Learning. Closely related to our work is that of Zhang et al. (2020). There, the authors suggest a notion of imitability, showing when observational data can help identify a policy under some partially observed structural causal model. Our work provides an alternative perspective on the problem. In contrast to their work, we rely on concurrent imitation approaches (i.e. state-action frequency matching techniques) and importantly, allow access to the online environment. Furthermore, we provide guarantees and practical algorithms for both the imitation and RL settings. We refer the reader to Appendix D for an interpretation of our framework in CI terminology, from a perspective of stochastic interventions in a Structural Causal Model.
Another intersection with causal inference discusses the problem of causal confusion in imitation (de Haan et al., 2019). Causal confusion is concerned with the problem of nuisances in observed confounded data due to an unknown causal structure. These “causal misidentifications” can lead to spurious correlations and catastrophic failures in generalization. In contrast, our work discusses the orthogonal problem of hidden confounders with possible covariate shift.
Intersection of Causal Inference and Reinforcement Learning. Previous work has analyzed the problem of optimal control from logged data with unobserved confounders (Lattimore et al., 2016), as well as utilizing (non-expert) confounded data for online interactions (Tennenholtz et al., 2021b). Much work has revolved around the reinforcement learning setup with access to (nonexpert) confounded data (Zhang & Bareinboim, 2019; Wang et al., 2020). Other work has considered the problem of off-policy evaluation from confounded data (Tennenholtz et al., 2020; Oberst & Sontag, 2019; Kallus & Zhou, 2020). Our work is focused on leveraging expert data with hidden confounders and possible covariate shift in both the imitation and the RL settings.

heading: 7 CONCLUSION
text: This work presented and analyzed the problem of using expert data with hidden confounders for both the imitation and RL settings. We showed that covariate shift of hidden confounders between the expert data and the online environment can result in learning catastrophic policies, rendering imitation learning hard, or even impossible (Theorem 2). In addition, we showed that when a reward is provided, using the expert data is still possible under arbitrary hidden covariate shift (Theorem 5), and proposed new algorithms for tackling this problem using corrective trajectory sampling (CTS).

heading: APPENDIX

heading: A BACKGROUND: DISTRIBUTION MATCHING FOR IMITATION LEARNING
text: A common approach used in (non-confounded) imitation learning is matching the policy’s state-action frequency dπρo to the offline target distribution d π∗
ρe . Consider a source distribution p ∈ ∆N and target distribution q ∈ ∆N . GAIL (Ho & Ermon, 2016) uses the distribution ratio objective log(p/q), which can estimated using a GAN-like objective DR(p||q) = supg:Z7→(0,1) Ep[log(g(z))] + Eq[log(1− g(z))], to match the distribution p to q.
This technique can be generalized to f -divergences (Csiszár & Shields, 2004; Liese & Vajda, 2006; Kostrikov et al., 2019; Ke et al., 2020). Specifically, we wish to minimize a discrepancy measure from p to q, namely minp∈KD(p||q). For a convex function f : [0 :∞) 7→ R, the f -divergence of p from q is defined by Df (p||q) = Eq [ f ( p q )] . DICE (Kostrikov et al., 2019) uses the variational representation of the f -divergence,
Df (p||q) = sup g:Z7→R Ep[g(z)]− Eq[f∗(g(z))],
where f∗ is the Fenchel conjugate of f defined by f∗(y) = supx xy − f(y). The convex conjugate has closed form solutions for the total variation distance, KL-divergence, χ2-divergence, Squared Hellinger distance, Le Cam distance, and Jensen-Shannon divergence. Using the variational representation of the f -divergence we can estimate Df using samples from p and q. Table 1 presents examples of various f -divergences and their respective dual formulation. We also add the distribution ratio for comparison to the table, though it is not an f -divergence.

heading: B CONFOUNDED IMITATION - ALGORITHM AND CONVERGENCE GUARANTEES

heading: B.1 A TOY EXAMPLE
text: To gain intuition, we start with a simple toy example. Consider the three-state example depicted in Figure 6. Here, the environment initiates at state A w.p. 1, after which the agent can choose to (deterministicaly) transition to state B or C. The agent then receives a reward depending on the context. The optimal policy is given by π∗(a|s, x) = 1{a = aB , x = x1}+ 1{a = aC , x = x2} for s = A, and any action is optimal for s 6= A. Without loss of generality we assume π∗(aB |B, x) = π∗(aC |C, x) = 1. We turn to analyze the marginalized stateaction frequency, which uniquely defines the set of optimal policies (Puterman, 2014). Denoting ρe(x1) = ρ, we have that dπ ∗ ρe (s, a) = ρd π∗(s, a|x1) + (1− ρ)dπ ∗ (s, a|x2). Then, dπ ∗
ρe (s, a) = (1− γ)1{s = A}+ ργ1{s = B, a = aB}+ (1− ρ)γ1{s = C, a = aC}.
Algorithm 3 Confounded Imitation 1: input: Expert data with missing context D∗ ∼ dπ∗ρe , λ > 0, sensitivity bound δ ≥ 0. 2: init: Υ = ∅ 3: for n = 1, . . . do 4: Sample u(s, a) ∼ U [0, δ],∀s, a 5: L∗(π; g0) := Es,a∼dπρo (s,a)[g0(s, a)]− Es,a∼dπ∗ρe (s,a)+u(s,a)[g0(s, a)] 6: Li(π; gi) := Ex∼ρo,s,a∼dπ(s,a|x)[gi(s, a, x)]− Ex∼ρo,s,a∼dπi (s,a|x)[gi(s, a, x)] , i ≥ 1 7: Compute πn by solving
min π∈Πdet max |g0|≤ 12 ,|gi|≤ 1 2
{ L∗(π; g0(s, a))− λmin
i Li(π; gi(s, a, x))
} (4)
8: if πn ∈ Υ then 9: Terminate and return π̄(a|s, x) = ∑n−1 i=1 d
πi (s,a,x)∑n−1 i=1 ∑ a′ d πi (s,a′,x)
10: else 11: Υ = Υ ∪ {πn} 12: end if 13: end for
No Covariate Shift. Suppose ρo = ρe, and ρ = 12 . Trivially d π∗ ρe (s, a) = d π∗
ρo (s, a). We define the (suboptimal) policy
π0(a|A, x) = 1− π∗(a|A, x) , a ∈ A, x ∈ X . (3)
It can be verified that dπ ∗
ρe (s, a) = d π0 ρo (s, a) still holds, yet π0 is catastrophic (Equation (1)) with
value zero. A question arises: can we show that π0 is a suboptimal policy given access to the expert data (i.e., access to dπ ∗
ρe (s, a)) and a forward model P (s ′|s, a, x)?
Unfortunately, one cannot prove that π0 is suboptimal. Informally, notice that π0 is an optimal policy for an alternative reward function, r0(s, a, x) = 1− r(s, a, x), yet is catastrophic w.r.t. the true reward r. Indeed, since r is unknown and dπ0ρo (s, a) = d π∗
ρo (s, a), we cannot reject r0 (i.e., we cannot conclude that r0 is not the true reward). In other words, one cannot use the data to differentiate which of {π0, π∗} is the optimal policy.
With Covariate Shift. Next, assume ρo 6= ρe, and define π0 as in Equation (3). Let ρ̃e = 1− ρe and recall that ρe(x1) = ρ. Then, we have that
dπ0ρ̃e (s, a) = (1− ρ)d π0(s, a|x1) + ρdπ0(s, a|x2) = (1− ρ)dπ
∗ (s, a|x2) + ρdπ ∗ (s, a|x1) = dπ ∗
ρe (s, a).
Indeed, the expert data is incapable of distinguishing π0 and π∗, since dπ0ρ̃e = d π∗
ρe , and ρe is unknown. Unfortunately, as we’ve shown previously, π0 achieves value zero. Notice that, unlike the previous section, one cannot distinguish π∗ from the catastrophic policy π0 for any choice of ρo.

heading: B.2 A PRACTICAL ALGORITHM
text: The ambiguity set of Theorem 1 may contain suboptimal policies. Instead of randomly selecting a policy from this set, we can choose the average policy. The following proposition shows that such a selection is favorable.
Proposition 2. Define the mean policy π̄(a|s, x) = ∑ π∈Υπ∗
dπ(s,a,x)∑ π∈Υπ∗ ∑ a′ d π(s,a′,x) , and denote
α∗ = |Π∗M
⋂ Υπ∗ |
|Υπ∗ | ∈ [0, 1]. Then, vM(π̄) ≥ α∗v∗ + (1− α∗) minπ∈Υπ∗ vM(π).
Remark 1. Note that π̄ is generally not the average policy 1|Υπ∗ | ∑ π∈Υπ∗ π(a|s, x).
Remark 2. In an episodic setting, π̄ can be estimated by uniformly sampling a policy π ∈ Υπ∗ at the beginning of the episode, and playing it until the environment terminates.
Algorithm 3 describes our method for calculating the ambiguity set of Theorem 1, and returns π̄ of Proposition 2. At every iteration of the algorithm, we find a new policy in the set by minimizing the total variation distance (written in variational form) between dπ ∗
ρo (s, a) and d π ρo(s, a), while regular-
izing it with the distance between π and all previously collected πi ∈ Υ. Algorithm 3 also uses a sensitivity parameter δ ≥ 0 (defined formally in Appendix C) whenever bounded covariate shift is present. For this section we assume δ = 0.
In practice, the functions L∗ and Li in lines 4 and 5 are estimated using samples from trajectories of π, πi, and D∗. We then solve the min-max problem of Equation (4) using a parametric representations of gi and online gradient decent. The following proposition states that Algorithm 3 indeed retrieves the set Υπ∗ .
Proposition 3. Assume ρe ≡ ρo and |Υπ∗ | < ∞. Then there exists λ∗ > 0 such that for any λ ∈ (0, λ∗), Algorithm 3 (with δ = 0 sensitivity) will return π̄ of Proposition 2 after exactly |Υπ∗ | iterations.
B.3 IMITATION WITH CONTEXT-FREE REWARD
We tested Algorithm 3 on both the RecSim environment as well as a four-rooms environment with random instantiations of walls. Experiments for the RecSim environment are readily provided in Section 5. Here we describe our simple four-rooms environment and show experiments w.r.t. Theorem 3.
The four-rooms environment, as depicted in Figure 7, is a 15× 15 grid-world in which an agent can take one of four actions: LEFT, RIGHT, UP, or DOWN. Each action moves the agent in the specified direction whenever no obstacle is present. The agent (shown in blue) must reach the (green) goal while avoiding the (red) mine. When the goal is reached the agent receives a reward of +1 and the episode terminates. In contrast, if the agent reaches the mine, she receives a reward of −1 and the episode terminates. The state space of the environment consists of the agent’s (row, col) position in the world. The rest of the information in the environment is defined by the context x. Particularly, the context is defined by the position of the green goal, the position of the red mine, and the specific instantiation of walls (two instantiations are depicted in Figure 7).
We trained an agent with full information (i.e., observed context, including goal location, mine location, and walls). We generated expert data w.r.t. the trained agent. To demonstrate the result of Theorem 3 we executed Algorithm 3 with both a shifted distribution and the default distribution of walls. We did not change the distribution of goal and mine. Note that since the distribution of walls only affects the transition function and not the reward, we expect, by Theorem 3, the optimal solution to remain the same. Indeed, as shown in Figure 7 after training an agent with no access to the contextual information of the walls in the expert data, the agent achieved comparable results both with and without covariate shift on the distribution of walls.
This result seem surprising at first, as the walls are essential for solving the task at hand. Nevertheless, since the distribution of wall is observed in the online environment, the partially observed expert data suffices to obtain an optimal policy. This settles with Theorem 3 which indeed states that this information is not needed in the expert data in order to obtain an optimal policy.

heading: C BOUNDED HIDDEN CONFOUNDING
text: In this section we discuss the imitation learning problem under bounded hidden confounders. There are several ways to define boundness of unobserved confounders. In Section 3 we showed that, under arbitrary covariate shift and context-free transitions, the imitation learning problem is impossible, i.e., one cannot rule out a catastrophic policy. We begin by considering the effect of bounded covariate shift, i.e., ρoρe ≤ C. We then consider almost-context-free rewards, showing a tradeoff w.r.t. the hardness of the imitation problem.
A Sensitivity Perspective. A common approach in causal inference is to bound the bias of unobserved confounding through sensitivity analysis (Hsu & Small, 2013; Namkoong et al., 2020; Kallus & Zhou, 2021). In our setting, this confounding bias occurs due to a covariate shift of the unobserved covariates. As we’ve shown in Theorem 2, though these covariates are observed in the online environment, their shifted and unobserved distribution in the offline data can render catastrophic results. Therefore, we consider the odds-ratio bounds of the sensitivity in distribution between the online environment and the expert data, as stated formally below.
Assumption 2 (Bounded Sensitivity). We assume that Supp(ρe) ⊆ Supp(ρo) and that there exists some Γ ≥ 1 such that for all x ∈ Supp(ρe)
Γ−1 ≤ ρo(x)(1− ρe(x)) ρe(x)(1− ρo(x)) ≤ Γ.
Next, we define the notion of δ-ambiguity, a generalization of the ambiguity set in Definition 1.
Definition 5 (δ-Ambiguity Set). For a policy π ∈ Π, we define the set of all deterministic policies that are δ-close to π by
Υδπ = { π′ ∈ Πdet : ∣∣∣dπ′ρo(s, a)− dπρe(s, a)∣∣∣ < δ, s ∈ S, a ∈ A}. Similar to Definition 1, the δ-ambiguity set considers all deterministic policies with a marginalized state-action frequency of distance at most δ from π. The following results shows that ΥΓ−1π∗ is a sufficient set of candidate optimal policies, as long as Assumption 2 holds for some Γ ≥ 1. Theorem 6. [Sufficiency of ΥΓ−1π∗ ] Let Assumption 2 hold for some Γ ≥ 1. Then π∗ ∈ ΥΓ−1π∗ .
The above result suggests that Algorithm 3 can be executed over ΥΓ−1π∗ by adding δ = Γ − 1 additive uniform noise to dπ ∗
ρe (s, a) (see Line 4 of Algorithm 3), and executing the algorithm for a finite number of iterations, finally selecting an average policy from the approximate set.
Context Reconstruction. When bounded covariate shift is present, one might attempt to learn an inverse mapping of contexts from observed trajectories in the data.
We denote by Pπρ the probability measure over contexts x ∈ X and trajectories τ = (s0, a0, s1, a1, . . . sH) as induced by the policy π and context distribution ρ. That is,
Pπρ (x, τ) = ρ(x)ν(s0|x) H−1∏ t=0 P (st+1|st, at, x)π(at|st, x).
As the true context is observed in the online environment, we can calculate for any π the quantity Pπρo(x, τ). As the expert data distribution was generated by the marginalized distribution Pπ ∗ ρo (τ) = ∑ x∈X P π∗ ρo (x, τ), it is unclear if knowledge of P π ρo(x, τ) is beneficial.
Fortunately, whenever Assumption 2 holds, a high probability of reconstructing a context in the online environment induces a high probability of reconstructing it in the expert data. To see this, assume that there exists δ ∈ [0, 1] such that for all π ∈ ΥΓ−1π∗ , τ ∈ Supp(Pπρe(τ)), there exists x ∈ X such that
Pπρo(x|τ) ≥ min{(1− δ)(ρo(x) + Γ(1− ρo(x))), 1}. (5)
That is, we assume that for any policy that δ-ambiguous to π∗, and any induced trajectory of x ∈ X , one can with high probability identify x in the online environment. Importantly, this property can be verified in the online environment. When Assumption 2 and 5 hold, we get that
Pπρe(x|τ) = Pπρe(τ |x)ρe(x) Pπρe(τ) ≥ Pπρe(τ |x) Pπρe(τ)
ρo(x)
ρo(x) + Γ(1− ρo(x)) = Pπρo(x|τ) ρo(x) + Γ(1− ρo(x)) ≥ 1− δ.
In other words, we can reconstruct x with probability 1 − δ for any trajectory τ which satisfies the above. This allows us to deconfound essential parts of the expert data, rendering it useful for the imitation problem, even when reward is not provided. We leave further analysis of this direction for future work.
Context-Dependent Reward. In Theorem 3 we showed that whenever the reward is independent of the context then the imitation problem is easy, in the sense that any policy π0 ∈ Υπ∗ is also an optimal policy. Here, we relax the assumption on the reward, and instead assume bounded dependence of the reward on the context. The following definition upper bounds the confounding effect of the reward w.r.t. the context.
Definition 6. Let : X 7→ R such that
min r0:S×A7→R
|r(s, a, x)− r0(s, a)| ≤ (x) , s ∈ S, a ∈ A, x ∈ X
Using the above definition, we can now show that any policy in Υπ∗ is still approximately optimal, as shown by the following result. Theorem 7. [Context Dependent Reward] Let : X 7→ R of Definition 6. Denote oe = Ex∼ρo(x)[ (x)] + Ex∼ρe(x)[ (x)]. Then for any π∗ ∈ Π∗M, π0 ∈ Υπ∗
v(π0) ≥ v(π∗)− oe.
A direct corollary for the above result states that for : X 7→ R of Definition 6, if (x) = , for all x ∈ X , then for π0, π∗ of Theorem 7, it holds that v(π0) ≥ v(π∗) − 2 . That is, π0 is an approximately optimal policy.

heading: D RELATION TO CAUSAL INFERENCE
text: Our work is focused on the problem of hidden confounders in expert data for imitation and reinforcement learning. We have chosen to write the paper in terminology famililar to the RL community. In this section we address and formalize the problem in Causal Inference (CI) terminology. We begin by defining Structural Causal Models (SCM, Pearl (2009a)) – a basic building block of our framework. We then show how the confounded imitation problem can be formalized as an intervention over a specific SCM. Generally speaking, the causal view casts the environment, namely the expert environment generating the offline data, and the online environment, as confounders. Definition 7 (Structural Causal Models). A Structural Causal Model (SCM) is a tuple M = (U, V,F , P (U)) where U is a set of exogenous variables and V is a set of endogenous variables. F is a set of functions such that fi ∈ F are functions mapping a set of endogenous variables Pai ⊆ V \{Vi} and a set of exogenous variables Ui ⊆ U to the domain of Vi, i.e., Vi = fi(Pai, Ui). Finally, P (u) is a probability distribution over the set of exogenous variables U . We assume that the SCM is recursive, i.e., that the causal diagram associated with it is acyclic.
Every SCMM is associated with a causal diagram G, as depicted in Figure 8. Our framework relies largely on the formulation of stochastic interventions, as proposed in Correa & Bareinboim (2020).
We consider stochastic, conditional (non-atomic) interventions, defined by regime indicators σZ (Pearl, 2000; Correa & Bareinboim, 2020), defined formally below.
Definition 8 (Non-Atomic Interventions). Given a SCM M = (U, V,F , P (U)) and a subset Z ⊆ V , an intervention σZ = {σZ1 , . . . , σZn} defines a new SCM MσZ = (U, V,F∗, P (U)) in which the set of functions F is changed to F∗ = {f∗i }i:Vi∈{Zj}nj=1 ⋃ {fi}i:Vi∈V \{Zj}nj=1 .
Non-atomic interventions are a generalization of the classic atomic do(X = x) interventions, defined by the SCMMz and causal diagram GZ in which all edges incoming into Z are removed. We have that
P (y|do(Z = z)) = P (y|z;σZ = do(Z = z)).
Atomic interventions replace function in F by constant functions, whereas non-atomic interventions use general functions. For notational simplicity, when a single intervention is applied to some fi ∈ F , we denote it by σZ = d(fi ← f∗i ), indicating that in the interventional distribution f∗i is used instead of fi. Next we define the identifiability of a causal effect under an intervention, as follows.
Definition 9 (Identifiability). LetX,Y, Z ⊆ V with Y ∩Z = ∅ in some SCM with causal diagram G. Given an intervention σZ = {σZ1 , . . . , σZn}, the causal effect P (y|x, σZ) is said to be identifiable from V ′ ⊆ V if it can be uniquely computed from P (V ′) for every assignment (y, x) in every model that induces G and P (V ′).
Following the model definitions of Section 2, we define the contextual MDP SCM as follows.
Definition 10. A contextual MDP SCM is defined by the causal diagram of Figure 8. For some horizon H > 0, the SCM is defined by the set of endogenous variables V = {si}Hi=0 ∪ {ai} H i=0 ∪ {x} ∪ {ri} H i=0 (denoting the states, actions, context and rewards, respectively), a set of exogenous variables U , and functions F = {fsi , fai , fri , fρe , fν0}, where fsi correspond to the transition function, fai the expert policy, fri the reward function, fρe the context expert distribution, and fν0 the initial context-dependent state distribution.
Relating to our formal definition of our model in Section 2, with slight abuse of notations, the functions {fs, fg, fa, fρe , fν0} adhere to the following relations
P (si+1 = s ′|si = s, ai = a, x) = P (fsi(s, a, x, U)) P (ri = r|si = s, ai = a, x) = δ(fri(s, a, x) = r) π∗(ai = a|si = s, x) = P (fai(a, s, x, U) ρe(x) = P (fρe(x, U)) ν0(s0|x) = P (fν0(s0, x, U)),
where δ(·) indicates the Dirac delta distribution. We are now ready to define the confounded imitation problem. We define the (non-atomic) intervention σx = do(fρe ← fρo) which replaces fρe with fρo in the contextual MDP SCM defined above. The goal of imitation learning is then to identify the quantities
P (ai|si, x, σx = do(fρe ← fρo)) , 0 ≤ i ≤ H − 1, (6)
where, importantly, we assume we only have access to P (si, ai), P (si+1|si, ai, x), P (s0|x), and P (x|σx = do(fρe ← fρo)). Notice that in our setting, P (si+1|si, ai, x), P (s0|x), and P (x|σx = do(fρe ← fρo)) correspond to known quantities of the online environment, whereas P (si, ai) corresponds to the (partially observed) offline expert data. We also emphasize that P (si+1|si, ai, x) is not dependent on the intervention σZ . That is,
P (si+1|si, ai, x, σx = do(fρe ← fρo)) = P (si+1|si, ai, x).
Remark. Our work studies a slightly different version of the identifiability problem in Equation (6), as we only wish to identify an optimal policy from the set Π∗M, as opposed to the single specific policy π∗. This requirement can be formalized by defining an extended SCM which includes all optimal policies in Π∗M, with the assumption that only one is observed (corresponding to the expert data).
Algorithm 4 RL using Expert Data with Unobserved Confounders (Complete Algorithm) 1: input: Expert data with missing context D∗, λ, α,B,N,M > 0, policy optimization algorithm
ALG-RL 2: init: Policy π0, global bonus reward network g∗θ 3: for k = 1, . . . do 4: Generate dataset of rolloutsRk ∼ d πk−1 ρo (s, a) 5: Initialize local networks gmθm ← gθ,m ∈ [M ] 6: for m = 1, . . .M do 7: Sample weight vector wm uniformly from ∆n 8: for e = 1 . . . N do 9: Sample batch uniformly fromRk, i.e., {si, ai}Bi=1
U∼ Rk 10: Sample batch according to weights wm from D∗, i.e., {sei , aei} B i=1
wm∼ D∗ 11: Update gmθm according to
∇θmLm(θm) = 1
B B∑ i=1 ∇θm [ (1− α)f∗(gmθm(s e i , a e i )) + αf ∗(gmθm(si, ai))− g m θm(si, ai) ] 12: end for 13: m∗ ∈ arg minm∈[M ] Lm(θm) 14: Update global parameters from the selected local network g∗θ ← gm ∗
θm∗
15: πk ← ALG-RL(r(s, a, x)− λg∗θ(s, a)) 16: end for 17: end for
E IMPLEMENTATION DETAILS
Our experiments were based off of the recently proposed assistive-gym (Erickson et al., 2020) and recsim (Ie et al., 2019) environments. In this section we discuss further implementation details, hyperparameters, context distributions, and generation of the expert data.
Algorithm Details. A complete description of Algorithm 2 is presented in Algorithm 4. Specific hyperparameters used are shown in Tables 2 and 3. We implemented the algorithm using the RLlib framework (Liang et al., 2018). We used PPO (Schulman et al., 2017) as our policy-optimization algorithm. All neural networks consisted of two-layer fully connected MLPs with 100 parameters in each layer. We used the same rollout buffer (of size 19200 samples) for both our PPO agent as well as our imitation module, which estimated the augmented reward.
Motivated by Kostrikov et al. (2019), we regularized the expert demonstrations with samples from dπ . Particularly, we let α ∈ (0, 1], such that 1 − α corresponds to the probabiilty of sampling an expert example and α corresponds to the probability of sampling from the replay. This leads to minimizing an augmented version of the f -divergence which can be written as
min g:S×A7→R Es,a∼dπρo (s,a|x)[g(s, a)− αf ∗(g(s, a))]− (1− α)Es,a∼dπ∗ρe (s,a)[f ∗(g(s, a))].
Our imitation module consisted of two networks gθ and hθ as proposed in Fu et al. (2018). The “done” signal was also added to the state for training the imitation module. For training CTS we used the Nevergrad optimization platform (Rapin & Teytaud, 2018) with a budget of 10000 and one worker. Here a copied version of the networks gθ and hθ were used to for initialization and then approximate the minimum Df .
For choosing λ we used an adaptive strategy which ensured λ balanced the RL objective with the imitation objective. Specifically, we used the following tradeoff between reward r and bonus g
(1− λadap)r(s, a) + λadapg(s, a), where λadap = rmeanrmean+gmean . Here rmean corresponds to the average reward in the replay buffer and gmean to the average bonus in the replay buffer. By averaging the two, we maintained a similar scale to effectively use the expert data in all the evaluated environments without optimizing for λ.
Context Distribution. For each environment we used a varying context distribution in the expert data, with increasing distance to that of the online environment. The context distribution for the RecSim environment is formally described in Section 5. For the assistive-gym environment the context was defined by the following features: gender, mass, radius, height, patient impairment, and patient preferences. The patient’s mass, radius, and height distributions were dependent on gender. The patient’s impairment was given by either limited movement, weakness, or tremor (with sporadic movement). Finally, the patient’s preferences were affected by the velocity and pressure of touch forces applied by the robot. We used default average values that were provided with the simulator. Particularly, we used the following distributions for each feature
gender ∼ Bern(pmale) mass(gender) ∼ N (µmass(gender), σ2mass) radius(gender) ∼ N (µradius(gender), σ2radius) height(gender) ∼ N (µheight(gender), σ2height) velocity weight ∼ Unif([`vel, uvel]) force nontarget weight ∼ Unif([`target, utarget]) high forces ∼ Unif([`high forces, uhigh forces]) food hit weight ∼ Unif([`hit, uhit]) food velocity weight ∼ Unif([`food vel, ufood vel]) high pressures weight ∼ Unif([`high pressure, uhigh pressure]) impairment ∼ Multinomial(pnone, plimits, pweakness, ptremor).
The values for each distribution are provided in Table 4. For setting covariate shift, we used a a set of distributions that were shifted w.r.t. the default context distribution. We then sampled a shifted distribution w.p. β and the default distribution w.p. 1 − β. That is, when β = 1, the user sampled a context only from the shifted distribution. Table 5 shows an example of one of the shifted distribution that were used.
Expert Data Generation. For the assistive-gym experiments we a dense reward function for generating the expert data and a sparse one for our experiments using the expert data. Specifically, the dense reward function used the environment’s default reward function, defined by
w1 · distance to goal + w2 · action + w3 · task specific reward + w4 · preference score,
where the preferences were weighted according to the context features. Specific weights are provided in the implementation of assistive-gym (Erickson et al., 2020). The sparse reward function did not use the distance to goal (i.e., w1 = 0).

heading: F MISSING PROOFS

heading: F.1 PROOFS FOR SECTION 3
text: We begin by proving two auxilary lemmas.
Lemma 1. Let π2 ∈ Υπ1 . Then, Υπ1 = Υπ2 .
Proof. We show that Υπ1 ⊆ Υπ2 and Υπ2 ⊆ Υπ1 . Let π ∈ Υπ1 , then dπ(s, a) = dπ1(s, a). By our assumption, π2 ∈ Υπ1 , then dπ2(s, a) = dπ1(s, a). Hence, dπ(s, a) = dπ2(s, a). That is, π ∈ Υπ2 . This proves Υπ1 ⊆ Υπ2 . Similarly, let π ∈ Υπ2 , then dπ(s, a) = dπ2(s, a). By our assumption, π2 ∈ Υπ1 , then dπ2(s, a) = dπ1(s, a). Hence, dπ(s, a) = dπ1(s, a). That is, π ∈ Υπ1 . This proves Υπ2 ⊆ Υπ1 , completing the proof.
Lemma 2. Let π0 be a deterministic policy and let M0 = (S,A,X , P, r0, γ) such that r0(s, a, x) = 1{a = π0(s, x)}. Then π0 is the unique, optimal policy inM0.
Proof. By definition of π0 and r0,
r0(s, π0(s, x), x) = 1,∀s ∈ S, x ∈ X .
In particular, Eπ0 [r0(st, at, x)] = 1. Then
V ∗M0 ≤ (1− γ) ∞∑ t=0 γt = Eπ0
[ (1− γ)
∞∑ t=0 γtr0(st, at, x) ] = V π0M0 .
This proves π0 is an optimal policy. To prove uniqueness, assume by contradiction there exists an optimal policy π1 6= π0. Then,
V π1 = Es,a,x∼dπ1 (s,a,x)[1{a = π0(s, x)}] = Es,x∼dπ1 (s,x) [ Ea∼π1(·|s,x)[1{a = π0(s, x)}] ] < 1 = V π0M0 .
In contradiction to π1 is optimal. Then, π0 is a unique optimal policy.
We are now ready to prove Theorem 1.
Theorem 1. [Sufficiency of Υπ∗] Assume ρe ≡ ρo. Let π∗ ∈ Π∗M and let π0 ∈ Υπ∗ . Then, Υπ∗ = Υπ0 . Moreover, if π0 6= π∗, then there exists r0 such that π0 ∈ Π∗M0 but π
∗ /∈ Π∗M0 , where M0 = (S,A,X , P, r0, ρo, ν, γ).
Proof. Let π∗ ∈ Π∗M and let π0 ∈ Υπ∗ . By Lemma 1, as π0 ∈ Υπ∗ , it holds that Υπ∗ = Υπ0 . Next, choosing r0(s, a, x) = 1{a = π0(s, x)}, by Lemma 2 we get that π0 is an optimal policy in M0. This proves π0 ∈ Π∗M0 . Finally, by Lemma 2, Π ∗ M0 = {π0}, proving π
∗ /∈ Π∗M0 if and only if π∗ 6= π0. Proposition 2. Define the mean policy π̄(a|s, x) = ∑ π∈Υπ∗
dπ(s,a,x)∑ π∈Υπ∗ ∑ a′ d π(s,a′,x) , and denote
α∗ = |Π∗M
⋂ Υπ∗ |
|Υπ∗ | ∈ [0, 1]. Then, vM(π̄) ≥ α∗v∗ + (1− α∗) minπ∈Υπ∗ vM(π).
Proof. Let π̃ as defined. Then by linearity of expectation
V π̃M = Es,a,x∼dπ̃ [r(s, a, x)] = 1 |Υπ∗ | ∑
π∈Υπ∗ Es,a,x∼dπ [r(s, a, x)] =
1 |Υπ∗ | ∑
π∈Υπ∗ V πM.
Denote B∗ = Π∗M ∩Υπ∗ , then
V π̃M = 1 |Υπ∗ | ∑ π∈B∗ V πM + 1 |Υπ∗ | ∑ Υπ∗\B∗ V πM
= |B∗| |Υπ∗ | V ∗M + 1 |Υπ∗ | ∑
Υπ∗\B∗ V πM
≥ |B ∗|
|Υπ∗ | V ∗M + |Υπ∗\B∗| |Υπ∗ | min π∈Υπ∗\B∗ V πM
≥ |B ∗|
|Υπ∗ | V ∗M + |Υπ∗\B∗| |Υπ∗ | min π∈Υπ∗ V πM,
completing the proof.
Theorem 2. [Catastrophic Imitation] Assume |X | ≥ |A|, and P (s′|s, a, x) = P (s′|s, a, x′) for all x, x′ ∈ X . Then ∃πe,1, πe,2 s.t. {πe,1, πe,2} are non-identifiable, catastrophic expert policies.
Proof. We first sketch the proof for the special case X = {x0, x1}, A = {a0, a1} and a singleton state space S = {s0}. The general proof follows similarly and is given below. By letting π1, π2 be the determinisic policies which choose opposite actions at opposite contexts, i.e., π1(xi) = ai, π2(xi) = a1−i, we can choose ρe(x) = d∗(π1(x)) and ρ̃e(x) = d∗(π2(x)) which yield
dπ1ρe (a) = 1∑ i=0 ρe(xi)1{a = π1(xi)}
= 2∑ i=1 d∗(π1(xi))1{a = π1(xi)}
= k∑ i=1 d∗(ai)1{ai = a} := d∗(a).
Similarly, dπ2ρ̃e (a) = d ∗(a).
For the second part of the proof choose r1(a, x) = 1{x = xi, a = ai} and r2(a, x) = 1{x = xi, a = a1−i}. Notice that πi is optimal for ri under any distribution of contexts, yet πi achieves zero reward for r1−i.
We now provide a complete proof for the general case.
Let ρo, d∗(a). Without loss of generality, let X = {x0, . . . , xm}, A = {a0, . . . , ak} with m ≥ k, and denote Xk = {x1, . . . , xk} ⊆ X . By definition there exists an injective function fromA into X . Define
f(x) =
{ ai , x = xi, i = 0, . . . , k
a0 , o.w.
g(x) =
{ ai+1 (mod k) , x = xi, i = 0, . . . , k
a0 , o.w.
Then we can select π1, π2, ρe, ρ̃e as follows
π1(a|x) = 1{a = f(x), x ∈ Xk}+ 1
k + 1 1{x /∈ Xk}
π2(a|x) = 1{a = g(x), x ∈ Xk}+ 1
k + 1 1{x /∈ Xk},
and ρe(x) = d
∗(f(x))1{x ∈ Xk}, ρ̃e(x) = d ∗(g(x))1{x ∈ Xk}.
We get that
dπ1ρe (a) = m∑ i=1 ρe(xi)π1(a|xi)
= k∑ i=1 d∗(f(xi))1{a = f(xi)}
= k∑ i=1 d∗(ai)1{ai = a} = d∗(a).
Similarly,
dπ2ρ̃e (a) = k∑ i=1 d∗(g(xi))1{a = g(xi)}
= k∑ i=1 d∗(ai+1 (mod k))1 { ai+1 (mod k) = a } =
k∑ i=1 d∗(ai)1{ai = a} = d∗(a).
This proves the first part of the theorem. For the other parts, choose r1, r2 as follows
r1(a, x) = 1{x = xi, a = ai, 0 ≤ i ≤ k} r2(a, x) = 1 { x = xi, a = ai+1 (mod k), 0 ≤ i ≤ k } .
Then, by definition, for any P (x) such that Supp(P ) ∩ Xk 6= ∅,
Ex∼P (x),a∼π1(·|x)[r1(a, x)] = 1 = max π∈Π Ex∼P (x),a∼π(·|x)[r1(a, x)],
Ex∼P (x),a∼π1(·|x)[r2(a, x)] = 0 = min π∈Π Ex∼P (x),a∼π(·|x)[r2(a, x)].
And similarly,
Ex∼P (x),a∼π2(·|x)[r1(a, x)] = 0 = min π∈Π Ex∼P (x),a∼π(·|x)[r1(a, x)],
Ex∼P (x),a∼π2(·|x)[r2(a, x)] = 1 = max π∈Π Ex∼P (x),a∼π(·|x)[r2(a, x)].
The condition on the support holds for ρe, ρ̃e by definition. If, Supp(ρo) ∩ Xk = ∅, then the result holds trivially as Ex∼ρo(x),a∼π(·|x)[r1(a, x)] = Ex∼ρo(x),a∼π(·|x)[r2(a, x)] = 0 for all π ∈ Π. This completes the proof.
Lemma 3. Assume Supp(ρo) ⊆ Supp(ρe). Then
arg max π Ex∼ρe(x),s,a∼dπ(s,a|x)[r(s, a, x)] ⊆ arg maxπ Ex∼ρo(x),s,a∼dπ(s,a|x)[r(s, a, x)]
Proof. For clarity we denote
Π∗ρe = arg maxπ Ex∼ρe(x),s,a∼dπ(s,a|x)[r(s, a, x)] Π∗ρo = arg maxπ Ex∼ρo(x),s,a∼dπ(s,a|x)[r(s, a, x)] Π∗Supp(ρe) = × x∈Supp(ρe) arg max π Es,a∼dπ(s,a|x)[r(s, a, x)].
To prove the lemma, we will show Π∗ρe = Π ∗ Supp(ρe) ⊆ Π∗ρo .
We begin by proving Π∗ρe = Π ∗ Supp(ρe) . Indeed, let π∗ ∈ Π∗Supp(ρe). Then, for any x ∈ Supp(ρe)
Es,a∼dπ∗ (s,a|x)[r(s, a, x)] = max π Es,a∼dπ(s,a|x)[r(s, a, x)].
In particular, Ex∼ρe(x),s,a∼dπ∗ (s,a|x)[r(s, a, x)] = Ex∼ρe(x) [ max π Es,a∼dπ(s,a|x)[r(s, a, x)] ] ≥ max π Ex∼ρe(x),s,a∼dπ∗ (s,a|x)[r(s, a, x)],
where we used Jensen’s inequality. This proves Π∗Supp(ρe) ⊆ Π ∗ ρe .
To see the other direction, let πe ∈ Π∗ρe and assume by contradiction that πe /∈ Π ∗ Supp(ρe) . Then, there exists x̃ ∈ Supp(ρe) such that
Es,a∼dπe (s,a|x̃)[r(s, a, x̃)] < max π Es,a∼dπ(s,a|x̃)[r(s, a, x̃)].
Define
π̃(·|s, x) = 1{x = x̃}πx̃(·|s, x̃) + 1{x 6= x̃}πe(·|s, x),
where πx̃ ∈ arg maxπ Es,a∼dπ(s,a|x̃)[r(s, a, x̃)]. Then, v(πe) = P (x = x̃)Es,a∼dπe (s,a|x̃)[r(s, a, x̃)] + ∑
x∈Supp(ρe)\{x̃}
P (x)Es,a∼dπe (s,a|x)[r(s, a, x)]
< P (x = x̃)Es,a∼dπ̃(s,a|x̃)[r(s, a, x̃)] + ∑
x∈Supp(ρe)\{x̃}
P (x)Es,a∼dπe (s,a|x)[r(s, a, x)] = v(π̃),
in contradiction to πe ∈ Π∗ρe . This proves Π ∗ ρe ⊆ Π ∗ Supp(ρe) . We have thus shown that Π∗ρe = Π∗Supp(ρe).
Finally, it is left to show that Π∗Supp(ρe) ⊆ Π ∗ ρo . Similar to before, let π ∗ ∈ Π∗Supp(ρe). Then, for any x ∈ Supp(ρe), by Jensen’s inequality
Ex∼ρo(x),s,a∼dπ∗ (s,a|x)[r(s, a, x)] = Ex∼ρo(x) [ max π Es,a∼dπ(s,a|x)[r(s, a, x)] ] ≥ max π Ex∼ρo(x),s,a∼dπ∗ (s,a|x)[r(s, a, x)].
This completes the proof.
Theorem 3. [Sufficiency of Context-Free Reward] Assume Supp(ρo) ⊆ Supp(ρe) and r(s, a, x) = r(s, a, x′) for all x, x′ ∈ X . Then Υπ∗ ⊆ Π∗M.
Proof. Let π0 ∈ Υπ∗ , we will show π0 ∈ Π∗M. Since r(s, a, x) = r(s, a, x′) for all x ∈ X we denote r(s, a) = r(s, a, x). By definition of Υπ∗ we have that.
dπ0ρo (s, a) = d π∗ ρe (s, a)
Then,
v(π0) = Ex∼ρo(x),s,a∼dπ0 (s,a|x)[r(s, a)]
= Ex∼ρo(x)  ∑ s∈S,a∈A dπ0(s, a | x)r(s, a)  =
∑ s∈S,a∈A r(s, a)Ex∼ρo(x)[d π0(s, a | x)]
= Es,a∼dπ0ρo (s,a)[r(s, a)] = Es,a∼dπ∗ρe (s,a)[r(s, a)] = Ex∼ρe(x),s,a∼dπ∗ (s,a|x)[r(s, a)] = max
π Ex∼ρe(x),s,a∼dπ(s,a|x)[r(s, a)]
Then, π0 ∈ arg maxπ Ex∼ρe(x),s,a∼dπ(s,a|x)[r(s, a)]. Applying Lemma 3
π0 ∈ arg max π Ex∼ρo(x),s,a∼dπ(s,a|x)[r(s, a)] = Π ∗ M,
completing the proof.

heading: F.2 PROOFS FOR SECTION 4
text: Proposition 1. [Trajectory Sampling Equivalence] Let ρ∗s which minimizes Problem (P2) for some π ∈ Π, g : S × A 7→ R, and assume Supp(ρo) ⊆ Supp(ρe). Then, there exists pn ∈ ∆n such that dπ ∗
ρ∗s (s, a) = lim n→∞ Ei∼pn
[ (1− γ) ∑∞ t=0 γ t1 { (sit, a i t) = (s, a) }] .
Proof. We can write
dπ(s, a | x) = (1− γ) ∞∑ t=0 γtP (st = s, at = a|x)
= (1− γ) ∑ τ ∞∑ t=0 γtP (st = s, at = a|x, τ)P (τ |x)
= (1− γ) ∑ τ ∞∑ t=0 γt1{τt = (s, a)}P (τ |x).
Then, denoting Pπρ∗s (τ) = Ex∼ρ∗s [P (τ | x) ], we get that
dπρ∗s (s, a) = (1− γ) ∑ τ ∞∑ t=0 γt1{τt = (s, a)}Pπρ∗s (τ)
= Eτ∼Pπ ρ∗s
[ (1− γ)
∞∑ t=0 γt1{τt = (s, a)}
] .
Since, Supp(ρo) ⊆ Supp(ρe), there exists pn ∈ ∆n such that Ei∼pn [ (1− γ) ∑∞ t=0 γ t1 { sit, a i t = (s, a) }] is an unbiased estimator of dπρ∗s (s, a). The result follows by the law of large numbers.
Theorem 5. Let ALG-RL be an approximate best response player that solves the RL problem in iteration k to accuracy k = 1√k . Then, Algorithm 1 will converge to an -optimal solution to
Problem (P2) in O ( 1 4 ) samples.
Proof. We begin by showing that h(P ) = minx∈∆n Df (P ||Ex[Qx]) is convex in P . We can write Df in its variational form, rewriting h(P ) as
h(P ) = min x∈∆n max g:Z7→R
Ez∼P [g(z)]− Ex,z∼Qx [f∗(g(z))],
where
f∗(w) = sup y {yw − f(y)}.
We have that Ez∼P [g(z)]−Ex,z∼Qx [f∗(g(z))] is affine in g and x. Therefore, strong duality holds, yielding
h(P ) = max g:Z7→R min x∈∆n
Ez∼P [g(z)]− Ex,z∼Qx [f∗(g(z))]
= max g:Z7→R
{ Ez∼P [g(z)] + ( max x∈∆n Ex,z∼Qx [f∗(g(z))] )}
We have that maxx∈∆n Ex,z∼Qx [f∗(g(z))] is convex in g as a maximum over convex (affine) functions in a compact set. Therefore h(P ) is also convex as a maximum over convex functions.
Then, the objective in Problem (P2) is convex in dπρo . Following the meta algorithm framework for convex RL in Zahavy et al. (2021), we write the gradient of Df (dπρo(s, a)||d π∗ ρe (s, a)). Notice
that for any general f -divergence Df (xi||yi) = Eyi [ f ( xi yi )] it holds that
∇xjDf (xi||yi) = 0, j 6= i,
and
∇xiDf (xi||yi) = ∇xiEyi [ f ( xi yi )] = Eyi [ 1 yi ∇zf(z) |z= xiyi ] .
Specifically, for the KL-divergence, DKL(pi||qi) = −Eqi [ log ( pi qi )] . Then,
∇piDKL(pi||qi) = Eqi [ 1
pi
] .
Applying Lemma 2 of Zahavy et al. (2021) with a Follow the Leader (FTL) cost player completes the proof.

heading: F.3 PROOFS OF ADDITIONAL RESULTS IN APPENDIX
text: Proposition 3. Assume ρe ≡ ρo and |Υπ∗ | < ∞. Then there exists λ∗ > 0 such that for any λ ∈ (0, λ∗), Algorithm 3 (with δ = 0 sensitivity) will return π̄ of Proposition 2 after exactly |Υπ∗ | iterations.
Proof. Denote
λ∗1 = max π∈Πdet,π 6∈Υπ∗ ,π′∈Υπ∗ dTV
( dπρo(s, a, x), d π′ ρo(s, a, x) ) ,
λ∗2 = min π∈Πdet,π 6∈Υπ∗ dTV (d π ρo(s, a), d
π∗ ρe (s, a)),
where dTV is the total variation distance. Let λ∗ = λ∗2 λ∗1 and λ ∈ (0, λ∗) and notice that λ∗ > 0.
To prove the result., we will show that at iteration n of the algorithm πn ∈ Υπ∗ and that either πn /∈ Υn−1 := {πj}n−1j=1 or Υn−1 = Υπ∗ .
Base case (n = 1). By the variational representation of the f -divergence,
max g0:S×A7→R Es,a∼dπρo (s,a)[g0(s, a)]− Es,a∼dπ∗ρe (s,a)[f ∗(g0(s, a))] = dTV (d π ρo(s, a), d
π∗ ρe (s, a)).
By definition Υπ∗ = arg minπ∈Πdet dTV (d π ρo(s, a)||d
π∗
ρe (s, a)). Then, π1 ∈ Υπ∗ . Finally since Υ0 = ∅, we have that π1 /∈ Υ0. Induction step. Suppose the claim holds for some n = k. We will show it holds for n = k + 1.
We begin by showing that πk+1 ∈ Υπ∗ . Assume by contradiction that πk+1 ∈ Πdet, πk+1 /∈ Υπ∗ . Using the variational form of the f -divergence,
max gi:S×A×X
Li(πk+1; gi) = dTV (d πk+1 ρo (s, a, x), d πi ρo(s, a, x)) ≤ λ ∗ 1,
max g0:S×A
L∗(πk+1; g0) = dTV (d πk+1 ρo (s, a), d
π∗
ρe (s, a)) ≥ λ ∗ 2.
We have that
max g0:S×A7→R, gi:S×A×X 7→R
L∗(πk+1; g0)− λmin i Li(πk+1; gi) ≥ λ∗2 − λλ∗1 > λ∗2 − λ∗λ∗1 = 0.
Next, let π̃k+1 ∈ Υπ∗ , then
max g0:S×A7→R, gi:S×A×X 7→R
L∗(π̃k+1; g0)− λmin i Li(π̃k+1; gi) ≤ 0,
where we used the fact that L∗(π̃k+1; g0) = 0 by definition of Υπ∗ , and Li ≥ 0. We have reached a contradiction to πk+1 being a solution to Equation (4). This proves that πk+1 ∈ Υπ∗ .
Finally, we show that πk+1 /∈ Υk if and only if Υk 6= Υπ∗ . First, notice that if Υk = Υπ∗ then Equation (4) will return πk+1 ∈ Υk by definition of the total variation distance. Next, assume Υk 6= Υπ∗ and assume by contradiction πk+1 ∈ Υk. Then, ∃i : maxgi Li(πk+1; gi) = 0, and maxg0:S×A7→R L ∗(πk+1; g0) = 0, by definition of Υπ∗ . Hence,
max g0:S×A7→R, gi:S×A×X 7→R
L∗(πk+1; g0)− λmin i Li(πk+1; gi) = 0.
In contrast, since Υk 6= Υπ∗ , there exists π̃ ∈ Υπ∗ such that π̃ /∈ Υk, and
max g0:S×A7→R, gi:S×A×X 7→R
L∗(π̃; g0)− λmin i Li(π̃; gi) ≤ λ∗1 < 0,
in contradiction to πk+1 being a solution Equation (4). This completes the proof.
Theorem 6. [Sufficiency of ΥΓ−1π∗ ] Let Assumption 2 hold for some Γ ≥ 1. Then π∗ ∈ ΥΓ−1π∗ .
Proof. Let π ∈ Π. We will show that π ∈ ΥΓ−1π . By elementary algebra, we have that, under Assumption 2,
ρo(x)(1− Γ−1) + Γ−1 ≤ ρo(x)
ρe(x) ≤ ρo(x)(1− Γ) + Γ.
Since Supp(ρe) ⊆ Supp(ρo),
dπρo(s, a) = Ex∼ρo(x)[d π(s, a | x)] = Ex∼ρe(x) [ ρo(x)
ρe(x) dπ(s, a | x) ] ≤ Ex∼ρe(x)[(ρo(x)(1− Γ) + Γ)d π(s, a | x)].
Subtracting dπρe from both sides we get that
dπρo(s, a)− d π ρe(s, a) ≤ Ex∼ρe(x)[(ρo(x)(1− Γ) + Γ− 1)d π(s, a | x)] = (Γ− 1)Ex∼ρe(x)[(1− ρo(x))d
π(s, a | x)] ≤ Γ− 1.
Similarly,
dπρo ≥ Ex∼ρe(x) [( ρo(x)(1− Γ−1) + Γ−1 ) dπ(s, a | x) ] .
Hence,
dπρo(s, a)− d π ρe(s, a) ≥ Ex∼ρe(x) [( ρo(x)(1− Γ−1) + Γ−1 − 1 ) dπ(s, a | x) ] = ( Γ−1 − 1 ) Ex∼ρe(x)[(1− ρo(x))d
π(s, a | x)] ≥ −(1− Γ−1) ≥ −(Γ− 1)
where the last two transitions hold since Γ ≥ 1. Then, we have that∣∣dπρo(s, a)− dπρe(s, a)∣∣ ≤ Γ− 1. This completes the proof.
Theorem 7. [Context Dependent Reward] Let : X 7→ R of Definition 6. Denote oe = Ex∼ρo(x)[ (x)] + Ex∼ρe(x)[ (x)]. Then for any π∗ ∈ Π∗M, π0 ∈ Υπ∗
v(π0) ≥ v(π∗)− oe.
Proof. Let π∗ ∈ Π∗M, π0 ∈ Υπ∗ . The proof follows similar steps to that of Theorem 3.
v(π0) = Ex∼ρo(x),s,a∼dπ0 (s,a|x)[r(s, a, x)] = Ex∼ρo(x),s,a∼dπ0 (s,a|x)[r(s, a, x)− r0(s, a)] + Ex∼ρo(x),s,a∼dπ0 (s,a|x)[r0(s, a)]
= Ex∼ρo(x),s,a∼dπ0 (s,a|x)[r(s, a, x)− r0(s, a)] + Ex∼ρo(x)  ∑ s∈S,a∈A dπ0(s, a | x)r0(s, a)  = Ex∼ρo(x),s,a∼dπ0 (s,a|x)[r(s, a, x)− r0(s, a)] +
∑ s∈S,a∈A r0(s, a)Ex∼ρo(x)[d π0(s, a | x)]
= Ex∼ρo(x),s,a∼dπ0 (s,a|x)[r(s, a, x)− r0(s, a)] + Es,a∼dπ0ρo (s,a)[r0(s, a)] = Ex∼ρo(x),s,a∼dπ0 (s,a|x)[r(s, a, x)− r0(s, a)] + Es,a∼dπ∗ρe (s,a)[r0(s, a)] = Ex∼ρo(x),s,a∼dπ0 (s,a|x)[r(s, a, x)− r0(s, a)] + Ex∼ρe(x),s,a∼dπ∗ (s,a|x)[r0(s, a)].
Then,
v(π∗)− v(π0) = Ex∼ρe(x),s,a∼dπ∗ (s,a|x)[r(s, a, x)− r0(s, a)]− Ex∼ρo(x),s,a∼dπ0 (s,a|x)[r(s, a, x)− r0(s, a)] ≤ Ex∼ρe(x)[ (x)] + Ex∼ρo(x)[ (x)],
completing the proof.

Review summaries:
The authors considered the problem of imitation and reinforcement learning in the setting of contextual MDP with latent confounders. They defined an ambiguity set, i.e., the set of all deterministic policies that match the marginalized stationary distributions of a given policy. In imitation learning, in the case of no covariate shift, no policy in the ambiguity set of optimal policy can be ruled out. Moreover, they showed that acting uniformly with respect to this set results in a policy that is better than the worst policy in the set. In the case of having covariate shift, the authors showed that imitating the policy of the expert might result in a catastrophic policy if the transition probabilities are context-free. However, if the reward function is independent of context, then the problem of imitation learning becomes feasible. Moreover, the case of bounded confounding is studied in the appendix. In particular, it is shown that under some conditions, the optimal policy is in a \delta-ambiguity set. Finally, in the problem of RL, the proposed algorithm converges to the optimal policy using corrective trajectory sampling.

The paper studies how to incoporate expert data with covariate shift, defined in a contextual MDP where expert data comes from a different context distribution and where the expert contexts are unobserved. The paper provides limited theoretical results for pure imitation based on state-action marginals only, in particular showing that matching state-aciton marginals alone can lead to arbitrarily bad policies with covariate shift in the contexts. They also propose a hybrid RL-imitation algorithm that utilizes the expert data without allowing it to bias the final solution.

This paper focuses on imitation learning with confounded expert data (with no contextual information) that might have covariate shifts. Theoretical results (upper and lower bounds) are obtained, supported by numerical experiments and real data applications. The main conclusions are: (1) with no access to reward in the online setting, it is impossible to learn the optimal policy with confounded expert data with covariate shifts; (2) but with access to reward, then it is possible to learn the optimal policy consistently.