Revisiting Distribution Correction Estimation for Offline Imitation Learning with Suboptimal Dataset
Abstract
Imitation Learning (IL) learns high-quality policies from expert demonstrations but degrades in low-expert-data regimes. To address this, recent work studies ``offline IL with supplementary data", augmenting expert data with trajectories from suboptimal policies. A prominent framework is Distribution Correction Estimation (DICE), which estimates density ratios via the dual of a divergence minimization problem between learned and expert visitation distributions. However, existing DICE-based methods either rely on a strict coverage assumption or introduce additional dataset regularization, limiting performance. We propose ReDICE, a new DICE-based method that addresses these issues through an objective-level reformulation. Our approach constructs a mixture-distribution objective that preserves the original expert-imitation objective while removing the coverage assumption, and its dual reduces to a stable Gumbel regression objective for efficient optimization. We further introduce a novel policy extraction mechanism that improves performance. Experiments on standard and real-world offline IL benchmarks show that ReDICE consistently outperforms prior methods and achieves state-of-the-art results.
Lay Summary
Imitation learning teaches an AI agent to act by copying expert demonstrations, but collecting many expert examples is often expensive or difficult. A practical alternative is to also use cheaper data from imperfect behavior, but existing methods can either imitate bad behavior or rely on assumptions that may not hold in real-world datasets. We propose ReDICE, a new method that learns from both expert and imperfect data while keeping the main goal of matching the expert. The key idea is to combine expert and imperfect data into a mixture that makes learning more stable without changing the original imitation objective. ReDICE also uses a stable optimization approach and a policy extraction step that helps the agent focus on useful actions. Across standard robotics benchmarks and a real-world maritime navigation task, ReDICE learns stronger policies than prior methods, especially when expert data is scarce.