On Covariate Shift of Latent Confounders in Imitation and Reinforcement Learning
Authors: Guy Tennenholtz
Assaf Hallak
Gal Dalal
Shie Mannor
Gal Chechik
Uri Shalit
Venue: ICLR 2022 Poster (2022)
Keywords: imitation learning
reinforcement learning
expert data
hidden confounding
causal inference
covariate shift

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. In addition, we discuss the problem of distribution shift between the expert data and the online environment when partial observability is present in the data. 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.

Published as a conference paper at ICLR 2022
ON COVARIATE SHIFT OF LATENT CONFOUNDERS IN IMITATION AND REINFORCEMENT LEARNING
Guy Tennenholtz
Assaf Hallak Nvidia Research
Gal Dalal
Nvidia Research Technion
Nvidia Research
Shie Mannor
Gal Chechik
Uri Shalit
Nvidia Research
Nvidia Research
Technion
Technion
ABSTRACT
We consider the problem of using expert data with unobserved confounders for imitation and reinforcement learning. We begin by defining the problem of learn- ing from confounded expert data in a contextual MDP setup. We analyze the limitations of learning from such data with and without external reward, and pro- pose 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 on- line 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 sampling procedure that addresses the unknown shift and prove convergence optimal solution. Finally, we validate our claims empirically on challenging assistive healthcare and recommender system simulation tasks.
INTRODUCTION
Reinforcement Learning (RL) is increasingly used across many fields to create agents that learn via interaction and reward feedback. Imitation Learning (II., Hussein et al. (2017) concerned with learning via expert demonstrations without access reward function. Similarly, RL settings often use expert data to boost performance, climinating 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 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 and Theorem
In this paper we define the tasks of imitation and reinforcement learning using expert data with un- observed confounders and possible covariate shift. We focus on contextual MDP setting, where a context sampled at every episode from some distribution, affecting both the reward and the transi- tion 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-leaming problem, (i.e., without access to reward) in Section 3 Under no covariate shift the unobserved context, we characterize sufficient and necessary set of optimal policies. In contrast, depends on the context, then the initation-learning problem is -ennidentifiable and prone to catas- trophic 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 deepicts a possible failure case of using confounded expert data with unknown covariate shift in dressing task. Unlike the imitation setting,
prove that in the presence of covariate shift, if the true reward

---

Published as a conference paper at ICLR 2022
Expert Context Distribution (unobserved and unknown)
Online Context Distribution (observed) 05% female,, 45% male Average ago 77 tremor impaiment
Distribttion Mamatch
40% female, 60% male
Average age 86 10%
Ilustration Trembiing Potlent
Falure Covariate
Sampling Correction
Fallure Arm misses slaave
Buocecoful: Arm through skoevb
Failure of using confounded expert data under context distribution mismatch between online envi- ronment and expert data. Caregiver does learn to perform well in dressing task when covariate shift of hidden confounders is present but not accounted for.
Figure 1:
we show that optimality can still be achieved in the RL setting while using confounded expert data with arbitrary covariate shift. We use corrective data sampling procedure and prove convergence to an optimal policy.
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) the RL setting, under arbitrary covariate shift, provide novel algorithm with convergence guarantees which uses corrective sampling technique to account for the unknown context distribution in the expert data. (3) Finally, conduct extensive experiments
Our contributions
recommender system (le et al.) (019) and assistive-healthcare (Erickson et al., 2020) environments, demonstrating our theoretical results, and suggesting that confounded expert data can be used in 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 = USS A, P.r. Po: M, 7). where is state space, X' is the context space, A the ac- tion space, P S x S x A x X + j0, 1] is the context dependent transition kernel, r x A x + [0, the context dependent reward function, and (0, 1) is the discount factor. We assume an initial distribution over contexts Po : X' + [0, [0,1] and an initial state distribution S X' [0. 1).
The environment initializes at some context x ~ Po(-), and state so ~ v(-1z). At time the environment is at state s, S and agent selects an action a, E A. The agent receives a reward ri r(8e, .) and the environment then transitions to state 8141 P(- x).
91+1
We define Markovian stationary policy as mapping v T:Sx X' x A + [0, 1], such that ((-1s, x)) is the action sampling probability. We define the value of policy by (M( (1
Figure 2: Causal Diagram Con- textual MDP
- 22 ) |xr ~ Po. ~ v(- x)) where E, denotes the expectation induced by the policy æ. We denote by the set policies and Ia the set of deterministic Markovian policies. We define the optimal value and policy by V'M max, (M))). and क MM E arg max (T). respectively. Whenever appropriate, simplify notation and write v n² We use II' to denote the set of optimal policies in M. i.c.,
Markovian
II' arg ma (IT).. We also define the catastrophic policies II the set
II' = arg min U. ( TEll
(1)
2

---

