直接暴涨十个点！地平线提出DOME：扩散和Occ双管齐下，打造超真实世界模型！...-优快云博客

点击下方卡片，关注“自动驾驶之心”公众号

今天自动驾驶之心为大家分享地平线和中科院软件研究所最新的世界模型工作—DOME！扩散和OCC双管齐下，打造世界模型。如果您有相关工作需要分享，请在文末联系我们！

自动驾驶课程学习与技术交流群事宜，也欢迎添加小助理微信AIDriver004做进一步咨询

>>点击进入→自动驾驶之心『世界模型』技术交流群

论文作者 | Songen Gu等

编辑 | 自动驾驶之心

写在前面 && 笔者理解

在任何一组东西中，最重要的只占其中一小部分，约20%，其余80%尽管是多数，却是次要的，这就是著名的“二八定律”。而自动驾驶发展至今，解决剩余20%的“重要的”长尾问题和极端案例一直是一件比较棘手的事情。即使是最先进的SOTA和最广泛的数据收集工作也难以解决。而解决这些挑战的一个有希望的方法在于世界模型。世界模型结合了历史上下文和其它智能体的行动来预测环境观察的未来演变。这允许自动驾驶模型更深入地预测未来，改进行动可行性的评估。

世界模型可以分为几种类型：包括基于2D视频的模型和基于3D表示的模型，比如利用LiDAR和占用框架的模型。前者基于视频的世界模型在维持跨视图和跨时间一致性方面会面临一些挑战，从而阻碍了它们在现实世界场景中的应用。而基于占用的世界模型避免了这个问题。这些模型以历史占用序列作为输入，并预测未来的占用观察，用原始的3D表示确保了内在的3D一致性。另外，占用标注相对容易获得，因为它们可以有效地从稀疏的LiDAR标注中学习，或者可能通过时间帧的自监督学习获得。基于占用的模型也是模态不确定的，表明它们可以从单目或环视相机生成，或者从LiDAR传感器生成。

现有的基于占用的世界模型可以分为两种类型：基于自回归的和基于扩散的：

基于自回归的方法以自回归的方式使用离散标记预测未来的占用。然而，由于这些方法依赖于离散标记器，量化过程导致信息丢失，限制了预测高保真占用的能力。此外，自回归方法难以生成真实的长时间占用序列。
基于扩散的方法将空间和时间信息展平为一维的标记序列，而不是单独处理它们，导致在有效捕获空间时间信息方面存在困难。因此，将历史占用信息整合到模型中变得困难，因为空间和时间数据被结合在一起。这种限制意味着模型可以生成输出，但不能预测，限制了其在现实世界场景中的适用性。此外，作者发现大多数占用世界模型对细粒度控制的探索不足，导致过度拟合特定场景，限制了它们对下游任务的适用性。

为了解决上述问题，作者提出了一种预测未来占用帧的新方法，称为DOME。具体来说，作者的方法包括两个组件：Occ-VAE和DOME。为了克服离散标记的限制，作者的Occ-VAE使用连续的潜在空间来压缩占用数据。这允许有效的压缩，同时保留高保真细节。作者世界模型展示了两个关键特征：

高保真度和长时生成。作者采用时空扩散变换器来预测未来的占用帧。通过利用上下文占用条件，作者将历史占用信息作为输入。时空架构有效地捕获了空间和时间信息，实现了精细细节，并能够生成长时预测（32秒）。
细粒度可控性。作者通过引入轨迹重采样方法来解决预测中的精确控制挑战，这显著提高了模型生成更精确和多样化占用预测的能力。

模型框架

作者介绍了DOME，一个基于扩散的占用世界模型。作者的方法由两个主要组件组成：Occ-VAE和DOME。为了使世界模型与轨迹条件对齐，作者提出了轨迹编码器和轨迹重采样技术，专门设计用于增强模型的可控性。

Occ-VAE

Occ-VAE是作者模型的核心组件，它使用变分自编码器（VAE）将占用数据压缩到潜在空间，这对于提高表示的紧凑性和世界模型预测的效率至关重要。注意到离散标记器通常无法保留占用帧的精细细节，作者提出将密集的占用数据编码到连续的潜在空间中，以更好地保留复杂的空间信息。如图2所示的提出的架构，细节如下：

占用数据：由于Occ-VAE专门设计用于占用数据，作者首先讨论这种3D场景表示。3D占用数据将自我车辆周围的环境体素化为一个体素网格，每个网格单元根据其所包含的物体分配语义标签。

