UniAD 论文学习

一、解决了什么问题？

当前的自动驾驶方案大致由感知（检测、跟踪、建图）、预测（motion、occupancy）和规划三个模块构成。

为了实现各种功能，智驾方案大致包括两种路线。一种是针对每个任务都部署一个模型，该范式能降低各团队间的研发困难度，但由于各个优化目标是孤立的，会引发模块之间信息丢失、错误累加和特征不对齐的问题。另一种是多任务的设计路线，多个任务 heads 共享一个特征提取器，该范式能节省边缘计算平台的资源消耗，并且扩展性强，但会带来“负迁移”的问题。

端到端运动规划

自从 Pomerleau 提出使用一个网络直接预测控制信号，端到端运动规划受到越来越多的关注。后续研究通过闭环仿真、多模态输入、多任务学习、强化学习以及专家模型蒸馏的方式，取得了长足进展。但是，考虑到鲁棒性和安全性，这些方法直接从传感器数据输出控制信号、从合成场景迁移到实际应用仍有问题。因此，学者们试图显式地构造出一个网络的中间表征，预测场景是如何变化的。

二、提出了什么方法？

本文认为应该围绕着规划这一最终的目标来设计整体架构，于是提出了 UniAD。UniAD 在一个网络中包括了全栈的自动驾驶任务，它能充分利用各模块的优势，从全局的角度为各 agents 之间的交流提供互补的特征抽象，并且它以统一的 query 接口来连接所有的模块，推动各模块向规划这一目标对齐。

UniAD 遵循一切为规划服务的原则，核心构成就是 query-based 的接口设计，连接各个模块。与经典的边框表征相比，queries 得益于更大的感受野，能减轻上游任务预测所带来的复合错误的影响，而且能灵活地编码和建模 agents 之间的交流。

任务定义

检测和跟踪

检测和跟踪是自动驾驶领域两个重要的任务，在 3D 空间对它们做表征以支持下游任务。3D 检测负责定位每个时刻周围的物体（坐标、长宽高、偏航角等）；跟踪目的是找到不同时刻物体之间的对应关系，在时域内将它们串联起来。本文使用多目标跟踪来表示检测和跟踪的过程。最终输出是每帧内的一组 3D 框，它们对应的特征 $Q_A$ 会输入运动模块。此外，有一个特殊的自车 query 用于下游任务，在预测框和 ground-truth 框的匹配过程中不会包括自车 query。

在线建图

地图体现了环境的几何和语义信息。在线建图是利用车载传感器数据，分割出有价值的道路元素，以替代离线标注的高精地图。在 UniAD，在线地图建模了四种元素：车道线、可行驶区域、间隔物、人行横道。在 BEV 视角下对它们做分割。运动预测模块使用 map queries $Q_M$ 来建模 agent-map 关系。

运动预测

运动预测将感知和规划连接起来，在整个自动驾驶系统扮演重要角色，确保最终的安全。通常，运动预测是一个独立模块，利用高精地图和检测到的边框来预测 agent 未来的轨迹。现有的运动数据集，边框都是 ground-truth 标注，这在车载场景不现实。本文，运动预测模块将之前编码的稀疏 queries($Q_A$和$Q_M$)和密集 BEV 特征 $B$ 作为输入，预测各 agent 在未来 $T$ 时刻的 $\mathcal{K}$ 个可能轨迹，这些预测的轨迹是各 agent 的相对当前位置的偏移。Agent 特征编码了过去和未来的信息，会输入占用模块来理解未来的场景。

自车 query 预测未来时刻自车的运动，它实际上是比较粗糙的规划估计，planner 会利用该特征来生成最终的目标路径点。

占用预测

