DFIT-OccWorld：通过解耦动态流和图像辅助训练实现高效占用世界模型

24年12月来自香港中文大学深圳分校、香港大学和华为诺亚的论文“An Efficient Occupancy World Model via Decoupled Dynamic Flow and Image-assisted Training”。

自动驾驶领域对世界模型的兴趣日益浓厚，该模型旨在根据历史观察预测潜在的未来场景。本文介绍 DFIT-OccWorld ，这是一种高效的 3D 占用世界模型，它利用解耦动态流和图像辅助训练策略，大幅提高 4D 场景预测性能。为了简化训练过程，摒弃之前的两阶段训练策略，将占用预测问题重新表述为一个解耦的体素扭曲过程。模型通过体素流扭曲现有观测来预测未来的动态体素，而静态体素则可以通过姿势变换轻松获得。此外，该方法结合图像辅助训练范式来提高预测可靠性。具体而言，采用可微体渲染，通过预测的未来体去生成渲染的深度图，这些深度图用于基于渲染的光度一致性。实验证明该方法的有效性，展示其在 nuScenes 和 OpenScene 基准测试中针对 4D 占用预测、端到端运动规划和点云预测的性能领先。具体而言，与现有的 3D 世界模型相比，它实现最先进的性能，同时计算成本大幅降低。

3D 语义占用率对 3D 体内的占用状态和语义信息进行编码，成为描述自动驾驶 3D 场景的有效表示。基于这一概念，Zheng 提出 OccWorld [51]，这是一个 3D 世界模型，可以同时预测未来的占用率并根据过去的占用率观测规划本车的轨迹。具体来说，OccWorld 采用两阶段训练策略，如图 (a) 所示：在第一阶段，矢量量化变分自编码器 (VQ-VAE) [34] 充当占用token化器，以自监督的方式学习离散场景tokens。这可以从观察占用率输入中获得高级表示。在第二阶段，4D 占用预测问题被重定义为预测码本索引的分类任务，其中时空生成式transformer被用作预测未来占用的世界模型。虽然展示有希望的结果，但 OccWorld 面临着影响其现实世界灵活性和可扩展性的限制：

1. • 第一阶段 VQ-VAE 训练严重影响效率，并为第二阶段引入性能瓶颈。为场景token化器找到合适的超参在平衡重建和预测性能方面仍然具有挑战性 [51]。

2. • OccWorld 仅以自回归方式依赖隐特征。它要么忽略相邻场景中的显式结构一致性，要么忽略图像中随时可用的纹理信息。这些阻碍模型完全捕捉环境动态的能力。

为了解决 OccWorld 中发现的缺点，本文在自动驾驶中引入一种端到端 3D 占用世界模型。该框架建立在单阶段视频预测的优势和潜力之上，能够同时预测多个未来体和图像，如上图（b）所示。这种设计有利于未来预测之间的信息共享，提高结果的准确性和合理性。此外，直接预测每帧的占用率会导致不理想的性能，因为大多数体素是空的。为了解决这个问题，用占用网络预测的语义信息将体素解耦为动态和静态类别。然后，世界模型只预测动态目标的体素流并相应地扭曲这些体素。对于静态目标，由于它们的全局位置保持不变，可以通过姿势变换轻松获得它们。此外，提出一种图像辅助训练策略，在训练阶段将图像作为辅助输入。

问题定义

在 OccWorld [51] 中，世界模型 W 将具有 N/h 个三维语义占用率的历史序列 O/T-N/h:T 作为输入，以及过去的自我轨迹序列 P/T-N/h:T ，以预测未来的占用率和自我位置。它采用两阶段训练方案，以自回归方式预测未来帧，如上图 (a) 所示。

该框架不像两阶段训练策略那样灵活和可扩展，而自回归预测策略容易受到误差累积问题的影响。本文提出一种端到端的 3D 占用世界模型，该模型采用单阶段训练范式来模拟周围场景的演变。

