Offline Imitation Learning with Variational Counterfactual Reasoning

In offline imitation learning (IL), an agent aims to learn an optimal expert behavior policy without additional online environment interactions. However, in many real-world scenarios, such as robotics manipulation, the offline dataset is collected from suboptimal behaviors without rewards. Due to the scarce expert data, the agents usually suffer from simply memorizing poor trajectories and are vulnerable to variations in the environments, lacking the capability of generalizing to new environments. To automatically generate high-quality expert data and improve the generalization ability of the agent, we propose a framework named underlineOffline underlineImitation underlineLearning with underlineCounterfactual data underlineAugmentation (OILCA) by doing counterfactual inference. In particular, we leverage identifiable variational autoencoder to generate textitcounterfactual samples for expert data augmentation. We theoretically analyze the influence of the generated expert data and the improvement of generalization. Moreover, we conduct extensive experiments to demonstrate that our approach significantly outperforms various baselines on both textscDeepMind Control Suite benchmark for in-distribution performance and textscCausalWorld benchmark for out-of-distribution generalization. Our code is available at urlhttps://github.com/ZexuSun/OILCA-NeurIPS23.