Robust Offline Imitation Learning from Diverse Auxiliary Data

TL;DR

ROIDA addresses offline imitation learning with unlabeled, diverse auxiliary data by learning a PU-based reward model to identify high-quality transitions and applying weighted behavioral cloning on those samples, while also using TD learning with the learned rewards to exploit lower-quality samples. This two-pronged approach enables robust utilization of heterogeneous demonstrations without prior data-quality assumptions, achieving consistent improvements across seven D4RL environments for both locomotion and manipulation. The method demonstrates strong robustness to the relative mix of expert and non-expert data and outperforms existing offline IL baselines that rely on specific data compositions. The work offers practical value for leveraging large, unlabeled offline datasets in real-world settings and provides open-source code for replication.

Abstract

Offline imitation learning enables learning a policy solely from a set of expert demonstrations, without any environment interaction. To alleviate the issue of distribution shift arising due to the small amount of expert data, recent works incorporate large numbers of auxiliary demonstrations alongside the expert data. However, the performance of these approaches rely on assumptions about the quality and composition of the auxiliary data, and they are rarely successful when those assumptions do not hold. To address this limitation, we propose Robust Offline Imitation from Diverse Auxiliary Data (ROIDA). ROIDA first identifies high-quality transitions from the entire auxiliary dataset using a learned reward function. These high-reward samples are combined with the expert demonstrations for weighted behavioral cloning. For lower-quality samples, ROIDA applies temporal difference learning to steer the policy towards high-reward states, improving long-term returns. This two-pronged approach enables our framework to effectively leverage both high and low-quality data without any assumptions. Extensive experiments validate that ROIDA achieves robust and consistent performance across multiple auxiliary datasets with diverse ratios of expert and non-expert demonstrations. ROIDA effectively leverages unlabeled auxiliary data, outperforming prior methods reliant on specific data assumptions. Our code is available at https://github.com/uditaghosh/roida.

Figures (7)

• Figure 1: Robustness to composition of auxiliary data. Performance of existing offline IL algorithms, such as DWBC xu2022discriminator, varies significantly depending on the amount of high-quality transitions present in the auxiliary data (given expert set is kept fixed). In contrast, ROIDA is more robust, highlighting its ability to extract information even from low-quality transitions. The setup shown here is on the Hopper environment; refer to Sec. \ref{['sec:experiments']} for details.
• Figure 2: Framework overview. ROIDA first learns a reward function using PU learning. It then identifies high-reward expert-like transitions and combines them with the expert data for weighted BC (Sec. \ref{['sec:reward']},\ref{['sec:BC']}). To extract value from lower quality samples, ROIDA applies TD learning, steering the policy towards high reward states (Sec. \ref{['sec:TD']}). By combining weighted BC and TD learning, ROIDA effectively leverages uncurated offline data.
• Figure 3: Scalability to the size of the auxiliary dataset. We visualize the performance of ROIDA on the Hopper environment as the number of random transitions is varied. Here, we show 3 scenarios corresponding to different proportions of the D4RL random set. This highlights ROIDA’s ability to learn policies even when the expert and noisy data ratio is quite imbalanced.
• Figure 4: rliable evaluation on the Locomotion benchmark.
• Figure 5: rliable evaluation on the Adroit benchmark.