Causal Imitation for Markov Decision Processes: a Partial Identification Approach

Authors: kangrui  ruan
junzhe  zhang
xuan  di
elias  bareinboim
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 05:11:24

Abstract: Imitation learning enables an agent to learn from expert demonstrations when the performance measure is unknown and the reward signal is not specified. Standard imitation methods do not generally apply when the learner and the expert's sensory capabilities mismatch and demonstrations are contaminated with unobserved confounding bias. To address these challenges, recent advancements in causal imitation learning have been pursued. However, these methods often require access to underlying causal structures that might not always be available, posing practical challenges.
In this paper, we investigate robust imitation learning within the framework of canonical Markov Decision Processes (MDPs) using partial identification, allowing the agent to achieve expert performance even when the system dynamics are not uniquely determined from the confounded expert demonstrations. Specifically, first, we theoretically demonstrate that when unobserved confounders (UCs) exist in an MDP, the learner is generally unable to imitate expert performance. We then explore imitation learning in partially identifiable settings --- either transition distribution or reward function is non-identifiable from the available data and knowledge. Augmenting the celebrated GAIL method (Ho \& Ermon, 2016), our analysis leads to two novel causal imitation algorithms that can obtain effective policies guaranteed to achieve expert performance.

Main paper: 
 introduction: 
 Children often learn how to behave in an unfamiliar environment by imitating adults. Imitation learning (IL) enables a learning agent to behave in an unknown environment by observing expert demonstrations. It provides a viable approach for policy learning from demonstrations when the reward function is not fully known and reward signals are not specified [2, 30, 8, 20, 31] . Imitation learning has been widely applied across disciplines, such as autonomous driving [13, 39] , robotics [18] , natural language processing [10, 11, 40, 41] , and chronic disease management [52, 44] . It has been acknowledged in the literature that imitation learning could face significant challenges when unobserved confounding bias in expert demonstrations cannot be ruled out a priori [17, 57, 27, 42] . For illustration with simplicity, consider a Multi-Armed Bandit (MAB) model [28] described in Fig. 1 ; X ∈ {0, 1} is a binary action, andNon-Identfiable RewardStandard IRL (e.g., GAIL [19] ) CAIL-R (Alg. 1 in Sec. 3.1)Non-Identfiable Transition CAIL-T (Alg. 2 in Sec. 3.2) Inimitable (Thm. 1 in Sec. 2.1)Table 1 : Summary of main contributions in this paper, including the analysis and proposed algorithms.Y is the reward; U is a latent covariate (to the imitator) uniformly drawn over a binary domain {0, 1}.Values of the reward are decided by a reward function Y ← X ⊕ U where ⊕ is a "xor" operator. An expert demonstrator, having access to covariate U , selecting action based on an expert policy X ← ¬U . Evaluating the expert's performance gives E[Y ] = E[¬U ⊕ U ] = 1. On the other hand, an imitator, mimicking the expert's behavior, will follow a policy π(X) = P (X) = 0.5, selecting action uniformly at random. Evaluating the imitator's performance gives E π [Y ] = x π(x)E[x⊕U ] = 0.5, far from the expert's performance E[Y ] = 1.Causal Inference (CI) addresses the challenges of unobserved confounding bias within the observational data [32, 47, 53] . It leverages causal knowledge integral to the data generation process, typically represented as a causal diagram or potential outcomes [32, 43, 5] . More recently, incorporating causal inference methods into the imitation learning paradigm, causal imitation learning has evolved into a critical area of research [16, 57, 27, 7, 49, 42] . To compensate for the presence of unobserved confounding bias, these methods rely on additional structural or parametric knowledge about causal relationships among variables in the environment. By utilizing such domain knowledge, the imitator is able to recover the underlying system dynamics (i.e., causal effect) from confounded expert demonstrations and, in turn, obtain an imitating policy that can achieve the expert's performance.By and large, the combination of causal knowledge and observational data does not always allow one to point-identify the causal effect, called the non-identifiable. That is, more than one parametrization of the target effect is compatible with the same observational data and model assumptions [32, Def. 3.2.2] . For instance, in the MAB environment described previously, the imitator's performanceE π [Y ]is not identifiable from the confounded observational distribution P (X, Y ) [32, Thm. 3.4.1] . Partial identification methods concerned with inferring about target causal effects in non-identifiable settings, and has been a target of growing interest in the domains of causal inference [3, 12, 37, 14, 58] , econometrics [21, 35, 38, 48] , and more recently, in machine learning [25, 24, 23] . Among these works, two approaches are often employed: (1) bounds are derived for the target effect under minimal assumptions, or (2) additional untestable assumptions are invoked under which the causal effect is identifiable, and then sensitivity analysis is conducted to assess how the target causal effect varies as the untestable assumptions are changed. Despite their effectiveness in addressing data bias, partial identification has still been rarely explored in the context of imitation learning.This paper studies the partial identification for imitation learning in a generalized sequential decisionmaking environment of Markov Decision Processes (MDPs, [36] ). The imitator must determine a sequence of actions, while unobserved confounders cannot be ruled out a priori in expert demonstrations. We discuss the solutions case-by-case, depending on the identifiability of the underlying system dynamics from the confounded data, including the reward function R and the transition distribution T . Specifically, our contributions can be summarized as follows. (1) We theoretically prove that when unobserved confounders generally exist, it is infeasible to learn a robust policy that is guaranteed to achieve expert performance from the demonstration data. (2) When only the transition distribution T is identifiable, we propose a novel imitation algorithm that leverages the bounds over the non-identifiable reward R; by matching the weighted occupancy measure, the imitator is able to obtain a policy that can outperform the expert. (3) We propose an alternative algorithm when the reward R is identifiable, but there is unobserved confounding affecting the transition T . Our proposed algorithms could be implemented by augmenting the celebrated generative adversarial imitation learning framework (GAIL, [19] ). Table 1 briefly summarizes this paper's main contributions. Due to space constraints, all proofs are provided in Appendices A and B.This section introduces the basic notations and definitions used throughout the paper. We use capital letters to denote random variables (X), small letters for their values (x), and D X for the domain of X. For an arbitrary set X, let |X| be its cardinality. Fix indices i, j ∈ N. Let X i:j stand for a sequence of variables {X i , X i+1 , . . . , X j }; for consistency, the sequence X i:j = ∅ if j < i. We denote by P (X) represents a probability distribution over variables X, and will consistently use P (x) as abbreviations for probabilities P (X = x). Finally, 1{Z = z} is an indicator function that returns 1 if event Z = z holds true; otherwise, it returns 0.The basic semantical framework of our analysis rests on structural causal models (SCMs) [32, 4] .An SCM M is a tuple ⟨V , U , F, P (U )⟩, where V is a set of endogenous variables and U is a set of exogenous variables. F is a set of functions s.t. each f V ∈ F decides values of an endogenous variable V ∈ V taking as argument a combination of other variables in the system. That is,V ← f V (PA V , U V ), PA V ⊆ V , U V ⊆ U .Exogenous variables U ∈ U are mutually independent, values of which are drawn from the exogenous distribution P (U ). Naturally, M induces a joint distribution P (V ) over endogenous variables V , called the observational distribution. An intervention on a subset X ⊆ V , denoted by do(x), is an operation where values of X are set to constants x, replacing the functions {f X : ∀X ∈ X} that would normally determine their values.For an SCM M , let M x be a submodel of M induced by intervention do(x). For a set Y ⊆ V , the interventional distribution P x (Y ) induced by do(x) is defined as the joint distribution over Y in the submodel M x , i.e., P x (Y ; M ) ≜ P (Y ; M x ). We leave M implicit when it is obvious from the context. For a detailed survey on SCMs, we refer readers to [32, Ch. 7 ] and [4] .Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?Answer: [Yes] Justification: The major claims made in the abstract and introduction has accurately reflect the paper's contributions and scope. Specifically, we summarized our contributions in Sec. 1, e.g., the theoretical findings and the propose innovative algorithms. Additionally, Table 1 provides a brief summary of this paper's main contributions.Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 
