Consistent Zero-Shot Imitation with Contrastive Goal Inference

Zero-shot imitation learning requires an agent to reproduce expert behavior from a single demonstration without additional environment interaction or gradient updates at test time. We introduce Contrastive Inverse Reinforcement Learning (CIRL), a self-supervised framework for pre-training zero-shot imitation agents. Our methods rests on a key observation that many useful tasks can be summarized by a single goal state. We can thus convert the multi-task inverse RL problem into a more tractable goal-inference problem, and utilize state-of-the-art goal-conditioned RL methods to recover a policy that reaches the goal. During pre-training, CIRL jointly employs three components to learn without any rewards or demonstrations: (1) a variant of contrastive RL designed to learn maximum-entropy goal-conditioned policies, (2) an automatic goal proposal mechanism (GoalKDE) that drives exploration, and (3) a mean-field variational model that performs amortized goal inference from trajectories. We prove that this procedure consistently recovers the demonstrator's intent by accounting for the relative difficulty of reaching different states and show how structurally similar prior work may otherwise fail to infer the correct reward. Experiments on goal-conditioned and standard reward-maximizing control tasks show that CIRL outperforms prior zero-shot imitation methods, supporting the expressiveness of goals as a compact summary of behavior.

Imagine you have a household robot that you want to quickly teach to do new chores with just one demonstration each. This is the challenge of "zero-shot imitation learning": can an AI agent watch a single demonstration and instantly and robustly reproduce that behavior? We built CIRL (Contrastive Inverse Reinforcement Learning), a framework that lets an AI agent learn entirely through self-directed play, no rewards or labelled data during training. The key insight is that most tasks can be summarized by a goal state (where the expert is trying to end up), so instead of inferring a complex reward function from a demonstration, the agent simply asks: "what destination/result was this expert trying to reach?" We also prove that a popular alternative approach systematically gives the wrong answer to this question, because it ignores relative difficulty of possible goals. At test time, the agent watches one demonstration, infers the goal, and immediately navigates toward it. Our experiments show this outperforms prior methods across a range of robotic control tasks. Beyond robotics, this "imagine and practice" approach to learning points toward AI agents that can acquire new skills more efficiently. These agents are then less dependent on collection of expensive human demonstrations during training, and can quickly pick up new skills from limited demonstrations at test time.