Provably Efficient Policy-Reward Co-Pretraining for Adversarial Imitation Learning
Abstract
Adversarial imitation learning (AIL) achieves high-quality imitation compared to behavioral cloning (BC), but demands substantial online environment interaction. Recent empirical work has explored initializing AIL algorithms with BC-pretrained policies to address this limitation, yet a rigorous theoretical understanding of pretraining's role in AIL remains elusive. This paper provides a systematic theoretical analysis and introduces principled pretraining algorithms for accelerating AIL. We begin by analyzing AIL with policy pretraining alone, identifying reward error as the dominant source of suboptimality. This reveals a critical and previously overlooked gap: the absence of reward pretraining. Motivated by this finding, we develop a principled policy–reward co-pretraining approach grounded in a reward-shaping analysis. Our analysis uncovers a fundamental connection between expert policies and shaping rewards, which naturally gives rise to CoPT-AIL, an approach that jointly pretrains both policy and reward through a single BC procedure. We prove that CoPT-AIL achieves an improved imitation gap bound over standard AIL, establishing the first theoretical guarantee for the benefits of pretraining in AIL. Experimental results confirm CoPT-AIL's superior performance over existing AIL methods.
Lay Summary
Teaching an AI agent to mimic expert behavior, like a robot learning from a skilled gymnast, typically requires enormous amounts of costly trial-and-error practice. A natural shortcut is to first load the agent with the expert's recorded actions, then let it practice. Researchers have long observed that this barely helps: the agent's skill quickly collapses once practice begins. But no one had a rigorous explanation for why. We identified the missing piece: loading the expert's actions is not enough. The agent also needs a pre-loaded internal sense of which actions are good: a scoring function. We proved mathematically that both can be derived from the same expert recordings in a single step. Our algorithm, CoPT-AIL, simply reuses how likely the expert found each action as its initial scoring signal; no extra data is needed. This is the first result to mathematically guarantee that pre-loading both actions and scoring leads to faster learning. Experiments on a suite of robot control tasks show CoPT-AIL reaches expert-level performance with significantly fewer practice runs than existing methods.