SlowGuess/NeurIPS_2023_OCR huggingface dataset:
and finally the agents join. Intuitively, the map influences the interpretation of traffic lights but not vice versa, whereas map and traffic lights together influence the behavior of agents but not vice versa. After the inter-class Transformer decoders, the all-to-all Transformer encoder in Figure 3c builds a full attention matrix allowing tokens to attend to each other irrespective of their class. Finally, each agent token is concatenated with  $N_{\mathrm{AC}}$  learnable anchors, which are shared among each type of agent. Since the task is marginal motion prediction, we batch over agents and anchors, i.e. now the number of anchor tokens is  $N_{\mathrm{AG}} \cdot N_{\mathrm{AC}}$ . Then the anchor-to-all Transformer decoder in Figure 3d lets each of these anchor tokens attend to all tokens so as to aggregate more contextual information. The final output is denoted as  $\hat{\mathbf{z}}_i^{\mathrm{AG}}, i \in \{1, \dots, N_{\mathrm{AG}} \cdot N_{\mathrm{AC}}\}$ , and it is used to generate the multi-modal future trajectories. We set the number of neighbors to  $K$  for the intra-class and all-to-all Transformers. This number  $K$  is multiplied by  $\gamma_{\mathrm{TL}}$ ,  $\gamma_{\mathrm{AG}}$  and  $\gamma_{\mathrm{AC}}$  respectively for the enhance-TL, enhance-AG and AC-to-all Transformers, such that these Transformers can have a larger receptive field.
HPTR organizes the Transformers in a hierarchical way such that some of the intermediate results can be cached and reused during the online inference; for example the outputs of the intra-MP Transformer, i.e. the static map features. This allows tokens from different classes to be updated asynchronously, which significantly reduces the online inference latency. We can further improve the efficiency without sacrificing the performance by trimming the redundant attentions. The attention matrices shown in Figure 3a, 3b and 3c are overlapping, which means some relationships are repeatedly modeled, such as the agent-to-agent self-attention. We propose to remove the intra-TL, intra-AG and all-to-all Transformer, while keeping the intra-MP, enhance-TL and enhance-AG Transformer. These three Transformers together build a lower triangular attention matrix which is necessary and sufficient to model the relationship between map, traffic lights and agents. Our approach can be seen as an extension to Wayformer [36] which does not introduce the inter-class Transformer decoders. We can trim Figure 3 differently and cast HPTR into the Wayformer. Specifically, Wayformer with late fusion corresponds to HPTR with diagonal attention matrix (only intra-class TF), Wayformer with early fusion corresponds to HPTR with full attention matrix (only all-to-all TF) and Wayformer with hierarchical fusion corresponds to HPTR with the diagonal followed by the full attention matrix.
3.4 Output representation and training strategies
We follow the common practice [36, 47, 52] to represent the outputs as a mixture of Gaussians and train with hard assignment. The multi-modal future trajectories for each agent are generated by decoding  $\hat{\mathbf{z}}_i^{\mathrm{AG}}, i \in \{1, \dots, N_{\mathrm{AG}} \cdot N_{\mathrm{AC}}\}$  via two MLPs; the confidence head and the trajectory head. The confidence head predicts a scalar confidence of each trajectory. Besides the Gaussian parameters of 2D positions ( $\mu_x, \mu_y, \sigma_x, \sigma_y, \rho$ ) at each future time step, our trajectory head predicts also the yaw angles, speeds and 2D velocities. We use the cross entropy loss for confidences, negative log-likelihood loss for 2D positions, negative cosine loss for yaw angles and Huber loss for speeds and 2D velocities. The final training loss is the unweighted sum of all losses. Following the hard-assignment strategy, for each agent we optimize only the predicted trajectory that is closest to the ground truth in terms of 2D average displacement error. Please refer to the appendix for more details.
**Benchmarks.** We benchmark our method on the two most popular datasets: the Waymo Open Motion Dataset (WOMD) [16] and the Argoverse-2 motion forecasting dataset (AV2) [60]. Both datasets have an online leaderboard for marginal motion prediction. However, their task descriptions are slightly different. For each scenario, WOMD evaluates the predictions of up to 8 agents, whereas AV2 evaluates only one agent. Both datasets require exactly 6 futures to be predicted for each target agent. The sampling time is 0.1 seconds for both datasets. The history length is 11 steps for WOMD and 50 steps for AV2, whereas the future length is 80 steps for WOMD and 60 steps for AV2. We use the official evaluation tool of the leaderboards to compute the metrics. For the WOMD leaderboard [58], the ranking metric is soft mAP; for the AV2 leaderboard [1], it is brier-minFDE. Please refer to the leaderboard homepages for more details about the dataset and evaluation metrics.
Implementation details. We use Transformer with pre-layer normalization [62] for HPTR. For each episode, we consider  $N_{\mathrm{MP}} = 1024$  map polylines,  $N_{\mathrm{TL}} = 40$  traffic light stop points and  $N_{\mathrm{AG}} = 64$  agents. Each polyline contains up to  $N_{\mathrm{node}} = 20$  one-meter-long segments. The base number of neighbors considered by KNARPE is  $K = 36$ . This number is multiplied by  $\gamma_{\mathrm{TL}} = 2$ ,  $\gamma_{\mathrm{AG}} = 4$  and  $\gamma_{\mathrm{AC}} = 10$  respectively for the enhance-TL, enhance-AG and AC-to-all Transformer. We set  $N_{\mathrm{AC}} = 6$  to predict exactly 6 futures as specified by the leaderboard. We do not apply ensembling or expensive post-processing such as trajectory aggregation. Our post-processing manipulates only the confidences. To improve the soft mAP, we use the non-maximum suppression of MPA [31] for the WOMD leaderboard. To improve the brier-minFDE, we set the softmax temperature to 0.5 for the AV2 leaderboard. More details are provided in the appendix.
Training details. Thanks to the viewpoint invariance of the pairwise-relative representation, no data augmentation or input permutation is needed for the training of HPTR. We use AdamW optimizer with an initial learning rate of 1e-4 and decaying by 0.5 every 25 epochs. We train with a total batch size of 12 episodes on 4 RTX 2080Ti GPUs. For WOMD, we randomly sample  $25\%$  from all training episodes at each epoch; for AV2 we use  $50\%$ . Our final models are trained for 120 epochs for WOMD and 150 epochs for AV2. The complete training takes 10 days. For WOMD, the SDV agents, agents of interest and agents to be predicted are used for optimization. For AV2, the SDV agents, scored agents and focal agents are used for optimization. To address the imbalance between agent types, for WOMD we additionally optimize for pedestrians and cyclists which are tracked for at least 4 seconds.
In Table 1 we benchmark our approach on the public leaderboards of WOMD and AV2. From the leaderboard we observe that using ensemble and predicting redundant trajectories can increase the prediction diversity and hence improve the soft mAP, brier-minFDE and miss rate. For example, MTR-Adv-ens aggregates the outputs of 7 ablation models, each predicting 64 futures. Besides the exaggerated large number of redundant predictions, it is also non-trivial to aggregate them into 6 predictions. Some works use K-means clustering while the others use non-maximum suppression; both involve heuristic parameter fine-tuning. Since our framework is dedicated to real-world autonomous driving, applying these computationally expensive techniques is contradictory to our motivation. For a fair comparison, we focus on end-to-end methods which do not apply these techniques. On the WOMD dataset, we achieve SOTA performance among the end-to-end methods, including MTR-e2e, the end-to-end version of MTR, and MPA, the end-to-end version of MultiPath++. Specifically, HPTR outperforms the pairwise-relative HDGT and the scene-centric SceneTransformer by a large margin in all metrics. We also show competitive performance on the AV2 leaderboard. We achieve better performance in all metrics except the brier-minFDE compared to GoRela, which uses GNNs to tackle the pairwise-relative representation. Since there exists a performance gap while adapting from one dataset to another and we choose WOMD to be our main benchmark, our performance on the AV2 leaderboard could be further improved by fine-tuning the hyper-parameters for the AV2 dataset.
In Table 2 we ablate different input representations and hierarchical architectures. The scene-centric baseline, HPTR SC, uses the architecture of our HPTR and the input representation of
Table 1: Results on the marginal motion prediction leaderboards of WOMD and AV2. Both tables are sorted according to the ranking metric such that the best performing method is on the top. The ranking metric is soft mAP for WOMD, and brier-minFDE $_6$  for AV2.  $\dagger$  denotes ensemble. * denotes predicting more futures than required. SC: scene-centric. AC: agent-centric. PR: pairwise-relative.
WOMD test	repr.	soft mAP ↑	mAP ↑	minADE ↓	minFDE ↓	miss rate ↓
*†MTR-Adv-ens [47]	AC	0.4594	0.4492	0.5640	1.1344	0.1160
*†Wayformer [36]	AC	0.4335	0.4190	0.5454	1.1280	0.1228
*MTR [47]	AC	0.4216	0.4129	0.6050	1.2207	0.1351
*†MultiPath++ [52]	AC	N/A	0.4092	0.5557	1.1577	0.1340
HPTR (Ours)	PR	0.3968	0.3904	0.5565	1.1393	0.1434
MPA [31] (MultiPath++)	AC	0.3930	0.3866	0.5913	1.2507	0.1603
HDGT [28]	PR	0.3709	0.3577	0.7676	1.1077	0.1325
Gnet [18]	AC	0.3396	0.3259	0.6207	1.2391	0.1718
SceneTransformer [37]	SC	N/A	0.2788	0.6117	1.2116	0.1564
WOMD valid	repr.	soft mAP ↑	mAP ↑	minADE ↓	minFDE ↓	miss rate ↓
HPTR (Ours)	PR	0.4222	0.4150	0.5378	1.0923	0.1326
MTR-e2e	AC	N/A	0.3245	0.5160	1.0404	0.1234
AV2 test	repr.	brier-minFDE6↓	minFDE6↓	minFDE1↓	minADE6↓	miss rate6↓
*ProphNet [55]	AC	1.88	1.33	4.74	0.68	0.18
†Gnet [18]	AC	1.90	1.34	4.40	0.69	0.18
†TENET [56]	AC	1.90	1.38	4.69	0.70	0.19
*MTR [47]	AC	1.98	1.44	4.39	0.73	0.15
GoRela [11]	PR	2.01	1.48	4.62	0.76	0.22
HPTR (Ours)	PR	2.03	1.43	4.61	0.73	0.19
THOMAS [19]	AC	2.16	1.51	4.71	0.88	0.20
HDGT [28]	PR	2.24	1.60	5.37	0.84	0.21
SceneTransformer [37]. The agent-centric baseline, WF baseline, is our reimplementation of the Wayformer [36] with multi-axis attention and early fusion. Both baselines achieve their expected performances compared to their original implementations. The large performance gap between HPTR SC and other models confirms the disadvantage of scene-centric representation. The agent-centric baseline requires more training iterations. After convergence, its performance is on par with our HPTR. To ablate the hierarchical architecture, we implement three variations of HPTR; each corresponds to a fusion strategy investigated by Wayformer. The full attention corresponds to the early fusion, diagonal corresponds to late fusion, and HPTR with diagonal followed by full attention corresponds to Wayformer with hierarchical fusion. While the early fusion performs the best for Wayformer, for HPTR the lower triangular attention we proposed outperforms both the early and the late fusion by a significant margin. The performance of hierarchical fusion is slightly worse than ours, but its inference latency is significantly longer because of the redundancy in its attention matrices.
4.4 Efficiency analysis and qualitative results
In Figure 4 we compare the computational efficiency of the ablation models presented in Table 2. As discussed in prior works [36, 47], one of the major drawbacks of agent-centric approaches is the poor scalability, which is reflected by the large slope of the memory and latency curves of our WF baseline. On the RTX 2080Ti, it can handle only up to 48 agents while the inference latency is  $140\mathrm{ms}$ . On the contrary, the scene-centric baseline is extremely efficient, but it suffers from poor accuracy. Our HPTR takes the best of both approaches. In terms of efficiency, our HPTR is comparable to the scene-centric approaches, while our performance is on par with the agent-centric approaches. Compared to the scene-centric baseline, the GPU memory consumption and the inference latency of HPTR are slightly higher because for each new agent, we have to compute its relative pose to all existing tokens. Compared to the hierarchical architectures introduced by Wayformer, HPTR with our lower triangular attention achieves the best trade-off between accuracy and latency. By reusing the static map features during the online inference, our HPTR predicts the multi-modal futures for 64 agents in  $37\mathrm{ms}$  without the use of inference libraries. Narrowing the receptive field, for example by
Table 2: Ablation on the valid split of WOMD. The table is sorted according to the soft mAP such that the best-performing method is on the top. Performances are reported as the mean plus-minus 3 standard deviations over 3 training seeds. Models are trained for 60 epochs if not otherwise mentioned. WF: Wayformer. SC: scene-centric, AC: agent-centric, PR: pairwise-relative.
WOMD valid	input repr.	intra MP	intra TL/AG	enhance TL/AG	all 2all	minFDE ↓	soft mAP ↑
HPTR (Ours)	PR	✓	×	✓	×	1.145 ± 0.016	0.399 ± 0.010
WF baseline (100-epoch)	AC	×	×	×	✓	1.161 ± 0.006	0.397 ± 0.007
HPTR diag+full (WF hier.)	PR	✓	✓	×	✓	1.156 ± 0.014	0.391 ± 0.002
HPTR diag (WF late)	PR	✓	✓	×	×	1.169 ± 0.013	0.387 ± 0.012
HPTR full (WF early)	PR	×	×	×	✓	1.158 ± 0.093	0.386 ± 0.041
WF baseline	AC	×	×	×	✓	1.212 ± 0.019	0.378 ± 0.014
HPTR SC	SC	✓	×	✓	×	1.687 ± 0.046	0.246 ± 0.005
Figure 4: HPTR is as efficient as scene-centric methods in terms of GPU memory consumption and inference latency, while being as accurate as agent-centric methods. We use standard Ubuntu, Python and Pytorch without optimizing for real-time deployment. We predict one scenario at each inference time on one 2080Ti, i.e. we batch over scenarios and the batch size is 1. The number of context agents is the same as the number of predicted agents for all experiments. During offline inference, every inference starts from scratch, while during online inference, we cache and reuse the static map features. Specifically, we repeat the inference of the same scenario for 100 times to simulate online inference, where the static map features are computed at the first step and reused for the next 99 steps.
reducing  $\lambda_{\mathrm{AC}}$  from 10 to 8, can improve the inference speed but the accuracy might be affected in some cases. Using half precision at inference time can reduce the latency to  $25\mathrm{ms}$  (40 fps) without affecting the accuracy. Compared to the most efficient agent-centric method Wayformer, we reduce the memory consumption and the online inference latency by  $80\%$  without sacrificing the accuracy.
In Figure 5 we compare the computational efficiency of our method with the GNN-based pairwise relative methods. Currently there are two such methods, GoRela [11] and HDGT [28]. While HDGT does not match the prediction accuracy of GoRela or our method, it is worth noting that GoRela is not open-sourced, whereas HDGT is. Therefore, we use HDGT in this efficiency comparison. The left plot of Figure 5 confirms the good scalability of HDGT in terms of GPU memory. This is expected because it uses the pairwise-relative representation. In the middle plot, we can observe that HDGT is slower than both our HPTR and our agent-centric Wayformer baseline in terms of offline inference speed. To confirm that HDGT runs correctly on our setup, in the right plot we reproduce the inference time of HDGT on the complete WOMD validation split with different validation batch size and we compare the reproduce numbers with the reported number in the HDGT paper. The slow inference speed of GNN-based methods such as HDGT is mainly because GNN libraries cannot utilize the GPU as efficiently as the basic matrix operations do. Our KNARPE is implemented with the most basic matrix operations, hence it is better suited for real-time and on-board applications.
Figure 6 illustrates the qualitative results of our HPTR. For each type of agent, we select a successful case where the most confident prediction matches the ground truth, and a failure case to show the limitation of our method. In the successful cases, we observe the vehicle stops at the red light, the pedestrian walks along the road edge with another person, and the cyclist rides across the road via the
Figure 5: Efficiency comparison with HDGT. We run their official repository on our machine with a single 2080Ti. The left plot shows that HDGT achieves good scalability in terms of GPU memory consumption as expected. The middle plot shows that the offline inference latency (with batch size 1) of HDGT scales well, but it is significantly larger than that of other Transformer-based methods. The right plot shows the inference times for the complete WOMD validation split (64 agents per episode) with different validation batch sizes. It confirms the inference speed reported in the original HDGT paper is correctly reproduced on our setup. Our setup is faster because it has a more powerful CPU.
Figure 6: Qualitative results of HPTR. For each type of agent we show a successful case (top) and a failure case (bottom). The intersections are cluttered because we visualize traffic lights by overlaying the lanes they control with their color. The ground truth is in orange. The target agent and the predictions are in cyan. The most confident prediction has the least transparent color, the thickest line and the biggest cross. Please refer to the appendix for a detailed explanation of the visualization.
crosswalk. Although in the failure cases the most confident prediction deviates from the ground truth, we are encouraged to see that our predictions are still reasonable. The failure cases can be addressed by improving the prediction diversity via goal-conditioning, which is orthogonal to our contributions.
In this paper we introduce a novel attention module, KNARPE, that allows the pairwise-relative representation to be used by Transformers. Based on KNARPE, we present a pure Transformer-based framework called HPTR, which uses hierarchical architecture to enable asynchronous token update and avoid redundant computations. While agent-centric methods suffer from poor scalability and scene-centric methods suffer from poor accuracy, our HPTR gets the best from both worlds. Experiments on the two most popular benchmarks show that our approach achieves superior performance, while satisfying the real-time and on-board requirements of real-world autonomous driving.
Limitations. The poses in this work reside on a 2D plane, which can be extended to the 3D space in the future. For simplicity, we use L2 distance to obtain the K-nearest neighbors. Future works can explore more sophisticated distances which involve map topology. Currently we use the most basic anchor-based decoder which has limited diversity. This can be improved by using more advanced decoding techniques, such as goal-conditioning. In this paper we focus on marginal motion prediction. It would be interesting to investigate the potential of our approach in other prediction tasks.
This work is funded by Toyota Motor Europe via the research project TRACE-Zürich.
[62] Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tieyan Liu. On layer normalization in the transformer architecture. In ICML, 2020.
[63] Ekim Yurtsever, Jacob Lambert, Alexander Carballo, and Kazuya Takeda. A survey of autonomous driving: Common practices and emerging technologies. IEEE access, 2020.
[64] Fangao Zeng, Bin Dong, Yuang Zhang, Tiancai Wang, Xiangyu Zhang, and Yichen Wei. Motr: End-to-end multiple-object tracking with transformer. In ECCV, 2022.
[65] Wenyuan Zeng, Ming Liang, Renjie Liao, and Raquel Urtasun. Lanercnn: Distributed representations for graph-centric motion forecasting. In IROS, 2021.
[66] Wei Zhan, Liting Sun, Hengbo Ma, Chenran Li, Xiaosong Jia, Masayoshi Tomizuka, et al. Interpret: Interaction-dataset-based prediction challenge. http://challenge.interaction-dataset.com/leader-board, 2021. Accessed: 2023-05-10.
[67] Zhejun Zhang, Alexander Liniger, Dengxin Dai, Fisher Yu, and Luc Van Gool. End-to-end urban driving by imitating a reinforcement learning coach. In ICCV, 2021.
[68] Zhejun Zhang, Alexander Liniger, Dengxin Dai, Fisher Yu, and Luc Van Gool. Trafficbots: Towards world models for autonomous driving simulation and motion prediction. In ICRA, 2023.
[69] Hang Zhao, Jiyang Gao, Tian Lan, Chen Sun, Benjamin Sapp, Balakrishnan Varadarajan, Yue Shen, Yi Shen, Yuning Chai, and Cordelia Schmid. Tnt: Target-driven trajectory prediction. In CoRL, 2021.
[70] Zikang Zhou, Luyao Ye, Jianping Wang, Kui Wu, and Kejie Lu. Hivt: Hierarchical vector transformer for multi-agent motion prediction. In CVPR, 2022.
A Output representation and training strategies
For each anchor token  $\hat{\mathbf{z}}_i^{\mathrm{AG}}, i \in \{1, \dots, N_{\mathrm{AG}} \cdot N_{\mathrm{AC}}\}$ , the confidence head predicts the logits  $p_k, k \in \{1, \dots, 6\}$ , whereas the trajectory head predicts 6 trajectories, each of which is represented as  $(\mu_x^t, \mu_y^t, \log \sigma_x^t, \log \sigma_y^t, \rho^t, v_x^t, v_y^t, \theta^t, s^t)$ ,  $t \in \{1, \dots, T_f\}$ , i.e. the mean of the Gaussian in  $x, y$ , the log standard deviation of the Gaussian in  $x, y$ , the correlation of the Gaussians, the velocity in  $x, y$ , the heading angle and the speed. We denote the ground truth as  $(\hat{x}, \hat{y}, \hat{v}_x, \hat{v}_y, \hat{\theta}, \hat{s})$ . The negative log-likelihood loss for position is formulated as
$$
L _ {\text {p o s}} = - \log \mathcal {N} (\hat {x}, \hat {y} \mid \mu_ {x}, \mu_ {y}, \sigma_ {x}, \sigma_ {y}, \rho). \tag {6}
$$
The negative cosine loss for the heading angle is formulated as
$$
L _ {\text {r o t}} = - \cos (\hat {\theta} - \theta). \tag {7}
$$
The Huber loss for velocities and speed is formulated as
$$
L _ {\mathrm {v e l}} = \mathcal {L} _ {\delta} \left(\hat {v} _ {x} - v _ {x}\right) + \mathcal {L} _ {\delta} \left(\hat {v} _ {y} - v _ {y}\right) + \mathcal {L} _ {\delta} (\hat {s} - s), \tag {8}
$$
where  $\mathcal{L}_{\delta}$  is the Huber loss. We use  $\delta = 1$  for all Huber losses. The final regression loss for a trajectory is the unweighted sum
$$
L _ {\text {t r a j}} = L _ {\text {p o s}} + L _ {\text {r o t}} + L _ {\text {v e l}}, \tag {9}
$$
which is averaged over the future time steps where the ground truth is available. We use a hard assignment strategy, i.e. among the 6 predictions of each agent we select the one that is closest to the ground truth in terms of average displacement error and optimize only for that prediction. Denoting the index of this prediction as  $\hat{k}$ , we train the confidence head via the cross entropy loss by taking  $\hat{k}$  as the ground truth:
$$
L _ {\text {c o n f}} = - \log \frac {\exp \left(p _ {\hat {k}}\right)}{\sum_ {i = 1} ^ {6} \exp \left(p _ {i}\right)}. \tag {10}
$$
The final training loss for the complete model is the unweighted sum
$$
L = L _ {\text {t r a j}} + L _ {\text {c o n f}}. \tag {11}
$$
Our output representation and training strategies are the same as prior works [36, 47, 52], except for the auxiliary losses on velocities, speeds and heading angles.
B.1 KNARPE implementation
Figure 7 shows how KNARPE is implemented with the most basic matrix operations, i.e. matrix indexing, summation and element-wise multiplication.
B.2 Network architectures
We use ReLU activation and set the hidden dimension  $D$  to 256. Our KNARPE is implemented with multi-head attention with 4 heads. We use Transformer with pre-layer normalization [62] with a dropout rate of 0.1. The feed-forward hidden dimension of Transformers is set to 1024. The base frequency  $\omega$  of the sinusoidal positional encoding is set to 1000. We train a single model for all types of agents, while each type of agent has its own anchors. The polyline-line level PointNet and MLPs have 3 layers, the intra-MP Transformer encoder has 6 layers, and the inter-class as well as the AC-to-all Transformer decoders, have 2 layers. Our HPTR has 15.2M trainable parameters in total. The same setup is used for both the WOMD dataset and the AV2 dataset.