Consistent Zero-Shot Imitation with Contrastive Goal Inference

Abstract

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.