Deconfounding Imitation Learning with Variational Inference
Authors: Risto Vuorio
Pim De Haan
Johann Brehmer
Hanno Ackermann
Daniel Dijkman
Taco Cohen
Venue: Accepted by TMLR ()

Abstract: Standard imitation learning can fail when the expert demonstrators have different sensory inputs than the imitating agent. This is because partial observability gives rise to hidden confounders in the causal graph. In previous work, to work around the confounding problem, policies have been trained using query access to the expert’s policy or inverse reinforcement learning (IRL). However, both approaches have drawbacks as the expert’s policy may not be available and IRL can be unstable in practice. Instead, we propose to train a variational inference model to infer the expert’s latent information and use it to train a latent-conditional policy. We prove that using this method, under strong assumptions, the identification of the correct imitation learning policy is theoretically possible from expert demonstrations alone. In practice, we focus on a setting with less strong assumptions where we use exploration data for learning the inference model. We show in theory and practice that this algorithm converges to the correct interventional policy, solves the confounding issue, and can under certain assumptions achieve an asymptotically optimal imitation performance.

Published in Transactions on Machine Learning Research (09/2024)
Deconfounding Imitation Learning with Variational Inference
Risto Yuorio-f:12 1.2, Pim de Haan 2.3, Johann Brehmer', Hanno Ackermann', Daniel Dijkman', and Taco Cohen
Univeesiity of Oxford
?Qualcomm AI Research. Qualcomm Al Research an initiative of Qualcomm Technologies, Inc. BQUVA Lab, University of Amsterdam
Reviewed on OpenReview: http::///paaareeessseeeeeerrrrrs
Abstract
Standard imitation learning can fail when the expert demonstrators have different sensory imputs than the imitating agent. This is because partial observability gives rise to hidden confounders in the causal graph. Previously, to work around the confounding problem, policies have been trained by accessing the expert's policy or using inverse reinforcement learning (IRL). However, both approaches have drawbacks as the expert's policy may not be available and IRL can be unstable in practice. Instead, we propose to train a variational in- ference model to infer the expert's latent information and use it to train a latent-condiitonall policy. We prove that using this method, under strong assumptions, the identification of the correct imitation learning policy is theoretically possible from expert demonst rations alone. In practice, we focus on a setting with less strong assumptions where we use exploration data for Icarning the inference model. We show in theory and practice that this algorithm converges the correct interventional policy, solves the confounding issue, and can under certain assumptions achieve asymptotically optimal imitation performance.
1 Introduction
Successful training of policies via behavioral cloning (BC) requires a high quality expert dataset. The conditions under which the dataset is collected have to exactly match those encountered by the imitator. Sometimes such data collection may not be feasible. For example, imagine collecting data from human drivers for training a driving policy for a self-driving car. The drivers are aware of the weather forecast and lower their speed in icy conditions, even when the ice is not visible on the road. However, the driving policy in the self-driving car does not have access to the same forecast, it is unwware of the ice on the road and may thus be unable to adapt to the conditions. In this paper, we focus on such imitation learning in settings where the expert knows more about the world than the imitator.
4n
68
61
(a) Expert
(b) Ignorant policy
Figure 1: Bayes nets for (a) an expert trajectory and (b) an imitator trajectory. The expert action depends on the latent variable (red arrows) whereas the imi- tator action does not.
"Equal contribution.
1 *Corresponding author: risto. vuorio@gnail com Work done during internship Qualcomm AT Research.

---

Published in Transactions on Machine Learning Research (09/2024)
The information the expert observes about the world but the imitator does not can be modeled a latent variable that affects the dynamics of the environment. Tmitating expert who observes the latent variable, with a policy that does not, results in the learned policy marginalizing its uncertainty about the latent variable. This can result in poor performance, as the agent always acts randomly in states where the expert acts on its knowledge about. the latent. To fix this, we assume that the imitator policy is dependent on the entire history of interaction with the environment instead of just the current observation. This enables the agent to make an inference about the value of the latent variable based on everything it has observed so far, which eventually allows it to break ties over the actions in states where the latent variable matters. However, this introduces another problem where the latent variable of the environment acts as causal confounder in the graph modeling the interaction of the agent with its environment as illustrated in Figure
In such a situation, an imitating agent may take their own past actions as evidence for the values of the confounder. The self-driving car, for instance, could conclude that as it is driving fast there can be no ice on the road. This issue of causal delusion was first pointed out in Ortega Braun (2010a b) and studied in more depth by Ortega et al. (2021). This problem can be a source of "delusions" in generative models including large language models Ortega et al. (2021). Instead of general generative modeling, we focus a control sctting where we have access to the truc environment where the agent is going to be deployed.
Ortega et al. 2021 | show that using the classic dataset aggregation (DAgger) algorithm (Ross et al., 2011 solves this problem by querying the expert policy in the new situations the agent encounters to provide new supervision for learning the correct behavior. However, this kind of query access to the expert may not be available in practice. As another solution the confounded imitation learning problem, [wamyy et al.) (2022b) propose to use inverse reinforcement learning (IRL) (Russell 1998
Ng et al. 2000). However, IRL. typically requires adversarial methods (Ho & Ermon 2016 Fu et al. 2017 which are not well behaved
and scalable as supervised learning.
To get around the confounding problem without query access to the expert or IRL, we propose a practical algorithm based on variational inference. In our approach, an inference model for the latent variable is learned from exploration data collected using the imitator's policy. This inference model is used for inferring the latent variable on the expert trajectories for training the imitator and for inferring the latent variable online when the imitator is deployed in the environment. We show that, in theory, this algorithm converges to the expert's policy. In other words, we show that the expert's 's policy can be identified. Furthermore we show that under strong assumptions this identification can be carried out purely from offline data. Finally, we validate its performance empirically in a set of confounded imitation learning problems with high-dinensional o observation and action spaces. These contributions can be summarized as follows:
We propose practical method based on variational inference to address confounded IT. without expert queries. We verify its performance empirically in various confounded control problems, out- performing naive BC and IRI.
We propose theoretical method for confounded IL, purely from offline data, without expert queries nor exploration.
We provide theoretical insight into the proposed methods by proving under strong assumptions that the expert's policy can be identified.
2 Related work
Imitation learning Learning from demonstration has long history with applications in autonomous
driving (Pomerlcau) 1988 Lu et al., 2022) and robotics (Schaal 1999 Padalkar et al. 2023 Standard
algorithrs include BC, IRL Russell
1998 Ng et al. 2000 Ziebart et al. 2008), and ad adversarial methods (Ho
& Ermon 2016 Fu et al. 2017).
Imitation lcarning can suffer from a mismatch between the distributions faced by the expert and imitator due to the accumulation of errors when rolling out the imitator policy. This is commonly addressed by querying experts during the training (Ross et al. 2011) by noise insertion in the expert actions Laskey al. 2017
2

---

Published in Transactions on Machine Learning Research (09/2024)
Note that this issue is qualitatively different from the one we discuss in the paper: it is consequence of the limited support of the expert actions and occurs even in the absence of the latent confounders.
Kumar et al. (2021) and Peng et al. (2020) consider a similar setting to ours where the environment has a latent variable, which explains the dynamics. They use privileged information to learn the dynamics encoder
via supervised learning. In contrast, only the expert has access to privileged information in our setting.
Causality-aware imitation learning Ortega Braun (2010a by and Ortega et al. (2021) pointed out the issue of latent confounding and causal delusions that we discuss in this paper. In particular, Ortega al.) (2021) propose training algorithm that learns the correct interventional policy. Unlike our algorithm, their approach requires querying experts during the training. However, we discuss in section 4 has weaker assumptions and also applies to non- Mon-Markovian dynamics.
their solution
Most similar to our work is Swamy et al. 2022b). which finds theoretical bounds on the gap between expert behavior and an imitator in the confounded setting, when imitating via BC, DAgger (Ross et al. 2011), which requires expert queries, or inverse RL. Inverse RL suffers from two key challenges: (1) requires reinforcement learning in inner loop with a non-stationary reward, and (2) the reward is typically learned via potentially unstable adversarial methods (Ho & Ermon 2016; Fu et al. 2017). In contrast, our met trains the behavior policy using well-behaved BC objective, which often enjoys better robustness and scalability compared to inverse RL.
There are various other works on the intersection of causality and IL which differ in setup, De Haan et al. (2019) consider the confusion of an imitator when the expert's decisions can be explained from causal and non-causal features of the state. It differs from our work they assume the state to be fully observed, meaning it does not apply to situation in which there is latent confounder which needs to be inferred. This problem has been also discussed in Codevilla et al. (2019); Wen et al. (2020) 2022) and Spencer et al. 2001)) under various names. Rezende et al. (2020) point out that the same problem appears in partial models that only use a subset, of the state and find a minimal set of variables that avoid confounding. Swamy et al. (2022a) consider imitation learning with latent variables that affect the expert policy, but not the state dynamics, which is different from our case, in which the latent affects both the state and the expert's actions. Kumor et al. (2021) study the case in which a graphical model of the partially observed state is known and find which variables can be adjusted for that conditional BC is optimal. An extension, Ruan et al. (2022), also handles sub-optimal experts.
Meta-lcarning bchaviors Our problem is related to meta-IL (Duan et al. 2017 Beck et al. 2023 where the aim is to train an imitation learning agent that can adapt to new demonstration data elficiently. Differently from our problem, the tasks in meta-IL can vary in the reward function. For imitation learning to work with the new reward functions, demonstrations of policies maximizing the new rewards are required. While meta-IL also considers a distribution of MDPs, the motivations and methods are different from ours. Our work does not consider adapting to new demonstrations, whereas meta-IL, does not consider the confounding problem in the demonstrations.
Furthermore, our problem is related to meta-reinforcement learning (RL) @Duan et al. 2016 Wang et al. 2016 Beck et al., 2023), where an adaptive agent. is trained to learn new tasks quickly via. RL. Rakelly et al. (2199) and Zintgraf et al. (2020) propose meta-RL algorithms that consist of a task encoder and a cask-oonditional policy, similar to our inference model and latent-conditional policy. Zhou et al. (2019) propose agents that learn to probe the environment determine the latent variables explaining the dynamics. Differently from our problem, the true reward function of the task is known.
3 Background
We begin by introducing confounded imitation learning. Following Ortega al. (002)) we discuss how BC fails in the presence of latent confounders. We then define the interventional policy, the ideal (but priori intractable) solution to the confounding problem.
3

