Learning Distilled Collaboration Graph
for Multi-Agent Perception

Feature encoder. The functionality is to extract informative features from raw measurements for each agent. Let $\mathcal{X}_{i}$ be the 3D point cloud collected by the $i$th agent ($M$ agents in total), the feature of the $i$th agent is obtained as $\mathbf{F}_{i}^{\rm s}\leftarrow\Theta^{\rm s}(\mathcal{X}_{i})$, where $\Theta^{\rm s}(\cdot)$ is the feature encoder shared by all the agents and the superscript ${\rm s}$ reflects the student mode. To implement the encoder, we convert a 3D point cloud to a bird’s-eye-view (BEV) map, which is amenable to classic 2D convolutions. Specifically, we quantize the 3D points into regular voxels and represent the 3D voxel lattice as a 2D pseudo-image, with the height dimension corresponding to image channels. Such a 2D image is virtually a BEV map, whose basic element is a cell that is associated with a binary vector along the vertical axis. Let $\mathbf{B}_{i}\in\{0,1\}^{K\times K\times C}$ be the BEV map of the $i$th agent associated with $\mathcal{X}_{i}$. With this map, we can apply four blocks composed of 2D convolutions, batch normalization and ReLU activation to gradually reduce the spatial dimension and increase the number of channels for the BEV, to obtain the feature map $\mathbf{F}_{i}^{\rm s}\in\mathbb{R}^{\bar{K}\times\bar{K}\times\bar{C}}$ with $\bar{K}\times\bar{K}$ the spatial resolution and $\bar{C}$ the number of channels.

Feature compression. To save the communication bandwidth, each agent could compress its feature map prior to transmission. Here we consider a $1\times 1$ convolutional autoencoder [24] to compress/decompress the feature maps along the channel dimension. The autoencoder is trained together with the whole system, making the system work with limited collaboration information.

Collaboration graph process. The functionality is to update the feature map through data transmission among the agents. The core component here is a collaboration graph¹¹1We consider a fully-connected bidirectional graph, and the weights for both directions are distinct. $\mathcal{G}_{\Pi}(\mathcal{V},\mathcal{E}_{\Pi})$, where $\mathcal{V}$ is the fixed node set and $\mathcal{E}_{\Pi}$ is the trainable edge set. Each node in $\mathcal{V}$ is an agent with the real-time pose information; for instance, $\bm{\xi}_{i}\in\mathfrak{se}(3)$ is the $i$th agent’s pose in the global coordinate system; and each edge in $\mathcal{E}_{\Pi}$ is trainable and models the collaboration between two agents, with $\Pi$ an edge-weight encoder, reflecting the trainable collaboration strength between agents. Let $\mathcal{M}_{\mathcal{G}_{\Pi}}(\cdot)$ be the collaboration process defined on the collaboration graph $\mathcal{G}_{\Pi}$. The feature maps of all the agents after collaboration are $\{\mathbf{H}_{i}^{\rm s}\}_{i=1}^{M}\leftarrow\mathcal{M}_{\mathcal{G}_{\Pi}}\left(\{\mathbf{F}_{i}^{\rm s}\}_{i=1}^{M}\right)$. This process has three stages: neural message transmission (S1), neural message attention (S2) and neural message aggregation (S3).

