Equivariant Map and Agent Geometry for Autonomous Driving Motion Prediction

Autonomous vehicle motion prediction has been a focal area of research, with diverse methodologies and frameworks being proposed with a number of them utilizing deep neural networks[4, 5, 6]. LSTM-based models like Social-LSTM[7] have been utilized for capturing temporal dependencies. Convolutional Neural Networks (CNN) have been employed to process rasterized map data[3]. More recently, attention mechanisms like Transformer models have been adapted for motion prediction[8]. Graph-based approaches like Graph Neural Networks have been explored for multi-agent interaction[9, 10, 11]. VectorNet[2] has introduced the use of vectorized representations, highlighting the importance of preserving geometric relationships. These diverse approaches lay a rich foundation for ongoing innovations in the field.

Map information encoding is a critical aspect of motion prediction, with two main approaches: rasterized and vectorized representations. Rasterized encoding, which converts map information into pixel grids, has been employed in many CNN-based model[3, 12]. It offers simplicity and compatibility with standard image processing techniques but suffers from loss of geometric precision and can be computationally expensive. In contrast, vectorized representations, like those used in VectorNet[2], preserve the geometric structure of the map, enabling accurate encoding of spatial relationships. Vectorized encoding ensures scalability and fidelity but may require more complex processing to fully exploit its potential. Recent advancements in transformer-based models have demonstrated the efficacy of vectorized encoding in autonomous driving applications[13]. The choice between these approaches hinges on specific requirements, balancing simplicity with geometric accuracy.

The concept of equivariance has become particularly prominent in the realm of 2D image analysis. Owing to the susceptibility of CNN structures to rotational alterations, there has been a pursuit of designs that embrace rotation-equivariance, such as the implementation of directional convolutional filters[14]. In Graph Neural Networks, [15] achieves both rotation and translation equivariance by using tensor field neural networks. In these works, equivariant models have shown their advantage of an efficient training process by avoiding data augmentation. A myriad of research initiatives have proposed equivariant layers tailor-made for specific tasks, including protein structure decoding[16]. However, many remain limited to state prediction and overlook sequence data. Few attempts, like one that focuses on coordinate processing, strive for motion equivariance. Our model innovates an equivariant map processor on top of the equivariant sequence predictor, EqMotion[1], to predict the trajectories for autonomous vehicles.

Map features provide the context for our ego motion prediction. To account for equivariant nature between the map and ego agent, we first translate the map coordinate frame to align with the current location at time $t$ of the ego agent $a$, which is denoted as $x_{a}^{t}$. For all the center lines and center line points in the map $M_{l}^{q}$, $l\in[0,L]$ and $q\in[0,Q]$, we have:

$$M_{\text{centered},l}^{q}=M_{l}^{q}-x_{a}^{t}.$$

(1)

In addition to translating the map, we also rotate the map to align with the ego agent heading which is computed from the current velocity:

$$v_{a}^{t}=x_{a}^{t}-x_{a}^{t-1},$$

(2)

