Deconfounded Imitation Learning
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.
We break down the space of confounded imitation learning problems and identify three settings with different data requirements in which the correct imitation policy can be identified.
We then introduce an algorithm for deconfounded imitation learning, which trains an inference model jointly with a latent-conditional policy.
At test time, the agent alternates between updating its belief over the latent and acting under the belief.
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.
Introduction
In imitation learning (IL), an agent learns a policy directly from expert demonstrations without requiring the specification of a reward function.
This paradigm could be essential for solving realworld problems in autonomous driving and robotics where reward functions can be difficult to shape and online learning may be dangerous.
However, standard IL requires that the conditions under which the agent operates exactly match those encountered by the expert.
In particular, they assume that there are no latent confounders —variables that affect the expert behavior, but that are not observed by the agent.
This assumption is often unrealistic.
Consider a human driver who is aware of the weather forecast and lowers its speed in icy conditions, even if those are not visible from observations.
An imitator agent without access to the weather forecast will not be able to adapt to such conditions.
In such a situation, an imitating agent may take their own past actions as evidence for the values of the confounder.
A self-driving car, for instance, could conclude that it is driving fast, thus there can be no ice on the road.
This issue of causal delusion was first pointed out in Ortega and Braun [2010a,b] and studied in more depth by Ortega et al.
The authors analyze the causal structure of this problem and argue that an imitator needs to learn a policy that corresponds to a certain interventional distribution.
They then show that the classic DAgger algorithm [Ross et al., 2011], which requires querying experts at each time step, solves this problem and converges to the interventional policy.
In this paper, we present a solution to a confounded IL problem, where both the expert policy and the environment dynamics are Markovian.
The solution does not require querying experts.
We first present a characterization of confounded IL problems depending on properties of the environment and expert policy (Section 3).
We then show theoretically that an imitating agent can learn behaviors that approach optimality when the above Markov assumptions and a recurrence property hold.
We then introduce a practical algorithm for deconfounded imitation learning that does not require expert queries (Section 4).
An agent jointly learns an inference network for the value of latent variables that explain the environment dynamics as well as a latent-conditional policy.
At test time, the agent iteratively samples latents from its belief, acts in the environment, and updates the belief based on the environment dynamics.
An imitator steering a self-driving car, for instance, would learn how to infer the weather condition from the dynamics of the car on the road.
This inference model can be applied both to its own online experience as well as to expert trajectories, allowing it to imitate the behavior adequate for the weather.
[Figure]
Finally, our deconfounded imitation learning algorithm is demonstrated in a multi-armed bandit problem.
We show that the agent quickly adapts to the unobserved properties of the environment and then behaves optimally (Section 5).
Imitation learning and latent confounders
We begin by introducing the problem of confounded imitation learning.
Following Ortega et al.
[2021], we discuss how behavioral cloning fails in the presence of latent confounders.
We then define the interventional policy, which solves the problem of confounded imitation learning.
2.1 Imitation learning
Imitation learning learns a policy from a dataset of expert demonstrations via supervised learning.
The expert is a policy that acts in a (reward-free)
Markov decision process (MDP) defined by a tuple M = ( S , A , P ( s ′ | s, a ) , P ( s 0 )) , where S is the set of states, A is the set of actions, P ( s ′ | s, a ) is the transition probability, and P ( s 0 ) is a distribution over initial states.
The expert’s interaction with the environment produces a trajectory τ =
( s 0 , a 0 , . .
. , a T − 1 , s T ) .
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 et al.
In the simplest form of imitation learning, a behavioral cloning policy π η ( a | s ) parametrized by η is learned by minimizing the loss − (cid:80) s,a ∈D log π η ( a | s ) , where D is the dataset of state-action pairs collected by the expert’s policy.
2.2 Confounded imitation learning
We now extend the imitation learning setup to allow for some variables θ ∈ Θ 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 ( θ ) , and for each θ ∈ Θ , a reward-free MDP M θ = ( S , A , P ( s ′ | s, a, θ ) , P ( s 0 | θ )) .
We assume there exists an expert policy π exp ( a | s, θ ) for each MDP.
When it interacts with the environment, it generates the following distribution over trajectories τ :
[Equation]
The imitator does not observe the latent θ .
It may thus need to implicitly infer it from the past transitions, so we take it to be a non-Markovian policy π η ( a t | s 1 , a 1 , . . .
, s t ) , parameterized by η .
The imitator generates the following distribution over trajectories:
[Equation]
The Bayesian networks associated to these distributions are shown in Figure 1.
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
[Equation]
If the expert solves some task (e. g. maximizes some reward function), this amounts to solving the same task.
2.3 Naive behavioral cloning
If we have access to a data set of expert demonstrations, one can learn an imitator via behavioral cloning on the expert’s demonstrations.
At optimality, this learns the conditional policy :
[Equation]
Following Ortega et al.
[2021], consider the following example of a confounded multi-armed bandit with A =
Θ = { 1 , . .
, 5 } and S = { 0 , 1 } :
[Equation] [Equation]
The expert knows which bandit arm is special (and labeled by θ ) and pulls it with high probability, while the imitating agent does not have access to this information.
If we roll out the naive behavioral cloning policy in this environment, shown in Figure 2, we see the causal delusion at work.
At time t , the latent that is inferred by p cond takes past actions as evidence for the latent variable.
This makes sense on the expert demonstrations, as the expert is cognizant of 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 to repeat that action, as it mistakenly takes the first action to be evidence for the latent.
2.4 Interventional policy
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 defines the following imitator policy:
[Equation]
In a causal framework, that corresponds to treating the choice of past actions as interventions.
In the notation of the do-calculus [Pearl, 2009], this equals p ( a t | s 1 , do( a 1 ) , s 2 , do( a 2 ) , . . .
, s t ) .
The policy in Equation (5) is therefore known as interventional policy [Ortega et al., 2021].
[Figure]
We now present our theoretical results on how imitation learning can be deconfounded.
We first show that the interventional policy is optimal in some sense, before analyzing in which settings it can be learned.
3.1 Optimality of the interventional policy
Under some reasonable assumptions, the interventional policy approaches the expert’s policy, as we prove in the Appendix B.
Theorem 3.1 (Informal) .
If the interventional inference p int ( θ | τ <t ) approaches the true latent of the environment as t → ∞ on the rollouts of π int , and if the expert maximises some reward that is fixed across all environments, then as t → ∞ , the imitator policy π int ( a t | s ) approaches the expert policy.
Proof.
See Lemma 2.1 in the appendix.
The requirement here means that the transition dynamics must be informative about the latent — we consider latent confounders that manifest in the dynamics, not those that affect only the agent behavior.
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 only guarantees asymptotic optimal performance.
However, it is not guaranteed to be the policy that adapts to a given test environment in the fastest possible way.
Efficient exploration may require learning behavior that the expert never demonstrated Zhou et al.
3.2 Identification
The rest of this section concerns the question: under which assumptions on the model and data availability can we identify the interventional policy?
Under the typical rules of do-calculus, identification should be impossible without observing the latent.
However, we discern three tiers of decreasing strength of model assumptions and increasing data requirements in which identification is nevertheless possible without observing the latent.
In this way, we split the confounding problem discussed by Ortega and Braun [2010a,b], Ortega et al.
[2021] into three separate classes, provide conditions for when each class applies, and weaken the requirements for identifiability for two of these classes.
Tier 1: identifiability from demonstrations
In the first tier, we make the strongest assumption and show that the interventional policy is identifiable only from the expert’s demonstrations.
Theorem 3.2 (Informal) .
Let the expert’s demonstrations be recurrent, meaning that each trajectory contains every state and action pair infinitely often.
Then the interventional policy is identifiable only from expert demonstrations.
Proof.
See Theorem A.2 in the appendix.
The assumption of recurrence is for instance satisfied if for any state, the expert policy has full support over actions.
In that case, on one trajectory, with one value of the latent θ , we are able to accurately estimate the distributions over both the dynamics p ( s ′ | s, a, θ ) and the expert π exp ( a | s, θ ) .
Then from the distribution over trajectories, we estimate the distribution over distributions and thus identify p ( θ ) and the θ -conditional dynamics and expert distributions.
The interventional policy can then be constructed from these models via Equation (5).
In practice, this method suffices when the demonstrations are of sufficient length and the expert and dynamics are sufficiently noisy to sufficiently cover the state and action space in each trajectory.
In case the expert’s demonstrations are not recurrent, but we are able to execute an exploratory policy in the environment, the interventional policy may still be identifiable:
Theorem 3.3 (Informal) .
Consider the setting in which in addition to the expert demonstration data, we can collect trajectories by rolling out an exploration policy.
Let the exploration data be recurrent.
Let the true latents be identifiable from the expert trajectories given the true state-transition model p ( s ′ | s, a, θ ) .
Then the interventional policy is identifiable from expert demonstrations and exploration data.
Proof.
See Theorem A.4 in the appendix.
How can we identify the interventional policy from expert data and exploration data?
Under the assumption of recurrence of exploration trajectories, which is satisfied if the exploration policy has full support and any state is reachable from any state in the environment, the interventional policy can be identified in the following way.
First, we estimate the distribution over the state transitions from the exploratory data.
This allows us to construct the interventional inference model p int ( θ | τ ) from equation (5).
With that model, we can infer the latent for any expert trajectory and learn a policy π η ( a | s, θ ) with behavioral cloning on the demonstrations.
We need to make the crucial assumption that the inference collapses to the true latent of the dynamics — in other words, the state transitions encountered by the expert are informative about the latent variables.
(If the expert only visits states where the latent variables do not influence the transition dynamics, it stands to reason that we cannot deconfound the imitator.)
Under this assumption, this method identifies the interventional policy, as we prove in Theorem A.4 in the appendix.
We will demonstrate the identifiability of the interventional policy in this tier in practice in the following section.
Tier 3: identifiability from expert queries
If the dynamics of the environment are such that no exploratory policy leads to recurrent trajectories, then the interventional policy can only be identified if we can interact in the environment and query the expert for its decisions.
The procedure is to learn a recurrent policy π η ( a t | s 1 , a 1 , . . .
, s t ) , execute it in the enviroment, query the expert — aware of the true latent — for its actions, and train the imitator with maximum likelihood to match those actions.
This is similar to the setup in DAgger [Ross et al., 2011].
In Ortega et al.
[2021], it is shown that at optimality this converges to the interventional policy.
Examples for problems that fall into tier 3 include cases where the environment dynamics is not
Markov given the latent, where the expert policy is not Markov given the latent, or where the expert only visits regions of the state space in which the dynamics are insensitive to the latent.
Practical algorithm
We now introduce a practical algorithm for training an agent from expert data in the presence of latent confounders.
For simplicity, we focus on tier 2 of the settings discussed in Section 3, i. e. assume access to the expert data as well as the ability to gather more data interactively.
In Appendix C, we describe an algorithm that applies to tier 1 (i. e. does not require the ability to gather more data interactively), but faces a more difficult learning problem in practice.
As outlined in Section 3, we use the access to the family of MDPs defined by the latent θ to learn an inference model for the latent variable, infer the latents on the expert trajectories using the learned model, and learn a policy conditioned on the inferred latent, which imitates the expert based on the demonstrations.
At test time, the agent uses posterior sampling (or Thompson sampling)
[Thompson, 1933]: it alternates between updating its belief about which MDP it is in and acting optimally under its current belief.
Components The agent consists of: (1) an inference model q ϕ , which maps trajectories τ = ( s 0 , a 0 , . . .
, s n )
to a belief over a latent variable q ϕ
(ˆ θ | τ ) ; (2) a dynamics model p ψ , which maps latent, state, and action to a distribution over the next state p ψ ( s ′ | s, a, ˆ θ ) ; (3) a prior over latents p
Training the inference model
The inference model q ϕ , dynamics model p ψ , and prior p (ˆ θ ) form a latent variable model for trajectory data.
Similarly to a variational autoencoder (VAE) [Kingma and
Welling, 2013], they are trained by minimizing a variational inference objective 4 given by
[Equation]
At timestep t , the encoder q ϕ takes as input the trajectory observed until timestep t , τ : t = ( s 0 , a 0 , . .
. , a t − 1 , s t ) , and predicts a distribution of the belief over the latent ˆ θ .
The decoder p ψ is a dynamics model, which predicts the next state s t +1 given the state s t , action a t , and a sample from the belief distribution.
Unlike a VAE, the proposed latent variable model is not an autoencoder, since instead of decoding the input data distribution, the dynamics model p ψ decodes the future states.
To minimize the loss in Equation (6), the inference model needs to encode in the learned latent ˆ θ all of the information in the ground-truth latent θ that is important to predict the state transition.
Crucially, the training data for the inference model is collected from from the interactions of an exploration policy with the environment.
As exploration policy, we use the latent-conditional policy π η , sampling ˆ θ from the prior p (ˆ θ ) .
Other exploration policies may be used as long as they explore sufficiently diverse trajectories and do not depend on the true latent θ .
However, we cannot simply train the inference model on expert data, as those trajectories where collected from the θ -dependent expert policy.
An inference model trained on that data would exhibit the problem of causal delusions that we are trying to solve.
learned inference model is used for inferring the latents in the expert data distribution.
These inferences are valid despite the distribution shift from exploratory data to expert data, as the transition model is the same.
Standard imitation learning can then be used to learn a policy conditional on the inferred latents from the expert demonstrations.
We use behavioral cloning as the imitation learning method because it is simple and reasonably effective.
We show the full training algorithm in algorithm 1.
Test time At test time, the agent faces an environment with an unknown latent and needs to adapt to the correct expert behavior.
We solve this problem by alternating between updating a posterior belief over the latent and acting under the current belief, i. e. sampling a latent from the current belief and acting like the corresponding expert.
Concretely, the agent initially samples a latent from the prior ˆ θ ∼ p (ˆ θ ) and an action a ∼ π η ( a | s, ˆ θ ) 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.
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 algorithm 2.
In practice, using posterior sampling successfully requires that the imitator conditioned on a randomly sampled latent does not end up in states the true expert would not have visited.
For example, in a robotic manipulation task, after accidentally pushing the object outside the reach of the robot, none of the expert policies could continue the task.
A simple way of avoiding this problem is allowing multiple episodes of interaction with the task, where the environment is reset to a state sampled from the initial state distribution after each episode.
Solving this problem for arbitrary environments without resets requires methods beyond behavioral cloning; we leave this for future work.
Experiments
To test our method in practice, we conduct an experiment in the multi-armed bandit problem proposed by Ortega et al.
[2021] and described in Section 2.
With this experiment, we aim to answer the following questions.
First, how big is the effect of confounding on naive imitation learning — large enough to justify the use of specialized methods?
Second, is our algorithm capable of identifying the interventional policy?
Finally, does the interventional policy converge to the expert policy?
Algorithm 1 Behavioral cloning with latent inference
Require: Initial parameters of the imitation policy η , inference model ϕ , and dynamics model ψ , an expert dataset { τ je } , an MDP M , true latent distribution p ( θ ) , learning rates α 1 , α 2 , and α 3 .
while not done do θ ∼ p ( θ ) s 0 ∼ p 0 ( s 0 ) τ
= s 0 , t =for t ≤ H do a t ∼ π ( a t | s t )
▷ Sample action from a Markov policy, e.g. a random policy s t +1 = p
( s t +1 | s t , a t , θ ) ▷ True dynamics
Append ( a t , s t +1 ) to τ .
t =
t +end for ϕ =
ϕ − α 1 ∇ ϕ
ˆ L ( τ ) with ˆ L ( τ ) being a sample estimate of Equation 6.
ψ = ψ − α 2 ∇ ψ ˆ L ( τ ) with ˆ L ( τ ) being a sample estimate of Equation 6.
Infer latents ˆ θ jH ∼ q ϕ
(ˆ θ jH | τ je ) for expert trajectories { τ je } .
η ← η − α 3 ∇ η (cid:80) j (cid:80) s,a ∈ τ je log π η ( a | s, ˆ θ jH ) .
end while
Algorithm 2 Trained agent imitating an expert
Require:
Trained parameters of the imitation policy η , inference model ϕ an MDP M , and a latent θ .
s 0 ∼ p 0 ( s 0 ) τ
= s 0 , t =for t ≤ H do ˆ θ t ∼ q ϕ
(ˆ θ t | τ ) ▷ Infer the latent for trajectory a t ∼ π η ( a t | s t , ˆ θ t )
▷ Condition on the inferred latent s t +
= p ( s t +1 | s t , a t , θ ) ▷ True dynamics
Append ( a t , s t +1 ) to τ .
t =
t +end for
Setup In the experiments, the expert policy is defined in Equation (4).
We consider episodes of length 100.
We are only interested in the effects of confounding on imitation learning, and not issues arising from optimization.
Therefore, in order to avoid overfitting to a finite expert dataset, we resample new 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. Naive imitation learning and the conditional policy To answer the first question, we compare a naive imitation learner to the true conditional policy described in Section 2.
Even the naive learner needs a memory to be able to adapt to a new instance of the bandit defined by the latent.
To provide such a memory, we implement the imitation learner as a recurrent neural network (RNN).
From Figure 3a, we see that the naive behavioral cloning agent implemented with an RNN learns a policy matching the true conditional policy closely.
This results in problems for the policy learned with naive imitation learning when it is deployed in the environment and has to choose the actions itself, as shown in Figure 3b.
In the figure, it can be seen that the naive imitation learner does not converge to the expert policy during the episode.
Furthermore, it closely tracks the action probability of the true conditional policy, suggesting that it has suffered the full impact of the confounding problem.
Given that the naive imitator closely matches the true conditional policy at convergence, we can conclude that naive imitation learning is unlikely to be enough to solve confounded imitation learning problems.
Deconfounded imitation learning and the interventional policy To answer the second question, we implement the deconfounded imitation learner as described in Section 4.
While in principle the inference model can be learned from the data collected by the imitator, we use a random policy for [Figure] E v a l u a t i o n i m i t a t i o n l o ss
[Figure] the data collection. Figures 3a and 3b show that the proposed method (labeled “Deconfounded”) closely matches the true interventional policy both on expert trajectories and online.
From this we can conclude that in this simple environment, the proposed method produced an inference model and a conditional policy that together match the interventional policy.
Figure 3c shows the imitation loss during training.
Notice that the deconfounded algorithm achieves a lower imitation loss than naive BC even though the policy trained with naive BC collapses on the best arm more rapidly on the expert trajectories as shown in Figure 3a.
This happens because the deconfounding step in the algorithm conditions the imitator policy on the latent inferred for the whole episode at every step of the episode.
Figure 3d shows the number of times the policies chose the best arm during an episode.
This is a convenient proxy for how closely a policy matches the interventional policy.
From the figure, it can be seen that the deconfounded imitation learning algorithm closely matches the counterfactual policy, while naive BC converges to lower best arm count.
We conclude that our deconfounded imitation learning algorithm is indeed able to learn the interventional policy, solve the issue of causal delusions faced by naive behavioral cloning, and achieve near-perfect imitation performance.
This comes at the price of requiring exploration data to train the inference model as well as an increased number of training iterations needed for convergence.
Related work
Imitation learning Imitation learning (or learning from demonstration) has a long history with applications in autonomous driving [Pomerleau, 1988] and robotics [Schaal, 1999].
Standard algorithms include behavioral cloning as well as inverse reinforcement learning Russell [1998], Ng et al.
Scaling imitation learning to high-dimensional continuous control problems is more challeng- ing and has been solved through adversarial methods
[Ho and Ermon, 2016, Fu et al., 2017].
Imitation learning 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 issue is commonly addressed by querying experts during the training [Ross et al., 2011] or by noise insertion in the expert actions.
Note that this issue is qualitatively different from the one we discuss in the paper: it is a consequence of the limited support of the expert actions and occurs even in the absence of the latent confounders.
Causality-aware imitation learning De Haan et al.
[2019] diagnose the issue of causal confusion in imitation learning, in which the imitator draws causally wrong conclusions from its inputs because such patterns may be easier to learn.
This problem is different from the issue of latent confounders we discuss and may occur even if the expert and imitator have access to the same information.
The works most closely related to our paper are Ortega and Braun [2010a,b], Ortega et al.
[2021], which point out the issue of latent confounding and causal delusions that we discuss in this paper.
In particular, Ortega and Braun [2010b] propose a training algorithm that learns the correct interventional policy.
Unlike our algorithm, their approach requires querying experts during the training.
However, as we discuss in Section 3, their solution has weaker assumptions and in particular also applies to non-Markovian environment dynamics.
Kumor et al.
[2021] study imitation learning on partially observed structural causal models.
They characterize when behavioral cloning conditional on observed adjustment variables can maximize the reward.
The environments we consider do not meet these criteria, so behavioral cloning fails to imitate the expert, while our method works under additional assumptions.
In an extension, Anonymous [2022] propose that an optimal policy can be recovered from a sub-optimal expert whenever the causal effect from the policy on a hypothetical reward would be identifiable from the observables through the
I DENTIFY algorithm
[Tian and Pearl, 2002], which is not the case for the environments we consider.
Rezende et al.
[2020] point out that the same problem appears in partial models, i. e. world models that only use a subset of the (in principle fully observable) state.
The part of the state that is not modeled then acts as a confounder for the model predictions.
The authors then identify a minimal set of variables that a partial model needs to include in order to not suffer from this confounding issue.
While their and our works describe closely related problems, the solutions differ as in our problem we cannot redefine the inputs to include the confounder θ .
Meta-reinforcement learning Our problem is related to meta-reinforcement learning [Duan et al., 2016, Wang et al., 2016].
Here an agent is trained to perform well in a variety of tasks.
Both Rakelly et al.
[2019] and Zintgraf et al.
[2019] propose meta-RL algorithms that consist of a task encoder and a task-conditional policy, similar to our inference model and latent-conditional imitator policy.
One difference to our work is that these tasks can vary through the reward function, which is not compatible with imitation learning because the reward function is not available at test time.
Zhou et al.
[2019], like our work, consider multiple environments that differ through some latents.
They propose agents that learn to probe the environment to determine these latents.
The difference to our work is that they work in an RL setting with given reward functions, while our algorithm only requires expert demonstrations.
Another closely related research topic is that of meta-imitation learning [Duan et al., 2017], where the aim is to make imitation learning possible from a small number of demonstrations.
Despite both problems centering on imitation learning in a distribution of MDPs, the motivations and methods are different.
Our work does not consider adapting to new demonstrations, whereas meta-imitation learning does not consider the confounding problem in the demonstrations.
Conclusion
Naive imitation learning algorithms fail in the presence of latent confounders — for instance when the expert has access to more information than the demonstrator.
We presented a breakdown of this confounding problem based on when the interventional policy, which solves the confounding issue, is identifiable without query access to the expert.
We proposed two algorithms that provably converge to the true interventional policy.
In a multi-armed bandit experiment, we demonstrated for one of them that it is able to learn the correct interventional policy, solves the confounding problem that limits naive imitation learning, and converges to the expert behavior.
While our work makes progress theoretically, the empirical demonstration is limited to a simple toy problem.
In future work we aim to scale up the experiments to simulated robotics environments.
We also plan to adapt the algorithm to other imitation learning algorithms and to study learning objectives that incentivize efficient exploration.
A scaled-up algorithm for deconfounded imitation learning may be an important stepping stone on the way to general learning algorithms for control problems, with potentially large impact in the context of robotics and autonomous driving A Identifiability A.1 Tier Lemma A.1. Given a countable set X .
Let M ( X ) =
[ X, PX ] the set of stochastic maps from
X to X .
Note that given an initial distribution, each stochastic map gives a Markov chain.
Let
M ( X ) r ⊂ M ( X ) denote the subset of recurrent maps, meaning that for each state x ∈ X ,
its
Markov chain starting with x visits x infinitely often with probability one.
Let ϕ λ : M ( X ) → P ( X N ) denote the distribution over infinite chains starting with an initial distribution λ ∈ PX .
Let
C : X N → M ( X ) denote the map that gives a maximum likelihood estimate of the transition probabilities given an infinite chain.
Let P ( M ( X ) r ) denote the set of probabilities over M ( X ) r .
Then for all m ∈ M ( X ) r , λ ∈ PX , with probability one, τ ∼ ϕ λ ( m ) , C ( τ ) = m .
Furthermore, ∀ p ∈ P ( M ( X )) r ) , λ ∈ PX , m ∼ p, τ ∼ ϕ λ ( m ) , C ( τ ) is distributed as p .
Proof.
The first result is shown via the ergodic theorem in [Norris, 1997, Theorem 1.10.2].
The second result follows from the fact that (cid:82) A dδ a ( a ′ )
= a , where δ a is the Dirac measure around a (alternatively, the second result follows from a probability monad law).
Theorem A.2.
Given countable state and action spaces S , A , let the set Θ Π =
[ S, PA ] denote the set of stochastic policies, Θ F =
[ S × A, PS ] the set of stochastic dynamics, Θ = Θ F × Θ Π the product.
Note that each θ ∈ Θ gives a Markov chain with state S × A .
Let Θ r be the subset of recurrent Markov chains.
Let p ∈ P (Θ r ) be a joint distribution dynamics and policy giving recurrent chains.
Then p can be identified from the distribution over infinite sequences P (( S × A ) N ) given by rolling out p .
Proof.
Following lemma A.1, we can identify p ′ ∈ P ([ S × A,
P
( S × A )]) that generates the chain.
The map Θ → [ S × A, P ( S × A )] is injective, so we can recover p ∈ P (Θ r ) from p ′ .
Corollary A.3.
In Tier 1, assuming the chains are recurrent, we can identify all distributions, thus we can construct the interventional policy.
A.2 Tier
Theorem A.4.
Given p ( θ F , θ Π ) ∈ P (Θ) and p ′
( θ F , θ Π ) ∈ P (Θ) r , such that the marginal distributions over dynamics agree p ( θ F )
=
p ′ ( θ F ) .
Define on a trajectory τ , p do ( A ) ( θ F | τ ) ∝ p ( θ F ) (cid:81)
t p ( s t +1 | s t , a t , θ F ) .
Assume that with probability one, for policy and dynamics ( θ Π , θ F ) ∼ p and infinite rollout τ ∼ ϕ λ ( θ Π , θ F ) , that p do ( A ) (ˆ θ F | τ ) =
δ θ F .
Then, given trajectory distributions ϕ
Pλ ( p ) , ϕ Pλ ′ ( p ′ ) ∈ P (( S × A ) N ) , we can identify the interventional policy
[Equation]
Proof.
We identify the dynamics p ( θ F ) using the recurrent model p ′ via theorem A.2.
Then for a model p ψ ( a | s, ˆ θ F ) with parameters ψ , maximise the likelihood of
[Equation]
which by assumption is equal to
[Equation]
which is maximised by p ψ ( a | s, θ F )
[Equation]
Then: p ( a t | s 1 , do ( a 1 ) , ..., s t ) =
E ˆ θ F ∼ p do ( A )
[Equation]
τ
[Equation]
p
[Equation]
a | s, ˆ θ F ) .
B Optimality
Lemma B.1.
Given expert distribution p ( θ F , θ Π ) ∈ P (Θ) and agent distribution p ′
( θ F , θ Π ) ∈ P (Θ) , such that the marginal distributions over dynamics agree p ( θ F )
=
p ′ ( θ F ) .
Assume the expert’s policy is uniquely determined by the environment dynamics (e.g. as it is trained with RL on a single reward function on each environment), so we can write p ( a | s, θ F ) for the expert policy.
Assume that with probability one, for policy and dynamics ( θ Π , θ F ) ∼ p ′ and infinite rollout τ ∼ ϕ λ ( θ Π , θ F ) , that p ′ do ( A ) (ˆ θ F | τ ) =
δ θ F .
Then, on rollouts of p ′ , the interventional policy approaches the expert policy:
[Equation]
Proof.
[Equation]
where in the last line we used that in the limit to infinite sequences, that p do ( A ) (ˆ θ F | τ ) =
δ θ F .
C Tier 1 algorithm
The paper presents the proof that at tier 1, when suitable demonstrations from the expert are available, the interventional policy is identifiable from the expert data alone, without access to 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 as the tier 2 algorithm presented in the paper.
Training an inference model directly on the expert demonstration faces the same problem as naive imitation learning, i. e., the trained model takes the expert’s actions as evidence for the latent.
However, by the assumption stated in appendix B, that the expert is uniquely determined by the environment dynamics, we can directly learn the conditional policy and dynamics model explaining the expert trajectories from the demonstrations because a latent that explains the dynamics also explains the expert.
To train such models, we use variational inference to learn a trajectory encoder q ϕ off , which infers the latent for the expert trajectories, and a factorized decoder, which reconstructs the dynamics of the environment and the expert’s policy using networks p ψ and π η respectively.
The variational inference objective is given by
[Equation]
which, unlike the objective in the main paper, represents a VAE where the encoder q ϕ off takes as input the full expert trajectory τ ie , 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 q ϕ off for implementing the interventional policy, because it takes the expert’s actions as evidence for the latent.
The tieralgorithm works by separately learning an inference model from interactions with the environment and using that inference model for deconfounding the expert trajectories.
In tier 1, where we do not have sampling access to the environment, we cannot learn the inference model directly.
Instead, we observe that one factor of the decoder used for training the inference model is a 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 q ϕ on 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 in tier 2.
The full tier 1 algorithm is presented in algorithm 3.
Algorithm
Require:
The initial parameters of the imitation policy η , offline inference model ϕ off , online inference model ϕ on , dynamics model ψ , prior distribution for the learned latent space p (˜ θ ) , a dataset of expert trajectories { τ ie } , an MDP ( S , A , p, p 0 , H ) , learning rates α 1 , α 2 , α 3 , and α 4 .
while not done do ˜ θ ∼ p (˜ θ )
▷ Sample latent from the prior for the learned latent space s 0 ∼ p 0 ( s 0 )
▷ Sample first state from a learned model or expert data τ synth =
s 0 , t =for t ≤ H do a t ∼ π ( a t | s t )
▷ Sample action from a Markov policy s t +1 =
p ψ ( s t +1 | s t , a t , ˜ θ )
▷ Dynamics model
Append ( a t , s t +1 ) to τ synth .
t =
t +end for ϕ on =
[Equation]
on
[Equation]
[Equation]
▷ Train the online inference model to minimize equation 6.
[Equation]
▷ Train the offline inference model to minimize equation 7.
[Equation]
▷ Train the dynamics model to minimize equation 7.
[Equation]
▷ Train the imitator policy to minimize equation 7.
end while
D Experimental setup
Implementation details The inference model q ϕ is implemented as an RNN with GRU architecture Cho et al.
[2014] with a hidden layer of 256 units.
Before the RNN, the observation is preprocessed by an MLP with two hidden layers of size 256 units and output size 32.
The action is preprocessed by a linear transformation to a 32 dimensional vector.
The outputs of the RNN are processed by a linear transformation to a vector which parametrizes the latent distribution.
The latent distribution is a 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.
It 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 as the other networks and then has two hidden layers with 256 units each.
The naive behavioral cloning baseline uses the same network architecture as the deconfounded algorithm, except it does not represent the belief as a 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 and Ba, 2014] with otherwise default settings from pytorch [Paszke et al., 2019] apart from the learning rate.
Hyperparameter settings The hyperparameters used for the learning algorithms are presented in table 1.
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.
[Table]
In practice, the true belief over theta can be computed for any trajectory as follows
[Equation]
The true interventional policy can then be computed by sampling a belief ˆ θ t ∼ log p (ˆ θ t ) , and sampling an action from π exp ( a t | s t , ˆ θ t ) .
This can be seen as a Thompson sampling policy [Thompson, 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
[Equation]