编码器：受基于图像的VAE方法的启发，作者提出了一个特别为占用数据设计的连续VAE。为了处理由离散语义ID组成的3D占用数据，作者首先将其转换为鸟瞰图（BEV）风格的张量，通过索引一个可学习的类别嵌入。这个过程将占用数据展平为一致的特征维度。随后，一个编码器网络将转换后的数据编码为压缩表示。这个表示然后沿通道维度分割为和，其中和表示编码数据的空间维度，表示通道维度。编码后，连续潜在变量使用重参数化技巧进行采样，遵循VAEs的方法：，其中是从标准正态分布中采样的噪声向量，表示逐元素乘法。

编码器包括2D卷积层和注意力块。类别嵌入是随机初始化的，并与Occ-VAE一起训练。

解码器：解码器网络负责从采样的潜在变量重建输入占用。它使用3D反卷积层对潜在表示进行上采样，确保改善时间一致性。上采样的特征然后重塑为。通过与类别嵌入的点积计算逻辑分数，逻辑分数的arg max确定最终的类别预测。

训练损失：在Occ-VAE的训练中，作者的损失函数由两个组件组成：重建损失和KL散度损失，遵循标准的VAE框架。作者使用交叉熵损失作为重建损失。此外，为了解决预测中的类别不平衡问题，作者额外加入了Lovasz-softmax损失，这有助于缓解不平衡问题。总损失定义如下：

其中和分别是Lovasz-softmax损失和KL散度损失的损失权重。训练完成后，Occ-VAE模型被冻结，其编码器作为特征提取器用于DOME训练，其解码器从DOME重建潜在表示以生成占用数据。

DOME：基于扩散的占用世界模型

占用世界模型基于智能体的历史数据（）预测未来的占用观察，其中表示占用观察，表示智能体的行动。为了实现这一点，作者采用了一个具有时间感知层的潜在扩散模型，使模型能够有效地从时间变化中学习。使用时间掩码整合历史占用观察，鼓励模型基于条件帧预测未来帧。此外，为了给世界模型提供增强的运动先验和可控性，作者的轨迹编码器整合了自我车辆的行动，允许精确的下一帧预测，由给定的摄像机姿态控制。具体来说，作者的模型以编码的潜在和自我车辆的轨迹作为输入，其中表示对应于4D占用数据中帧数的时间维度。潜在部分被掩码遮盖，只允许帧（）可见，模型被训练以预测剩余的掩码帧。

时空扩散变换器：为了预测具有时间感知的未来占用，作者采用了一个受基于视频方法启发的时空潜在扩散变换器。作者首先将潜在表示分割为帧的序列标记，每个序列包含个标记，其中表示patch大小。然后向空间和时间维度添加位置嵌入。如图2所示，作者的模型由两种基本类型的块组成：空间块和时间块。空间块捕获共享相同时间索引的帧之间的空间信息，而时间块在固定的空间索引处沿时间轴提取时间信息。这些块以交错的方式排列，有效地捕获空间和时间依赖性。

历史占用条件：为了使模型能够预测未来的占用特征，必须根据历史占用数据对生成进行条件化。这是通过条件掩码实现的。给定占用数据的多帧上下文和表示上下文帧数的超参数，从历史占用观察中编码潜在。然后作者构建一个条件掩码，确保模型根据可用的上下文帧进行预测。在训练期间，噪声标记根据条件掩码部分替换为上下文潜在，对于任何使用上下文帧的训练迭代：。

为了使模型能够在没有条件的情况下生成，作者应用了dropout机制，其中的比例时间内，模型在没有上下文帧的情况下进行训练。

损失函数：作者将普通的扩散损失扩展为时空版本，使其与上下文占用条件兼容。由于作者预测一系列特征占用，总体损失是跨所有帧计算的。在上下文占用条件下，噪声潜在被替换为真实值（如上所述），因此，这些帧的损失使用条件掩码忽略。训练扩散模型的损失函数定义如下：

其中是第个扩散时间戳的第帧，是去噪网络，特别是作者的DOME模型。

轨迹作为条件

轨迹条件注入：对于世界模型来说，动作条件注入是必不可少的，因为世界观察应该根据智能体在的最后动作合理且连贯地变化。作者将轨迹信息注入模型以进行条件生成。具体来说，给定自我车辆的姿态，作者首先计算相对平移和相对旋转。从，作者提取，从，作者获得偏航角，代表自我车辆的朝向。然后作者对应用位置编码，使用线性层将编码值投影到隐藏大小，并与时间嵌入结合。这些组合值随后被传递到自适应层归一化（adaLN）块。

