A Model-Based Approach to Imitation Learning through Multi-Step Predictions

TL;DR

This work addresses imitation learning under distribution shift and measurement noise by introducing Predictive Imitation Learning (PIL), a model-based framework that leverages multi-step predictors and a consistency loss to anticipate long-horizon consequences without costly full-horizon unrolls. It provides finite-sample performance guarantees for linear time-invariant systems and demonstrates superior long-horizon robustness compared to Behavior Cloning and rollout-based IL across linear, nonlinear, and MuJoCo control tasks. The results show that horizon-aware, predictive modeling improves stability and generalization in imitation, highlighting a principled alternative to purely model-free or purely rollout-based methods. The approach has potential to impact safety-critical control and robotics by offering scalable, predictive imitation with provable guarantees.

Abstract

Imitation learning is a widely used approach for training agents to replicate expert behavior in complex decision-making tasks. However, existing methods often struggle with compounding errors and limited generalization, due to the inherent challenge of error correction and the distribution shift between training and deployment. In this paper, we present a novel model-based imitation learning framework inspired by model predictive control, which addresses these limitations by integrating predictive modeling through multi-step state predictions. Our method outperforms traditional behavior cloning numerical benchmarks, demonstrating superior robustness to distribution shift and measurement noise both in available data and during execution. Furthermore, we provide theoretical guarantees on the sample complexity and error bounds of our method, offering insights into its convergence properties.

Figures (4)

• Figure 1: Illustration of the Predictive Imitation Learning \ref{['eq:mpc-proposed']} for horizon $H=3$.
• Figure 2: Maximum trajectory discrepancy \ref{['eq:im-gap']} relative to $\widehat{K}_{\mathrm{BC}}$, for varying prediction orders under high state measurement noise (left) and high action measurement noise (right). Results are averaged over $100$ seeds, with the shaded area indicating half standard deviation.
• Figure 3: Maximum trajectory discrepancy \ref{['eq:im-gap']} for BC, Rollout-based approach and PIL, for varying prediction orders. Results are averaged over $50$ seeds, with the shaded area indicating half standard deviation.
• Figure 4: Maximum trajectory discrepancy \ref{['eq:im-gap']} and episode returns (normalized by expert returns) for BC, PIL, and rollout-based approach across different training trajectory counts. Results are averaged over 5 seeds, with shaded regions indicating half a standard deviation.

Theorems & Definitions (18)

• Theorem 1: learning error, sketched
• proof : Proof Sketch
• Theorem 2: comparison of $\widehat{K}_{\mathrm{PIL}}$ against $\widehat{K}_{\mathrm{BC}}$, sketched
• proof : Proof Sketch
• Lemma 3
• proof
• Lemma 4
• proof
• Lemma 5
• proof