Hydra-MDP: End-to-end Multimodal Planning with Multi-target Hydra-Distillation

Zhenxin Li, Kailin Li, Shihao Wang, Shiyi Lan, Zhiding Yu, Yishen Ji, Zhiqi Li, Ziyue Zhu, Jan Kautz, Zuxuan Wu, Yu-Gang Jiang, Jose M. Alvarez

cs.CV

Introduction

End-to-end autonomous driving, which involves learning a neural planner with raw sensor inputs, is considered a promising direction to achieve full autonomy. Despite the promising progress in this field , recent studies have exposed multiple vulnerabilities and limitations of imitation learning (IL) methods, particularly the inherent issues in open-loop evaluation, such as the dysfunctional metrics and implicit biases . This is critical as it fails to guarantee safety, efficiency, comfort, and compliance with traffic rules. To address this main limitation, several works have proposed incorporating closed-loop metrics, which more effectively evaluate end-to-end autonomous driving by ensuring that the machine-learned planner meets essential criteria beyond merely mimicking human drivers.

Therefore, end-to-end planning is ideally a multi-target and multimodal task, where multi-target planning involves meeting various evaluation metrics from either open-loop and closed-loop settings. In this context, multimodal indicates the existence of multiple optimal solutions for each metric.

Existing end-to-end approaches often try to consider closed-loop evaluation via post-processing, which is not streamlined and may result in the loss of additional information compared to a fully end-to-end pipeline. Meanwhile, rule-based planners struggle with imperfect perception inputs. These imperfect inputs degrade the performance of rule-based planning under both closed-loop and open-loop metrics, as they rely on predicted perception instead of ground truth (GT) labels.

To address the issues, we propose a novel end-to-end autonomous driving framework called Hydra-MDP (Multimodal Planning with Multi-target Hydra-distillation). Hydra-MDP is based on a novel teacher-student knowledge distillation (KD) architecture. The student model learns diverse trajectory candidates tailored to various evaluation metrics through KD from both human and rule-based teachers. We instantiate the multi-target Hydra-distillation with a multi-head decoder, thus effectively integrating the knowledge from specialized teachers. Hydra-MDP also features an extendable KD architecture, allowing for easy integration of additional teachers.

The student model uses environmental observations during training, while the teacher models use ground truth (GT) data. This setup allows the teacher models to generate better planning predictions, helping the student model to learn effectively. By training the student model with environmental observations, it becomes adept at handling realistic conditions where GT perception is not accessible during testing.

Our contributions are summarized as follows:

We propose a universal framework of end-to-end multimodal planning via multi-target hydra-distillation, allowing the model to learn from both rule-based planners and human drivers in a scalable manner.

Our approach achieves the state-of-the-art performance under the simulation-based evaluation metrics on Navsim.

Solution

Let $O$ represent sensor observations, $\hat{P}$ and $P$ denote ground truth and predicted perceptions (e.g. 3D object detection, lane detection), $\hat{T}$ be the expert trajectory, and $T^{*}$ be the predicted trajectory. $\mathcal{L}_{im}$ represents the imitation loss. We first introduce the two prevailing paradigms and our proposed paradigm (Fig. 1) in this section:

A. Single-modal Planning + Single-target Learning. In this paradigm , the planning network directly regresses the planned trajectory from the sensor observations. Ground truth perceptions can be used as auxiliary supervision but does not influence the planning output. Perception losses are not included in the formula for simplicity. The whole processing can be formulated as:

where $\mathcal{L}_{im}$ is usually an L2 loss.

B. Multimodal Planning + Single-target Learning. This approach predicts multiple trajectories $\{T_{i}\}_{i=1}^{k}$ , whose similarities to the expert trajectory are computed:

where $\mathcal{L}_{im}$ can be KL-Divergence or the max-margin loss . Perception outputs $P$ are explicitly used to post-process suitable trajectories via a cost function $f(T_{i},P)$ . The trajectory with the lowest cost is selected:

which is a non-differentiable process based on imperfect perception $P$ .

Here, we only consider one cost function $f$ for clarity. The trajectory with the lowest predicted cost is selected:

