On Covariate Shift of Latent Confounders in Imitation and Reinforcement Learning

1 INTRODUCTION . Reinforcement Learning ( RL ) is increasingly used across many fields to create agents that learn via interaction and reward feedback . In many such cases , we rely on experts to perform certain tasks , integrating their knowledge to improve learning efficiency and overall performance . Imitation Learning ( IL , Hussein et al . ( 2017 ) ) is concerned with learning via expert demonstrations without access to a reward function . Similarly , RL settings often use expert data to boost performance , eliminating the need to learn from scratch . In this work we consider the IL and RL paradigms in the presence of partially observable expert data . While expert demonstration data is useful , in many realistic settings such data may be prone to hidden confounding ( Gottesman et al. , 2019 ) , i.e. , there may be features used by the expert which are not observed by the learning agent . This can occur due to , e.g. , privacy constraints , continually changing features in ongoing production pipelines , or when not all information available to the human expert was recorded . As we show in our work , covariate shift of unobserved factors between the expert data and the real world may lead to significant negative impact on performance , frequently rendering the data useless for imitation ( see Figure 1 and Theorem 2 ) . In this paper we define the tasks of imitation and reinforcement learning using expert data with unobserved confounders and possible covariate shift . We focus on a contextual MDP setting ( Hallak et al. , 2015 ) , where a context is sampled at every episode from some distribution , affecting both the reward and the transition between states . We assume that the agent has access to additional expert data , generated by an optimal policy , for which the sampled context is missing , yet is observed in the online environment . We begin by analyzing the imitation-learning problem , ( i.e. , without access to reward ) in Section 3 . Under no covariate shift in the unobserved context , we characterize a sufficient and necessary set of optimal policies . In contrast , we prove that in the presence of a covariate shift , if the true reward depends on the context , then the imitation-learning problem is non-identifiable and prone to catastrophic errors ( see Section 3.2 and Theorem 2 ) . We further analyze the RL setting ( i.e. , with access to reward and confounded expert data ) in Section 4 . Figure 1 depicts a possible failure case of using confounded expert data with unknown covariate shift in a dressing task . Unlike the imitation setting , we show that optimality can still be achieved in the RL setting while using confounded expert data with arbitrary covariate shift . We use a corrective data sampling procedure and prove convergence to an optimal policy . Our contributions are as follows . ( 1 ) We introduce IL and RL with hidden confounding and prove fundamental characteristics w.r.t . covariate shift and the feasibility of imitation . ( 2 ) In the RL setting , under arbitrary covariate shift , we provide a novel algorithm with convergence guarantees which uses a corrective sampling technique to account for the unknown context distribution in the expert data . ( 3 ) Finally , we conduct extensive experiments on recommender system ( Ie et al. , 2019 ) and assistive-healthcare ( Erickson et al. , 2020 ) environments , demonstrating our theoretical results , and suggesting that confounded expert data can be used in a controlled manner to improve the efficiency and performance of RL agents . 2 PRELIMINARIES Online Environment . We consider a contextual MDP ( Hallak et al. , 2015 ) defined by the tuple M = ( S , X , A , P , r , ρo , ν , γ ) , where S is the state space , X is the context space , A is the action space , P : S × S × A ×X 7→ [ 0 , 1 ] is the context dependent transition kernel , r : S × A ×X 7→ [ 0 , 1 ] is the context dependent reward function , and γ ∈ ( 0 , 1 ) is the discount factor . We assume an initial distribution over contexts ρo : X 7→ [ 0 , 1 ] and an initial state distribution ν : S × X 7→ [ 0 , 1 ] . The environment initializes at some context x ∼ ρo ( · ) , and state s0 ∼ ν ( ·|x ) . At time t , the environment is at state st ∈ S and an agent selects an action at ∈ A . The agent receives a reward rt = r ( st , at , x ) and the environment then transitions to state st+1 ∼ P ( ·|st , at , x ) . We define a Markovian stationary policy π as a mapping π : S × X × A 7→ [ 0 , 1 ] , such that π ( ·|s , x ) is the action sampling probability . We define the value of a policy π by vM ( π ) = Eπ [ ( 1− γ ) ∑∞ t=0 γ tr ( st , at , x ) | x ∼ ρo , s0 ∼ ν ( · | x ) ] , where Eπ denotes the expectation induced by the policy π . We denote by Π the set of all Markovian policies and Πdet the set of deterministic Markovian policies . We define the optimal value and policy by v∗M = maxπ∈Π vM ( π ) , and π ∗ M ∈ arg maxπ∈Π vM ( π ) , respectively . Whenever appropriate , we simplify notation and write v∗ , π∗ . We use Π∗M to denote the set of optimal policies inM , i.e. , Π∗M = arg maxπ∈Π vM ( π ) . We also define the set of catastrophic policies Π † M as the set Π†M = arg min π∈Π vM ( π ) . ( 1 ) We later use this set to show impossibility of imitation under arbitrary covariate shift and a contextindependent transition function . Expert Data with Unobserved Confounders . We assume additional access to a confounded dataset consisting of expert trajectories D∗ = { ( si0 , a i 0 , s i 1 , a i 1 , . . . , s i H , a i H ) } n i=1 , where aij ∼ π∗ ∈ Π∗M . The trajectories in the dataset were sampled i.i.d . from the marginalized expert distribution ( under possible context covariate shift ) P ∗ ( s0 , a0 , s1 , a1 , . . . , sH ) = ∑ x ρe ( x ) ν ( s0|x ) ∏H−1 t=0 P ( st+1|st , at , x ) π∗ ( at|st , x ) , where ρe is some distribution over contexts . Importantly , ρe does not necessarily equal ρo – the distribution of contexts in the online environment . Notice that it is assumed that π∗ that generated the data had access to the context xi ( i.e. , π∗ is context-dependent ) , though it is missing in the data . In this work , we consider two settings : 1 . Confounded Imitation Learning ( Section 3 ) : The agent has access to confounded expert data ( with context distribution ρe ) as well as real environment ( S , X , A , P , ρo , ν , γ ) , without access to reward . 2 . Reinforcement Learning with Confounded Expert Data ( Section 4 ) : The agent has access to confounded expert data ( with context distribution ρe ) as well as real environment M = ( S , X , A , P , r , ρo , ν , γ ) , with access to reward . In both settings we aim to find a context-dependent policy which maximizes the cumulative reward . The confounding factor here is w.r.t . the unobserved context and distribution ρe in the offline data . Marginalized State Action Distribution . We denote the state-action frequency of a policy π ∈ Π given context x ∈ X by dπ ( s , a|x ) = ( 1 − γ ) ∑∞ t=0 γ tPπ ( st = s , at = a|x , s0 ∼ ν ( ·|x ) ) , where Pπ denotes the probability measure induced by π . Similarly , given a distribution over contexts , we define the marginalized state-action frequency of a policy π under the corresponding context distribution by dπρo ( s , a ) = Ex∼ρo [ d π ( s , a | x ) ] ( online environment ) , dπ ∗ ρe ( s , a ) = Ex∼ρe [ dπ ∗ ( s , a | x ) ] ( offline expert data ) . A Causal Perspective . Our work sits at an intersection between the fields of RL and Causal Inference ( CI ) . We believe it is essential to bridge the gap between these two fields , and include an interpretation of our model using CI terminology in Appendix D , where we equivalently define our objective as an intervention over the unknown distribution ρe in a specified Structural Causal Model , as depicted in Figure 2 . 3 IMITATION LEARNING WITH UNOBSERVED CONFOUNDERS . In this section , we analyze the problem of confounded imitation learning , namely , learning from expert trajectories with hidden confounders and without reward . Similar to previous work , we consider the task of imitation learning from expert data in the setting where the agent is allowed to interact with the environment ( Ho & Ermon , 2016 ; Fu et al. , 2017 ; Kostrikov et al. , 2019 ; Brantley et al. , 2019 ) . In the first part of this section we assume no covariate shift between the online environment and the data is present , i.e. , ρe = ρo . We lift this assumption in the second part , where we focus on the covariate shift of the hidden confounders . 3.1 NO COVARIATE SHIFT : ρo = ρe We first consider the scenario in which no covariate shift is present between the offline data and the online environment , i.e. , ρo = ρe . We begin by defining the marginalized ambiguity set , a central component of our work . Definition 1 . For π ∈ Π , we define the set of deterministic policies that match the marginalized state-action frequency of π by Υπ = { π′ ∈ Πdet : dπ ′ ρo ( s , a ) = d π ρe ( s , a ) ∀ s ∈ S , a ∈ A } . Recall that , in general , π∗ ∈ Π∗M may depend on the context x ∈ X . Therefore , the set Υπ∗ corresponds to all deterministic policies that can not be distinguished from π∗ based on the observed expert data . The following theorem shows that for any policy π∗ ∈ Π∗M and any policy π0 ∈ Υπ∗ , one could design a reward function r0 , for which π0 is optimal , while the set Υπ∗ is indiscernible from Υπ0 , i.e. , Υπ∗ = Υπ0 ( see Appendix F for proof ) . In other words , Υπ∗ is the smallest set of candidate optimal policies . Theorem 1 . [ Sufficiency of Υπ∗ ] Assume ρe ≡ ρo . Let π∗ ∈ Π∗M and let π0 ∈ Υπ∗ . Then , Υπ∗ = Υπ0 . Moreover , if π0 6= π∗ , then there exists r0 such that π0 ∈ Π∗M0 but π ∗ /∈ Π∗M0 , where M0 = ( S , A , X , P , r0 , ρo , ν , γ ) . The above theorem shows that any policy in Υπ∗ is a candidate optimal policy , yet without knowing the context the expert used . Such ambiguity can result in selection of a suboptimal or even catastrophic policy . We provide a practical algorithm in Appendix B which calculates the ambiguity set Υπ∗ , and returns an average policy , with computational guarantees . In the next subsection we analyze a more challenging scenario , for which ρo 6= ρe . In this case , Υπ∗ may not be sufficient for the imitation problem . 3.2 WITH COVARIATE SHIFT : ρo 6= ρe Next , we assume covariate shift exists between the online environment and the expert data , i.e. , ρo 6= ρe . Particularly , without further assumptions on the extent of covariate shift , we show two extremes of the problem . In Theorem 2 we prove that whenever the transitions are independent of the context , the data can not in general be used for imitation . In contrast , in Theorem 3 we prove that , whenever the reward is independent of the context , the imitation problem can be efficiently solved . Clearly , if Supp ( ρo ) 6⊆ Supp ( ρe ) 1 then there exists x ∈ Supp ( ρo ) for which π∗ is not identifiable from the expert data2 . We therefore assume throughout that Supp ( ρo ) ⊆ Supp ( ρe ) . We begin by defining the set of non-identifiable policies as those that can not be distinguished from their respective state-action frequencies without information on ρe . Definition 2 . We say that { πi } ki=1 are non-identifiable policies if there exist { ρi } k i=1 such that dπiρi ( s , a ) = d πj ρj ( s , a ) for all i 6= j . Next , focusing on catastrophic policies ( recall Equation ( 1 ) ) , we define catastrophic expert policies as those which could be either optimal or catastrophic under ρo for different reward functions . Definition 3 . We say that { πi } ki=1 are catastrophic expert policies if there exist { ri } k i=1 such that for all i , πi ∈ Π∗Mi , and ∃j ∈ [ k ] , j 6= i such that πi ∈ Π † Mj , whereMj = ( S , X , A , P , rj , ρo , ν , γ ) . Using the fact that both ρe and r are unknown , the following theorem shows that whenever P ( s′|s , a , x ) is independent of x , one could find two policies which are non-identifiable , catastrophic expert policies ( see Appendix F for proof ) . In other words , in the case of context-independent transitions , without further information on ρe or r the expert data is useless for imitation . Furthermore , attempting to imitate the policy using the expert data could result in a catastrophic policy . Theorem 2 . [ Catastrophic Imitation ] Assume |X | ≥ |A| , and P ( s′|s , a , x ) = P ( s′|s , a , x′ ) for all x , x′ ∈ X . Then ∃πe,1 , πe,2 s.t . { πe,1 , πe,2 } are non-identifiable , catastrophic expert policies . 1For a distribution P we denote by Supp ( P ) the support of P. 2We define non-identifiability in Definition 2 . We use a similar notion of identifiability as in Pearl ( 2009b ) While Theorem 2 shows the impossibility of imitation for context-free transitions , whenever the reward is independent of the context , the imitation problem becomes feasible . In fact , as we show in the following theorem , for context-free rewards , any policy in Υπ∗ is an optimal policy . Theorem 3 . [ Sufficiency of Context-Free Reward ] Assume Supp ( ρo ) ⊆ Supp ( ρe ) and r ( s , a , x ) = r ( s , a , x′ ) for all x , x′ ∈ X . Then Υπ∗ ⊆ Π∗M . Theorems 2 and 3 suggest that the hardness of the imitation problem under covariate shift lies on a wide spectrum ( as depicted in Figure 3 ) . While dependence of the transition P ( s′|s , a , x ) on x provides us with information to identify x in the expert data , the dependence of the reward r ( s , a , x ) on x increases the degree of confounding in the imitation problem . Both of these results are concerned with arbitrary confounding . Bounded Confounding : A Sensitivity Perspective . A common approach in causal inference is to bound the bias of unobserved confounding through sensitivity analysis ( Hsu & Small , 2013 ; Namkoong et al. , 2020 ; Kallus & Zhou , 2021 ) . In our setting , this confounding bias occurs due to a covariate shift of the unobserved covariates . As we ’ ve shown in Theorem 2 , though these covariates are observed in the online environment , their shifted and unobserved distribution in the offline data can render catastrophic results . Therefore , we consider the odds-ratio bounds of the sensitivity in distribution between the online environment and the expert data , as stated formally below . Assumption 1 ( Bounded Sensitivity ) . We assume that Supp ( ρe ) ⊆ Supp ( ρo ) and that there exists some Γ ≥ 1 such that for all x ∈ Supp ( ρe ) , Γ−1 ≤ ρo ( x ) ( 1−ρe ( x ) ) ρe ( x ) ( 1−ρo ( x ) ) ≤ Γ . Next , we define the notion of δ-ambiguity , a generalization of the ambiguity set in Definition 1 . Definition 4 ( δ-Ambiguity Set ) . For a policy π ∈ Π , we define the set of all deterministic policies that are δ-close to π by Υδπ = { π′ ∈ Πdet : ∣∣∣dπ′ρo ( s , a ) − dπρe ( s , a ) ∣∣∣ < δ , s ∈ S , a ∈ A } . Similar to Definition 1 , the δ-ambiguity set considers all deterministic policies with a marginalized state-action frequency of distance at most δ from π . The following results shows that ΥΓ−1π∗ is a sufficient set of candidate optimal policies , as long as Assumption 2 holds for some Γ ≥ 1 . Theorem 4 . [ Sufficiency of ΥΓ−1π∗ ] Let Assumption 2 hold for some Γ ≥ 1 . Then π∗ ∈ ΥΓ−1π∗ . For the interested reader , we further analyze the case of bounded confounding in Appendix C. We also demonstrate the effect of bounded confounding in Section 5 . In the following section , we show that , while arbitrary confounding may result in catastrophic results for the imitation learning problem , when coupled with reward , one can still make use of the expert data .

Summary: 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.