challenges of unobserved confounding: 
 We focus on the sequential decision-making setting of an agent operating in a MDP environment [36] over a series of interventions :t-1 , Ȳ1:t-1 ⊥ ⊥ Xt:∞ , St+1:∞ , Ȳt:∞ | S t .It follows from Def. 1 that for any horizon T , the joint distribution P π X1:T , S1:T , Ȳ1:T generated by a policy π(X | S) factorizes as follows,2 P π ( x1:T , s1:T , ȳ1:T ) = P (s 1 )EQUATIONwhere the transition distribution T and the reward distribution R are interventional queries given by For analytical clarity, we define reward function R(s, x) as the expected value y yR(s, x, y). Fix the discounted factor γ ∈ [0, 1]. A common objective for an agent to optimize is the cumulative returnEQUATIONR(s t , x t , y t ) = P xt (y t | s t ) = ut 1 {y = f Y (s t , x t , u t )} P (u t ) (3) S 1 X 1 Y 1 S 2 X 2 Y 2 S 3 X 3 Y 3 (a) MDPexp S 1 X 1 Y 1 S 2 X 2 Y 2 S 3 X 3 Y 3 (b) MDPobsR t = Y t + γY t+1 + γ 2 Y t+2 + • • • = ∞ k=0 γ k Y t+k .Imitation Learning. When a detailed parametrization of the transition distribution T and the reward function R is available, the agent can obtain an optimal policy using standard planning algorithms [36, 6] . However, in many practical applications, complete knowledge of these parametrizations is often unavailable, necessitating a learning process. In this paper, we consider the imitation learning setting, where the agent has access to observed trajectories generated by the expert. More specifically, at each time step t, the expert selects an action X t ← f X (s t , u t ) based on the current state S t = s t and latent noise U t = u t . Fig. 2b shows the graphical representation of the datagenerating process of the expert; the highlighted bi-directed arrows, e.g., X t ↔ Y t , indicate the presence of an unobserved confounder U ∈ U t affecting both the action X t and outcome Y t . We summarize the expert trajectories using the observational distribution P (X, S, Y ) over sequences of variables X = {X 1 , X where T and R are the expert's nominal transition distribution and reward function computed from the observational distribution as follows:EQUATIONBy convention in imitation learning, we assume the rewards Y t are generally unobserved to the learner; instead, it has access to a parametric family R containing the expert's nominal reward functionE [Y t | s t , x t ].Given the expert demonstrations D sampled from P (X 1 , X 2 , . . . , S 1 , S 2 , . . . ) and the parametric reward family R, the imitator attempts to learn policy π that can achieve expert performance, i.e., E π∞ t=1 γ t-1 Y t ≥ E ∞ t=1 γ t-1 Y t .Standard imitation methods focus on the identifiable setting where the imitator's transition distribution T and reward function R is consistent with the expert's nominal transition T and reward R. Formally, Definition 2 (Causal Consistency). For an interventional distribution P π and an observational distribution P satisfying the Markov property (Def. 1), Causal Consistency is said to hold with respect to P π and P if the following statement is true, for every time step t = 1, 2, . . . ,EQUATIONWhen the invariances of Def. 2 hold, the learner could recover the parametrization of the transition distribution T from observational data P (X, S) and infer about the reward function R from the parametric family R. An imitating policy π is obtainable by solving the following minimax program,EQUATIONwhere the imitator's ρ π and the expert's ρ occupancy measures are defined as, respectively, ρ π (s) = ∞ t=0 γ t P π (S t = s) and ρ(s) = ∞ t=0 γ t P (S t = s). The solution π is guaranteed to achieve expert performance when the gap ν * ≤ 0. This means that the imitator following policy π performs as well as the expert, even in the worst-case environment instance compatible with the demonstration data and model assumption. Several imitation learning algorithms have been proposed to solve the optimization problem in Eq. ( 7), including [1, 50, 19] .Graphical criteria exist [33, 46, 34] to examine whether causal consistency (Def. 2) holds from causal knowledge of the environment, including the celebrated backdoor condition [32, Def. 3.3.1] , [15] . In MDPs, this means that the causal links between the latent noise U t and action X t are not effective -the graphical representation of the imitator's (Fig. 2a ) and the expert's (Fig. 2b ) data-generating process coincide. However, in practice, causal consistency could be fragile and does not necessarily hold due to the presence of unobserved confounders in the demonstration data [57, 39] . The remainder of this paper studies imitation learning when violations occur in the invariance relationships of Eq. ( 6).We first consider the imitation setting described in Fig. 2b where unobserved confounders generally exist in the expert demonstrations; both the transition distribution T and reward function R are not identifiable from Eq. ( 6). Here, we will show that expert performance is not imitable by constructing worst-case MDP instances where the expert always outperforms the imitator.The state value function V π (s) is defined as the expected return given the imitator's starting state S t = s following a policy π, i.e., V π (s) = E π [R t | S t = s].For any policy π, the imitator's performance can be written as E π [R 1 ] = s1 P (s 1 )V π (s 1 ). The value function of any state s can thus be recursively defined using the celebrated Bellman Equation [6] :EQUATIONwhere γ denotes the discount factor. While the transition distribution T and the reward function R are not uniquely discernible from the observational distribution due to the unobserved confounding, it is still possible to learn about them from demonstrations using partial identification. Without loss of generality, the reward Y t is normalized in a real interval [0, 1]. Through rigorous adaptation of the bounding strategies established in [29, 55] , we successfully derive the bounds for the transition distribution T and reward function R, for every realization (s, x, s ′ ) ∈ S × X × S, T (s, x, s ′ ) ∈ T (s, x, s ′ ) P (x | s), T (s, x, s ′ ) P (x | s)EQUATIONAmong the above quantities, T and R are the expert's nominal transition distribution and reward function in Eq. (5); P (x | s) stands for the propensity score P (X t = x | S t = s) and P (¬x | s) = 1 -P (x | s). We can then construct a worst-case MDP for any policy π at state s by solving the following optimization program: minimize the Bellman's equation in Eq. ( 8) as the objective function, subject to the observational constraints in Eqs. ( 9) and (10) . Solving this program enables a valid MDP construction since the transition distribution T and the reward function R are independent components induced by the underlying model and can be optimized separately.Theorem 1. Given any positive observational distribution P (X, S, Y ) > 0, there exists an MDP model M compatible with the causal graph of Fig. 2b such that P (X, S, Y ; M ) = P (X, S, Y ) and for any policy π, any time step t = 1, 2, . . . , any state s ∈ S,EQUATIONIn other words, there always exists a candidate MDP instance M compatible with the demonstration data such that an imitator is always unable to achieve expert performance (r.h.s. in Eq. ( 11)), regardless of the deployed policy π. It follows from Thm. 1 that there is no policy π learnable from confounded demonstrations that is guaranteed to perform at least as the expert in all possible scenarios. This means that expert performance is not imitable when unobserved confounding generally exists. The following example demonstrates the challenges of unobserved confounding in a single-stage MDP.Example 1 (Single-Stage MDP). Consider a 1-stage MDP model with horizon T = 1. For any policy π(X 1 | S 1 ), the imitator's expected return isE π [Y 1 ] = s1,x1 R(s 1 , x 1 )π(x 1 | s 1 )P (s 1 ). It follows S 1 X 1 Y 1 S 2 X 2 Y 2 S 3 X 3 Y 3 (a) MDPobs-Y S 1 X 1 Y 1 S 2 X 2 Y 2 S 3 X 3 Y 3 (b) MDPobs-SFigure 3 : (a) MDP obs-Y shows a data-generating process for expert demonstrations where only the reward Y t is confounded with the action X t ; (b) MDP obs-S shows a data-generating process for expert demonstrations where only the next state S t+1 is confounded with the action X t .from the tight lower bound in Eq. ( 10) that there exists an worst-case MDP model M compatible with the observational distributionP (X 1 , S 1 , Y 1 ) such that R(s 1 , x 1 ) = E [Y 1 | s 1 , x 1 ] P (x 1 | s 1 ).In this MDP instance M , the imitator's expected return can be further written asEQUATIONThe last step holds since probabilities of the policy π(x 1 | s 1 ) ∈ [0, 1] and x1 π(x 1 | s 1 ) = 1.Marginalizing the above equation givesE π [Y 1 ] < E[Y 1 ]-the imitator is unable to achieve expert performance regardless of the deployed policy π. This analysis applies analogously to the MAB model in Fig. 1 , which can be thought of as a 1-stage MDP with no initial state S 1 = ∅. We refer the readers to Appendix F for more examples about the 2-stage MDP.Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes] Justification: We have discussed the requirements and limitations of the work in Sec. 1 and Sec. 2. To address infinite horizon decision-making challenges, we utilize the Markov Property, as outlined in Def. 1. However, our study generalizes standard imitation methods by focusing on scenarios where causal consistency does not universally hold true.• The answer NA means that the paper does not include experiments.• If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not. • If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable. • Depending on the contribution, reproducibility can be accomplished in various ways.For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed. • While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. , with an open-source dataset or instructions for how to construct the dataset). (d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility.In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results. • At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable). • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
