Wayve：从源头讲起，如何实现以对象为中心的自监督感知方法？（附代码）¶

1. 摘要¶

以对象中心的表示使自主驾驶算法能够推理大量独立智能体和场景特征之间的交互。传统上，这些表示是通过监督学习获得的，但会使感知与下游驾驶任务分离，可能会降低模型的泛化能力。在这项工作中，我们设计了一个以对象为中心的自监督视觉模型，仅使用RGB视频和车辆姿态作为输入来实现进行对象分割。我们在Waymo公开感知数据集上证明了我们的方法取得了令人满意的结果。我们发现我们的模型能够学习一种随时间推移融合多个相机姿势的表示，并在数据集中成功跟踪大量车辆和行人。我们介绍了该方法的起源和具体实现方法，并指明了未来的发展方向，为了帮助大家更好地复现代码，我们将详细地参数列入附表。

论文地址：https://arxiv.org/abs/2307.07147

模型代码：https://github.com/wayveai/SOCS。

2. 方法起源¶

人类和机器人有一种倾向，即认为对象的行为是单一连贯的，这是一种与生俱来的感知法则。对象在人类视觉中发挥着核心作用。我们根据特征将对象分组，用它们来描述我们周围的环境，同时，为我们不熟悉的对象寻找语义标签。当使用视觉表示进行下游任务时，如机器人技术等，对象中心模型是令人满意的：因为它们比端到端模型更容易被人类理解——这对验证安全性和赢得人类对视觉系统的信任非常重要。除此之外，以对象为中心的表示还提供了一套多样而强大的推理真实世界的工具，如物理理解模型、多智能体预测和规划模型以及因果推理模型。支持这类模型的表示可能对自动驾驶至关重要，因为在自动驾驶中，使用这些表示对大量相互作用的智能体和物理因素进行推理，可以获得车辆的最佳运动轨迹。

传统上，以对象中心的表示是通过训练监督对象检测模型，并从中提取对象属性(如位置和速度)来实现的。这种方法有两大缺点。首先，它需要与检测对象相匹配的带标签数据集，而大规模获取带标签数据集的成本很高，而且可能会引入不必要的偏差。另外，为了使系统能够处理新的对象或新的环境，必须收集新的带标签的数据。没有足够正确标注的数据会影响这类方法的泛化能力，但这正是以对象为中心的自监督模型的关键优势之一。

其次，根据有监督的视觉模型预测创建的对象表示，会使感知和决策组件之间脱节。例如，骑自行车的人应该被视为一个对象还是两个对象？两人乘双人自行车又如何？他们在向人行道上的行人挥手重要吗？如果他们正在打转向手势呢？这些问题的正确答案取决于如何使用信息做出决策。理想情况下，感知行为的结果应该反馈并改善感知本身，利用端到端学习，可以找到比手工设计更好的以对象为中心的表示。

这些考虑激发了以对象为中心的自监督感知模型的设计，这种模型将图像编码到一个将相关信息划分为多个“槽”的潜在空间中，然后这些槽中的信息被解码为自监督目标，如例如重建原始RGB输入，以及用于下游任务。已经有不少论文提出了各种方法，来鼓励模型在单个槽中编码关于不同对象的信息，例如通过使用使槽在像素上竞争注意力的编码器，或者通过不同的自动编码器损失来鼓励槽解码。然而，这些方法在复杂的现实世界数据中难以获得良好的结果。最近，SAVi++算法在Waymo真实世界驾驶视频公开数据集上显示了其分割能力。然而，这些结果需要额外的深度监督，而且需要初始对象与周围环境边界的槽，才能够获得最佳性能。

基于上述考虑，我们提出了一种仅使用RGB视频和相机运动信息进行自监督分割的模型，并在真实驾驶视频上获得了良好的的结果。相机运动信息在自动驾驶车辆中容易获得(例如来自同时定位与建图(SLAM)或轮组测距)，而且不需要专用激光雷达等传感器，因此是一种特别经济的方法。

3. 方法实现¶