Published as a conference paper at ICLR 2022
We later use this set to show impossibility of imitation under arbitrary covariate shift and context- independent transition function.
con- Expert Data with Unobserved Confounders. We assumee additional access founded dataset consisting of expert trajectories D {(s). a/. si .i..... S'y ) where a}
nh" E M The trajectories in the dataset were sampled i.i.d. from the marginalized expert distribution (under possible context covariate shift) Pa (so, ag, $1, C1.............................. (H) = En (e(x) I P(s++1|$;, z)) 2),, where Pe some distribution over contexts. Importantly, Pe. does not necessarily equal Po the distribution of contexts in the online environment. Notice that it is assumed that x² that generated the data had access to the context r' (i.e., π² is context-dependent though it is missing in the data.
In this work, consider two settings:
1. Confounded Imitation Learning (Section The agent has access confounded ex- pert data (with context distribution pe) as well as real environment (S, X, A, P, Po: 7), without access to reward.
2. Reinforcement Learning with Confounded Expert Data (Section The agent has ac- cess confounded expert data (with context distribution Pa) as well as real environment M = (S, X, A, P. Pos 7), with access 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 Pe in the offline data.
Marginalized State Action Distribution. We denote the state-action frequency of policy a e II given context x E X' by d" (s,a(z) (1 (1-r) Ei= 7² pM pr denotes the probability measure induced by n. Similarly, given distribution over contexts, we define the marginalized state-action frequency of a policy æ under the corresponding context distribution by
(Pr( (st = 8, az = a|2r, 80 ~ v(-(a=), where
d= (s, a) = Emapp (d æ)]
(online environment),
dh (s,a) a) = Exapo adTh (s, 2:) (offline expert data).
A Causal Perspective. Our work sits at intersection between the fields of RL and Causal In- ference (CI). We believe essential to bridge the gap between these two fields, and include an interpretation of our model using CI terminology in Appendix D) where equivalently define objective as an intervention over the unknown distribution Pe in a specified Structural Causal Model, as depicted in Figure 2
IMITATION LEARNING WITH UNOBSERVED CONFOUNDERS
In this section, analyze the problem of confounded imitation learning, namely, learning from pert 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 interact with the environment (Ho & Ermon 2016 Fu et al., 2017 Kostrikov et al. 22019 Brantley et al.) 2019). In the first part of this section we assume covariate shift between the online environment and the data is present, i.e., Pe Pa. We lift this assumption in the second part, where we focus on the covariate shift of the hidden confounders.
3.1 No COVARIATE SHIFT: Po = Pe
We first consider the scenario in which covariate shift is present between the offline data and the online environment, i.c., Po = Pc- We begin by defining the marginalized ambiguity set, a central component of our work.
Definition 1. For E HI, we define the set of deterministic policies that match the marginalized stale-action frequency of π by Tx 4 Mae d (s, a) = d7 (8,a) V&ES, A}.
3

---

Recall that, in general, x² E. 11 may depend on the context z E X. Therefore, the set T corresponds to all deterministic policies that cannot be distinguished from π based on the observed expert data. The following theorem shows that for any policy π² E I and any policy To E T one could design a reward function ro. for which To is optimal, while the T. is indiscernible (Fadxxx/foo for proof). In other words, T,- is the smallest set of Published as a conference paper at ICLR 2022
Impossible
Hard
Easy
Context- Indeeedendd
Context- lonexxt-Indeeendent Reward
Bounded Confounding (Appendix C
Transition
(Theorem 2)
(Theorem
Figure 3: A spectrum for the difficulty confounded imitation with covariate shift. Context-Indepenee transitions may result in non-identifiabl and catastrophic candidates, whereas context-independen rewards reduce the problem standard imitation learning.
from Y .o i.e., T_- Yz, (see Appendix candidate optimal policies.
Theorem 1. /Sufficiency of T Assume Pe = Po Let x¹ E I1'M and let to Yar. Then, T, = TYoo Moreover, if no + n,, then there exists ro such that Tu E [I'M but T & II where Mo = AS, To, Po: v. v.7).
The above theorem shows that any policy in T, is candidate optimal policy, yet without knowing the context the expert used. Such ambiguity can result in selection of suboptimal or even catas- trophic policy. We provide practical algorithm in Appendix B which calculates the ambiguity set Y,", and returns an average policy, with computational guarantees. In the next subsection we ana- lyze more challenging scenario, for which Po # Pe. In this case, Y,- may not be sufficient for the imitation problem.
3.2 WITH COVARIATE SHIFT: Po +
Next, we assume covariate shift exists between the online environment and the expert data, i.c., Po F Pe- Particularly, without further assumptions on the extent of covariate shift, we show two extremes of the problem. Theorem / 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 prove that, whenever the reward independent of the context, the imitation problem can efficiently solved.
Clearly, if Supp(p.) Ø Supp(pe then there exists E Supp(p.) for which n² is not identifiable from the expert data We therefore assume throughout that Supp(p.) S Supp(pe). We begin by defining the set of non-identifiable policies those that cannot be distinguished from their respective state-action frequencies without information on Pe-
Definition 2. We say that { { are non-identifiable policies if there exist {pi) such that (s, a) = d33 (s, a) for all i $ $..
Next, focusing on catastrophic policies (recall Equation (I). we define catastrophic expert policies as those which could be either optimal or catastrophic under po for different reward functions.
Definition 3. We say that J are catastrophic expert policies there exist (ri) such that for all ग. E Mi and 3j E [k]. # such that π. E EII where M, = Ar, u.)
Using the fact that both Pe and r are unknown, the following theorem shows that whenever P(s' [s, a, is independent of v. could find two policies which non-identifiable, catastrophic expert policies (see Appendix for proof). In other words, the case of context-independent tran- sitions, without further information on Pe or r the expert data useless for imitation. Furthermore, attempting to imitate the policy using the expert data could result in a catastrophic policy.
For a distribution P we denote by Supp(P) the support of P.
We define con-identifubiility in Definition 2
We similar notion of identifiability in Pearl 2009b
4

---

Published as a conference paper at ICLR 2022
Theorem 2. /Catastrophic Imitation] Assume [2] 2 |A), and P(s'|s, a) P(s'ls, a.x'), for all I, E X. Then 3Te,1 Te,2 s.t. {Te,, xe.2) (e,2) are non-identifiable, catastrophic expert policies.
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 Y an optimal policy.
Theorem 3. /Sufficiency of Context- Reward]
Assume Supp(po) C Supp(pe) and
r(s, x) = r(,,, a, x') for all æ, x' E X. Then Y,- - 1
Theorems andd 3 suggest that the hardness of the imitation problem under covariate shift lies on a wide spectrum (as depicted in Figure While dependence of the transition P(s"|s, a, on pro- vides us with information to identify x in the expert data, the dependence of the reward r(s, a, x) on increases the degree of confounding in the imitation problem. Both of these results concerned with arbitrary confounding.
Bounded Confounding: A Sensitivity Perspective. 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, 202T). In our setting, this confounding bias occurs to 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 (Bounded Sensitivity).
We assume that Supp(pe) S Supp(po) and that there exists -o(=)(1-P((2)) ST.
some r 2 1 such that for all x E Supp(pe), r-1 Pa(x [1)](
Next, we define the notion of 5-ambiguity, generalization of the ambiguity set in Definition D Definition 4 (ō-Ambiguity Set). For policy II, define the set of all deterministic policies that are 5-close to by rę = n &T d (s, a) d,, (s, a) < 5, s E S.S.a E A].
Similar to Definition ! the 6-ambiguity set considers all deterministic policies with a marginalized state-action frequency of distance at most 8 from и. The following results shows that r is sufficient set of candidate optimal policies, as long as Assumption 2 holld for some [ 2>1 1.
Theorem 4. /Sufficiency of - Let Assumption 2 hold for some r 1. Then
rf-
For the interested reader, we further analyze the case of bounded confounding in Appendix C We also demonstrate the effect of bounded confounding 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.
Section 5 In the following section, we
4 USING EXPERT DATA WITH UNOBSERVED CONFOUNDERS FOR RL
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, assume access to the true reward signal in the online environment and wish to leverage the offline expert data aid the agent in converging 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, TV. distance, x² -ddivergence) of state- action frequencies in Tx.,
max E Egr (mole)) a, a) (D (s, (s,a)lld"" (s, (s,
(P1)
TEIl
5

---

Published as a conference paper at ICLR 2022
Algorithm RL using Expert Data with Unob- served Confounders (Follow the Leader)
Algorithm 2 RL using Expert Data with Unob- served Confounders (Online Gradient Descent)
1: input: Expert data with missing context D",
input: Expert data with missing context, A, B,N N > 0. policy optimization alg. ALG-RL
À >> O, policy optimization algorithm ALG-RI.
2: init: Policy x²
2: init: Policy x² bonus reward network
3: for k 1,... do
3: for l: 1....d
4: Pa arg min, DKL (dF= (s.a)lld² (s, a))
4:
< arg min, (s, (s, 5: eore=1.... ...NNo do
5: g" - g s,amarh-1
6: Sample batch {si, ai ~ d (s,
d7 (mra)
6: n² ALG A1G--R((((s,,,)) -o (s,
7:
Sample batch fst, ai ~ dz (8.a)
8:
Update 90 according to VoL(O)
7: end for
B DB (go
90 (st,
9: end for
10: x <- ALG ALG-RI(r(s,,,2) as Agn(s,
11:
end for
Here, À 0 and D is the f-divergence, where is convex function solution to Problem (PI) is an optimal policy x¹ E Ii2 as long as Po = Pe. Rewriting D, using its variational form (see Appendix for background on the variational form of f-divergences), we get the following equivalent optimization problem, motivated by Nachum et al. (2019):
(0, + R. The
max min Papppass [r(s, a, x) + Ag(s, a)] AB (s.a) 1 (1(), (PIb) wEll x
where f" is the convex conjugate of f. i.e., f (v) sup, ay ((y).
Unfortunately, when covariate shift exists (i.e., Po # Pe); Problems [P] and Pib) are not ensured to converge to optimal policy (Theorem 2). Instead, propose to reformulate Problem (PIb)
using a distribution Ps which minimizes the 7-divergence, as follows,
max min (arnoassoodddtttppppsss.s. + Ag(s, a)] AE ABxandy (s.a) [f (g(s, a))
(P2)
rEll AARR P.EB(X))
Here, ((X)) denotes the set of probability measures on the Borel sets of X, and
d (s, a) Ezup, [d= (s. | =)]. Indeed, whenever Supp(p.) C Supp(pa), we have that Ps
(T, P.) (T"", Po) is a solution Problem (P2). That is, unlike Problems (PI) and (PIb). Prob-
lem (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, " (s, a). Fortunately, d,," (s, a) can be equivalently written as an expectation over trajectories in DD, rather than expectation over unobserved contexts, as shown by the following proposition (see Appendix for proof):
Proposition 1. (Trajectory Sampling Equivalence) Let P which minimizes Problem P2 for some E 11, S x A + R, and assume Supp(pp)) C Supp(pe)- Then, there exists p" E such that (s. lim Eingo ((1-)) Er= 7 7'1 (s) a;) = (s
100000
Proposition I allows us to estimate the inner minimization problem over Ps in Problem P2 using trajectory samples. Particularly, we uniformly sample & distributions p....-pi -ph where p" An. and then estimate
{D/( do (s,a) min JE(1....k)
(1 (s, (1.00) ((.6))))((0))))) 1)}
min Dy D1(d,, ~M
Eappy
E
Σ 1=0
(2)
which can estimated by using the variational form of Dy (see Appendix Aj. Corrective Trajectory Sampling (CTS), as it uses complete trajectory samples to account for the unknown context distribution Pe.
We call this procedure
Solving Problem P2
Algorithm I provides iterative procedure for solving the optimization uses alternative updates of a cost player (line 5) and policy player (line
problem in Problem 12 6). In line 5 the gradient of DEL. w.r.t. d= is taken using Follow the Leader (FTL.) cost player

---

Published as a conference paper at ICLR 2022
FeedingSawyer
DrwssingSawyer
Expert
Export
reward
reward
R2
-150
-200
-200
-190
steps
$1ep
SainkrngSawyer
BedBathingPR2
133
Expert
Expert
roword
RL
3178
RL
With CTS (ours, -
With (ours, 0.3)
232
No CTS ($ 0.81
3300006 3800000
300000
430000
5000
3000006 3300000
3000000
4000000
5000000
step
step
Figure 4: Plots compare training curves using CTS vs. normal sampling expert and large (B = 0.8) covariate shift bias in four assistve-healthhare tasks. Dashed black lines show expert and RL (without data) scores. Runs were averaged over seeds. Legend is shared across all plots.
small (8 0.3)
to estimate the next bonus iterate. Finally, in line 6, an efficient, approximate policy optimization algorithm ALG-RL executed using an augmented reward. The following theorem, provides con- vergence guarantees for Algorithm (see Appendix F for proof based on Zahavy et al. (202T).
Theorem 5. Let ALG-RL be an approximate best response player that solves the RL problem in iteration to accuracy 6g J.
vi Then, Algorithm 7 will converge to an e-optimal solution to
Problem P2) in (²)) samples.
Notice that, while Theorem 5 shows Algorithm converges optimal policy, does not deter- mine 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 show that such data can indeed improve overall performance on various tasks.
drawback of Algorithm D is that it needs estimate the state-action frequencies. A practical implementation of Algorithm | using online gradient descent (OGD) is provided in Algorithm 2 the algorithm does not require approximate estimates the state-action frequencies. but rather, only the ability sample from them. Similar to Algorithm [ we use CTS (see Equation 2 to estimate Ps in line according to some f-divergence. Here, samples are drawn from the current policy as well as samples from D* (with CTS). We write D, in its variational form, and use neural network representation for ga. We then use the aforementioned samples to minimize the f-divergence using OGD. Finally, the policy is updated using ALG-RL and augmented reward.
5 EXPERIMENTS
We tested our proposed approach for using expert data with hidden confounding in recommender- system and assistive-healthert
environments. For all our experiments used x -diveergee as 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. (017) implemented in RLlib (Liang et al.) 2018) for both the imitation as well as RL settings. We include specific choice hyperparameters and an exhaustive overview of further implementation details in Appendix

---

Published as a conference paper at ICLR 2022
Divergence Comparison
RecSim (Confounded Imitation)
RecSim (RL with Confounded Data)
Du bugion fobi
1 vergonce
HOUTH, 21
NL ML-diorgaece
10. 10
180000 200000
140000
1006000 720/800
Figure 5: Left plot shows comparison of different choices of f f-divergences for pure imitation (without reward and without covariate shift) on the BedBathingPR2 environment. Middle plot depicts execution of imitation with hidden confounding (without reward) for different levels of covariate shift. Right plot compares our CTS correction the RecSim environment with strong covariate shift bias. All runs were averaged over 5 seeds.
Assistive Healthcare. Consider the challenge of providing physical assistance to disabled persons. simulating autonomous robots as versatile
A recently proposed set of tasks for 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 Po we used the default values provided by the original environment. To enforce a distribution shift in the expert data, we shifted each distri- bution randomly with an additive factor B dz. where 3 E [0, 1]. and de was a random distribution the covariate shift strength. The expert data was generated by a fixed policy trained using a dense reward function. sparse reward signal was used for executing our experiments with the confounded expert data, For further details, we refer the reader to Appendix E
chosen from set of shifting distributions (see Appendix [E). Here 3 corresponds
Figure I depicts results for executing Algorithm 2 on four assistive-gyn 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 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) al., 2019). Leveraging expert data collected using, e.g., surveys to users, may greatly benefit fu- ture solutions. Because features are repeatedly added to these systems, full information in the data rarely available. Here, we use the recently proposed RecSim Interest Evolution environment (Ie) et al. 2019) simulating sequential user interaction with a slate- recommender system. The environment consists of 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 session length or visit frequency).
We used slate of 10 documents and a user context of dimension 20. To test the severity of the implications of Theorem i in the confounded imitation setting, we used a user-model sampled from a distrbiitribttion. Specifically, for the expert data the user context features I = (Zo: 3.3 were sampled from Beta-distribution where 2; Beta 4), and a; 1.5 + 19 i. In contrast, the online environment features were sampled from a shifted Beta-distribution with a; ===(1 (1- B) 8.5 + ((10- i)) where B [0, 1] defined the shift strength. While the
219)
8

---

Published as a conference paper at ICLR 2022
original environment used uniform distribution to generate user contexts, the Beta- Beta-distribution let us analyze severe forms of covariate shifts, testing the limits of our results in Sections 3 and 4
Figure | (a) depicts the effect of increased covariate shift 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 (3 0) optimal score is achieved, and as B increases, perfor- mance decreases. Particularly, as the mirrored distribution is reached (33 = 1). catastrophic policy reached. While the imitator "believes" have reached an optimal policy, it has in fact reached catastrophic one, as shown by the orange plot. Conversely, Figure 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.
6 RELATED WORK
Imitation I Learning. The imitation learning problem has been extensively studied in both the fully
offiine (Pomerleau 1989 Bratko al.] 1995) well online setting (Ho & Ermon 2016 Fu ct al. 2018 Kim & Park 2018 Brantley et al. 2019). Specific to our work are GAIL (Ho & Ermon, 2016 AIRL (Fu al. 2017). and DICE (Kostrikov 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.] [02T] Tennenholtz et al., [02 la Fujimoto & Gu 202T). In the online RI. setting, the combination of offline data to improve RL efficiency has shown great success (Nair et al., (020). KI.-regularized techniques (Peng et al., 2019 Siegel al.)
2019) as well as DICE-based algorithms (Nachum et al. 2019h 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 am 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 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 misidentification" 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 (non- expert) confounded data (Zhang & Barcinboim 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 leveraging expert data with hidden confounders and possible covariate shift in both the imitation and the RL settings.
7 CONCLUSION
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, 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).

---

Published as a conference paper at ICLR 2022
REFERENCES
Kianté Brantley, Wen Sun, and Mikael Henaff. Dissgreemeent-rggulliiied imitation learning. In International Conference on Learning Representations, 2019.
Ivan Bratko, Tanja Urbančič, and Claude Sammut. Behavioural cloning: phenomena, results and problems. IFAC Proceedings Volumes, 28(21):143-144 1995.
Juan Correa and Elias Bareinboim. A calculus for stochastic interventions: Causal effect identifica-
tion and surrogate experiments. Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pp. 1009000000000 10100, 2020.
Imre Csiszár and Paul C Shields. Information theory and statisties: A tutorial. 2004.
Pim de Haan, Dinesh Jayaraman, and Sergey Levine. Causal confusion in imitation learning. Ad- vances in Neural Information Processing Systems, 32: 1698-11709, 2019.
Zackory Erickson, Vamsee Gangaram, Ariel Kapusta, C Karen Liu, and Charles C Kemp. Assis- tive gym: A physics simulation framework for assistive robotics. In 2020 IEEE International Conference on Robotics and Automation (ICRA), pp. 10169-10176. IEEE, 2020.
Justin Fu, Katie Luo, and Sergey Levine. Learning robust rewards with adversarial inverse rein- forcement learning. arXiv preprint arXiv:1710.11248, 2017.
Justin Fu, Katie Luo, and Sergey Levine. Learning robust rewards with adverserial inverse rein-
forcement learning. In International Conference Learning Representations, 2018.
Scott Fujimoto and Shixiang Shane Gu. minimalist approach offline reinforcement learning.
arXiv preprint arXiv:2106.06860 2021.
Omer Gottesman, Fredrik Johansson, Matthicu Komorowski, Aldo Faisal, David Sontag, Finale Boshi-Velez, and Leo Anthony Celi. Guidelines for reinforcement learning healthcare. Nature medicine, 25(1):16-188 2019.
Assaf Hallak, Dotan Di Castro, and Shie Mannor. Contextual markov decision processes. arXiv preprint arXiv:1502.0223: 2015.
Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning. Advances in neural information processing systems, 29:4565-4573, 2016.
Jesse Y Hsu and Dylan S Small. Calibrating sensitivity analyses to observed covariates in observa- tional studies. Biometrics, 69(4):803-811, 2013.
Ahmed Hussein, Mohamed Medhat Gaber, Eyad Elyan, and Chrisina Jayne. Imitation learning: A survey of learning methods. ACM Computing Surveys (CSUR), 50(2):1-35, 2017.
Eugene le, Chih-wei Hsu, Martin Mladenov, Vihan Jain, Sanmit Narvekar, Jing Wang, Rui Wu, and Craig Boutilier. Recsim: A configurable simulation platform for recommender systems. arXiv preprint arXiv:1909.0458 2019.
Nathan Kallus and Angela Zhou. Confounding-robus policy evaluation in anfinite-horizon rein- forcement learning. arXiv preprint arXiv:2002.0451/. 2020.
Nathan Kallus and Angela Zhou. Minimax-optimal policy learning under unobserved confounding. Management Science, 67(5):2870-2890, 2021.
Liyiming Ke, Sanjiban Choudhury, Matt Barnes, Wen Sun, Gilwoo Lee, and Siddhartha Srinivasa. Imitation learning as f-divergence minimization. In International Workshop on the Algorithmic Foundations of Robotics, pp. 313-329, Springer, 2020.
Kee-Eung Kim and Hyun Soo Park. Imitation learning via kernel mean embedding, In Thirty-Second AAAI Conference on Artificial Intelligence, 2018.
Ilya Kostrikov, Ofir Nachum, and Jonathan Tompson. Imitation learning via off off-policy distribution matching. arXiv preprint arXiv: 19/v 2019.
10

---

Published as a conference paper at ICLR 2022
Iya Kostrikov, Rob Fergus, Jonathan Tompson, and Ofir Nachum. Offline reinforcement learning with fisher divergence critic regularization. In International Conference on Machine Learning, pp. 5774--788.. PMLR, 2021.
Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. arXiv preprint arXiv:2006. 2007704777 2020.
Finnian Lattimore, Tor Lattimore, and Mark D Reid. Causal bandits: learning good interventions via causal inference. Proceedings of the 30th International Conference on Neural Information Processing Systems, pp. 1189-1197, 2016.
Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tuto- rial, review, and perspectives on open problems. arXiv preprint arXiv:2005.01665) 2020.
Eric Liang, Richard Liaw, Robert Nishihara, Philipp Moritz, Roy Fox, Ken Goldberg, Joseph Gon- zalez, Michael Jordan, and lon Stoica. Rilib: Abstractions for distributed reinforcement learning. In International Conference on Machine Learning, pp. 3053-3062. PMLR, 2018.
Friedrich Liese and Igor Vajda. On divergences and informations in statistics and information theory. IEEE Transactions on Information Theory, 52(10))433444412, 2006.
Ofir Nachum, Bo Dai, lya Kostrikov, Yinlam Chow, Lihong Li, and Dale Schuurmans. Algaedice: Policy gradient from arbitrary experience. arXiv preprint (ttxii:/9/2.00000 2019.
Ashvin Nair, Murtaza Dalal, Abhishek Gupta, and Sergey Levine. Accelerating online reinforcement learning with offline datasets. arXiv preprint arXiv:2006.09333 09359, 2020.
Hongseok Namkoong, Ramtin Keramati, Steve Yadlowsky, and Emma Brunskill. Off-policy policy evaluation for sequential decisions under unobserved confounding. arXiv preprint arXiv:2003.0500 056233. 2020.
Michael Oberst and David Sontag. Counterfactual off-policy evaluation with gumbel-max structural causal models. In International Conference on Machine Learning, pp. 4881 4890. PMLR, 2019.
Judea Pearl. Causality: models, reasoning and inference, volume 29, Cambridge University Press, 2000.
Judea Pearl. Causality: Models, Reasoning and Inference. Cambridge University Press, New York, NY, USA, 2nd edition, 2009a. ISBN 052189560X, 9780521895606.
Judea Pearl. Causal inference in statistics: An overview. Statistics surveys, 3:96-146, 2009b.
Xue Bin Peng, Aviral Kumar, Grace Zhang, and Sergey Levine. weeinte regression: Simple and scalable off- off-policy reinforcement learning. arXiv preprint arXiv: 17000:1910.00177, 2019.
Dean A Pomerleau. Alvinn: An autonomous land vehicle in a neural network. Technical report, CARNEGIE-MELLON UNIV PITTSBURGH PA ARTIFICIAL INTELLIGENCE AND PSY- CHOLOGY 1989.
Martin L Puterman. Markov decision processes: discrete stochastic dynamic programming. John Wiley & Sons, 2014.
Rapin and O. Teytaud. Nevergrad A gradient- freee optimization platform. httts:://ittHu 2018.
com/ Foom/Face aaa.eeooooeeeeaareeeeeeeeeeeeeee
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint 1707.06377 2017.
Noah Siegel, Jost Tobias Springenberg, Felix Berkenkamp, Abbas Abdolmaleki, Michael Neunert, Thomas Lampe. Roland Hafner, Nicolas Heess, and Martin Riedmiller. Keep doing what worked: Behavior modelling priors for offline reinforcement learning. International Conference on Learning Representations, 2019.
11

---

Published as a conference paper at ICLR 2022
Guy Tennenholtz, Uri Shalit, and Shie Mannor. Off-policy evaluation in partially observable vironments. Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pp. 10276-10283, 2020.
Guy Tennenholtz, Nir Baram, and Shie Mannor. Gelato: Geometrically enriched latent model for offline reinforcement learning. arXiv preprini 1.Xi::2102.1132,, 2021a.
Guy Tennenholtz, Uri Shalit, Shie Mannor, and Yonathan Efroni. Bandits with partially observable confounded data. In Conference on Uncertainty Artificial Intelligence. PMLR, 2021b.
Lingxiao Wang, Zhuoran Yang, and Zhaoran Wang. Provably efficient causal reinforcement learning with confounded observational data. arXiv preprint arXiv:2006. /2311 2020.
Tom Zahavy, Brendan O' O'Donguue, Guillaume Desjardins, and Satinder Singh. Reward is enough for convex mdps. arXiv preprint arXiv:2106.00660 00661,, 2021.
Junzhe Zhang and Elias Bareinboim. Near-optimal reinforcement learning in dynamic treatment regimes. Advances in Neural Information Processing Systems, 32:13401-13411, 2019.
Junzhe Zhang, Daniel Kumor, and Elias Bareinboim. Causal imitation learning with unobserved confounders. Advances in neural information processing systems, 33, 2020.
12

---

Published as a conference paper at ICLR 2022
Distribution Matching
Equivalent Representation
Comments
Distribution Ratio
Sx E sz log(g(s, a))l Bnard,, [log(1 g(s,a)))
GAIL
HH & Ermon
2016)
suP A- Be,andy [g(s, a)] log E F.anddd (eolnaa)
Kl-divergence
Donsker- Donsker-Varadhan Representation
(Kostrikov et al.
2019
x" divergence
sup an [g(s,, a)]
s,ands [g (s)
Variational Representation of f-Divergence
TV Distance
1 Es,andy [g( E B, [g(
Variational Representation of f D-Divergence
Table Different distribution matching techniques and their equivalent representations.
APPENDIX
A BACKGROUND: DISTRIBUTION MATCHING FOR IMITATION LEARNING
A common approach used in (non-confounded) imitation learning is matching the policy's
state-action frequency d=, to the offline target distribution d= Consider a source distribu- tion p AN and target distribution q E AN- GAIL (Ho & Ermon 2016) uses the distribu- tion ratio objective log(p/g). which can estimated using a GAN-like objective Dx(pllg) = suPp:z-(00.1) [gg()((g(>)) (=))] + (l,, g())),, to match the distribution p to q.
This technique be generalized to f-divergences (Csiszár & Shields 2004
Liese & Vajda 2006
Kostrikov et al.) 2019; Ke et al., 2020). Specifically, we wish to minimize a discrepancy measure from p q, namely minpex (p For convex function ʃ [0: oc) IR, the f-divergence of p from q is defined by Dy E, [t())] representation of the f-divergence,
DICE (Kostrikov et al. 2019) uses the variational
Dy(plla) sup - (,({a) 3:
where f* is the Fenchel conjugate of f defined by has closed form solutions for the total variation distance, KL-divergence, -ivverggeee Squared Cam distance, and Jensen-Shannon divergence. Using the variational repre- sentation of the f-divergence we can estimate D, using samples from p and q. Table [ presents examples of various f-divergences and their respective dual formulation. We also add the distribu- tion ratio for comparison to the table, though it is not an f-divergence.
sup xy f(y) The convex conjugate
Hellinger distance,
B CONFOUNDED IMITATION ALGORITHM AND CONVERGENCE GUARANTEES
B.1 A TOY EXAMPLE
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 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,
1(a an, = x1 1(a ac, 22) for a = A, and any action is optimal for s + A. Without loss of generality we as- ('c]((,2)
1. We turn analyze the marginalized state- action frequency, which uniquely defines the set of optimal policies (Puterman) 2014). De- noting (z1) = P. we have that dr (s,a) =dd" (s, a)r1) + (1 - -p)dd (s, a (s,)).)). Then, d= d (s, (1 (s) == A) 1> B, a an) + 1s C, =.))
sume -" (aglB,
13

---

Published as a conference paper at ICLR 2022
X
X
a
ac
&B
ac
B
A
C
B
A
C
m=1
T=0 0
T= 0
r=1 1
Figure 6: A contextual MDP with state space {A, B, C.. action space A = {ap, ac} and context space X' { x2). We assume v(A[z) for all X. The actions ag ac transition the agent to states B, C, respectively, after which the agent receives reward E {0, 1} depending on the context. We assume B, C' are sink states.
Algorithm Confounded Imitation
1: input: Expert data with missing context D d= A sensitivity bound 8 2
2: init: Y /
3: for = =,.... do
4: Sample u(s, U10, Vs,
5: ( 9o) == Es,andz (m.0 (00 (s) Braands (s.a) (90
6: (:(T):)) Expp.... (s.l))) (9) a) B .221
7:
Compute by solving
min
max
(TTT00((ssaa) Amin -
(4)
TEll laol S Sillns
8:
if Wn T then
y (8,a,x)
9: Terminate and return #(c)s,z)
E En d' (s,af,,
10:
else
11:
= {
12:
end if
13:
end for
No Covariate Shift. (suboptimal) policy
Suppose Po = Pe. and 2. Trivially d% (s, (s, = (s, We define the
2, (a a A
(3)
It can be verified that di (s, (s, a) still holds, yet is catastrophic (Equation (1) with value zero. question arises: can show that no suboptimal policy given access the expert data (i.e., access d= (s,a)) and forward model P(s'|s, a,);?
Unfortunately, one cannot prove that mp is suboptimal. Informally, notice that To is an optimal policy for alternative reward function, ro(s, a, 1 r(s, a, yet is catastrophic w.r.t. the true reward Indeed, since is unknown and d (s, a) dz (s,a), cannot reject ro (i.e., we cannot conclude that тo is not the true reward). In other words, one cannot use the data to differentiate which of {To, π²) is the optimal policy.
With Covariate Shift. Next, assume Po + Per and define mo as Equation 3 Let Pe and recall that (21) p. Then, we have that
Pe
To (s. (1
+ pd (s, al22) (1-p)d" (s,)) + pd" (s,a)11) dz (s,
Indeed, the expert data is incapable of distinguishing To and n,, since d = d= and pe is un- known. Unfortunately, as we've shown previously, To achieves value zero. Notice that, unlike the previous section, cannot distinguish =² from the catastrophic policy no for any choice of po-
14

---

Published as a conference paper at ICLR 2022
B.2 A PRACTICAL ALGORITHM
The ambiguity set of Theorem may contain suboptimal policies. Instead of randomly selecting policy from this set, we can choose the average policy. The following proposition shows that such selection is favorable.
EnEr ('(a))))
Proposition 2. Define the mean policy A(a|s,z) = Erer,, Eardaaa'a)) and denote
a- 11 nr E [0, 1]. Then, (M( 2 a²v" + (1- a²) min UM(T).
Remark 1. Note that ā is generally not the average policy 1 ExEr,- A(a(s,)).
Remark 2. In an episodic setting, ā can be estimated by uniformly sampling a policy π E Yg at the beginning of the episode, and playing it until the environment terminates.
Algorithm 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 d5. (s, a) and d (s, while regular- izing with the distance between and all previously collected π. Y. Algorithm 3 also uses sensitivity parameter 5 2 0 (defined formally in Appendix C whenever bounded covariate shift is present. For this section we assume
0.
In practice, the functions L and L; in lines 4 and 5 are estimated using samples from trajectories of π, Te, and D'. We then solve the min-max problem Equation (4) using parametric represen- tations of 9i and online gradient decent. The following proposition states that Algorithm indeed retrieves the set g-.
Proposition 3. Assume Pe Po and |T < co. Then there exists X > O such that for any E (0, )*), Algorithm 3 (with 8 = O sensitivity) will return ß of Proposition after exactly [Ta.] iterations.
B.3 IMITATION WITH CONTEXT-FREE REWARD
We tested Algorithm 3 on both the RecSim environment as well as four-rooms environment with random instantiations of walls. Experiments for the RecSim environment are readily provided in Section Here we describe our simple four-rooms environment and show experiments w.r.t. Theo- rem
3
The four-rooms environment, as depicted in Figure is 15 x 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 blue) must reach the (green) goal while avoiding the (red) mine. When the goal is reached the agent receives 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 æ. Particularly, the context 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
We trained 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 with both 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, shown in Figure 7 after training agent with 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. Never- theless, 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 which indeed states that this information is not needed in the expert data in order to obtain optimal policy.
15

---

Published as a conference paper at ICLR 2022
73
40
62-
Context Free Reward
2.0
1.5
reward
No Covariate Shift
With Covariate Shift
200000
400000
600000
800000
1000000
step
Figure 7: Results for the rooms environment with covariate shift affecting only the distribution of walls. is evident that whenever the reward context-free comparable performance obtained. Runs averaged over seeds.
C BOUNDED HIDDEN CONFOUNDING
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 impos- sible, i.e., one cannot rule out a catastrophic policy. We begin by considering the effect of bounded covariate shift, i.e., Po < C. We then consider almost-context -ree rewards, showing tradeoff w.r.t. Ps the hardness of the imitation problem.
Sensitivity Perspective. common approach causal inference bound the bias of unob- served confounding through sensitivity analysis (Hsu & Small 2013 Namkoong al. 2020 Kallus & Zhou (021). our setting, this confounding bias occurs due to a covariate shift of the unobserved covariates. As we've shown in Thcorem 2 though these covariates ronment, 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.
observed in the online envi-
16

---

Published as a conference paper at ICLR 2022
Assumption 2 (Bounded Sensitivity). We assume that Supp(pe) C Supp(po) and that there exists some 2 1 such that for all x E Supp(pe)
< (o(x)(1 (e-))) S. Pa(x)) Po(x))
Next, we define the notion of d-ambiguity, generalization of the ambiguity set in Definition 1 Definition (6-Ambiguity Set). For a policy x E II, we define the set of all deterministic policies that are 8-close by
r = { D''TTder'
1.% (s,
d2 (8, <
Similar to Definition ! the 6-ambiguity considers all deterministic policies with a marginalized
state-action frequency of distance at most 5 from и. The following results shows that - is sufficient set of candidate optimal policies, as long as Assumption 2
hoolds for some r 2 1.
Theorem 6. [Sufficiency of T- Let Assumption 2 hold for some I 1. Then N² r
The above result suggests that Algorithm can be executed over r by adding
additive uniform noise to d (s, a) (see Line of Algorithm 3), and executing the algorithm for finite number of iterations, finally selecting 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 PT the probability measure over contexts x E X and trajectories (5o, do, .1, 3H) as induced by the policy x and context distribution That is,
H-1
PF =() (x)
II t=
ar, 2))
As the context is observed in the online environment, we can calculate for any ı the quan- tity po (I. As the expert data distribution was generated by the marginalized distribution pr (T)=EEENT Pr. (z, T). unclear if knowledge of PT (z, r) is beneficial.
Fortunately, whenever Assumption 2 holds, a high probability of reconstructing a context in the online environment induces high probability of reconstructing it in the expert data. To see this, assume that there exists å [0, 1] such that for all rf-1, E Supp( (P. (r)), there exists X such that
(P,(())
(5) That is, we assume that for any policy that 5-ambiguous to n², and any induced trajectory of x E X, one can with high probability identify x in the online environment. Importantly, this property can be verified in the online environment. When Assumption and hold, we get that
p (x|T) 2 min (e. + - p()x))))
PT (T(x))((z) Pd (x/T)
Pc (T(x) (T) Po(z) + r(1 po(x)) po(z)+ po(x)) 2 6.
(x(
P (x|T)
PT (T)
other words, we can reconstruct x with probability - å for any trajectory T 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-Dependee 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 To E Y,- is also an optimal policy. Here, relax the assumption the reward, and instead assume bounded depen- dence of the reward on the context. The following definition upper bounds the confounding effect of the reward w.r.t. the context.
17

---

Published as a conference paper at ICLR 2022
Ug Si
U
U
aj
Pe
Si
C
m
F1
a;
8i+1
Figure 8: Contextual MDP Causal Diagram.
Definition 6. Let : X + R such that
min A-RR
|m(s, a, ro(s, s c(z) ,.E,
Using the above definition, we can now show that any policy in xx still approximately optimal, as shown by the following result.
Theorem 7. /Contexi Dependent Reward] Let X + R of Desinition <oe (xg (<(x)) + (Exxp(())(((t)). Then for any x² 2y no E Y.
Denote
v(mu) 2 v(xx() Coe.
A direct corollary for the above result states that for X' + R of Definition all x E X, then for To. π²" of Theorem 7 it holds that v(TO) 2 v(x") 2c. That is, To is an approximately optimal policy.
6 if e(x) for
D RELATION TO CAUSAL INFERENCE
Our work is focused the problem of hidden confounders in expert data for imitation and reinforce- ment 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, |Paar| (2009a)) a basic building block of our frame- work. We then show how the confounded imitation problem can be formalized an intervention over 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 tuple M = (U, V,J, P(U)) where is set of exogenous variables and V is set of endogenous vari- ables. F.
of functions such that fs E F are functions mapping a set of endogenous variables Pa; C VV VV} and set of exogenous variables U, S U to the domain V,, i.e.. V; = fe(Pa;, U.). 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 SCM M is associated with a causal diagram G, as depicted in Figure largely the formulation of stochastic interventions, proposed in Correa & Bareinboim 2020
8 Our framework relies
18

---

Published as a conference paper at ICLR 2022
We consider stochastic, conditional (non-atomic) interventions, defined by regime indicators 02 Pearl 2000; Correa & Bareinboim 2020), defined formally below.
Definition (Non-Atomic Interventions). Given a SCM M = =U, P(U)) and a subset Z CVV, intervention ø {oz..... ,Z } defines new SCM Moz (U,V,F",P(T) in which the set functions F changed to F {fi ( Ulf
Non-atomic interventions are a generalization of the classic atomic do(X X = x) interventions, de- fined by the SCM M, and causal diagram 2z in which all edges incoming into are removed. We have that
Poy = 2)) = P = do 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 E F, we denote it by oz d(f; < f:). indicating that in the interventional distribution f; is used instead of ft. Next we define the identifiability of causal effect under an intervention, as follows.
Definition (Identifiability), Ler X, Y, S with Ynz 0 some SCM with causal diagram g. Given an intervention oz = =fo,..... o2, the causal effect Plylæ, oz) is said to be identifiable from V V if it can be uniquely computed 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 For some horizon (s
> O, the SCM is defined by the set of endogenous variables H U {ail; f ( (r:); (denoting the states, actions, context and rewards, respec- lively), set of exogenous variables U, and functions F = S Saz Sras Spe Suy where Je, corre- spond to the transition function, Sa: the expert policy, Jr; the reward function, So. the context expert distribution, and fwg the initial context-der dependent state distribution.
Relating to our formal definition of our model in Section 2 with slight abuse of notations, the functions (ss, foy fa Sp J adhere to the following relations
P(s 1 s'|s; s, a. = (ff,,s, U)) P(r; = r|s:: = =sa; = - (s, =r)
n" Ta = =, (a,
pa(x) = (r,,
vo(so|>) P(fm (S0. (5. U)),
where 8(.) indicates the Dirac delta distribution.
We are now ready to define the confounded imitation problem. We define the (non-atomic) interven- tion o = do(fp- < fp.)) which replaces fp. with fp. in the contextual MDP SCM defined above. The goal of imitation learning is then to identify the quantities
P(a: 18, = do(Sp. < sc.)) -0SiSH-1,
(6)
where, importantly, we assume we only have access to P(si, az), P(s 02, r). P(so (r). and P(= do(Sp, + Sp.)). Notice that in our setting, P(8 (8;; ai x), P(so|z), and P(z|a, do(fe. + fp.)) correspond to known quantities of the online environment. whereas P(s,, a.) corresponds to the (partially observed) offline expert data. We also emphasize that P(s41 x) is not dependent the intervention oz. That is,
P(s141|s; ai, o2 o2 do(fe. fp.)) P(54418),,0,,)))
Remark. Our work studies slightly different version of the identifiability problem in Equation (6) as we only wish to identify an optimal policy from the set IT as opposed to the single specific extended SCM which includes all
policy n¹. This requirement optimal policies in Wy with the assumption that only one is observed (corresponding to the expert data).
be formalized by defining
19

---

Published as a conference paper at ICLR 2022
Algorithm 4 RL using Expert Data with Unobserved Confounders (Complete Algorithm)
imput: Expert data with missing context 2", A, a, B,N, M 0, policy optimization algorithm
ALG-RL
2: init: Policy n global bonus reward network 9ø
3: for == =....... do
4: Generate dataset of rollouts Rx: ~ dp. d (8, a)
5: Initialize local networks 90... 90, [M]
6: for m = 1.... M do
7: Sample weight vector Wmm. uniformly from An
8:
for e = 1... N do
9:
Sample batch uniformly from Ru ie., {s,, a; 1:-1 Rs
10:
Sample batch according weights Win from D", i.e., f8;, a]) ne Dr Update 88.. according to
11:
Von (0m)
E
om af (5), af)) af (s:. I6 (5; a:)
B
end for
13:
m² arg maimm(MMM Lm (0m)
14: Update global parameters from the selected local network g6 D
15: T² + ALG-RL(r(s,, a, x) (s, a))
16: end for
17: end for
Name
Value
Comments
Batch size
128
Learning rate
5e-5 5
Rollout size
19, 200
Total timesteps
5e6
Num epochs
50
How many training epochs to do after each rollout
0.95
Discount factor
kl
0.2
Initial coefficient KL divergence
kl target
Target value KI. divergence
GAE
The GAE (lambda) parameter
Num workers
Table 2: Hyper-parameters used to train the PPO agent.
E IMPLEMENTATION DETAILS
Our experiments were based off of the recently proposed assistive-gym (Erickson 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 Algorithm 2 is presented in Algorithm 4 Specific hyperparameters used are shown in Tables 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 imitation module, which estimated the augmented reward.
20

---

Published as a conference paper at ICLR 2022
Name
Value
Comments
Batch size
128
Learning rate
1e-1
Imitation Method
X-divergence
Num epochs
50
How many training epochs to do after each rollout
0.9
Dj regularization coefficient
M
10000
Budget for CTS optimizer
Table Hyper-parameters used for imitation and CTS.
Motivated by Kostrikov al. (2019), we regularized the expert demonstrations with samples from d Particularly, we let a E 1I, such that corresponds to the probabillty of sampling an expert example and corresponds the probability of sampling from the replay. This leads to minimizing augmented version of the f-divergence which can be written as
min E Es (s,0la) (( a) - af (g(s, a))l (1 a (s.a) [f (g(s,
g:5 x AFoR
Our imitation module consisted two networks go and hg as proposed in Fu et al. (018). The "done" signal was also added the state for training the imitation module. For training CTS used the Nevergrad optimization platform (Rapin & Teytaud (0118) with a budget of 10000 and one worker. Here copied version of the networks go and ho used for initialization and then approximate the minimum D,
For choosing 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
(1- (r, a) + >adapg( a),
Here V'mean corresponds to the average reward in the replay buffer and Imean to the average bonus in the replay buffer. By averaging the two, we maintained similar scale effectively use the expert data in all the evaluated environments without optimizing for X.
where Aadap
Tewon
Tf Brean
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 gender. The patient's 's impairment was given by either limited movement, weakness, or tremor (with sporadic movement). Finally, the patient's '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, used the following distributions for each feature
gender ~ Bern
mass(gender) ~ N (Amass (pmass (gender),
radius(gender) N (Pradiun (gender), o
height(gender) ~N ((Mnsigh((gnndee o neight) velocity weight Unif([ uvel))
force nontarget weight Unif (eenrge M(arger))
high forces Unif([ forces: Mhigh forces))
food hit weight ~ Unif([Amm, Whil) food velocity weight Unif( (Crood vel: Udood vell)
high pressures weight Unif( [Rigg pressure Uhigh pressure]
impairment Multinomial (rones Plimits Pweakness: Prremor)
21

---

Published as a conference paper at ICLR 2022
Name
Value
Name
Value
Name
Value
Pinale
0.3
Evel
0.225
Enigh presaure
0.009
Mmuas (male)
78.4
U/vll
0.275
Mhigh pressure
0.011
Mmass (female)
62.5
Parget
0.009
Phone
0.1
mass
10
Murget
0.011
Plimits
0.4
Mendius (male)
Enigh focces
0.045
Pueskness
0.3
Mendius (female)
Unigh foroes
0.055
Puemor
0.2
0.1
Phit
0.9
Pheigh (male)
Ubit
1.1
Pheigh (female)
Cjood vel
0.9
Thight
Rrood vel
Table 4: Parameters for context distribution used in assistive-gym
Name
Value
Name
Value
Name
Value
Pinale
0.8
Evel
0.225
Enigh
0.009
Manss (male)
88.4
Uvel
0.275
Whigh pressure
0.111
Panas (female)
72.5
Rurget
0.007
Prose
0.1
2
20
Nargot
0.016
Plinits
0.1
Pradios A
0.9
Rnigh forces
0.035
Poeskness
0.1
Pradios (female)
0.9
Whigh
0.06
Phemor
0.7
Oadius
0.2
fnic
0.4
Mheight (malc)
1.1
Ubit
2.1
Mueight (female)
1.1
fsood
0.4
0.2
2.1
Table Parameters for one of the shifted context distribution used in assistive-gyn gymm
The values for each distribution are provided in Table For setting covariate shift, we used set of distributions that were shifted w.r.t the default context distribution. We then sampled shifted distribution w.p. B and the default distribution w.p. 1 - B. That is, when 3 = 1, the user sampled context only from the shifted distribution. Table shows an example of one of the shifted distribution that were used.
Expert Data Generation. For the assistive-gym experiments we dense reward function for gen- erating the expert data and 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 goal + W2 action + Wg task specific reward + w. preference score,
where the preferences were weighted according to the context features. Specific weights are pro- vided in the implementation of assistive-gym (Erickson et al.) 2020). The sparse reward function did the distance goal (i.e., wi 0).
22

---

Published as a conference paper at ICLR 2022
F MISSING PROOFS
F. PROOFS 3
We begin by proving auxilary lemmas.
Lemma 1. Let iT2 E Tx. Then, xx, T Tz
Proof. We show that x, S T2 and n2 C T.m
Let E then d= (s, a) = dmt (s, a). By our assumption, 22 Y m1 then d (8, d (8, a). Hence, (s, a) = (s, That is, x E Tz. This proves Tx, T..
Similarly, let E my. then d (s,a) (s. a). By our assumption, π2 Y m then d (s.a) dth (8, a). Hence, d= (s, a) = d (s, a). That is, π E Y.. This proves Ynz C Y TT completing the proof.
Lemma 2. Let To be a deterministic policy and let Ma X, such that ro(s, a)) 1fa {)(s,T}} Then To is the unique. optimal policy in Mo.
Proof. By definition of no and
1, ro = E
In particular, E. (ru(s), x), 1. Then
VM, S (1
E (1 (1-7) 7)
Σ
r'ross,, = Mo'
1=0
E
This proves nn is an optimal policy. prove uniqueness, assume by contradiction there exists an optimal policy m # Then,
= Es (s.a {a = To(s, x)){ = Es (a.xx
VMo
In contradiction to is optimal. Then, is unique optimal policy.
We ready prove Theorem 1
Theorem [Sufficiency of ] Assume Pe = Po. Let n E II' and let To Ty Then, T-- T... Moreover, if no $ TT,, then there exists ro such that wa E II, Ma but I & My where Mo = U. V.).
Proof. Let n" E IT and let mo E T... By Lemma 1 as E T.., it holds that T,. T
Next, choosing ro(s, a, x) 1fa = To(s, x)), by Lemma
2 we get that no is an optimal policy in
Mo. This proves To E M Finally, by Lemma 2 l1;, {mo). proving x 11;, if and only n" # mo.
Ener ('(2))))
Proposition 2. Define the mean policy #(a(s,x) Σver, Ea and denote
at I1 nr. E (0, 1). Then, (M( 2 a²v" (1 a1)) a*) minner, (M(x).
Proof. Let के as defined. Then by linearity of expectation
VM = Eg (r(s, = [
Σ
Esa..d (,,,
Σ
T TET,
TET
23

---

Published as a conference paper at ICLR 2022
Denote = IM 0 Ta then
VM
VM
+
1
Σ
VM
[T
T xt*
xC
E
T, \B
\B*
|B'
VM
VM
Y
Tx Te
B
[x. |B=
VM
min VM
T
|T
NET. /B
|
\B*
VM
min VM
[T
[T
TET,
completing the proof.
Theorem 2. /Catastrophic Imitation] Assume |X] 2 |A], and P(s'|s, a, = P(s'ls, a, for all
x, X. Then BT,,1 Te,2 s.t. {Te,, (e,2) are non non-identifiable, catastrophic expert policies.
Proof. We first sketch the proof for the special case X' = {zo, x1). A {ao, 01} and singleton state space S { The general proof follows similarly and is given below.
By letting T1, T2 be the determinisic policies which choose opposite actions at opposite contexts, ie., Ti (2;) = aq, (2; = a1-i, we can choose Pe(x) = d' (I1 (x)) and pe pe(x) = d' (#2(z)) which yield
de (a)
Σ (a
Σ (m) x
Σ =
= : da)a
Similarly, d
For the second part of the proof choose (1, 1fz z₁ a a,} and r2(a, 1fx= F, a = a1-s}. Notice that म. is optimal for r; under any distribution of contexts, yet T. achieves zero reward for ri-i-
We now provide a complete proof for the general case.
Let Po. dl (a). Without loss of generality, let A' {ro. ... Fm A {au..... ag) with k and denote X4 = {I1..... xx} C X. By definition there exists an injective function from A into
Define
ai
I = Vi, = O,...,k: A
f(x)
, O.O.W
[itt1 (mod
x = x;, O,...,k
g(x)
O.W.
Then we can select #1, n2, Pe, Pe as follows
1
(ala)) 1 (a f(x), Xx] +
1a Xx}
k
1
<2(a(ar) = ( E {x) + ++
1 for $ Xz).
and
(x(x) = d' d'( Xx},
(x(x) d' ('(())) Xx).
24

---

Published as a conference paper at ICLR 2022
We get that
de (a) (a) = (e(
Σ f(x=))
Σ (a,) = a) da).
Similarly,
T (a) d' (((z;)))(a== g(z))
d' (ai+1 (mad R.) 2))1{( ai+1 (mad k) a}
k
Σ =
This proves the first part of the theorem. For other parts, choose r1,72 as follows
(a) 1 1 24, a,, S k) 22(a, 1 { z4, a a,+1 (mod 4)) i <s.
Then, by definition, for any P(z) such that Supp(P) n Xx + 0
max E ((z) (x)
E E (1)) min E M, Epo.
#ell
And similarly,
Ben x:]
min E, Ey- (
=)],
#ell
P(=) (22((, x)] max E.- (1)) =)]
mell
The condition on the support holds for Pe. Pe by definition. If. Supp(p.) 0 X₂ 0. then the result holds trivially as (e (a, x]] = (a) x)] = 0 for all æ E II. This completes the proof.
Lemma 3. Assume Supp(p.) S Supp(pe). Then
arg max (tw (1), S arg max Ex (appe
Proof. For elarity we denote
Pe arg (eat
11" arg max (a
IT Supp(e_)
X
arg max E Eg (s [ms,,,,))
To prove the lemma, will show II; II pp)) C II
We begin by proving pe [TSpp(()) Indeed, let n- E II Supp( Then, for any x Supp(pe)
Es (s.als) (s,,,)) max E. Be a,))).
25

---

Published as a conference paper at ICLR 2022
In particular,
(.), (s,a (r)ss)) (x)) max E (s (sal))) a, x) 2 max (r,)o)))
(rss,
where we used Jensen's inequality. This proves Suppt..) C p.
To see the other direction, let me IT, and assume by contradiction that Te & Supplee) Then, there exists E Supp(p.) such that
E ge (e.a))) #)] < max (eaoa
Define
# #-|s,2) 1f2 (tsss + 1fx * 8)
where TT arg max, E (snn R) Then,
v(TT))= = P(x (sd
Σ
P(x) (x dl
eESupp(o)\\
< P(x = a
Σ
Pio =((t),
* =ESupp(P))M())
in contradiction me E I. This proves IT; Supp(p.)'
We have thus shown that II* Pc C I Supp(ee)
Finally, it is left to show that IT Supp(pe) S - Similar to before, x E IT Supp(e.) Then, for any
E Supp(pe), by Jensen's 's inequality
a. (e,alz) (rss, 2)] (x) max aa
2 max (ewpppeeeeessoorddaaaaaaaaaa
This completes the proof.
Theorem 3. |Sufficiency of Context-Free Free Reward] Assume Supp(po) S Supp(pe) and
a(s,a, x) = rrss, a,,')) x') for all x, x' E X. Then Ta- S IIIM
Proof. Let no E Tx will show no M- Since a(s,a, r(8,a,x') for all x X denote r(s, a) = rrss a, a, x). By definition of Y we have that.
1 (s. a) d (s,a)
Then,
v(To)) = (unppp
Ezxpo(a)
Σ (
sES,aEA
Σ NES
a (aspa))ann (s) x))
aa (s,a) [r(s
a (s.a) [r)s
Ex-pa()) Ba (aam)))sss)))
max (ode
Then, To E arg max, Exwps (n) [r a) Applying Lemma 3
To arg max E Bx (s.))) IT.
completing the proof.
26

---

Published as a conference paper at ICLR 2022
F.2 PROOFS FOR SECTION 4
Proposition 1. (Trajectory Sampling Equivalence] Let P which minimizes Problem P2) for some
I1 : A R2 and assume (app(po) C Supp(pe)- Then, there exists p" such that (s,a))]
:; (s. a) lim Ei-p" 10000
7'1 (4
Proof. We can write
d' ((s,alz)=(1--
7 P(s = 3,0 = a/
1=0
(1 - 7) Σ
y P((s = s, =
(1 7) Σ - (
Then, denoting PT (T) Exxup P(T [2)], we get that
(s, = (1- -1-1) 7) Σ
1'fr (s,a))
1=0
r 1 = (s,0)) Σ 1=
= ET-Pt 19
(1 (1-
Since, Supp(po) C Eippr [[1 -)) Ei ( at (s, a}}]
Supp(pe), there exists
p" E An such that
is an unbiased estimator d (s. 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 V Then, Algorithm 7 will converge to an e-oplimal solution to Problem P2) in (++) samples.
Proof. We begin by showing that h(P) = min (P(IBB,)Q.)) is convex in P. We can write Dy in its variational form, rewriting h(P)
((P) min max E [g()]] -e..- =E 9:2-2
where
('(w) = sup f(y)-
We have that [s>)] - R l (g()) is affine in g and x. Therefore, strong duality holds, yielding
h(P)
max min Egwp(g(2)] Ex.r If*
Exwrlo(2)]
max E. o. xEA,
g:2-+R
We have that max A., 2 [ (g(())))) is convex in g as maximum convex (affine) func-
tions in compact set. Therefore h( is also convex as a maximum over convex functions.
Then, the objective in Problem
P2 is convex in d= Following the meta algorithm framework
for convex RL in Zahavy et al. 2021 we write the gradient of (d (s, aslad" (s)))). Notice
that for any general f-divergence Dy (2illus Ey, [s(#). 11
it holds that
V D)(2z1l1s) 0, ioii+00 ti,
27

---

Published as a conference paper at ICLR 2022
and
Vz,Di(Dillui) = Vx, Byg 5
(s(())
Ey:
V f(()
y/s
Specifically, for the K divvergence, DRI(Pille:) -BBy |oo
R1 q:
Then.
(pI = Ey
Pi
Applying Lemma 2 of Zahavy et al. (202T) with a Follow the Leader (FTL) cost player completes the proof.
F.3 PROOFS OF ADDITIONAL RESULTS IN APPENDIX
Proposition 3. Assume Pr. = and | < co. Then there exists X > O such that for any
Ae (0, Algorithm (with ö O sensitivity) will return of Proposition after exactly [Ta' iterations.
Proof. Denote
max TT drv ( (8, x), de d)) (s,,,))) (3,,,,)))
X1
X/ min ael
drv (s, d (8,a)),
where drv is the total variation distance. Let X
and X E (0) and notice that A > 0. X X
To prove the result., we will show that at iteration of the algorithm Tn E Ty- and that either 1n1 or
Tn & Tn-1 = {
Base case (n 1). By the variational representation of the f f-divergence,
max E (s.)) a)] Ex (s.a) (go(s, A-
(f a))l (s, dz (s, (s,
By definition T = arg min drv (d%, (8, a)ll (s, a)) Then, mi E Y=-. Finally since Yo = 0. we have that Ti o Yp.
Induction step. Suppose the claim holds for some k. We will show it holds for + 1.
We begin by showing that TH-1 E Yg. Assume by contradiction that T&+1 E Hldes Th+1 & Tx.. Using the variational form of the f-divergence,
max L<(m+11g)) drv (s, (s, S si SxXxXXX
max L' (1(T+11990)) = Sx
(++++ (sv de (s) Xz.
We have that
max S01S A-
L'(Tx+1:g0) Amin (. 2 fYr X₂ X 0.
9 3
Next, let #+11 E Tr then
max go: A-
LA [7 Xmin (. S 0,
9 A
where we used the fact that L ( = 0 by definition of Y,., and L, 2 0. We have reached contradiction To+1 being solution Equation (4). This proves that mt+1 Y...
28

---

Published as a conference paper at ICLR 2022
Finally, show that Th+1 & Yx if and only if Ta T. First, notice that if Tx xx then Equation (4) will return Tk-1 E T3 by definition of the total variation distance. Next, as- sume xx # xx and assume by contradiction Th+11 E Тę- Then, Ei maxg. L((Tx+1:g) 0, and max max I." (TT+11 9o) = 0, by definition of n- Hence,
L' (T6+11900 Amin (T O.
9o A 9.3 A
In contrast, since T # Y there exists स E Tx such that # & Ya and
max 9o.S A- 9:8 x
go: min L, 9b) S Ao < 0,
in contradiction to Th+ being a solution Equation 4 This completes the proof.
Theorem 6. [Sufficiency of T Let Assumption hold for some r
Then n² rF-
Proof. Let E II. We will show that a Assumption 2
r By elementary algebra, we have that, under
Po(x) Pc(x)
(o(=)(1---)++0)
S po(x)(1-1)) +r.
Since Supp(pe) S Supp(pa),
d% (s, a) Ex (d
(s, a
x]]
Polx) d= (s, x)
Exmpp
(c(x)
(x) (Po (x) - D) + T)d (s, ()ddd(sss x)).
Subtracting d, from both sides we get that
d. (s.a) -dp (s,a) a) s -
x)]
= (T 1) 1)I ( x)]
< r 1
Similarly,
dh 2 Benn [ d"
(()) [
Hence,
dr. (s. - do (s, > (| -1)d x))
= ( - 1) Ex Po (o(=))dd'ss, (s, x)]
2 r
2 -Q
where the last two transitions hold since r 2 1. Then, we have that
[g, (8. - d; (s, S -T
This completes the proof.
Theorem 7. [Context Dependent Reward] Let e : X + R of Definition 6
Denote
for Ex-po(t) (<(x) + Ex=on(<))((())) Then for any x² I1 no E Tş-
v( 2 v(xT) Coc-
29

---

Published as a conference paper at ICLR 2022
Proof. Let 'T" M To E x.. The proof follows similar steps that of Theorem 3
(TT)) = Br a()
da (r(,, ro(s, a)] + brpp
ae (1,, ro(s,a)] + Bz)po(a)
Σ ( x))o(s, a)
sES,aEA
(s (r(, ro(s,a)] +
Σ
ro(s, a) Bx ( x)]
sES.eEA
(r) ro(s, a)] + E, (1 [r)(s,)]
( 2)) ro(s, a)] + E sand (s,))))(ss
p (1(,0,,)) ro(s.a)] + (od
Then,
v(T²) - (TT) (oadd ro(s,a)] Bx aa 20)osss.))))) (w) (<(x)) + Eewpo(s) ((x))]
completing the proof.
30


Reviews:
[{"review_id": "ejo4SZ-w3Qr", "reviewer": "Reviewer_JYzA", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n\nThis paper tackles an important and relevant problem of RL in practice -- imitation learning with access to expert data that might be different from the actual experimental/online environment, and more importantly, that lacks contextual (confounders) information. Results in this paper cover settings that one may encounter in practice. Abstractly speaking, this problem is very similar to a problem that the statistics community starts to get interested in: that is, how to use data from different resources (e.g. a small trial data and a large, but potentially confounded observational data) to improve the qualify and confidence of decision making. Therefore, the authors should be applauded to make connections to causality by writing how their notation and setup can be translated into causal inference language in the appendix.\n\nProofs of the theoretical claims are not technically-dense and relatively self-complete, guarded by numerical experiments. The paper is quite well written and accessible to people who are not so familiar with RL but have backgrounds in a related subject (e.g. causality).\n\nWeaknesses/minor comments:\n\nThe theoretical claims in this paper are all as expected and can be guessed by experts. Results in this paper are not a huge advancement of our understanding of RL but I do think this direction should be explored at some point. But as admitted by the authors, the lack of theoretical support that expert data improves learning efficiency in Section 4 is a main weakness of the paper, leaving a very important gap between theory and practice. Apart from this, I have the following minor comments.\n\n1. In Appendix D, it might be helpful to further relate atomic/non-atomic interventions to terminologies like hard/soft interventions that the causal inference community is more familiar with.\n\n2. The terminology \"stationary distribution of a policy $\\pi$\" before Section 3 is confusing. It sounds like a probability distribution over policies. It might be helpful if the authors could choose a new terminology or make some further clarification when it was first introduced. For example, \"stationary distribution of states and actions under a policy $\\pi$'' might be a more accurate terminology.\n\n3. Any figures in the paper should be self-contained. But one cannot get much information from reading Figure 2 and its legend alone. The authors should expand the legend a bit to explain the figure in more detail.\n\n4. After Theorem 1, \"the imitation solution is uniquely defined by the set $\\Gamma_{\\pi_{\\ast}}$\". Shouldn't this sentence be written as \"the imitation solution is uniquely defined up to the set $\\Gamma_{\\pi_{\\ast}}$\"? Also \"in selection of a suboptimal ...\" should instead be \"in the selection of a suboptimal ...\".\n\n5. Is it standard in the field of RL to assume experts data always follow the optimal policy? Could the authors provide some discussions on this assumption? To me, this assumption looks quite suspicious and might cast doubt on how the theoretical results and algorithms under this assumption could be useful to guide practice.\n\n", "summary_of_the_paper": "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.", "summary_of_the_review": "Overall, this is a technically sound paper. Both the theoretical and empirical results in this paper are expected. But it is good to have a paper to confirm the related theory. Furthermore, considering that the authors did not manage to give a theorem on when expert data improves learning efficiency, I think it is fair to evaluate their contribution as marginally above the acceptance threshold but I would not fight for the acceptance of this paper at this stage.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "bBfYuLTk3YJ", "reviewer": "Reviewer_BBsG", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n- This work tries to categorize the cases that using expert data is helpful in the problem of imitation learning in the presence of latent confounders. It provides cases that is infeasible to solve this problem from expert data. Moreover, it shows some cases this problem becomes feasible. Therefore, it is quite interesting to see the impossibility results.\n- The authors proposed an algorithm in the case of the RL problem that can obtain an optimal policy in the case of covariate shift.\n\nWeaknesses:\n- The case of bounded confounding is more appealing and it makes sense to happen in practical settings. It is unclear why the authors put it in the appendix. Moreover, it is better to elaborate on Assumption 1 on page 16. How much is the set in Theorem 5 big?\n- As mentioned by the authors, it is not determined whether using expert data improves the learning efficiency. It was shown only by some experiments. In fact, we know that we can converge to optimal policy without using expert data if we observe enough samples. Therefore, the advantage of using the proposed algorithm is not justified theoretically.", "summary_of_the_paper": "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.", "summary_of_the_review": "In overall, the paper has some interesting results to characterize the cases that the problem of imitation learning becomes feasible. However, it seems that these are some sufficient conditions and it does not determine the boundaries of feasible cases. Moreover, in the case of the RL problem, the advantage of using the proposed algorithm is not studied theoretically.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "ZvOwSbpE4p", "reviewer": "Reviewer_26tA", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The setting considered in this paper is interesting and practical, and the authors give a theoretical analysis about the limitation of imitation learning in such a setting. Also, the solutions given in the paper are simple and easy to follow. However, there are still some weak points in this paper:\n- It would be better to show a causal diagram, which would be easier for readers to understand the problem definition.\n- Considering that there are some methods try to solve the confounded imitation learning problem [1][2], it would be better to compare the performance of the proposed methods with theirs.\n- The authors only perform experiments on assistant health/ Recommendation systems, which is less considered in the evaluation of RL methods. Can authors perform more experiments on robotics tasks? e.g. half-cheetah, hopper. As far as I know, there are some offline data is available to perform imitation learning.\n\n[1] Kumor, D., Zhang, J., & Bareinboim, E. (2021). Sequential Causal Imitation Learning with Unobserved Confounders.\n\n[2] Zhang, J., Kumor, D., & Bareinboim, E. (2020). Causal imitation learning with unobserved confounders. Advances in neural information processing systems, 33.", "summary_of_the_paper": "This paper considers a very interesting setting:  the expert data in the imitation learning is confounded by the context, and the distribution of context may vary between expert data and online data. The authors give a analysis about the limitation of imitation learning methods under the counfounded imitation learning, and proposed a solution for it even if the counfounder distribution varies.", "summary_of_the_review": "In summary, I think that the setting considered this paper is quite interesting, but it would be better to perform more experiments to evaluate the performance of the proposed method.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "74XOrY4aaMx", "reviewer": "Reviewer_KtAT", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "**Regarding the ambiguity set (definition 1)**, it is claimed that this set consists of all deterministic policies that cannot be distinguished from $\\pi^*$ based on the observed data. However, the condition only checks the marginal state/action distributions match that of the expert, but it should be possible that using the joint distributions along observed trajectories could reveal additional distinguishing information.\n\nConsider the following example. Suppose we have horizon 2, two states: $s_0, s_1$, two actions: $a, b$, we always start at $s_0$, and transitions always go from $s_0 \\to s_1$.\nSuppose we have two contexts $x_1, x_2$  sampled uniformly and two policies $\\pi_1, \\pi_2$ that map contexts to action sequences as follows:\n\n$\\pi_1(x_1) = (a, a)$ (takes action a at timestep 0, and action a at timestep 1)\n\n$\\pi_1(x_2) = (b, b)$ (takes action b at timestep 0, and action b at timestep 1)\n\n$\\pi_2(x_1) = (a, b)$ (takes action a at timestep 0, and action b at timestep 1)\n\n$\\pi_2(x_2) = (b, a)$ (takes action b at timestep 0, and action a at timestep 1)\n\nThe state-action marginal for either policy is uniform over ((s0, a), (s0, b), (s1, a), (s1, b)), but the two are distinguishable if we look at the trajectories themselves instead of just state-action marginals.\nAs such, I don't believe the ambiguity set is actually the minimal set of candidate optimal policies, though I don't think this distinction is ultimately that important for the claims in the paper. \n\n**Regarding proposition 1:** while the average policy is certainly better than the worst policy in the set, I disagree that this really means it is \"robust\" to the ambiguity set. Viewing \\alpha* as simply the probability of getting an optimal policy if we selected a candidate policy at random, the inequality in prop 1 is (essentially) a statement that the mean policy has better value than uniformly guessing a candidate policy and running it (lower bounding the values of all non-optimal candidates with the value of the worst candidate), without quantifying how much better. I'm not convinced this is a very interesting result, and it's likely that stronger assumptions will be needed to obtain any stronger results.\n\nThe result in theorem 2 is unsurprising, but it is definitely good to formalize that imitating state-action marginals is insufficient and can be arbitrarily bad.\n\n**Theorem 3:** The statement of theorem 3 intuitively makes a lot of sense, but doesn't appear to actually say that this makes the imitation problem actually easy, as no algorithm is given to actually find the candidate policies that are proven to be optimal in the theorem. Indeed, the definition of ambiguity sets might not even be reasonable when the context distribution shifts, as even applying the same expert policy with different context distribution would likely get different state-action marginals. If different context distributions induce different distributions over the initial state s_0, and supposing there's no way to return to s_0 after the intial timestep, then the ambiguity set can't contain any policies due to the context mismatch. In this case, the set of these candidate policies is actually empty, in which case theorem 3 is trivially true.\n\nIn **section 4**, the authors show that running RL regularized with the expert policy can lead to a suboptimal policy with context mismatch, and propose an algorithm to do RL regularized by the expert policy, using a learned context distribution to reweight the expert data to allow for consistency. While the proposed algorithm is interesting, the theoretical analysis is not very useful, since (as noted by the authors) it does not show that the expert data is actually useful for policy learning. The paper could be improved substantially by quantifying how the expert can improve the RL learning. For example, it would be good to establish sample complexity for the proposed algorithm in the cases where there actually is no covariate shift and with bounded cova\nriate shift.\n\nTo summarize, I do not believe the theoretical results presented are currently strong contributions. I would particularly appreciate additional results on how the expert data is useful for the proposed RL algorithm, though the algorithm is potentially interesting even without any formal analysis. I suggest heavily reducing the pure imitation sections, as the sufficiency results appear to be uninteresting in those sections. I would suggest only leaving theorem 2 and discussing how matching the expert's state-action distribution is insufficient, as it directly motivates the modifications made in the proposed algorithm.\n\nRegarding the actual proposed algorithm, I do believe it to be an interesting way to incoporate the expert data under covariate shift. Empirical results show that the proposed algorithm can outperform a baseline that directly uses the expert data distrib\nution without adaptively reweighting the context distribution. If the main issue with the naively imitating the expert is simply that it will not converge the optimal policy, I think it would also be important to compare to a slightly more complex baseline that simply places decaying weight on the expert imitaiton over time, allowing the policy to use the expert in initial exploration without being limited at convergence.\nMore generally, it would also be important to reference and compare to existing work on incorporating imperfect expert data in imitation and reinforcement learning settings (for example [1]).\n\nI also have concerns over how applicable the proposed setting is. Is it a realistic assumption that our learned policy can observe the context, but we don't have access to it for the logged expert data? The experiments are synthetically constructed and don't provide a convincing example of why this setting is important. I would appreciate additional discussion on when the proposed setting is relevant. \n\n[1] Gao, Yang, et al. \"Reinforcement learning from imperfect demonstrations.\" arXiv preprint arXiv:1802.05313 (2018).\n", "summary_of_the_paper": "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.", "summary_of_the_review": "I do not believe the paper is currently ready for acceptance. I believe the theoretical contributions are very limited. The proposed algorithm is potentially interesting, but I have concerns about the relevance of the proposed problem setting and how it compares to some simpler baselines. The paper would benefit from more thorough analysis (both theoretical and empirical) of the proposed algorithm as well as additional discussion about the problem setting.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The authors consider the problem of using expert data with unobserved confounders for
both imitation and reinforcement learning settings. They showed how latent confounders 
negatively affect the learning process and proposed a sampling algorithm that mitigates the 
impact and delivers good empirical results.

I agree with the reviewers, this is a borderline paper but with a preference to accept. 
The most salient concern was the lack of clear contribution. While the algorithm is interesting 
with good experimental results that attract interest, it lacks actual theoretical backbone.  
That being said, the authors put in solid effort and addressed concerns sufficiently in the rebuttal stage.
Thus I would prefer to see it accepted. The proposed research direction should be explored in the future.
Decision: Accept (Poster)