In the neural message transmission stage (S1), each agent transmits its BEV-based feature map to the other agents. Since the BEV-based feature map summarizes the information of each agent, we consider it as a neural message. In the neural message attention stage (S2), each agent receives others’ neural messages and determines the matrix-valued edge weights, which reflect the importance of the neural message from one agent to another at each individual cell. Since each agent has its unique pose, we leverage the collaboration graph to achieve feature transformation across agents. For the $i$th agent, the transformed BEV-based feature map from the $j$th agent is then $\mathbf{F}_{j\to i}^{\rm s}=\Gamma_{j\to i}(\mathbf{F}_{i}^{\rm s})\in\mathbb{R}^{\bar{K}\times\bar{K}\times\bar{C}}$, where the transformation $\Gamma_{j\to i}(\cdot)$ is based on two ego poses $\bm{\xi}_{j}$ and $\bm{\xi}_{i}$. Now $\mathbf{F}_{j\to i}^{\rm s}$ and $\mathbf{F}_{i}^{\rm s}$ are supported in the same coordinate system. To determine the edge weights, we use the edge encoder $\Pi$ to correlate the ego feature map and the feature map from another agent; that is, the matrix-valued edge weight from the $j$th agent to the $i$th agent is $\mathbf{W}_{j\to i}=\Pi(\mathbf{F}_{j\to i}^{\rm s},\mathbf{F}_{i}^{\rm s})\in\mathbb{R}^{\bar{K}\times\bar{K}}$, where $\Pi$ concatenates two feature maps along the channel dimension and then uses four $1\times 1$ convolutional layers to gradually reduce the number of channels from $2\bar{C}$ to $1$. Meanwhile, there is a softmax operation applied at each cell in the feature map to normalize the edge weights across multiple agents. Note that previous works [19, 20, 34] generally consider a scalar-valued edge weight to reflect the overall collaboration strength between two agents; while we consider a matrix-valued edge weight $\mathbf{W}_{j\to i}$, which models the collaboration strength from the $j$th agent to the $i$th agent with a $\bar{K}\times\bar{K}$ spatial resolution. In this matrix, each element corresponds to a cell in the BEV map, indicating a specific spatial region; thus, this matrix reflects the spatial attention at a cell-level resolution. In the neural message aggregation stage (S3), each agent aggregates the feature maps from all the agents based on the normalized matrix-valued edge weights. The updated feature map of the agent $i$ is ${\mathbf{H}}_{i}^{\rm s}=\sum^{M}_{j=1}\mathbf{W}_{j\to i}\odot\mathbf{F}_{j\to i}^{\rm s},$ where $\odot$ denotes the dot product with channel-wise broadcasting.

Remark. The proposed collaboration graph is trainable because each matrix-valued edge weight is a trainable matrix to reflect the agent-to-agent attention in a cell-level spatial resolution; it is pose-aware, empowering all the agents to work with the synchronized coordinate system; furthermore, it is dynamic at each timestamp as each edge weight would adapt to the real-time neural messages. According to the proposed collaboration graph, the agents can discover the region requiring collaboration on the fly, and strategically select appropriate partners to request supplementary information.

Decoder and header. After collaboration, each agent decodes the updated BEV-based feature map. The decoded feature map is $\mathbf{M}_{i}^{\rm s}\leftarrow\Psi^{\rm s}({\mathbf{H}}_{i}^{{\rm s}})$. To implement the decoder $\Psi^{\rm s}(\cdot)$, we progressively up-sample ${\mathbf{H}}_{i}^{{\rm s}}$ with four layers, where each layer first concatenates the previous feature map with the corresponding feature map in the encoder and then uses a $1\times 1$ convolutional operation to halve the number of channels. Finally, we use an output header to generate the final detection outputs, $\widehat{\mathbf{Y}}_{i}^{\rm s}\leftarrow\Phi^{\rm s}(\mathbf{M}_{i}^{\rm s})$. To implement the header $\Psi^{\rm s}(\cdot)$, we use two branches of convolutional layers to classify the foreground-background categories and regress the bounding boxes.

During the training phase, an early-collaboration model, as the teacher, is introduced to guide the intermediate-collaboration model, which is a student. Similar to the student model, the teacher’s pipeline has the feature encoder $\Theta^{\rm t}$, feature decoder $\Psi^{\rm t}$ and the output header $\Phi^{\rm t}$. Note that all the agents share the same teacher model to guide one student model; however, each agent provides the inputs with its own pose, and its inputs to both the teacher and student models should be well aligned.

