Uncertainty-Aware Module for Trajectory Prediction and Ego-Planning of Challenging Agents in Autonomous Driving

In autonomous driving, trajectory prediction requires modeling the interactions between agents and maps, hence enhancing the ego-planning. Most methods adopt an encoder-decoder framework, where the encoder extracts features and the multi-layer decoder predicts trajectories. However, these approaches typically treat all agents equally, overlooking the varying levels of prediction difficulty among different agents. To address this limitation, we propose a Uncertainty-Aware Module (UAM) to capture and quantify the prediction uncertainty for each agent. By incorporating uncertainty-aware decoding, the module fixes the trajectories of agents with lower prediction uncertainty early, allowing the model to focus on more challenging agents. This approach not only improves both accuracy and computational efficiency but can also be seamlessly integrated as a plug-and-play module into state-of-the-art networks. By adding UAM to Gameformer, we outperform Gameformer in both accuracy and efficiency on nuPlan Dataset. Additionally, we propose a Mixture of Experts planning framework to integrate rule-based method and Gameformer with UAM. Experiments on nuPlan show state-of-the-art performance.