On the Sample Efficiency of Inverse Dynamics Models for Semi-Supervised Imitation Learning

TL;DR

The paper investigates semi-supervised imitation learning (SSIL) where a small labeled dataset and a large unlabeled dataset are available. It shows that VM-IDM and IDM labeling converge to the same IDM-based policy in the infinite unlabeled-data limit, and that with sufficient labeled data the IDM-based policy recovers the expert. The authors argue that IDM learning is more sample-efficient than BC because the ground-truth IDM is typically simpler and less stochastic than the expert policy, and they support this with maze, ProcGen, and manipulation experiments. They also propose improvements to latent-action policies (LAPO/LAPO+) and demonstrate benefits of current architectures like UVA in IDM-based SSIL, highlighting practical gains for sample-efficient imitation in complex tasks.

Abstract

Semi-supervised imitation learning (SSIL) consists in learning a policy from a small dataset of action-labeled trajectories and a much larger dataset of action-free trajectories. Some SSIL methods learn an inverse dynamics model (IDM) to predict the action from the current state and the next state. An IDM can act as a policy when paired with a video model (VM-IDM) or as a label generator to perform behavior cloning on action-free data (IDM labeling). In this work, we first show that VM-IDM and IDM labeling learn the same policy in a limit case, which we call the IDM-based policy. We then argue that the previously observed advantage of IDM-based policies over behavior cloning is due to the superior sample efficiency of IDM learning, which we attribute to two causes: (i) the ground-truth IDM tends to be contained in a lower complexity hypothesis class relative to the expert policy, and (ii) the ground-truth IDM is often less stochastic than the expert policy. We argue these claims based on insights from statistical learning theory and novel experiments, including a study of IDM-based policies using recent architectures for unified video-action prediction (UVA). Motivated by these insights, we finally propose an improved version of the existing LAPO algorithm for latent action policy learning.

Figures (6)

• Figure 1: Varying the environment complexity. Comparing BC and VM$^*\!$-IDM, with different architectures for $\hat{\pi}$ and $\hat{h}$, respectively: LC is a linear classifier, 5L MLP is a 5-layer multilayer perceptron and $n$L CNN is an $n$-layer convolutional neural network. See Appendix \ref{['app:maze_exp']} for details. Averaging over 5 seeds.
• Figure 2: Varying the number of goals. BC$_G$ is behavior cloning with goal conditioning, VM$^*_G$-IDM is VM$^*\!$-IDM where the VM is goal conditioned and the IDM is not; and VM$^*_G$-IDM$_G$ is when both the VM and the IDM are goal conditioned. See Appendix \ref{['app:maze_exp']} for details. Averaging over 5 seeds.
• Figure 3: Varying the stochasticity of the expert. First row shows the state visitation distributions for different experts. Comparing the average reward of BC and IDM labeling. See Appendix \ref{['app:maze_exp']} for details. Averaging over 10 seeds.
• Figure 4: Normalized maximum return achieved during training by BC, IDM labeling, LAPO and LAPO+ in the 16 Atari-like ProcGen environments cobbe2019procgen, varying the number of action-labeled transitions in ${\mathcal{D}_{L}}$. We show the average of 3 seeds, with ${\mathcal{D}_{L}}$ being randomly sampled from the whole datasets each time. IDM labeling, LAPO and LAPO+ use all 2.6M transitions without action labels as ${\mathcal{D}_{U}}$.
• Figure 5: We study the complexity and stochasticity factors from Section \ref{['sec:stat_advantage_idm']} across the 16 ProcGen environments. As an aggregate score, we consider the mean return gap$r(\hat{\pi}_{IDM}) - r(\hat{\pi}_{BC})$ between IDM labeling and BC, which we average over ${\mathcal{D}_{L}}$ sizes (Figure \ref{['fig:lapo_procgen']}). Complexity Analysis. We divide environments based on their dynamics (Appendix \ref{['app:procgen_complexity']}): simple environments are those the agent moves strictly because of an input action while the agent can move due to other factors in complex environments (e.g., momentum, gravity, moving platform). We expect the IDM to be simpler under simple dynamics and find the mean return gap to be higher in those environments. Stochasticity Analysis. As a proxy for the entropy of the ground truth policy and IDM in each environment, we consider the average conditional entropies on a test set of $\hat{\pi}_{BC}$ and the IDM $\hat{h}$ learned on all 2.6M transitions, allowing us to study an approximate entropy gap $H(\hat{\pi}_{BC}) - H(\hat{h})$. We first note that $H(\hat{\pi}_{BC}) > H(\hat{h})$ across all environments, as hinted in Section \ref{['sec:stochastic_expert']}. We then observe a positive, although imperfect correlation with the mean return gap.

Theorems & Definitions (6)

• Claim 1
• Claim 2
• Remark 3.1
• Remark 3.2
• Remark 3.3
• Remark 3.4