Feature encoder. Let $\mathcal{X}=\mathcal{A}(\bm{\xi}_{1}\circ\mathcal{X}_{1},\bm{\xi}_{2}\circ\mathcal{X}_{2},...,\bm{\xi}_{M}\circ\mathcal{X}_{M})$ be a holistic-view 3D point cloud that aggregates all the points from all $M$ agents in the global coordinate system, where $\mathcal{A}(\cdot,...,\cdot)$ is the aggregation operator of multiple 3D point clouds, and $\bm{\xi}_{i}$ and $\mathcal{X}_{i}$ are the pose and the 3D point cloud of the $i$th agent, respectively. To ensure the inputs to the teacher model and the student model are aligned, we transform the holistic-view point cloud $\mathcal{X}$ to an agent’s own coordinate based on the pose information. Now, for the $i$th agent, the input to the teacher model $\bm{\xi}^{-1}_{i}\circ\mathcal{X}$ and the input to the student model $\mathcal{X}_{i}$ are in the same coordinate system. Similarly to the feature encoder in the student model, we convert the 3D point cloud $\bm{\xi}^{-1}_{i}\circ\mathcal{X}$ to a BEV map and use 2D convolutions to obtain the feature map of the $i$th agent in the teacher model, $\mathbf{H}_{i}^{\rm t}\in\mathbb{R}^{\bar{K}\times\bar{K}\times\bar{C}}$. Here we crop the BEV map to ensure it has the same spatial range and resolution with the BEV map in the student model.

Decoder and header. Similarly to the decoder and header in the student model, we adopt $\mathbf{M}_{i}^{\rm t}\leftarrow\Psi^{\rm t}({\mathbf{H}}_{i}^{{\rm t}})$ to obtain the decoded BEV-based feature map and $\widehat{\mathbf{Y}}_{i}^{\rm t}\leftarrow\Phi^{\rm t}(\mathbf{M}_{i}^{\rm t})$ to obtain the predicted foreground-background categories and regressed bounding boxes.

Teacher training scheme. As in common teacher-student frameworks, we train the teacher model separately. We employ the binary cross-entropy loss to supervise foreground-background classification and the smooth $L_{1}$ loss to supervise the bounding-box regression. Overall, we minimize the loss function $\mathcal{L}^{\rm t}=\sum_{i=1}^{M}\mathcal{L}_{\rm det}(\mathbf{Y}_{i}^{\rm t},\widehat{\mathbf{Y}}_{i}^{\rm t})$, where classification and regression losses are collectively denoted as $\mathcal{L}_{\rm det}$ and $\mathbf{Y}_{i}^{\rm t}=\mathbf{Y}_{i}^{\rm s}$ is the ground-truth detection in the perception region of the $i$th agent.

Given a well-trained teacher model, we use both detection loss and knowledge distillation loss to supervise the training of the student model. We consider minimizing the following loss

The detection loss $\mathcal{L}_{\rm det}$ is similar to that of the teacher, including both foreground-background classification loss and the bounding box regression loss, pushing the detection result of each agent to be close to its local ground-truth. The second and third terms form a knowledge distillation loss, regularizing the student model to generate similar feature maps with the teacher model. The hyperparameter $\lambda_{\rm kd}$ controls the weight of the knowledge distillation loss $\mathcal{L}_{\rm kd}$ defined as follows

where $D_{KL}(p(\mathbf{x})||q(\mathbf{x}))$ denotes the Kullback-Leibler (KL) divergence of distribution $q(\mathbf{x})$ from distribution $p(\mathbf{x})$, $\sigma(\cdot)$ indicates the softmax operation of the feature vector along the channel dimension, and $\left(\mathbf{H}_{i}^{\rm s}\right)_{n}$ and $\left(\mathbf{H}_{i}^{\rm t}\right)_{n}$ denote the feature vectors at the $n$th cell of the $i$th agent’s feature map in the student model and teacher model, respectively. Similarly, the loss on decoded feature maps $\mathcal{L}_{\rm kd}(\mathbf{M}_{i}^{\rm s},\mathbf{M}_{i}^{\rm t})$ can be introduced to enhance the regularization.

Remark. As mentioned in Section 2.1, the feature maps output by the collaboration graph process in the student model is computed by $\{\mathbf{H}_{i}^{\rm s}\}_{i=1}^{M}\leftarrow\mathcal{M}_{\mathcal{G}_{\Pi}}\left(\{\mathbf{F}_{i}^{\rm s}\}_{i=1}^{M}\right)$. Intuitively, the feature map of the teacher model $\{\mathbf{H}_{i}^{\rm t}\}_{i=1}^{M}$ would be the desired output of the collaboration graph process $\mathcal{M}_{\mathcal{G}_{\Pi}}(\cdot)$. Therefore, we constrain all the post-collaboration feature maps in the student model to match the correspondences in the teacher model through knowledge distillation. This constraint would further regularize the upfront trainable components: i) the distilled student encoder $\Theta^{\rm s}$, which abstracts the features from raw measurements and produces the input to $\mathcal{M}_{\mathcal{G}_{\Pi}}(\cdot)$, and ii) the edge-weight encoder $\Pi$ in the distilled collaboration graph. Consequently, through knowledge distillation and back-propagation, the distilled student encoder would learn to abstract informative features from raw data for better collaboration; and the distilled edge-weight encoder would learn how to control the collaboration based on agents’ features. In a word, our distilled collaboration network (DiscoNet) can comprehend feature abstraction and fusion via the proposed knowledge distillation framework.