partial identification for robust imitation: 
 The impossibility results in Thm. 1 imply that robust imitation cannot be guaranteed when unobserved confounders generally exist in the demonstration data. This means we must explore alternative assumptions to learn an imitating policy guaranteed to achieve expert performance. Meanwhile, standard imitation methods apply when causal consistency of Def. 2 holds, and no unobserved confounder affects the transition or reward function. A natural question at this point arises: whether robust imitation is feasible for settings between the unconfounded (Fig. 1b ) and fully confounded cases (Fig. 1a ), where unobserved confounding bias affects only either the transition distribution or reward function? This section aims to answer this question.We first examine the setting graphically described in Fig. 3a where the reward function is confounded, while the transition distribution is identifiable from the demonstration data. In this case, the first equation of Def. 2 holds while the second one fails. To initiate the discussion, we write the expected return of a candidate policy π in an MDP environment as follows [36] ,EQUATIONAmong quantities in the above equation, the state occupancy measure ρ π (s) = ∞ t=0 γ t P π (S t = s) is a function of the initial state distribution P (s) = P (S 1 = s) and the transition distribution T . Specifically, ρ π (s) can be recursively written as ρ π (s) = P (s)+γ s ′ ,x T (s ′ , x, s) π(x | s ′ )ρ π (s ′ ).When the transition distribution is unconfounded (Fig. 3a ), one could recover its parametrization T (s, x, s ′ ) following the first formula of Def. 1. Therefore, what remains undetermined in Eq. ( 13) is the non-identifiable reward function R. It follows from Eq. ( 10) that parametrization of R(s, x) can be bounded from the observational distribution. The imitator's expected return could thus be lower bounded asE π [R 1 ] ≥ s,x R (s, x) P (x | s)π(x | s)ρ π (s), where R is the nominal reward function defined in Eq. ( 5). Similarly, the expert's expected return could be decomposed asE [R 1 ] = x,s R(s, x)P (x | s)ρ(s), where ρ(s) = ∞ t=0 γ t P (S t = s)is the expert's occupancy measure. Optimizing the worst-case gap between the imitator E π [R 1 ] and expert E [R 1 ] leads to a minimax optimization problem, the solution of which leads to a possible imitating policy.Algorithm 1: Causal GAIL with Confounded Reward R (CAIL-R) 1: Input: Expert demonstrations D = {(S i , X i )} N i=1 2: for iteration k = 0, 1, 2, . . . do 3:Collect expert trajectories from DCollect imitator trajectories based on the policy π k (x | s)5:Update the parameters w of discriminator D k with gradientEQUATIONUpdate the policy π k+1 = arg min π E π [P (x | s) log(1 -D(s, x))] using any forward RL algorithm 7: end for Algorithm 2: Causal GAIL with Confounded Transition T (CAIL-T )1:Input: Expert demonstrations D = {(S i , X i )} N i=1 2: for iteration k = 0, 1, 2, . . . do 3:Collect expert trajectories from DCollect imitator trajectories based on the policy π k (x | s) from the worst-case occupancy measure by solving the optimization problem presented in Eq. ( 19) and Eq. ( 20)Update the parameters w of discriminator D k with gradientEQUATIONUpdate the policy π k+1 = arg min π E π [log(1 -D(s, x)); T ] with any forward RL algorithm 7: end for Theorem 2. Given an MDP M compatible with the causal graph of Fig. 3a , let R be a parametric family containing the conditional reward E[Y t | s t , x t ]. Consider the following optimization program,EQUATIONWhen the gap ν * ≤ 0, the solution π * is an imitating policy satisfyingE π * [R 1 ] ≥ E[R 1 ].In other words, Thm. 2 computes an imitating policy within the environment depicted in Fig. 3a by finding a policy maximizing the worst-case reward function compatible with the demonstration data and the expert's nominal reward. Later in Sec. 4, we will demonstrate that such a solution exists and robust imitation learning is feasible in Fig. 3a .The optimization program in Thm. 2 could be solved by augmenting some standard imitation learning such as GAIL [19] . To make the argument more precise, let the parametric family R be a set of reward function R(s, x) taking values in the real space R. We penalize the complexity of a reward function R by subtracting a convex regularization function ψ(R) from Eq. ( 14); the detailed definition of ψ(R) is given by [19, Eq. 13] . Solving the optimization program of Eq. ( 14) is equivalent to matching weighted occupancy measures between the imitator and the expert, shown in Appendix B,EQUATION= minEQUATIONwhereψ * = max R a ⊤ R -ψ(R) is a conjugate function of ψ; function D ∈ S × X → (0, 1) is a discriminator classifier (e.g, a neural network). The above optimization problem is in the form of two neural networks competing against each other in a zero-sum game. The detailed implementation of our proposed algorithm, called CAIL-R, is provided in Alg. 1. Compared to the standard GAIL algorithm, Alg. 1 adds weight to the signal generated by the discriminator for the imitator and then attempts to match the distribution between the weighted samples and expert demonstrations.In this section, we examine the MDP obs-S environment as graphically depicted in Fig. 3b , where the reward function is unconfounded, but UCs affect the action X t and the next state S t+1 simultaneously. In this setting, the second equation of Causal Consistency (Def. 2) is satisfied, aligning the reward function R with the expert's nominal reward function. However, the first equation of Def. 2 does not generally hold due to confounding bias, making the transition distribution T not identifiable from demonstrations. Despite these challenges, we utilize partial identification techniques to bound the transition function T , and subsequently estimate the imitator's performance.More precisely, consider again the expected return decomposition in Eq. ( 13). The identifiable reward function R must be contained in the parametric space of the expert's nominal reward R. The transition distribution T can be bounded from the demonstration data using Eq. ( 9). One could thus obtain a lower bound over the imitator's performance by reasoning about the worst-case occupancy measure compatible with demonstrations. Formally, with the fixed reward function R and the fixed policy π, the imitator's return E π [R 1 ] is bounded by:EQUATIONEQUATIONThe above optimization problem is similar to the classical linear program for planning in MDPs [36] . The main difference is that the transition distribution T is no longer fixed but bounded in a convex space T specified from the observational data. Therefore, we develop an imitating policy by minimizing the performance gap between the imitator and the expert in the worst-case environment compatible with the observational data and prior knowledge.Theorem 3. Given an MDP M compatible with the causal graph of Fig. 3b , let R be a parametric family containing the conditional reward E[Y t | s t , x t ], and T be a parametric family over conditional probabilities P (s t+1 | s t , x t ) defined in Eq. (20) . Consider the following program,EQUATIONWhen the gap ν * ≤ 0, the solution π * is an imitating policy satisfyingE π * [R 1 ] ≥ E[R 1 ].We solve the optimization program in Thm. 3 by augmenting GAIL, a standard imitation method [19] . By penalizing the complexity of a reward function R using a convex regularization function ψ(R) from Eq. ( 14), Eq. ( 21) is reducible to the following distribution matching problem,EQUATIONWe present the step-by-step implementation of our imitation method, CAIL-T , in Alg. 2. It is similar to the standard GAIL [19] ; however, a significant distinction arises at step 4, where the imitator collects trajectories from the worst-case occupancy measure as presented in Eq. ( 19) and Eq. ( 20), which is obtainable by iteratively solving a series of linear programs. We refer readers to Appendix C for a more detailed discussion, where we propose an iterative algorithm designed to find the worst-case occupancy measure efficiently. 
