O N C OVARIATE S HIFT OF L ATENT C ONFOUNDERS IN I MITATION AND R EINFORCEMENT L EARNING
Assaf Hallak
Nvidia Research
Gal Dalal
Nvidia Research Technion Gal Chechik Nvidia Research
Uri Shalit Technion A BSTRACT We consider the problem of using expert data with unobserved confounders for imitation and reinforcement learning.
We begin by deﬁning 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 ﬁt this setup.
We then discuss the problem of distribution shift between the expert data and the online environment when the data is only partially observable.
We prove possibility and impossibility results for imitation learning under arbitrary distribution shift of the missing covariates.
When additional external reward is provided, we propose a sampling procedure that addresses the unknown shift and prove convergence to an optimal solution.
Finally, we validate our claims empirically on challenging assistive healthcare and recommender system simulation tasks
I NTRODUCTION
Reinforcement Learning (RL) is increasingly used across many ﬁelds to create agents that learn via interaction and reward feedback.
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 signiﬁcant negative impact on performance, frequently rendering the data useless for imitation (see Figure 1 and Theorem 2).
In this paper we deﬁne the tasks of imitation and reinforcement learning using expert data with unobserved confounders and possible covariate shift.
We focus on a contextual MDP setting, 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 sufﬁcient 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-identiﬁable 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, [Figure] 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.
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 efﬁciency and performance of RL agents.
P RELIMINARIES
[Figure]
Online Environment.
We consider a contextual MDP (Hallak et al., 2015) deﬁned 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 (cid:55)→ [0 , 1] is the context dependent transition kernel, r : S ×
A × X (cid:55)→ [0 , 1] is the context dependent reward function, and γ ∈ (0 , 1) is the discount factor.
We assume an initial distribution over contexts ρ o : X (cid:55)→ [0 , 1] and an initial state distribution ν : S × X
(cid:55)→ [0 , 1] .
The environment initializes at some context x ∼ ρ o ( · ) , and state s 0 ∼ ν ( ·| x ) .
At time t , the environment is at state s t ∈ S and an agent selects an action a t ∈ A .
The agent receives a reward r t = r
( s t , a t , x ) and the environment then transitions to state s t +1 ∼
P ( ·| s t , a t , x ) .
We deﬁne a Markovian stationary policy π as a mapping π : S × X × A (cid:55)→ [0 , 1] , such that π ( ·| s, x ) is the action sampling probability.
We deﬁne the value of a policy π by v M
[Equation]
π
[Equation]
=
[Equation]
, 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 deﬁne the optimal value and policy by v ∗M = max π ∈ Π v M ( π ) , and π ∗M ∈ arg max π ∈ Π v M ( π ) , respectively.
Whenever appropriate, we simplify notation and write v ∗ , π ∗ .
We use Π ∗M to denote the set of optimal policies in M , i.e., Π ∗M = arg max π ∈ Π v M ( π ) .
We also deﬁne the set of catastrophic policies Π †M as the set
[Equation]
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 ∗
[Equation]
where a ij ∼ π ∗ ∈ Π ∗M .
The trajectories in the dataset were sampled i.i.d.
from the marginalized expert distribution (under possible context covariate shift)
[Equation]
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 x
i (i.e., π ∗ is context-dependent), though it is missing in the data.
In this work, we consider two settings:
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.
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 ﬁnd 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 ofﬂine data.
Marginalized State Action Distribution.
We denote the state-action frequency of a policy π ∈ Π given context x ∈ X by d π ( s, a | x ) =
(1 − γ ) (cid:80) ∞
t =
γ t P π ( s t = s, a t = a | x, s 0 ∼ ν ( ·| x )) , where
P π denotes the probability measure induced by π .
Similarly, given a distribution over contexts, we deﬁne the marginalized state-action frequency of a policy π under the corresponding context distribution by
[Equation]
A Causal Perspective.
Our work sits at an intersection between the ﬁelds of RL and Causal Inference (CI).
We believe it is essential to bridge the gap between these two ﬁelds, and include an interpretation of our model using CI terminology in Appendix D, where we equivalently deﬁne our objective as an intervention over the unknown distribution ρ e in a speciﬁed Structural Causal Model, as depicted in Figure 2.
I MITATION L EARNING WITH U NOBSERVED C ONFOUNDERS
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 ﬁrst 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.
[Equation] o = ρ e
We ﬁrst consider the scenario in which no covariate shift is present between the ofﬂine data and the online environment, i.e., ρ o = ρ e .
We begin by deﬁning the marginalized ambiguity set, a central component of our work.
Deﬁnition 1.
For π ∈ Π , we deﬁne the set of deterministic policies that match the marginalized state-action frequency of π by Υ π
=
(cid:110) π (cid:48)
∈ Π det : d π (cid:48)
ρ o ( s, a ) =
d π ρ e ( s, a ) ∀ s ∈ S , a ∈ A (cid:111) .
[Figure]
Recall that, in general, π ∗ ∈
Π ∗M may depend on the context x ∈ X .
Therefore, the set Υ π ∗ corresponds to all deterministic policies that cannot be distinguished from π ∗ based on the observed expert data.
The following theorem shows that for any policy π ∗ ∈ Π ∗M and any policy π 0 ∈ Υ π ∗ , one could design a reward function r 0 , for which π 0 is optimal, while the set Υ π ∗ is indiscernible from Υ π 0 , i.e., Υ π ∗ =
Υ π
(see Appendix F for proof).
In other words, Υ π ∗ is the smallest set of candidate optimal policies.
Theorem 1.
[Sufﬁciency of Υ π ∗ ] Assume ρ e ≡ ρ o .
Let π
∗ ∈ Π ∗M
and let π 0 ∈ Υ π ∗ .
Then, Υ π ∗ =
Υ π 0 .
Moreover, if π 0 (cid:54) =
π ∗ , then there exists r 0 such that π 0 ∈ Π ∗M 0 but
π ∗ / ∈ Π ∗M 0 , where
M 0 = ( S , A , X , P, r 0 , ρ 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 (cid:54) =
ρ e .
In this case, Υ π ∗ may not be sufﬁcient for the imitation problem.
[Equation]
Next, we assume covariate shift exists between the online environment and the expert data, i.e., ρ o (cid:54) =
ρ e .
Particularly, without further assumptions on the extent of covariate shift, we show two extremes of the problem.
In Theorem 2 we prove that whenever the transitions are independent of the context, the data cannot in general be used for imitation.
In contrast, in Theorem 3 we prove that, whenever the reward is independent of the context, the imitation problem can be efﬁciently solved.
Clearly, if Supp ( ρ o ) (cid:54)⊆ Supp ( ρ e ) 1 then there exists x ∈ Supp ( ρ o ) for which π ∗ is not identiﬁable from the expert data 2 .
We therefore assume throughout that Supp ( ρ o ) ⊆
Supp ( ρ e ) .
We begin by deﬁning the set of non-identiﬁable policies as those that cannot be distinguished from their respective state-action frequencies without information on ρ e .
Deﬁnition 2.
We say that { π i } ki =1 are non-identiﬁable policies if there exist { ρ i } ki =1 such that d π i ρ i ( s, a ) =
d π j ρ j ( s, a ) for all i (cid:54) =
j .
Next, focusing on catastrophic policies (recall Equation (1)), we deﬁne catastrophic expert policies as those which could be either optimal or catastrophic under ρ o for different reward functions.
Deﬁnition 3.
We say that { π i } ki =1 are catastrophic expert policies if there exist { r i } ki =1 such that for all i , π
i ∈
Π ∗M i , and ∃ j ∈ [ k ] , j (cid:54) =
i such that π i ∈
Π †M j , where M j = ( S , X , A , P, r j , ρ o , ν, γ ) .
Using the fact that both ρ e and r are unknown, the following theorem shows that whenever
P ( s (cid:48) | s, a, x ) is independent of x , one could ﬁnd two policies which are non-identiﬁable, 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 (cid:48) | s, a, x )
= P ( s (cid:48) | s, a, x (cid:48) ) for all x, x (cid:48) ∈ X .
Then ∃ π e, 1 , π e, 2 s.t. {
π e, 1 , π e, 2 } are non-identiﬁable, 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 Υ π ∗ is an optimal policy.
Theorem 3.
[Sufﬁciency of Context-Free Reward]
Assume Supp ( ρ o ) ⊆
Supp ( ρ e ) and r ( s, a, x ) =
r ( s, a, x (cid:48) ) for all x, x (cid:48) ∈ 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 (cid:48) | 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 ofﬂine 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 deﬁne the notion of δ -ambiguity, a generalization of the ambiguity set in Deﬁnition 1.
Deﬁnition 4 ( δ -Ambiguity Set) .
For a policy π ∈ Π , we deﬁne the set of all deterministic policies that are δ -close to π by Υ
[Equation]
Similar to Deﬁnition 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 sufﬁcient set of candidate optimal policies, as long as Assumption 2 holds for some Γ ≥ 1 .
Theorem 4.
[Sufﬁciency 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.
U SING E XPERT D ATA WITH U NOBSERVED C ONFOUNDERS FOR RL
In the previous section we showed sufﬁcient 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.
Speciﬁcally, we assume access to the true reward signal in the online environment and wish to leverage the ofﬂine expert data to aid the agent in converging to an optimal policy.
To do this, we deﬁne an optimization problem that maximizes the cumulative reward, while minimizing an f -divergence (e.g., KL-divergence, TVdistance, χ 2 -divergence) of state-action frequencies in Υ π ∗ ,
[Equation]
(P1)
[Table]
[Equation]
[Table] [Table]
[Table]
λ > 0 and D f is the f -divergence, where f is a convex function f : (0 , ∞ ) (cid:55)→ R .
The solution to Problem (P1) is an optimal policy
π ∗ ∈ Π ∗M as long as ρ o ≡ ρ e .
Rewriting D f using its variational form (see Appendix A for background on the variational form of f -divergences), we get the following equivalent optimization problem, motivated by Nachum et al.
[Equation]
where f ∗ is the convex conjugate of f , i.e., f ∗ ( y )
=
sup x xy − f ( y ) .
Unfortunately, when covariate shift exists (i.e., ρ o (cid:54) = ρ e ), Problems (P1) and (P1b) are not ensured to converge to an optimal policy (Theorem 2).
Instead, we propose to reformulate Problem (P1b) using a distribution ρ s which minimizes the f -divergence, as follows,
[Equation]
Here, B ( X ) denotes the set of probability measures on the Borel sets of X , and
d π ∗ ρ s ( s, a ) =
E x ∼
ρ s (cid:2)
d π ∗ ( s, a | x ) (cid:3) .
Indeed, whenever Supp ( ρ o ) ⊆
Supp ( ρ e ) , we have that ( π, ρ
[Equation]
[Equation]
, ρ o ) is a solution to Problem (P2).
That is, unlike Problems (P1) and (P1b), Problem (P2) can achieve an optimal solution to the RL problem which still uses the expert data.
Corrective Trajectory Sampling (CTS).
Solving Problem (P2) involves an expectation over an unknown distribution, d π ∗
ρ s ( s, a ) .
Fortunately, d π ∗ ρ s ( s, a ) can be equivalently written as an expectation over trajectories in D ∗ , rather than expectation over unobserved contexts, as shown by the following proposition (see Appendix F for proof):
Proposition 1.
[Trajectory Sampling Equivalence] Let ρ ∗ s which minimizes Problem (P2) for some π ∈ Π , g : S × A (cid:55)→ R , and assume Supp ( ρ o ) ⊆
Supp ( ρ e ) .
Then, there exists p n ∈ ∆ n such that
[Equation]
Proposition 1 allows us to estimate the inner minimization problem over ρ s in Problem (P2) using trajectory samples.
Particularly, we uniformly sample k distributions p n 1 , . . .
p nk , where p nj ∈ ∆ n , and then estimate
[Equation]
which can be estimated by using the variational form of D
f (see Appendix A).
We call this procedure
Corrective Trajectory Sampling (CTS), as it uses complete trajectory samples to account for the unknown context distribution ρ e .
Solving Problem (P2).
Algorithm 1 provides an iterative procedure for solving the optimization problem in Problem (P2).
It uses alternative updates of a cost player (line 5) and policy player (line 6).
In line 5 the gradient of D KL w.r.t.
d π is taken using a Follow the Leader (FTL) cost player to estimate the next bonus iterate.
Finally, in line 6, an efﬁcient, approximate policy optimization algorithm ALG-RL is executed using an augmented reward.
The following theorem, provides convergence guarantees for Algorithm 1 (see Appendix F for proof based on Zahavy et al.
Theorem 5.
Let ALG-RL be an approximate best response player that solves the RL problem in iteration k to accuracy (cid:15)
k
= 1 √ k .
Then, Algorithm 1 will converge to an (cid:15) -optimal solution to
Problem (P2) in O (cid:0) 1 (cid:15) 4 (cid:1) samples.
Notice that, while Theorem 5 shows Algorithm 1 converges to an optimal policy, it does not determine whether the expert data improves overall learning efﬁciency.
We leave this theoretical question for future work.
Nevertheless, in the following section we conduct extensive experiments to show that such data can indeed improve overall performance on various tasks.
A drawback of Algorithm 1 is that it needs to estimate the state-action frequencies.
A practical implementation of Algorithm 1 using online gradient descent (OGD) is provided in Algorithm 2 – the algorithm does not require approximate estimates of the state-action frequencies, but rather, only the ability to sample from them.
Similar to Algorithm 1, we use CTS (see Equation (2)) to estimate ρ s in line 4 according to some f -divergence.
Here, samples are drawn from the current policy as well as samples from D ∗ (with CTS).
We write D f in its variational form, and use a neural network representation for g θ .
We then use the aforementioned samples to minimize the f -divergence using
OGD.
Finally, the policy is updated using ALG-RL and an augmented reward.
E XPERIMENTS
We tested our proposed approach for using expert data with hidden confounding in recommendersystem and assistive-healthcare environments.
For all our experiments we used χ 2 -divergence as our choice of f -divergence, as we found it to work best.
Comparison to other divergences is provided in Figure 5 (left).
We used PPO (Schulman et al., 2017) implemented in RLlib (Liang et al., 2018) for both the imitation as well as RL settings.
We include speciﬁc choice of hyperparameters and an exhaustive overview of further implementation details in Appendix E.
Assistive Healthcare.
Consider the challenge of providing physical assistance to disabled persons.
A recently proposed set of tasks for assistive-healthcare, simulating autonomous robots as versatile caregivers (Erickson et al., 2020).
Each task has a unique goal, affected by both the physical world as well as the patient speciﬁc preferences and disabilities.
We tested our algorithm on four tasks: feeding, dressing, bathing, and drinking.
In these, we used the following features to deﬁne user context: gender, mass, radius, height, patient impairment, and patient preferences.
The patient’s mass, radius, and height distributions were dependent on gender.
The patient’s impairment was given by either limited movement, weakness, or tremor (with sporadic movement).
Finally, the patient’s preferences were affected by the velocity and pressure of touch forces applied by the robot.
For the context distribution ρ o we used the default values as provided by the original environment.
To enforce a distribution shift in the expert data, we shifted each distribution randomly with an additive factor β · ˜ d x , where β ∈ [0 , 1] , and ˜ d x was a random distribution chosen from a set of shifting distributions (see Appendix E).
Here β corresponds to the covariate shift strength.
The expert data was generated by a ﬁxed policy trained using a dense reward function.
A sparse reward signal was used for executing our experiments with the confounded expert data.
For further details, we refer the reader to Appendix E.
Figure 4 depicts results for executing Algorithm 2 on four assistive-gym environments with various covariate shift strengths.
As evident in most of the enviornments, covariate shift strongly affected overall performance.
Particularly in the feeding, drinking, and dressing environments, the success of reaching the goal (i.e., spoon to mouth, cup to mouth, and sleeve to hand) was highly affected by the degree of covariate shift.
This is due to the changing distribution of size, movement, and preferences of the patient, and thus of the goal.
Nevertheless, in all environments, using the expert data (with and without CTS) was found to help induce better policies than executing the same RL algorithm without expert data.
This suggests that expert data can assist in improving overall RL performance, yet correcting for covariate shift may signiﬁcantly improve it in these domains.
Recommender Systems.
In practical recommender systems, sequential interaction with users presents a great challenge for optimizing user long-term engagement and overall satisfaction (Ie et al., 2019).
Leveraging expert data collected using, e.g., surveys to users, may greatly beneﬁt future solutions.
Because features are repeatedly added to these systems, full information in the data is rarely available.
Here, we use the recently proposed RecSim Interest Evolution environment (Ie et al., 2019), simulating sequential user interaction with a slate-based recommender system.
The environment consists of a document model for sampling documents, a user model for deﬁning a distribution over user context features, and a user choice model, which deﬁnes the intent of the user based on observable document features and the user’s sampled context (e.g., personality, satisfaction, interests, demographics, and other behavioral features such as session length or visit frequency).
We used a slate of 10 documents and a user context of dimension 20.
To test the severity of the implications of Theorem 2 in the confounded imitation setting, we used a user-model sampled from a Beta-distribution.
Speciﬁcally, for the expert data the user context features x = ( x 0 , . . .
, x 19 ) were sampled from a Beta-distribution, where x i ∼ Beta ( α i , 4) , and α i = 1 . 5 + 8 .
519 i .
In contrast, the online environment features were sampled from a shifted Beta-distribution with α i =
(1 − β ) (cid:0) 1 . 5 + 8 .
5 19 i (cid:1) + β (cid:0) 10 − 8 .
5 19 i (cid:1) , where β ∈ [0 , 1] deﬁned the shift strength.
While the original environment used a uniform distribution to generate user contexts, the Beta-distribution let us analyze severe forms of covariate shifts, testing the limits of our results in Sections 3 and 4.
Figure 5(a) depicts the effect of increased covariate shift on imitation in the RecSim environment with a dataset of 100 expert trajectories (generated by an optimal policy that had access to the full context).
Without covariate shift ( β = 0 ) an optimal score is achieved, and as β increases, performance decreases.
Particularly, as the mirrored distribution is reached ( β = 1 ), a catastrophic policy is reached.
While the imitator “believes” to have reached an optimal policy, it has in fact reached a catastrophic one, as shown by the orange plot.
Conversely, Figure 5(b) depicts the beneﬁt 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.
R ELATED W ORK
Imitation Learning.
The imitation learning problem has been extensively studied in both the fully ofﬂine (Pomerleau, 1989; Bratko et al., 1995) as well as online setting (Ho & Ermon, 2016; Fu et al., 2018; Kim & Park, 2018; Brantley et al., 2019).
Speciﬁc to our work are GAIL (Ho & Ermon, 2016), AIRL (Fu et al., 2017), and DICE (Kostrikov et al., 2019), which use distribution matching methods.
Our work generalizes these settings to imitation with hidden confounders.
Reinforcement Learning with Expert Data.
Much work has revolved on leveraging ofﬂine data for RL.
Recently, ofﬂine RL (Levine et al., 2020) has shown great improvement over regular ofﬂine imitation techniques (Kumar et al., 2020;
Kostrikov et al., 2021; Tennenholtz et al., 2021a; Fujimoto & Gu, 2021).
In the online RL setting, the combination of ofﬂine data to improve RL efﬁciency has shown great success (Nair et al., 2020).
KL-regularized techniques (Peng et al., 2019; Siegel et al., 2019) as well as DICE-based algorithms (Nachum et al., 2019) have also shown efﬁcient utilization of ofﬂine 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.
There, the authors suggest a notion of imitability, showing when observational data can help identify a policy under some partially observed structural causal model.
Our work provides an alternative perspective on the problem.
In contrast to their work, we rely on concurrent imitation approaches (i.e. state-action frequency matching techniques) and importantly, allow access to the online environment.
Furthermore, we provide guarantees and practical algorithms for both the imitation and RL settings.
We refer the reader to Appendix D for an interpretation of our framework in CI terminology, from a perspective of stochastic interventions in a Structural Causal Model.
Another intersection with causal inference discusses the problem of causal confusion in imitation (de Haan et al., 2019).
Causal confusion is concerned with the problem of nuisances in observed confounded data due to an unknown causal structure.
These “causal misidentiﬁcations” can lead to spurious correlations and catastrophic failures in generalization.
In contrast, our work discusses the orthogonal problem of hidden confounders with possible covariate shift.
Intersection of Causal Inference and Reinforcement Learning.
Previous work has analyzed the problem of optimal control from logged data with unobserved confounders (Lattimore et al., 2016), as well as utilizing (non-expert) confounded data for online interactions (Tennenholtz et al., 2021b).
Much work has revolved around the reinforcement learning setup with access to (nonexpert) confounded data (Zhang & Bareinboim, 2019; Wang et al., 2020).
Other work has considered the problem of off-policy evaluation from confounded data (Tennenholtz et al., 2020; Oberst &
Sontag, 2019; Kallus & Zhou, 2020).
Our work is focused on leveraging expert data with hidden confounders and possible covariate shift in both the imitation and the RL settings.
C ONCLUSION
This work presented and analyzed the problem of using expert data with hidden confounders for both the imitation and RL settings.
We showed that covariate shift of hidden confounders between the expert data and the online environment can result in learning catastrophic policies, rendering imitation learning hard, or even impossible (Theorem 2).
In addition, we showed that when a reward is provided, using the expert data is still possible under arbitrary hidden covariate shift (Theorem 5), and proposed new algorithms for tackling this problem using corrective trajectory sampling (CTS).
A PPENDIX A B ACKGROUND : D ISTRIBUTION M ATCHING FOR I MITATION L EARNING
A common approach used in (non-confounded) imitation learning is matching the policy’s state-action frequency d πρ o to the ofﬂine target distribution d π ∗ ρ e .
Consider a source distribution
p ∈ ∆ N and target distribution q ∈ ∆ N .
GAIL (Ho & Ermon, 2016) uses the distribution ratio objective log ( p/q ) , which can estimated using a GAN-like objective D R ( p || q )
=
sup g : Z(cid:55)→ (0 , 1) E p [log( g ( z ))]
+ E q [log(1 − g ( z ))] , to match the distribution p to q .
This technique can be generalized to f -divergences (Csisz ´ ar & Shields, 2004; Liese & Vajda, 2006;
Kostrikov et al., 2019; Ke et al., 2020).
Speciﬁcally, we wish to minimize a discrepancy measure from p to q , namely min p ∈K D ( p || q ) .
For a convex function f : [0 : ∞ ) (cid:55)→ R , the f -divergence of p from q is deﬁned by D f ( p || q ) =
E q (cid:104) f (cid:16) pq (cid:17)(cid:105) .
DICE (Kostrikov et al., 2019) uses the variational representation of the f -divergence,
[Equation]
where f ∗ is the Fenchel conjugate of f deﬁned by f ∗ ( y )
=
sup x xy − f ( y ) .
The convex conjugate has closed form solutions for the total variation distance, KL-divergence, χ 2 -divergence, Squared
Hellinger distance, Le Cam distance, and Jensen-Shannon divergence.
Using the variational representation of the f -divergence we can estimate D f using samples from p and q .
Table 1 presents examples of various f -divergences and their respective dual formulation.
We also add the distribution ratio for comparison to the table, though it is not an f -divergence.
B C ONFOUNDED I MITATION - A LGORITHM AND C ONVERGENCE G UARANTEES B. A T OY E XAMPLE
To gain intuition, we start with a simple toy example.
Consider the three-state example depicted in Figure 6.
Here, the environment initiates at state A w.p. 1, after which the agent can choose to (deterministicaly) transition to state B or C .
The agent then receives a reward depending on the context.
The optimal policy is given by π ∗ ( a | s, x ) = 1 { a = a B , x = x 1 } + 1 { a = a C , x = x 2 } for s = A , and any action is optimal for s (cid:54) =
A .
Without loss of generality we assume π ∗ ( a B | B, x ) =
π ∗ ( a C | C, x ) = 1 .
We turn to analyze the marginalized stateaction frequency, which uniquely deﬁnes the set of optimal policies (Puterman, 2014).
Denoting ρ e ( x 1 )
=
ρ , we have that d π ∗ ρ e ( s, a ) =
ρd π ∗ ( s, a | x 1 ) + (1 − ρ ) d
π ∗ ( s, a | x 2 ) .
Then, d π ∗ ρ e ( s, a ) =
(1 − γ ) 1 { s = A } + ργ 1 { s = B, a = a B } + (1 − ρ ) γ 1 { s = C, a = a C } .
1: input: Expert data with missing context D ∗ ∼ d π ∗ ρ e , λ > 0 , sensitivity bound δ ≥ 0 .
2: init: Υ = ∅ 3: for n = 1 , . . .
do 4:
Sample u ( s, a ) ∼ U [0 , δ ] , ∀ s, a 5:
L
[Equation]
π
[Equation] [Equation]
E
[Equation]
a
[Equation]
L i ( π ; g i ) := E x
[Equation]
x
[Equation] [Equation]
a
[Equation]
x
[Equation] [Equation]
Compute solving
[Equation]
(4) 8: if π n ∈ Υ then 9:
Terminate and return ¯ π ( a | s, x ) =
[Figure]
,x ) 10: else 11: Υ =
Υ ∪ { π n } 12: end if 13: end for
No Covariate Shift.
Suppose ρ o = ρ e , and ρ = 12 .
Trivially d π ∗ ρ e ( s, a ) =
d π ∗ ρ
o ( s, a ) .
We deﬁne the (suboptimal) policy
[Equation]
It can be veriﬁed that d π ∗ ρ e ( s, a ) =
d π 0 ρ
o ( s, a ) still holds, yet π 0 is catastrophic (Equation (1)) with value zero.
A question arises: can we show that π 0 is a suboptimal policy given access to the expert data (i.e., access to d π ∗ ρ e ( s, a ) ) and a forward model P ( s (cid:48) | s, a, x ) ?
Unfortunately, one cannot prove that π 0 is suboptimal.
Informally, notice that π 0 is an optimal policy for an alternative reward function, r 0 ( s, a, x ) = 1 − r ( s, a, x ) , yet is catastrophic w.r.t.
the true reward r .
Indeed, since r is unknown and d π 0 ρ
o ( s, a ) =
d π ∗ ρ
o ( s, a ) , we cannot reject r 0 (i.e., we cannot conclude that r 0 is not the true reward).
In other words, one cannot use the data to differentiate which of { π 0 , π ∗ } is the optimal policy.
With Covariate Shift.
Next, assume ρ o (cid:54) =
ρ e , and deﬁne π 0 as in Equation (3).
Let (cid:101) ρ e = 1 − ρ e and recall that ρ e ( x 1 ) =
ρ .
Then, we have that
[Equation]
Indeed, the expert data is incapable of distinguishing π 0 and π ∗ , since d π 0 (cid:102) ρ e = d π ∗ ρ e , and ρ e is unknown.
Unfortunately, as we’ve shown previously, π 0 achieves value zero.
Notice that, unlike the previous section, one cannot distinguish π ∗ from the catastrophic policy π 0 for any choice of ρ o .
B. A P RACTICAL A LGORITHM
The ambiguity set of Theorem 1 may contain suboptimal policies.
Instead of randomly selecting a policy from this set, we can choose the average policy.
The following proposition shows that such a selection is favorable.
Proposition 2.
Deﬁne the mean policy ¯ π ( a | s, x )
[Equation] [Equation]
) , and denote
[Equation]
∈ [0 , 1] .
Then, v M (¯ π ) ≥ α ∗ v ∗ + (1 − α ∗ ) min π ∈ Υ π ∗ v M ( π
[Equation]
Remark 1.
Note that ¯ π is generally not the average policy
[Equation]
Remark 2.
In an episodic setting, ¯ π can be estimated by uniformly sampling a policy π ∈
[Equation]
at the beginning of the episode, and playing it until the environment terminates.
Algorithm 3 describes our method for calculating the ambiguity set of Theorem 1, and returns ¯ π of
Proposition 2.
At every iteration of the algorithm, we ﬁnd a new policy in the set by minimizing the total variation distance (written in variational form) between d π ∗ ρ o ( s, a ) and d πρ o ( s, a ) , while regularizing it with the distance between π and all previously collected π i ∈ Υ .
Algorithm 3 also uses a sensitivity parameter δ ≥
(deﬁned formally in Appendix C) whenever bounded covariate shift is present.
For this section we assume δ = 0 .
In practice, the functions L ∗ and L i in lines 4 and 5 are estimated using samples from trajectories of π, π i , and D ∗ .
We then solve the min-max problem of Equation (4) using a parametric representations of g i and online gradient decent.
The following proposition states that Algorithm 3 indeed retrieves the set Υ π ∗ .
Proposition 3.
Assume ρ e ≡ ρ o and | Υ π ∗ | < ∞ .
Then there exists λ ∗ > 0 such that for any λ ∈ (0 , λ ∗ ) , Algorithm 3 (with δ = 0 sensitivity) will return ¯ π of Proposition 2 after exactly | Υ π ∗ | iterations.
B. I MITATION WITH C ONTEXT -F REE R EWARD
We tested Algorithm 3 on both the RecSim environment as well as a four-rooms environment with random instantiations of walls.
Experiments for the RecSim environment are readily provided in
Section 5.
Here we describe our simple four-rooms environment and show experiments w.r.t.
Theorem 3.
The four-rooms environment, as depicted in Figure 7, is a 15 × 15 grid-world in which an agent can take one of four actions: LEFT, RIGHT, UP, or DOWN.
Each action moves the agent in the speciﬁed direction whenever no obstacle is present.
The agent (shown in blue) must reach the (green) goal while avoiding the (red) mine.
When the goal is reached the agent receives a reward of +1 and the episode terminates.
In contrast, if the agent reaches the mine, she receives a reward of − 1 and the episode terminates.
The state space of the environment consists of the agent’s ( row , col ) position in the world.
The rest of the information in the environment is deﬁned by the context x .
Particularly, the context is deﬁned by the position of the green goal, the position of the red mine, and the speciﬁc instantiation of walls (two instantiations are depicted in Figure 7).
We trained an agent with full information (i.e., observed context, including goal location, mine location, and walls).
We generated expert data w.r.t.
the trained agent.
To demonstrate the result of Theorem 3 we executed Algorithm 3 with both a shifted distribution and the default distribution of walls.
We did not change the distribution of goal and mine.
Note that since the distribution of walls only affects the transition function and not the reward, we expect, by Theorem 3, the optimal solution to remain the same.
Indeed, as shown in Figure 7 after training an agent with no access to the contextual information of the walls in the expert data, the agent achieved comparable results both with and without covariate shift on the distribution of walls.
This result seem surprising at ﬁrst, as the walls are essential for solving the task at hand.
Nevertheless, since the distribution of wall is observed in the online environment, the partially observed expert data sufﬁces to obtain an optimal policy.
This settles with Theorem 3 which indeed states that this information is not needed in the expert data in order to obtain an optimal policy.
C B OUNDED H IDDEN C ONFOUNDING
In this section we discuss the imitation learning problem under bounded hidden confounders.
There are several ways to deﬁne boundness of unobserved confounders.
In Section 3 we showed that, under arbitrary covariate shift and context-free transitions, the imitation learning problem is impossible, i.e., one cannot rule out a catastrophic policy.
We begin by considering the effect of bounded covariate shift, i.e., ρ o ρ e ≤ C .
We then consider almost-context-free rewards, showing a tradeoff w.r.t.
the hardness of the imitation problem.
A Sensitivity Perspective.
A common approach in causal inference is to bound the bias of unobserved confounding through sensitivity analysis (Hsu & Small, 2013; Namkoong et al., 2020; Kallus & Zhou, 2021).
In our setting, this confounding bias occurs due to a covariate shift of the unobserved covariates.
As we’ve shown in Theorem 2, though these covariates are observed in the online environment, their shifted and unobserved distribution in the ofﬂine data can render catastrophic results.
Therefore, we consider the odds-ratio bounds of the sensitivity in distribution between the online environment and the expert data, as stated formally below.
Assumption 2 (Bounded Sensitivity) .
We assume that Supp ( ρ e ) ⊆
Supp ( ρ o ) and that there exists some Γ ≥ 1 such that for all x ∈ Supp ( ρ e )
[Equation]
Next, we deﬁne the notion of δ -ambiguity, a generalization of the ambiguity set in Deﬁnition 1.
Deﬁnition 5 ( δ -Ambiguity Set) .
For a policy π ∈ Π , we deﬁne the set of all deterministic policies that are δ -close to π by
[Equation]
Similar to Deﬁnition 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 sufﬁcient set of candidate optimal policies, as long as Assumption 2 holds for some Γ ≥ 1 .
Theorem 6.
[Sufﬁciency of Υ Γ − 1 π ∗ ] Let Assumption 2 hold for some Γ ≥ 1 .
Then π ∗ ∈
Υ Γ − 1 π ∗ .
The above result suggests that Algorithm 3 can be executed over Υ Γ − 1 π ∗ by adding δ = Γ −additive uniform noise to d π ∗ ρ e ( s, a ) (see Line 4 of Algorithm 3), and executing the algorithm for a ﬁnite number of iterations, ﬁnally selecting an average policy from the approximate set.
Context Reconstruction.
When bounded covariate shift is present, one might attempt to learn an inverse mapping of contexts from observed trajectories in the data.
We denote by P πρ the probability measure over contexts x ∈
X and trajectories τ =
( s 0 , a 0 , s 1 , a 1 , . . .
s H ) as induced by the policy π and context distribution ρ .
That is,
[Equation]
As the true context is observed in the online environment, we can calculate for any π the quantity P
πρ o ( x, τ ) .
As the expert data distribution was generated by the marginalized distribution
P π ∗ ρ o ( τ ) =
(cid:80) x ∈X P π ∗ ρ
o ( x, τ ) , it is unclear if knowledge of P πρ o ( x, τ ) is beneﬁcial.
Fortunately, whenever Assumption 2 holds, a high probability of reconstructing a context in the online environment induces a high probability of reconstructing it in the expert data.
To see this, assume that there exists δ ∈ [0 , 1] such that for all π ∈ Υ Γ − 1 π ∗ ,
τ ∈ Supp ( P πρ e ( τ )) , there exists x ∈ X such that
[Equation]
That is, we assume that for any policy that δ -ambiguous to π ∗ , and any induced trajectory of x ∈ X , one can with high probability identify x in the online environment.
Importantly, this property can be veriﬁed in the online environment.
When Assumption 2 and 5 hold, we get that
[Equation]
In other words, we can reconstruct x with probability 1 − δ for any trajectory τ which satisﬁes the above.
This allows us to deconfound essential parts of the expert data, rendering it useful for the imitation problem, even when reward is not provided.
We leave further analysis of this direction for future work.
Context-Dependent Reward.
In Theorem 3 we showed that whenever the reward is independent of the context then the imitation problem is easy, in the sense that any policy π 0 ∈ Υ π ∗ is also an optimal policy.
Here, we relax the assumption on the reward, and instead assume bounded dependence of the reward on the context.
The following deﬁnition upper bounds the confounding effect of the reward w.r.t.
the context.
Deﬁnition 6.
Let (cid:15) :
X (cid:55)→ R such that
[Equation]
Using the above deﬁnition, we can now show that any policy in Υ π ∗ is still approximately optimal, as shown by the following result.
Theorem 7.
[Context Dependent Reward]
Let (cid:15) :
X (cid:55)→ R of Deﬁnition 6.
Denote (cid:15) oe = E x ∼ ρ o ( x ) [ (cid:15) ( x )] + E x ∼ ρ e ( x ) [ (cid:15) ( x )] .
Then for any π ∗ ∈ Π ∗M , π 0 ∈
Υ
π ∗ v ( π 0 ) ≥ v ( π ∗ ) − (cid:15) oe .
A direct corollary for the above result states that for (cid:15) : X (cid:55)→ R of Deﬁnition 6, if (cid:15) ( x ) = (cid:15) , for all x ∈ X , then for π 0 , π ∗ of Theorem 7, it holds that v ( π 0 ) ≥ v ( π ∗ ) − 2 (cid:15) .
That is, π 0 is an approximately optimal policy.
D R ELATION TO C AUSAL I NFERENCE
Our work is focused on the problem of hidden confounders in expert data for imitation and reinforcement learning.
We have chosen to write the paper in terminology famililar to the RL community.
In this section we address and formalize the problem in Causal Inference (CI) terminology.
We begin by deﬁning Structural Causal Models (SCM, Pearl (2009a)) – a basic building block of our framework.
We then show how the confounded imitation problem can be formalized as an intervention over a speciﬁc SCM.
Generally speaking, the causal view casts the environment, namely the expert environment generating the ofﬂine data, and the online environment, as confounders.
Deﬁnition 7 (Structural Causal Models) .
A Structural Causal Model (SCM) is a tuple
M = ( U, V, F , P ( U )) where U is a set of exogenous variables and V is a set of endogenous variables.
F is a set of functions such that f i ∈ F are functions mapping a set of endogenous variables
Pa i ⊆
V \{ V i } and a set of exogenous variables
U i ⊆
U to the domain of V
i , i.e., V i =
f i ( Pa i , U i ) .
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 8.
Our framework relies largely on the formulation of stochastic interventions, as proposed in Correa & Bareinboim (2020).
We consider stochastic, conditional (non-atomic) interventions, deﬁned by regime indicators σ Z (Pearl, 2000; Correa & Bareinboim, 2020), deﬁned formally below.
Deﬁnition 8 (Non-Atomic Interventions) .
Given a SCM M =
( U, V, F , P ( U )) and a subset
Z ⊆ V , an intervention σ Z = { σ Z 1 , . . .
, σ Z n } deﬁnes a new SCM M σ Z =
( U, V, F ∗ , P ( U )) in which the set of functions F is changed to F ∗ =
{ f ∗ i } i : V i ∈{ Z j } nj =1 (cid:83) { f i } i : V
i ∈
V \{ Z j } nj =1 .
Non-atomic interventions are a generalization of the classic atomic do ( X = x ) interventions, deﬁned by the SCM M z and causal diagram G Z in which all edges incoming into Z are removed.
We have that
[Equation]
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 f i ∈
F , we denote it by σ Z = d ( f i ← f ∗ i ) , indicating that in the interventional distribution f ∗ i is used instead of f i .
Next we deﬁne the identiﬁability of a causal effect under an intervention, as follows.
Deﬁnition 9 (Identiﬁability) .
Let X, Y, Z ⊆ V with Y ∩ Z = ∅ in some SCM with causal diagram G .
Given an intervention σ Z = { σ Z 1 , . . .
, σ Z n } , the causal effect P ( y | x, σ Z ) is said to be identiﬁable from V (cid:48) ⊆
V if it can be uniquely computed from P ( V (cid:48) )
for every assignment ( y, x ) in every model that induces G and P ( V (cid:48) ) .
Following the model deﬁnitions of Section 2, we deﬁne the contextual MDP SCM as follows.
Deﬁnition 10.
A contextual MDP SCM is deﬁned by the causal diagram of Figure 8.
For some horizon H > 0 , the SCM is deﬁned by the set of endogenous variables
V = {
s i } Hi =0 ∪ { a i } Hi =0 ∪ { x } ∪ { r i } Hi =
(denoting the states, actions, context and rewards, respectively), a set of exogenous variables U , and functions F = { f s i , f a i , f r i , f ρ e , f ν 0 } , where f s i correspond to the transition function, f a i the expert policy, f r i the reward function, f ρ e the context expert distribution, and f ν 0 the initial context-dependent state distribution.
Relating to our formal deﬁnition of our model in Section 2, with slight abuse of notations, the functions { f s , f g , f a , f ρ e , f ν 0 } adhere to the following relations
[Equation]
where δ ( · ) indicates the Dirac delta distribution.
We are now ready to deﬁne the confounded imitation problem.
We deﬁne the (non-atomic) intervention σ x = do ( f ρ e ← f ρ o ) which replaces f ρ e with f ρ o in the contextual MDP SCM deﬁned above.
The goal of imitation learning is then to identify the quantities
[Equation]
where, importantly, we assume we only have access to P ( s i , a i ) , P (
s i +
| s i , a i , x ) , P ( s 0 | x ) , and P ( x | σ x = do ( f ρ e ← f ρ o )) .
Notice that in our setting, P ( s i +
| s i , a i , x ) , P ( s 0 | x ) , and
P ( x | σ x = do ( f ρ e ← f ρ o )) correspond to known quantities of the online environment, whereas
P ( s i , a i ) corresponds to the (partially observed) ofﬂine expert data.
We also emphasize that
P ( s i
[Equation]
1 | s i , a i , x ) is not dependent on the intervention σ Z .
That is,
[Equation]
Remark.
Our work studies a slightly different version of the identiﬁability problem in Equation (6) , as we only wish to identify an optimal policy from the set Π ∗M , as opposed to the single speciﬁc policy π ∗ .
This requirement can be formalized by deﬁning an extended SCM which includes all optimal policies in Π ∗M , with the assumption that only one is observed (corresponding to the expert data).
1: input: Expert data with missing context D ∗ , λ, α, B, N, M > 0 , policy optimization algorithm
ALG-RL 2: init: Policy π 0 , global bonus reward network g ∗ θ 3: for k = 1 , . . .
do 4:
Generate dataset of rollouts R k ∼ d π k − 1 ρ
o
( s, a ) 5:
Initialize local networks g mθ m ← g θ , m ∈
[ M ] 6: for m = 1 , . . .
M do 7:
Sample weight vector w m uniformly from ∆ n 8: for e = 1 . . .
N do 9:
Sample batch uniformly from R k , i.e., { s i , a i } Bi =1 U ∼ R k 10:
Sample batch according to weights w m from D ∗ , i.e., { s ei , a ei } Bi =
w m ∼ D ∗ 11:
Update g mθ m according to ∇ θ m L m ( θ m )
= 1 B
B (cid:88) i = ∇ θ m (cid:2) (1 − α ) f ∗ ( g mθ m ( s ei , a ei ))
+ αf ∗ ( g mθ m
( s i , a i ))
− g mθ m ( s i , a i ) (cid:3) 12: end for 13: m ∗ ∈ arg min m ∈
[ M ] L m ( θ m ) 14:
Update global parameters from the selected local network g ∗ θ ← g m ∗ θ m ∗ 15: π k ← ALG-RL ( r ( s, a, x ) − λg ∗ θ ( s, a )) 16: end for 17: end for
Name
[Table]
E I MPLEMENTATION D ETAILS
Our experiments were based off of the recently proposed assistive-gym (Erickson et al., 2020) and recsim (Ie et al., 2019) environments.
In this section we discuss further implementation details, hyperparameters, context distributions, and generation of the expert data.
Algorithm Details.
A complete description of Algorithm 2 is presented in Algorithm 4.
Speciﬁc hyperparameters used are shown in Tables 2 and 3.
We implemented the algorithm using the RLlib framework (Liang et al., 2018).
We used PPO (Schulman et al., 2017) as our policy-optimization algorithm.
All neural networks consisted of two-layer fully connected MLPs with 100 parameters in each layer.
We used the same rollout buffer (of size 19200 samples) for both our PPO agent as well as our imitation module, which estimated the augmented reward.
[Table]
Motivated by Kostrikov et al.
(2019), we regularized the expert demonstrations with samples from d π .
Particularly, we let α ∈ (0 , 1] , such that 1 − α corresponds to the probabiilty of sampling an expert example and α corresponds to the probability of sampling from the replay.
This leads to minimizing an augmented version of the f -divergence which can be written as
[Equation]
Our imitation module consisted of two networks g θ and h θ as proposed in Fu et al.
The “done” signal was also added to the state for training the imitation module.
For training CTS we used the Nevergrad optimization platform (Rapin & Teytaud, 2018) with a budget of 10000 and one worker.
Here a copied version of the networks g θ
and h θ were used to for initialization and then approximate the minimum D f .
For choosing λ we used an adaptive strategy which ensured λ balanced the RL objective with the imitation objective.
Speciﬁcally, we used the following tradeoff between reward r and bonus g
[Equation]
where λ adap =
r mean r mean + g mean .
Here r mean corresponds to the average reward in the replay buffer and g mean to the average bonus in the replay buffer.
By averaging the two, we maintained a similar scale to effectively use the expert data in all the evaluated environments without optimizing for λ .
Context Distribution.
For each environment we used a varying context distribution in the expert data, with increasing distance to that of the online environment.
The context distribution for the
RecSim environment is formally described in Section 5.
For the assistive-gym environment the context was deﬁned by the following features: gender, mass, radius, height, patient impairment, and patient preferences.
The patient’s mass, radius, and height distributions were dependent on gender.
The patient’s impairment was given by either limited movement, weakness, or tremor (with sporadic movement).
Finally, the patient’s preferences were affected by the velocity and pressure of touch forces applied by the robot.
We used default average values that were provided with the simulator.
Particularly, we used the following distributions for each feature
gender ∼ Bern ( p male ) mass ( gender ) ∼ N ( µ mass ( gender ) , σ 2 mass ) radius ( gender ) ∼ N ( µ radius ( gender ) , σ 2 radius ) height ( gender ) ∼ N ( µ height ( gender ) , σ 2 height ) velocity weight ∼ Unif ([ (cid:96) vel , u vel ]) force nontarget weight ∼ Unif ([ (cid:96) target , u target ]) high forces ∼ Unif ([ (cid:96) high forces , u high forces ]) food hit weight ∼ Unif ([ (cid:96) hit , u hit ]) food velocity weight ∼ Unif ([ (cid:96) food vel , u food vel ]) high pressures weight ∼ Unif ([ (cid:96) high pressure , u high pressure ]) impairment ∼ Multinomial
( p none , p limits , p weakness , p tremor ) .
[Table] [Table]
The values for each distribution are provided in Table 4.
For setting covariate shift, we used a a set of distributions that were shifted w.r.t.
the default context distribution.
We then sampled a shifted distribution w.p. β and the default distribution w.p. 1 − β .
That is, when β
= 1 , the user sampled a context only from the shifted distribution.
Table 5 shows an example of one of the shifted distribution that were used.
Expert Data Generation.
For the assistive-gym experiments we a dense reward function for generating the expert data and a sparse one for our experiments using the expert data.
Speciﬁcally, the dense reward function used the environment’s default reward function, deﬁned by w 1 · distance to goal + w 2 · action + w 3 · task speciﬁc reward + w 4 · preference score , where the preferences were weighted according to the context features.
Speciﬁc weights are provided in the implementation of assistive-gym (Erickson et al., 2020).
The sparse reward function did not use the distance to goal (i.e., w 1 = 0 ).
F M ISSING P ROOFS F. P ROOFS FOR S ECTION
We begin by proving two auxilary lemmas.
Lemma 1.
Let π 2 ∈ Υ π 1 .
Then, Υ π 1 = Υ π 2 .
Proof.
We show that Υ π 1 ⊆
Υ π 2 and Υ π 2 ⊆
Υ π 1 .
Let π ∈ Υ π 1 , then d π ( s, a ) =
d π 1 ( s, a ) .
By our assumption, π 2 ∈ Υ π 1 , then d π 2 ( s, a ) =
d π 1 ( s, a ) .
Hence, d π ( s, a ) =
d π 2 ( s, a ) .
That is, π ∈ Υ π 2 .
This proves Υ π 1 ⊆
Υ π 2 .
Similarly, let π ∈ Υ π 2 , then d π ( s, a ) =
d π 2 ( s, a ) .
By our assumption, π 2 ∈ Υ π 1 , then d π 2 ( s, a ) =
d π 1 ( s, a ) .
Hence, d π ( s, a ) =
d π 1 ( s, a ) .
That is, π ∈ Υ π 1 .
This proves Υ π 2 ⊆
Υ π 1 , completing the proof.
Lemma 2.
Let π 0 be a deterministic policy and let M
= ( S , A , X , P, r 0 , γ ) such that r 0 ( s, a, x ) = 1 { a = π 0 ( s, x ) } .
Then π 0 is the unique, optimal policy in M 0 .
Proof.
By deﬁnition of π 0 and r 0 ,
[Equation]
In particular, E π
[ r 0 ( s t , a t , x )] = 1 .
Then
[Equation]
This proves π 0 is an optimal policy.
To prove uniqueness, assume by contradiction there exists an optimal policy π 1 (cid:54) =
π 0 .
Then,
[Equation]
In contradiction to π 1 is optimal.
Then, π 0 is a unique optimal policy.
We are now ready to prove Theorem 1.
Theorem 1.
[Sufﬁciency of Υ π ∗ ] Assume ρ e ≡ ρ o .
Let π
∗ ∈ Π ∗M
and let π 0 ∈ Υ π ∗ .
Then, Υ π ∗ =
Υ π 0 .
Moreover, if π 0 (cid:54) =
π ∗ , then there exists r 0 such that π 0 ∈ Π ∗M 0 but
π ∗ / ∈ Π ∗M 0 , where
M 0 = ( S , A , X , P, r 0 , ρ o , ν, γ ) .
Proof.
Let
π ∗ ∈ Π ∗M and let π 0 ∈ Υ π ∗ .
By Lemma 1, as π 0 ∈ Υ π ∗ , it holds that Υ π ∗ =
Υ π 0 .
Next, choosing r 0 ( s, a, x ) = 1 { a = π 0 ( s, x ) } , by Lemma 2 we get that π 0 is an optimal policy in
M 0 .
This proves π 0 ∈ Π ∗M 0 .
Finally, by Lemma 2, Π ∗M
= { π 0 } , proving π ∗ / ∈ Π ∗M
if and only if π ∗ (cid:54) =
π 0 .
Proposition 2.
Deﬁne the mean policy ¯ π ( a | s, x )
[Equation]
,x
[Equation]
, and denote α ∗ =
| Π ∗M (cid:84)
Υ π ∗ | | Υ π ∗ | ∈ [0 , 1] .
Then, v M (¯ π ) ≥ α ∗ v ∗ + (1 − α ∗ ) min π ∈ Υ π ∗ v M ( π ) .
Proof.
Let ˜ π as deﬁned.
Then by linearity of expectation
[Equation]
Denote B ∗ =
Π ∗M ∩ Υ π ∗ , then
[Equation]
completing the proof.
Theorem 2.
[Catastrophic Imitation] Assume |X| ≥ |A| , and P ( s (cid:48) | s, a, x )
= P ( s (cid:48) | s, a, x (cid:48) ) for all x, x (cid:48) ∈ X .
Then ∃ π e, 1 , π e, 2 s.t. {
π e, 1 , π e, 2 } are non-identiﬁable, catastrophic expert policies.
Proof.
We ﬁrst sketch the proof for the special case X = { x 0 , x 1 } ,
A
= { a 0 , a 1 } and a singleton state space S
=
{ s 0 } .
The general proof follows similarly and is given below.
By letting π 1 , π 2 be the determinisic policies which choose opposite actions at opposite contexts, i.e., π 1 ( x i )
=
a i , π 2 ( x i ) =
a 1 − i , we can choose ρ e ( x ) =
d ∗ ( π 1 ( x )) and (cid:101)
ρ e ( x )
=
d ∗ ( π 2 ( x )) which yield
[Equation]
Similarly, d π 2 (cid:101) ρ e ( a ) =
d ∗ ( a ) .
For the second part of the proof choose r 1 ( a, x ) = 1 { x = x i , a = a i } and r 2 ( a, x ) = 1 { x = x i , a = a 1 − i } .
Notice that π i is optimal for r i under any distribution of contexts, yet π i achieves zero reward for r 1 − i .
We now provide a complete proof for the general case.
Let ρ o , d ∗ ( a ) .
Without loss of generality, let X = { x 0 , . . .
, x m } , A = { a 0 , . . .
, a k } with m ≥ k , and denote X k = { x 1 , . . .
, x k } ⊆ X .
By deﬁnition there exists an injective function from A into X .
Deﬁne
[Equation]
Then we can select π 1 , π 2 , ρ e , (cid:101) ρ e as follows
[Equation]
and
[Equation]
We get that
[Equation]
Similarly,
[Equation]
This proves the ﬁrst part of the theorem.
For the other parts, choose r 1 , r 2 as follows
[Equation]
Then, by deﬁnition
[Equation]
for any P ( x ) such that Supp ( P ) ∩
X k (cid:54) = ∅ ,
[Equation]
And similarly,
[Equation]
The condition on the support holds for ρ e , (cid:101) ρ e by deﬁnition.
If, Supp ( ρ o ) ∩ X k = ∅ , then the result holds trivially as E x ∼ ρ o ( x ) ,a ∼ π ( ·| x ) [ r 1 ( a, x )]
=
E x ∼ ρ o ( x ) ,a ∼ π ( ·| x ) [ r 2 ( a, x )] = 0 for all π ∈ Π .
This completes the proof.
Lemma 3.
Assume Supp ( ρ o ) ⊆
Supp ( ρ e ) .
Then
[Equation]
Proof.
For clarity we denote
[Equation]
To prove the lemma, we will show Π ∗ ρ e =
Π ∗ Supp ( ρ e ) ⊆
Π ∗ ρ o .
We begin by proving Π ∗ ρ e =
Π ∗ Supp ( ρ e ) .
Indeed, let π ∗ ∈
In particular,
[Equation]
where we used Jensen’s inequality.
This proves Π ∗ Supp ( ρ e ) ⊆
Π ∗ ρ e .
To see the other direction, let π e ∈ Π ∗ ρ e and assume by contradiction
that π e / ∈ Π ∗ Supp ( ρ e ) .
Then, there exists ˜ x ∈ Supp ( ρ e ) such that
[Equation]
Deﬁne
[Equation]
where π ˜ x ∈ arg max π E s,a ∼ d π ( s,a | ˜ x )
[ r ( s, a, ˜ x )] .
Then,
[Equation]
) E
[Equation] [Equation]
[Equation]
s,
[Equation]
, x )] =
[Equation]
˜ π ) , in contradiction to π e ∈
Π ∗ ρ e .
This proves Π ∗ ρ e ⊆
Π ∗ Supp ( ρ e ) .
We have thus shown that Π ∗ ρ e =
Π ∗ Supp ( ρ e ) .
Finally, it is left to show that Π ∗ Supp ( ρ e ) ⊆
Π ∗ ρ
o .
Similar to before, let
π
∗
∈
Π ∗ Supp ( ρ e ) .
Then, for any x ∈ Supp ( ρ e ) , by Jensen’s inequality
[Equation]
This completes the proof.
Theorem 3.
[Sufﬁciency of Context-Free Reward]
Assume Supp ( ρ o ) ⊆
Supp ( ρ e ) and r ( s, a, x ) =
r ( s, a, x (cid:48) ) for all x, x (cid:48) ∈ X .
Then Υ π ∗ ⊆
Π ∗M .
Proof.
Let π 0 ∈
Υ π ∗ , we will show π 0 ∈ Π ∗M .
Since r ( s, a, x ) =
r ( s, a, x (cid:48) ) for all x ∈ X we denote r ( s, a )
=
r ( s, a, x ) .
By deﬁnition of Υ π ∗ we have that.
[Equation]
Then,
[Equation]
Then, π 0 ∈ arg max π E x ∼ ρ e
[Equation]
x
[Equation]
s,a ∼ d π
[Equation]
s,a | x
[Equation]
[ r ( s
[Equation]
a )] .
Applying Lemma
[Equation]
completing the proof.
F. P ROOFS FOR S ECTION
Proposition 1.
[Trajectory Sampling Equivalence] Let ρ ∗ s which minimizes Problem (P2) for some π ∈ Π , g : S × A (cid:55)→ R , and assume Supp ( ρ o ) ⊆
Supp ( ρ e ) .
Then, there exists p n ∈ ∆ n such that d
[Equation]
, a )
[Equation]
lim
[Equation]
a it
[Equation]
, a
[Equation]
Proof.
We can write
[Equation]
P
[Equation]
Then, denoting P π ρ ∗ s ( τ ) =
E x ∼ ρ ∗ s [ P ( τ | x ) ] , we get that
[Equation]
Since,
Supp ( ρ o ) ⊆
Supp ( ρ e ) , there exists p n ∈ ∆ n such that E i ∼
[Equation]
γ
[Equation]
∞ t =0 γ t 1 (cid:8) s it , a it =
( s, a ) (cid:9)(cid:3) is an unbiased estimator of d πρ ∗ s ( s, a ) .
The result follows by the law of large numbers.
Theorem 5.
Let ALG-RL be an approximate best response player that solves the RL problem in iteration k to accuracy (cid:15)
k
= 1 √ k .
Then, Algorithm 1 will converge to an (cid:15) -optimal solution to
Problem (P2) in O (cid:0) 1 (cid:15) 4 (cid:1) samples.
Proof.
We begin by showing that h ( P ) =
min x ∈ ∆ n D f ( P || E x [ Q x ]) is convex in P .
We can write
D f in its variational form, rewriting h ( P ) as
[Equation]
where
[Equation]
We have that E z ∼ P [ g ( z )] − E x,z ∼ Q x
[ f ∗ ( g ( z ))] is afﬁne in g and x .
Therefore, strong duality holds, yielding
[Equation]
We have that max x ∈ ∆ n E x,z ∼ Q x
[ f ∗ ( g ( z ))] is convex in g as a maximum over convex (afﬁne) functions in a compact set.
Therefore h ( P ) is also convex as a maximum over convex functions.
Then, the objective in Problem ( P 2) is convex in d πρ o .
Following the meta algorithm framework for convex RL in Zahavy et al.
(2021), we write the gradient of D f ( d πρ o ( s, a ) ||
d π ∗ ρ e ( s, a )) .
Notice that for any general f -divergence D f
( x i
|| y i ) =
E y
i
(cid:104)
f (cid:16) x i y i (cid:17)(cid:105) it holds that
[Equation]
and ∇
[Equation]
Speciﬁcally, for the KL -divergence, D KL ( p i || q i ) =
− E
[Equation]
(cid:104)
[Equation]
(cid:16) p
[Equation]
(cid:17)(cid:105) .
Then,
[Equation]
Applying Lemma 2 of Zahavy et al.
(2021) with a Follow the Leader (FTL) cost player completes the proof.
F. P ROOFS OF A DDITIONAL R ESULTS IN A PPENDIX
Proposition 3.
Assume ρ e ≡ ρ o and | Υ π ∗ | < ∞ .
Then there exists λ ∗ > 0 such that for any λ ∈ (0 , λ ∗ ) , Algorithm 3 (with δ = 0 sensitivity) will return ¯ π of Proposition 2 after exactly | Υ π ∗ | iterations.
Proof.
Denote
[Equation]
where d TV is the total variation distance.
Let λ ∗ =
λ ∗ 2 λ ∗ 1 and λ ∈ (0 , λ ∗ ) and notice that λ ∗ > 0 .
To prove the result.
, we will show that at iteration n of the algorithm π n ∈ Υ π ∗ and that either π n / ∈ Υ n − 1 :=
{ π j } n − 1 j =1 or Υ n − 1 = Υ π ∗ .
Base case ( n = 1 ).
By the variational representation of the f -divergence,
[Equation] [Equation]
s, a )) .
By deﬁnition Υ π ∗ = arg min π ∈ Π det d TV
( d πρ o ( s, a ) ||
d π ∗ ρ e ( s, a )) .
Then, π 1 ∈ Υ π ∗ .
Finally since Υ 0 = ∅ , we have that π 1 / ∈ Υ 0 .
Induction step.
Suppose the claim holds for some n = k .
We will show it holds for n = k + 1 .
We begin by showing that π k +1 ∈ Υ π ∗ .
Assume by contradiction that π k +1 ∈ Π det , π k +1 / ∈
Υ π ∗ .
Using the variational form of the f -divergence,
[Equation]
We have that max
[Equation]
Next, let ˜ π k +1 ∈ Υ π ∗ , then
[Equation]
where we used the fact that L ∗ (˜ π k +1 ; g 0 )
= 0 by deﬁnition of Υ π ∗ , and L i ≥ 0 .
We have reached a contradiction to π k +1 being a solution to Equation (4).
This proves that π k +1 ∈ Υ π ∗ .
Finally, we show that π k +1 / ∈ Υ k if and only if Υ k (cid:54) =
Υ π ∗ .
First, notice that if Υ k =
Υ π ∗ then Equation (4) will return π k +1 ∈ Υ k by deﬁnition of the total variation distance.
Next, assume Υ k (cid:54) =
Υ π ∗ and assume by contradiction
π k +1 ∈ Υ k .
Then, ∃ i : max g i L i ( π k +1 ; g i ) = 0 , and max g 0 : S×A(cid:55)→ R L ∗ ( π k +1 ; g 0 ) = 0 , by deﬁnition of Υ π ∗ .
Hence,
[Equation]
In contrast, since Υ k (cid:54) =
Υ π ∗ , there exists ˜ π ∈ Υ π ∗ such that ˜ π / ∈ Υ k , and
[Equation]
in contradiction to π k +1 being a solution Equation (4).
This completes the proof.
Theorem 6.
[Sufﬁciency of Υ Γ − 1 π ∗ ] Let Assumption 2 hold for some Γ ≥ 1 .
Then π ∗ ∈
Υ Γ − 1 π ∗ .
Proof.
Let π ∈ Π .
We will show that π ∈ Υ Γ − 1 π .
By elementary algebra, we have that, under
Assumption 2,
[Equation]
Since Supp ( ρ e ) ⊆
Supp ( ρ o ) ,
[Equation]
Subtracting d πρ e from both sides we get that
[Equation]
Similarly,
[Equation]
Hence
[Equation]
where the last two transitions hold since Γ ≥ 1 .
Then
[Equation]
we have that
[Equation]
This completes the proof.
Theorem 7.
[Context Dependent Reward]
Let (cid:15) :
X (cid:55)→ R of Deﬁnition 6.
Denote
[Equation]
Proof.
Let
π ∗ ∈
Π ∗M , π 0 ∈
Υ π ∗ .
The proof follows similar steps to that of Theorem
[Equation]
Then,
[Equation]
completing the proof
[Equation]