💻 computer science

Provably Efficient Policy-Reward Co-Pretraining for Adversarial Imitation Learning

This paper introduces CoPT-AIL, a principled policy-reward co-pretraining algorithm that provides the first theoretical guarantee for accelerating adversarial imitation learning by jointly pretraining both the policy and reward through behavioral cloning to overcome the limitations of standard pretraining approaches.

Original authors: Tian Xu, Zexuan Chen, Zhilong Zhang, Yi-Chen Li, Chenyang Wang, Lei Yuan, Yang Yu

Published 2026-06-23

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Tian Xu, Zexuan Chen, Zhilong Zhang, Yi-Chen Li, Chenyang Wang, Lei Yuan, Yang Yu

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 walk like a professional dancer. You have a video of the dancer (the "expert"), but you don't have a manual explaining why they move the way they do. You just have the video.

This is the problem of Imitation Learning. There are two main ways robots try to learn from these videos:

The "Copycat" Method (Behavioral Cloning): The robot watches the video and tries to memorize every move.
- The Flaw: If the robot makes a tiny mistake early on, it ends up in a weird position it never saw in the video. It panics, makes a bigger mistake, and the whole routine falls apart. It's like a student trying to solve a math problem by memorizing the steps but failing the moment the numbers change slightly.
The "Game Show" Method (Adversarial Imitation Learning - AIL): The robot plays a game against a "Judge." The Judge tries to figure out what the robot is doing wrong compared to the dancer, and the robot tries to fool the Judge.
- The Flaw: This works much better than the Copycat method, but it's incredibly slow. The robot has to practice millions of times in the real world (or a simulation) to learn the rules of the game. It's like trying to learn to play chess by playing millions of games without a coach.

The Problem: The "Warm Start" Trap

Researchers tried to fix the slowness of the Game Show method by giving the robot a head start. They said, "Let's use the Copycat method first to get the robot almost good, then switch to the Game Show method to polish it."

But here's the catch: In practice, this often backfired. The moment they switched to the Game Show, the robot would get confused, forget what it learned, and perform worse than if it had started from scratch. It was like a student who studied hard for a test, got a B, and then when the teacher started a new, harder unit, the student forgot everything and got an F.

The Paper's Discovery: The Missing "Judge"

The authors of this paper dug into the math to find out why the "Copycat then Game Show" approach failed. They found the problem wasn't the robot (the policy); it was the Judge (the reward function).

The Robot: The Copycat method did a great job of teaching the robot the moves.
The Judge: When they switched to the Game Show, they gave the robot a brand new, random Judge who had no idea what "good" looked like. The robot was trying to impress a judge who was making up the rules as it went along.

The paper argues that the robot failed because the Judge was untrained, not because the robot was untrained.

The Solution: Co-Pretraining (The "Coach" and the "Judge" Together)

The authors propose a new method called CoPT-AIL. Instead of just training the robot, they train both the robot and the Judge at the same time, using the same video data.

Here is the creative analogy:
Imagine you are training a soccer player.

Old Way: You watch a pro player's highlight reel. You teach the rookie player to copy the moves (Robot training). Then, you hire a random person off the street to be the referee (Random Judge) and tell them to start the game. The referee doesn't know the rules, so the game is chaotic.
CoPT-AIL Way: You watch the highlight reel. You teach the rookie player to copy the moves. Crucially, you also teach the referee to watch that same highlight reel. You show the referee, "See how the pro moves? That's what 'good' looks like."

Now, when the game starts, the referee already knows what a good play looks like. The rookie player and the referee are on the same page from the very first second.

How They Did It (The Magic Trick)

The paper reveals a clever mathematical trick. They realized that the "score" a robot gets for doing a good job is mathematically related to how likely the expert was to do that move.

So, they didn't need a separate, complicated process to train the Judge. They simply took the robot's "Copycat" brain and said, "Okay, you are now also the Judge."

If the robot thinks a move is 90% likely to be what the expert did, the Judge gives it a high score.
If the robot thinks a move is 10% likely, the Judge gives it a low score.

This creates a perfect "warm start" for the Game Show. The robot and the Judge are already aligned before the expensive, real-world practice even begins.

The Results

The paper proves two things:

Mathematically: They showed that by training the Judge this way, the robot learns much faster and makes fewer mistakes than previous methods. It's the first time anyone has mathematically proven why this specific type of head-start works.
Practically: They tested this on 8 different robot tasks (like walking, running, and balancing). The new method (CoPT-AIL) learned to perform these tasks faster and more stably than all the other top methods. It reached expert-level performance with significantly fewer practice attempts.

Summary

This paper solves a puzzle that has confused researchers for years: Why does giving a robot a head start often make it worse?
Answer: Because you were only training the student, not the teacher.
Solution: Train the student and the teacher together using the same video. This aligns them perfectly, allowing the robot to learn complex skills much faster and more reliably.

Technical Summary: Provably Efficient Policy-Reward Co-Pretraining for Adversarial Imitation Learning

Problem Statement

Adversarial Imitation Learning (AIL) is a prominent method for learning complex behaviors by mimicking expert demonstrations, often outperforming Behavioral Cloning (BC) by mitigating compounding errors through a minimax optimization framework. However, AIL typically requires extensive online environment interactions, making it data-inefficient and costly in scenarios where interaction is expensive.

