Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Efficient Imitation Learning with Conservative World Models

Victor Kolev, Rafael Rafailov, Kyle Hatch, Jiajun Wu, Chelsea Finn

TL;DR

This work tackles reward-free imitation learning by reframing it as offline pre-training with online fine-tuning using conservative model-based methods. It introduces CMIL, a conservative model-based adversarial imitation learning algorithm that leverages variational world models, a state-action discriminator, and an uncertainty-penalized model rollout reward to stabilize training and improve sample efficiency. The authors derive a theoretical bound combining distribution-matching error and model-mismatch error, and provide empirical evidence on high-dimensional manipulation tasks showing state-of-the-art performance with limited demonstrations (e.g., Franka Kitchen from images with 10 demos). The approach yields more robust performance and practical applicability for reward-free, pixel-based manipulation, offering a principled path for reward-free fine-tuning of world-model agents.

Abstract

We tackle the problem of policy learning from expert demonstrations without a reward function. A central challenge in this space is that these policies fail upon deployment due to issues of distributional shift, environment stochasticity, or compounding errors. Adversarial imitation learning alleviates this issue but requires additional on-policy training samples for stability, which presents a challenge in realistic domains due to inefficient learning and high sample complexity. One approach to this issue is to learn a world model of the environment, and use synthetic data for policy training. While successful in prior works, we argue that this is sub-optimal due to additional distribution shifts between the learned model and the real environment. Instead, we re-frame imitation learning as a fine-tuning problem, rather than a pure reinforcement learning one. Drawing theoretical connections to offline RL and fine-tuning algorithms, we argue that standard online world model algorithms are not well suited to the imitation learning problem. We derive a principled conservative optimization bound and demonstrate empirically that it leads to improved performance on two very challenging manipulation environments from high-dimensional raw pixel observations. We set a new state-of-the-art performance on the Franka Kitchen environment from images, requiring only 10 demos on no reward labels, as well as solving a complex dexterity manipulation task.

Efficient Imitation Learning with Conservative World Models

TL;DR

This work tackles reward-free imitation learning by reframing it as offline pre-training with online fine-tuning using conservative model-based methods. It introduces CMIL, a conservative model-based adversarial imitation learning algorithm that leverages variational world models, a state-action discriminator, and an uncertainty-penalized model rollout reward to stabilize training and improve sample efficiency. The authors derive a theoretical bound combining distribution-matching error and model-mismatch error, and provide empirical evidence on high-dimensional manipulation tasks showing state-of-the-art performance with limited demonstrations (e.g., Franka Kitchen from images with 10 demos). The approach yields more robust performance and practical applicability for reward-free, pixel-based manipulation, offering a principled path for reward-free fine-tuning of world-model agents.

Abstract

We tackle the problem of policy learning from expert demonstrations without a reward function. A central challenge in this space is that these policies fail upon deployment due to issues of distributional shift, environment stochasticity, or compounding errors. Adversarial imitation learning alleviates this issue but requires additional on-policy training samples for stability, which presents a challenge in realistic domains due to inefficient learning and high sample complexity. One approach to this issue is to learn a world model of the environment, and use synthetic data for policy training. While successful in prior works, we argue that this is sub-optimal due to additional distribution shifts between the learned model and the real environment. Instead, we re-frame imitation learning as a fine-tuning problem, rather than a pure reinforcement learning one. Drawing theoretical connections to offline RL and fine-tuning algorithms, we argue that standard online world model algorithms are not well suited to the imitation learning problem. We derive a principled conservative optimization bound and demonstrate empirically that it leads to improved performance on two very challenging manipulation environments from high-dimensional raw pixel observations. We set a new state-of-the-art performance on the Franka Kitchen environment from images, requiring only 10 demos on no reward labels, as well as solving a complex dexterity manipulation task.
Paper Structure (21 sections, 3 theorems, 20 equations, 3 figures)

This paper contains 21 sections, 3 theorems, 20 equations, 3 figures.

Table of Contents

  1. Introduction
  2. Related Work
  3. World models and RL
  4. Conservative MBRL
  5. Model-based imitation learning
  6. Preliminaries
  7. POMDP
  8. Variational Dynamics Models
  9. Imitation learning as divergence minimization
  10. Adversarial Imitation Learning
  11. Conservative Model-Based Adversarial Imitation Learning
  12. Extension to POMDPs
  13. Our Method
  14. Adversarial Formulation
  15. Actor Optimization
  16. ...and 6 more sections

Key Result

proposition 1

Let $V^{{\pi}}_{\mathcal{M}}$ denote the expected return of a policy ${\pi}$ in $\mathcal{M}$. We we can then bound the sub-optimality of any policy ${\pi}$ as: where $R_{\max}$ is the maximum reward in the underlying MDP and $\mathbb{D}_{TV}$ is total variation distance.

Figures (3)

  • Figure 1: The suite of environments. Left: ShadowHand Baoding Balls. Right: Franka Kitchen.
  • Figure 2: Training curve of success rate of our approach, CMIL, vs. baselines.
  • Figure 3: Empirical estimates of elements from performance gap bound.

Theorems & Definitions (5)

  • proposition 1
  • theorem 1
  • proof
  • theorem 2
  • proof