占用网格图是离散化的 BEV 表征，每个网格表示它是否被占用的置信度。占用预测任务用于发现网格图在未来 $T_o$ 时刻是如何变化的。运动预测依赖于稀疏的 agents，占用预测则是全场景密集表征的。为了研究场景和稀疏 agents 是如何变化的，占用模块的输入是 BEV 特征 $B$ 和 agent 特征 $G^t$。完成多步骤 agent-scene 交流后，对占用特征和密集场景特征做矩阵乘，得到实例级的概率图 ${\hat{O}}_A^t\in {\mathbb{R}}^{N_a\times H\times W}$。然后使用逐像素的 $\mathrm{argmax}$ 操作，将各时刻的概率图融合，得到保留了 agent ID 的全场景占用 ${\hat{O}}^t\in {\mathbb{R}}^{H\times W}$。

规划

规划模块是最终的目的，输入是上述模块的输出。传统的规划方法都是 rule-based，由各种 if-else 状态机组成，用检测和预测的结果来描述各类场景。而本文提出的 learning-based 模型则以上游的自车 query、密集的 BEV 特征 $B$ 作为输入，预测全部 $T_p$ 时刻的轨迹 $\hat{\tau }$。然后，用上游预测的未来占用 $\hat{O}$ 来优化预测轨迹 $\hat{\tau }$，以防发生碰撞，确保安全。

概览

UniAD 包括四个基于 transformer decoder 的感知与预测模块，以及一个 planner。Queries $Q$ 将整个流程串联起来，建模驾驶场景中各实例间的不同关系。下图展示了 UniAD 的流程，设计上非常精妙，遵循规划导向的思想。作者探索了感知和预测各模块的作用，充分发挥节点协同优化的优势。

• 首先，将多相机图像输入特征提取器，通过 BEVFormer 的 BEV 编码器将透视视角的特征变换为 BEV 特征 $B$。

• 然后在 TrackFormer 中，track queries 从 $B$ 中查询各 agents 的信息，进行检测和跟踪。

• 在 MapFormer 中，map queries 作为道路元素（车道线、间隔物）的语义抽象，并进行 maps 的全景分割。

• MotionFormer 获取各 agents 和 maps 之间的关系，预测每个 agent 未来的轨迹。在场景里面，各 agent 的行为会彼此影响，所以 MotionFormer 对所有的 agents 做联合预测。

• OccFormer 以 BEV 特征 $B$ 作为 queries，将 agent 的信息作为 keys 和 values，预测未来多步的占用网格图（保留 agents 的 IDs）。

• 最终，Planner 利用 MotionFormer 给出的自车 query 来预测出规划结果，并让自车远离 OccFormer 预测的被占用区域，以防碰撞发生。

1. 感知：跟踪和建图

1.1 TrackFormer

它协同完成检测和多目标跟踪任务，没有不可微的后处理操作。除了目标检测任务用到的检测 queries，它还包含了跟踪用的 track queries。在每一时刻，新初始化的检测 queries 负责检测首次被感知到的 agents，而 track queries 则对之前帧检测到的 agents 持续建模。检测 queries 和 track queies 都是通过关注 BEV 特征 $B$ 来获取各 agents 的信息。随着场景的变化，当前时刻的 track queries 通过一个自注意力模块与之前时刻的 queries 进行交流，聚合时域信息，直到相应的 agents 完全消失在画面中。TrackFormer 包括 $N$ 层，最终的输出 $Q_A$ 提供 $N_a$ 个有效 agents 的信息，供下游任务使用。除了编码自车周围 agents 的 queries，作者还增加了一个自车 query，显式地建模自车本身，规划模块会用到。

1.2 MapFormer

作者基于 2D 全景分割方法 Panoptic SegFormer 设计 MapFormer。将道路元素稀疏地表征为 map queries，编码了位置和结构信息，从而帮助下游的运动预测。本文将车道线、间隔物和人行横道设为 things，将可行驶区域设为 stuff。MapFormer 也有 $N$ 层，最后一层的 queries $Q_M$ 会输入 MotionFormer 做 agent-map 交流。

