Neural Density Imitation Learning Based on Contrastive Representation

Generative Adversarial Imitation Learning (GAIL) achieves superior performance on certain tasks by leveraging expert data to train policy networks for intelligent agents. However, on pixel-based tasks, GAIL requires extensive environment interactions because the feature extractor struggles to efficiently capture high-dimensional features from image pixels. To address this issue, we propose Neural Density Imitation Learning based on Contrastive Learning (NDICL). NDICL utilizes Convolutional Neural Networks (CNNs) to extract high-dimensional features and enhances network training through contrastive learning. It also improves sample efficiency by employing image augmentation, while Inverse Reinforcement Learning (IRL) supports the training of the feature extractor. Experimental results demonstrate that NDICL achieves state-of-the-art performance and sample efficiency in the DMControl suite across 5 out of 6 environments, excelling in pixel-based tasks with limited samples, where it reaches approximately 98.3% of expert efficiency.