experiments: 
 In this section, we validate the theoretical findings presented in Thm. 1 and illustrate the applications of the proposed CAIL algorithms (Alg. 1 and Alg. 2) on various causal imitation learning tasks. Such tasks range from synthetic causal models to real-world scenarios. To summarize, when both the transition and the reward are confounded, there always exists a worst-case MDP instance M compatible with the expert demonstrations, but the imitator consistently fails to match expert Fig. 4b shows the convergence plot for CAIL, GAIL, and BC performance. Fig. 4c shows the final performance, where y-axis represents the expected return.performance, aligning with the proof provided in Sec. 2.1. When either the transition or the reward is confounded, we systematically evaluate our algorithms against the standard BC and GAIL methods, highlighting the importance of optimizing within the worst-case SCM. Standard BC mimics the expert's nominal behavior policy P (X|S) via supervised learning; standard GAIL learns a policy by solving a min-max game [19] . We provide in Appendix D more details on the experiment setup.MDP obs -Random Instances. This experiment aims to empirically validate the theoretical findings discussed in Thm. 1. Consider SCM instances compatible with Fig. 2b including binary observed variables S t , X t , Y t ∈ {0, 1}. 1000 random discrete MDPs are sampled, in other words, the reward functions and the transition probabilities are generally different among these models. The expert is able to observe the state S t , the unobserved variable U t . However, the imitator, lacking access to both U t or the reward E π [Y t ], makes decisions solely on S t . As shown in Fig. 4a , imitators consistently failed to match expert performance. Specifically, prevalent negative performance gaps indicate that most of imitators were consistently worse than experts; only in rare cases did the performance gaps near -0.5, supporting our theoretical insights in Thm. 1. In summary, imitators fail to achieve the expert's performance when both the reward and the transition are confounded.MDP obs-Y -Driving. To demonstrate the proposed framework, as outlined in Alg. 1, we consider a scenario when an autonomous vehicle ('ego vehicle') aims to learn optimal driving strategies from expert demonstrations. The state S t contains some critical driving information, e.g., the velocities of the ego vehicle and the leading vehicle and the spatial distance between them. The action X t represents acceleration or deceleration decisions the ego vehicle makes. The unobserved variable U t represents some information accessible to the expert but inaccessible to the imitator, e.g. slippery road conditions [26] . The reward Y t is designed to reflect multiple realistic driving objectives, e.g., safety, comfort, efficiency, and so on. U t has an effect on the reward Y t . Unlike the scenarios described in [39, 42] , due to UCs between X t and Y t at each step t, it is impossible to find a π-backdoor admissible set. BC, GAIL, and CAIL utilize the same policy space π(x | s). The major difference between CAIL and GAIL lies in that CAIL optimizes the imitator by the weighted reward generated from the discriminator -P (x | s) log(1 -D(s, x)). As illustrated in Fig. 4b , where means and standard deviations are computed over 100 trajectories, CAIL consistently outperforms BC and GAIL.MDP obs-S -Medical Treatment. Consider the challenge of providing medical treatment to acutely ill patients, where the primary goal is to learn a policy so that the morality rate can be decreased. We utilize the real-world medical treatment dataset, i.e., Medical Information Mart for Intensive Care III (MIMIC-III) dataset [22] . MIMIC-III consists trajectories of clinical information (e.g., heart rate, oxygen saturation, and so on) recorded at various time intervals. However, due to privacy concerns, certain essential variables are masked or not properly recorded [45] , e.g., socioeconomic status or the experience levels of caregivers [9, 56] . Specifically, the state S t encapsulates the critical health information for the patients, e.g., prolonged elevated heart rate (peHR). The action X t represents whether to treat the medicine or not. The reward Y t is designed to represent the intent of the doctor as much as possible, e.g., avoiding the patient's mortality. The unobserved confounded U t simultaneously affects the action X t and the next state S t+1 . Simulation results are illustrated in Fig. 4c , which shows that the proposed framework performs the best among all strategies. BC and IRL fail to obtain an imitating policy that could match expert performance. 