2. 预测：运动预测

有了 TrackFormer 和 MapFormer 分别提供的动态 agents 的 queries $Q_A$ 和静态图 $Q_M$，MotionFormer 就可以预测所有 agents 未来的多模态运动了，即 top-k 个可能的轨迹。同时，作者将 TrackFormer 中的自车 query 传入 MotionFormer，让自车和其它 agents 产生交流。输出的运动状态表示为 $\{{\hat{\mathbf{x}}}_{i,k}\in {\mathbb{R}}^{T\times 2}|i=1,...,N_a;k=1,...,\mathcal{K}\}$，其中 $i$ 表示 agent 的索引，$k$ 表示轨迹模态的索引，$T$ 是预测的长度。

2.1 MotionFormer

包括 $N$ 层，每层都获取到三种交互关系：agent-agent, agent-map, agent-goal point。对于每个 motion query $Q_{i,k}$，它和其它 agents $Q_A$ 或地图元素 $Q_M$ 的关系可以表示如下：

$$Q_{a/m}=\text{MHCA}(\text{MHSA}(Q),Q_A/Q_M)$$

上面的 $\text{MHCA}$ 和 $\text{MHSA}$ 表示多头跨注意力和多头自注意力。同时，我们也要关注目标点（goal point），优化预测轨迹，本文基于可变形注意力设计了一个 agent-goal point 注意力：

$$Q_g=\text{DeformAttn}(Q,{\hat{\mathbf{x}}}_T^{l-1},B)$$

其中 ${\hat{\mathbf{x}}}_T^{l-1}$ 是上一层预测轨迹的路径点。$\text{DeformAttn}(q,r,x)$ 是可变形注意力，输入为 query $q$、参考点 $r$ 和空间特征 $x$。它对参考点周围的空间特征应用稀疏注意力。这样，预测轨迹能进一步感知到路径点周围的环境。

这三种交互关系的建模是同时进行的，然后将生成的 $Q_a,Q_m,Q_g$ concat 到一起，输入一个 MLP，得到 query context $Q_{ctx}$。然后将 $Q_{ctx}$ 输入后续的层做优化，或者在最后一层就解码为预测结果。

2.2 Motion queries

MotionFormer 每一层的输入 queries 记作 motion queries，包括两个部分：前一层输出的 query context $Q_{ctx}$ 和 query position $Q_{pos}$。$Q_{pos}$ 整合了四重的位置信息：

• 场景级 anchor $I^s$ 的位置；
• agent 级 anchor $I^a$ 的位置；
• 第 $i$ 个 agent 的当前位置；
• 预测的 goal point。

$$Q_{pos}=\text{MLP}(\text{PE}(I^s))+\text{MLP}(\text{PE}(I^a))+\text{MLP}(\text{PE}({\hat{\mathbf{x}}}_0))+\text{MLP}(\text{PE}({\hat{\mathbf{x}}}_T^{l-1}))$$

这里的正弦位置编码 $PE(\cdot )$ 后跟着一个 MLP 用于编码位置点，第一层的 $I^s$ 设为 ${\hat{\mathbf{x}}}_T^0$。场景级 anchor 代表了全局视角下之前时刻的运动统计，agent 级 anchor 则在局部坐标捕捉可能的意图。它们都通过 k-means 算法对 ground-truth 轨迹的路径点做聚类，从而缩小预测结果的搜索空间。起始点提供每个 agent 的位置编码，而预测出的路径点则作为动态 anchor，逐层不断地优化。

2.3 非线性优化

直接从一个不准确的检测位置或偏航角回归 ground-truth 的路径点会产生不真实的轨迹预测，曲率和加速度可能会非常大。于是作者采用了一个非线性平滑方法，来调节目标的轨迹，使它们更加合理。该过程如下：

$${\tilde{\mathbf{x}}}^{\ast }=\underset{\mathbf{x}}{\mathrm{argmin}}c(\mathbf{x},\tilde{\mathbf{x}})$$

