DIAMOND: Diffusion-Augmented Imitation Learning for Motion Planning in Autonomous Driving

Imitation-based planners for autonomous driving are prone to error accumulation due to the lack of recovery trajectories for out-of-distribution scenarios in expert datasets. This limitation, further compounded by the causal confusion problem, fundamentally constrains their performance in interactive, closed-loop driving environments. To tackle these challenges, we propose Diffusion-Augmented Imitation Learning for Motion Planning in Autonomous Driving (DIAMOND). DIAMOND leverages a diffusion-based data augmentation strategy to synthesize realistic and interaction-rich scenarios that lie off the expert data distribution, which occurred from the rollouts of an early-stage imitation policy in closed-loop environments. This data augmentation strategy enables DIAMOND to learn recovery behaviors from unseen states, thereby mitigating distribution shift. Furthermore, explicit causal auxiliary losses are incorporated to address causal confusion. Evaluated on the nuPlan benchmark, DIAMOND makes significant improvement in route completion rates (+13.4%). DIAMOND achieves the state-of-the-art closed-loop driving performance in interactive driving scenarios among imitation-based planners. The code is available at https://github.com/asddgfhzvxc/DIAMOND.