AI Research 技术调研

OccSora: 4D Occupancy Generation Models as World Simulators for Autonomous Driving

12分钟阅读

一句话定位

扩散式 4D occupancy 生成世界模型：用 3D VQ-VAE 式 tokenizer 把 nuScenes-Occupancy 的长序列体素占据压缩成时空 token，再用条件扩散 Transformer（DiT-XL/2）在轨迹提示下去噪生成，用”整段时空 diffusion”替换”逐帧自回归 next-token”，把 OpenAI Sora 的思路搬到 3D occupancy 领域。

背景与定位

自动驾驶世界模型的主线是：GAIA-1、DriveDreamer 等在图像/BEV 空间做条件生成，OccWorld 把这条路线搬进 3D occupancy 空间但仍用自回归 next-token 预测未来占据和自车 token。OccSora 直接对标 OccWorld（同一作者 Wenzhao Zheng 团队的前作，论文正文引用为 zheng2023occworld，库内暂无独立页面故不做 wikilink），指出自回归框架在长时序生成上效率低、误差易累积，转而借鉴 OpenAI Sora（视频扩散世界模拟器）的思路，提出用扩散 Transformer 直接对整段时空占据 token 做去噪生成，条件是任意自车轨迹而非历史帧或 3D 检测框。论文自称是”the first 4D occupancy generation model”（4D occupancy 生成模型的首个工作）。与其后的 Drive-OccWorld、DOME 等 occupancy world model 相比，OccSora 更早验证”扩散替代自回归”这条路线，但代价是压缩率极高、重建保真度大幅下降（见评测部分）。

模型架构

两阶段：(1) 4D occupancy scene tokenizer（VQ-VAE 式 3D 卷积编解码器 + codebook），(2) diffusion-based world model（DiT 结构扩散 Transformer）。

4D occupancy scene tokenizer

  • 输入 R_in ∈ R^{B×D×H×W×T}（体素占据栅格，D 为高度维、H×W 为 BEV 平面、T 为帧数）。先对每个语义类别做可学习的 category embedding b∈R^{c′}，把 R_in reshape 为 R_in′ ∈ R^{B×(D·c′)×T×H×W}。
  • 3D 编码器：一系列 3D 下采样卷积层，论文原文表述为对输入做”three times”下采样，累计把时间 T 与空间 H×W 都压缩到 1/8、通道提升到 8·D·c′，得到 R_in″ ∈ R^{B×(8·D·c′)×T/8×H/8×W/8}（论文未逐轮给出单轮压缩比）；下采样后在前馈/注意力层间加 dropout 正则，并引入 cross-channel attention——把 8·D·c′ 通道维切段后做跨段注意力，用于捕捉不同轴特征之间的关系，reshape 回原形状得到 R_mi。
  • 量化：codebook ζ_token ∈ R^{N×D}，逐位置对特征做最近邻量化（L2 距离），得到离散场景 token R_mi。
  • 3D 解码器：cross-channel attention + 残差块 + 一系列 3D 反卷积，做时间/空间上采样，重建体素占据 R_o。
  • 具体配置（nuScenes-Occupancy）：对 32 连续帧做三轮压缩，通道维提升到 128；对应消融表里 Input Size 32×200×200 → Token Size 128×4×25×25（时间 32→4、空间 200→25，均 ×8 压缩，通道 128）。

