AI Research 技术调研

NVIDIA OmniDreams: Real-Time Generative World Model for Closed-Loop Autonomous Vehicle Simulation

18分钟阅读

一句话定位

OmniDreams 是 NVIDIA 从 Cosmos-Predict 2.5 mid/post-train 出来的一个 2B 参数、动作条件化的自回归 causal diffusion 驾驶世界模型：靠 streaming KV cache 维持长时一致性,靠 Self Forcing + DMD 蒸馏把生成压到 2 步扩散,单卡 GB300 上单视角可达 68 effective FPS(720p),16 卡 GB300 上四视角联合生成可达 105 effective FPS/相机;接入 AlpaSim + Alpamayo 1 闭环仿真栈后,论文进一步证明从同一 backbone fine-tune 出的 2B World-Action Model(WAM)策略,在碰撞率上超过参数量约 5 倍的 VLA 策略 Alpamayo 1.5(全部碰撞 6.9%→4.2%)。

背景与定位

闭环仿真要求策略动作实时改变仿真状态、进而改变下一步的传感器观测。以 3D Gaussian Splatting/NeRF 为代表的重建式仿真器(NVIDIA NuRec 等)锚定在原始采集数据上,虽然能在被重建的场景内做”what-if”测试,但难以外推到未采集的视角、天气或 agent 行为——这是论文反复强调的核心动机。

生成式世界模型路线试图用大规模视觉先验绕开这一限制。论文把自己放在两条脉络的交叉点:一是通用视频生成”世界模拟器”路线(Sora、Movie Gen、Wan、Veo 3、Genie),二是专门面向自动驾驶的生成式世界模型路线(DriveGAN、DriveDreamer、GAIA-1/2、Vista、MagicDrive、Drive-WM、GenAD、Waymo World Model)。OmniDreams 的定位是:直接从 Cosmos-Predict 2.5 这个通用视频基础模型出发,加上动作条件化(policy 轨迹)与 world-scenario-map 控制,在与真实 policy(Alpamayo 1)、真实仿真编排器(AlpaSim)组成的闭环系统里满足每 chunk 的延迟预算——这是它与前述通用视频生成模型的关键区别:后者不面向逐 chunk 实时交互。训练数据(RDS-HQ-1M)是 Cosmos-Drive-Dreams 所用 RDS-HQ 数据集的扩展版。

模型架构

两个变体:OmniDreams-SV(单视角,前视相机,每 step 生成 8 帧/2 个 latent 帧,120° 前视);OmniDreams-MV(四视角联合:front-wide、cross-left、cross-right、front-telescope,每 step 生成 16 帧/4 个 latent 帧)。两者均为 2B 参数(HF model card:Omniverse-Dreams-v1-2B)。

Backbone:继承 Cosmos-Predict 2.5 的 causal transformer 架构,用 rectified-flow/flow-matching 目标训练:L = E_{x,t}[‖u_θ(x_t,t)-v_t‖²],x_t=(1-t)x+tε,v_t=ε-x,t 采自 logit-normal 分布。

三路条件输入(见论文 Fig.3):

  • First-frame RGB:仿真会话首帧,编码为 clean latent 初始化生成;
  • Text prompt:描述天气/光照/时段等环境属性,用 Cosmos text encoder 编码,经 cross-attention 注入;
  • Abstract world-scenario map:HD map(车道线、路缘、停止线、人行横道、交通灯/标志等)+ 动态 agent 3D bbox 投影渲染成的 pixel-aligned 条件视频。轻量 control branch——用一个小 MLP 把结构化 world-scenario 编码为 control token,与 latent token 对齐后直接拼接送入 transformer,而非另起一个 ControlNet 式独立网络,开销很小且训练更稳定;
  • Memory cache:streaming KV cache,复用历史生成的 attention key/value 以维持长时一致性。

多视角一致性:每个视角配一个可学习 view embedding,通过 AdaLN(与 timestep embedding 同一注入机制、zero-init)注入;在每个 transformer block 的 text cross-attention 之后加一层 cross-view attention(输出投影 zero-init),让不同视角的 token 互相 attend 以对齐共享场景结构。朴素的跨视角/跨时间全自注意力复杂度是 O(N²T²)(N=视角数,T=时间长度),该设计把它因式分解为时间维内部 attention(走 causal KV cache)+ 视角维跨视角 attention,复杂度降为 O(NT²)+O(N²)。实践中支持最多 7 个同步视角的环视仿真。

数据