Since there is no public dataset to support the research on multi-agent 3D object detection in self-driving scenarios, we build such a dataset named as V2X-Sim $1.0$²²2V2X (vehicle-to-everything) denotes the collaboration between a vehicle and other entities such as vehicle (V2V) and roadside infrastructure (V2I). V2X-Sim dataset is maintained on https://ai4ce.github.io/V2X-Sim/, and the first version of V2X-Sim used in this work includes the LiDAR-based V2V scenario., based on both SUMO [15], a micro-traffic simulation, and CARLA [6], a widely-used open-source simulator for autonomous driving research. SUMO is firstly used to produce numerically-realistic traffic flow and CARLA is employed to get realistic sensory streams including LiDAR and RGB data, from multiple vehicles located in the same geographical area. The simulated LiDAR is rotated at a frequency of 20Hz and has 32 channels; the range which denotes the maximum distance to capture is 70 meters. We use Town05 with multiple lanes and cross junctions for dense traffic simulation; see Fig. 3.

V2X-Sim $1.0$ follows the same storage format of nuScenes [2], a widely-used autonomous driving dataset. nuScenes collected real-world single-agent data; while we simulate multi-agent scenarios. Each scene includes a 20 second traffic flow at a certain intersection of Town05, and the LiDAR streams are recorded every 0.2 second, so each scene consists of 100 frames. We generate 100 scenes with a total of 10,000 frames, and in each scene, we randomly select 2-5 vehicles as the collaboration agents. We use 8,000/900/1,100 frames for training/validation/testing. Each frame has multiple samples, and there are 23,500 samples in the training set and 3,100 samples in the test set.

Implementation and evaluation. We crop the points located in the region of $[-32,32]\times[-32,32]\times[-3,2]$ meters defined in the ego-vehicle $XYZ$ coordinate. We set the width/length of each voxel as 0.25 meter, and the height as 0.4 meter; therefore the BEV map input to the student/teacher encoder has a dimension of $256\times 256\times 13$. The intermediate feature output by the encoder has a dimension of $32\times 32\times 256$, and we use an autoencoder (AE) for compression to save bandwidth: the sent message is the embedding output by AE. After the agents receive the embedding, they will decode it to the original resolution $32\times 32\times 256$ for intermediate collaboration. The hyperparameter $\lambda_{\rm kd}$ is set as $10^{5}$. We train all the models using NVIDIA GeForce RTX 3090 GPU. We employ the generic BEV detection evaluation metric: Average Precision (AP) at Intersection-over-Union (IoU) threshold of 0.5 and 0.7. We target the detection of the car category and report the results on the test set.

Baselines and existing methods. We consider the single-agent perception model; called Lower-bound, which consumes a single-view point cloud. The teacher model based on early collaboration with a holistic view is considered as the Upper-bound. Both lower-bound and upper-bound can involve late collaboration, which aggregates all the boxes and applies a global non-maximum suppression (NMS). We also consider three intermediate-collaboration methods, Who2com [20], When2com [19] and V2VNet [34]. Since the original Who2com and When2com do not consider pose information, we consider both pose-aware and pose-agnostic versions to achieve fair comparisons. All the methods use the same detection architecture and conduct collaboration at the same intermediate feature layer.