diffusion-based world model

  • 把 R_mi flatten 成 R_re，用 sin/cos 函数做位置编码得到 R_re^(emb)；轨迹条件 T_r ∈ R^{B×t×2}（自车 x/y 绝对坐标序列）先 reshape 再经 MLP δ(T_r) 编码，与去噪时间步嵌入 ν(t) 相加得到条件 g，与 token 嵌入一起送入 Transformer block。
  • 标准 DDPM 前向加噪 q(R_re^g|R_re) = N(√σ̄ᵍ R_re, (1-σ̄ᵍ)I)，反向去噪 p_θ(R_re^{g-1}|R_re^g) 由网络预测均值/协方差；训练先用简单噪声 MSE（L_simple）预训练，再联合完整 KL 变分下界 L 训练协方差，遵循 Dhariwal & Nichol (2021) 的两阶段做法。
  • 主干为 DiT-XL/2：论文正文未给出具体层数/隐藏维等超参，但官方 GitHub 训练脚本 torchrun ... train_2.py --model DiT-XL/2 明确采用 Facebook DiT 代码库的 XL/2 规格（该细节来自代码仓库而非论文正文）。
  • 推理：从纯高斯噪声 R_re^g ~ N(0,I) 出发，配合任意轨迹条件迭代采样 R_re^{g-1}，最终解码得到 4D occupancy 视频 R_o。

数据

  • 数据集：nuScenes-Occupancy（Occ3D 语义占据标注），train/val 切分沿用 TPVFormer 仓库提供的 pickle 索引文件（nuscenes_infos_train/val_temporal_v3_scene.pkl）。
  • 序列长度：每次处理 32 连续帧；摘要、GitHub README、项目页三处均一致表述为可生成”16s-videos”（32 帧与 16 秒的换算关系论文未在正文中显式给出，此处不做推算，仅如实转述三处一致的表述）。
  • 空间分辨率：BEV 平面 200×200（消融表 Input Size 直接给出），高度维 D 的具体取值论文未单独列出。
  • 未披露：训练/验证场景数、总体素/帧总量、是否有数据增强或额外过滤、动作（轨迹）标签是否经平滑/插值处理。
  • 单一数据来源，无仿真数据、无跨数据集混合或跨具身泛化设计。

训练方法

  • 两阶段串行、非联合训练：先训 tokenizer（VQ-VAE），固定后用其把训练集全部编码为离散 token，再单独训扩散模型。
  • Tokenizer 训练：AdamW，初始学习率 1×10⁻⁵，weight decay 0.01；8×NVIDIA A100，每卡 batch size 2；每卡显存约 42GB；训练 150 epoch，耗时 50.6 小时。
  • 扩散模型训练：同样 8×A100、每卡 batch size 2（论文用同一超参描述覆盖两阶段）；每卡显存约 47GB；训练 1,200,000 步，耗时 108 小时。
  • 训练目标：tokenizer 阶段为 VQ 最近邻量化 + 编解码重建（具体重建/commitment loss 权重未披露）；扩散阶段为 DDPM 噪声预测 MSE（L_simple）+ 完整 KL 变分下界 L 两阶段训练。
  • 条件信号为真实自车轨迹（ego trajectory），推理时可替换为任意轨迹（直行/右转/静止等）以实现轨迹可控生成；未使用强化学习或额外蒸馏。

Infra（训练 / 推理工程）

  • 训练硬件：8×NVIDIA A100（80GB，GitHub README 明确要求”on A100 with 80G GPU memory”）。
  • 并行方式：GitHub 训练命令为单机 8 卡 torchrun --nnodes=1 --nproc_per_node=8，即数据并行；未披露张量并行/FSDP 等高级并行策略。
  • 显存占用：tokenizer 阶段每卡约 42GB，扩散阶段每卡约 47GB。
  • 训练时长：tokenizer 150 epoch / 50.6 小时；扩散模型 1,200,000 步 / 108 小时（均为 8×A100 累计墙钟时间，论文未给出总 GPU-hours）。
  • 精度（FP16/BF16/FP32）：未披露。
  • 推理 FPS / 控制频率 / 端上延迟：未披露——论文仅在消融中给出去噪步数（10/100/1000）与去噪完成比例对 FID 的影响（见评测），未报告生成一段 16s 视频的实际采样耗时。

评测 benchmark

4D occupancy 重建（Table 1，vs OccWorld，同为 nuScenes-Occupancy）

Method压缩比IoUmIoU
OccWorld1662.265.7
OccSora51237.027.4