To address this, researchers have explored initializing AIL algorithms with policies pre-trained via BC. Empirical studies, however, have consistently shown that this strategy yields minimal benefits; the performance of the pre-trained policy often degrades rapidly during the early stages of AIL training, negating the advantages of the initial BC phase. While recent works have attempted to incorporate reward pretraining to overcome this, a rigorous theoretical understanding of why policy pretraining alone fails and how reward pretraining should be conducted remains elusive.

Methodology

Theoretical Diagnosis: The Role of Reward Error

The authors first provide a systematic theoretical analysis of AIL with policy pretraining alone (Algorithm 1). They decompose the imitation gap (the performance difference between the expert and the learner) into two primary components: policy error and reward error.

Policy Error: Arises from the discrepancy between the expert policy and the initial policy. The authors demonstrate that BC pretraining effectively reduces this error.
Reward Error: Arises from the discrepancy between the true reward function and the initial reward used in AIL. The analysis reveals that when the reward is randomly initialized (as is standard), the reward error remains a dominant bottleneck. This error persists regardless of how well the policy is pre-trained, explaining the empirical failure of policy-only pretraining.

Proposed Solution: Policy-Reward Co-Pretraining (CoPT-AIL)

Motivated by the identification of reward error as the critical bottleneck, the authors propose CoPT-AIL, a principled approach that jointly pretrains both the policy and the reward function.

Reward Shaping Insight: The authors leverage the theory of reward shaping (Ng et al., 1999), which posits that adding a potential-based term to a reward function preserves the optimal policy. They prove that for the purpose of reducing the relative policy evaluation error (the gap between the expert and learner values), it is sufficient to learn a reward function that approximates a shaping reward of the true reward, rather than the true reward itself.
The Connection: A key theoretical insight is that the log-probability of the expert policy, $\log(\pi^E_h(a|s))$ , is exactly a shaping reward of the true reward $r^*$ under specific potential functions.
The Algorithm:
- Stage 1 (Co-Pretraining): A standard Behavioral Cloning (BC) procedure is performed on expert demonstrations to obtain a policy $\pi_{BC}$ . This policy is then used to initialize the AIL policy. Crucially, the reward function is initialized as the log-likelihood of this BC policy: $r_1(s, a) = \log(\pi_{BC}(a|s))$ . This reuses the single BC procedure to generate both a high-quality policy and a high-quality reward initialization.
- Stage 2 (Online Fine-Tuning): The co-pretrained policy and reward are fine-tuned using standard online AIL interactions.

Key Contributions

Theoretical Diagnosis of Policy-Only Pretraining: The paper establishes a sharp theoretical baseline showing that while BC pretraining reduces policy error, it leaves reward error unaddressed. This reward error, induced by random reward initialization, is identified as the primary source of suboptimality in standard AIL with policy pretraining.
Principled Co-Pretraining Algorithm (CoPT-AIL): The authors derive a method to jointly pretrain policy and reward based on reward-shaping analysis. They prove that inferring a shaping reward (via the log-probability of the BC policy) is sufficient to reduce the relative policy evaluation error, circumventing the reward ambiguity problem.
First Theoretical Guarantee for Pretraining in AIL: The paper provides the first rigorous theoretical guarantee demonstrating that pretraining improves AIL efficiency. Theorem 1 proves that CoPT-AIL achieves an improved imitation gap bound compared to standard AIL without pretraining, specifically showing that the reward error term decays rapidly with the number of expert trajectories $N$ .
Empirical Validation: Extensive experiments on 8 continuous control tasks from the DMControl benchmark demonstrate that CoPT-AIL consistently outperforms existing state-of-the-art AIL methods (including IQLearn, PPIL, FILTER, HyPE, and OLLIE) and achieves near-expert performance with significantly fewer online interactions.

Results

Performance: In experiments, CoPT-AIL consistently matches or exceeds the convergence rates of prior SOTA methods. On tasks such as Cartpole Swingup, Hopper Hop, and Finger Spin, CoPT-AIL achieves near-expert performance with significantly fewer environment interactions than existing approaches.
Ablation Studies: Comparisons with "pure AIL" (no pretraining) and "policy-pretrained AIL" (only policy pretraining) confirm that policy pretraining alone offers limited improvement and can be unstable. In contrast, the joint pretraining of policy and reward in CoPT-AIL leads to faster and more stable convergence.
Sensitivity Analysis: The method remains robust across varying numbers of expert demonstrations (5 to 20 trajectories) and different regularization coefficients.
Gradient Consistency: Analysis of the gradient cosine similarity between the BC-based reward pretraining objective and the AIL fine-tuning objective shows high alignment ( $\ge 0.8$ ), indicating that the two stages operate synergistically rather than adversarially.

Significance and Claims

The paper claims to bridge the gap between theory and practice in Adversarial Imitation Learning by providing the first systematic theoretical treatment of reward pretraining. By identifying reward error as the dominant bottleneck in policy-pretrained AIL and proposing a method to jointly initialize the reward via the log-probability of the BC policy, the authors establish a new standard for data-efficient imitation learning.

The work asserts that CoPT-AIL offers a provably efficient path to AIL, reducing the reliance on costly online environment interactions. Theoretically, it provides the first formal guarantee that pretraining yields efficiency gains in AIL. Practically, it offers a computationally efficient solution that requires no separate reward learning step, as both components are derived from a single BC procedure. The authors note that while the current theoretical analysis focuses on the tabular setting, the empirical success suggests potential for extension to function approximation and more complex robot learning tasks.