其中 $\tilde{\mathbf{x}}$ 和 ${\tilde{\mathbf{x}}}^{\ast }$ 表示 ground-truth 和平滑后的轨迹，$\mathbf{x}$ 通过 multiple-shooting 产生，代价函数为：

$$c(\mathbf{x},\tilde{\mathbf{x}})={\lambda }_{xy}{\Vert \mathbf{x},\tilde{\mathbf{x}}\Vert }_2+{\lambda }_{goal}{\Vert {\mathbf{x}}_T,{\tilde{\mathbf{x}}}_T\Vert }_2+\sum _{\varphi \in \Phi }\varphi (\mathbf{x})$$

其中，${\lambda }_{xy}$ 和 ${\lambda }_{goal}$ 是超参数，动力学函数集合 $\Phi$ 有五项，包括 jerk, curvature, curvature rate, acceleration, lateral acceleration。这个代价函数对目标轨迹起到正则的作用，使其遵守动力学约束条件。目标轨迹优化只在训练时进行，不影响推理。

3. 预测：占用预测

占用网格图是离散化的 BEV 表征，每个格子都有一个置信度，表示该格子是否被占用。占用预测任务用于预测网格图在未来是如何变化的。OccFormer 从两个方面融合了场景级和 agent 级的语义信息：

• 通过一个精心设计的注意力模块，从密集场景特征学到 agent 级的特征；
• 对 agent 级特征和密集场景特征做矩阵乘法，输出实例占用。

OccFormer 由 $T_o$ 个序列模块组成，$T_o$ 表示预测的长度。由于占用网格图过于密集，这里的 $T_o$ 通常要小于运动任务中的预测长度 $T$。每个模块的输入包括丰富的 agent 特征 $G^t$ 和前一层的状态（密集特征）$F^{t-1}$，然后输出 $t$ 时刻的状态 $F^t$。为了得到 agent 特征 $G^t$，我们在模态维度对 MotionFormer 的 motion queries 做最大池化，记作 $Q_X\in {\mathbb{R}}^{N_a\times D}$，$D$ 是特征维度。然后通过一个时域 MLP 将它与上游的 track query $Q_A$ 及当前位置编码 $P_A$ 融合：

$$G^t=\text{MLP}([Q_A,P_A,Q_X]),t=1,...,T_o$$

$[\cdot ]$ 表示 concat 操作。对于场景级信息，出于计算效率考虑，BEV 特征 $B$ 会缩小到 $1/4$ 分辨率，作为第一个模块的输入 $F^0$。为了进一步节约训练时内存占用，每个模块都遵循下采样-上采样的方式，在中间有一个注意力模块，在 $1/8$ 大小的特征（记作 $F_{ds}^t$）上进行 pixel-agent 交流。

3.1 Pixel-agent 交流

在预测未来占用网格图时，Pixel-agent 交流用于统一对场景和 agents 的理解。将密集特征 $F_{ds}^t$ 作为 queries，实例级特征作为 keys 和 values 不断更新密集特征。$F_{ds}^t$ 输入一个自注意力层，建模网格间的响应，然后用一个跨注意力层建模 agent 特征 $G^t$ 和各网格的特征之间的关系。为了对齐 pixel-agent 的对应关系，作者用一个注意力 mask 来约束跨注意力，每个像素只关注于 $t$ 时刻占据它的 agent。密集特征的更新过程如下：

$$D_{ds}^t=\text{MHCA}(\text{MHSA}(F_{ds}^t),G^t,\text{attn\_mask}=O_m^t)$$