论文原话：尽管压缩率是 OccWorld 的 32 倍，OccSora 仍”maintains over half of / nearly 50% of”OccWorld 的重建精度（实际 27.4/65.7 ≈ 41.7%，为论文原文表述，未做数值修正）。

4D occupancy 生成 FID 对比（Table 2，跨数据集/模态，作者自陈仅供参考，非严格同基准对比）

MethodTypeDimDatasetFID
DiTImage2DImageNet12.03
MagicDriverVideo3DnuScenes14.46
DriveDreamerVideo3DnuScenes14.9
DriveGANVideo3DnuScenes27.8
SemCityOccupancy3DKITTI40.63
OccSoraOccupancy Video4DnuScenes8.348

组件消融（Table 3）

Input SizeToken SizeChannelT embed.TrajectoryIoUmIoUFID
32×200×200128×4×25×258✓✓37.0327.428.34（基线）
32×200×200128×4×25×258✗✓37.0327.4287.26
32×200×200128×4×25×258✓✗37.0327.4217.48
32×200×200128×4×25×254✓✓29.6723.2134.24
32×200×200128×8×50×508✓✓32.9124.472.32
12×200×20064×3×50×508✓✓26.7314.12187.78
12×200×20064×3×25×258✓✓22.429.27270.23
12×200×20032×3×25×258✓✓13.603.85465.18

去掉时间步嵌入 FID 从 8.34 骤增到 87.26；去掉轨迹条件 FID 升到 17.48，且论文描述”生成场景失去运动控制，呈现近乎线性的、由数据分布主导的运动模式”；class 通道数 8→4 时 mIoU 27.42→23.21（论文正文称重建精度”下降约 18%”）、FID 恶化到 34.24；降低压缩率（token 变大到 8×50×50）时 IoU/mIoU 略升但 FID 反而恶化到 72.32；进一步缩小输入/加大压缩（12×200×200 系列）时 IoU/mIoU/FID 全面恶化。

去噪步数与去噪比例（Table 4）：在 10/100/1000 三种总去噪步数下，FID 随”去噪完成比例”（10%→100%）大幅下降，例如 token 128×4×25×25、1000 步配置下 FID 从 32171（10%）降到 8.94（100%）；64×3×50×50 配置同样趋势但绝对值更差（1000 步 100% 时 FID=353）。论文结论：去噪步数对结果影响相对较小，去噪比例与 token 规模/通道数才是生成质量的主导因素。

无闭环规划或下游控制任务评测（论文将此列为未来工作）。

创新点与影响

  • 提出（作者自称）首个扩散式 4D occupancy 生成模型，把 occupancy world model 从”自回归 next-token”（OccWorld 路线）换成”扩散去噪整段时空 token”，无需先验 3D 框、地图或历史帧输入，仅用随机噪声 + 任意轨迹即可生成 16 秒时长、具时空一致性的 3D occupancy 视频。
  • 具体方法设计：(1) tokenizer 中的 cross-channel attention，用于建模压缩后不同通道段（对应语义/空间/时间轴）之间的关系；(2) 轨迹条件经 MLP 编码后与去噪时间步嵌入相加，一并作为 DiT 的条件输入，实现轨迹可控场景生成。
  • 论文自陈局限（原文 Limitations）：受体素粒度限制，无法构建更精细的 4D 场景；生成结果对运动物体的细节表现不一致，推测与训练数据规模较小有关。未来方向：更精细的 4D occupancy 世界模型，以及闭环设定下的端到端自动驾驶探索。
  • 客观影响：重建精度相对 OccWorld 大幅下降（mIoU 约打四折），换来生成 FID 数值上的优势和”非逐帧自回归”的生成范式——是”生成效率/质量 vs. 重建保真度”权衡的一个早期数据点，为后续 occupancy world model（如 Drive-OccWorld、DOME 等）提供了扩散路线的参照。

原始链接

一手源存档（sources/）

关系图谱

反向链接