🤖 machine learning

Consistent Zero-Shot Imitation with Contrastive Goal Inference

This paper introduces Contrastive Inverse Reinforcement Learning (CIRL), a self-supervised pre-training framework that converts multi-task inverse reinforcement learning into a tractable goal-inference problem to enable agents to consistently reproduce expert behavior in zero-shot imitation scenarios without rewards or test-time updates.

Original authors: Kathryn Wantlin, Chongyi Zheng, Benjamin Eysenbach

Published 2026-06-25

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Kathryn Wantlin, Chongyi Zheng, Benjamin Eysenbach

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a robot to do a new task, like stacking blocks or walking across a room. Usually, you have to either show the robot exactly what to do step-by-step (imitation) or give it a strict set of rules and points for doing things right (rewards). But what if you want the robot to learn on its own, without you showing it anything or giving it a scorecard, and then later, just by watching one video of a human doing the task, the robot figures out how to do it perfectly?

That is the challenge this paper tackles. The authors, from Princeton University, introduce a new method called CIRL (Contrastive Inverse Reinforcement Learning). Here is how it works, broken down into simple concepts and analogies.

The Big Idea: "Goals" are the Secret Language

The core insight of the paper is that almost any complex task can be summarized by a single goal state.

The Analogy: Think of a cooking show. You don't need to memorize every single chop, stir, and flip the chef does. You just need to know the final dish (the goal). If you know the goal is "a perfect omelet," you can figure out the steps to get there.
The Problem: Most AI tries to guess the "recipe" (the reward function) directly, which is like trying to guess the chef's exact thoughts. It's messy and often wrong.
The CIRL Solution: Instead of guessing the recipe, CIRL guesses the final dish (the goal). Once the robot knows the goal, it uses a pre-trained skill to reach that goal.

How CIRL Learns (The "Playground" Phase)

Before the robot ever sees a human, it goes through a "pre-training" phase where it learns entirely on its own, with no human help, no rewards, and no demonstrations.

The Curious Explorer (GoalKDE):
Imagine a child in a giant playground. To learn how to move, the child doesn't just wander randomly; they specifically look for places they haven't visited yet. CIRL has a mechanism called GoalKDE that acts like this child. It looks at all the places the robot has been and says, "You've been here a million times, but you've never been there." It picks a new, unexplored spot as a "goal" and tells the robot to try to get there. This forces the robot to explore the whole environment and learn how to move in all directions.
The "What-If" Trainer (Contrastive RL):
While the robot is exploring, it learns a special kind of map. Instead of learning "this action gives me 10 points," it learns a sense of distance. It learns: "If I am here and I want to get to that goal, how hard is it?" It learns to distinguish between "easy to reach" states and "hard to reach" states. This is crucial because it helps the robot understand that some goals are naturally harder to achieve than others.
The Memory Bank:
The robot saves all these "trips" to different goals in a memory bank. It learns a general policy: "Here is how I move to reach any goal I might be given."

The Test: The "One-Shot" Imitation

Now, the real test begins. You show the robot a single video of an expert doing a task (e.g., a human walking to a specific chair). The robot has never seen this task before and cannot ask for help.

The Detective (Goal Inference):
The robot watches the video and asks, "What was the expert trying to reach?" It uses a statistical model (a "mean-field" model, which is like a smart shortcut) to guess the goal.
- The Smart Trick: The paper proves that simply looking at where the expert ended up isn't enough. The robot must also consider how hard it was to get there. If the expert took a very difficult path to reach a spot, it's a strong signal that this spot was the intended goal. If they took an easy path, maybe they were just wandering. CIRL accounts for this difficulty, making it a much better detective than previous methods.
The Execution:
Once the robot guesses the goal (e.g., "The expert wanted to sit in that chair"), it activates its pre-trained "playground skills" to reach that specific spot. It doesn't need to relearn how to walk; it just applies its existing skills to the new target.

Why This is Better Than Old Methods

The paper compares CIRL to other methods (like "Forward-Backward" or FB representations).

The Flaw in Old Methods: Imagine a map that says, "The place people visit most is the most important." If a robot wanders into a room and gets stuck because the door is jammed, it visits that room a lot. Old methods might think, "Oh, they must have wanted to stay in that jammed room!"
The CIRL Advantage: CIRL knows that the jammed room was hard to get to and hard to stay in. It realizes, "They didn't want to be there; they were just stuck." By accounting for the difficulty of reaching a state, CIRL correctly infers the true goal, even if the expert had a hard time getting there.

The Results

The authors tested this on various robot tasks, from moving a robotic arm to pushing objects and walking.

The Outcome: CIRL consistently beat other "zero-shot" methods (methods that try to learn from a single example without retraining).
The Proof: They mathematically proved that their method is "consistent," meaning it will reliably find the correct goal, whereas other similar methods can be tricked by the difficulty of the task.

Summary

In short, CIRL is a way to teach a robot to be a self-sufficient explorer. It spends time playing in a sandbox, learning how to get to any corner of the room on its own. Then, when it sees a human do something once, it doesn't try to copy the human's exact movements. Instead, it figures out where the human was trying to go, and uses its sandbox skills to get there. It works because it understands that just because a place is visited often, it doesn't mean it was the goal—sometimes, the goal is the place that was hardest to reach.

Technical Summary: Consistent Zero-Shot Imitation with Contrastive Goal Inference

Problem Statement

