AI Research 技术调研

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

14分钟阅读

一句话定位

把自动驾驶 3D 占据（occupancy）预测建模成连续潜空间上的时空扩散过程：用连续 VAE（Occ-VAE）取代离散 token 化以保留细节，再用时空分离的 Diffusion Transformer 支持历史帧条件化和长时程 rollout，并靠轨迹重采样数据增强解决控制多样性不足的问题。

背景与定位

自动驾驶世界模型大致分两条路线：2D 视频类（如 GAIA-1、DriveDreamer-2）生成前视/环视画面，但存在跨视角、跨时刻一致性差的问题；3D 表征类则直接对 LiDAR 点云或 3D 占据网格建模，天然保持 3D 一致性，且占据标注可以从稀疏 LiDAR 语义（Occ3D）或自监督获得，对输入模态（相机 / LiDAR）不敏感。

在占据世界模型内部，此前存在两条技术路线，各有明显短板：

  • 自回归离散 token 路线，如 occworld（OccWorld）、OccLLaMA：先用离散 tokenizer 压缩占据数据再做类 GPT 自回归预测，量化过程丢信息、长序列训练也不稳定；
  • 扩散路线的 OccSora：把时空信息拉平成一维 token 序列整体生成，导致历史占据信息难以作为条件注入，本质上”只能生成、不能预测”，限制了作为世界模型的实用性。

DOME 的定位是把上述两条路线的优点结合：用连续潜空间（而非离散码本）保留细节，同时采用时空分离（而非拉平）的 Diffusion Transformer 架构（借鉴视频生成的 Latte）来支持历史帧条件化和长时程生成，并额外引入轨迹重采样解决 nuScenes 数据集里”车辆几乎都直行”（约 87%）导致的可控性不足问题。论文在两项任务上都与 OccWorld、OccSora、OccLLaMA 直接对比并声称刷新 SOTA。

模型架构

DOME 由两个组件构成：Occ-VAE（占据数据的连续 tokenizer）+ 时空扩散 Transformer（预测未来占据的隐变量）。

Occ-VAE：

  • 输入 3D 占据体素网格 x∈R^{H×W×D}（每格是 17 类语义标签之一），先通过一个可学习的类别嵌入 E_cls∈R^{n×C_emb} 把离散语义 ID 转成 BEV 风格张量 x_bev∈R^{H×W×(D·C_emb)}；
  • 编码器（2D 卷积 + attention block）输出 μ、σ∈R^{n_h×n_w×C}，用重参数化技巧采样连续潜变量 z=μ+σ⊙ε；
  • 解码器用 3D 反卷积层上采样潜表征以保持时间一致性，重塑回 H×W×D×C_emb 后与类别嵌入做点积得到 logits，argmax 得到语义类别；
  • 训练损失 = 交叉熵重建损失 + β·KL 散度 + λ·Lovasz-softmax 损失（缓解类别不均衡）；
  • 压缩率 64×（OccWorld/OccLLaMA 为 16×，OccSora 为 512× 但代价是重建精度大幅下降——OccSora 额外做了 8× 的时间维压缩，DOME 选择不做时间压缩以平衡压缩率与重建精度）。

时空扩散 Transformer（DOME 主体）：

  • 借鉴视频生成的 Latte（Ma et al., 2024）设计，潜变量 z∈R^{n_f×n_h×n_w×C} 切分为 n_f 帧、每帧 n_t=(n_h/p)×(n_w/p) 个 token（p 为 patch size），空间块（spatial block）与时间块（temporal block）交替堆叠：空间块在同一时间步内做跨空间 attention，时间块在固定空间位置上沿时间轴做 attention；
  • 位置编码：ViT 风格 2D sin-cos 固定空间位置编码 + 1D sin-cos 固定时间位置编码；
  • 历史占据条件化：构造条件掩码 M=[t<n_c]，训练时用真实历史潜变量 z_c 替换被掩码位置的加噪 token（ẑ_i=M·z_c+(1-M)·z_i），扩散损失只在未被掩码的（待预测）帧上计算；同时以比例 p_inj 的训练迭代完全丢弃上下文帧，使模型也能无条件生成（p_inj 具体数值论文未披露）；
  • 轨迹条件化：由自车相邻帧位姿计算相对平移 Δt_t 与相对旋转 ΔR_t，提取 [x,y,θ_yaw]，用 NeRF 式正弦-余弦位置编码（Mildenhall et al., 2020）编码后投影到隐藏维度，与扩散时间步嵌入相加后送入 adaLN 模块；
  • 采样：DDPM，训练 1000 步 / 推理 20 步；
  • 长时程 rollout：每次仅生成 n_f 帧窗口，用上一窗口最后预测帧作为下一窗口的上下文帧迭代滚动，配合轨迹切片对齐，从而支持任意长度生成。