Results and comparisons. Table 2 shows the comparisons in terms of AP (@IoU = 0.5/0.7). We see that i) early collaboration (first row) achieves the best detection performance, and there is an obvious improvement over no collaboration (last row), i.e., [email protected] and [email protected] are increased by 38.2$\%$ and 42.3$\%$ respectively, reflecting the significance of collaboration; ii) late collaboration improves the lower-bound (second last row), but hurts the upper-bound (second row), reflecting the unreliability of late collaboration. This is because the final, global NMS can remove a lot of noisy boxes for the lower-bound, but would also remove many useful boxes for the upper-bound; and iii) among the intermediate collaboration-based methods, the proposed DiscoNet achieves the best performance. Comparing to the pose-aware When2com, DiscoNet improves by 31.9$\%$ in [email protected] and 29.3$\%$ in [email protected]. Comparing to V2VNet, DiscoNet improves by 6.2 $\%$ in [email protected] and 6.3 $\%$ in [email protected]. Even with $16$-times compression by the autoencoder, DiscoNet(16) still outperforms V2VNet.

Performance-bandwidth trade-off analysis. Fig. 4 thoroughly compares the proposed DiscoNet with the baseline methods in terms of the trade-off between detection performance and communication bandwidth. In Plots (a) and (b), the lower-bound with no collaboration is shown as a dashed line since the communication volume is zero, and the other method is shown as one dot in the plot. We see that i) DiscoNet with feature-map compression by an autoencoder achieves a far-more superior trade-off than the other methods; ii) DiscoNet(64), which is DiscoNet with $64$-times feature compression by autoencoder, achieves 192 times less communication volume and still outperforms V2VNet for both [email protected] and 0.7; and iii) feature compression does not significantly hurt the detection performance.

Comparison of training times with and without KD. Incorporating KD will increase the training time a little bit, from ~1200s to ~1500s per epoch (each epoch consists of 2,000 iterations). However, during inference, DiscoNet does not need the teacher model anymore and can work alone without extra computations. Note that our key advantage is to push the computation burden into the training stage in pursuit of efficient and effective inference which is crucial in the real-world deployment.

Visualization of edge weight. To understand the working mechanism of the proposed collaboration graph, we visualize the detection results and the corresponding edge weights; see Fig. 5 and Fig. 7. Note that the proposed edge weight is a matrix, reflecting the collaboration attention in a cell-level resolution, which is shown as a heat map. Fig. 5 (a), (b) and (c) show three exemplar detection results of Agent 1 based on the lower-bound, upper-bound and the proposed DiscoNet, respectively. We see that with the collaboration, DiscoNet is able to detect boxes in those occluded and long-range regions. To understand why, Fig. 5 (d) and (e) provide the corresponding ego edge weight and the edge weight from agent 2 to agent 1, respectively. We clearly see that with the proposed collaboration graph, Agent 1 is able to receive complementary information from the other agents. Take the first row in Fig. 5 as an example, we see that in Plot (d), the bottom-left spatial region has a much darker color, indicating that Agent 1 has less confidence about this region; while in Plot (e), the bottom-left spatial region has a much brighter color, indicating that Agent 1 has much stronger demands to request information from Agent 2.

Comparison with When2com and V2VNet. Fig. 6 shows two examples to compare the proposed DiscoNet with When2com, V2VNet. We see that DiscoNet is able to detect more objects. The reason is that both V2VNet and when2com employ a scalar to denote the agent-to-agent attention, which cannot distinguish which region is more informative; while DiscoNet can adaptively find beneficial region in a cell-level resolution; see the visualization of matrix-valued edge weights in Fig. 6 (d)-(f).

Effect of collaboration strategy. Table 3 compares the proposed method with five baselines using different intermediate collaboration strategies. Sum simply sums all the intermediate feature maps. Average calculates the mean of all the feature maps. Max selects the maximum value of all the feature maps at each cell to produce the final feature. Cat concatenates the mean of others’ features with its own feature, and then uses a convolution layer to halve the channel number. Weighted Average appends one convolutional layer to our edge encoder $\Pi$, to generate a scalar value as the agent-wise edge weight. We use a matrix-valued edge weight to model the collaboration among agents. Without knowledge distillation, Max has achieved the second best performance probably because it ignores the noisy information with lower value. We see that with the proposed matrix-valued edge weight and knowledge distillation, DiscoGraph significantly outperforms all the other collaboration choices.