训练数据横跨 15 个国家(US、DE、JP、KR、GB、FR、ES、SE、PT、DK、FI、PL、IT、AT、BE)的真实驾驶日志:

  • RDS(mid-training):16,600 小时,3M 条 20 秒 clip,7 个同步相机(front-wide/front-telescope/front-left/front-right/rear-left/rear-right/rear-tele)@ 30fps、1080p,兼作 Cosmos-Predict 2.5-AV 预训练集的一部分。
  • RDS-HQ-1M(finetune/post-training):4,944 小时,1,142,285 条 clip(504,488 条 10s ≈44% + 637,797 条 20s ≈56%),world-scenario 标注质量更高,是 RDS-HQ 数据集的扩展版。
  • 训练分辨率 704×1280,7 个视角中 4 个用于训练(SV 用 1 个,MV 用 4 个)。
  • Held-out 5,000 条 clip 用于评测/测试,按道路使用者(骑行者、行人、摩托车、婴儿车)、大型车辆、天气/光照(雨雪雾夜)、罕见基础设施(隧道、铁路道口、施工区、事故现场)等场景类别平衡采样(而非按训练分布比例采样);其中 300 条额外采集了 60 秒长版本用于长时一致性评测。
  • 条件信号提取:HD map 标注来自预构建的城市级地图并严格匹配驾驶路线;动态目标 3D 检测/追踪跑在 10 FPS、插值到 30 FPS 匹配视频帧率。文本 caption 用 Qwen2.5-VL-7B 在 10 秒窗口上对 7 个相机独立生成,每个窗口生成 short(≈40 词)/medium(≈80 词)/long(≈200 词)三种长度,训练时按 0.1/0.2/0.7 概率采样以增强模型对 prompt 长度的鲁棒性。
  • 数据质量过滤:剔除 ego 轨迹跳变/自动标注高不确定性/预测分歧的序列;用 VLM 检测色差等视觉伪影并剔除受影响 clip;基于 ego 轨迹与视觉特征做去重,降采样直行高速等重复场景。
  • 用 NVIDIA SIL 的视频检索/编目平台 SIL-Wheel 承担三个角色:构建 finetuning 切片(上调罕见天气、施工区、VRU、复杂多 agent 交互等自然分布下欠采样的场景权重)、构建评测切片(支撑长尾/分场景报告)、数据质检(校验 auto-label 质量、triage 过滤管线筛出的 clip)。
  • HF model card 用标准 schema 给出一个粗粒度口径:约 400 万条多视角 20 秒驾驶录像,视频训练数据规模档位为”10,000–1,000,000 小时”——这与论文正文 16,600+4,944≈21,544 小时(呼应摘要中 “21k hours of driving scenarios”)的精确数字不完全一致,两者出处不同(论文正文精确统计 vs. HF 卡片标准化分档),如实并列记录。

训练方法

多阶段训练策略,从 Cosmos-Predict 2.5 出发:

  1. World-Scenario Control + Multi-view 适配(bidirectional 阶段):先在 RDS 上 mid-train 提升 AV 能力;再加 cross-view attention 层做多视角适配(view embedding 与 cross-view attention 输出投影均 zero-init 以保证收敛稳定),用 text-to-video : image-to-video = 1:1 混合训练;随后加入 zero-init 的 world-scenario control branch,先在 93 帧片段训练至收敛,再扩展到 189 帧片段学习更长时一致性,得到 bidirectional teacher。
  2. Diffusion Forcing mid-training(自回归适配):用 Flex-Attention 实现 causal masking,把 p(x_1:T)=Π p(x_i|x_<i) 参数化为 flow-matching 模型;先在不带 world-scenario control 的权重上、用 RDS 训练(经验发现这一步收益超过现有 world-scenario 标注数据的规模),再加入 control branch、用 RDS-HQ-1M 继续训练。
  3. Self Forcing 蒸馏(few-step,消除自回归误差累积):自回归 self-rollout,K=2 步扩散、timestep schedule [1000, 450];每次迭代只对一个随机采样的 denoising step 反传梯度,且 KV cache 梯度被 detach,只让梯度流经当前帧;用 rolling KV cache(固定窗口 L)把推理复杂度从 O(TL²) 降到 O(TL);损失是 DMD(Distribution Matching Distillation)holistic 视频级分布匹配目标,ℒ_DMD=E[½‖x̂-sg(x̂-f_ψ(x̂_t,t)-f_φ(x̂_t,t))‖²](f_φ 冻结的真实 score 网络,f_ψ 学习中的 fake score 网络)。
  4. Progressive teacher 策略:先训练一个更长时间上下文的 bidirectional 教师(无 rolling-cache 伪影),再用它继续蒸馏此前用短上下文教师蒸馏出的模型,显著减少 rolling KV cache 超出训练上下文窗口时产生的漂移伪影(见评测 Table 6)。蒸馏阶段用一个精选的 58k 视频高质量子集(聚焦复杂城市场景)。