关键配置数字：n_f=11（每窗口帧数），n_c=4（默认上下文帧数，用于 4D 预测；控制性优先场景下 n_c=1）、n_h=n_w=25（Occ-VAE 输出的空间潜分辨率）、patch size p=1、DiT 层数 14 层空间块 + 14 层时间块。

数据

  • 数据集：nuScenes（Caesar et al., 2019），占据标注来自 Occ3D（Tian et al., 2023）；沿用 occworld 的划分，训练 700 个场景/序列，验证 150 个；每个序列约 40 帧，采样率 2Hz。
  • 体素网格：分辨率 [0.4, 0.4, 0.4] m，感知范围 [-40m,-40m,-1m, 40m,40m,5.4m]，网格大小 [200,200,16]，每格 17 类语义标签（LiDAR 语义标注得到）。
  • 两种输入设置：DOME-O 用真值 3D 占据作为输入（占据侧的上界设置）；DOME-F 用 FB-OCC（相机输入的 3D 占据预测器）产出的占据作为输入，构成端到端纯视觉 pipeline。
  • 轨迹重采样（数据增强，论文称是首个针对世界模型预测的占据数据增强方法）：
    1. 多帧点云聚合——把自车坐标系下的占据序列转成 3D 点云，用自车位姿变换到世界坐标系，按语义标签过滤掉动态物体（车辆、行人等）；
    2. 生成可行驶区域——聚合后按道路类别过滤并体素化成俯视 BEV 地图；
    3. 生成多样且可行的轨迹——在 BEV 地图上随机采样起点/终点，用平滑 A*（Hart et al., 1968）连接生成轨迹，转换为 R^{4×4} 位姿（z 坐标置 0）；
    4. 用新轨迹位姿从聚合点云重新提取占据真值（沿用 Occ3D 式提取方法），得到与新轨迹配对的新占据序列。
  • 该方法只需要位姿和占据数据（不需要原始 LiDAR 或 3D 框），因此可推广到机器标注、LiDAR 采集或自监督产生的任意占据数据源。动机数据：nuScenes 训练集里约 87% 的场景是直行（见论文 Fig. 4c），且每个场景自车只经过一次，导致同一场景下轨迹条件的占据样本缺乏多样性、模型易过拟合到具体场景。

训练方法

两阶段训练：Occ-VAE 先单独训练，训完冻结，编码器作为特征提取器、解码器用于把 DOME 输出的潜变量还原成占据数据；随后训练时空扩散 Transformer。

  • Occ-VAE：AdamW，lr=1×10⁻³，cosine scheduler，输入形状 200×200×16，batch size 10/GPU，200 epoch，8× RTX 4090。
  • DOME（扩散 Transformer，论文称”第二、三阶段”）：14 层空间块 + 14 层时间块，AdamW + EMA，batch size 8/GPU，2000 epoch，32× RTX 4090；使用 xformers、混合精度、梯度检查点降低显存占用。
  • 扩散目标：标准 DDPM 噪声预测损失，按条件掩码 (1-M_t) 只在待预测帧上计算（式 3）。
  • 推理时的两种上下文帧数设置：n_c=4（更多历史帧提供更明确的运动/场景信息，利于 4D 预测精度）与 n_c=1（减少历史占据运动先验的干扰，更利于轨迹控制精度，仅保留一帧起始观测）。
  • 未采用 RL；无蒸馏；无预训练大模型权重初始化的说明。

Infra（训练 / 推理工程）

  • Occ-VAE 训练：8× RTX 4090，batch 10/GPU，200 epoch。
  • DOME 扩散 Transformer训练：32× RTX 4090，batch 8/GPU，2000 epoch。
  • 显存/效率优化：xformers 高效 attention、混合精度、梯度检查点。
  • 未披露：总 GPU-hours、具体并行策略（是否用 FSDP/张量并行，论文只说明了 GPU 数和 batch/GPU，隐含数据并行）、混合精度具体位宽（fp16/bf16）、推理侧 FPS / 控制频率 / 端到端延迟、部署硬件。
  • 采样步数：训练 1000 步 DDPM，推理 20 步（加速采样），但未给出这 20 步对应单帧或整窗口生成的实际耗时。