---

Published in Transactions on Machine Learning Research (09/2024)
3.1 Imitation learning
Imitation learning learns policy from dataset of experi demonstrations via supervised learning. The expert is policy that acts in (reward-free) Markov decision process (MDP) defined by tuple M = (S,A, A. P(s' 8, a). P(su)), where S is the set of states, A is the set of actions, P(s' s, a) is the transition probability, and P(su) is a distribution over initial states. The expert's interaction with the environment produces a trajectory r = (so, Ao,..., (r) The expert may maximize the expectation over some reward function, but this is not necessary (and some tasks cannot be expressed through Markov rewards (Abel al, 2021)). In the simplest form of imitation learning, BC policy ay(a s) parametrized by 1 is learned by maximizing the likelihood of the expert data, i.e., minimizing the loss Ee log 7 8), where DD is the dataset of state-action pairs collected by the expert's policy.
3.2 Confounded imitation learning
We now extend the imitation learning setup to allow for some latent variables 0 E θ that are observed by the expert, but not the imitator. We define a family of Markov decision processes as a latent space θ a distribution P(O), and for each 0 E θ, a reward-free MDP Mo = (S, P(s' 8,a,0), P(so 0)). Here, we make the crucial assumption that the latent variable is constant in each trajectory. We assume there exists an expert policy Fexp(a s, () for each MDP. When it interacts with the environment, generates the following distribution over trajectories T:
T P(142 [&,
Pexp(1 (T = P(so 0) II f=0
at;a;; az; Texp (atl sc; 0):
In this setting, the trajectories from the expert distributions are called confounded, because the states and actions have an common ancestor, the latent variable 0. The imitator does not observe this latent variable. It may thus need to implicitly infer from the past transitions, so we take it to be non-Markovian policy (a $1, St).. S;), parameterized by n. The imitator generates the following distribution over trajectories:
Pr(T 10) P(so 0) II P( 8t, ae; 0) ny(at 80, ..,...,,:)) .8).
The Bayesian networks associated to these distributions are shown in Figure
The goal of imitation learning in this setting is to learn imitator parameters such that when the imitator is executed in the environment, the imitator agrees with the expert's decisions. This means we wish to maximise Ea- Pr,,,, ( log eaxp st, f). the expert solves some task (c.g. g. maximizes some reward function), this amounts to solving the same task. The latent variable stays fixed during the entire interaction of the agent with the environment.
3.3
Naive behavioral cloning
If we have access to a data set of expert demonstrations, we can use BC on the demonstrations to learn the maximum likelihood estimate of the expert's policy. At optimality, this learns the conditional policy:
Teand (as 81, (d......8))
=
E Texp (at sc, 0),
(1)
Pre (~Peona((IT)
Peond (0 ()) x p(0) II p p(s st, a, 0 Texp stc 0)
Following Ortega et al. 2021
consider the following example of a confounded multi-armed bandit with
= θ I..... 5) and S = [0, 1}:
1 p(0) 5' Texp(a1 | 0)
6
if a, =
3
if 0, =
10
P(8111 se, 0z, 0)
4
2
1
a, # #0*
a, +
10
4

---

Published in Transactions on Machine Learning Research (09/2024)
Expert
Conditional
Interventional
Figure 2: Actions from rollouts from bandit environment defined by Equation The z-axis episode time, In the y-axis five roll-outs are shown from the expert and policies Equation and Equation 3 Colors denote actions, with the correct arm labelled green. The interventional imitator tends to the expert policy, while the conditional policy tends to repeat itself.
The expert knows which bandit arm is special (and labeled by 0) and pulls it with high probability, while the imitating agent does not have access to this information. We define the reward this bandit environment as тz St+1- However, note that we compare imitation learning algorithms that do not access the reward of the environment and only use the rewards for comparing the learned behaviors evaluation time.
If we roll out the naive BC policy in this environment, shown in Figure 2 we see the causal delusion at work. At time i, inferring the latent by Poond takes past actions as evidence for the latent variable. This makes sense on the expert demonstrations, as the expert knows the latent, variable. However, during an imitator roll-out, the past actions are not evidence of the latent, as the imitator is blind to it. Concretely, the imitator will take its first action uniformly and later tends repeat that action, as mistakenly takes the first action to be evidence for the latent.
3.4 Interventional policy
A solution to this issue to only take as evidence the data that was actually informed by the latent, which is the dynamics distribution p(5++1 St: de, 0). This defines the following imitator policy:
(a $1, ................................ st) (0)1))
Texp (ar st. 6), Pint 7) x p(0) II p(s st. 0),
(3)
In a causal framework, this corresponds treating the choice of past actions as interventions. To see this, consider the distribution p(a,, $1:...; sel...,,s/d(aa.... seldo(a1,...,, .x-11) By classic rules of the do-caleulus (Pearl 2009) this is equal to
1 p(0)p((c) 0) II ps4 a,,
Taking the conditional p(a./s1..... ....:,............. do(a1, d 0c-1) of that expression leads to equation (3). The policy in Equation 3 is therefore known interventional policy (Ortega et al. 2021
3.5 Failure of the conditional policy
Whether the conditional policy of Equation able to imitate expert behaviors under latent confounding depends the relative contribution of evidence due to the dynamics p(s++1 82, o,, 6) and the policy Texp(a; 0) when inferring the latent in the posterior Poond() 7). If most evidence comes from the dynamics, we expect the conditional and the interventional policies to be similar. On the other hand, when most evidence comes from the policy decisions, we expect large delusion. This is for instance the case when the expert policy is deterministic and the dynamics stochastic, or when the latent influences the state transitions only parts of the state space, but always affects the expert behavior. In Appendix Appendixx A} we quantitatively study the behavior of the conditional and interventional policies on the bandit example for varying amounts of dynamics and expert stochastieity.
5

---

Published in Transactions on Machine Learning Research (09/2024)
3.6 Optimality of the interventional policy
To understand why the interventional policy may tend the expert's policy consider (f Equation 3 Applying it to the bandit example, each subsequent observation 8t41 after taking action d, provides some evidence for the latent 0: if we find 8 = 1, then the latent 0 is more likely to equal d,, if we find 8t41 O, then 0 is less likely to equal dg. In Figure 2) we see that this "true interventional" indeed approaches the expert's policy. This is further illustrated in ssubscction 6.1. In this case, the interventional policy thus presents a solution to the confounding problem. In the rest of this paper, we focus on the question if and how it can be learned from data.
Note that the interventional policy may only achieve optimal performance after obscrving many time steps. However, it is not guaranteed to be the policy that adapls identifies the latent as fast as possible. This may require form of active exploration (hoou al. 2019) that is beyond the scope of this work.
4 Deconfounding imitation learning
In this section, we discuss how, instead of the conditional policy, we can learn the interventional policy and thus fix the confounding problem. This improvement we call deconfounding the imitation learning. First, we illustrate how this is provably possible in theory under strong assumptions. We then show how to learn it in practice, even when these assumptions don't ' fully hold.
4.1 Identifying the interventional policy in theory
Based on the graphical model shown in Figure and applying the rules of do-calculus, we can identify the interventional policy Tmm | $1, 01........ .......... s,) p _ $1...., .......,,,, St, do(oi...., ax-1) by back-door adjustment Pearl
(2009) if we observe the latent variable 0. This results in Equation 3 But since the latent is unobserved, the interventional policy is not identifiable without further assumptions on the distribution.
However, will now show that under some assumptions, the interventional policy does become identifiable. First, assume that the latent variable 0 of the MDP is static, in other words stays fixed during the interaction of the policy with the MDP. This makes the dynamics and expert policy stationary in cach trajectory. If we combine this with the strong assumption of recwrence, commonly used in theorelical reinforcement learning [Watkins& Dayan 1992 Thm. 1), the interventional policy becomes identifiable.
Assumption 1. We assume that the MDP is recurrent (Norris 1997 Sec. 1.5), meaning that all state- action pairs 8, are reached infinitely often in each trajectory.
If the MDP is recurrent with the expert policy, it is possible to identify the interventional policy in theory from an infinite dataset. D of expert demonstrations of infinite length. On each individual trajectory т with index i, we are able to count the state and action transitions to estimate the transition probability (r(5'ss, a) and expert policy (a(s) of that trajectory. We can then compute the likelihood of a datasct trajectory given a state-action sequence (s1, "1,,,,)..,8 based on the estimated environment transitions, H. fi(Set 188, at). The corresponding belief over dataset trajectories i follows from Bayes' theorem as
Hl Te 192, az)
Pin oc, 8t) En II. de 08141115
Then we can estimate
Mime (a.ls1, (1, 01,
81) Eingina Bin (1 ( (o
(4)
This successfully identifies the interventional policy:
Theorem 1. From infinitely many demonstrations of infinite length from a MDP that is recurrent with the expert policy Texps have that Wint from Equatiin 4 equals wint from Equation 3
6

---

Published in Transactions on Machine Learning Research (09/2024)
Proof. Because of the recurrence assumption, correctly estimate (ss5 a( and A+(a|s) (alp 0z). Then, suppressing normalization factor that is constant in ar, we write:
Io o)} (oxp(ar|oT, 0)d0
Tint (aT 81 ST) x p(0)
E p(01))(1))
I1lsers1ss a)
Texp (ar (orp(aaTTT 1sr, 0)d0
Il p( 0.) Taxp(arToT, 0.)
Σn(o (()
α E0(0 p())
I (1 ac) ?((aT(oT)
x iint (ar (0rr|811..... ($1,..,,8)) ST)
where we recognized p(0li) as the (Dirac) delta distribution around 0,, and p(i) is the uniform distribution over samples from the dataset D.
In this proof, the stationarity of the dynamics and expert make it possible to collect statistics across the time steps into distributions. Furthermore, the proof crucially relies on the recurrence assumption, as it allows us to evaluate the likelihood of any state-action sequence on each of the trajectories' distributions. Without this assumption, we will in general only be able to learn an approximation of the likelihood from state-action pairs observed in the trajectory.
The interventional policy estimated in Equation 4 requires marginalization over the entire demonstration dataset and thus isn't very practical. Instead, we propose to match the expert data with learned latent variable model Pa.
(T) = (0( (P(0))
II Po ($141lss a, () Od00.
(5)
Similar to the argument in Theorem under the same assumptions, if we match the ground-truth demon- stration likelihood p(7) with the latent variable model Pv.: then the interventional policy derived from the latent variable model Pwin matches the ground truth interventional policy Tin
4.2 Learning from demonstrations and explorations
If each expert demonstration covers the entire stale-action space, we have shown that we can estimate state dynamics and expert policy, and thus the interventional policy using Equation 4or the latent variable model in Equation
What if this strong assumption on the expert data does not hold? We can make the problem easier by allowing our agent to interact with the environment during training. That allows us to see more samples of the environment dynamics, making up for the lack of recurrence guarantees. Allowing interactions with the environment a common modification to the imitation learning problem, used for example in IRL.
In that case, we explore with policy (alp) the state distribution
(T) = =fo p(0) II क
ar, 0)d0
(6)
The corresponding latent model is given by
Po.expl (T) = / Pw (0)
II Texpi (s ar 0)d0
(7)
where we assume that we can evaluate the likelihood of the exploration policy.

---

Published in Transactions on Machine Learning Research (09/2024)
Training the inference model
Training the imitation pollcy
~~ P,(c)
-
L
* Pexp(t)
0
L LBC
VI
BC
Fesp
9p10 p,l5's,,,0)
9p(01
x,(als,0)
Figure Overview of the method. On the left, the inference model consisting of the encoder de and decoder Pé is Irained using variational inference. Training data is sampled from the environment using exploratory policy. The dotted arrow depicts how the exploratory policy can optionally be the imitation policy Ty trained on the right. On the right, the imitation policy is trained with behavioral cloning on expert data. The expert does not need to interact with the environment at training time. Instead, stored expert trajectories can be used. To learn the deconfounded policy, the encoder trained on the left. used for inferring the latent variable on the expert trajectories.
We propose to learn the latent model of Equation via amortized variational inference, using an inference model (0] (Kingma & Welling, 2014) that infers the latent variable from a trajectory. This involves maximizing the evidence lower bound (ELBO), lower bound of the log likelihood, with respect to vp and
log pepexpl( (T)
E Poppo ( (0)
log log 2o (017) + Σ log (aels)) log Pv (s
TroPexpl
E
log log q0 (017) + Σ log loggo const Lv1
prepecopl (1)=)
(8)
where can ignore the exploration policy likelihood it does not depend on the parameters. At optimality, and assuming sufficient expressivity, the inference model learns the model posterior
()( = (e)) x P (0) II (8441188
Crucially, the learned inference model (017) does not depend on the likelihood of the policy selecting actions. This is because we explore with a policy that unlike the expert does not depend on the latent 6. Therefore, we learn exactly the interventional latent inference model from Equation 3) with the learned dynamics.
In next step, we the inference model (o(7)) together with learned dynamics model (sc4 de, to learn an imitation policy (a|s,, 0), that approximates the expert policy Texp (a)s,a) f). Again, we use variational inference and train the model Pain from Equation 5 with samples from the demonstrations.
Thus, we maximize 7) in the ELBO
E.pp log (o., 2 Brap [og (a.) se 0)] const LBC
which effectively becomes a behavioral cloning loss. Here omitted terms constant in the policy parameters
7]
The training based on these two objectives illustrated in Figure 3 By optimizing the loss in Equation 8 we learn model infer the latent variable just basod on the dynamics of trajectory; by optimizing the loss in Equation 9 we learn a policy conditional on that latent. Combined, we learn to approximate the interventional policy of Eqqution
8

---

Published in Transactions on Machine Learning Research (09/2024)
5 Practical algorithm
We now present practical algorithm for training an agent from expert data in the presence of latent confounders. As outlined infsection 4 we learn to infer the latent variable by interacting with the environment using an exploratory policy, and learn to clone the expert on the demonstrations conditioned on the inferred latent. At test time, the agent alternates between updating its belief about the latent variable and imitating an expert dictated by its current belief. In [ppendix C we describe an algorithm that does not require the ability to gather more data interactively, but faces a more difficult learning problem in practice.
The agent consists of: 1) an inference model go, which maps trajectories T = (50, , to belief over latent variable (0(0) r); 2) dynamics model p& mapping latent, state, and action to a distribution over the next state (s 8, a. f): and 3) a latent-conditional policy 7n(a s, s, An overview of the training algorithm is presented in Figure 3
Components
Training the model for online inference The inference model training closcly follows the outline given in Section 4 However, since at Lest time, the agent needs to infer the latent online from partial trajectories, we adjust the model and the ELBO of Equation slightly. At timestep t, the encoder qo takes as input the trajectory observed until timestep t, and predicts a distribution of the belief over the latent ô.
We follow Zintgraf et al.] (2000) by defining the prior at timestep t as the inferred latent distribution from timestep & t- I, starting with a diagonal Gaussian with unit variance as the initial prior. This process is similar to Bayesian filtering, where beliefs about the state of a process are updated sequentially in response to new evidence. As such, each prior evolves based the latest observed data, similar to the updating mechanisms in Kalman filters and other Bayesian state estimators (Kalman 1960
Krishnan et al. 2015
The modified ELBO can then be written as
Lvi = Bynpeepp ( log
8t, st, 0] 3Dxl (20 1 TT)) go(0 [(--i)
(10)
where DKL is the KL divergence andT, denotes the trajectory until timestep t. Following prior work VAEs, we use a coefficient B3 for the prior regularization (Higgins et al. 2017].
As exploration policy, we use the latent-conditional policy Ty conditioned on i inferred by gs. This is convenient choice because using TH means we do not have to train multiple policies. Furthermore, using Tn for exploration biases the training data distribution for the inference model toward data that the policy encounters when it is deployed in the environment potentially improving the generalization of the inference model. In practice, we condition the policy on the mean of the inferred distribution instead of a sample from it. We find that using either does not make big difference. Other exploration policies may be used as long as they explore sufficiently diverse trajectories and do not depend on the true latent 0.
As described in section 4 the learned inference model is used for inferring the latents on the expert trajectories r.. The policy T. is then trained to minimize Equation 9 We show the pscudocode for the full training
algorithm in Appendix B
Test time At test time, the agent faces an enviromment with an unknown latent and needs to adapt to the correct expert behavior. We solve this problem by alternating between updating posterior belief over the latent and acting under the current belief. Concretely, the agent initially samples a latent from the prior ~ (O, 11) and an action a ~ (a O) to imitate the expert corresponding to that latent. It observes the state transition and computes the posterior belief with the inference network. Another latent is sampled from the updated belief, and so on. In practice, as during exploration, we do not sample from the inferred distribution but instcad condition the policy on the mean. Once the inference has converged to match the true latent for the environment, the true expert for the environment will be imitated consistently. We summarize the test-time behavior in pseudocode in Appendix B
9

---

Published in Transactions on Machine Learning Research (09/2024)
Deconfounded
Naive BC
True interventional
True conditional
Expert
0.6
50
0.5
Best amm count
0.4
0.4
0.3
0.3
0.2
20
60
80
100
20
40
100
500
1000
1500
2000
Episode step
Episode step
Training iteratonns
(a) Best arm probability on expert data
(b) Best arm probability on online data
(c) Learning curves
Figure 4: Imitation learning a multi-armed bandit problem. The shading shows the standard error of the mean. The left panel compares the policies when evaluated on trajectories sampled by the expert policy. The x-axis the slep on the trajectory and the y-axis the probability of choosing the best arm. The middle panel is otherwise the same, except run on trajectories sampled online with the policies themselves. The right panel shows the learning curves. The x-axis shows training iterations and the y-axis shows the number of times the best arm is chosen by the policy under training during trajectory with 100 time steps. The curves are averages for sliding window of length 10 training iterations.
6 Experiments
To test our method practice, we conduct experiments the multi-armed bandit problem from Ortega et al. (2021) and in multiple control environments. We aim answer three questions: 1) How big the effect of confounding naive BC- large enough to justify the use of specialized methods? 2) Is our algorithm capable of identifying the interventional policy? 3) How well does the interventional policy imitate the expert?
6.1 Investigating deconfounding a multi-armed bandit
We begin the empirical study by experimenting with the multi-armed bandit problem proposed by (Otega et al. (2021) and described in iufsection The expert policy defined in Equation 2) We consider episodes of length 100. As are only interested in the effects of confounding arising from over-fitting a small dataset, generate new training data from the expert for each update of the learning algorithms. Each learning algorithm is run for ten independent seeds and the results are averaged. The hyperparameters for the algorithms are provided in Appendix D
imitation learning, and not in effects
Naive BC and the conditional policy To answer our first question, we compare naive BC to the true conditional policy deseribed in Section 3] As the latter policy is non-Markovian, we also allow the BC policy to observe the history. One way to enable this adaptation to equip the agent with a memory, which the agent can learn use for representing its belief about the latent variable. To provide such a memory, we implement the imitation learner as a recurrent neural network (RNN). Figure 4 a) shows the probability the different policies assign to choosing the best arm when evaluated on data collected by the expert policy. We see that the naive BC agent learns policy that matches the true conditional policy closely. This results in problems for the policy learned with naive BC when it is deployed in the environment and has to choose the actions itself, as shown in Figure 4 b). The naive BC closely tracks the action probability of the true conditional policy, which performs much worse than the expert policy. These results suggests that it has suffered the full impact of the confounding problem.
Deconfounded imitation learning and the interventional policy To answer our second question, we implement the dcconfounded imitation learner as described in (section 5 Figure 4 a) and b) show that the proposed method closely matches the true interventional policy both on expert trajectories and online, Figure 4 c) shows the number of times the policies chose the best during an episode. Our deconfounded
10

---

Published in Transactions on Machine Learning Research (09/2024)
DAgger
Deconfounded
GAIL
Naive BC
Oracle
200
-500
-500
Episode return
Episode retum
-1000
-200
Episode
-1500
-400
-1500
-2000
-600
100
200
100
200
Training iterations
Training Iterations
Training Iterations
(a) Lunar Lander
(b) Half Chectah
(c) Ant
Experiments in our confounded, stochastic environments. We show the episodic return each agent over the course of training. The curves for the are averages sliding window of length 5. The shading shows the standard error of the mean.
Figure 5:
imitator converges to the true interventional policy, showing that it learned to optimally imitate the expert in the presence of latent confoundrr, and answering our third question. This near-perfect imitation perfor- mance comes at the price of requiring exploration data to train the inference model well as an increased number of training iterations needed for convergence.
6.2 Demonstrating deconfounding confounded MDPs
Next, we demonstrate that the confounding issue affects MDPs with considerably more complex dynamics than the bandit and that the answers to the three questions do not change with the increased complexity, For LumarLander-v2 [Brockman al., 2016), we consider modified version with unknown key bindings: latent 0 specifies a permutation in the map between two of the the agent's actions and the behaviors (firing the left and right engines of the space craft). This permutation is known to the expert, but not the imitator. For HalfCheetahBulletEnrrr
(Coumans & Bai, 2016 2021), modify the environment similarly to Swany et al. (2022b) by varying the target speed. The expert observes the true target speed while the imitator only observes a noisy indicator, which shows whether it running faster or slower than the target speed. For AntGoal-vo (Todorov et al.) 2012 we consider a version, where the task is to run to a goal randomly sampled from within a circle around the starting point. The expert observes the true goal coordinates, while the imitator only observes a noisy indicator of the goal direction and distance. In all enviromments, in addition to providing the imitator with less information about the latent variable than the expert, make the enviroument somewhat stochastic. These changes make the naive imitators try to infer the latent from the more deterministic expert actions rather than the stochastic dynamics, which results in the confounding problem. For full details about the environments, and how the confounding problem may arise in each of them, please see the Appendix E.
The expert policies are trained with proximal policy optimization (Schulman et al. 2017). In order to avoid fiite-sampll-sizee effects, we use an infinite-size training dataset by generating expert trajectories on the fly. Each learning algorithm is run for or 6 independent seeds and the results are averaged. Additional details are provided in Appendix El
Naive BC Like in the bandit example, we first analyze how strongly confoundigg affects the naive BC policy. The naive learner is again implemented as an RNN. In Figure 5) we see that the naive BC agent fails to imitate the expert behavior and performs much worse than an oracle imitator, which otherwise exactly the same sctting but it is trained with knowledge of the latents. This failure to generalize again evidence for causal delusions: the agent learned to infer the latent from the expert behavior rather than the noisy dynamics of the environment.
11

---

Published in Transactions on Machine Learning Research (09/2024)
Deconfounded imitation learning To solve the confounding issue in these environments, we first test the DAgger algorithm (Ross et al. (011), which queries the experts during training. We find that it indeed solves the confounding problem: the agent quickly approaches the performance of the oracle imitator in all three environments.
The DAgger performance serves as an unachievable upper bound for our method, as it not only solves the confounding issue but also reduces the more common distribution shift issue present in imitation learning. However, recall that the expert policy may be defined by, e.g., human expert, who would be expensive impossible to query at imitation learning time making DAgger a potentially difficult method to use in practice.
Can our deconfounded imitator also solve the confounding problem? We test. the deconfounded imitation learning algorithm described in (section 3 on the control environments. We find it beneficial to modify the described algorithm in two ways: 1) extending the roconstruction loss in Equation 10 to multi-stop predictions (Hafner et al. 2020 : 2) conditioning the policy on the inferred latent 02 at the current timestep instead of the last time step O₂ when training the policy. See Appendix E for details.
In Figure 5, we see that the deconfounded learner clearly outperforms the naive BC baseline, matching the performance of the policy learned by DAgger in two out of three environments and coming close the third. While it is hard to interpret what functions high-dimensional neural networks have learned exactly, matching the performance of DAgger, which we know can recover the interventional policy, is encouraging. It suggests that the deconfounded agent learned to infer the latent from environment transitions and learn a policy that acts like the expert under the inference. In the Ant environment, it is impossible to infer the direction from the noisy indicalor from a single observation. Therefore, the agent may start moving the wrong way in the beginning of the episode. Recovering from such a false start may be much easier with the help of extra supervision from the expert, giving DAgger a particularly strong advantage in this environment.
Inverse reinforcement learning In theory, IRL is capable of recovering the expert performance even in confounded environments (Swamy et al. 2022b). We test this theory in practice, by comparing to generative adversarial imitation learning (GAIL) (Ho & Ermon, 2016). We implemented GAIL closely following a popular publicly available implementation" and using recurrent PPO by Raffin et al. (2021) | as the RL algorithm. To stabilize GAIL, we sample four times more data for every training step compared to the other algorithms.
In LunarLander, GAIL is able to achieve higher returns than Naive BC, but does not recover the performance of the proposed method. In HalfCheetah and Ant, GAIL does not recover the performance of Naive BC. In Appendix D we provide results for GAIL in the bandit environment from |subsection 6.1] We found that in the bandit, GAIL does not recover the interventional policy and therefore does not converge to the expert behavior during an episode. This is because the reward function defined by GAIL can always be maximized by a deterministic policy, even in the case when the expert is a stochastic policy, like in the bandit. From these results, we conclude that while IRL is in theory capable of solving confounded imitation learning problems, where the expert is deterministic policy, getting it to work well can be difficult in practice.
7 Limitations
While the proposed method recovers the policy DAgger learns under ideal conditions, it does not address the original challenge DAgger is designed to solve. That is, if the expert. data does not cover the stale-action space sufficiently, the proposed method may not be enough to learn the optimal policy. Furthermore, while the proposed method does not require access to the expert, it does require sampling access to the environment similarly to IRL. To find the deconfounded policy fully ofllin,, an algorithın outlined in the in section 4 could be used. For the methods to work provably, we need to make strong assumptions of recurrence.
Following Ortega et al. (2021) we model the confounded II setting with a CMDP where the latent stays fixed during the entire interaction of the policy with the environment. This may be a limiting assumption,
12

---

Published in Transactions on Machine Learning Research (09/2024)
for example in driving scenario, where the environment may have latent factors that govern the dynamics over small section of the road but the agent is expected to be able to traverse many different sections.
Finally, while the environments we experiment in are commonly used to evaluate RL and IL policies, we acknowledge that these are fairly limited domains and do not reflect the full complexity of training useful policies for the real world.
The components of our algorithm, variational inference and behavioral cloning, are commonly used techniques elsewhere in machine learning. Therefore, scaling the proposed method to real world problems should be possible. However, we did not optimize for sample complexity of either learning algorithm, which would limit the application of our method to real world problems. Developing a more sample efficient version of the algorithm and lesting it in real world domains is an exciting avenue of future work.
8 Conclusion
Naive imitation learning algorithums can fail the presence of latent confounders -ooeens--oorr for instance when the expert has access to more information than the imitator. This work presents a breakdown of this confounding problem. First, we studied under which conditions latent confounding impacts the performance of BC. We demonstrated that this issue is more relevant when there is substantial stochasticity in environment transitions or the expert policy is nearly deterministic. We then analyzed in which settings the interventional policy, which solves the confounding issue, is identifiable without query access to the expert.
Informed by the theoretical results, we proposed a practical algorithm for deconfounding imitation learning with variational inference that provably converges to the interventional policy. While this problem has been previously studied in theory with inverse RL, we propose a novel variational inference solution, which we analyze theoretically and evaluate in practice. Finally, we evaluated the proposed method with experiments in multi-armed bandit and confounded MDPs. We found it was able to learn interventional policies in all of the settings, alleviating the confounding problem that limits naive imitation lcarning.
References
David Abel, Will Dabney, Anna Harutyunyan, Mark K Ho, Michael Littman, Doina Precup, and Satinder Singh. On the expressivity of markov reward. Advances in Neural Information Processing Systems, 34: 7799-7812, 2021.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey Hinton. arXiv:1607.06443 2016.
Layer normalization.
arXiv preprint
Jacob Bock, Risto Vuorio, Evan Zheran Liu, Zhong Xiong, Luisa Zintgraf, Chelsea Finn, and Shimon White- son. A survey of mela-reiifforcement learning. arXie preprint urXiv:2301.088 0080088 2023.
Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojeiech Zaremba. Openai gym. arXio preprint 1rXiv:1606.0111 01540, 2016.
Kyunghyun Cho, Bart Van Merriënboer, Dzmitry Bahdanau, and Yoshua Bengio. On the properties of neural machine translation: Encoder-decoder approaches. In Workshop on Synlaz, Semantics and Structure in Statistical Translation, 2014.
Djork-Arné Clevert, Thomas Unterthiner, and Sepp Hochreiter. Fast and accurate deep network learning by exponential linear units (elus). orXiv preprint arXiv:1511.07989, 2015.
Felipe Codevilla, Eder Santana, Antonio M López, and Adrien Gaidon. Exploring the limitations of behavior cloning for autonomous driving. April 2019. URL http:///rxxi.org//bbar wwt.:///r/aab./9.....o. http:/9994.088880 08980
Erwin Coumans and Yunfei Bai. Pybullet, a python module for physics simulation for games, robotics and machine learning. http://pybuulett. org 0016-2021.
Pim de Haan, Dinesh Jayaraman, and Sergey Levine. Causal confusion in imitation learning. Advances in Neural Information Processing Systems, 32, 2019.
13

---

Published in Transactions on Machine Learning Research (09/2024)
Yan Duan, John Schulman, Xi Chen, Peter L Bartlett, Ilya Sutskever, and Pieter Abbeel. RI?: Fast reinforcement learning via slow reinforcement, learning. arXiv preprint arxiv:661 (0001.0000 1.02799, 2016.
Yan Duan, Marcin Andrychowicz, Bradly C Stadie, Jonathan Ho, Jonas Schneider, Ily Sutskever, Picter Abbeel, and Wojeiech Zaremba. One-Shot imitation learning. 2017.
Justin Fu, Katie Luo, and Sergey Levine. Learning robust rewards with adversarial inverse reinforcement learning. 2017.
Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. "International Conference Learning Representations " 2020.
Irina Higgins, Loie Maithey, Arka Pal, Christopher P Burgess, Xavier Glorot, Matthew M Botvinick, Shakir Mohamed, and Alexander Lerchner. beta-vae: Learning basic visual concepts with constrained variational framework. ICLR (Poster), 3, 2017.
Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning. Advances in neural information processing systems, 29, 2016.
Rudolph Emil Kalman. A new approach to linear filtering and prediction problems. 1960.
Diederik Kingma and Jimmy Ba. Adam: method for stochastic optimization. In International Conference on Learning Representations, 2015.
Diederik P Kingma and Max Welling. Auto-encoding variational bayes. In International Conference on Learning Representations, 2014.
Rahul G Krishnan, Uri Shalit, and David Sontag. Deep kalman filters. arXiv preprint arXiv:1511.05191. 2015.
Ashish Kumar, Zipeng Fu, Deepak Pathak, and Jitendra Malik. Rma: Rapid motor adaptation for legged robots. 2021.
Daniel Kumor, Junzhe Zhang, and Elias Bareinboim. Sequential causal imitation learning with unobserved confounders. 2021. URL https://opeerrrieet wttp://ww..eeeoooorooooooooooo00
Michael Laskey, Jonathan Lee, Roy Fox, Auca Dragan, and Ken Goldberg. Dart: Noise injection for robust imitation learning. In Conference on robot learning, pp. 143-156. PMLR, 2017.
Nya Loshchilov and Frank Hutter. Decoupled weight decay regularization. ar preprint urXiv:1711.0511 2017.
Yiren Lu, Justin Fu, George Tucker, Xinlei Pan, Eli Bronstein, Becca Roelofs, Benjamin Sapp, Brandyn White, Aleksandra Faust, Shimon Whiteson, al. Tmitation not enough: Robustifying imitation with reinforcement learning for challenging driving scenarios. arXio preprint arXiv:2212. 11419, 2022.
Andrew Y Ng, Stuart Russell, et al. Algorithms for inverse reinforcement learning. In International Confer- ence on Machine Learning, 2000.
R Norris. Markov Chains. Cambridge Series in Statistical and Probabilistic Mathematics. Cambridge University Press, 1997.
PA Ortega and DA Braun. A bayesian rule for adaptive control based on causal interventions, In Third Conference on Artificial General Intelligence, 2010a.
Pedro A Ortega and Daniel A Braun. A minimum relative entropy principle for learning and acting. Journal of Artificial Intelligence Research, 38:475-511, 2010b.
Pedro A Ortega, Markus Kunesch, Grégoire Delétang, Tim Genewein, Jordi Grau-Moya, Joel Veness, Jonas Buchli, Jonas Degrave, Bilal Piot, Julien Perolat, Tom Everitt, Corentin Tallec, Emilio Parisotto, Tom Erez, Yutian Chen, Scott Reed, Marcus Hutter, Nando de Freitas, and Shane Legg, Shaking the founda- tions: delusions in sequence models for interaction and control. October 2021.
14

---

Published in Transactions on Machine Learning Research (09/2024)
Abhishek Padalkar, Acorn Pooley, Ajinkya Jain, Alex Bewley, Alex Herzog, Alex Irpan, Alexander Khaz- atsky, A Anant Rai, Anikait Singh, Anthony Brohan, et al. Open x-embodiment: Robotic learning datasets and rt-x models. arXiv preprint S.kv:2310008666, 2023.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Autiga, el al. Pytorch: An imperative style, high-perlormance deep learning library. Advances in neural information processing systems, 32, 2019.
Judea Pearl. Causality. Cambridge University Press, September 2009. URL https: hetps://wwk. http://wwww.abbiddg..... org/
core/books/cauaa /tt///////uuaaaaa/00000660//0/00
Xue Bin Peng, Erwin Coumans, Tingnan Zhang, Tsang- Wei Edward Lee, Jie Tan, and Sergey Levine. Learning agile robotic locomotion skills by imitating animals. In Robolics: Science and Systems, 07 2020.
doi: 10.15607 (016607/RSS.0000 htt./RRSS.200.X//00 2020.XVT.064.
Pomerleau. Alvinn: An autonomous land vehicle in a neural network. Advances in Neural Information Processing Systems, 1988.
Antonin Raffin. RI baselines3 zoo. htttp:///ggthhh..om///////////es
2020.
Antonin Raffin, Ashley Hill, Adam Gleave, Anssi Kanervisto, Maximilian Ernestus, and Noah Dormann. Stable-baselines:: Reliable reinforcement learning implementations. Journal of Machine Learning Re- search, 22(268):1 8, 2021. URL http://ww...orrpppperrrs//////6a
Kate Rakelly, Aurick Zhou, Chelsca Finn, Sergey Levine, and Deirdre Quillen. Efficient off-policy meta- reinforcement learning via probabilistic context variables. In International conference on machine learning. PMLR, 2019.
Danilo J Rezende, Ivo Danihelka, George Papamakarios, Nan Rosemary Ke, Ray Jiang, Theophane We- ber, Karol Gregor, Hamza Merzie, Fabio Viola, Jane Wang, et al. Causally correct partial models for reinforcement learning. arXiv preprinl ar Xiv:2000 02836, 2020.
Stephane Ross, GeolTrey Gordon, and Drew Bagnell. A reduction of imitation learning and structured prediction to No-Regret online learning. In Geoffrey Gordon, David Dunson, and Miroslav Dudfk (eds.), International Conference on Artificial Intelligence and Statistics, volume 15 of Proceedings of Machine Learning Reseorch, Fort Lauderdale, FL, USA, 2011. PMLR.
Kangrui Ruan, Junzhe Zhang, Xuan Di, and Elias Bareinboim. Causal imitation learning via inverse rein- forcement learning. In International Conference Learning Representations, 2022.
Stuart Russell. Learning agents for uncertain environments (extended abstract). In The Eleventh Annual Conference on Computational Learning Theory, 1998.
S Schaal. Is imitation learning the route to humanoid robots? Trends in Cognitive Sciences, 3(6):233-245 1999.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov, Proximal policy optimiza- tion algorithms. arXio preprint Xrv:: 1eXi::177.003:: 100.00347 2017.
Jonathan Spencer, Sanjiban Choudhury, Arun Venkatraman, Brian Ziebart, and J Andrew Bagnell. Feedback in imitation learning: The three regimes of covariate shift. February 2021. URL http://arxiv.. http:/awxwvw.r.or/aeemmo abs/
2102. 02872
Gokul Swamy, Sanjiban Choudhury, J. Andrew Bagnell, and Zhiwei Steven Wu. Causal imitation learning under temporally correlated noise, 2022a.
Gokul Swamy, Sanjiban Choudhury, J Andrew Bagnell, and Zhiwei Steven Wu. Sequence model imitation learning with unobserved contexts. Advances in Newral Information Processing Systems, 35, 2022b.
15

---

Published in Transactions on Machine Learning Research (09/2024)
William R Thompson. On the likelihood that unknown probability exceeds another in view of the evidence of two samples. Biomelrika, 25(3-4):285-290 1933.
Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physies engine for model- based control. In 2012 TEEE/RSJ indernational conference intelligend robots and systems, pp. 5026 5033. IEFF, 2012.
Jane X Wang, Zeb Kurth-Nelson, Dhruva Tirumala, Hubert Soyer, Joel 7 Leibo, Remi Munos, Charles Blundell, Dharshan Kumaran, and Matt Botvinick. Learning to reinforcement learn. In Annual Meeting of the Cognitive Science Society, 2016.
Christopher J C Н Watkins and Peter Dayan. Q-learning. Mach. Learn., 8(3):279 292, May 1992. URL https://doi.co 100.0.000000//00555050
Chuan Wen, Jierui Lin, Trevor Darrell, Dinesh Jayaraman, and Yang Gao. copycat agents in behavioral cloning from observation histories. Ade. Nearal Jnf. Pro-
Fighting
cess. Syst., 33:2564 2575, 2020. URL https://proceeddnngs newris.nnrppppp. pctppppp.///ppp////// 1 Abstract.ttattttooooooaaet
Chuan Wen, Jianing Qian, Jierui Lin, Jiaye Teng, Dinesh Jayaraman, and Yang Gao. Fighting fire with fire: Avoiding DNN shortculs through priming. June 2022. URL
10816
Wenxuan Zhou, Lerrel Pinto, and Abhinav Gupta, Enviroument probing interaction policies. International Conference on Learning Representations, 2019.
B D Ziebart, A L Maas, J A Bagnell, and A K Dey. Maximum entropy inverse reinforcement learn- ing. Aaai, 2008. URL https: //t::///wwww/ /wwp//wwaaooooooo aaai org.//Appre//AA1//20//AAA 227 http:/AAAAAA--22..fffooorrreeeep pdffssurcccoppstt page----a
Luisa Zintgraf, Kyriacos Shiarlis, Maximilian Tgl, Sebastian Schulze, Yarin Gal, Katja Hofmann, and Shimon Whiteson. Varibad: A very good method for bayes-adaptive deep rl via meta-learning, In International Conference on Learning Representations, 2020.
16

---

Published in Transactions on Machine Learning Research (09/2024)
Expert
Interventional
Conditorrl
0.96
0.50
Expett stochuskily
0.72
0.56
0.4
0.40
0.2
0.22
0.0
0.4
0.6
0.8
0.0
0.2
0.4
0.6
0.8
0.9
0.2
0.4
0.6
0.8
1.0
Environment stochasticty
Environment stochastictty
Environmant stochastichty
Figure Conditional and interventional policies in the bandit environment Equation 2 function of environment stochasticity and expert stochasticity. We show the expected success probability for the expert (left), the interventional policy (middle), and the conditional policy (right). Along the axes, we vary the expert policy from deterministic (bottom) to stochastic (top) and the enviromment transitions from determin- istic (left) to stochastic (right); see text for details. The interventional policy performs similar to the expert, while the conditional policy performs worse due to latent confounding when the environment stochasticity is larger than the expert stochasticity.
A Confounding and stochasticity
The behavior success of the conditional and interventional policies depend the stochasticity of the envi- ronment and the expert, this affects the evidence for the inference of the latent variable. If most evidence comes from the dynamics, which would g. be the case when the expert policy is very noisy and the dy- namics deterministie, we expect the conditional and the interventional policies to be similar. However, when most evidence comes from the policy decisions, we expect large delusion.
See figure Figure 6 for plot of the rewards of the various policies given varying amounts stochasticity. Instead of the fixed coefficients of Equation 2 we vary Texp(a; 6) between (deterministic expert) and (maximally stochastic expert) and P(s++1 = 1/a, 0) between 1 (dcterministic environment) and 2 (maximally stochastic environment).
In the left panel, we show the expected success probability of the expert policy: as expected, it is highest when both the expert and the environment are deterministic. The interventional policy of Equation 3 performs similarly. However, the conditional policy, shown in the right panel, only able to match the expert behavior if the stochasticity When the environment is stochastic and the expert comparably deterministic, the conditional policy (and thus naive BC) infers latents largely from past actions, suffers from delusions, and fails to imitate.
the expert policy outweighs the stochasticity in the environment.
B Training deconfounded imitators
Algorithm 1 [shows how agent is trained in practice from expert data in the presence of latent confounders. Algorithm shows the test time implementation of the deconfounded imitation learners, which alternate between updating a posterior belief over the latent and acting under the current belief.
C Deconfounded imitation learning from expert data alone
As discussed in |eccion 4 when suitable demonstrations from the expert are available, the interventional policy identifiable from the expert data alone, without access the environment or the expert. In this section, we present a practical algorithm for learning the deconfounded imitation policy from such an expert dataset. At test time, the agent works the same way the Algorithm ] presented in the paper.
17

---

Published in Transactions on Machine Learning Research (09/2024)
Data: Initial parameters of the imitation policy n, inference model Ø, and dynamics model di; expert dataset (rd), MDP family Mo, true latent distribution p(0), learning rates d1, 022 03-
while not done do
p(0); Su (s();
su.
for t H do
0, (0l r:)):
Infer latent for trajectory // Sample action from Markov policy // True dynamics
at (at st. f);
8(+1 p(8++1 82, u2 0); Append (at, S141) to T; + +-1+l. 1;
end
an V VoLvi(T);
// Sample estimate of Equation 10
1 t v a2 V (T)(
P d 0 rz);
1/ Infer latent from j-th expert trajectory Algorithm 1: Training deconfounded imitators
5 as Vn En Σ Exdert log (a 6.0%)):
end
Data: Parameters of the policy and inference model Φ; MDP Mo with ground-truth latent 0.
~ (s():
+ 8n, 1 1-+0 O;
for t t H do
0, ~ go(0, (0: TTp)
// Infer the latent for trajectory
at ~ ny se, 0.)
Condition on inferred latent
set1 ~ p( 81, at, 0)
// True dynamica
Append (ar, 3(+1) to т;
t + 1;
end
Algorithm 2: Deconfounded imitators at test time
18

---

Equation 11 Published in Transactions on Machine Learning Research (09/2024)
Training inference model directly on the expert demonstration faces the same problem naive imitation learning, i. e.e the trained model takes the expert's actions as evidence for the latent. However, by the assumption, that the expert is uniquely determined by the environment dynamics, we can directly learn the conditional policy and dynamics model explaining the expert trajectorics from the demonstrations because a latent that explains the dynamics also explains the expert. To Irain such models, we use variational inference learn a trajectory encoder Teur which infers the latent for the expert trajectoriis, and a factorized decoder, which reconstructs the dynamics of the environment and the expert's policy using networks pé and An respectively. The variational inference objective is given by
H
| 8t, ar, 0) + log (ay(aass, 0) -
3Dxr, (ooaa (0 Tz) (22(|100) v(0)).
Lott.i = BB (0
Po P ($
(11)
which, unlike the objective in the main paper, represents VAE where the encoder Yaur takes input the full expert trajectory reo and the decoder decodes both the action and transition probabilities throughout the trajectory.
This gives us a way for training the conditional policy imitating the experts in the demonstrations and a dynamics model. However, we cannot directly use the learned inference model lsoa for implementing the interventional policy, because it takes the expert's actions as evidence for the latent. The Algorithm | works by separately learning an inference model from interactions with the environment and using that inference model for deconfounding the expert trajectories. When we do not have sampling access to the environment, cannot learn the inference model directly. Instead, we observe that one factor of the decoder used for training the inference model is dynamics model of the environment conditional on the predicted latent. Therefore, we can use it to generate synthetic trajectories for training an online inference model oden to minimize the online variational inference objective given in the main paper. The online inference model can then be used for implementing the interventional policy similarly as Algorithm 2 The full offline dynamics learning algorithm is presented in Algorithm
Data:
The initial parameters of the imitation policy offline inference model datl online inference
model oon dynamics model v prior distribution for the learned latent space p(0), a dataset of expert trajectories {r²), an MDP (S.A, p. Po. H), learning rates 01, a2 dg, and a4
while not done do
θ 0~ p~p(0)=
// Sample from the prior // Sample from learned model or expert data
8g (S()
Tsynth 80, 0;
for 1 S do
a, π xac st)
// Sample action from Markov policy // Dynamics model
8++1 = (8++1 8t, or 0);
Append (az, 8141) to Tsynthi
t=
end
Don = don - a1 V oou (( (Tayyth))
// Train the online model on
Equation 10
Doff
doff a2 V aua E, LoT(
// Train the offline model Equation 11
1b = + 03 v E; (o)
// Train the dynamics model on
= 11 ao v Σ, (T)(
// Train the policy on Equation
end
Algorithm 3: Training deconfounded imitators, offline variant.
Multi-armed bandit experiment
Implementation details The inference model qa is implemented as an RNN with GRU architecture Cho et al. 2014) with hidden layer of 256 units. Before the RNN, the observation is preprocessed by
19

---

Published in Transactions on Machine Learning Research (09/2024)
Hyperparameter
Value
Episode length
100
Imitation training steps
5000
Dynamics model training batch size (full episodes)
100
Imitation training batch size (full episodes)
100
Behavioral cloning learning rate
0.001
Variational inference learning rate
0.0001
KL coefficient (3)
0.001
Table 1: Hyperparameters for the deconfounded behavioral cloning and naive behavioral cloning algorithms
MLP with two hidden layers of size 256 units and output size 32. The action is preprocessed by a linear transformation to 32 dimensional vector. The outputs of the RNN processed by linear transformation to vector which parametrizes the latent distribution. The latent distribution is 256 dimensional Gaussian. One half of the predicted vector represents the mean of the latent distribution and the other half, after softplus activation has been applied to it represents the variance.
The decoder is an MLP with two hidden layers of size 256 and a linear output layer. uses the same input preprocessing networks for the observations and actions as the inference model. The policy is an MLP, which takes the latent sample, and an observation as inputs. It uses the same observation embedding network the other networks and then has two hidden layers with 256 units each. The naive BC bascline uses the same nelwork archilecture as the deconfounded algorithm, except. it does not represent the belief as probabilistic latent variable and therefore there is no sampling step. It just directly passes output of the trajectory encoder as the input to the policy network. All of the MLPs use ReLU activations. All networks are optimized using the Adam optimizer (Kingma Ba 2015) with default settings from PyTorch (Paszke
et al.
2019 except for the learning rate.
Hyperparameter settings The hyperparameters used for the learning algorithms are presented in Ta- ble Table
Computing the ground truth policies All of the probabilities relevant to the bandit problem are known exactly from the definition of the problem and the conditional and interventional policies given in the main paper. Using these probabilities we can compute the true conditional and interventional policies, allowing us to compare the learned algorithms to the relevant optimal policies.
In practice, the true belief over theta can be computed for any trajectory as follows
(0. [Atl) + 8t log 3 log log 5. log logp(0(+1[A1) lo p [A,]) + st log + (1 st) log otherwise
i + (1 5t) log 1
if A. az
(12)
The true interventional policy can then be computed by sampling a belief 0, ~ log (().) and sampling an action from Taxp(a, (a, (x, dz). This can be seen Thompson sampling policy Thhompson 1933), which acts
optimally given its current belief of the task. The true conditional policy is computed similarly, except taking the actions as evidence for the latent is added to the update
(ogr((.)[A)) + 5t log 3 + (1 st) log 1 + log 10
if A otherwise
log ((0++1(AA))
(13)
logg((,)) + log / + (1 s,) log ] + log TO
E Confounded MDPs
Environment descriptions For LunarLander-v2 Brockman et al. 2016). we consider modified version with unknown key bindings: latent specifies a permutation in the map between two of the the agent's actions and the behaviors (firing the left and right engines of the space craft). This permutation is known
20

---

Published in Transactions on Machine Learning Research (09/2024)
-400
Episode return
-600
-800
-1000
-1200
Naive BC
Oracle
-1400
50
100
150
200
250
Training iterations
Figure 7: Comparing naive BC and oracle in a deterministic variant of HaLfCheettabBauletBbvv-... Most of the difference in the episodic returns here compared to the returns reported in (waamy et al. (202bb) follows from their results adding constant 1000 to the episodic return for plotting.
the expert, but not the imitator. In addition, we make the environment stochastic: at each step, with probability of 0.15, a uniformly chosen random action performed instead of the action chosen by the agent. The confounding problem arises because the expert's actions are temporally correlated, for example it may select the same action multiple times to adjust the lander's attilude. A naive recurrent BC policy may thus use past actions to predict future actions causing confounding if the controls are swapped, as a wrong initial guess influences future decisions.
(Counans & Bai 2016 2021), we modily the enviromment similarly to (Swamy et al. (2022b) by varying the target speed. The expert observes the true target speed while the imitator only observes an indicator, which shows whether it running faster or slower than the target speed. The expert smoothly tends to the target speed. This leads to correlated actions, in which the trend of past speeds can be extrapolated. A confounded naive BC imitator may attempt to extrapolate random trend of its past speeds. Contrary to Swanny et al. (2022b), we find that the deterministic version of this environment does not necessarily differentiate between naive BC and DAgger. Therefore, we make the indi- cator stochastic by randomizing it 20% of the time. For a detailed discussion about the effect of stochasticity on HallcChhetahBBlleeeeeeeeaea see below.
For Hal.CChettaBulllteEns
In AntGoal-vo Todorov al. 2012); ; we consider version of the classic ant robot locomotion environment, where the task is to run to goal randomly sampled from within circle around the starting point. This is similar to the ant environment considered by Zinngraf et al. (2020). We make the goal radius 10 times larger to make the locomotion task more challenging and remove the control cost. The imitators receive noisy observations of the length and angle of the vector between the ant and the goal in the global coordinate frame. Gaussian noise with scale 0.1 is added to the indicator. The expert knows the true goal location and starts moving toward it immodiately, making its actions correlated across time. The confounding arises because the expert trajectories, the goal direction and distance be inferred from the expert actions and states without considering the noisy indicator.
Deterministic vs stochastic variant of HalfCheetah. We test the different imitation learning algo- rithmns on the faLlfhheetabBBlleeEEn--00 environment from Swamy al. (20226). We found that running our implementation of naive BC matches the performance of our implementation of DAgger, which is incon- sistent with the results of Swanny et al. (2022)), where they report that DAgger outperforms naive BC, See figure Figure for learning curves. One reason for why their results show difference between the algorithms may be that their implementation of BC may sufler from covariate shift resulting from it having been trained on a small number of samples compared to DAgger. However, if that is the case, it is not an example of the confounding problem we are interested in. To make sure that we are testing for the ability of the algorithms to deal with the confounding problem in this environment, we make the indicator variable stochastic. The results for the stochatic variant reported in the main paper.
21

---

Published in Transactions on Machine Learning Research (09/2024)
Hyperparameter
Value
Top-level training itcrations
300
Initation inference training steps per top-level iteration
100
Dynamics model training batch size (full episodes)
16
Imitation training batch size (full episodes)
16
BC learning rate
0.0003
Variational inference learning rale
0.0003
KL coefficient (3)
1.0
Table 2: Hyperparameters for the deconfounded BC, DAgger, and naive BC algorithms for LunarLander-vZ HalfCheetabbulletti---a and AntGoal-vO environments. We fix the episode length for LunarLander-v2 to 500 steps, for HalfCheetalBulletEnrraa to 1000 steps, and for AntGoal-vO to 200 steps.
Implementation details For LunarLander, HalfChoetah, and Ant use the same network architecture as the bandit. experiments with the following changes. The output of the encoder paramelrizes the mean and the log-variance of the latent distribution. There is no parameter sharing between the different input processing networks in the encoder, policy, and decoder. The RNN has a hidden size of 500. The MLPs have two hidden layers of size 128. The action embedding has size 128. The latent distribution is a 64 dimensional Gaussian for LanarLander and HalfCheetah, and 8 dimensional Gaussian for Ant. In practice, the latent dimensionality does not matter very much. The MLPs use ELU (Clevert al. 2015)) activations followed by LayerNorm (Ba et al.) 2016 The networks are optimized with AdamW (Loshchilov & Hutter 2017).
The experts are trained using a PPO implementation by Raffin et al. 2021 with hyperparameters from 2021 is applied to the observations
Raflin 20200) Additionally, a running normalization layer (Rafiin et al.
and fixed after expert training to make the distribution the observations easier to learn the decoder on.
On each iteration of the training algorithm, a batch of new episodes are collected with the appropriate policy or policies. Naive BC only collects data using the expert policy. DAgger collects data using policy defined by mixture of the expert. and imitator policies, which linearly annealed from the expert policy to the imitator policy during the training. Deconfounded algorithm collects batch episodes with the expert policy for BC and batch episodes with the imitation policy for training the encoder. These trajectories are saved in circular bulfers that hold 2 * 106 transitions. After each data collection step, the imitalor (and encoder) are updated for 100 steps with the update function corresponding to the algorithm.
We implemented GAIL closely following popular publicly available implementation and using recurrent PPO from stable-baseline) Raflin et al. (2021) the RT algorithm. We the same data sampling pipeline the deconfounded imitation learning for GAIL. Deconfounded algorithm collects a batch episodes with the expert policy for BC and batch episodes with the imitation policy for training the encoder and stores those samples in separate buffers. The network architectures are the same as for the other policies, where applicable. We found that getting GAIL to work at all in our environments, we had to increase the number of episodes sampled between cach update from 16 to 64, Furthermore, we found that GAIL did not work with reward function parameterized as function (s, a, 8') in our environments, but worked better as function of just (s').
Hyperparameter settings The hyperparameters used for the learning algorithms are presented in ta- ble Table 2
22

---

Published in Transactions on Machine Learning Research (09/2024)
Hyperparamcter
Value
Top-level training iterations
300
PPO training iterations per algorithm iteration
1
PPO batch size (timesteps) PPO gradient max norm
400
0.5
PPO normalize advantage
True
PPO epochs per training iterations
10
PPO discount rate PPO GAE-A
0.99
0.95
PPO learning rate
0.0003
PPO episodes per training iteration
64
Discriminator training sleps per top-level iteration
20
Discriminator batch size (full episodes)
16
Discriminator maximum gradient norm
100
Table 3: Hyperparameter settings for GAIL.
23


Reviews:
[{"review_id": "nmnDsmJOmr", "reviewer": "Reviewer_ZAFY", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "1.\tThis paper studies the challenging problem where naïve imitation learning algorithms fail in the presence of latent confounders. To address the challenge, this work first studied under which conditions latent confounding impacts the performance of BC. Informed by theoretical results, this work proposed a practical algorithm for deconfounding imitation learning with amortized variational inference. The experiments in multi-armed bandits and multiple control environments demonstrated that the proposed method can learn interventional policies. \n\n2.\tThe proposed algorithm can get around the confounding problem without query access to the expert or IRL. In addition, the algorithm converges to the expert’s policy in theory, which means the expert’s policy can be identified under strong assumptions. The latent variable was inferred using variational inference and used for training imitation learning policies.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper considers imitation learning problems when the expert demonstrators have different sensory inputs than the imitating agent. In other words, the expert knows more about the world than the imitator. The main challenge comes from the hidden confounders in the casual graph caused by partial observability. To address this challenge, this paper proposes to train a variational inference model to train a latent-conditional policy. The proposed method is evaluated in the multi-armed bandit problem and multiple control environments.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1.\tThis work makes a strong assumption that the latent variable is constant in each trajectory. This assumption makes the dynamics and expert policy stationary in each trajectory and the interventional policy becomes identifiable. However, this strong assumption may limit the generalization of the proposed algorithm. \n\n2.\tThere is no clear definition of Confounded repeated multi-armed bandit and Confounded sequential decision making problems. \n\n3.\tThe experiments were only conducted in low-dimensional environments.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "81C5m91vmP", "reviewer": "Reviewer_QK3p", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n1. Although the derivation of of the algorithm is based on strong assumptions, the authors validate the main assumptions in the experiment section.\n2. The motivation of this method is well-explained, particularly with the example in section 3.3\n3. The proposed method is intuitive and supported by experiments.\n\nWeaknesses:\n1. Some details of the method are not clear: \n    1) As described in 4.1, the identifiability of the interventional policy is based on several crucial assumptions. One assumption is that the latent variable theta is static. Can the authors elaborate on what that means? Static with respect to the environment or static with respect to a single trajectory?\n    2) How is the exploration policy obtained? \n    3) How many expert demonstrations are required to learn the approximated dynamics model for each environment? \n    4) In the experiment section, how are the true interventional policy and the true conditional policy obtained?\n2. In Figure 5, the proposed method is generally worse than DAgger. Can the authors include an analysis of the benefits of this method over DAgger?\n3. The authors discussed this method's limitations in section 7. A more detailed discussion about how this method would involve solving real-world tasks with confounding variables is encouraged.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes to learn an imitation learning policy with partial observability of the environment. The authors first explain the naive BC algorithm, and how it would fail under partial observability. Then, it gives a theoretical analysis of how to learn an interventional policy and a practical implementation of it. The key is to approximate the state dynamics using an expert demo and learn an inference model that maps exploration trajectories to the confounding variable theta. In the experiment section, the authors show two sets of simulated experiments that show that the proposed algorithm recovers the interventional policy.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. Figure 3 is confusing. What's the difference between the dotted line and the solid line? Why is the exploration policy coming from the blue blocks on the right? Could the authors explain each component more clearly?\n2. A more detailed discussion of the abovementioned weaknesses could strengthen the paper.", "broader_impact_concerns": "NA", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "pPm16r6M9c", "reviewer": "Reviewer_robv", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "### Strengths\n- In example of the self-driving car in the Introduction was insightful and provides a good grounding to understand the core problem that the paper tries to tackle.\n- Sufficient summary of the literature.\n- The theoretical justification (if correct -> I was not able to check the correctness of the theory) followed by experimental work is a nice progression.\n- Running experiments on 10 random seeds is good.\n### Weaknesses\nSee requested changes.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper tackles the issues of confounding latent variables in imitation learning. In particular, this is where there are unknown facts that the expert is aware of that affect their policy, but the imitator is not. The authors propose a solution by learning a _interventional policy_ using variational inference. They make a theoretical case for this idea (under some assumptions) and showcase it in practice (where the assumptions are relaxed). They run experiments on a multi-armed bandit and various other more complicated environments.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "- I think this paper requires a pass for grammar and typos. Further, the writing could be made clearer, for example, by breaking some long sentences up. For example:\n    - Page 1: _This enables the agent to make an inference about the value of the **latent** based on everything it has observed so far, which eventually allows it to break ties over the actions in states where the **latent** matters_; latent what? Do you mean **latent variable**?\n    - Page 4: _The goal of imitation learning in this setting is to learn imitator parameters η such that when the imitator is executed in the environment, the imitator agrees with the expert’s decisions, meaning we wish to maximise Eθ∼P(θ) Eτ∼Pη(τ;θ)[Pst,at∈τ logπexp(at | st,θ)]._ I think this sentence should be broken up.\n    - Page 4: _We say **that current** state $s_t$_:  that **the** current...\n    - Page 4: _next expert’s action at are confounded**,** because the distribution_: remove the comma.\n    - In some equations with multiple terms (e.g. Equation 2 and Equation 3, can the spacing after commas be increased for readability?\n    - Page 10: _The expert policy is defined in Eq. **equation** 2._: Remove **equation**.\n    - Page 10: _The hyperparameters for the algorithms are provided in **App**. D._: In a previous part, you used **Appendix**. Be consistent! Similarly, for Figure(s) vs Fig(s).\n\n- The following statements need to be supported better, either by another reference or a broader discussion:\n     - Page 2: _This problem can be a source of “hallucinations” in generative models including large language models._\n     - Page 2: _Acting upon a mistaken belief about the latent variable propels the agent out of the support of the expert data._ \n     - Page 5: _In a causal framework, this corresponds to treating the choice of past actions as interventions._\n     - Page 10: _Even the naive learner needs a memory to be able to adapt to a new instance of the bandit defined by the latent._\n\n### Suggested changes for clarity\n\n- Page 2: _query access to the expert policy_: what does **query access to the expert policy** mean exactly? Since you are saying that your solution doesn't require this **query access** I think it would be good to describe this more. Query according to what?\n- Page 4: _distribution over at when we intervene on $s_t$_: mathematically, what does it mean to **intervene$ on $s_t$? \n- In Section 3.3, Equation 1, is $p(\\theta)$ assumed to be known? Is this the same as $P(\\theta)$ or is this a prior over the latent variables? If it is the latter, I think this should be stated in the text.\n- Page 5: _A solution to this issue is to only take as evidence the data that was actually informed by the latent, which are the transitions._; This sentence is unclear. Could you expand on it?\n- Page 6: _First, we assume that the latent variable θ of the MDP is **static**_: what exactly does **static** mean here? Do you mean that the latent variable is fixed within an episode? I think writing it out more explicitly would be better.\n- Page 6: What is the definition of **recurrent** here? Does Assumption 1 state a definition for a **recurrent MDP** or is it stating an implication: if all pairs $s, a$ are reached infinitely often in each trajectory, then the MDP is recurrent; is this implication the definition of a recurrent MDP? I think I am confused by this statement being an **Assumption**. \n- Section 4: What are the key practical takeaways from the theory?\n- Page 10: _As we are only interested in the effects of confounding on imitation learning, and not in **finite-data effects**,_: what are some examples of **finite-data** effects?\n- Figure 4: What do the sub-captions mean? It is confusing.\n\n- Section 6: I think the experiments need to be described better. What exactly are the differences between the setups in Figure 4(a) and Figure 4(b)?\n\n- In Figure 5, the legend should be moved outside of the figure bounding boxes (for example to be above all of them) since it is applied to all the sub-figures.\n\n- It would be good to consistently use either **Naive BC** or **Naive Imitation Learning**.\n\n\n\n### Minor changes\n- In Figure 1, can you change the last sentence to: _The expert action depends on the latent variable θ **(red arrows)** whereas the imitator action does not_.\n- Section 3.3 Heading: I think it would better to expand out BC to Behavioural Cloning.", "broader_impact_concerns": "N/A.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The reviewers provided detailed suggestions for revisions, which have been addressed by the authors' responses. My recommendation is based on the fact that no further revisions were suggested in the final recommendations.
Decision: Accept (poster)