🤖 machine learning

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

This paper demonstrates that inverse dynamics models (IDMs) achieve superior sample efficiency in semi-supervised imitation learning by operating within a lower-complexity, less stochastic hypothesis class than expert policies, leading to the proposal of an improved latent action policy learning algorithm validated across multiple benchmarks.

Original authors: Sacha Morin, Moonsub Byeon, Alexia Jolicoeur-Martineau, Sébastien Lachapelle

Published 2026-07-03

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Sacha Morin, Moonsub Byeon, Alexia Jolicoeur-Martineau, Sébastien Lachapelle

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a robot how to play a video game or move a robotic arm. You have two types of data available:

The "Gold Standard" Data: A small pile of videos where you can see exactly what the robot did (the actions) to get from point A to point B. This is expensive and hard to get because it requires a human expert to record every move.
The "Wild" Data: A massive mountain of videos showing the robot moving from A to B, but with no labels on what specific buttons were pressed or how the arm moved. This is easy to get (just record the screen), but you don't know how the robot did it.

This paper tackles a problem called Semi-Supervised Imitation Learning. The goal is to use that tiny pile of "Gold Standard" data combined with the huge mountain of "Wild" data to teach the robot better than if we only used the small pile.

The Old Way vs. The New Way

The Old Way (Behavior Cloning):
Think of this like a student memorizing a textbook. The robot looks at a state (e.g., "I see a wall") and tries to memorize the action ("Turn left"). It learns by copying the expert directly.

The Problem: If you only have a few pages of the textbook (small data), the student memorizes poorly and fails when they see a slightly different wall.

The New Way (Inverse Dynamics Models - IDM):
Instead of asking "What should I do now?", the robot asks a different question: "If I was here, and now I am there, what must I have done to get here?"
This is called an Inverse Dynamics Model (IDM). It's like a detective looking at a crime scene (the "before" and "after" states) to figure out what the culprit (the action) did.

The paper shows two main ways to use this detective:

The Detective + The Movie Maker (VM-IDM): You train a "Movie Maker" (Video Model) to predict where the robot will go next. Then, you use the Detective to figure out what action caused that move.
The Detective as a Labeler (IDM Labeling): You use the Detective to guess the actions for the huge mountain of "Wild" data. Once the mountain is labeled, you train the robot on this new, massive dataset.

The Big Discovery: They Are the Same Team

The authors proved a surprising mathematical fact: If you have an infinite amount of "Wild" data and perfect models, both methods above end up teaching the robot the exact same thing. They call this the IDM-based policy.

Why is the Detective Better than the Student?

The paper asks: Why does the Detective method (IDM) learn faster and better than the direct copying method (Behavior Cloning), even though the Detective has to solve a harder math problem (predicting actions from two states instead of one)?

The authors argue it comes down to two reasons, using some great analogies:

1. The "Simpler Puzzle" Theory (Complexity)

The Expert (The Goal): The expert's strategy (the policy) can be incredibly complex. In a maze, the expert has to remember a long, winding path to the exit.
The Detective (The IDM): The detective only needs to figure out the immediate step. If you are at (x,y) and end up at (x+1, y), the action was simply "Right."
The Analogy: Imagine trying to memorize a whole novel (the expert's path) versus just memorizing the rule "If you move right, you go right" (the detective's rule). The rule is much simpler to learn. Because the "Detective's rule" is simpler, the robot can learn it accurately with far fewer examples.

2. The "Less Guessing" Theory (Stochasticity)

The Expert: Sometimes, an expert might be a bit random. Maybe they turn left 50% of the time and right 50% of the time because both work. This randomness makes it hard for the robot to learn the "correct" answer.
The Detective: Even if the expert is random, the result is often clear. If the robot ends up in the "Left" room, the action must have been "Turn Left." The detective doesn't have to guess the expert's mood; it just looks at the physics of the movement.
The Analogy: It's easier to learn a language from a speaker who always says the exact same word for "Apple" (low randomness) than from someone who sometimes says "Apple," sometimes "Red Fruit," and sometimes "Crunchy Thing" (high randomness). The IDM filters out the expert's randomness, making the learning signal cleaner.

