VectorFlow: Combining Images and Vectors for Traffic Occupancy and Flow Prediction

We process the input into a rasterized image and a set of vectors. To obtain the image, we create a rasterized grid at each time step in the past given the observed agent trajectories and map data, with respect to the local coordinate of the self-driving car (SDC)²²2More detail can be found at https://github.com/waymo-research/waymo-open-dataset/blob/master/tutorial/tutorial_occupancy_flow.ipynb. To obtain the vectorized input that is consistent with the rasterized image, we follow the same transformation by rotating and shifting the input agent and map coordinates with respect to SDC’s local view.

The encoder includes two parts: a VGG-16 model [10] that encodes the rasterized representation, and a VectorNet [1] model that encodes the vectorized representation. We fuse the vectorized features with the features of the last two stages of VGG-16 by cross-attention modules. The fused features are upsampled to the original resolution as the input rasterized features by an FPN-style network. More detail can be found in Fig. 1 and Sec. 4.3.

The decoder is a single 2D convolution layer that maps the output of the encoder to the occupancy flow fields prediction, which includes a sequence of 8 grid maps representing the predicted occupancy and flow at each time step over the next 8 seconds.

We follow the loss function as in [8], including a cross entropy loss on observed occupancy prediction $\mathcal{L}_{O,b}$, a cross entropy loss on occluded occupancy prediction $\mathcal{L}_{O,c}$, and an L2 loss on flow prediction $\mathcal{L}_{F}$. The total loss is:

We train, evaluate, and test our model on the Waymo Open Motion Dataset (WOMD) based on the standard split and the filtered scenarios.

We follow the metrics proposed by [8], which include AUC and Soft IoU for observed occupancy, occluded occupancy, and flow-grounded occupancy, as well as the end-point error (EPE) that measures the error of flow prediction.

Our model is illustrated in Fig 1. We use the standard VGG-16 model from torchvision.models as our rasterized encoder, and follow the implementations of VectorNet as in [2]³³3Code available at https://github.com/Tsinghua-MARS-Lab/DenseTNT. The input to the VectorNet includes i) a set of road element vectors with a shape of $B\times N_{R}\times 9$, where $B$ is the batch size, $N_{R}=10000$ is the maximum number of road element vectors, and the last dimension of 9 represents the positions ($x,y$) and headings ($\cos\theta,\sin\theta)$ of two end points in each vector and the vector id; ii) a set of agent vectors with a shape of $B\times 1280\times 9$, including the vectors of up to 128 agents in a scene, where each agent has 10 vectors from the observed positions. We follow VectorNet by first running a local graph over each traffic element based on their ids and second running a global graph over all local features to obtain a vectorized feature with a shape of $B\times 128\times N$, where $N$ is the total number of traffic elements, including road elements and agents. We further quadrupled the size of the feature through an MLP layer to obtain a final vectorized feature $V$ with a shape of $B\times 512\times N$, so that its feature size is consistent with the channel size of the image feature, as discussed in the following paragraph.

We denote the output features of each VGG stage as $\{C_{1},C_{2},C_{3},C_{4},C_{5}\}$, and they have strides of $\{1,2,4,8,16\}$ pixels with respect to the input image and a hidden dimension of 512. The vectorized feature $V$ is fused with the rasterized image feature $C_{5}$ with a shape of $B\times 512\times 16\times 16$ by a cross-attention module to obtain $F_{5}$ with the same shape. The query term of the cross-attention is the image feature $C_{5}$ flattened into $B\times 512\times 256$ with 256 tokens, and the key and value term is the vectorized feature $V$ with $N$ tokens. We then concatenate $F_{5}$ and $C_{5}$ in the channel dimension and pass it through two $3\times 3$ conv layers to obtain $P_{5}$ with a shape of $B\times 512\times 16\times 16$. $P_{5}$ is upsampled and concatenated with $C_{4}$ ($B\times 512\times 32\times 32$) by an FPN-style $2\times 2$ upsampling module to generate $U_{4}$ with the same shape as $C_{4}$. Next, we perform another round of fusion between $V$ and $U_{4}$ to obtain $P_{4}$ ($B\times 512\times 32\times 32$) following the same procedure, including cross attention. At the end, $P_{4}$ will be gradually upsampled by the FPN-style network and concatenated with $\{C_{3},C_{2},C_{1}\}$ to generate $P_{1}$ with a shape of $B\times 512\times 256\times 256$. We pass $P_{1}$ through two $3\times 3$ conv layers to obtain the final output feature with a shape of $B\times 128\times 256\times 256$.

The decoder is a single 2D convolution layer with an input channel size of 128 and an output channel size of 32 (8 waypoints $\times$ 4 output dimensions).

We train our model on the full training set of WOMD with a batch size of 32 for 16 epochs on 8 Nvidia A10 GPUs. We use an Adam optimizer and a learning rate scheduler that decays the learning rate by 50% every 5 epochs, with an initial value of 1e-3. The loss coefficients are $\alpha=\beta=1000$, and $\gamma=1$, as customary in [8].

We present the results of our model and other entries in the Waymo Challenge in Table 1. Our model achieves the best performance in three metrics, including AUC and Soft IoU scores for occluded occupancy predictions, and Soft IoU score for observed occupancy predictions. The gap in AUC scores is partially due to the choice of our loss function, compared to other entries that use a focal loss. Our model ranks 3rd place on the Waymo Open Dataset Occupancy and Flow Prediction Challenge.

Furthermore, we compare the performance of our model with a variant that only includes the VGG encoder. The results in Table 2 show that fusion helps improve the prediction performance, especially in occluded metrics. More specifically, the Occluded AUC and Occluded Soft IoU scores improve by 24.46% and 47.06%, respectively.