Imitation from Observations with Trajectory-Level Generative Embeddings

TL;DR

This paper introduces TGE, a trajectory-level generative embedding that creates a smooth surrogate reward for offline imitation learning, effectively handling scarce and suboptimal data by capturing long-term dynamics.

Contribution

The paper proposes TGE, a novel trajectory-level generative embedding using diffusion models to improve offline imitation learning from limited and imperfect data.

Findings

TGE outperforms prior methods on D4RL benchmarks.

It effectively captures long-horizon dynamics.

It robustly handles distributional gaps in offline data.

Abstract

We consider the offline imitation learning from observations (LfO) where the expert demonstrations are scarce and the available offline suboptimal data are far from the expert behavior. Many existing distribution-matching approaches struggle in this regime because they impose strict support constraints and rely on brittle one-step models, making it hard to extract useful signal from imperfect data. To tackle this challenge, we propose TGE, a trajectory-level generative embedding for offline LfO that constructs a dense, smooth surrogate reward by estimating expert state density in the latent space of a temporal diffusion model trained on offline trajectory data. By leveraging the smooth geometry of the learned diffusion embedding, TGE captures long-horizon temporal dynamics and effectively bridges the gap between disjoint supports, ensuring a robust learning signal even when offline data is distributionally distinct from the expert. Empirically, the proposed approach consistently matches or outperforms prior offline LfO methods across a range of D4RL locomotion and manipulation benchmarks.