What Did They Do to Prove It?

The authors didn't just talk about it; they ran experiments:

Maze Games: They showed that in complex mazes, the "Detective" method learned to solve the maze perfectly with very little data, while the "Student" method struggled.
Video Games (Procgen): They tested this on 16 different Atari-style games. The "Detective" methods (IDM Labeling and a new improved version they called LAPO+) consistently beat the standard "Student" method, especially when data was scarce.
Robotics (Push-T & LIBERO): They applied this to real-world robot tasks (pushing a T-shaped block and moving objects). Again, the IDM-based methods were more efficient at learning from limited demonstrations.

The Takeaway

The paper concludes that Inverse Dynamics Models are a super-efficient tool for learning from videos.

By shifting the focus from "What should I do?" to "What happened between these two moments?", the robot can learn the underlying rules of the world much faster. It's like teaching a child to drive: instead of memorizing every single turn on a specific road (Behavior Cloning), it's better to teach them the physics of steering and braking (Inverse Dynamics), which allows them to drive on any road with much less practice.

The authors also propose a new, improved version of an existing algorithm called LAPO+, which uses this "Detective" logic to decode hidden actions, showing that this approach is the key to unlocking better robot learning with less data.

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

Problem Definition

The paper addresses Semi-Supervised Imitation Learning (SSIL), a setting where an agent must learn a control policy from a small dataset of action-labeled expert trajectories ( $D_L$ ) and a significantly larger dataset of action-free (state-transition only) trajectories ( $D_U$ ). While Behavior Cloning (BC) is a standard supervised learning approach for imitation, it is data-hungry and struggles when labeled data is scarce. Recent SSIL methods utilize an Inverse Dynamics Model (IDM) to predict actions from state transitions $(s, s')$ . These methods generally follow two paradigms:

VM-IDM: Coupling an IDM with a Video Model (VM) trained on $D_U$ to form a policy.
IDM Labeling: Using a trained IDM to synthetically label $D_U$ , followed by standard BC on the augmented dataset.

Despite empirical success, the theoretical reasons for the superior sample efficiency of IDM-based methods over BC remain partially hypothesized (e.g., IDM being "non-causal" or simpler). This work seeks to unify these methods and provide a rigorous statistical learning theory explanation for their performance.

Methodology and Theoretical Framework

Unification of VM-IDM and IDM Labeling

The authors first establish a theoretical equivalence between VM-IDM and IDM labeling. They prove that in the limit of infinite unlabeled data ( $D_U$ ) and sufficient model capacity, both methods converge to the same policy, termed the IDM-based policy ( $\hat{\pi}_{v^*, \hat{h}}$ ).

VM-IDM: Samples a next state $\hat{s}'$ from a Video Model $\hat{v}(s' | s)$ and then an action $\hat{a}$ from an IDM $\hat{h}(a | s, \hat{s}')$ .
IDM Labeling: Trains a policy $\hat{\pi}$ to match the distribution of actions sampled from the IDM given transitions in $D_U$ .
Equivalence: Both converge to $\hat{\pi}(a|s) = \int \hat{h}(a|s, s') v^*(s'|s) ds'$ , where $v^*$ is the ground-truth video model.

Theoretical Analysis of Sample Efficiency

The paper argues that the sample efficiency advantage of IDM learning over BC stems from two primary factors regarding the ground-truth IDM ( $h^*$ ) versus the ground-truth expert policy ( $\pi^*$ ):

Reduced Complexity: The hypothesis class required to represent $h^*$ is often lower complexity than that required for $\pi^*$ .
- Evidence: In deterministic maze environments, the expert policy $\pi^*$ may require complex path-planning logic (high complexity), whereas the IDM $h^*$ can often be expressed as a simple linear classifier based on the state difference $s' - s$ .
- Implication: Learning a simpler function ( $h^*$ ) allows for better bias-variance tradeoffs, enabling high performance with fewer labeled samples.
Reduced Stochasticity: The ground-truth IDM is often less stochastic than the expert policy.
- Evidence: Even if an expert policy $\pi^*$ is stochastic (e.g., choosing between two optimal paths with equal probability), the transition $(s, s')$ often uniquely determines the action taken. Thus, $h^*(a | s, s')$ can be deterministic while $\pi^*(a | s)$ is not.
- Implication: Lower stochasticity in the target distribution reduces the variance of the estimator, leading to faster convergence with limited data.

Proposed Algorithm: LAPO+

Motivated by the finding that IDM learning is more sample-efficient than BC, the authors propose LAPO+, an improved version of the Latent Action Policy Learning (LAPO) algorithm.

Original LAPO: (1) Train Latent IDM (LIDM) and Latent Forward Dynamics Model (LFDM); (2) Train a latent policy $\tilde{\pi}$ to predict LIDM outputs; (3) Decode $\tilde{\pi}$ to actions using a head $\phi$ .
LAPO+ Modification: Inverts the order of the last two stages. Instead of decoding a latent policy, LAPO+ learns a decoding head $\phi$ directly on top of the LIDM ( $\tilde{h}$ ). This effectively treats the LIDM as a pretrained backbone for IDM learning, followed by a final BC step to map latent actions to true actions.

Experimental Results

The authors evaluate their claims and proposed methods across three benchmarks: Procgen (16 Atari-like environments), Push-T (robotic manipulation), and LIBERO-10 (long-horizon robot tasks).

1. Theoretical Validation (Maze/Grid Worlds)

Complexity: In maze environments, low-capacity VM-IDM models achieved perfect test accuracy with limited data, while low-capacity BC failed. This confirmed that $h^*$ is simpler than $\pi^*$ .
Stochasticity: In environments with stochastic experts, BC performance degraded significantly as expert stochasticity increased. In contrast, IDM-based policies remained robust because the induced IDM remained deterministic.

2. Procgen Benchmark

IDM vs. BC: IDM labeling outperformed BC by orders of magnitude in 9 out of 16 environments, particularly those with "simple" dynamics (e.g., bigfish, maze).
LAPO+ Performance: LAPO+ consistently outperformed both standard BC and the original LAPO in low-data regimes. This validates the hypothesis that leveraging the sample efficiency of IDM learning (via the LIDM backbone) improves latent action policies.

3. Manipulation Tasks (Push-T & LIBERO-10)

VM-IDM with UVA: The authors utilized the Unified Video-Action (UVA) architecture. They demonstrated that sampling actions via a VM-IDM configuration (using UVA as both VM and IDM) outperformed or matched standard policy sampling, even with very few labeled demonstrations (e.g., 10 for Push-T, 2 per task for LIBERO).
Correlation: Performance gaps between IDM-based methods and BC correlated with the complexity and stochasticity factors identified in the theoretical analysis.

Significance and Claims

The paper claims to provide a unified theoretical framework explaining why IDM-based SSIL methods often outperform Behavior Cloning. The key contributions are:

Unification: Proving that VM-IDM and IDM labeling converge to the same optimal policy under infinite unlabeled data.
Explanation of Efficiency: Attributing the sample efficiency of IDM learning to the lower complexity and lower stochasticity of the ground-truth IDM compared to the expert policy. This offers a framework for predicting when IDM-based methods will succeed.
Algorithmic Improvement: Proposing LAPO+, which restructures the latent action learning pipeline to capitalize on IDM sample efficiency, demonstrating superior performance in low-data regimes.
Empirical Validation: Extensive experiments across discrete (Procgen) and continuous (Push-T, LIBERO) control tasks confirming that IDM properties correlate with performance gains.

The authors conclude that while IDM-based policies are powerful, their success is not universal but depends on the specific statistical properties of the environment (dynamics and expert stochasticity). They acknowledge limitations, such as the assumption of curated expert data in $D_U$ and the computational cost of VM-IDM sampling, suggesting these as avenues for future work.