注意力 mask $O_m^t$ 语义上类似于占用网格图，用一个额外的 agent 级特征和密集特征 $F_{ds}^t$ 相乘得到，我们将这个 agent 级特征叫做 mask 特征 $M^t=\text{MLP}(G^t)$。经过上述交流过程，$D_{ds}^t$ 就上采样到了 $B$ 的 $1/4$ 大小。我们将 $D_{ds}^t$ 通过残差连接加到模块输入 $F^{t-1}$ 上，得到的结果 $F^t$ 再输入进下一模块。

3.2 实例级占用

它表示的是保留了每个 agent ID 的占用网格图。它可以通过简单的矩阵乘法提取。为了得到 BEV 特征 B （原始大小为 $H\times W$ ）的预测占用，场景级特征 $F^t$ 通过一个卷积解码器上采样为 $F_{dec}^t\in {\mathbb{R}}^{C\times H\times W}$，其中 $C$ 是通道维度。

对于 agent 级特征，我们通过另一个 MLP 进一步将粗糙的 mask 特征 $M^t$ 更新为占用特征 $U^t\in {\mathbb{R}}^{N_a\times C}$。实验表明，$U^t$ 要比 $G^t$ 带来更优的表现。最终 $t$ 时刻的实例级占用表示为：

$${\hat{O}}_A^t=U^t\cdot F_{dec}^t$$

4. 规划

不带高精地图的规划一般需要高层级的指令来表示往哪个方向走。作者将原始的导航信号（左转、右转、保持前进）转换为三个可学习的 embeddings，叫做 command embeddings。由于 MotionFormer 的自车 query 已经表达了多模态意图，作者用 command embeddings 补充它，得到 plan query。然后将 plan query 关注到 BEV 特征 $B$，使它感知周围环境，然后将其解码，得到未来的路径点 $\hat{\tau }$。

为了避免碰撞，只在推理时基于牛顿法来优化 $\hat{\tau }$：

$${\tau }^{\ast }=\underset{\tau }{\mathrm{argmin}}f(\tau ,\hat{\tau },\hat{O})$$

其中，$\hat{\tau }$ 是原始的规划预测，${\tau }^{\ast }$ 表示优化后的规划，最小化代价函数 $f(\cdot )$ 得到。$\hat{O}$ 是经典的二值占用网格图，从 OccFormer 的实例占用预测融合得到。代价函数如下：

$$f(\tau ,\hat{\tau },\hat{O})={\lambda }_{coord}{\Vert \tau ,\hat{\tau }\Vert }_2+{\lambda }_{obs}\sum _t\mathcal{D}({\tau }_t,{\hat{O}}^t)$$

$$\mathcal{D}({\tau }_t,{\hat{O}}^t)=\sum _{(x,y)\in \mathcal{S}}\frac{1}{\sigma \sqrt{2\pi }}\exp (-\frac{{\Vert {\tau }_t-(x,y)\Vert }_2^2}{2{\sigma }^2})$$

这里，${\lambda }_{coord},{\lambda }_{obs}$ 和 $\sigma$ 是超参数，$t$ 是未来时刻的索引。考虑到周围的位置受到 $\mathcal{S}=\left\{(x,y)|{\Vert (x,y)-{\tau }_t\Vert }_2<d,{\hat{O}}_{x,y}^t=1\right\}$ 的限制，$l_2$ 代价函数将轨迹拉向原来预测的位置，而碰撞项 $\mathcal{D}$ 则将其推离开被占用的网格。

5. 学习

UniAD 训练包括两个阶段。首先协同训练感知部分，即跟踪和建图模块，训练 $6$ 个 epochs。然后端到端训练感知、预测和规划模块共 $20$ 个 epochs。

5.1 共享匹配

UniAD 包括实例建模，所以在感知和预测任务上，需要将预测结果和 ground-truths 配对。与 DETR 相似，在跟踪和在线建图阶段它使用了二分匹配算法。至于跟踪，检测 queries 的候选框会和新出现的 ground-truth 物体做配对，track queries 的预测则会继承之前帧的配对结果。跟踪模块的匹配结果会在运动和占用节点复用，从而持续地对历史跟踪的 agents 和未来的运动之间做建模。

