Consistent Zero-Shot Imitation with Contrastive Goal Inference

TL;DR

Zero-shot imitation is pursued by pretraining interactive agents with self-supervised exploration so they can mimic unseen demonstrations at test time without any rewards or demonstrations. The core approach CIRL couples maximum-entropy inverse reinforcement learning with contrastive RL by inferring a latent goal $g$ under which trajectories maximize $\sum_t r_g(s_t,a_t)$, while the agent learns a goal-conditioned policy; goal inference uses a variational posterior $q_\xi(g|\tau)$ and a forward KL objective, with the posterior $p^\star(g|\tau) \propto p(g) p^\star(\tau|g)$ and the partition function $Z_g$ playing a key role. The method includes automatic goal sampling via GoalKDE and a mean-field variant of goal inference, and a theoretical consistency proof showing CIRL recovers the true goal distribution. Empirically, CIRL outperforms zero-shot baselines on URLB and standard GCRL benchmarks, demonstrating effective self-supervised pretraining for rapid adaptation to new tasks.

Abstract

In the same way that generative models today conduct most of their training in a self-supervised fashion, how can agentic models conduct their training in a self-supervised fashion, interactively exploring, learning, and preparing to quickly adapt to new tasks? A prerequisite for embodied agents deployed in real world interactions ought to be training with interaction, yet today's most successful AI models (e.g., VLMs, LLMs) are trained without an explicit notion of action. The problem of pure exploration (which assumes no data as input) is well studied in the reinforcement learning literature and provides agents with a wide array of experiences, yet it fails to prepare them for rapid adaptation to new tasks. Today's language and vision models are trained on data provided by humans, which provides a strong inductive bias for the sorts of tasks that the model will have to solve (e.g., modeling chords in a song, phrases in a sonnet, sentences in a medical record). However, when they are prompted to solve a new task, there is a faulty tacit assumption that humans spend most of their time in the most rewarding states. The key contribution of our paper is a method for pre-training interactive agents in a self-supervised fashion, so that they can instantly mimic human demonstrations. Our method treats goals (i.e., observations) as the atomic construct. During training, our method automatically proposes goals and practices reaching them, building off prior work in reinforcement learning exploration. During evaluation, our method solves an (amortized) inverse reinforcement learning problem to explain demonstrations as optimal goal-reaching behavior. Experiments on standard benchmarks (not designed for goal-reaching) show that our approach outperforms prior methods for zero-shot imitation.

Figures (8)

• Figure 1: Zero shot imitation learning. Assuming access to a multi-task environment, our generalist agent must imagine and practice its own tasks to effectively imitate unknown task demonstrations at test time.
• Figure 2: Zero-shot imitation learning with CIRL via goal inference. CIRL combines goal-conditioned contrastive RL pre-training, automatic goal sampling for exploration, and a mean field goal inference model to imitate expert demonstrations. Here we see how an Ant's imitation policy and posterior distribution over goal states evolve across timesteps toward a final maximum a posteriori (MAP) estimate.
• Figure 3: Value of self-supervised RL pre-training CIRL consistently outperforms the alternative FB representation zero-shot imitation method as well as the naive 1-NN policy baseline.
• Figure 4: Summarizing behavior via goals yields better imitation than reward-based explanations. When using the last expert demonstration state as the goal, CIRL achieves high imitation scores on goal-conditioned environments while FB struggles to infer goal-conditioned reward functions from online learning.
• Figure 5: Mean field goal inference models outperform alternative full $\tau$ input models For all environments, the Mean Field model should contain the true posterior and is computationally easier to train, making it the superior choice for CIRL.

Theorems & Definitions (6)

• Definition 1
• Corollary 1
• Lemma 1
• Lemma 2
• proof
• proof