Optimizing Diffusion Models for Joint Trajectory Prediction and Controllable Generation

Yixiao Wang¹ Chen Tang¹² Lingfeng Sun¹ Simone Rossi³ Yichen Xie¹ Chensheng Peng¹ Thomas Hannagan³ Stefano Sabatini³ Nicola Poerio⁴ Masayoshi Tomizuka¹ Wei Zhan¹

¹University of California, Berkeley, USA ²The University of Texas at Austin, USA ³Stellantis, France ⁴Stellantis, Italy

Accepted to ECCV 2024

Paper Code

Abstract

Diffusion models are promising for joint trajectory prediction and controllable generation in autonomous driving, but they face challenges of inefficient inference time and high computational demands. To tackle these challenges, we introduce Optimal Gaussian Diffusion (OGD) and Estimated Clean Manifold (ECM) Guidance. OGD optimizes the prior distribution for a small diffusion time $T$ and starts the reverse diffusion process from it. ECM directly injects guidance gradients to the estimated clean manifold, eliminating extensive gradient backpropagation throughout the network. Our methodology streamlines the generative process, enabling practical applications with reduced computational overhead. Experimental validation on the large-scale Argoverse 2 dataset demonstrates our approach's superior performance, offering a viable solution for computationally efficient, high-quality joint trajectory generation and controllable generation for autonomous driving.

Overview of Optimal Gaussian Diffusion and Estimated Clean Manifold Guidance. (a) OGD uses the mean and variance of the data distribution to calculate the optimal prior distribution at a small T. It can largely reduce the diffusion time compared with vanilla diffusion. (b) ECM directly injects the gradient of guidance into the clean data manifold to save computation time at the inference stage.

Joint Trajectory Prediction

We propose optimizing the learnable Gaussian prior $\mathcal{N}(\mu,\Sigma)$ to enhance diffusion model performance with smaller steps, improving computational efficiency for diffusion-based autonomous driving applications.

Joint trajectory prediction results on Argoverse 2. Best result in bold, second best in underline. $T=40$ is the minimal diffusion time where OGD (Ours) outperforms VD (Vanilla Diffusion) on all metrics.

Controllable Generation

We reformulate controllable generation as multi-objective optimization, propose Estimated Clean Manifold (ECM) Guidance to avoid gradient propagation, and use reference points to improve guided sampling efficiency, realism and effectiveness.

Evaluation on controllable generation: route set U and Deceleration. Magenta diamonds represent goal points. In the first (second) row, goal points are set at the fork lane (right lane). Baselines (NNM & SF) struggle to drag samples from one modal to another. Our methods can achieve better guidance effectiveness and realism.

Visualizations on controllable generation: route set GT and Acceleration. Magenta diamonds represent goal points. Dodgerblue curves represent the predicted joint trajectory from 0s to 5s. Blue curves represent the predicted joint trajectory from 5s to 6s.

Visualizations on controllable generation: route set U and Acceleration. Magenta diamonds represent goal points. Dodgerblue curves represent the predicted joint trajectory from 0s to 5s. Blue curves represent the predicted joint trajectory from 5s to 6s.

BibTeX

@article{wang2024optimizing,
  title={Optimizing Diffusion Models for Joint Trajectory Prediction and Controllable Generation},
  author={Wang, Yixiao and Tang, Chen and Sun, Lingfeng and Rossi, Simone and Xie, Yichen and Peng, Chensheng and Hannagan, Thomas and Sabatini, Stefano and Poerio, Nicola and Tomizuka, Masayoshi and others},
  journal={arXiv preprint arXiv:2408.00766},
  year={2024}
}