6. 实现细节

6.1 检测和跟踪

继承了 BEVFormer 的大多数检测设计，通过一个 BEV 编码器将图像特征变换为 BEV 特征 $B$，再使用一个可变形 DETR 头对 $B$ 做检测。为了避免繁琐的匹配后处理，作者引入了一组 track queries，持续地跟踪之前检出的实例。跟踪过程细节如下：

6.1.1 训练阶段

训练开始时，所有的 queries 默认为检测 queries，预测新出现的目标，这和 BEVFormer 一样。通过匈牙利算法将检测 queries 和 ground-truths 匹配起来。将它们保存起来，在下一时刻通过 query interaction 模块（QIM）更新为 track queries。在下一时刻，track queries 会根据对应的 track ID 直接匹配到部分的 ground-truth 目标，而检测 queries 会匹配到其余的 ground-truths（新出现的目标）。为了让训练稳定，采用 3D IOU 来过滤匹配到的 queries。只保存和更新那些与 ground-truth 框的 3D IOU 大于一定阈值的预测框。

6.1.2 推理阶段

推理时，序列帧按顺序送入网络，track queries 存在的时间可能要长于训练时的。另一区别就是 query 更新，推理时使用分类得分来过滤 queries，而非 3D IOU（因为没有 ground-truths 了）。此外，为了避免由遮挡引发的短时间轨迹中断的情况，在推理阶段使用了生命周期机制。对于每个 track query，若它的分类得分低连续在 $2s$ 内于 $0.35$，它就被认为完全消失，则被移除。

6.2 在线建图

Map queries 被分为 thing queries 和 stuff queries。Thing queries 建模实例级的地图元素（即车道线、边界、人行横道），通过二分匹配来关联到 ground-truths，stuff queries 只负责语义元素（即可行驶区域），通过固定类别分配来处理。Thing queries 个数为 $300$，stuff query 个数为 $1$。堆叠了 $6$ 个位置解码层和 $4$ 个 mask 解码层。选取位置解码器后的 thing queries 作为 map queries $Q_M$ 供下游任务用。

6.3 运动预测

MotionFormer 用 $I_T^a,I_T^s,{\hat{x}}_0,{\hat{x}}_T^{l-1}\in {\mathbb{R}}^{\mathcal{K}\times 2}$ 来编码 query 位置，用 $Q_{ctx}^{l-1}$ 作为 query context。通过 k-means 算法对训练数据中所有的 agents 做聚类，得到 anchors，$\mathcal{K}=6$ 与输出模态个数一样。为了编码场景先验，根据各 agent 的当前位置和偏航角，将 anchor $I_T^a$ 旋转和平移到世界坐标系下，记作 $I_T^s$：

$$I_{i,T}^s=R_i{I}_T^a+T_i$$

其中 $i$ 是 agent 的索引。作者也使用了前一层预测的 goal point ${\hat{x}}_T^{l-1}$，使得更加准确。同时，将 agent 当前的位置广播到其它模态，记作 ${\hat{x}}_0$。然后，对每个先验位置信息应用 MLP 和正弦位置编码，记作 query position $Q_{pos}\in {\mathbb{R}}^{\mathcal{K}\times \mathcal{D}}$，形状与 $Q_{ctx}$ 一样。$Q_{ctx}$ 和 $Q_{pos}$ 一起构建了 motion query。在 MotionFormer 中，$\mathcal{D}=256$。