The paper addresses Zero-Shot Imitation Learning (IL), a setting where an agent must reproduce expert behavior from a single demonstration without additional environment interaction, gradient updates, or access to reward functions at test time. Current approaches often rely on mimicking demonstrations directly (Behavioral Cloning) or inferring reward functions via Inverse Reinforcement Learning (IRL). However, Behavioral Cloning suffers from distributional shift, and standard IRL methods often struggle with the under-specified nature of reward inference, particularly when failing to account for the relative difficulty of reaching different states (the partition function).

The authors propose a framework where agents are pre-trained purely through self-supervised interaction (no demonstrations, no rewards, no human data) to develop a foundation of knowledge. The core challenge is to enable these agents to infer a demonstrator's intent from a single trajectory and execute a policy to achieve that intent without further training.

Methodology: Contrastive Inverse Reinforcement Learning (CIRL)

The proposed framework, Contrastive Inverse Reinforcement Learning (CIRL), reframes the multi-task IRL problem as a goal-inference problem. It operates on the observation that many tasks can be summarized by a single goal state. The method consists of three primary components integrated during pre-training:

Maximum Entropy Contrastive Reinforcement Learning (MaxEnt CRL):
- The agent learns a goal-conditioned policy using a self-supervised objective.
- It extends Contrastive Reinforcement Learning (CRL) to the Maximum Entropy (MaxEnt) setting. Instead of just learning discounted future rewards, the agent optimizes for the sum of discounted future rewards and entropy.
- The critic function $f_{\phi,\psi}(s, a, g)$ estimates the discounted state occupancy measure using Noise Contrastive Estimation (NCE), distinguishing between the current state-action-goal tuple and negative samples.
- The actor maximizes an entropy-regularized objective, learning a soft Q-function $\tilde{Q}_g(s, a)$ .
Automatic Goal Proposal (GoalKDE):
- To drive exploration without external rewards, the agent employs GoalKDE, an automatic goal sampling mechanism.
- It fits a Gaussian Kernel Density Estimator (KDE) to the states visited in the replay buffer.
- The agent samples goals from states with the lowest probability under the KDE (under-visited states), encouraging the agent to explore the state space autonomously.
Variational Goal Inference:
- At test time, given a demonstration trajectory $\tau$ , the agent must infer the expert's goal $\hat{g}$ .
- The authors prove that under MaxEnt IRL with goal-conditioned rewards, inferring the reward is equivalent to inferring the latent goal.
- They employ a mean-field variational model $q_\xi(g|\tau) = \prod_{t=0}^T q_\xi(g|s_t, a_t)$ to approximate the posterior distribution of the goal. This amortized inference model is trained to maximize the likelihood of the goal given the trajectory, effectively learning to "read" the intent from the demonstration.

Key Theoretical Contributions

The paper provides a theoretical proof of consistency for the CIRL method:

Consistent Goal Inference: The authors prove that CIRL consistently recovers the demonstrator's intent by accounting for the relative difficulty of reaching different states (the partition function $Z_g$ ).
Inconsistency of Prior Work: They demonstrate that structurally similar methods, such as Forward-Backward (FB) representations, are inconsistent. FB methods infer rewards based on state visitation counts without considering the partition function. The authors provide a counterexample (a 2-state MDP) showing that FB can incorrectly identify a demonstrator's policy because it fails to distinguish between a state that is frequently visited due to ease of access versus a state that is the true, difficult-to-reach goal.

Experimental Results

The authors evaluate CIRL on the JaxGCRL and Unsupervised Reinforcement Learning Benchmark (URLB) environments, including Reacher, Pusher, Ant, and PointMaze tasks.

Performance vs. Baselines: CIRL consistently outperforms prior zero-shot imitation methods, including:
- 1-Nearest Neighbor (1-NN): A naive baseline that copies the action of the nearest state in the expert trajectory.
- Forward-Backward (FB) Representations: A state-of-the-art method for zero-shot RL and imitation.
- HILP: A latent variable approach for unsupervised pre-training.
Ablation Studies:
- Pre-training: CIRL with self-supervised pre-training significantly outperforms methods without it.
- Goal Inference: The mean-field approximation for goal inference universally outperforms models that use the full trajectory as input, offering better computational efficiency and performance.
- Exploration: While GoalKDE achieves near-perfect imitation scores in some environments (e.g., Reacher), the gap between GoalKDE and an "oracle" goal sampler widens in higher-dimensional spaces (e.g., Ant, PointMaze), suggesting room for improvement in exploration strategies for complex state spaces.
Generalization: CIRL successfully imitates policies trained on standard reward-maximizing tasks (non-goal-conditioned) by inferring goals related to the agent's 3D position and velocity, demonstrating that the goal-space hypothesis can capture complex reward structures.

Significance and Claims

The paper claims that goals serve as a compact and expressive summary of behavior. By reducing the hypothesis space of reward functions to goal-conditioned rewards, the authors show that:

Efficiency: Imitation becomes more efficient when projecting behavior onto the restricted space of goal-conditioned reward functions.
Theoretical Soundness: The method is theoretically consistent, correctly inferring intent where other methods fail due to ignoring the partition function.
Self-Supervised Foundation: It is possible to build a foundation for interactive agents through purely self-supervised exploration (proposing goals, attempting to reach them, and learning from the outcome) without any human demonstrations or reward signals during pre-training.

The authors conclude that this "imagine-and-practice" loop enables agents to learn transferable policies across diverse task distributions, supporting a re-evaluation of goals as a fundamental mechanism for zero-shot imitation.