我们的模型架构建立在SIMONe的视图监督变体上。该模型的目标是将场景分割为一组K个对象槽，对场景中每个对象的信息进行编码。通过以下步骤可以获得这些槽。首先，输入X(F帧图像序列，可选择从多相机视点)由标准卷积神经网络(CNN)并行处理，得到一组特征补丁。补丁与位置嵌入连接，位置嵌入在每个补丁源图像内的位置，以及与源图像相关的时间和视点变换矩阵进行编码。然后，它们作为仅解码transformer的输入信号。输出标记在整个图像维度上的平均值。在原始SIMONe模型中，这个维度对应于单个相机在几个时间点上拍摄的图像；但在这项工作中，我们使用了三个不同的姿势的相机，在时间和视点上进行取平均值。最后，通过MLP将每个信号解码为m维向量$z^k$和$\sigma^k$，这两个向量分别包含单个槽中潜预测平均值和方差。如果输入信号的数量与所需的槽的数量不同，则在transformer层的中途可选地进行跨信号特征的空间池化操作。具体模型如下图所示。

3.1. 训练和损失¶

为了鼓励模型在不同的槽中存储不同对象的信息，在训练期间我们应用三个互补的损失。 第一项是每个槽潜在向量与单位正态分布之间的KL散度之和，对所有槽求和:

$L_{KL}(X) = \sum_{k} D_{KL}(q_k(X)||p(\cdot))$

其中$X$是输入帧，$q_k(X) = N(z_k(X)， \sigma_k(X)I)$是槽 $k$的m维正态分布，均值和方差由模型预测，先验分布$p(\cdot)$是一个单位球面正态分布。直观地，这个损失鼓励模型避免使用多个槽表示同一对象，因为与使用单个槽表示对象并让其他槽保持接近单位正态分布相比，这样做会导致更大的惩罚。这个损失还鼓励潜在向量的每个维度之间解耦。

第二个损失基于模型执行对象级重建任务的能力。首先，根据每个槽 $k$的潜在分布获取对象的潜在向量$o_k-q_k(X)$。然后，每个$o_k$被独立解码到空间广播解码器进行逐像素预测。为了满足模型计算要求和内存需求，在训练期间每个输入序列中仅对N个随机像素进行解码。每个槽 $k$对每个像素的预测结果用RGB分布表示。$p(x^{(n)}|o_k)$和逻辑单元 $\hat{\alpha}_k$(经过槽归一化)表示槽 $k$表示像素的可能性。为了获得每个像素的最终预测，我们取每个槽预测的加权混合:

$p(x^{(n)}|o^1，...，o_k) = \frac{1}{K}\sum_{k}\bar{\alpha}_k[p(x^{(n)}|o_k)]$

其中针对像素$x^{(n)}$的每槽分布经softmax后的$\hat{\alpha}_k$值加权:

$\begin{align}\bar{\alpha}*k = \frac{exp(\hat{\alpha}_{k})}{\sum{k'}exp(\hat{\alpha}_{k'})}\end{align}$

$p(x^{(n)}|o_k)$的分布在第3.3节中进行了更详细的讨论。

最后，重建损失是每个像素在混合分布下的真实RGB值的对数概率:

$L_{recon}(X) = -\frac{1}{N}\sum_{n}logp(x^{(n)}|o^1，...，o_k)$

$o_k \sim q_k(X)$

由于每个槽的对象潜变量被独立解码，模型被迫在预测每个像素的RGB值时一次只使用单个槽中编码的信息。因此，直观地，这个损失鼓励模型将预测像素$n$的颜色所需的所有信息存储在单个槽中。

另外，学习到的槽表示还可以用于各种辅助任务。在本文中，受到对象与自动驾驶环境中相关的表示与用于预测良好驾驶动作的表示之间的密切联系的启发，我们实验使用预测车辆自身的未来路径作为辅助任务。在图像池化步骤之后，槽信息经过两个Transformer 解码器层，求平均值，并通过单层MLP解码成一系列预测偏移量$\hat{s}$，在自我参考帧的xy平面上。然后我们应用以下任务损失:

$L_{task}(s) = \sum_{t'}||s_{t'} - \hat{s}_{t'}||_1$

其中汇总是对每个未来时间点。我们使用以10Hz频率开始于最后图像帧之后0.1秒的自我参考帧中的16个未来路径。

最终训练损失与负ELBO损失类似，增加了辅助任务损失:

$L_{total}(X， s) = L_{recon}(X) + \omega_{task}L_{task}(s) + \beta L_{KL}(X)$

其中$\beta$和$\omega_{task}$超参数平衡不同的损失项。

3.2. 附加的模型输出¶

除了参数化加权混合像素分布，权重$\bar{\alpha}_k$还作为每一个槽 的α 掩码，使我们可以非常直观地看到每个槽关注场景的哪些像素。 为每个槽取最大值可以得到场景的预测分割。这种分割可以帮助模型调试和解释。例如，如果无法用掩码跟踪特定车辆，表明模型没有从场景的其他特征中区分出该对象的特征，因此没有独立表示其运动状态。

对象的槽或潜在向量也可以解码为图像重建或轨迹预测以外的其他输出。其他潜在下游任务可能包括视频预测、生产系统模型或有运动条件世界模型。在端到端机器学习中，哪些辅助任务可以协同提升性能，是一个激动人心的开放问题。

3.3. 对象槽解码分布¶

原始SIMONe模型采用正态分布预测的像素RGB值$p(x^{(n)}|o_k)$。(注意，在本节中，我们将RGB元组称为正态，但在现实中，R、G和B通道是独立对待的。)我们发现，在进行场景中分割时，这个分布会导致模型过度依赖颜色差异。这会导致一些失败的案例，如无法分割车身和挡风玻璃、无法识别出与背景颜色相似的对象等。我们猜测这是因为在两个不同颜色的区域边界附近，模型不确定为给定像素分配什么颜色。 为了用正态分布表示这种不确定性，模型被迫将不同颜色的区域分配给不同的槽，并使用每个槽的$\hat{\alpha}_k$权重给出每种颜色的可能性 。

对于我们的体系结构，我们使用多头正态分布替换SIMONe中的正态分布，来减轻这个问题。 定性的说，我们发现这能够更好地反映对象运动的分割。 对于每个像素和每个槽 k，解码器输出H个模式，其中预测平均RGB元组$\hat{\mu}^{(h)}_k$和逻辑单元$\alpha^{(h)}_k$ 决定每个模式的权重。 (注意，此外还有一个单独的“全局”$\hat{\alpha}_k$ 回归，它控制第k个槽对总混合分布的贡献，如方程2所示。) 因此，每个槽的分布是：

$p(x^{(n)}|o_k) \sim \frac{1}{H}\sum_{h} exp(\alpha^{(h)}*k)\frac{\sum*{h'}exp(\alpha^{(h')}_k)}{N(\hat{\mu}^{(h)}_k， \sigma_x)}$

其中正常分布的方差$\sigma_x$是一个超参数。当$H=1$时，这简化为SIMONe中的解码分布:

$p(x^{(n)}|o_k) = N(\hat{\mu}_k， \sigma_x)$

在我们的实验中，我们使用$H=3$和$\sigma_x=0.08$。

最后，模型的分割结果和轨迹预测如下图所示：

4. 结论¶

最近，以对象为中心的自我监督表示学习方法，在具有明确定义对象的人工数据集上表现出了很强的性能，但在具有复杂纹理和模糊对象的复杂真实世界数据上仍然举步维艰。在本文中，我们已展示的结果表明，通过使用相机姿态作为附加输入，有可能在RGB驾驶视频中获得合理的动态、以对象中心的表示。与3D深度传感器不同，姿态估计是自动驾驶汽车的一个普遍特征，因此我们认为我们的方法是在自动驾驶领域实现可扩展的、实用的，以对象中心的表示学习的一个很有前途的途径。此外，我们的研究结果表明，预测车辆自身的未来姿态是一项协同任务，它不会阻碍学习表达的质量。这对端到端驾驶模型来说尤其令人兴奋，因为它打开了驾驶性能和表示学习共同建立良性循环的可能性，同时保留了以对象为中心的表示的关键优势，如可解释性。

我们认为仍有可能进一步提高对象分割质量，例如，通过扩大模型规模和采取数据增强策略(这两点对SAVi++的性能非常重要)。我们还注意到，Waymo 公开感知数据集包含三个前向摄像头的480，000帧图像，对于理想的表示学习而言相比，该数据集的规模与其复杂性可能不够大。相比之下，最近的对象场景表示transformer模型是在1000万帧合成数据集上训练的。

最后，我们注意到，我们模型中的KL-发散损失鼓励学习理顺的对象潜在特征。更详细地研究这些特征是未来工作的一个令人兴奋的方向。

5. 论文中的超参数¶