Epona: Autoregressive Diffusion World Model for Autonomous Driving
Paper: arXiv:2506.24113 Code: Kevin-thu/Epona Code reference:
main@69b24c55(2025-07-22) Project / Checkpoints: project page · Hugging Face
1. Motivation (研究动机)
Epona 关注的是自动驾驶里的生成式世界模型:给定历史前视相机画面和自车轨迹,模型不仅要“看懂过去”,还要能连续想象未来道路、车辆、建筑和交通规则,并且把这种想象能力转成可用的运动规划。论文把这个问题放在 Physical AI 语境下理解:传统自动驾驶管线通常把感知、预测、规划拆成多个监督模块,需要 3D box、车道线、地图、占用等昂贵标注;世界模型则试图直接从大规模驾驶视频中自监督地学习场景动力学,让 agent 在内部模拟未来,再据此规划。
现有生成式驾驶世界模型主要有两类弱点。第一类是视频扩散模型,如 Vista、DriveDreamer 一类,它们视觉质量强,但通常把一段固定长度视频作为全局联合分布来建模。这个设定天然适合一次性生成 8 帧、16 帧或 25 帧,却不适合在测试时自由延长到几十秒甚至分钟级;如果强行 roll out,模型训练时看到的是 ground-truth 条件,推理时看到的是自己生成的历史,误差会逐步积累,导致画面漂移、道路结构破坏、车辆细节崩坏。更关键的是,全局固定窗口扩散不容易把“规划一个未来轨迹”作为可独立调用的模块接入,因为视觉生成和动作生成都被包在同一个固定长度视频分布里。
第二类是 GPT-style 自回归世界模型,它们很自然地支持逐步预测,但通常依赖离散视觉 token。离散化对驾驶场景尤其敏感:小车、车道线、红绿灯、路缘、建筑边缘都可能在 tokenization 中丢细节;而规划任务又恰恰需要这些几何和语义细节。论文的核心判断是:自动驾驶世界模型既需要自回归的因果时间结构,又需要扩散/flow 在连续视觉 latent 上的细节保真能力。因此 Epona 不是简单把视频扩散模型 roll out,也不是把视频压成离散 token 后做 next-token prediction,而是把“时间因果建模”和“空间细节生成”拆开。
Figure 1 解读:teaser 把 Epona 的四个目标放在一起:A 是分钟级高分辨率驾驶视频生成,B 是给定不同轨迹时进行可控世界模拟,C 是从视频自监督中学到类似“红灯停车”的交通知识,D 是只调用轨迹分支时作为实时端到端 planner。这个图也说明本文不是单纯追求 FVD,而是把世界生成、动作控制和规划放进同一个模型。
论文要解决的实际矛盾是:如果模型只追求视觉生成,就很容易变成一个离线 driving video generator;如果只追求 planning,就又会退化成 camera-to-trajectory 的判别式 planner,丢掉世界模型能模拟未来的优势。Epona 的研究动机正是在两者之间建立一个共享历史表征:同一个历史 latent 既可以驱动 next-frame diffusion 生成未来视觉,也可以驱动 trajectory diffusion 生成未来 3 秒轨迹;推理时二者可以联合使用,也可以单独关闭视觉分支只做规划。
2. Idea (核心思想)
Epona 的核心思想可以概括为“自回归的时间骨架 + 连续 latent 上的扩散细节 + 可异步调用的轨迹/视觉双分支”。它首先把驾驶世界建模从固定窗口的视频联合分布改写成逐时间步的局部分布:给定过去 帧观测 和历史自车动作 ,模型预测未来轨迹策略 以及在动作条件下的下一帧视觉分布 每个动作 被表示为 ,即朝向变化和相对位移。这样一来,Epona 的视觉生成不是无条件地“续写视频”,而是可以显式受轨迹控制;轨迹生成也不是后处理,而是和视觉分支共享历史因果表征。
Figure 2 解读:method comparison 图对应本文的关键立场。传统视频扩散建模固定长度序列的 joint distribution,长时预测时只能重复 roll out;GPT-style 方法有因果顺序,但通常牺牲连续视觉细节;Epona 将历史时间动态压到 compact latent,再用两个 diffusion transformer 生成下一帧 latent 和未来轨迹,既保留自回归结构,也避免离散 token 带来的视觉损失。
具体设计上,论文提出三个互相配合的模块化想法。第一是 Decoupled Spatiotemporal Factorization:用 Multimodal Spatiotemporal Transformer(MST)专门处理历史帧与历史轨迹的时间因果关系,把它们编码成 compact latent ;空间细节不在 MST 里直接生成,而交给后续 VisDiT。这样,时间建模和未来画面细节生成不再被一个大视频扩散窗口绑死。
第二是 Asynchronous Multi-modal Generation:轨迹规划 TrajDiT 和视觉预测 VisDiT 共享 ,但各自以 diffusion/rectified-flow 方式独立去噪。推理时可以先采样轨迹,再把轨迹第一个动作作为 VisDiT 的 action modulation;也可以给用户指定轨迹来做 controllable simulation;还可以完全跳过 VisDiT,只运行 MST+TrajDiT 做实时 planning。
第三是 Continuous Visual Tokenization + Temporal-aware DCAE:Epona 用 DCAE 将图像压成连续 latent,避免离散 token 量化导致的小目标和几何结构损失;同时在 DCAE decoder 前加入时序交互模块,减少逐帧解码带来的 flickering。
Figure 3 解读:pipeline 图给出了完整数据流。前 帧经 DCAE encoder 变成 latent,历史轨迹和 latent 一起进入 MST;MST 输出 compact latent 后,TrajDiT 预测未来 帧 pose trajectory,VisDiT 生成 帧 latent,再由 temporal-aware DCAE decoder 还原图像。Chain-of-Forward 训练在这个闭环上模拟推理时的自生成历史,专门缓解 autoregressive drift。
这个想法的优点是接口清晰:世界模型不再是一个单一的 video generator,而是一个可组合系统。若目标是生成长期视频,循环执行 MST → TrajDiT → VisDiT → DCAE decoder;若目标是 motion planning,执行 MST → TrajDiT 即可,论文报告 10-step 采样时 MST+TrajDiT 约 ,对应接近 20 Hz 的 planner 速率;若目标是轨迹可控仿真,则把外部给定轨迹替代 TrajDiT 预测输入 VisDiT。
3. Method (方法)
3.1 Rectified flow 与自回归世界建模
论文的生成目标建立在 diffusion/rectified flow 上。一般 diffusion 可以写成从噪声分布 到数据分布 的微分方程 Rectified flow 选择一条直线路径,把数据 和噪声 混合成 并优化速度场 Epona 将这个 flow 目标同时用于轨迹和视觉 latent。给定历史图像编码 与历史相对动作,MST 得到历史上下文 。随后 TrajDiT 在未来轨迹空间采样,VisDiT 在下一帧 latent 空间采样。这个“下一步局部分布”被循环使用,于是模型天然支持可变长度预测:生成一帧后,将预测结果并入历史窗口,再预测下一帧。
3.2 MST:把历史视觉和动作压成因果 compact latent
MST 的输入是连续视觉 latent 和离散化后的 pose/yaw token。代码中 models/stt.py 的 SpatialTemporalTransformer 使用空间块与因果时间块:SpaceBlock/CausalSpaceBlock 在同一时间步内处理 pose token、yaw token 与图像 token 的空间关系,CausalTimeBlock/CausalTimeSpaceBlock 保证时间维只能看过去。这个结构对应论文中的“GPT-style transformer with causal attention handles temporal dynamics in compressed latent space”。MST 输出的不是未来图像,而是一个条件表示 ;这让后续的轨迹 diffusion 和视觉 diffusion 可以共享同一历史理解,同时各自保留专门结构。
3.3 TrajDiT:未来 3 秒轨迹的 diffusion policy
TrajDiT 预测未来轨迹 ,其中每个点包含相对 位移和 yaw。论文给出的轨迹损失为
代码中 models/traj_dit.py 的 training_losses 会采样噪声,构造 x_t = t * traj + (1 - t) * noise,目标是 traj - noise,再用 MSE 训练预测速度。这里和论文 Eq. (2) 的 只是时间方向约定不同:代码把 越大越靠近数据,目标速度仍是 data-noise。sample 函数按时间步循环,用预测速度更新轨迹;models/model.py 的 step_eval 会先调用 MST 得到 stt_features,再调用 self.traj_dit.sample(...) 得到 predict_traj,最后 denormalize 成真实尺度。
3.4 VisDiT:动作条件下的下一帧 latent diffusion
VisDiT 和 TrajDiT 结构相似,但额外有 action modulation 分支,条件是 以及未来动作 。论文中的视觉损失为
总损失是
代码对应 models/flux_dit.py 的 FluxDiT.training_losses:它对目标 image latent 加噪,预测 img - noise,再取 mean_flat((target - pred) ** 2)。models/model.py 的 model_forward 同时调用 self.traj_dit.training_losses(...) 与 self.dit.training_losses(...),返回 loss_all、loss_diff 和 loss_yaw_pose。在 step_eval 中,如果 self_pred_traj=True,模型把 TrajDiT 预测轨迹的第一步转成 pose/yaw embedding,再送给 VisDiT;如果用户提供轨迹,则 VisDiT 可以直接按外部 action 控制视频。
Figure 4 解读:supplementary 的 DiT 结构说明了 TrajDiT/VisDiT 为什么适合条件生成。dual-stream 阶段让条件 latent 和待去噪样本先分流处理,只通过 attention 交互;single-stream 阶段再拼接融合。这种设计继承了 FLUX/HunyuanVideo 一类 DiT 的条件生成范式,区别在于这里的条件不是文本,而是驾驶历史和动作。
3.5 Chain-of-Forward:训练时暴露推理误差
自回归模型的典型问题是 teacher forcing gap:训练时下一步条件来自 ground truth,推理时条件来自模型自己生成的历史。Epona 的 Chain-of-Forward 每隔 10 个训练 step 触发一次,每次执行 3 次 forward pass,用模型估计的未来 latent/trajectory 替代一部分 ground-truth 条件,从而让训练分布更接近推理分布。为了避免真正从纯噪声采样带来的巨大开销,论文使用一阶估计
其中 可以是 image latent 或 trajectory。这个估计值被塞回历史条件,继续后面的训练 forward。代码 config 中 forward_iter=3、multifw_perstep=10 与论文“每 10 步,3 个 forward passes”一致;models/model.py 的 step_train 会把 condition latents 与 target latents 拼成 latents_total,再把 targets = torch.cat([latents, latents_gt], dim=1)[:, 1:] 作为下一步监督。
Figure 5 解读:training 图强调 Chain-of-Forward 不是额外模型,而是训练调度。它在原本 rectified-flow loss 外,把模型自己的去噪预测放入下一步条件,使 MST/DiT 在训练阶段就见到“有噪声的自生成历史”。这对长视频尤其关键,因为 600 帧 roll-out 中任何局部错误都会被后续条件反复放大。
Figure 6 解读:有/无 Chain-of-Forward 的 qualitative 对比显示,长时 roll-out 里无 CoF 版本更容易出现场景结构漂移和视觉质量下降;有 CoF 版本能在更长时间保持车道、建筑和车辆轮廓。这个结果支持论文对 teacher-forcing gap 的诊断。
3.6 Temporal-aware DCAE:在连续 latent 压缩与时序一致性之间折中
Epona 使用 DCAE 以 下采样把 图像压到紧凑 latent。原始 DCAE 是图像 autoencoder,不显式建模相邻帧关系,因此逐帧解码时可能 flicker。论文的改动是在 decoder 前插入 spatiotemporal self-attention 层,encoder 固定,只 fine-tune decoder 前的 temporal interaction 模块。这样做保留了 DCAE 的压缩效率,又让最终还原视频具备跨帧交互。
Figure 7 解读:qualitative 图对比 Epona 与 Vista 在长时生成中的结构保持。论文指出 Vista 是固定 25 帧扩散模型,长视频需要滚动生成;Epona 的自回归 flow 结构直接面向下一帧局部分布,因此在道路边界、建筑、车辆和整体视角上更稳定。
3.7 代码锚点与实现映射
Code reference for mapping:
main@69b24c55(2025-07-22). 下面的实现映射只对应这个公开仓库状态;论文中 48 A100、600k iteration 的完整训练 launch 未在仓库中以一键脚本完全复现,仓库 README 给的是训练/finetune 示例命令。
- 总体封装:
models/model.py::TrainTransformersDiT。初始化时创建SpatialTemporalTransformer、FluxDiT和TrajDiT;model_forward负责构造相对 pose/yaw token、未来轨迹 target、调用两个 diffusion loss,并返回loss_all = diff_loss + yaw_pose_loss。 - 历史表征:
models/stt.py::SpatialTemporalTransformer。包含空间 attention、因果时间 attention、pose/yaw embedding 和evaluate接口,对应论文 MST。 - 视觉分支:
models/flux_dit.py::FluxDiT。training_losses实现 image latent 的 rectified-flow MSE;sample按离散 timestep 迭代更新 latent。 - 轨迹分支:
models/traj_dit.py::TrajDiT。training_losses和sample与 VisDiT 对称,但输入/输出是轨迹 token;配置中traj_len=15,n_layer_traj=[1,1],体现 tiny planner。 - 训练入口和配置:
scripts/train_deepspeed.py、configs/dit_config_dcae_nuplan.py、configs/dit_config_dcae_nuscenes.py。公开 config 中image_size=(512,1024)、condition_frames=10、diffusion_model_type="flow"、num_sampling_steps=100、forward_iter=3、multifw_perstep=10与论文实现细节相吻合。
伪代码化后,Epona 的推理循环如下:
history_latents, history_poses = DCAE.encode(history_frames), history_actions
for each future step:
F, pose_context = MST.evaluate(history_latents, history_poses)
future_traj = TrajDiT.sample(noise_traj, cond=F)
if user_control:
action = user_action[step]
else:
action = future_traj[0]
next_latent = VisDiT.sample(noise_img, cond=F, action=action)
next_frame = TemporalAwareDCAE.decode(next_latent)
append(next_latent, action) to history训练阶段则是联合 flow loss:
latents_total = concat(condition_latents, target_latents)
pose_tokens, yaw_tokens = relative_pose_tokens(rot_matrix)
traj_targets = relative_future_trajectory(rot_matrix)
F = MST(latents_total, pose_tokens, yaw_tokens)
L_traj = TrajDiT.flow_loss(traj_targets, cond=F)
L_vis = VisDiT.flow_loss(next_latents, cond=F, action=target_action)
if step triggers chain-of-forward:
replace part of future context with one-step denoised self predictions
optimize L = L_traj + L_vis4. Experimental Setup (实验设置)
论文实验覆盖三类能力:视频世界生成、轨迹可控仿真、端到端 motion planning。世界模型总规模为 2.5B 参数:12-layer MST 约 1.3B,12-layer next-frame prediction diffusion transformer 约 1.2B,2-layer trajectory planning diffusion transformer 约 50M。训练数据来自 NuPlan 公开视频和 NuScenes 的 700 个 scenes,所有图像 resize 到 。训练使用 rectified flow 同时监督视频生成与轨迹规划,48 张 NVIDIA A100 训练接近两周,总计 600k iterations,batch size 96;AdamW 学习率 ,weight decay 。Chain-of-Forward 每 10 step 应用一次,每次 3 个 forward pass。推理时所有实验设置 DiT sampling steps 为 100。
公开仓库中的 config 说明了若干实现细节:NuPlan config 使用 downsample_fps=5,NuScenes config 使用 downsample_fps=10;两者都设 condition_frames=10、image_size=(512,1024)、vae_embed_dim=32、downsample_size=32、n_layer=[12,6,6]、n_embd=2048、num_sampling_steps=100、traj_len=15、n_embd_dit_traj=1024。NuPlan config 中 add_decoder_temporal=True、temporal_patch_size=6,与 temporal-aware DCAE 的论文动机一致;NuScenes config 中该开关为 false,说明仓库配置和论文完整训练设置需要按 checkpoint/experiment 解释,不能把 README 示例命令误当作论文主实验 launch。
评估指标方面,视频生成使用 FVD 和 FID;长时生成还报告“能保持 plausible 结果的最大时长/帧数”。规划在 NuScenes 上用 L2 displacement 和 collision rate,按 1s、2s、3s 以及平均值报告;在 NAVSIM 上使用 NC、DAC、TTC、Comfort、EP 和 PDMS。NAVSIM 对 Epona 很关键,因为它检验模型是否能从历史前视图中直接输出可驾驶轨迹,而不仅是生成好看的视频。
推理速度在单张 NVIDIA 4090 上测量,生成 3 秒轨迹和一张 图像:10-step 采样时 MST 约 0.02s、TrajDiT 约 0.03s、VisDiT 约 0.3s;100-step 采样时 MST 仍约 0.02s、TrajDiT 约 0.3s、VisDiT 约 2s。这解释了为什么论文强调“real-time motion planner”:只调用 MST+TrajDiT 时,低步数采样开销约 0.05s;而完整世界视频生成则主要受 VisDiT 采样成本限制。
Figure 8 解读:trajectory-controlled video generation 展示了同一历史上下文下,不同输入轨迹会导致不同未来视角和道路运动。这证明视觉分支不是无动作条件的视频续写,而确实接受 trajectory/action 控制,因此可以作为可控驾驶模拟器使用。
Figure 9 解读:supplementary 的 long-video 可视化给出 140 秒生成示例。注意这和主表中的 120s/600 frames 最大时长指标互相补充:表格给出和 baselines 的规范比较,长视频图则展示分钟级 roll-out 里的道路和建筑是否还保留局部一致性。
5. Experimental Results (实验结果)
5.1 视频生成:FVD 最优且可扩展到两分钟以上
在 NuScenes validation set 上,Epona 的核心定量结果是 FVD 82.8、FID 7.5、最大时长 120s/600 frames。和最强固定长度 baseline Vista 相比,Vista 的 FID 更低一点(6.9),但 FVD 为 89.4,最大时长只有 15s/150 frames;Epona 的 FVD 相对 Vista 下降约 7.4%,同时将可 plausible 生成时长扩展到 120s。和 DrivingWorld 相比,DrivingWorld 为 FID 7.4、FVD 90.9、40s/400 frames,Epona 在 FVD 和时长上都更强。更早的 DriveGAN、DriveDreamer、WoVoGen、Drive-WM、GenAD 在 FVD 或最大时长上明显落后,其中 Drive-WM 虽 FVD 122.7,但只有 8s/16 frames。
这个结果需要分开看。FID 是单帧图像质量,Vista 的 FID 6.9 略优于 Epona 的 7.5,说明 Epona 并不是在所有视觉指标上绝对领先;但 FVD 更重视视频时序分布,Epona 82.8 最优,且长时生成能力大幅领先。对世界模型来说,这比单帧 FID 更重要,因为模型最终要在未来时间轴上稳定滚动,而不是生成一个短窗口。
5.2 Chain-of-Forward 与 temporal-aware DCAE 的消融
Figure 10 解读:CoF FID 曲线展示的是长时滚动中误差积累的速度。随着自回归步数增加,无 Chain-of-Forward 的模型在 FID 上更快变差;有 CoF 的模型把训练阶段暴露给自生成历史,因此在更长 horizon 下保持更稳定。这和 Figure 6 的 qualitative 漂移现象一致。
Temporal-aware DCAE 的消融更直接。无 temporal module 时,NuPlan 上 FVD、FVD、FVD 分别为 52.95、76.46、100.11;加入 temporal-aware decoder 后,结果变为 50.77、61.46、74.88。越长的 40-frame 评估收益越大,说明该模块主要缓解时序 flickering 和长时累积不一致,而不只是改善短帧视觉质量。
条件帧数量也验证了历史上下文的重要性。使用 2 个 condition frames 时,FVD、FVD、FVD 为 59.85、81.58、103.70;5 帧时为 55.46、71.28、86.76;10 帧时为 50.77、61.46、74.88。这个趋势说明 MST 需要足够历史来估计道路几何和自车运动,尤其在 25/40 帧长 horizon 下更明显。
5.3 Motion planning:世界模型分支能超过强端到端 planner,但不是所有指标都最佳
在 NAVSIM 上,Epona 使用 camera 输入,NC 97.9、DAC 95.1、TTC 93.8、Comfort 99.9、EP 80.4、PDMS 86.2。PDMS 高于 DRAMA 的 85.5、PARA-Drive/LAW 的 84.6、TransFuser 的 84.0 和 UniAD 的 83.4。与 DrivingGPT 的 supplementary 对比中,DrivingGPT PDMS 82.4,而 Epona 86.2;DrivingGPT 的 NC 98.9 和 TTC 94.9 更高,但 Epona 在 DAC、Comfort、EP 和最终 PDMS 上更强。论文还比较了 joint training 的作用:没有 joint training 时 PDMS 78.1,完整 Epona 达到 86.2,说明视觉世界建模和轨迹规划共享训练确实提升 planner,而不是两个任务互相干扰。
在 NuScenes planning 上,Epona 的结果更需要谨慎解读。它只使用 front camera 且无 auxiliary supervision,L2 为 0.61/1.17/1.98,平均 1.25;collision rate 为 0.01/0.22/0.85,平均 0.36。它的 collision 平均低于 GenAD 的 0.43,也低于 VAD-Base 的 0.53 和 ST-P3/OccNet/VAD-Tiny 的 0.71/0.72/0.72;但 L2 平均不如 GenAD 的 0.91,也不如 UniAD 的 1.03。换言之,Epona 更像是从世界模型中学到保守且符合交通规则的规划,碰撞低,但未必在几何 L2 上绝对最优。这一点对论文结论很重要:Epona 证明了生成式世界模型可作为 planner,但还没有完全取代专门规划模型。
5.4 交通知识与可控仿真的证据
论文强调 Epona 从视频自监督中学到“红灯停车”等真实交通知识。这里的证据主要来自 teaser 和可控生成可视化,而不是一个大规模交通规则 benchmark:同一历史上下文下,给不同 future trajectory,模型会生成相应的道路视角变化;在 traffic-light 场景中,模型能表现停止行为。作为读者需要区分两层结论:一是模型具备一定 implicit traffic-world knowledge,这是可视化和 planning 指标共同支持的;二是它是否可安全部署到真实闭环驾驶,还需要更严格的交互式仿真、罕见事件评测和真实车测试,论文没有完成这一层证明。
5.5 局限与可复现性判断
Epona 的主要局限来自四方面。第一,训练成本很高:2.5B 参数、48 A100、近两周、600k iterations,这使得复现门槛显著高于普通 camera-only planner。第二,公开仓库提供了推理、训练/finetune 脚本、config 和 checkpoint 下载,但没有把论文完整 48-GPU 主实验 launch 以完全可复现脚本形式封装;因此笔记中的实现细节应以论文表述和公开 config 共同对齐,不能把 README 的 4 nodes × 8 GPUs、batch size 2、lr 2e-5 示例当作主实验。第三,实验主要基于前视 camera 和公开数据集,虽有 NuScenes/NAVSIM,但仍缺少闭环真实部署、corner cases、恶劣天气和多传感器设置下的系统性验证。第四,planning 的 L2 指标并非全面最优,说明生成式世界模型学到的轨迹分布仍可能偏保守或和 expert 轨迹有偏差。
整体上,Epona 的贡献不在于提出一个更大的视频扩散模型,而在于重新定义 driving world model 的因果接口:用 MST 负责历史因果压缩,用 TrajDiT 负责可实时调用的 trajectory policy,用 VisDiT 负责连续 latent 空间的下一帧视觉生成,再用 Chain-of-Forward 把自回归误差提前放进训练。这个设计使它同时覆盖 long-horizon video generation、trajectory-controlled simulation 和 motion planning 三个任务;从结果看,最强证据是 NuScenes FVD 82.8 与 120s/600-frame 时长、NAVSIM PDMS 86.2、以及 CoF/DCAE/condition-frame 消融的方向一致性。