轨迹重采样：这个问题源于训练数据集中的不平衡和有限多样性。例如，在nuScenes数据集中，训练集包含700个场景，但大多数涉及车辆直行（大约87%，见图4（c）），突出了不平衡问题。此外，在同一场景中，车辆只通过一次，导致在相同场景下不同轨迹条件下缺乏多样化的3D占用样本。这导致模型过度拟合场景，仅根据上下文观察学习基于真实特征的观察。原始轨迹分布如图4（a）所示。

为了解决这个问题，作者提出了一种轨迹重采样方法，如图3（a）所示。作者的目标是多样化自我车辆的动作和每个场景中采样的占用。该过程包括以下步骤：（1）多帧点云聚合：作者首先将自我框架中的占用序列转换为3D点云，然后使用自我姿态将其转换为世界框架。通过选择点云的语义标签来过滤潜在的动态对象（例如，汽车，行人）。（2）获取可行驶区域：为了生成多样化的观察，作者根据场景的可行驶区域创建各种可行的轨迹。在聚合所有点云到世界框架后，作者过滤道路类别，并从俯视图体素化道路点云以产生鸟瞰图（BEV）地图（见图3（b））。（3）生成多样化和可行的轨迹：使用BEV地图，作者随机采样两个点代表起始和目标位置。作者应用平滑的A*算法生成连接这些点的轨迹，模拟自我车辆的驾驶轨迹。得到的轨迹被转换为姿态，z坐标设置为0。（4）提取重采样占用：使用轨迹姿态，作者应用类似于Tian et al.（2023）的占用真实值提取方法，从点云中重新采样占用。

作者的重采样轨迹分布如图4（b）所示。与图4（a）相比，它填补了轨迹分布的空白，表明作者的方法增强了多样性并减轻了不平衡。图4（c）所示的驾驶方向直方图进一步支持了这种改进。

总之，作者的轨迹重采样方法既简单又有效。据作者所知，作者是第一个探索世界模型预测的占用数据增强。这种方法具有很高的通用性，可以应用于所有类型的占用数据，包括机器标注的、LiDAR收集的或自监督的数据。它只需要姿态和占用数据，而不需要LiDAR数据或3D边界框。

世界模型的应用

4D占用预测：在推理过程中，作者从对应于帧缓冲区大小的随机噪声开始（要预测的帧数），并通过Occ-VAE编码上下文占用帧以获得上下文潜在。作者将随机噪声中的帧替换为这些上下文潜在，然后将输入传递给作者的时空DiT（见图2底部）。在去噪循环过程中，上下文潜在保持不变，因为它们在每次迭代中都被重新引入。获得去噪潜在后，作者将其传递给Occ-VAE的解码器以生成最终的占用预测。超参数可以根据不同的要求进行调整。作者设置 = 4以进行精确的占用预测，因为更长的历史帧提供了更多的场景和运动信息。当需要更大的可控性时，如轨迹信号所要求的，作者设置 = 1以减少占用运动信息的影响，同时保持可控的起始观察。

长期生成的滚动：由于计算资源和内存限制，作者的模型在训练和推理过程中仅处理帧占用数据。为了生成更长期的占用预测，作者实施了类似于自回归方法的滚动策略。具体来说，在生成前帧后，作者重用最后一个预测帧作为上下文帧以预测接下来的帧。一个偏移量将相应的轨迹切片与上下文帧对齐。这种策略可以迭代应用以实现长期占用预测。

实验结果

实验设置

作者在nuScenes数据集上进行实验，使用IoU（交并比）和mIoU（平均交并比）指标来评估占用重建和4D占用预测。更高的IoU和mIoU值表明在压缩过程中信息丢失较少，反映了更好的重建性能，并展示了对未来周围环境更准确的理解。

占用重建

尽可能压缩的同时精确重建占用对于下游任务如预测和生成至关重要。在这里，作者将Occ-VAE与使用占用标记器的现有方法进行比较，并评估它们的重建精度。占用重建的定量结果如表1所示。作者在IoU和mIoU指标上都实现了最先进的重建性能，分别为83.1%的mIoU和77.3%的IoU。

4D占用预测