World-Action Model(WAM)后训练(把 OmniDreams-SV backbone fine-tune 成 policy):backbone 结构不变,每帧新增两路输入——DINOv2(dinov2_vitb14)编码的 patch feature + 一个 30° 前视远焦相机的特征,经线性投影加到 noisy video latent token 上;以及一个 history token(小 MLP 嵌入过去 1.6 秒的 ego motion),与 video patch 沿时间轴交错插入,attention mask 保持原有 video-to-video causal 结构、并把 history token 设为单向耦合(只 attend 当前和过去 video token,不 attend 其他 history token)。用联合 flow-matching 目标同时训练 video 和 trajectory latent 的去噪(各自独立采样 flow-matching 时间步)。推理时丢弃交错 noise schedule,一次 causal forward 只跑一次 backbone,输出接一个 12 层 U-Net 形状 MLP、用 4 步 flow-matching 采样出未来 **6.4 秒(64 个 waypoint,10Hz)**轨迹。

Infra(训练 / 推理工程)

  • 训练算力:论文未披露 mid-training / post-training / 蒸馏各阶段的 GPU 数量或 GPU-hours(仅 GitHub post-training quickstart 提及在 8×H100 80GB 上验证过复现流程,这是示例复现环境,非原始训练算力披露)。
  • 推理(训练后免训练优化,Sec.5):2 步扩散蒸馏模型,704×1280 分辨率。
    • OmniDreams-SV(8 帧/chunk,local-window 6 个 latent 帧=24 RGB 帧):NVIDIA GB300 上 1×GPU 118ms/chunk(68 effective FPS);2×GPU 102ms(78 FPS);4×GPU 80ms(100 FPS);8×GPU 78ms(103 FPS)(Table 2,含 world-scenario 编码/Diffusion DiT/RGB 解码/KV-cache 更新四段耗时拆分)。
    • OmniDreams-MV(16 帧/chunk,4 视角联合,local-window 8 个 latent 帧=32 RGB 帧):1×GPU 1,289ms(12 FPS);4×GPU 330ms(48 FPS);8×GPU 209ms(74 FPS);16×GPU 151ms(105 effective FPS/相机)(Table 3)。16-GPU 配置用 hierarchical context-parallel:V(视角)=4、T(时间)=4、HW(空间)=1,rank 顺序 V→T→HW,自研 ring-attention 实现 context-parallel attention。
    • 优化手段:local-window attention;streaming static-shape KV cache(预分配定长张量,独立线程做 cache 更新以不阻塞主线程);torch.compile + CUDA Graph(首 chunk 懒捕获,后续 chunk 复用);用 LightX2V 的 LightVAE/LightTAE 替换原 Cosmos VAE 做轻量编解码(MV 用 pixel-shuffle 技术编码条件视频,延迟 <1ms);把 RoPE 频率计算、patchify/unpatchify 等与 step 无关的运算提到 denoising loop 外只算一次。
    • FlashDreams(已开源的通用流式推理/服务框架,打包上述优化,不改动模型权重):在 Wan2.1-based backbone 上验证——Self Forcing(Wan2.1 T2V)单张 GB200 上比官方实现快达 1.95×;Lingbot-World(Wan2.1-14B-based camera-control 世界模型)4×H100 上快达 2.49×。提供 gRPC + WebRTC server/client 协议做端到端流式服务。
    • HF model card 对已发布的 distilled checkpoint(2b_res720p_30fps_i2v_hdmap_distilled.pt)给出的数字是 “~52 FPS @720p”(1×GB300),与论文正文 Table 2 的 68 FPS(1×GPU)存在差异——两者出处不同(论文正文基准测试 vs. HF 卡片对特定发布 checkpoint 的表述),如实并列记录,不做调和。

评测 benchmark

训练阶段对比(Table 4,OmniDreams-SV,评测集为 RDS-HQ-1M held-out 5,000-clip 中采样的 1,000 条,用 off-the-shelf BEVFormer 做 3D 检测、LATR 做车道线检测作判别式指标):

训练阶段FVD↓Temporal Sampson↓LET-AP↑LET-APL↑LET-APH↑Lane F1↑x-err(far)↓Cat.Acc↑
Bidirectional(AV adapted)26.80.8231.830.3780.2400.3660.3370.957
Causal(Diffusion Forcing)31.70.7751.870.2210.1360.2140.4180.941
Distilled(Self Forcing)24.80.8281.900.4000.2550.3880.3130.961