MotionFormer 有三个 transformer 模块，agent-agent, agent-map, agent-goal point 关系模块。Agent-agent 和 agent-map 模块用标准的 transformer 解码层构建，包括一个多头自注意力层、一个多头跨注意力层和一个前馈网络，内部还有多个归一化层和残差连接。作者也在 $Q_A$ 和 $Q_M$ 中加入了正弦位置编码，然后跟着 MLPs 层。Agent-goal 模块用可变形跨注意力层构建，将之前预测轨迹的 goal point （$R_i{\hat{x}}_{i,T}^{l-1}+T_i$）作为参考点使用，如下图所示。每条轨迹的采样点点个数为 $4$，每个 agent 有 $6$ 条轨迹。将每个关系模块的输出特征 concat 到一起，用 MLP 层映射成维度 $\mathcal{D}=256$。然后，使用高斯混合模型构建每个 agent 的轨迹，其中 ${\hat{x}}_l\in {\mathcal{R}}^{\mathcal{K}\times \mathcal{T}\times 5}$。预测时长 $T=12$，约 $6$ 秒。最终输出轨迹是最后一个维度的前两个值，即 $x,y$。此外，也要预测每个模态的得分，$score({\hat{x}}_l)\in {\mathcal{R}}^{\mathcal{K}}$。将该模块堆叠 $N=3$ 次。

6.4 占用预测

给定 BEV 特征，首先用卷积层将其下采样 （$/4$），然后输入 OccFormer。OccFormer 由 $T_o$ 个序列模块组成，如下图所示。$T_o=5$ 是时间长度（包括当前和未来帧），每个模块负责生成一帧的占用。该方法融合了密集场景特征和稀疏的 agent 特征。密集场景特征来自于最后一个模块的输出，用卷积层进一步下采样（$/8$），降低 pixel-agent 的计算量。将 track query $Q_A$, agent positions $P_A$ 和 motion query $Q_X$ concat 到一起，输入一个时域 MLP。计算像素级的自注意力，对剧烈变化的场景所需的长期依赖关系做建模；然后将每个像素点关注到对应的 agent，做 scene-agent 融合。为了增强 agents 和像素之间的位置对齐，用一个注意力 mask 来约束跨注意力，该注意力 mask 通过计算 mask 特征和下采样后的场景特征的矩阵乘得到，用一个 MLP 来编码 agent 特征以得到 mask 特征。然后，将密集特征上采样到与输入 $F^{t-1}$ 相同的分辨率（$/4$），用残差连接将它和 $F^{t-1}$ 相加。得到的特征 $F^t$ 输入下一模块和卷积解码器，以预测占用。复用 mask 特征，输入另一个 MLP，得到占用特征。对占用特征和解码的密集特征 $F_{dec}^t$ 做矩阵乘，得到实例级的占用。注意，在所有的 $T_o$ 模块中，共享 mask 特征的 MLP 层、占用特征的 MLP 层和卷积解码器，其它的组件则是独立的。在 OccFormer 中，所有的密集特征和 agent 特征的维度都是 $256$。

6.5 规划

如下图，planner 的输入是跟踪和运动模块预测的自车 query，用蓝色三角和黄色矩形表示。用 MLP 层编码这俩 queries 和 command embedding，然后是一个对模态维度做的 max-pool 操作，选取并聚合模态最显著的特征。BEV 特征模块用标准的 transformer 解码层构建，堆叠 $N=3$ 层。它用聚合后的 plan query 对密集的 BEV特征做跨注意力。为了嵌入位置信息，作者在 plan query 中加入了学到的位置编码，在 BEV 特征中加入了正弦位置编码。然后用 MLP 层回归轨迹，记作 $\hat{\tau }\in {\mathcal{R}}^{T_p\times 2}$。这里，$T_p=6$（$3$ 秒）。此外，根据预测占用 $\hat{O}$ 和轨迹 $\hat{\tau }$，作者设计了碰撞机制。

$Q_A^{ego}$ 和 $Q_{ctx}^{ego}$ 是跟踪预测模块和运动预测模块。用 MLP 层对它们和 command 做编码，然后是一个 max-pool 层，选取并聚合最显著的模态特征。

6.6 训练细节

协同学习

