DenseTNT: Waymo Open Dataset Motion Prediction Challenge $1^{st}$ Place Solution

Scene context modeling is the first step in behavior prediction. It extracts the features of the lanes and the agents and captures the interactions among them. Sparse encoding methods [5, 8] (also called vectorized methods) were proposed recently. Compared to dense encoding methods which rasterize the lanes and the agents into images and use CNNs to extract features, sparse encoding methods abstract all the geographic entities (e.g. lanes, traffic lights) and vehicles as polylines, and better capture the structural features of high-definition maps.

After scene context encoding, we perform probability estimation for goals on the map. TNT [14] defined discretized sparse anchors on the roads and then assigned probability values upon them. Our key observation is that sparse anchors are not a perfect approximation of real probability distributions on the roads, because (1) one anchor can only generate one goal, we cannot make multi-trajectory predictions around one anchor; (2) there are many ordinary points on the roads (those away from the lane centers or boundaries) that were not well modeled: different ordinary points on the same road have different local information, i.e. the relative distance to the nearest lane boundary.

Therefore, we perform dense goal probability estimation on the map instead. Concretely, a dense goal encoding module is used to extract the features of all the locations on the road under a certain sampling rate. Then, the probability distribution of the dense goals is predicted.

The dense goal encoding module uses an attention mechanism to extract the local information between the goals and the lanes. We denote the feature of the $i_{th}$ goal as $F_{i}$, which is obtained by a 2-layer MLP, and the input of the MLP is 2D coordinates of the $i_{th}$ goal. The local information between the goals and the lanes can be obtained by attention mechanism:

$$Q=FW^{Q},K=LW^{K},V=LW^{V}$$

(1)

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

(2)

where $W^{Q},W^{K},W^{V}\in\mathbb{R}^{d_{h}\times d_{k}}$ are the matrices for linear projection, $d_{k}$ is the dimension of query / key / value vectors, and $F,L$ are feature matrices of the dense goals and all map elements (i.e., lanes or agents), respectively.

The predicted score of the $i_{th}$ goal can be written as:

$$\phi_{i}=\frac{\exp(g(F_{i}))}{\sum^{N}_{n=1}\exp(g(F_{n}))},$$

(3)

where the trainable function $g(\cdot)$ is also implemented with a 2-layer MLP. The loss term for training the scene context encoding and the dense probability estimation is the binary cross-entropy between the predicted goal scores and the ground truth goal scores:

$$\mathcal{L}_{goal}=\sum_{i}\mathcal{L}_{CE}(\phi_{i},\psi_{i}),$$

(4)

where $\psi_{i}$ is the ground truth score of the $i_{th}$ goal. The ground truth score of the goal closest to the final position is $1$, and the others are $0$.

After the dense probability estimation, we use non-maximum suppression (NMS) algorithm to select goals. NMS iteratively selects the goal with the highest probability and removes the goals which are close to the selected goal. The first K selected goals are the predicted goals.

Similar to TNT, the last step is to complete each trajectory conditioned on the selected goals. We only have one ground truth trajectory, so we apply a teacher forcing technique [12] by feeding the ground truth goal during training. The loss term is the offset between the predicted trajectory $\hat{s}$ and the ground truth trajectory $s$:

$$\mathcal{L}_{\rm completion}=\sum\limits^{T}_{t=1}\mathcal{L}_{\rm reg}(\hat{s}_{t},s_{t})$$

(5)

where $\mathcal{L}_{\rm reg}$ is the smooth $l1$ loss between two points.

The previous steps can already achieve good performance in short-term (e.g. 3s) motion prediction tasks. However, long-term prediction is still challenging since the probability distribution may diverge into the long future. Inspired from sentence generation in natural language processing, we generate the probability distribution of the goals in an auto-regressive manner, at 3s, 5s and 8s respectively.

Since we aim to roll out dense probability estimation in 3 steps, we develop three branches in our model architecture.

The three branches share the same weights for the subgraph module in scene context encoding and have independent weights for other parts, e.g. the global graph module in scene context encoding and the dense probability estimation.

With $N$ goal selection at 3s, 5s and 8s auto-regressively, we obtain $N^{3}$ goal sets. We sort the top K goal sets according to their probability scores, and then complete them to obtain K trajectories.

More concretely, for each goal set, we use the above dense goal encoding module to get the features of the 3 goals. Then the features are passed to the trajectory completion module that is a 2-layer MLP. The output is a full trajectory $[\hat{s}_{1},\hat{s}_{2},...,\hat{s}_{T}]$.

In DenseTNT, the goal candidates are densely distributed on the map. We visualize the probabilities of the dense goals and the predicted trajectories based on the selected goals. As shown in Figure 3, DenseTNT gives diverse predictions such as going straight, left/right turns and U-turns.