Table of Contents
Fetching ...

Offline Imitation Learning with Model-based Reverse Augmentation

Jie-Jing Shao, Hao-Sen Shi, Lan-Zhe Guo, Yu-Feng Li

TL;DR

A novel model-based framework, called offline Imitation Learning with Self-paced Reverse Augmentation (SRA), which can efficiently generate trajectories leading to the expert-observed states in a self-paced style and guides maximizing long-term returns on these states, ultimately enabling generalization beyond the expert data.

Abstract

In offline Imitation Learning (IL), one of the main challenges is the \textit{covariate shift} between the expert observations and the actual distribution encountered by the agent, because it is difficult to determine what action an agent should take when outside the state distribution of the expert demonstrations. Recently, the model-free solutions introduce the supplementary data and identify the latent expert-similar samples to augment the reliable samples during learning. Model-based solutions build forward dynamic models with conservatism quantification and then generate additional trajectories in the neighborhood of expert demonstrations. However, without reward supervision, these methods are often over-conservative in the out-of-expert-support regions, because only in states close to expert-observed states can there be a preferred action enabling policy optimization. To encourage more exploration on expert-unobserved states, we propose a novel model-based framework, called offline Imitation Learning with Self-paced Reverse Augmentation (SRA). Specifically, we build a reverse dynamic model from the offline demonstrations, which can efficiently generate trajectories leading to the expert-observed states in a self-paced style. Then, we use the subsequent reinforcement learning method to learn from the augmented trajectories and transit from expert-unobserved states to expert-observed states. This framework not only explores the expert-unobserved states but also guides maximizing long-term returns on these states, ultimately enabling generalization beyond the expert data. Empirical results show that our proposal could effectively mitigate the covariate shift and achieve the state-of-the-art performance on the offline imitation learning benchmarks. Project website: \url{https://www.lamda.nju.edu.cn/shaojj/KDD24_SRA/}.

Offline Imitation Learning with Model-based Reverse Augmentation

TL;DR

A novel model-based framework, called offline Imitation Learning with Self-paced Reverse Augmentation (SRA), which can efficiently generate trajectories leading to the expert-observed states in a self-paced style and guides maximizing long-term returns on these states, ultimately enabling generalization beyond the expert data.

Abstract

In offline Imitation Learning (IL), one of the main challenges is the \textit{covariate shift} between the expert observations and the actual distribution encountered by the agent, because it is difficult to determine what action an agent should take when outside the state distribution of the expert demonstrations. Recently, the model-free solutions introduce the supplementary data and identify the latent expert-similar samples to augment the reliable samples during learning. Model-based solutions build forward dynamic models with conservatism quantification and then generate additional trajectories in the neighborhood of expert demonstrations. However, without reward supervision, these methods are often over-conservative in the out-of-expert-support regions, because only in states close to expert-observed states can there be a preferred action enabling policy optimization. To encourage more exploration on expert-unobserved states, we propose a novel model-based framework, called offline Imitation Learning with Self-paced Reverse Augmentation (SRA). Specifically, we build a reverse dynamic model from the offline demonstrations, which can efficiently generate trajectories leading to the expert-observed states in a self-paced style. Then, we use the subsequent reinforcement learning method to learn from the augmented trajectories and transit from expert-unobserved states to expert-observed states. This framework not only explores the expert-unobserved states but also guides maximizing long-term returns on these states, ultimately enabling generalization beyond the expert data. Empirical results show that our proposal could effectively mitigate the covariate shift and achieve the state-of-the-art performance on the offline imitation learning benchmarks. Project website: \url{https://www.lamda.nju.edu.cn/shaojj/KDD24_SRA/}.
Paper Structure (25 sections, 6 theorems, 31 equations, 6 figures, 2 tables, 1 algorithm)

This paper contains 25 sections, 6 theorems, 31 equations, 6 figures, 2 tables, 1 algorithm.

Key Result

Proposition 3.1

Given expert dataset $D^E$ and offline dataset $D^O$, the policy could imitate the expert behavior on the expert-observed states $s\in D^E$ and take the transition information in offline data to maximize the policy-dependent transition probability towards expert-observed states on the expert-unobser where $Pr^{\pi'}_{s}(s_t=s')= Pr(s_t =s'|\pi',s_0=s)$.

Figures (6)

  • Figure 1: Maze2d-Medium Task: The performance of offline imitation learning (DWBC) degrades significantly as the number of expert trajectories in the unlabeled data decreases. More details are available in the Appendix.
  • Figure 2: An illustration of our main idea: we make use of transition information in the offline data to guide the agent from expert-unobserved states to expert-observed states, ultimately ensuring a long-term return (reaching the target diamond).
  • Figure 3: Comparision under varying number of expert trajectories.
  • Figure 4: Quantitative analysis on expert-unobserved states.
  • Figure 5: Comparision under varying number of expert trajectories.
  • ...and 1 more figures

Theorems & Definitions (9)

  • Proposition 3.1: Expert-State Distribution Maximization
  • Theorem 3.2
  • Lemma 3.3
  • Lemma A.1: Lower Bound on Expert-Observed States
  • Proof A.1
  • Theorem A.1
  • Proof A.2
  • Lemma A.2
  • Proof A.3