Efficient Imitation Learning with Conservative World Models

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Kolev, Victor, Rafailov, Rafael, Hatch, Kyle, Wu, Jiajun, Finn, Chelsea
Format: Preprint
Published: 2024
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866916359450918912
author Kolev, Victor
Rafailov, Rafael
Hatch, Kyle
Wu, Jiajun
Finn, Chelsea
author_facet Kolev, Victor
Rafailov, Rafael
Hatch, Kyle
Wu, Jiajun
Finn, Chelsea
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2405_13193
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle Efficient Imitation Learning with Conservative World Models
Kolev, Victor
Rafailov, Rafael
Hatch, Kyle
Wu, Jiajun
Finn, Chelsea
Machine Learning
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.
title Efficient Imitation Learning with Conservative World Models
topic Machine Learning
url https://arxiv.org/abs/2405.13193