conclusion: 
 This paper investigates imitation learning in Markov Decision Processes where the unobserved confounding bias cannot be ruled out a priori. We establish theoretically that when such unobserved confounders generally exist, it is infeasible to obtain a robust imitating policy that can perform at least as well as the expert across all possible environments compatible with the demonstration data and prior knowledge. Departing from this critical realization, our research diverges into two distinct problem settings -one where only the transition distribution is unconfounded, but the reward function is non-identifiable due to unobserved confounding; and the other where the reward function is unconfounded and the transition distribution is non-identifiable. We then propose novel imitation learning algorithms using partial identification techniques, which allow the imitator to obtain effective policies that can achieve expert performance for both problem settings. Through extensive experiments, we empirically validate the theoretical findings and systematically evaluate our algorithms on different scenarios, ranging from simulated causal models to real-world datasets. 
experimental setting/details: 
 Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes]Justification: All the important training and test details in this paper can be found Sec. 4 and Appendix D.Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 
experiment statistical significance: 
 Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [Yes]Justification: Error bars and other appropriate information about the statistical significance the experiments could be found in Sec. 4 and Appendix D.Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors). • It should be clear whether the error bar is the standard deviation or the standard error of the mean. • It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified. • For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). • If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [Yes] Justification: Detailed information on the computer resources can be found in Sec. 4 and Appendix D.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [Yes] Justification: Societal impacts are discussed in Sec. 1 and Appendix E. Our framework can be applied to various fields in reality, including autonomous driving, robotics, industrial automation, medical decisions modeling and so on. One of positive impacts of this work is that we discuss the potential risk of training IRL algorithms when demonstrations are contaminated by unobserved confounders, and how to utilize partial identification techniques to make the imitator robust.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: In this section, we provide proofs for the theoretical claims delineated in the paper. Throughout this paper, it is important to note that detailed parametrizations of the underlying SCM are not known to the agent. Instead, the agent has access to the expert's demonstrations, which are summarized as the observational distribution P (X, S, Y ).We begin by revisiting the distribution of state visitation. Specifically, ρ π (s) can be calculated by:EQUATIONwhere P (s) represents the initial state distribution, γ represents the discount factor, T represents the transition probabilities for the imitator. Subsequently, we are able to develop the occupancy measure for the policy π:EQUATIONIt is important to note that, although the format of the occupancy measure ρ π (s, x) shares a formal resemblance to the one presented in GAIL [19] , ρ π (s, x) specifically represents an interventional distribution with policy do(π). The identifiability of the transition T (s, x, s ′ ) directly impacts the identifiability of P π (s t ). If P π (s t ) is not identifiable, ρ π (s) and ρ π (s, x) are consequently not identifiable.Theorem 1. Given any positive observational distribution P (X, S, Y ) > 0, there exists an MDP model M compatible with the causal graph of Fig. 2b such that P (X, S, Y ; M ) = P (X, S, Y ) and for any policy π, any time step t = 1, 2, . . . , any state s ∈ S,EQUATIONProof. Without loss of generality, the reward Y is normalized so that it has a range of [0, 1] Based on the value function defined in Eq. ( 8), we first show how to expand it into a recursive version:EQUATIONEQUATION)where γ is the discount factor, P π (s t+1 | s t ) denotes the transition probability when executing policy π.From the second last line to the last line is justified by the experimental markovian property, as discussed in Sec. 2, following the graph Fig. 2a . More details could be found in [54] .E π [Y t | s t ] = E[Y t | s t , do(π)] denotes the expected reward received by the agent when executing policy π.Similarly, the transition probabilityEQUATIONand P xt (s t+1 | s t ) = P (s t+1 | s t , do(x t )) = T (s t , x t , s t+1 ). Generally speaking, when any unobserved confounder exists between S t+1 and X t , the causal query P xt (s t+1 | s t ) is not identifiable [32, 43, 5, 51] . Building on the previous derivations, we arrive at the recursive formulation of the value function under policy π:EQUATIONNext, to establish the validity of the preceding claim, we proceed by applying the technique of mathematical induction. Let |S| denote the number of distinct states for S.Base case t = T . For the final timestep T , for each state index j where ∀j, 1 ≤ j ≤ |S|, the value function V π (s (T,j) ) cane be defined as follows:EQUATIONwhere s (T,j) refers to the scenario where the state at the final timestep S T is equal the specific state j.In order to obtain the worst-case SCM M , we need to minimize V π (s T ) -V (s T ) compatible with the observational distribution, by establishing its lower bound. To this end, we directly employ the natural bound [29] , which has been discussed in Sec. 2.1:min M V π (s (T,j) ; M ) -V (s (T,j) ; M ) = xt E xt Y T | S T = s (T,j) ; M π(x t | s (T,j) ) -V (s (T,j) ; M ) = xt E Y T | s (T,j) , x t P (x t | s (T,j) )π(X T = x t | s (T,j) ) - xt E Y T | s (T,j) , x t P (x t | s (T,j) )< 0 (34) The last step is justified because P (X, S, Y ) > 0 and 0 ≤ π(X T = x t | s (T,j) ) ≤ 1. Intuitive examples illustrating this conclusion are provided in Sec. 2.1 and Appendix F. Therefore, this confirms the validity of the inequality for the base case. Specifically, in certain degenerate cases where there is only one possible action, the imitator has no choice but to follow that single option. Consequently, unobserved confounders are less likely to introduce significant effects in these scenarios. However, under such conditions, pursuing imitation learning is not meaningful, as there is no variability in choice for the imitator to learn from. Therefore, such cases are of limited relevance to the scope of this analysis.Induction case. Suppose at t + 1, V π (s t+1 ) < V (s t+1 ), we need to prove V π (s t ) < V (s t ). V π (s t ) = E π [Y t | s t ] + γ j P π (s (t+1,j) | s t ) V π (s (t+1,j) )<V (s (t+1,j) ) (35) Without loss of generality, we assume that the state with the minimal value at t + 1 is denoted as s (t+1,|S|) . Our approach is founded on the premise that in obtaining the worst-case SCM M , it is strategic to allocate the lowest possible transition probabilities to the state with the highest value, while preferentially assigning higher probabilities to states demonstrating smaller values. Specifically, one starts with the estimate P (S t+1 = s (t+1,1) , x t | S t ) for P xt (S t+1 = s (t+1,1) | S t ). Following this logic, we systematically allocate probability masses for indices 1 ≤ j ≤ |S| -1 as follows:P xt (S t+1 = s (t+1,j) | S t ) ← P (S t+1 = s (t+1,j) , x t | S t )In accordance with the established properties of probability distributions, it follows that:|S| j=1 P xt (S t+1 = s (t+1,j) | S t ) = 1.Considering the state s (t+1,|S|) , the corresponding probability can be assigned as:P xt (S t+1 = s (t+1,|S|) | S t ) = 1 - |S|-1 j=1 P xt (S t+1 = s (t+1,j) | S t ).By substituting the assigned values, we are able to derive the following expression:P xt (S t+1 = s (t+1,|S|) | S t ) ← 1 -P (s (t+1,1) , x t | S t ) -P (s (t+1,2) , x t | S t ) • • • -P (s (t+1,|S|-1) , x t | S t ),where the right-hand side simplifies to:P (s (t+1,1) , x t | S t ) + P (s (t+1,2) , x t | S t ) • • • + P (s (t+1,|S|-1) , x t | S t ) =   |S|-1 j=1 P s (t+1,j) , x t | S t   = P (x t | S t ) -P (s (t+1,|S|) , x t | S t ).It is established that the expression 0 ≤ 1 -P (x t | S t ) + P (s (t+1,|S|) , x t | S t ) ≤ 1 holds true. This inequality is supported by the following equation:|S| j=1 P (s (t+1,j) , x t | S t ) = P (x t | S t ).To further analyze the expert policy, the associated value function V (s t ) can be expanded as follows:EQUATION)where γ is the discount factor, P (s (t+1,j) | s t ) denotes the observational transition probability. Notably, P (s (t+1,j) | s t ) and P π (s (t+1,j) | s t ) are generally different, because they reflect two distinct probabilities: P (s (t+1,j) | s t ) corresponding to the observational distribution and the other, P π (s (t+1,j) | s t ), representing the imitator's transition dynamics. In accordance with the established properties of probability distributions, it follows that:|S| j=1 P (S t+1 = s (t+1,j) | S t ) = 1.Without loss of generality, suppose the policy is a deterministic policy. Actually, the following proof holds true regardless of the choice of x t . Subsequently, we analyze the gap between V π (s t ) and V (s t ) as follows:V π (s t ) -V (s t ) =   E π [Y t | s t ] + γ |S| j=1 P π (s (t+1,j) | s t )V π (s (t+1,j) )   -   E[Y t | s t ] + γ |S| j=1 P (s (t+1,j) | s t )V (s (t+1,j) )   =   E π [Y t | s t ] + γ |S| j=1 xt P xt (s (t+1,j) | s t )π(x t | s t )V π (s (t+1,j) )   -   E[Y t | s t ] + γ |S| j=1 P (s (t+1,j) | s t )V (s (t+1,j) )   (37) min M V π (s t ; M ) -V (s t ; M ) = E π [Y t | s t ; M ] -E[Y t | s t ; M ] + γ |S|-1 j=1 P (s (t+1,j) , x t | s t ) -P (s (t+1,j) | s t ) V (s (t+1,j) ) + γ   1 -   |S|-1 j=1 P s (t+1,j) , x t | s t   -P (s (t+1,|S|) | s t )   V (s (t+1,|S|) ) = E π [Y t | s t ; M ] -E[Y t | s t ; M ] <0 + γ |S|-1 j=1 P (s (t+1,j) , x t | s t ) -P (s (t+1,j) | s t ) <0 V (s (t+1,j) ) -V (s (t+1,|S|) ) >0 < 0 (38) where min M E π [Y t | s t ; M ] -E[Y t | s t ; M ]< 0 follows a similar logic as previously introduced in the base case, and V (s (t+1,j) ) -V (s (t+1,|S|) ) > 0 is consistent with the ordering assumption, where the state s (t+1,|S|) represents the minimal value.In some degenerated cases when E π [Y t | s t ] = 0 and E[Y t | s t ] = 0, it might coincidentally follow that V π (s t ) = 0, which is equal to V (s t ) = 0. Another instance of degeneracy occurs when the value function V (s (t+1,j) ) remains the same across all states. Such occurrences are extremely unlikely in practical scenarios, especially when P (X, S, Y ) > 0. EQUATIONWhen the gap ν * ≤ 0, the solution π * is an imitating policy satisfyingE π * [R 1 ] ≥ E[R 1 ].Proof. Based on Eq. ( 13), we have:E π [R 1 ] = s,x R (s, x) π(x | s)ρ π (s)It follows from Eq. ( 10) that parametrization of R(s, x) can be bound from the observational distribution. The imitator's expected return could thus be lower bounded asE π [R 1 ] ≥ s,x R (s, x) P (x | s)π(x | s)ρ π (s)Note that the expert's expected return could be similarly decomposed asEQUATION= minEQUATIONwhich is the ultimate target expression.Next, we will show the derivation details for matching weighted occupancy measures between the imitator and the expert. Supposeψ * = max R a ⊤ R -ψ(R) is a conjugate function of ψ.Following a similar logic in [19] , we utilize a smiliar cost regularizer ψ GA , leading to the formulation of Alg. 1. Basically, Alg. 1 minimizes Jensen-Shannon divergence between P (x | s)ρ(s) andP (x | s)π(x | s)ρ π (s).First, we reformulate the equation into state-action occupancy measures:EQUATIONBased on the definition of ψ * , we have:EQUATIONEQUATIONEQUATIONwhere we make the change of variables R → -ϕ(R ′ ). Suppose D ∈ S × X → (0, 1) is a discriminator classifier (e.g, a neural network). Using the logistic loss ϕ(x) = log (1 + e -x ), we can get:EQUATIONEQUATION= maxEQUATIONwhich is the ultimate target expression.Theorem 3. Given an MDP M compatible with the causal graph of Fig. 3b , let R be a parametric family containing the conditional reward E[Y t | s t , x t ], and T be a parametric family over conditional probabilities P (s t+1 | s t , x t ) defined in Eq. (20) . Consider the following program,EQUATIONWhen the gap ν * ≤ 0, the solution π * is an imitating policy satisfyingE π * [R 1 ] ≥ E[R 1 ].Proof. Based on Eq. ( 13), we have:E π [R 1 ] = s,x R (s, x) π(x | s)ρ π (s) = s,x R (s, x) ρ π (s, x) Non-IDThe reward function R is identifiable and must be contained in the parametric space of the expert's nominal reward R. In other words,EQUATIONThe transition distribution T can be bounded from the demonstration data using Eq. ( 9). Therefore, we get:EQUATION= minEQUATION= minEQUATION= minEQUATIONwhich is the ultimate desired expression.Consider again the expected return decomposition in Eq. ( 13). The reward function R is identifiable and must be contained in the parametric space of the expert's nominal reward R. The transition distribution T can be bounded from the demonstration data using Eq. ( 9). One could thus obtain a lower bound over the imitator's performance by reasoning about the worst-case occupancy measure compatible with demonstrations. Formally, with the fixed reward function R and the fixed policy π, the imitator's return is bounded byEQUATIONsubject to: ρ π (s) ≥ 0, andEQUATIONThe above optimization problem is similar to the classic linear program for planning in MDPs [36] .The main difference is that the transition distribution T is no longer fixed but bounded in a convex space T specified from the observational data. Similar to the previous setting, we could solve an imitating policy by minimizing the performance gap between the imitator and the expert in the worst-case environment compatible with the observational data and prior knowledge.Next, we will provide a heuristic algorithm to solve the optimization program presented in Eq. ( 19) and Eq. (20) . Specifically, as discussed in Eq. ( 9), we are able to bound the transition distribution T by:EQUATIONThe intuition for Alg. 3 is: in order to find the worst case, we need to put as less transition probability mass as possible to the state with maximal values, and allocate higher transition probabilities to states with smaller values. Without loss of generality, suppose V xt (s (t+1,|S|) ) is found to have the smallest relative value. For all other states j ̸ = |S|, we need to allocate as less transition probability mass as possible. Therefore, we take the lower bound:EQUATIONEQUATIONwhereP xt (s t+1 | s t ) = P (s t+1 | s t , do(x t )) = T (s t , x t , s t+1 ), and P (s (t+1) | s t , x t ) =T (s t , x t , s t+1 ). For the state s (t+1,|S|) , we have:EQUATIONAlgorithm 3: Find Worst-Case Discounted Future Reward P xt (S t+1 = s (t+1,|S|) | s t ) := 1 -   |S|-1 j=1 P (S t+1 = s (t+1,j) , x t | s t )   5: returnFollowing a similar logic in Alg. 1: we reformulate the equation into state-action occupancy measures:EQUATIONBased on the definition of ψ * , we have:EQUATION= maxEQUATIONEQUATIONSuppose D ∈ S × X → (0, 1) is a discriminator classifier (e.g, a neural network). Using the logistic loss ϕ(x) = log (1 + e -x ), we can get:EQUATIONTherefore, we are able to obtain the ultimate target expression:EQUATIONIn this section, we provide a practical algorithm, Alg. 3, designed to solve the optimization problem formulated in Eq. ( 19) and Eq. (20) . The underlying rationale of Alg. 3 is to search for the worst-case scenario by allocating the minimal transition probability mass to the state with the highest value while assigning greater transition probabilities to states with lower values. The resulting solution should still adhere to a set of predefined observational constraints to ensure feasibility. This approach ensures that the most "adversarial" outcome is prioritized during the optimization process.To further clarify the approach above, consider the following numerical example. Suppose there are only two states. non-identifiability issue caused by the transition distribution P π1 (S 2 | S 1 ), as discussed in Eq. ( 9), we formulate the worst-case SCM by allocating f (S 2 = 0) with probability mass P (S 2 = 0, X 1 | S 1 ).In other words, we assign the lower bound P (S 2 = 0, x 1 | Z 1 ) to the non-identifiable query P x1 (S 2 = 0 | Z 1 ). As such, we are able to rewrite the expert's rewards as follows:EQUATION+ f (S 2 = 0) * P (S 2 = 0, X 1 = 1|Z 1 )P (Z 1 ) (76)EQUATIONEQUATIONand the imitator's reward can be written asEQUATIONEQUATIONEQUATIONEQUATIONwhere is E π1,π2 EQUATIONIn other words, the imitator is unable to learn a policy that can obtain the expert's performance in the worst-case 2-stage MDP compatible with the observational distribution P(X 1 , X 2 , S 1 , S 2 , Y 1 , Y 2 ).NeurIPS Paper ChecklistThe decomposition holds since state St blocks all backdoor paths from action Xt to nodes St+1 and Yt, i.e., path starting with arrow Xt ← St. It follows from Rule 2 of do-calculus[32, Theorem3.4.1] that Pπ (yt | st, xt) = Px t (st+1 | st) and Pπ (yt | st, xt) = Px t (yt | st).In this case, although Px t (s (t+1,2) | st) has a lower bound of P (s (t+1,2) , xt | st), it cannot be set exactly equal to this value. It is crucial to maintain the condition Px t (s (t+1,1) | st) + Px t (s (t+1,2) | st) = 1. 


