Mitigating Covariate Shift in Imitation Learning for Autonomous Vehicles Using Latent Space Generative World Models

2024-09-25

arXiv.org

AI Summary

To address poor policy generalization in autonomous driving imitation learning caused by distribution shift, this paper proposes a generative world model-based framework for mitigating covariate shift. Methodologically: (1) it constructs a latent-space world model that explicitly captures state dynamics and enforces policy alignment with human demonstration trajectories during training; (2) it designs a Transformer-based perception encoder integrating multi-view cross-attention and learnable scene queries to enhance robustness in complex, visually diverse scenarios. Evaluated on both CARLA and NVIDIA DRIVE Sim platforms, the approach achieves significant improvements over state-of-the-art methods in closed-loop evaluation. Moreover, it demonstrates superior resilience under strong perturbations—including sensor noise, abrupt weather changes, and adversarial occlusions—and exhibits strong cross-simulator generalization capability.

Technology Category

• Machine Learning: Imitation Learning & Inverse Reinforcement Learning
• Computer Vision: Diffusion Models for Vision
• Humans and AI: Human-Aware Planning and Behavior Prediction

Application Category

• Responsible Web: Machine-in-the-loop, human agency and autonomy
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: User modeling and simulation for interactive and conversational systems

Abstract

We propose the use of latent space generative world models to address the covariate shift problem in autonomous driving. A world model is a neural network capable of predicting an agent's next state given past states and actions. By leveraging a world model during training, the driving policy effectively mitigates covariate shift without requiring an excessive amount of training data. During end-to-end training, our policy learns how to recover from errors by aligning with states observed in human demonstrations, so that at runtime it can recover from perturbations outside the training distribution. Additionally, we introduce a novel transformer-based perception encoder that employs multi-view cross-attention and a learned scene query. We present qualitative and quantitative results, demonstrating significant improvements upon prior state of the art in closed-loop testing in the CARLA simulator, as well as showing the ability to handle perturbations in both CARLA and NVIDIA's DRIVE Sim.

Problem

• Mitigate covariate shift in autonomous driving imitation learning
• Use latent generative models to reduce training data needs
• Enable error recovery from out-of-distribution perturbations

Innovation

• Latent space generative world models mitigate covariate shift
• Transformer-based perception encoder with multi-view cross-attention
• End-to-end training aligns with human demonstration states