作者在各种设置下将其方法与现有的4D占用预测方法进行比较。这些设置包括使用地面实况3D占用数据（-O）作为输入，以及使用现成的3D占用预测器（-F）的预测结果作为输入。按照Wei et al.（2024）的实验设置，作者使用FB-OCC作为占用提取器，利用来自相机输入的预测。

定性结果如图5所示。定量结果如表2所示，表明作者的DOME-O实现了最先进的性能，mIoU为27.10%，IoU为36.36%。作者观察到与现有方法相比，在短期（1秒）和长期（3秒）预测方面都有显著改进，表明作者的模型有效地捕获了场景随时间的基本演变。DOME-F可以被认为是一种端到端的基于视觉的4D占用预测方法，因为它仅使用周围的相机捕获作为输入。尽管任务具有挑战性，作者的方法仍然实现了有竞争力的性能，进一步证明了DOME具有很强的泛化能力。

作者还展示了模型的长期生成能力，如图6所示，以及在给定相同起始帧的情况下，通过轨迹条件进行操作的能力，如图7所示。此外，作者还在表4中将其方法的生成能力与现有的占用世界模型进行了比较，作者的方法能够生成的持续时间是OccWorld的两倍，是OccSora的两倍。

消融研究

不同的轨迹条件：作者测试了轨迹条件的不同设置，结果如表3所示。Traj.表示是否使用姿态条件进行预测，Res.表示是否使用作者的轨迹重采样增强，Yaw表示是否添加偏航角嵌入。即使不使用任何姿态条件，作者发现其模型优于OccWorld。轨迹信息通过为模型提供场景变化的明确方向而不是要求它从多种可能性中推断出来，显著改善了预测。偏航角嵌入在IoU方面提供了轻微的改进。

上下文帧的数量：作者发现在预测过程中提供更多的上下文帧可以带来更好的预测（见表5），因为额外的帧为模型提供了关于其他车辆和场景的运动和变化的更明确信息。然而，作者也观察到增加帧的数量不如使用轨迹信息高效，因为模型必须导航模糊的帧历史来预测未来的运动。对于基于智能体确定的运动进行预测的世界模型来说，这种歧义是不必要的。

结论和展望

作者提出了一个基于扩散的世界模型DOME，它根据历史数据预测未来的占用帧。它整合了带有轨迹编码器和重采样技术的Occ-VAE，以增强可控性。DOME展示了高保真度生成，有效地预测了占用空间中未来场景的变化，并且可以生成比以前方法长两倍的占用序列。这种方法在自动驾驶的端到端规划中具有应用前景。不过作者发现训练作者的模型仍然需要大量的计算资源。在未来，作者将探索更轻量级和计算效率高的方法，或者采用微调范式以减少资源需求。

参考

[1] DOME: Taming Diffusion Model into High-Fidelity Controllable Occupancy World Model

『自动驾驶之心知识星球』欢迎加入交流！重磅，自动驾驶之心科研论文辅导来啦，申博、CCF系列、SCI、EI、毕业论文、比赛辅导等多个方向，欢迎联系我们！

① 全网独家视频课程

端到端自动驾驶、仿真测试、自动驾驶C++、BEV感知、BEV模型部署、BEV目标跟踪、毫米波雷达视觉融合、多传感器标定、多传感器融合、多模态3D目标检测、车道线检测、轨迹预测、在线高精地图、世界模型、点云3D目标检测、目标跟踪、Occupancy、CUDA与TensorRT模型部署、大模型与自动驾驶、NeRF、语义分割、自动驾驶仿真、传感器部署、决策规划、轨迹预测等多个方向学习视频（扫码即可学习）

网页端官网：www.zdjszx.com

② 国内首个自动驾驶学习社区

国内外最大最专业，近4000人的交流社区，已得到大多数自动驾驶公司的认可！涉及30+自动驾驶技术栈学习路线，从0到一带你入门自动驾驶感知（端到端自动驾驶、世界模型、仿真闭环、2D/3D检测、语义分割、车道线、BEV感知、Occupancy、多传感器融合、多传感器标定、目标跟踪）、自动驾驶定位建图（SLAM、高精地图、局部在线地图）、自动驾驶规划控制/轨迹预测等领域技术方案、大模型，更有行业动态和岗位发布！欢迎扫描下方二维码，加入自动驾驶之心知识星球，这是一个真正有干货的地方，与领域大佬交流入门、学习、工作、跳槽上的各类难题，日常分享论文+代码+视频