UniAD 训练分两个阶段。第一阶段，预训练感知任务，包括跟踪和在线建图，来稳定感知的预测。为了快速收敛，加载了 BEVFormer 的主干、FPN、BEV 编码器和检测解码器的权重，除了 object query embeddings。取消主干网络的梯度回传，以降低内存占用，训练 UniAD 共 $6$ 个 epochs，损失如下：

$$L_1=L_{track}+L_{map}$$

第二阶段，冻结图像主干和 BEV 编码器（负责图像视角变换到 BEV 视角），进一步降低内存消耗。这时，UniAD 训练包括了所有的损失，有跟踪、建图、运动预测、占用预测和规划，训练 $20$ 个 epochs。

$$L_2=L_{track}+L_{map}+L_{motion}+L_{occ}+L_{plan}$$

检测和跟踪损失

对于每对结果，采用匈牙利损失，即类别标签的 Focal loss 和 3D 边框定位的 $l_1$ 损失的线性组合。匹配策略是：用二分匹配将新出现的 queries 与 ground-truth 目标配对，track queries 的预测结果则继承其上一帧的 ground-truth 索引。

$$L_{track}={\lambda }_{focal}{L}_{focal}+{\lambda }_{l_1}{L}_{l_1}$$

其中，${\lambda }_{focal}=2,{\lambda }_{l_1}=0.25$。

在线建图损失

包括一个 thing 损失（车道线、间隔物）和一个 stuff 损失（可行驶区域）。分类损失用的 focal loss，thing 边框损失用的 $l_1$ 损失，分割用的 Dice loss 和 GIOU loss。

$$L_{map}={\lambda }_{focal}{L}_{focal}+{\lambda }_{l_1}{L}_{l_1}+{\lambda }_{giou}{L}_{giou}+{\lambda }_{dice}{L}_{dice}$$

其中，${\lambda }_{focal}={\lambda }_{giou}={\lambda }_{dice}=2,{\lambda }_{l_1}=0.25$.

运动预测损失

用高斯混合建模多模态轨迹，使用 multi-path loss，包括一个分类得分损失 $L_{cls}$ 和一个 negative log-likelihood 损失 $L_{nll}$。为了确保轨迹是连续平滑的，首先预测每一时刻各 agent 的速度，然后随着时间不断累积它们，得到最终的轨迹。

$$L_{motion}={\lambda }_{cls}{L}_{cls}+{\lambda }_{reg}{L}_{nll}$$

其中，${\lambda }_{cls}={\lambda }_{reg}=0.5$.

占用预测损失

输出的实例级的占用预测是每个 agent 的二值分割，因此采用了二值交叉熵和 Dice loss 作为占用损失。

$$L_{occ}={\lambda }_{bce}{L}_{bce}+{\lambda }_{dice}{L}_{dice}$$

其中，${\lambda }_{dice}=1,{\lambda }_{bce}=5$.

规划损失

安全是规划的最重要因素，因此，除了朴素的 imitation $l_2$ 损失，还用了一个碰撞损失，让预测轨迹远离障碍物：

$$L_{col}(\hat{\tau },\delta )=\sum _{i,t}\text{IOU}(box(\hat{\tau },w+\delta ,l+\delta ),b_{i,t})$$

$$L_{plan}={\lambda }_{imi}|\hat{\tau },\tilde{\tau }{|}_2+{\lambda }_{col}\sum _{(\omega ,\delta )}\omega L_{col}(\hat{\tau },\delta )$$

其中，${\lambda }_{imi}=1,{\lambda }_{col}=2.5$，$(\omega ,\delta )$是额外安全距离的权重值，$box(\cdot ,\cdot ,\cdot )$表示自车框在 $t$ 时刻变大，以保持较大的安全距离，$b_{i,t}$ 表示场景中预测的每个 agent。实际使用时，$(\omega ,\delta )$设为$(1,0),(0.4,0.5),(0.1,1)$。