Deconfounding Imitation Learning with Variational Inference

TL;DR

This paper addresses the failure of naive imitation learning when the expert observes latent information $\theta$ not available to the imitator, causing causal confounding. It introduces a variational inference framework that learns an inference model $q_\phi(\hat{\theta}|\tau)$ and a latent-conditional policy $\pi_\eta(a|s,\hat{\theta})$ to realize the interventional policy $\pi_{int}$, with theoretical identifiability under strong assumptions and a practical online-offline learning strategy using exploration data. The approach circumvents the need for expert queries or IRL by jointly training a dynamics model $p_\psi(s'|s,a,\hat{\theta})$ and leveraging ELBO-based objectives to recover the latent and align the imitator with the expert under the latent. Empirically, the method deconfounds imitation in a multi-armed bandit and several confounded MDPs, outperforming naive BC and competitive with or surpassing IRL-based baselines in key settings, thereby offering a scalable path to high-fidelity imitation when full observability is unavailable.

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.

Figures (7)

• Figure 1: Bayes nets for (a) an expert trajectory and (b) an imitator trajectory. The expert action depends on the latent variable $\theta$ (red arrows) whereas the imitator action does not.
• Figure 2: Actions from rollouts from bandit environment defined by \ref{['eq:bandit']}. The $x$-axis is episode time. In the $y$-axis five roll-outs are shown from the expert and policies \ref{['eq:conditional-policy']} and \ref{['eq:interventional-policy']}. 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.
• Figure 3: Overview of the method. On the left, the inference model consisting of the encoder $q_\phi$ and decoder $p_\psi$ is trained using variational inference. Training data is sampled from the environment using an exploratory policy. The dotted arrow depicts how the exploratory policy can optionally be the imitation policy $\pi_{\eta}$ 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 is used for inferring the latent variable on the expert trajectories.
• Figure 4: Imitation learning in 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 is the step on the trajectory and the y-axis is 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 a trajectory with 100 time steps. The curves are averages for sliding window of length 10 training iterations.
• Figure 5: Experiments in our confounded, stochastic environments. We show the episodic return of 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.

Theorems & Definitions (2)

• Theorem 1
• proof