Decision: The paper studies imitation learning for MDPs. The focus is on using causal methods to battle confounding. 

The contributions are:
- an impossibility result about unobserved confounders (Theorem 1)
- a result about how to construct a policy when the transition dynamics are identifiable but the reward is not (Theorem 2)
- a result about how to construct a policy when the reward is identifiable but the transition is not (Theorem 3)

While reviewers have voiced concerns, the meta-reviewer disagrees with their assessment. Specifically:
- the lack of extensive experiments shouldn't be a main reason to reject a theory paper.
- saying it's hard to validate assumptions in practice is a vacuous statement, which applies to (almost) all theory papers, especially to causality papers
- saying the paper isn't practical because it involves constrained optimisation is unfair - constrained optimisation is required by many NeurIPS papers.
- criticisms of reproducibility are exaggerated given the paper is theoretical and authors promised to release code upon acceptance
- some said the paper needs more baselines, but didn't provide what baselines to use (and an earlier suggestion turned out irrelevant)
- similarity in algorithms doesn't mean the theory is trivial. The algorithms are different from GAIL in ways which are crucial for battling confounding.

This paper is a solid contribution and pushes on what is possible in causal imitation learning. For this reason, I recommend acceptance.

Reviews: review: paper_summary: paper_summary: The paper addresses the challenges in imitation learning when expert demonstrations are contaminated with unobserved confounding bias. It proposes robust imitation learning methods within the framework of MDPs using partial identification techniques. The authors demonstrate theoretically that when unobserved confounders exist, learning a robust policy to achieve expert performance is generally infeasible. They introduce two novel causal imitation algorithms to handle settings where either the transition distribution or the reward function is non-identifiable from the available data. The proposed methods are validated through experiments, showing their effectiveness in achieving expert-like performance in the presence of unobserved confounders.
strength_weakness: strength_weakness: The experiment evaluation seems insufficient considering that there are already some causal imitation learning baselines such as [1], and some more experiment environments such as OpenAI gym for imitation learning. 

