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

arXiv: https://arxiv.org/abs/2409.16663 · HTML: https://mlsanity.com/paper/2409.16663/

## Key facts

- Cited by 22 papers in the map (data through 14 Oct 2026).
- Placed in World models for autonomous driving, one of 11 directions in World models.
- Builds on 32 of its 39 references that are in the map.
- First posted on arXiv 25 Sep 2024.

## Placed in

- World models / [World models for autonomous driving](https://mlsanity.com/wm/world-models-for-autonomous-driving/)
- End-to-end autonomous driving / [Generative simulation for driving](https://mlsanity.com/driving/generative-simulation-for-driving/)
- Offline RL and imitation learning from data / [Imitation and offline RL for driving](https://mlsanity.com/offlinerl/imitation-and-offline-rl-for-driving/)

## Builds on (most cited first, 30 of 32)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [DINOv2: Learning Robust Visual Features without Supervision](https://mlsanity.com/paper/2304.07193/) (2023)
- [End-to-End Object Detection with Transformers](https://mlsanity.com/paper/2005.12872/) (2020)
- [Mamba: Linear-Time Sequence Modeling with Selective State Spaces](https://mlsanity.com/paper/2312.00752/) (2023)
- [Analyzing and Improving the Image Quality of StyleGAN](https://mlsanity.com/paper/1912.04958/) (2019)
- [Elucidating the Design Space of Diffusion-Based Generative Models](https://mlsanity.com/paper/2206.00364/) (2022)
- [Stable Video Diffusion: Scaling Latent Video Diffusion Models to Large Datasets](https://mlsanity.com/paper/2311.15127/) (2023)
- [CARLA: An Open Urban Driving Simulator](https://mlsanity.com/paper/1711.03938/) (2017)
- [Dream to Control: Learning Behaviors by Latent Imagination](https://mlsanity.com/paper/1912.01603/) (2019)
- [World Models](https://mlsanity.com/paper/1803.10122/) (2018)
- [Align your Latents: High-Resolution Video Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2304.08818/) (2023)
- [Variational Inference: A Review for Statisticians](https://mlsanity.com/paper/1601.00670/) (2016)
- [Learning Latent Dynamics for Planning from Pixels](https://mlsanity.com/paper/1811.04551/) (2018)
- [Planning-oriented Autonomous Driving](https://mlsanity.com/paper/2212.10156/) (2022)
- [Lift, Splat, Shoot: Encoding Images From Arbitrary Camera Rigs by Implicitly Unprojecting to 3D](https://mlsanity.com/paper/2008.05711/) (2020)
- [NuPlan: A closed-loop ML-based planning benchmark for autonomous vehicles](https://mlsanity.com/paper/2106.11810/) (2021)
- [DayDreamer: World Models for Physical Robot Learning](https://mlsanity.com/paper/2206.14176/) (2022)
- [ChauffeurNet: Learning to Drive by Imitating the Best and Synthesizing the Worst](https://mlsanity.com/paper/1812.03079/) (2018)
- [Learning by Cheating](https://mlsanity.com/paper/1912.12294/) (2019)
- [Multi-Modal Fusion Transformer for End-to-End Autonomous Driving](https://mlsanity.com/paper/2104.09224/) (2021)
- [Exploring the Limitations of Behavior Cloning for Autonomous Driving](https://mlsanity.com/paper/1904.08980/) (2019)
- [Trajectory-guided Control Prediction for End-to-end Autonomous Driving: A Simple yet Strong Baseline](https://mlsanity.com/paper/2206.08129/) (2022)
- [Safety-Enhanced Autonomous Driving Using Interpretable Sensor Fusion Transformer](https://mlsanity.com/paper/2207.14024/) (2022)
- [End-to-End Urban Driving by Imitating a Reinforcement Learning Coach](https://mlsanity.com/paper/2108.08265/) (2021)
- [Model-Based Imitation Learning for Urban Driving](https://mlsanity.com/paper/2210.07729/) (2022)
- [Learning from All Vehicles](https://mlsanity.com/paper/2203.11934/) (2022)
- [Hidden Biases of End-to-End Driving Models](https://mlsanity.com/paper/2306.07957/) (2023)
- [ReasonNet: End-to-End Driving with Temporal and Global Reasoning](https://mlsanity.com/paper/2305.10507/) (2023)
- [Feedback in Imitation Learning: The Three Regimes of Covariate Shift](https://mlsanity.com/paper/2102.02872/) (2021)
- [Grasping with Chopsticks: Combating Covariate Shift in Model-free Imitation Learning for Fine Manipulation](https://mlsanity.com/paper/2011.06719/) (2020)

## Built on by (most cited first, 22 of 22)

- [A Comprehensive Survey on World Models for Embodied AI](https://mlsanity.com/paper/2510.16732/) (2025)
- [GaussianWorld: Gaussian World Model for Streaming 3D Occupancy Prediction](https://mlsanity.com/paper/2412.10373/) (2024)
- [A Survey of World Models for Autonomous Driving](https://mlsanity.com/paper/2501.11260/) (2025)
- [A Survey: Learning Embodied Intelligence from Physical Simulators and World Models](https://mlsanity.com/paper/2507.00917/) (2025)
- [SimScale: Learning to Drive via Real-World Simulation at Scale](https://mlsanity.com/paper/2511.23369/) (2025)
- [Vision-Language-Action Models for Autonomous Driving: Past, Present, and Future](https://mlsanity.com/paper/2512.16760/) (2025)
- [The Role of World Models in Shaping Autonomous Driving: A Comprehensive Survey](https://mlsanity.com/paper/2502.10498/) (2025)
- [GeoDrive: 3D Geometry-Informed Driving World Model with Precise Action Control](https://mlsanity.com/paper/2505.22421/) (2025)
- [A Smooth Sea Never Made a Skilled SAILOR: Robust Imitation via Learning to Search](https://mlsanity.com/paper/2506.05294/) (2025)
- [Data Scaling Laws for End-to-End Autonomous Driving](https://mlsanity.com/paper/2504.04338/) (2025)
- [Latent Chain-of-Thought World Modeling for End-to-End Driving](https://mlsanity.com/paper/2512.10226/) (2025)
- [Exploring the Interplay Between Video Generation and World Models in Autonomous Driving: A Survey](https://mlsanity.com/paper/2411.02914/) (2024)
- [Four Principles for Physically Interpretable World Models](https://mlsanity.com/paper/2503.02143/) (2025)
- [What Matters for Scalable and Robust Learning in End-to-End Driving Planners?](https://mlsanity.com/paper/2603.15185/) (2026)
- [WARPD: World model Assisted Reactive Policy Diffusion](https://mlsanity.com/paper/2410.14040/) (2024)
- [Scaling Self-Play for End-to-End Driving](https://mlsanity.com/paper/2606.19641/) (2026)
- [End-to-End Low-Level Neural Control of an Industrial-Grade 6D Magnetic Levitation System](https://mlsanity.com/paper/2509.01388/) (2025)
- [Self-Imitated Diffusion Policy for Efficient and Robust Visual Navigation](https://mlsanity.com/paper/2601.22965/) (2026)
- [OWMDrive: Causality-Aware End-to-End Autonomous Driving via 4D Occupancy World Model](https://mlsanity.com/paper/2606.30421/) (2026)
- [End-to-End Imitation Learning for Optimal Asteroid Proximity Operations](https://mlsanity.com/paper/2502.01034/) (2025)
- [CARIL: Confidence-Aware Regression in Imitation Learning for Autonomous Driving](https://mlsanity.com/paper/2503.00783/) (2025)
- [A Survey on End-to-End Autonomous Driving Training from the Perspectives of Data, Strategy, and Platform](https://mlsanity.com/paper/2610.00926/) (2026)

## Abstract (arXiv)

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.