Multi-class Agent Trajectory Prediction with Selective State Spaces for autonomous driving

Abstract

Understanding and predicting multi-class agents’ movement has become more critical and challenging in diverse applications such as autonomous driving and urban intelligent monitoring. The current research mainly focuses on the motion trajectory of single-class agents. However, due to real traffic scenarios’ complexity and interactive behaviors’ variability, the motion patterns displayed by various classes of agents show inherent randomness. In this paper, inspired by the linear-time sequence model Mamba, we propose a Multi-class Agent Trajectory Prediction with Selective State Spaces (MTPSS) to model the interaction between different agents and better predict the trajectory of an individual. Specifically, MTPSS includes modeling relationships in both temporal and spatial dimensions. When encoding the spatial correlation within the trajectory graph, we construct a category-based sorting approach, which puts large-size category nodes behind to enhance contextual access. Then the sorted nodes are bi-directionally scanned through Mamba blocks, which makes the model more robust to permutations. In terms of temporal, considering the highly dynamic nature of rapidly moving agents, we utilize Mamba’s remarkable performance on sequential data to effectively conduct temporal scans to capture long-range temporal dependencies. Finally, to compute physically feasible trajectories, MTPSS employs the Neural Ordinary Differential Equation to smooth the predicted trajectory of the agent. We conducted extensive experiments on two publicly available traffic datasets and compared our method with state-of-the-art methods. Quantitative experiments show that our performance metrics are superior to state-of-the-art methods, and qualitative experiments demonstrate that the predicted trajectories have good diversity, which shows its potential in real-world traffic scenarios.