最终蒸馏模型在 FVD 和条件保真度(3D 检测/车道线指标)上都最好,同时具备实时因果生成能力。

Decoder 权衡(Table 5):Distilled + 原始 VAE:FVD 24.8;Distilled + LightTAE(推理优化用): FVD 45.4——为换取推理速度接受生成质量下降。

长 rollout 分段 FVD(Table 6,20 秒 rollout 切 4 个 5 秒窗口):

训练教师0–5s↓5–10s↓10–15s↓15–20s↓Mean↓Δ(末窗-首窗)↓
Short-context teacher109.3183.0258.3409.2240.0299.9
Progressive long-context teacher95.5151.0202.5268.4179.4172.9

Progressive long-context teacher 大幅降低长 rollout 的漂移(Δ 从 299.9 降到 172.9)。

WAM 闭环对比(Sec.7.2,574-scene Physical AI Autonomous Vehicles NuRec 子集,10Hz replan,20 秒 rollout,排除 WAM 训练用过的场景):OmniDreams WAM(~2B 参数)vs Alpamayo 1.5(~10B 参数,VLA 基线)—— All Collision 6.9%→4.2%,Collision Front 1.0%→0.9%,Collision Lateral 0.6%→0.4%,Collision Rear 5.3%→3.0%,参数量约为 Alpamayo 1.5 的 1/5。

闭环仿真器保真度对比(Sec.9.4.1,501-scene subset,同一 NuRec 数据集但另取子集,只切换 AlpaSim 的 sensor simulator 为 NuRec 或 OmniDreams,4 类 policy 用同一 stack 评测,replan 节流到 OmniDreams 的 533ms chunk 频率,事故只在 ego 偏离 GT 轨迹 4m 内计入):headline All Incidents 指标下,四类 policy(OmniDreams WAM、Alpamayo 1.5 四相机、Alpamayo 1.5 双相机、Alpamayo 1.5 单相机)在切换 NuRec→OmniDreams 后排名保持不变(OmniDreams WAM 最强,依次到 Alpamayo 1.5 单相机最弱),论文以此论证 OmniDreams 是闭环策略评测的可靠代理(faithful proxy)。注:该图为柱状图,PDF 文本抽取顺序错乱导致各柱具体百分比无法可靠对应到各 policy,故此处只记录论文文字给出的排名结论,不编造逐项数值。

视觉真实感对比(Fig.14):按 rollout 轨迹与 GT 轨迹的偏离距离分 bin 计算 FVD,NuRec 随偏离增大迅速退化,OmniDreams 保持稳定视觉质量——该图仅有可视化结果,论文未给出具体 FVD 数值表。

创新点与影响

  • 动作条件化 + 轻量 control branch:用小 MLP 把 world-scenario map 编码为 control token 直接拼接(而非 ControlNet 式独立网络),开销小、训练稳定。
  • 可扩展的多视角一致性设计:cross-view attention 把多视角联合生成复杂度从 O(N²T²) 降到 O(NT²)+O(N²),支持最多 7 个同步视角。
  • Self Forcing + DMD + progressive long-context teacher:系统性解决自回归视频生成长 rollout 的 exposure bias 和漂移伪影,是把闭环仿真所需的”分钟级稳定 rollout”落地的关键工程组合。
  • World-Action Model(WAM):论文自称是较早展示”世界模型可直接后训练成优于专用 VLA 的驾驶 policy”的工作之一(呼应 Ye et al. 2026 在机器人操作领域的类似发现),用约 1/5 参数量在碰撞率上超过 Alpamayo 1.5。
  • FlashDreams:把 streaming KV cache、局部窗口 attention、CUDA Graph 等优化打包为通用开源推理框架,并在非本模型 backbone(Wan2.1 系)上验证了可迁移性。
  • 作者自陈的局限:(1)生成式仿真器计算需求远高于重建式仿真器,存在质量-算力权衡;(2)当前 chunk-based 生成(533ms/chunk)要求 policy/traffic 在 chunk 边界内的轨迹提前确定,无法逐帧级交互,论文把”缩小 chunk size、最终逐帧生成”列为未来方向;(3)世界模型与 policy 的联合训练(同一模型既 render 又 drive)论文只做了初步验证(WAM 独立后训练),联合训练本身留作未来工作。

原始链接

一手源存档(sources/)

关系图谱

反向链接