[1] P. de Haan, D. Jayaraman, and S. Levine. Causal confusion in imitation learning. In Advances in Neural Information Processing Systems, pages 11693–11704, 2019.
questions: questions: On what tasks are these approaches experimented on, what do these tasks look like, and how does the demonstration data look like? Will the model performance change with a varying number of expert demonstrations?
limitations: limitations: The limitations have been discussed in the paper.
paper_summary: paper_summary: The paper addresses challenges in imitation learning when the learner and expert have mismatched sensory capabilities and demonstrations are contaminated with unobserved confounding bias. The authors propose robust imitation learning within the framework of Markov Decision Processes (MDPs) using partial identification. They demonstrate that in the presence of unobserved confounders, learning a policy that guarantees expert performance is generally infeasible. The paper introduces two novel algorithms for imitation learning in partially identifiable settings—when either the transition distribution or the reward function is non-identifiable. These algorithms, based on augmentations of the Generative Adversarial Imitation Learning (GAIL) method, are designed to achieve expert-level performance even with confounded data.
strength_weakness: strength_weakness: 1. **Assumptions and Generalization**: The approach relies on certain assumptions about the partial identifiability of the MDP. In practice, these assumptions might not always hold. Even if these assumptions hold, one may not know the type of partial identification in advance, potentially limiting the generalizability of the methods.
2. **Complexity and Practicality**: The proposed method CAIL-$\mathcal{T}$ requires solving a complex constrained optimization problem in Eq. (18) and Eq. (19). It is difficult to solve this optimization problem in practice when neural networks are employed.
3. **Missing Experimental Details:** The paper does not provide details about the implementation of the proposed algorithm and the environments. Besides, the source code is also missing. As such, it is difficult to reproduce their experiments.
questions: questions: Please see the detailed review in the Weakness part.
limitations: limitations: The authors have thoroughly discussed the limitations of this paper.
paper_summary: paper_summary: This paper studies causal imitation learning from the perspective of partial identification. First, the authors show a hardness result that when both the transitions and the rewards are confounded, it is not possible to imitate or improve over the expert policy. Going forward, under only reward-confounding or transition-confounding, a new imitation learning objective is proposed using the partial identification bound of the corresponding non-identifiable unknowns, and it is proved that once the objective is optimized under $0$, the learned imitation policy would improve over the expert policy. Experimental result justifies the theoretical findings.
strength_weakness: strength_weakness: 1. The idea seems like a relatively direct extension of the standard imitation algorithm (e.g., GAIL) with the partial identification method from the causal inference literature.
2. Theoretically, it is still unknown when can we guarantee that the imitator can beat the expert. According to Theorem 2 (resp. Theorem 3), only when the optimization problem (14) (resp. (21)) has solutions with non-positive value would the imitator policy be guaranteed to achieve the same performance or improve over the expert policy. Is it possible to derive hardness results that in the reward-confounding-only case (or the transition-confounding-only case) there still exists instances such that it is impossible to improve over the expert?
questions: questions: 1. The derivation of the partial identification bound, which is central to the method proposed in this work, is missing. I suggest clearly include them in the appendix to make the argument self-content and more convincing, or at least make it clear how the conclusion is made given the existing literature like [27].
2. How to principally specify the (observational) reward function class $\mathcal{R}$​ under the unobserved confounding? That is, since the underlying SCM is not known to the agent a priori, how to equip the learning agent a suitable function class $\mathcal{R}$ to match the observational reward function of the unknown SCM? 
3. When the algorithm is applied to a non-confounded MDP (i.e., a standard MDP), how well would the algorithm perform when compared with prior arts designed for non-confounded MDPs? More generally, it seems unknown that how the performance of the proposed algorithms would change w.r.t. the confoundedness of the underlying MDP, e.g., how far the interventional probabilities deviate from the observational probabilities.
4. In Theorem 1, according to the proof, it seems essential to assume that *all* the observational probabilities $P(X,S,Y)$ are positive, which seems too strong an assumption. What if one drop this assumption? Does the conclusion change drastically?
5. Some minor typos:
   - In the last line of the footnote in the end of page 3, the first equation should be $P_{\pi}(s_{t+1}|s_t,x_t) = P_{x_t}(s_{t+1}|s_t)$.
   - The imitator's policy $\pi$ appearing in the subscript of $\rho$ is in a bold font but sometimes not, which is not consistent.