然而，直接回归未来占用而忽略大多数体素为空的证据是不太好的。因此，该世界模型可以表述为 W^∗ = W/p (W/f )，其中 W/f 预测相对于当前帧的未来流图，而未来占用可以通过扭曲操作 W/p 获得。

此外，考虑到图像是免费的，并且可用于自动驾驶中的占用预测，在训练过程中将它们用作辅助输入。

框架概述

如图所示，该方法首先将历史 3D 占用、图像序列和自我姿势输入预处理为时空tokens。然后将它们输入到跨模态时空编码器，以有效捕获模态之间的空间结构和局部时间依赖性。之后，利用各种解码器同时预测未来帧。为了促进学习过程，实施解耦动态流策略。此外，还提出一种图像辅助训练策略来增强占用预测，而不会在推理过程中产生计算负担。

Token化

3D 占用时空编码 。给定一系列历史观测的 N/h 帧，每个占用帧的 3D 占用表示为 O/i，H/0、W/0 和 D/0 表示以自我为中心的周围空间分辨率。每个体素被分配为 C 个类之一。为了将占用序列编码为时空tokens，首先使用可学习的类嵌入将 3D 占用映射到占用嵌入 y ˆ ，C/0 是嵌入通道数。然后，为了减少计算负担，按照先前的工作 [51] 将 3D 占用嵌入转换为 BEV 表示 y ̃。之后，将 BEV 嵌入分解为不重叠的 2D 块 y /p，H=H/0/P，W=W/0/P，C′=P^2·D/0·C/0，P 为每个图像块的分辨率。然后使用由多个 2D 卷积层组成的轻量级编码器（即 Conv2d-GroupNorm-SiLU）提取块嵌入，获得编码的历史占用时空token y 。

图像时空编码 。给定相同序列长度的历史图像 I/T −N/h :T ，还将 2D 图像模态处理为图像时空token x 。具体而言，将输入图像调整为与 BEV 表示相同的大小，并将 2D 图像分解为 2D 块。最后，使用独立的 2D 编码器提取每帧的图像块嵌入。

自我姿态编码 。与 [51] 一致，将自我姿态表示为 2D 地面上相邻帧之间的相对位移。给定历史自我姿态，利用多个线性层，然后使用 ReLU 激活函数来获取自我token e 。

跨模态时空编码器

跨模态时空编码器不仅负责捕获输入的空间结构和时间依赖性，还负责学习不同模态之间的相关性。

空间-觉察局部时间注意块 。受先前视频预测工作的启发 [28]，在编码器和解码器中采用 空间-觉察局部时间 (SALT) 注意块，如图 (a) 所示。它首先利用 2D 卷积层为时空token生成查询图和配对KV。这种空间觉察 CNN 操作很好地保留结构信息。然后，使用标准多头注意来捕获token之间的时间相关性。通过这种方式，它可以学习时间相关性并保留序列的空间信息。此外，它用 3D CNN 替换了 FFN（前馈网络）层，引入用于顺序建模的局部时间线索。

时间位置编码 。初始化一个全局嵌入 T，Nf 表示未来帧的数量。随后，合并自我姿态嵌入 e 并将其与占用和图像时空token集成在一起。在此集成之后，由多个掩码 SALT 块（下面会介绍）组成的编码器处理沿时间维度连接的多模态输入。该架构旨在有效地模拟各种输入之间的时空关系。

解码器

在获取编码token后，开发三个解码器：占用解码器、图像解码器和姿态解码器，预测未来的观察结果。为了增强学习过程，将未来占用预测任务重新定义为体素流估计问题，并结合扭曲过程。因此，占用解码器由流预测、扭曲过程和细化模块组成，它们共同促进精确有效的未来占用预测。

流预测 。流解码器由多个堆叠的 SALT 块组成，处理编码的历史 BEV 特征，根据当前自我坐标预测未来流。然后，用当前到未来的变换矩阵将绝对流图转换为未来帧。