JAM: Keypoint-Guided Joint Prediction after Classification-Aware Marginal Proposal for Multi-Agent Interaction

TL;DR

JAM is a two-stage multi-agent prediction framework that improves the accuracy of joint trajectory predictions in autonomous driving by classifying trajectory types and guiding predictions with key waypoints, achieving state-of-the-art results.

Contribution

The paper introduces a novel two-stage prediction framework with classification-aware marginal proposals and keypoint guidance for better multi-agent trajectory prediction.

Findings

Outperforms existing frameworks on Waymo dataset

Achieves state-of-the-art performance in interactive prediction

Effectively captures critical trajectory information

Abstract

Predicting the future motion of road participants is a critical task in autonomous driving. In this work, we address the challenge of low-quality generation of low-probability modes in multi-agent joint prediction. To tackle this issue, we propose a two-stage multi-agent interactive prediction framework named \textit{keypoint-guided joint prediction after classification-aware marginal proposal} (JAM). The first stage is modeled as a marginal prediction process, which classifies queries by trajectory type to encourage the model to learn all categories of trajectories, providing comprehensive mode information for the joint prediction module. The second stage is modeled as a joint prediction process, which takes the scene context and the marginal proposals from the first stage as inputs to learn the final joint distribution. We explicitly introduce key waypoints to guide the joint prediction module in better capturing and leveraging the critical information from the initial predicted trajectories. We conduct extensive experiments on the real-world Waymo Open Motion Dataset interactive prediction benchmark. The results show that our approach achieves competitive performance. In particular, in the framework comparison experiments, the proposed JAM outperforms other prediction frameworks and achieves state-of-the-art performance in interactive trajectory prediction. The code is available at https://github.com/LinFunster/JAM to facilitate future research.