limitations: limitations: Please see the weakness section and the question section above.

Comments: comment: We thank the reviewer for the response. Since the discussion phase is coming to an end, we would like to summarize below the logic behind our model assumptions and how they compare to standard imitation learning methods.

1. Standard imitation learning: Markov Property (Definition 1) + Unconofundedness (Definition 2);
2. Our approach: Markov Property (Definition 1).

As described above, our model only relaxes existing assumptions in the standard setting and, therefore, is more general, subsuming the standard imitation learning. When the imitator has strong knowledge about the MDP environment and is confident about unconfoundedness, it should follow the standard imitation algorithms and obtain a solution. When the imitator is unsure about the unconfoundedness assumption, it could still apply our method and obtain a robust policy with a performance guarantee. Given the wide applications of existing imitation learning methods and the relaxation of the critical assumption of unconfoundedness, our results could be generalized to many practical domains.
> _"More broadly, it is still unknown when can we guarantee that the imitator given by the algorithm can beat the expert, theoretically."_

This condition depends on the quality of the demonstration data and the prior knowledge of the latent reward. For example, when there is no unobserved confounding, one could show that the imitator is guaranteed to perform at least as well as the expert. We acknowledge that deriving closed-form solutions for the improvement condition is an exciting problem. However, we would like to note that our proposed algorithms do provide a numerical condition for policy improvement, similar to the standard inverse RL methods like GAIL. When our proposed algorithm returns the game $\nu^* < 0$, the learned policy is robust and guaranteed to outperform the expert. One could find such instances in Figure 4 (b-c).
Thank you to the authors for the detailed responses. I do recognize the differences in problem settings between this work and previous studies. However, I still believe it is necessary to include additional baselines, especially since the underlying SCM is not known in advance to the agent. Adding more baselines and discussing their relevance will enhance our understanding of how effectively the proposed approaches perform.
Dear AC,

Thank you for initiating this discussion. I acknowledge the differences in problem settings between this work and prior efforts (e.g., De Haan et al., 2019). However, I maintain that including additional baselines is essential, given that the underlying structural causal model is unknown a priori to the agent. Incorporating more baselines and discussing their relevance will provide a clearer understanding of the performance of the proposed approaches.
Thank you very much for your detailed responses!  However, I still remain concerned with the soundness of the proposed algorithm. For example, while I agree with the author that the reward function class is to approximate the observed rewards, I am concerned with how the unobserved hidden (unknown) causal structure would affect the performance. More broadly, it is still unknown when can we guarantee that the imitator given by the algorithm can beat the expert, theoretically. This is less discussed even under a suitable theoretical assumption. Given such concerns, I would maintain my score and still recommend rejection of this paper. I suggest further refinement of the work for publication.
Thanks a lot for the detailed responses.

Regarding the first question, it seems that the response from the authors does not address my concern. My concern here is that, in practice, we typically do not know in advance which type of assumption the underlying task satisfies. For instance, whether the underlying reward function is identifiable or not. In that case, it remains unclear how to choose their methods to solve the task.
Thanks for taking the time to read the paper and your reviews. We recognize the importance of integrating a theoretical framework based on causality into IRL, especially when either the transition dynamics, the reward function, or both, are confounded. Below, we provide further clarification on our contributions, theoretical contributions and the experimental design to enhance understanding and address your concerns effectively.