Imitation from Observations with Trajectory-Level Generative Embeddings

TL;DR

TGE is proposed, 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, ensuring a robust learning signal even when offline data is distributionally distinct from the expert.

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.

Figures (6)

• Figure 1: Overview of Trajectory-level Generative Embeddings (TGE). The framework employs a trajectory-level diffusion encoder to map trajectory segments into a latent embedding space. A surrogate reward is then computed using a Laplacian kernel over these embeddings to augment the reward-free suboptimal dataset, enabling offline RL training for policy learning.
• Figure 2: T-SNE Embedding Visualization. The latent embeddings produced by the generative planner naturally separate expert transitions from suboptimal ones. As shown, expert samples (colored in red) and suboptimal samples (colored in blue) form distinct clusters with a clear boundary, indicating that the trajectory-level embedding space is highly discriminative.
• Figure 3: Ablation of the Reward Signal Density. (Top) Walker2d-Medium. (Bottom) Hopper-Random. Distributions of synthetic rewards assigned to expert (red) and suboptimal (blue) data. Compared to Gaussian-kernel-based reward estimation, the Laplacian kernel preserves a heavier-tailed reward signal and is associated with improved performance and training stability.
• Figure 4: Ablation of the Temporal Horizon. The results demonstrate how performance varies with the context horizon $H$. We observe that larger $H$ generally leads to improved performance, highlighting the importance of temporal context.
• Figure 5: Training Dynamics of TGE + IQL. The curves display the mean normalized score and standard deviation (shaded region). TGE combined with IQL shows stable policy improvement, confirming that our geometric reward signal enables robust learning across different offline RL methods.