We stress that this framework is not restricted by non-differentiable post-processing. It can be easily scaled in an end-to-end fashion by involving more cost functions or leveraging imitation similarity in our implementation (Sec. 2.4).

2 Overall Framework

As shown in Fig. 2, Hydra-MDP consists of two networks: a Perception Network and a Trajectory Decoder.

Perception Network. Our perception network builds upon the official challenge baseline Transfuser , which consists of an image backbone, a LiDAR backbone, and perception heads for 3D object detection and BEV segmentation. Multiple transformer layers connect features from stages of both backbones, extracting meaningful information from different modalities. The final output of the perception network comprises environmental tokens $F_{env}$ , which encode abundant semantic information derived from both images and LiDAR point clouds.

Trajectory Decoder. Following Vadv2 , we construct a fixed planning vocabulary to discretize the continuous action space. To build the vocabulary, we first sample 700K trajectories randomly from the original nuPlan database . Each trajectory $T_{i}(i=1,...,k)$ consists of 40 timestamps of $(x,y,heading)$ , corresponding to the desired 10Hz frequency and a 4-second future horizon in the challenge. The planning vocabulary $\mathcal{V}_{k}$ is formed as K-means clustering centers of the 700K trajectories, where $k$ denotes the size of the vocabulary. $\mathcal{V}_{k}$ is then embedded as $k$ latent queries with an MLP, sent into layers of transformer encoders , and added to the ego status $E$ :

To incorporate environmental clues in $F_{env}$ , transformer decoders are leveraged:

Using the log-replay trajectory $\hat{T}$ , we implement a distance-based cross-entropy loss to imitate human drivers:

where $\mathcal{S}^{im}_{i}$ is the $i$ -th softmax score of $\mathcal{V}^{\prime\prime}_{k}$ , and $y_{i}$ is the imitation target produced by L2 distances between log-replays and the vocabulary. Softmax is applied on L2 distances to produce a probability distribution:

The intuition behind this imitation target is to reward trajectory proposals that are close to human driving behaviors.

3 Multi-target Hydra-Distillation

Though the imitation target provides certain clues for the planner, it is insufficient for the model to associate the planning decision with the driving environment under the closed-loop setting, leading to failures such as collisions and leaving drivable areas . Therefore, to boost the closed-loop performance of our end-to-end planner, we propose Multi-target Hydra-Distillation, a learning strategy that aligns the planner with simulation-based metrics in this challenge.

The distillation process expands the learning target through two steps: (1) running offline simulations of the planning vocabulary $\mathcal{V}_{k}$ for the entire training dataset; (2) introducing supervision from simulation scores for each trajectory in $\mathcal{V}_{k}$ during the training process. For a given scenario, step 1 generates ground truth simulation scores $\{<math xmlns="http://www.w3.org/1998/Math/MathML" display="block"><semantics><mrow><msubsup><mover accent="true"><mi mathvariant="script">S</mi><mo>^</mo></mover><mi>i</mi><mi>m</mi></msubsup></mrow><annotation encoding="application/x-tex">\hat{\mathcal{S}}^{m}_{i}</annotation></semantics></math>S^im|i=1,...,k\}_{m=1}^{|M|}$ for each metric $m\in M$ and the $i$ -th trajectory, where $M$ represents the set of closed-loop metrics used in the challenge. For score predictions, latent vectors $\mathcal{V}^{\prime\prime}_{k}$ are processed with a set of Hydra Prediction Heads, yielding predicted scores $\{<math xmlns="http://www.w3.org/1998/Math/MathML" display="block"><semantics><mrow><msubsup><mi mathvariant="script">S</mi><mi>i</mi><mi>m</mi></msubsup></mrow><annotation encoding="application/x-tex">\mathcal{S}^{m}_{i}</annotation></semantics></math>Sim|i=1,...,k\}_{m=1}^{|M|}$ . With a binary cross-entropy loss, we distill rule-based driving knowledge into the end-to-end planner:

For a trajectory $T_{i}$ , its distillation loss of each sub-score acts as a learned cost value in Eq. 4, measuring the violation of particular traffic rules associated with that metric.