$$\theta_{a}^{t}=\operatorname{arctan2}(v_{a}^{t}[1],(v_{a}^{t}[0]),$$

(3)

$$R_{\theta}=\begin{bmatrix}\cos(\theta_{a}^{t})&-\sin(\theta_{a}^{t})\\ \sin(\theta_{a}^{t})&\cos(\theta_{a}^{t})\end{bmatrix},$$

(4)

$$M_{\text{rotated},l}^{q}=R_{\theta}\cdot M_{\text{centered},l}^{q}.$$

(5)

We want to show our map transformation is equivariant to agent motion. This means for arbitrary agent movement, the map’s coordinates in the agent frame will shift the same amount. We denote each coordinate pair in the map $M$ as $m_{i}$ and $i\in[0,L\times Q]$. The agent’s current position is represented by $x_{a}^{t}$ and its heading represented by $\theta$.

Given the map feature is now equivariant to the agent motion, we now further expand its dimension space for downstream backbone. Specifically, we reshape the rotated map $M_{\text{rotated}}$ so that the lanes and waypoints per lane are vectorized into a vector $M_{\text{vectorized}}$ where $M_{\text{vectorized},n}$ is a single coordinate pair and $n\in[0,L\times Q]$, which we then encode $M_{\text{vectored}}$ with a Transformer [17] layer that is composed of the following:

1. 1.

   Multi-Head Self-Attention

   $$Attention(Q,K,V)=softmax\left(\frac{QK^{T}}{\sqrt{d_{k}}}\right)V,$$

   (9)

   where

   	$\displaystyle Q$	$\displaystyle=M_{\text{vectorized}}\cdot W_{Q}$		(10)
   	$\displaystyle K$	$\displaystyle=M_{\text{vectorized}}\cdot W_{K}$		(11)
   	$\displaystyle V$	$\displaystyle=M_{\text{vectorized}}\cdot W_{V},$		(12)

   and $W_{Q}$, $W_{K}$, $W_{V}$ are the matrices for linear transformation. $d_{k}$ is the dimensionality of the keys in the multi-head attention mechanism.

2. 2.

   Add and Norm

   $$Z=LayerNorm(X+Attention(Q,K,V))$$

   (13)

3. 3.

   Feed-Forward Neural Network (FFN)

   $$FFN(Z)=ReLU(ZW_{1}+b_{1})W_{2}+b_{2}$$

   (14)

   $$Y=LayerNorm(Z+FFN(Z))$$

   (15)

The resulting vector $Y$ is then concatenated with the agent features as fed to the downstream.

$$M_{\text{features}}=[X;Y]$$

(16)

We here employ the Eqmotion Network [1]. Specifically, given our $AgentMapFeatures$ from the above, we obtain the initial geometric $G^{0}\in\mathbb{R}^{N\times T_{out}}$ and pattern $H^{0}\in\mathbb{R}^{N\times hidden\_dim}$ features:

$$G^{0},H^{0}=\mathfrak{F}_{\text{FeatureInitLayer}}(M_{\text{features}}).$$

(17)

We then apply a graph convolution network to infer the relationships between the agents, and the map context:

$${e_{ij}}=\mathfrak{F}_{\text{GraphConvolutionLayer}}(G^{0},H^{0}).$$

(18)

With the above, we will sequentially apply two networks to learn the geometric and pattern features in an interleaving fashion. This step will repeat $P$ times and $P$ is a hyperparameter:

$$G^{p}=\mathfrak{F}_{\text{GeometricLayer}}(G^{p-1},H^{p-1},{e_{ij}}),$$

(19)

$$H^{p}=\mathfrak{F}_{\text{PatternLayer}}(G^{p-1},H^{p-1}).$$

(20)

For the detailed implementation of the above layers please refer to EqMotion[1].

We take the final geometric feature $G^{P}\in\mathbb{R}^{N\times hidden\_dim}$ and decode it with a 4-layer MLP which produces the output trajectory $\hat{y}\in\mathbb{R}^{T_{out}}$. We will also add the current position of the agent to it:

$$\hat{y}=MLP(G^{P}-\bar{G^{P}})+\bar{G^{P}}+\bar{X_{0}}.$$

(21)

We use the Average Displacement Error as the loss function:

$$\text{ADE}=\frac{1}{T}\sum_{t=1}^{T}\|\hat{y}_{t}-y_{t}\|_{2}.$$

(22)

In our work, we used the Argoverse motion forecasting dataset[18], a comprehensive and well-regarded collection specifically designed for autonomous vehicle applications. Argoverse 1 provides a training set of around 200k scenes, each includes high-definition maps and both ego and neighbor agent trajectories. By leveraging this dataset, we were able to validate and benchmark our model’s performance in realistic scenarios, ensuring that the findings are representative and applicable to real-world autonomous driving contexts. In this dataset, agent history is composed of coordinates from the past 2 seconds, and agent futures are the next 3 seconds, both sampled at 10Hz.

In our comprehensive evaluation, we utilized both Average Displacement Error (ADE) and Final Displacement Error (FDE) metrics at multiple time intervals (1s, 2s, and 3s) to thoroughly assess the predictive capabilities of our model. ADE measures the average difference between predicted and actual trajectories over specified time horizons, providing insights into prediction accuracy. On the other hand, FDE quantifies the endpoint difference between predicted and actual trajectories, offering a distinct perspective on model performance. By employing both ADE and FDE at different time intervals, we gained a comprehensive understanding of our model’s accuracy in both intermediate and long-term predictions. This multi-faceted evaluation approach enabled us to identify not only immediate accuracy but also the ability of our model to sustain accurate predictions over longer horizons. These metrics reinforce the rigor of our evaluation process and highlight our commitment to providing reliable and insightful autonomous vehicle motion predictions.

During training, we use the Adam optimizer with a learning rate of $1e^{-5}$. We train our network for 20 epochs and used a batch size of 512. Our training hardware is the NVidia RTX 3060Ti with 8GB of graphic memory. The training loss decay can be seen in Figure 2.

In our model, we used the configuration parameters as in Table I.

To measure the performance of our model, we developed a Long Short Term Memory (LSTM) based network to serve as the baseline. LSTM is a popular network that is able to learn a temporal sequence such as in [7] and [19], in which LSTM is used to predict vehicle behavior at intersections.

In our LSTM network, we rasterize the map into a $200\times 200$ boolean image with the lane centerlines represented as having a value of $1$ in the image while the others have $0$. The image then goes through two convolution neural networks (CNN), each has $3\times 3$ kernels followed by a MaxPool layer with stride $2$. The convoluted feature then goes through a 4-layer MLP to produce a feature vector, which we then concatenate to the feature vector produced by the LSTM. Lastly, we use a 4-layer MLP to produce the output trajectory.

To provide a better insight into the effectiveness of using the map feature processing and encoding logic via Transformer blocks, we performed experiments with these modifications.

1. 1.

   Skipping the map rotation logic in eq. 5 and only perform translation.

2. 2.

   Replacing the transformer block with a self attention layer[20] with a graph size of 64.

3. 3.

   Replacing the transformer block with the same CNN model as in the baseline.

4. 4.

   Ablating the map processing logic to measure the impact of the map information.

The quantitative results are shown in Table II. Rows 1 - 2 are the results for not using a map. Rows 3 - 7 are the results for using a map.

Our LSTM baseline has higher errors in ADE and FDE, signifying less accurate predictions. It has a moderate model size and reasonable training time. By replacing the LSTM backbone with EqMotion[1], we see an immediate improvement in both ADE and FDE compared to the baseline and reduced complexity in terms of the number of parameters. Lengthier training time may indicate a more intricate learning process. We observe that regardless of the map feature encoder used, we can see a performance improvement if we align the map with agent heading. Note that rotating the map does not incur an increase in the number of parameters. Thus it does not affect model size or training time. Using a transformer model than a simpler self attention one costs 1.2M more parameters, but also better performance. Lastly, applying a rasterized view of the map and supply it to the model performs worse than the our equivariant map processor, even with many more parameters and longer training time.

In conclusion, our models with map information generally outperform the non map models in prediction accuracy. Different architectural choices like map rotation and attention/transformer showed a noticeable improvement. The model with rotated map and transformer architecture appears to be a overall better choice given models at row 3-6 have similar sizes. Careful consideration is needed for the CNN variant, as it has higher complexity and significantly longer training time. Further experiments with hyperparameter tuning, different feature engineering, or additional data might lead to more robust insights.

Some visualizations of our model predictions are shown in Figure 4. We observe the predicted trajectories appear to stay close to the ground truth, suggesting our model is able to predict reasonably well. Additionally, the ground truth data suffers from sampling noise, represented as the spikes in the plot. Our prediction, however, produces a much smoother and realistic trajectory.