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