4 Inference and Post-processing

Given the predicted imitation scores $\{\mathcal{S}^{im}_{i}|i=1,...,k\}$ and metric sub-scores $\{<math xmlns="http://www.w3.org/1998/Math/MathML" display="block"><semantics><mrow><msubsup><mi mathvariant="script">S</mi><mi>i</mi><mi>m</mi></msubsup></mrow><annotation encoding="application/x-tex">\mathcal{S}^{m}_{i}</annotation></semantics></math>Sim|i=1,...,k\}_{m=1}^{|M|}$ , we calculate an assembled cost measuring the likelihood of each trajectory being selected in the given scenario as follows:

where $\{w_{i}\}_{i=1}^{4}$ represent confidence weighting parameters to mitigate the imperfect fitting of different teachers. The optimal combination of weights is obtained via grid search, which typically fall within the following ranges: $0.01\leq w_{1}\leq 0.1,0.1\leq w_{2},w_{3}\leq 1,1\leq w_{4}\leq 10$ , indicating the necessity to prioritize rule-based costs over imitation. Finally, the trajectory with the lowest overall cost is chosen.

4.2 Model Ensembling

We present two model ensembling techniques: Mixture of Encoders and Sub-score Ensembling. The former technique uses a linear layer to combine features from different vision encoders, while the latter calculates a weighted sum of sub-scores from independent models for trajectory selection.

Experiments

Dataset. The Navsim dataset builds on the existing OpenScene dataset, a compact version of nuPlan with only relevant annotations and sensor data sampled at 2 Hz. The dataset primarily focuses on scenarios involving changes in intention, where the ego vehicle’s historical data cannot be extrapolated into a future plan. The dataset provides annotated 2D high-definition maps with semantic categories and 3D bounding boxes for objects. The dataset is split into two parts: Navtrain and Navtest, which respectively contain 1192 and 136 scenarios for training/validation and testing.

Metrics. For this challenge, we evaluate our models based on the PDM score, which can be formulated as follows:

where sub-metrics $NC$ , $DAC$ , $TTC$ , $C$ , $EP$ correspond to the No at-fault Collisions, Drivable Area Compliance, Time to Collision, Comfort, and Ego Progress. For the distillation process and subsequent results, $DDC$ is neglected due to an implementation problem.https://github.com/autonomousvision/navsim/issues/14.

2 Implementation Details

We train our models on the Navtrain split using 8 NVIDIA A100 GPUs, with a total batch size of 256 across 20 epochs. The learning rate and weight decay are set to $1\times 10^{-4}$ and 0.0 following the official baseline. LiDAR points from 4 frames are splatted onto the BEV plane to form a density BEV feature, which is encoded using ResNet34 . For images, the front-view image is concatenated with the center-cropped front-left-view and front-right-view images, yielding an input resolution of $256\times 1024$ by default. ResNet34 is also applied for feature extraction unless otherwise specified. No data or test-time augmentations are used.

3 Main Results

Our results, presented in Tab. 1, highlight the absolute advantage of Hydra-MDP over the baseline. In our exploration of different planning vocabularies , utilizing a larger vocabulary $\mathcal{V}_{8192}$ demonstrates improvements across different methods. Furthermore, non-differentiable post-processing yields fewer performance gains than our framework, while weighted confidence enhances the performance comprehensively. To ablate the effect of different learning targets, the continuous metric EP (Ego Progress) is not considered in early experiments and we attempt the distillation of the overall PDM score. Nonetheless, the irregular distribution of the PDM score incurs performance degradation, which suggests the necessity of our multi-target learning paradigm. In the final version of Hydra-MDP- $\mathcal{V}_{8192}$ -W-EP, the distillation of EP can improve the corresponding metric.

4 Scaling Up and Model Ensembling

Previous literature suggests larger backbones only lead to minor improvements in planning performance. Nevertheless, we further demonstrate the scalability of our model with larger backbones. Tab. 2 shows three best-performing versions of Hydra-MDP with ViT-L and V2-99 as the image backbone. For the final submission, we use the ensembled sub-scores of these three models for inference.