评测 benchmark

评测数据集统一为 nuScenes + Occ3D 标注，指标为 mIoU / IoU。

Occupancy Reconstruction（Occ-VAE tokenizer 质量，Table 1）：

Method压缩率mIoUIoU
OccWorld16×65.762.2
OccSora512×27.437.0
OccLLaMA16×75.263.8
DOME (ours)64×83.177.3

4D Occupancy Forecasting（Table 2，Avg = 1s/2s/3s 平均）：

Method输入Recon mIoUAvg mIoURecon IoUAvg IoU
Copy&Paste3D-Occ66.3811.3362.2920.52
OccWorld-O3D-Occ66.3817.1462.2926.63
OccLLaMA-O3D-Occ75.2019.9363.7629.17
DOME-O (ours)3D-Occ83.0827.1077.2536.36
OccWorld-FCamera20.096.1635.6118.99
OccLLaMA-FCamera37.388.6638.9222.99
DOME-F (ours)Camera75.0018.2574.3128.84
  • 摘要中的提升幅度：占据重建相对 OccLLaMA 提升 mIoU +10.5%（(83.1-75.2)/75.2）、IoU +21.2%（(77.3-63.8)/63.8）；4D 预测相对 OccLLaMA-O 提升 mIoU +36.0%（(27.10-19.93)/19.93）、IoU +24.6%（(36.36-29.17)/29.17）——均以 OccLLaMA 为基线计算相对提升。

消融（Table 3，逐项叠加）：

Spatio-TempTrajResYawmIoUIoU
✗✗✗✗13.0823.10
✓✗✗✗18.6028.09
✓✓✗✗24.2434.28
✓✓✓✗27.0036.39
✓✓✓✓27.1036.36

时空架构本身贡献 +5.52 mIoU，轨迹条件贡献 +5.64 mIoU，轨迹重采样贡献 +2.76 mIoU，yaw 角嵌入贡献边际（+0.10 mIoU，IoU 反而略降 0.03）。

上下文帧数消融（Table 5）：无轨迹条件时，1/2/3/4 帧 mIoU 分别为 12.59/20.01/20.70/20.07（3 帧最优，4 帧反而略降）；有轨迹条件时 1/2/3/4 帧 mIoU 分别为 22.24/25.41/26.61/27.10（单调提升）——说明轨迹信息比堆叠更多历史帧更高效，模型不需要在含糊的历史帧中自行推断运动趋势。

生成时长对比（Table 4）：同为 2Hz 采样率下，OccWorld/OccLLaMA 生成 6 帧/3s，OccSora 生成 32 帧/16s，DOME 生成 64 帧/32s——是 OccWorld/OccLLaMA 的 10 倍、OccSora 的 2 倍。

创新点与影响

  • 首次把连续潜空间（而非离散 token）与时空分离的 Diffusion Transformer 结合用于占据世界模型，弥补了自回归离散 token 路线（OccWorld、OccLLaMA）因量化损失导致的细节丢失。
  • 解决了此前扩散路线占据世界模型（OccSora）“只能生成不能预测”的问题——通过帧级条件掩码机制，让扩散模型原生支持历史占据条件化。
  • 轨迹重采样是论文自称的首个占据数据增强方法：仅需位姿+占据数据（不依赖原始 LiDAR/3D 框）即可合成多样且物理可行的自车轨迹与配对占据，缓解 nuScenes 约 87% 直行、单场景轨迹多样性不足的问题，可推广到任意占据数据来源。
  • 在 nuScenes/Occ3D 上同时刷新占据重建与 4D 预测两项任务的 SOTA，且生成时长达到 OccWorld/OccLLaMA 的 10 倍、OccSora 的 2 倍（64 帧/32s @ 2Hz）。
  • 论文自陈的局限：训练仍需要相当规模的计算资源（8×RTX4090 训 Occ-VAE、32×RTX4090 训 2000 epoch 的扩散 Transformer），未来计划探索更轻量、计算效率更高的方法，或采用微调范式降低资源需求。

原始链接

一手源存档（sources/）

关系图谱

反向链接