Effect of knowledge distillation. In Table 3 we investigate the versions with and without knowledge distillation, We see that i) for our method, the knowledge distillation can guide the learning of collaboration graph, and the agents can work collaboratively to approach the teacher’s performance; ii) for max without a learnable module during collaboration, knowledge distillation has no impact; and iii) for cat and average, their performances are improved a little bit, as knowledge distillation can influence the feature abstraction process. Table 2 further shows the detection performances when we apply knowledge-distillation regularization at various network layers. We see that i) once we apply knowledge distillation to regularize the feature map, the proposed method starts to achieve improvement; and ii) applying regularization on four layers has the best performance.

We describe the architecture of the encoder below.

Sequential(
    Conv2d(13, 32, kernel_size=3, stride=1, padding=1)
    BatchNorm2d(32)
    ReLU()
    Conv2d(32, 32, kernel_size=3, stride=1, padding=1)
    BatchNorm2d(32)
    ReLU()
    Conv3D(64, 64, kernel_size=(1, 1, 1), stride=1, padding=(0, 0, 0))
    Conv3D(128, 128, kernel_size=(1, 1, 1), stride=1, padding=(0, 0, 0))
    Conv2d(32, 64, kernel_size=3, stride=2, padding=1)  ->(32,256,256)
    BatchNorm2d(64)
    ReLU()
    Conv2d(64, 64, kernel_size=3, stride=1, padding=1)
    BatchNorm2d(64)
    ReLU() ->(64,128,128)
    Conv2d(64, 128, kernel_size=3, stride=2, padding=1)
    BatchNorm2d(128)
    ReLU()
    Conv2d(128, 128, kernel_size=3, stride=1, padding=1)
    BatchNorm2d(128)
    ReLU() ->(128,64,64)
    Conv2d(128, 256, kernel_size=3, stride=2, padding=1)
    BatchNorm2d(256)
    ReLU()
    Conv2d(256, 256, kernel_size=3, stride=1, padding=1)
    BatchNorm2d(256)
    ReLU() ->(256,32,32)
    Conv2d(256, 512, kernel_size=3, stride=2, padding=1)
    BatchNorm2d(512)
    ReLU()
    Conv2d(512, 512, kernel_size=3, stride=1, padding=1)
    BatchNorm2d(512)
    ReLU() ->(512,16,16)
)

The input of the decoder is the intermediate feature output by each layer of the encoder. Its architecture is shown below.

Sequential(
    Conv2d(512 + 256, 256, kernel_size=3, stride=1, padding=1)
    BatchNorm2d(256)
    ReLU()
    Conv2d(256, 256, kernel_size=3, stride=1, padding=1)
    BatchNorm2d(256)
    ReLU() ->(256,32,32)
    Conv2d(256 + 128, 128, kernel_size=3, stride=1, padding=1)
    BatchNorm2d(128)
    ReLU()
    Conv2d(128, 128, kernel_size=3, stride=1, padding=1)
    BatchNorm2d(128)
    ReLU() ->(128,64,64)
    Conv2d(128 + 64, 64, kernel_size=3, stride=1, padding=1)
    BatchNorm2d(64)
    ReLU()
    Conv2d(64, 64, kernel_size=3, stride=1, padding=1)
    BatchNorm2d(64)
    ReLU() ->(64,128,128)
    Conv2d(64 + 32, 32, kernel_size=3, stride=1, padding=1)
    BatchNorm2d(32)
    ReLU()
    Conv2d(32, 32, kernel_size=3, stride=1, padding=1)
    BatchNorm2d(32)
    ReLU() ->(32,256,256)
)

The input whose dimension is $(c,w,h)=(512,32,32)$ of the edge encoder is the concatenation of the feature from ego agent and that from its neighbor agents. The output is a matrix-valued edge weight.

Sequential(
    Conv2d(512, 128, kernel_size=1, stride=1, padding=0)
    BatchNorm2d(128)
    ReLU()
    Conv2d(128, 32, kernel_size=1, stride=1, padding=0)
    BatchNorm2d(32)
    ReLU()
    Conv2d(32, 8, kernel_size=1, stride=1, padding=0)
    BatchNorm2d(8)
    ReLU()
    Conv2d(8, 1, kernel_size=1, stride=1, padding=0)
    BatchNorm2d(1)
    ReLU()
)