源码快速跳转：本文档中标注了
Ctrl+P快捷跳转路径。 在 VS Code 中按Ctrl+P，粘贴路径即可打开文件；带:行号的会直接定位到那一行。
| Ctrl+P 输入 | 说明 |
|---|---|
diffusion_planning_head.py |
迁移后的 DiffusionDrive 规划头（954行） |
e2e_diffusiondrive.py |
训练配置文件 |
navformer.py |
NavFormer 主模型（调用 planning head） |
traj_scoring_head.py |
已有的 HydraMDP/VADv2 规划头（参考） |
dense_heads/__init__.py |
模块注册 |
transfuser_model_v2.py |
原始 DiffusionDrive 模型 |
blocks.py |
原始工具函数 |
multimodal_loss.py |
原始损失函数 |
- 项目总览与架构
- Python 基础语法速查
- 原始 DiffusionDrive 代码解析
- WorldEngine NavFormer 框架解析
- 迁移后的 DiffusionPlanningHead 逐行解析
- 配置文件解析
- 数据流与调用链全流程
- 训练与推理流程对比
DiffusionDrive 是一篇自动驾驶论文提出的轨迹规划算法。它的核心思想是：
用扩散模型（Diffusion Model） 来生成自动驾驶车辆的未来行驶轨迹。
传统方法直接回归一条轨迹，而 DiffusionDrive 的做法是：
- 准备 20 条"锚点轨迹"（从训练数据中 K-means 聚类得到的典型轨迹模式）
- 对这些锚点加噪声（前向扩散）
- 用一个 DiT（Diffusion Transformer）解码器逐步去噪
- 最终输出 20 条候选轨迹 + 每条的置信度分数
- 选置信度最高的那条作为最终规划
WorldEngine 是一个自动驾驶后训练框架，它的核心模型叫 NavFormer。NavFormer 的结构是：
相机图像 → ResNet 骨干网络 → BEVFormer（生成鸟瞰图特征）→ Planning Head（规划头）→ 输出轨迹
WorldEngine 已经集成了两种规划头：
- HydraMDP：基于打分的方法（从 8192 条预计算轨迹中选最优）
- VADv2：类似的打分方法
现在我们要集成第三种：DiffusionDrive（基于扩散生成的方法）。
| 挑战 | 解决方案 |
|---|---|
| DiffusionDrive 有自己的骨干网络（ResNet34+LiDAR），WorldEngine 用 BEVFormer | 只迁移"规划头"部分，BEV 特征由 NavFormer 提供 |
| DiffusionDrive 的代码分散在多个文件 | 全部 inline 到一个文件 diffusion_planning_head.py |
BEV 特征形状不同（NavFormer: (H*W, B, C) vs DiffusionDrive: (B, C, H, W)） |
在 head 入口做 reshape |
| DiffusionDrive 没有 8192-vocab 的概念 | 添加 vocab 快照锚定，兼容 NavFormer 的评估流程 |
| 输出格式不同（DiffusionDrive: 8步, NavFormer 期望: 40步） | 用 _expand_to_40 把 8 步重复为 40 步 |
WorldEngine-main/
├── projects/AlgEngine/
│ ├── configs/navformer/
│ │ ├── e2e_diffusiondrive.py ← 新增：DiffusionDrive 的训练配置
│ │ ├── e2e_hydramdp.py ← 已有：HydraMDP 配置（参考）
│ │ └── e2e_vadv2.py ← 已有：VADv2 配置（参考）
│ └── mmdet3d_plugin/navformer/
│ ├── detectors/
│ │ └── navformer.py ← 主模型，调用 planning_head
│ └── dense_heads/
│ ├── __init__.py ← 注册导出（新增了 DiffusionPlanningHead）
│ ├── traj_scoring_head.py ← 已有：HydraMDP/VADv2 的规划头
│ └── diffusion_planning_head.py ← 新增：DiffusionDrive 规划头（954行）
│
├── DiffusionDrive-main/ ← 原始 DiffusionDrive 源码（参考用，最终可删除）
│ └── DiffusionDrive-main/navsim/agents/diffusiondrive/
│ ├── transfuser_model_v2.py ← 原始模型（559行）
│ ├── transfuser_config.py ← 原始配置
│ └── modules/
│ ├── blocks.py ← 原始工具函数
│ ├── multimodal_loss.py ← 原始损失函数
│ └── scheduler.py ← 学习率调度器（未迁移，WorldEngine 有自己的）
这一节帮你理解代码中频繁出现的 Python/PyTorch 语法。
class DiffusionPlanningHead(nn.Module):
def __init__(self, num_poses=8, d_model=256, ...):
super().__init__() # 调用父类 nn.Module 的初始化
self.d_model = d_model # 把参数存为实例属性
self.linear = nn.Linear(256, 256) # 创建一个可训练的线性层nn.Module是 PyTorch 所有神经网络模块的基类__init__里定义网络结构（哪些层、哪些参数）forward里定义前向传播（数据怎么流过这些层）super().__init__()必须调用，否则 PyTorch 无法追踪参数
@HEADS.register_module()
class DiffusionPlanningHead(nn.Module):
...@HEADS.register_module() 是 mmdetection 框架的注册机制。它的作用是：
- 把这个类注册到一个全局字典里
- 之后在配置文件中写
type="DiffusionPlanningHead"就能自动找到并实例化这个类
# 形状变换
bev_embed.permute(1, 2, 0) # 交换维度顺序：(H*W, B, C) → (B, C, H*W)
.contiguous() # 确保内存连续（某些操作要求）
.view(B, C, H, W) # 重塑形状（不改变数据，只改变"看法"）
# 索引与切片
tensor[:, 0] # 取第2维的第0个元素
tensor[..., :2] # 取最后一维的前2个元素（... 表示"前面所有维度"）
tensor[4::5] # 从第4个开始，每隔5个取一个
# 拼接与拆分
torch.cat([a, b], dim=-1) # 沿最后一维拼接
tensor.split([1, 30], dim=1) # 沿第1维拆成大小为1和30的两块# nn.Parameter: 一个可训练的张量（或设 requires_grad=False 变成常量）
self.plan_anchor = nn.Parameter(torch.from_numpy(data), requires_grad=False)
# nn.Linear(in, out): 全连接层，y = xW^T + b
self.proj = nn.Linear(256, 256)
# nn.Embedding(num_items, dim): 查找表，输入整数索引，输出对应向量
self._query_embedding = nn.Embedding(31, 256) # 31个可学习的256维向量# nn.Sequential: 按顺序执行的层（输入自动传递给下一层）
self.mlp = nn.Sequential(
nn.Linear(256, 512),
nn.ReLU(),
nn.Linear(512, 256),
)
# 使用：output = self.mlp(input)
# nn.ModuleList: 层的列表（需要手动循环调用）
self.layers = nn.ModuleList([MyLayer() for _ in range(3)])
# 使用：for layer in self.layers: x = layer(x)原始代码位于 DiffusionDrive-main/DiffusionDrive-main/navsim/agents/diffusiondrive/。
这是一个 Python dataclass，存储所有超参数：
@dataclass
class TransfuserConfig:
# 感知范围
lidar_max_x: float = 32.0 # BEV 地图 x 方向半径（米）
lidar_max_y: float = 32.0 # BEV 地图 y 方向半径（米）
pixels_per_meter: float = 4.0 # 每米对应多少像素
# Transformer 参数
tf_d_model: int = 256 # 特征维度
tf_d_ffn: int = 1024 # FFN 中间层维度
tf_num_layers: int = 3 # Transformer 解码器层数
tf_num_head: int = 8 # 多头注意力的头数
tf_dropout: float = 0.0 # Dropout 比率
# 检测
num_bounding_boxes: int = 30 # 最多检测 30 个周围物体
# 轨迹规划
ego_fut_ts: int = 8 # 预测未来 8 个时间步
ego_fut_mode: int = 20 # 20 条候选轨迹（锚点数）
# 损失权重
trajectory_cls_weight: float = 10.0 # 分类损失权重
trajectory_reg_weight: float = 8.0 # 回归损失权重关键理解：
lidar_max_x/y = 32表示 BEV 地图覆盖车辆周围 ±32 米的范围ego_fut_ts = 8表示预测未来 8 个时间步（每步 0.5 秒，共 4 秒）ego_fut_mode = 20表示同时预测 20 条可能的轨迹
def linear_relu_ln(embed_dims, in_loops, out_loops, input_dims=None):
"""
构建 [Linear → ReLU → ... → LayerNorm] 的层序列。
参数：
embed_dims: 输出维度（也是中间层维度）
in_loops: 每个 LayerNorm 前有几个 Linear+ReLU
out_loops: 重复几组（每组以 LayerNorm 结尾）
input_dims: 第一个 Linear 的输入维度（默认等于 embed_dims）
例子：linear_relu_ln(256, 1, 1, input_dims=512)
生成：[Linear(512→256), ReLU, LayerNorm(256)]
"""def gen_sineembed_for_position(pos_tensor, hidden_dim=256):
"""
把 2D 坐标 (x, y) 编码成高维向量。
原理：来自 DAB-DETR 论文，用不同频率的 sin/cos 函数
把一个标量位置编码成一个向量，让网络能区分不同位置。
输入：pos_tensor 形状 (..., 2)，最后一维是 (x, y)
输出：形状 (..., hidden_dim)
过程：
1. x 和 y 各自乘以 2π
2. 除以不同频率的基底 [10000^(0/128), 10000^(2/128), ...]
3. 对结果取 sin 和 cos
4. 拼接 x 的编码和 y 的编码
"""这个函数的直觉：就像把一个数字转换成"二进制"表示，不同位的变化频率不同， 低位变化快（高频），高位变化慢（低频），这样网络就能精确区分不同位置。
这是 DiffusionDrive 最核心的注意力机制之一：
class GridSampleCrossBEVAttention(nn.Module):
"""
功能：让轨迹上的每个点"看到"BEV地图上对应位置的特征。
原理：
1. 轨迹有 8 个路径点，每个点有 (x, y) 坐标
2. 把这些坐标归一化到 [-1, 1]（grid_sample 要求）
3. 用 F.grid_sample 在 BEV 特征图上双线性插值采样
4. 用注意力权重加权求和
5. 加上残差连接
直觉：轨迹上的每个点去"看"它经过的地方有什么（车道线？障碍物？）
"""
def __init__(self, embed_dims, num_heads, num_points=8,
bev_range_x=51.2, bev_range_y=51.2, in_bev_dims=256):
# bev_range_x/y: BEV 地图覆盖的物理范围（米）
# num_points: 轨迹上的采样点数（= 8 个时间步）
# in_bev_dims: BEV 特征图的通道数
self.attention_weights = nn.Linear(embed_dims, num_points)
# 输入 query 特征，输出每个采样点的注意力权重
self.value_proj = nn.Sequential(
nn.Conv2d(in_bev_dims, 256, kernel_size=3, padding=1),
nn.ReLU(),
)
# 对 BEV 特征做一次卷积投影
def forward(self, queries, traj_points, bev_feature, spatial_shape):
"""
queries: (B, 20, 256) — 20条候选轨迹的特征
traj_points: (B, 20, 8, 2) — 20条轨迹×8个点×(x,y)坐标
bev_feature: (B, 256, 200, 200) — BEV 特征图
"""
# 1. 坐标归一化：物理坐标 → [-1, 1]
normalized[..., 0] = traj_points[..., 0] / bev_range_y # y 方向
normalized[..., 1] = traj_points[..., 1] / bev_range_x # x 方向
normalized = normalized[..., [1, 0]] # 交换 x,y（grid_sample 要求 (x,y) 顺序）
# 2. 计算注意力权重
attention_weights = self.attention_weights(queries).softmax(-1)
# 形状：(B, 20, 8)，每条轨迹的 8 个点各有一个权重
# 3. 在 BEV 上采样
value = self.value_proj(bev_feature) # (B, 256, 200, 200)
sampled = F.grid_sample(value, grid) # (B, 256, 20, 8)
# 4. 加权求和 + 残差
out = (attention_weights * sampled).sum(dim=-1) # (B, 256, 20)
out = self.output_proj(out.permute(0,2,1)) # (B, 20, 256)
return self.dropout(out) + queries # 残差连接DiffusionDrive 预测 20 条轨迹，但 GT（Ground Truth）只有 1 条。怎么算损失？
策略：找到 20 条中离 GT 最近的那条（"赢家"），只对它算回归损失。
同时用分类损失让网络学会给"赢家"打高分。
class LossComputer(nn.Module):
def forward(self, poses_reg, poses_cls, target_traj, plan_anchor, target_mask):
"""
poses_reg: (B, 20, 8, 3) — 20条预测轨迹，每条8步，每步(x,y,heading)
poses_cls: (B, 20) — 20条轨迹的置信度 logit
target_traj: (B, 8, 3) — GT 轨迹
plan_anchor: (B, 20, 8, 2) — 20条锚点轨迹的 (x,y)
返回：(cls_loss, reg_loss)
"""
# 1. 找"赢家"：计算每条锚点到 GT 的平均距离
dist = ||target_traj[:,:2] - plan_anchor|| # (B, 20, 8)
dist = dist.mean(dim=-1) # (B, 20) 每条轨迹的平均距离
mode_idx = dist.argmin(dim=-1) # (B,) 最近的那条的索引
# 2. 分类损失：Focal Loss
# 目标：让 mode_idx 对应的那条得分最高
target_onehot = one_hot(mode_idx, 20) # (B, 20)，只有赢家位置是1
cls_loss = focal_loss(poses_cls, target_onehot)
# 3. 回归损失：只对赢家算 L1
best_reg = poses_reg.gather(1, mode_idx) # (B, 1, 8, 3)
reg_loss = L1_loss(best_reg, target_traj)
return cls_loss * 10.0, reg_loss * 8.0为什么用 Focal Loss？
- 20 条轨迹中只有 1 条是正样本，19 条是负样本（极度不平衡）
- Focal Loss 会降低"容易分类"的样本的权重，让网络专注于难分的样本
这是 DiffusionDrive 的主模型文件（559行）。我们只迁移其中的轨迹规划部分。
class V2TransfuserModel(nn.Module):
def __init__(self, config):
# ===== 感知部分（不迁移）=====
self._backbone = TransfuserBackbone(config) # ResNet34 + LiDAR 融合
self._bev_downscale = nn.Conv2d(512, 256) # BEV 特征降维
self._bev_semantic_head = ... # BEV 语义分割头
self._agent_head = AgentHead(...) # 物体检测头
# ===== Query 准备（迁移）=====
self._keyval_embedding = nn.Embedding(8**2 + 1, 256)
# 65 个位置编码：8×8=64 个 BEV token + 1 个 status token
self._query_embedding = nn.Embedding(1 + 30, 256)
# 31 个可学习 query：1 个 ego query + 30 个 agent query
self._tf_decoder = nn.TransformerDecoder(layer, 3)
# 3 层标准 Transformer 解码器
# ===== 轨迹规划（迁移）=====
self._trajectory_head = TrajectoryHead(...)def forward(self, features, targets):
# 1. 骨干网络提取 BEV 特征
bev_feature_high, bev_feature_low = self._backbone(features)
# bev_feature_high: (B, 512, 32, 32) — 高分辨率
# bev_feature_low: (B, 64, 128, 128) — 低分辨率
# 2. BEV 特征融合（原版特有，迁移时简化掉）
bev_proj = upscale(bev_feature_high) + bev_feature_low # 拼接两路
# 3. 准备 keyval（给 Transformer 解码器的"记忆"）
bev_downscaled = self._bev_downscale(bev_feature_high) # (B, 256, 8, 8)
bev_flatten = bev_downscaled.flatten(2).permute(0, 2, 1) # (B, 64, 256)
status_encoding = self._status_encoding(ego_status) # (B, 1, 256)
keyval = torch.cat([bev_flatten, status_encoding], dim=1) # (B, 65, 256)
keyval = keyval + self._keyval_embedding.weight # 加位置编码
# 4. Transformer 解码器：从 keyval 中提取 query
queries = self._query_embedding.weight.expand(B, -1, -1) # (B, 31, 256)
query_out = self._tf_decoder(queries, keyval) # (B, 31, 256)
# 5. 拆分 query
ego_query, agents_query = query_out.split([1, 30], dim=1)
# ego_query: (B, 1, 256) — 自车的全局表示
# agents_query: (B, 30, 256) — 30个周围物体的表示
# 6. 轨迹规划
trajectory = self._trajectory_head(
ego_query, agents_query, bev_proj, status_encoding, ...
)class TrajectoryHead(nn.Module):
"""这就是 DiffusionDrive 论文的核心：用扩散模型生成轨迹"""
def __init__(self, config):
# 扩散调度器（控制加噪/去噪的时间表）
self.diffusion_scheduler = DDIMScheduler(
num_train_timesteps=1000, # 总共 1000 步扩散
beta_schedule="scaled_linear", # 噪声增长方式
prediction_type="sample", # 网络直接预测干净样本（而非噪声）
)
# 锚点轨迹：20条从训练数据聚类得到的典型轨迹
self.plan_anchor = nn.Parameter(
torch.from_numpy(np.load("anchors.npy")), # (20, 8, 2)
requires_grad=False # 不参与训练
)
# 时间步嵌入
self.time_mlp = nn.Sequential(
SinusoidalPosEmb(256), # 把整数时间步编码为向量
nn.Linear(256, 1024),
nn.Mish(), # 激活函数（比 ReLU 更平滑）
nn.Linear(1024, 256),
)
# DiT 解码器（2层）
self.diff_decoder = CustomTransformerDecoder(
CustomTransformerDecoderLayer(...), num_layers=2
)这是 DiffusionDrive 论文的核心创新——一个 DiT（Diffusion Transformer）块：
class CustomTransformerDecoderLayer(nn.Module):
"""
一个 DiT 块包含 4 个子模块，按顺序执行：
┌─────────────────────────────────────────────────────┐
│ 输入：traj_feature (B, 20, 256) │
│ traj_points (B, 20, 8, 2) │
│ │
│ 1. cross_bev_attention（BEV 网格采样注意力） │
│ → 让轨迹"看到"它经过的 BEV 区域 │
│ │
│ 2. cross_agent_attention（多头交叉注意力） │
│ → 让轨迹"看到"周围 30 个物体 │
│ │
│ 3. cross_ego_attention（多头交叉注意力） │
│ → 让轨迹"看到"自车的全局状态 │
│ │
│ 4. FFN + time_modulation（前馈网络 + 时间调制） │
│ → 非线性变换 + 根据扩散时间步调整特征 │
│ │
│ 5. refinement_module（输出头） │
│ → 预测轨迹偏移量 + 置信度分数 │
│ │
│ 输出：poses_reg (B, 20, 8, 3) │
│ poses_cls (B, 20) │
└─────────────────────────────────────────────────────┘
"""
def __init__(self, num_poses, d_model, d_ffn, num_heads, ...):
# 三种交叉注意力
self.cross_bev_attention = GridSampleCrossBEVAttention(...)
self.cross_agent_attention = nn.MultiheadAttention(d_model, num_heads)
self.cross_ego_attention = nn.MultiheadAttention(d_model, num_heads)
# 前馈网络
self.ffn = nn.Sequential(Linear, ReLU, Dropout, Linear, Dropout)
# 时间调制（DiT 的核心）
self.time_modulation = ModulationLayer(d_model, condition_dims=d_model)
# 输出头
self.refinement_module = DiffMotionPlanningRefinementModule(d_model, 8, 20)
def forward(self, traj_feature, traj_points, bev_feature, ...,
agents_query, ego_query, time_embed, status_encoding):
# 1. BEV 注意力
traj_feature = self.cross_bev_attention(
traj_feature, traj_points, bev_feature, spatial_shape
)
# 2. Agent 注意力（标准多头注意力）
traj_feature = self.cross_agent_attention(
query=traj_feature, # (B, 20, 256) — 轨迹问
key=agents_query, # (B, 30, 256) — 物体答
value=agents_query,
)
# 3. Ego 注意力
traj_feature = self.cross_ego_attention(
query=traj_feature, # (B, 20, 256)
key=ego_query, # (B, 1, 256) — 自车状态
value=ego_query,
)
# 4. FFN + 时间调制
traj_feature = self.ffn(traj_feature)
traj_feature = self.time_modulation(traj_feature, time_embed, ...)
# 5. 预测轨迹
poses_reg, poses_cls = self.refinement_module(traj_feature)
# poses_reg 是相对于 noisy_traj_points 的偏移量
poses_reg[..., :2] += noisy_traj_points # 加上当前噪声位置
poses_reg[..., 2] = tanh(poses_reg[..., 2]) * π # heading 限制在 [-π, π]
return poses_reg, poses_clsclass ModulationLayer(nn.Module):
"""
DiT 的核心机制：根据扩散时间步 t 来调制特征。
原理（FiLM: Feature-wise Linear Modulation）：
output = feature * (1 + scale) + shift
其中 scale 和 shift 由时间步嵌入决定。
直觉：
- 在扩散早期（t 大，噪声多），网络需要做"粗略"的去噪
- 在扩散后期（t 小，噪声少），网络需要做"精细"的修正
- 时间调制让同一个网络在不同阶段表现不同
"""
def __init__(self, embed_dims, condition_dims):
self.scale_shift_mlp = nn.Sequential(
nn.SiLU(),
nn.Linear(condition_dims, embed_dims * 2), # 输出 scale 和 shift
)
def forward(self, traj_feature, time_embed, global_cond, global_img):
# time_embed: (B, 1, 256) — 当前时间步的嵌入
scale_shift = self.scale_shift_mlp(time_embed) # (B, 1, 512)
scale, shift = scale_shift.chunk(2, dim=-1) # 各 (B, 1, 256)
return traj_feature * (1 + scale) + shiftNavFormer 是 WorldEngine 的核心检测+规划模型，位于 detectors/navformer.py。
输入图像序列 → ResNet50 骨干 → BEVFormer（生成 BEV 特征）→ Planning Head → 输出轨迹
NavFormer 有两个主要入口：
forward_train()：训练时调用，返回 loss 字典forward_test()：推理时调用，返回预测结果列表
def forward_train(self, img, ..., command, sdc_planning, sdc_planning_past, sdc_status, ...):
# 第一步：感知（检测+跟踪）
outs_track = self.forward_track_train(img, gt_bboxes_3d, ...)
# 输出 outs_track['bev_embed']：形状 (H*W, B, C) = (40000, B, 256)
# 第二步：取出 BEV 特征
bev_embed = outs_track['bev_embed']
# 第三步：调用规划头的 forward
plan_results = self.planning_head.forward(
bev_embed, # (40000, B, 256) — BEV 鸟瞰图特征
command, # (B,) — 导航指令（左转/直行/右转）
sdc_planning_past, # (B, 1, 4, 3) — 过去4步的轨迹
sdc_status, # (B, 3) — 当前车辆状态（速度等）
sdc_planning_mask_past, # (B, 1, 4, 3) — 过去轨迹的有效性掩码
gt_pre_command_sdc, # (B, 1, 4, 1) — 预处理的导航指令
)
# 第四步：构建 PDM 评分字典（用于某些 head 的额外监督）
pdm_dict = {
"no_at_fault_collisions": no_at_fault_collisions,
"drivable_area_compliance": drivable_area_compliance,
"ego_progress": ego_progress,
...
}
# 第五步：计算损失
losses = self.planning_head.loss(
plan_results,
gt_pdm_score=pdm_dict,
sdc_planning=sdc_planning, # (B, 1, 8, 3) — GT 未来轨迹
sdc_planning_mask=sdc_planning_mask # (B, 1, 8, 3) — GT 掩码
)
return losses关键点：
bev_embed的形状是(H*W, B, C) = (40000, B, 256)，其中 H=W=200- Planning head 的
forward()不接收 GT 未来轨迹（GT 只在loss()中使用） - 这个调用签名对所有 planning head 都一样（HydraMDP、VADv2、DiffusionDrive）
def forward_test(self, img, ..., command, sdc_planning, ...):
# 感知
outs_track = self.forward_track_test(img, ...)
bev_embed = outs_track['bev_embed']
# 规划（注意：test 时输入多了 [0] 索引，因为 dataloader 多包了一层）
plan_results = self.planning_head.forward(
bev_embed,
command[0], # 去掉 dataloader 多包的那层
sdc_planning_past[0],
sdc_status[0],
sdc_planning_mask_past[0],
gt_pre_command_sdc[0],
)
# 取出结果
chosen_indices = plan_results['selected_indices'] # (B,) — 选中的 vocab 索引
# 逐 batch 构建输出
for batch_idx in range(b):
chosen_idx = chosen_indices[batch_idx]
pdm_dict = {
"trajectory": plan_results['trajectory'][batch_idx].cpu().numpy(),
# trajectory 形状：(40, 3)
"score": score[0][batch_idx, chosen_idx].item(),
...
}
# 计算 ADE/FDE
pred_traj = pdm_dict['trajectory'][4::5, :2] # (40,3) → 取每5步 → (8,2)
ade = mean(||pred - gt||)
fde = ||pred[-1] - gt[-1]||关键点：
plan_results['trajectory']必须是(B, 40, 3)形状plan_results['selected_indices']必须是(B,)的整数索引- 评估时用
[4::5]从 40 步中抽取 8 步（每 5 步取 1 步）
TrajScoringHead 是 HydraMDP/VADv2 使用的规划头。DiffusionDrive 的新 head 需要和它有相同的接口。
@HEADS.register_module()
class TrajScoringHead(nn.Module):
def __init__(self, num_poses, d_ffn, d_model, vocab_path, ...):
# 加载 8192 条预计算轨迹
self.vocab = nn.Parameter(torch.from_numpy(np.load(vocab_path)), requires_grad=False)
# vocab 形状：(8192, 40, 3) — 8192条轨迹，每条40步，每步(x,y,heading)
# Transformer 解码器：对每条 vocab 轨迹打分
self.transformer = build_transformer_layer_sequence(...)
# 打分头
self.heads = {'imi': MLP(d_model → 1)} # 输出每条轨迹的分数
def forward(self, bev_embed, command, sdc_planning_past, sdc_status, ...):
# 1. 编码状态信息
status_token = encode(command, sdc_planning_past, sdc_status, ...)
# 2. 对 8192 条轨迹打分
scores = self.transformer(vocab_features, bev_embed, status_token)
# 3. 选最高分的轨迹
selected_indices = scores.argmax(dim=-1) # (B,)
trajectory = self.vocab[selected_indices] # (B, 40, 3)
return {"trajectory": trajectory, "selected_indices": selected_indices}DiffusionDrive 与 TrajScoringHead 的本质区别：
- TrajScoringHead：从 8192 条固定轨迹中选择一条（分类问题）
- DiffusionDrive：从 20 条锚点出发生成新轨迹（生成问题）
def nerf_positional_encoding(tensor, num_encoding_functions=6):
"""
NeRF 风格的位置编码。把一个标量变成 12 维向量。
对输入 x，生成：
[sin(2^0 * x), cos(2^0 * x), sin(2^1 * x), cos(2^1 * x), ..., sin(2^5 * x), cos(2^5 * x)]
输入维度 d → 输出维度 d * 12（6个频率 × sin/cos 两个函数）
用途：把低维的状态信息（速度、位置等）编码成高维特征，
让网络更容易学习到细微的数值差异。
"""文件：projects/AlgEngine/mmdet3d_plugin/navformer/dense_heads/diffusion_planning_head.py（954行）
import copy, math
import numpy as np
import torch, torch.nn as nn, torch.nn.functional as F
from diffusers.schedulers import DDIMScheduler # HuggingFace 的 DDIM 调度器
from mmcv.runner import auto_fp16, force_fp32 # 混合精度训练装饰器
from mmdet.models.builder import HEADS # mmdet 的模块注册器
from .traj_scoring_head import nerf_positional_encoding # 复用已有的 NeRF 编码函数为什么从 traj_scoring_head 导入？
nerf_positional_encoding是一个纯函数（无状态），两个 head 都需要用- 避免代码重复，直接复用
这些函数直接从 DiffusionDrive 的 modules/blocks.py 搬过来，逻辑完全不变：
linear_relu_ln：构建 MLP 层（见 3.2.1 节）gen_sineembed_for_position：正弦位置编码（见 3.2.2 节）bias_init_with_prob：Focal Loss 的偏置初始化SinusoidalPosEmb：扩散时间步编码
class SinusoidalPosEmb(nn.Module):
# 把一个整数时间步 t 编码成 256 维向量
def forward(self, x):
# x: (B,) - 当前扩散时间步，如 [42, 42]
half_dim = self.dim // 2 # 128
emb = math.log(10000) / (128 - 1)
emb = torch.exp(torch.arange(128) * -emb)
# emb: [1.0, 0.93, 0.87, ..., 0.0001] - 128个递减的频率
emb = x[:, None] * emb[None, :]
# (B, 1) * (1, 128) -> (B, 128) - 时间步乘以每个频率
emb = torch.cat((emb.sin(), emb.cos()), dim=-1)
# (B, 256) - sin 和 cos 拼接
return emb直觉：不同的时间步 t 会产生不同的 256 维"指纹"，让网络知道当前处于去噪的哪个阶段。
与原版的区别：
- 原版：
self.config.lidar_max_x/self.config.lidar_max_y（从全局 config 对象读取） - 迁移版：
self.bev_range_x/self.bev_range_y（构造函数参数，值为 51.2）
为什么从 32 变成 51.2？
- DiffusionDrive 原版用 LiDAR，覆盖 +/-32 米
- WorldEngine 的 BEVFormer 覆盖 +/-51.2 米（point_cloud_range = [-51.2, ..., 51.2]）
- 归一化时除以的范围必须匹配实际 BEV 覆盖范围
从 modules/multimodal_loss.py 搬过来，做了两个改动：
-
LossComputer.init 签名变了：
- 原版：
__init__(self, config: TransfuserConfig)- 从 config 对象读权重 - 迁移版：
__init__(self, cls_loss_weight=10.0, reg_loss_weight=8.0)- 直接传参
- 原版：
-
LossComputer.forward 签名变了：
- 原版：
forward(self, poses_reg, poses_cls, targets, plan_anchor)- targets 是字典 - 迁移版：
forward(self, poses_reg, poses_cls, target_traj, plan_anchor, target_mask)- 直接传张量
- 原版：
每层 DiT 块的输出头：预测轨迹偏移 + 置信度分数。
class DiffMotionPlanningRefinementModule(nn.Module):
def __init__(self, embed_dims=256, ego_fut_ts=8, ego_fut_mode=20):
# 分类分支：预测每条轨迹的置信度
self.plan_cls_branch = [LN, Linear, ReLU, LN, Linear, ReLU, Linear(256->1)]
# 回归分支：预测轨迹偏移
self.plan_reg_branch = [Linear(256->256), ReLU, Linear(256->8*3)]
def forward(self, traj_feature):
# traj_feature: (B, 20, 256)
plan_cls = self.plan_cls_branch(traj_feature).squeeze(-1) # (B, 20)
plan_reg = self.plan_reg_branch(traj_feature).view(B, 20, 8, 3)
return plan_reg, plan_clsDiT 风格的时间步调制层（FiLM = Feature-wise Linear Modulation）。
公式：output = input * (1 + scale) + shift
其中 scale 和 shift 由时间步嵌入决定。
直觉：告诉网络"现在是去噪的第几步"，不同步骤需要不同程度的修正。
这是 DiT 块的核心 - 一个完整的解码器层：
数据流：
traj_feature (B, 20, 256)
-> cross_bev_attention（看 BEV 地图上轨迹经过的位置）
-> cross_agent_attention（看周围 30 个物体的特征）
-> cross_ego_attention（看自车状态特征）
-> FFN（前馈网络）
-> time_modulation（时间步调制）
-> refinement_module（输出轨迹偏移 + 分数）
forward 的最后几步：
# 轨迹偏移是相对于当前噪声轨迹的残差
poses_reg[..., :2] += noisy_traj_points # 加上当前位置（残差预测）
poses_reg[..., 2] = torch.tanh(poses_reg[..., 2]) * math.pi # heading 限制在 [-pi, pi]
return poses_reg, poses_cls堆叠 N=2 个 DiT 块，实现深度监督（每层都输出预测，都参与 loss 计算）：
class CustomTransformerDecoder(nn.Module):
def forward(self, traj_feature, noisy_traj_points, bev_feature, ...):
poses_reg_list = []
poses_cls_list = []
traj_points = noisy_traj_points
for layer in self.layers: # 2 层
poses_reg, poses_cls = layer(traj_feature, traj_points, ...)
poses_reg_list.append(poses_reg)
poses_cls_list.append(poses_cls)
# 关键：用当前层的输出更新下一层的输入坐标
traj_points = poses_reg[..., :2].clone().detach()
return poses_reg_list, poses_cls_list为什么 .detach()？
- 防止梯度从第2层回传到第1层的输出
- 每层独立优化，避免梯度爆炸
@HEADS.register_module()
class DiffusionPlanningHead(nn.Module):
def __init__(self,
num_poses=8, # 预测未来 8 个时间步
d_model=256, # 特征维度
d_ffn=1024, # FFN 中间维度
num_heads=8, # 注意力头数
dropout=0.0, # Dropout
num_bounding_boxes=30, # 最多 30 个周围物体
num_query_decoder_layers=3, # Query 准备用的 TF 解码器层数
query_keyval_size=8, # BEV 下采样到 8x8
num_anchors=20, # 20 条锚点轨迹
num_diff_decoder_layers=2, # DiT 解码器层数
plan_anchor_path=None, # 锚点轨迹文件路径（.npy）
vocab_path=None, # 8192-vocab 文件路径（.npy）
bev_h=200, bev_w=200, # BEV 特征图尺寸
bev_range_x=51.2, # BEV x 方向范围（米）
bev_range_y=51.2, # BEV y 方向范围（米）
odo_x_min=-1.2, # 轨迹 x 归一化最小值
odo_x_range=56.9, # 轨迹 x 归一化范围
odo_y_min=-20.0, # 轨迹 y 归一化最小值
odo_y_range=46.0, # 轨迹 y 归一化范围
num_train_timesteps=1000, # 扩散总步数
train_timestep_max=50, # 训练时最大时间步（截断）
inference_steps=2, # 推理时 DDIM 步数
trunc_timesteps=8, # 推理时起始噪声的时间步
cls_loss_weight=10.0, # 分类损失权重
reg_loss_weight=8.0, # 回归损失权重
use_nerf=True, # 是否用 NeRF 编码
):构造函数中创建的子模块（按顺序）：
(a) status_embed - 状态编码器
# 把车辆状态信息编码成 256 维向量
# 输入：command(4维 one-hot) + past_traj(经 NeRF 编码后 24维) + status(2维)
# 如果 use_nerf=True：输入维度 = 4 + 24 + 2 = 30
self.status_embed = nn.Sequential(nn.Linear(30, 256), nn.ReLU())(b) _query_embedding - 可学习的 Query
# 31 个可学习向量：1 个 ego query + 30 个 agent query
self._query_embedding = nn.Embedding(31, 256)(c) _keyval_embedding - Key-Value 位置编码
# 65 个位置编码：8x8=64 个 BEV token + 1 个 status token
self._keyval_embedding = nn.Embedding(65, 256)(d) _query_tf_decoder - Query 准备解码器
# 标准 Transformer 解码器（3层），用于从 BEV 中提取 ego/agent 信息
self._query_tf_decoder = nn.TransformerDecoder(
nn.TransformerDecoderLayer(d_model=256, nhead=8, dim_feedforward=1024, batch_first=True),
num_layers=3
)(e) diffusion_scheduler - DDIM 调度器
# 控制扩散过程的噪声调度
self.diffusion_scheduler = DDIMScheduler(
num_train_timesteps=1000, # 总共 1000 步
beta_schedule="scaled_linear", # 噪声增长方式
prediction_type="sample", # 网络直接预测干净样本（而非噪声）
)(f) plan_anchor - 锚点轨迹
# 从 .npy 文件加载 20 条 K-means 聚类中心轨迹
# 形状：(20, 8, 2) - 20条轨迹，每条8步，每步(x,y)
self.plan_anchor = nn.Parameter(torch.from_numpy(np.load(path)), requires_grad=False)
# 锚点编码器：把轨迹坐标编码成特征
self.plan_anchor_encoder = nn.Sequential(
*linear_relu_ln(256, 1, 1, input_dims=512), # 512->256
nn.Linear(256, 256)
)(g) time_mlp - 时间步编码器
# 把整数时间步编码成 256 维条件向量
self.time_mlp = nn.Sequential(
SinusoidalPosEmb(256), # int -> 256维
nn.Linear(256, 1024), # 256 -> 1024
nn.Mish(), # 激活函数
nn.Linear(1024, 256), # 1024 -> 256
)(h) diff_decoder - DiT 解码器
# 2 层 DiT 块，每层包含 BEV/agent/ego 注意力 + FFN + 时间调制
self.diff_decoder = CustomTransformerDecoder(
CustomTransformerDecoderLayer(...), num_layers=2
)(i) vocab_8192 - 评估用词汇表
# 8192 条预计算轨迹，仅用于 test 时的"快照锚定"
# 形状：(8192, 40, 3)
self.vocab_8192 = nn.Parameter(torch.from_numpy(np.load(vocab_path)), requires_grad=False)(j) loss_computer - 损失计算器
self.loss_computer = LossComputer(cls_loss_weight=10.0, reg_loss_weight=8.0)文件：e2e_diffusiondrive.py → Ctrl+P: e2e_diffusiondrive.py
_base_ = ["../_base_/default_runtime.py"] # 继承默认运行时配置
custom_imports = dict(imports=['mmdet3d_plugin']) # 导入自定义插件
# BEV 空间范围（单位：米）
point_cloud_range = [-51.2, -51.2, -5.0, 51.2, 51.2, 3.0]
# 含义：[x_min, y_min, z_min, x_max, y_max, z_max]
# 即覆盖车辆周围 102.4m x 102.4m 的区域
voxel_size = [0.2, 0.2, 8] # 体素大小：0.2m x 0.2m x 8m
# BEV 分辨率 = 102.4 / 0.2 = 512? 不，实际用 bev_h=bev_w=200
_dim_ = 256 # 特征维度（贯穿整个模型）
bev_h_ = 200 # BEV 特征图高度
bev_w_ = 200 # BEV 特征图宽度
# 所以 BEV 特征图有 200*200 = 40000 个位置model = dict(
type="NAVFormer", # 使用 NavFormer 模型
# === 图像骨干网络 ===
img_backbone=dict(type="ResNet", depth=50, ...),
# ResNet50，输出第2/3/4阶段的特征
# === BEV 编码器（BEVFormer）===
pts_bbox_head=dict(
type="BEVFormerTrackHead",
bev_h=200, bev_w=200,
# BEVFormer 的 Transformer 编码器
transformer=dict(
type="PerceptionTransformerV2",
encoder=dict(
# 6 层 BEV 编码器
transformerlayers=dict(
attn_cfgs=[
dict(type="TemporalSelfAttention", ...),
dict(type="SpatialCrossAttention",
deformable_attention=dict(
type="MSDeformableAttention3D", ...
))
]
)
)
),
),
# === 规划头（DiffusionDrive）===
planning_head=dict(
type="DiffusionPlanningHead", # 注册名，对应我们的新类
num_poses=8,
d_model=256,
d_ffn=1024,
num_heads=8,
num_bounding_boxes=30,
num_query_decoder_layers=3,
query_keyval_size=8,
num_anchors=20,
num_diff_decoder_layers=2,
bev_h=200,
bev_w=200,
bev_range_x=51.2,
bev_range_y=51.2,
plan_anchor_path="data/alg_engine/ckpts/plan_anchor_20_8.npy",
vocab_path="data/alg_engine/ckpts/test_8192_kmeans.npy",
num_train_timesteps=1000,
train_timestep_max=50,
inference_steps=2,
trunc_timesteps=8,
cls_loss_weight=10.0,
reg_loss_weight=8.0,
use_nerf=True,
),
)mmdetection 的配置系统工作方式：
配置文件 dict Python 代码
----------- ----------
planning_head=dict( class DiffusionPlanningHead:
type="DiffusionPlanningHead", def __init__(self,
num_poses=8, --> num_poses=8,
d_model=256, --> d_model=256,
plan_anchor_path=..., --> plan_anchor_path=...,
) ):
框架会自动：
- 根据
type字段找到注册的类 - 把 dict 中其余的 key-value 作为
__init__的参数传入 - 实例化对象
optimizer = dict(
type="AdamW",
lr=2e-4, # 学习率
paramwise_cfg=dict(
custom_keys={
"img_backbone": dict(lr_mult=0.1), # 骨干网络用 1/10 学习率
}
),
weight_decay=0.01,
)
lr_config = dict(
policy="CosineAnnealing", # 余弦退火学习率
warmup="linear", # 线性预热
warmup_iters=500, # 预热 500 步
warmup_ratio=1.0 / 3, # 预热起始为 1/3 学习率
min_lr_ratio=1e-3, # 最低学习率为初始的 1/1000
)
total_epochs = 8 # 训练 8 个 epochdataset_type = "NavSimOpenSceneE2E"
data_root = "data/raw/openscene-v1.1/"
# 训练集/验证集/测试集的标注文件
ann_file_train = "data/alg_engine/merged_infos_navformer/nuplan_openscene_navtrain.pkl"
ann_file_val = "data/alg_engine/merged_infos_navformer/nuplan_openscene_navtest.pkl"
# 预训练权重（只加载感知部分，规划头随机初始化）
load_from = "data/alg_engine/ckpts/track_map_nuplan_r50_navtrain_100pct_bs1x8.pth"┌─────────────────────────────────────────────────────────────────────┐
│ DataLoader │
│ 输出一个 batch： │
│ - img: (B, T, N_cam, 3, H, W) 多帧多相机图像 │
│ - command: (B,) 导航指令 [0=左转, 1=直行, 2=右转] │
│ - sdc_planning_past: (B, 1, 4, 3) 过去4步轨迹 [(x,y,heading)*4] │
│ - sdc_status: (B, 3) 当前状态 [vx, vy, heading_rate] │
│ - sdc_planning: (B, 1, 8, 3) GT未来8步轨迹 │
│ - score 等 PDM 评分张量: (B, 8192) │
└─────────────────────────────────────────────────────────────────────┘
|
v
┌─────────────────────────────────────────────────────────────────────┐
│ NavFormer.forward_train() │
│ │
│ 1. forward_track_train(img, ...) │
│ - ResNet50 提取多尺度图像特征 │
│ - BEVFormer 生成 BEV 特征 │
│ - 输出 bev_embed: (40000, B, 256) │
│ │
│ 2. planning_head.forward(bev_embed, command, ...) │
│ (见下方详细流程) │
│ │
│ 3. planning_head.loss(plan_results, sdc_planning=GT, ...) │
│ (见下方详细流程) │
└─────────────────────────────────────────────────────────────────────┘
输入：bev_embed (40000, B, 256)
command (B,), sdc_planning_past (B,1,4,3), sdc_status (B,3), ...
Step 1: BEV 形状转换
bev_embed (40000, B, 256)
-> permute(1,2,0) -> (B, 256, 40000)
-> view(B, 256, 200, 200) -> bev_feature
这就是 (B, C, H, W) 格式的 BEV 特征图
Step 2: 构建 status_token
command (B,) -> one-hot -> (B, 4)
sdc_planning_past (B,1,4,3) -> 取 [0] -> (B,4,3) -> flatten -> (B,12)
如果 use_nerf: 对 (B,12) 做 NeRF 编码 -> (B, 12*2*6=144)... 不对
实际：
past_flat (B, 2) [取最后一步的 x,y]
past_nerf = nerf(past_flat) -> (B, 24) [2维 * 12倍]
status_raw = cat(command_onehot, past_nerf, sdc_status[:,:2]) -> (B, 4+24+2=30)
status_token = Linear(30 -> 256) + ReLU -> (B, 1, 256)
Step 3: 准备 Query（_prepare_queries）
3a. BEV 下采样：(B, 256, 200, 200) -> AdaptiveAvgPool2d(8,8) -> (B, 256, 8, 8)
-> flatten -> (B, 64, 256)
3b. 拼接 status：cat([bev_8x8, status_token], dim=1) -> (B, 65, 256)
3c. 加位置编码：keyval = keyval + self._keyval_embedding.weight # (65, 256)
3d. 准备 query：self._query_embedding.weight -> (31, 256)
query = query.expand(B, 31, 256)
3e. TransformerDecoder（3层）：
query_out = self._query_tf_decoder(query, keyval) # (B, 31, 256)
3f. 拆分：ego_query, agents_query = query_out.split([1, 30], dim=1)
ego_query: (B, 1, 256) - 自车的全局特征
agents_query: (B, 30, 256) - 30个周围物体的特征
Step 4: 扩散过程（训练 vs 推理不同）
[训练] _forward_train:
4a. 取锚点轨迹：plan_anchor (20, 8, 2) -> expand -> (B, 20, 8, 2)
4b. 归一化：norm_odo_xy(plan_anchor) -> 值域 [-1, 1]
4c. 随机采样时间步：t ~ Uniform[0, 50)
4d. 加噪：noisy = scheduler.add_noise(anchor_norm, noise, t)
4e. 反归一化得到物理坐标：noisy_traj_points = denorm_odo_xy(noisy)
4f. 编码轨迹位置：
sin_embed = gen_sineembed_for_position(noisy_traj_points, hidden_dim=64)
traj_feature = plan_anchor_encoder(sin_embed) # (B, 20, 256)
4g. 编码时间步：time_embed = time_mlp(t) # (B, 1, 256)
4h. DiT 解码器（2层）：
poses_reg_list, poses_cls_list = diff_decoder(
traj_feature, noisy_traj_points, bev_feature,
agents_query, ego_query, time_embed, status_token
)
4i. 选最佳轨迹：argmax(poses_cls[-1]) -> best_reg (B, 8, 3)
4j. 扩展为 40 步：_expand_to_40(best_reg) -> (B, 40, 3)
4k. 快照到 vocab：_snap_to_vocab(best_reg) -> selected_indices (B,)
[推理] _forward_test:
4a. 取锚点 + 加截断噪声（t=8，很小的噪声）
4b. DDIM 迭代去噪（2步）：
for k in [8, 4]: # 两个时间步
- 编码当前轨迹位置
- 编码时间步 k
- DiT 解码器预测
- DDIM step 更新轨迹
4c. 选最佳 + 扩展 + 快照（同训练）
Step 5: 返回结果字典
{
"trajectory": (B, 40, 3), # 供 NavFormer 评估
"selected_indices": (B,), # 供 PDM 评分索引
"poses_reg_list": [...], # 供 loss() 使用（仅训练）
"poses_cls_list": [...], # 供 loss() 使用（仅训练）
"plan_anchor_expanded": (B,20,8,2) # 供 loss() 使用（仅训练）
}
输入：
result = forward() 的返回字典
sdc_planning: (B, 1, 8, 3) - GT 未来轨迹
sdc_planning_mask: (B, 1, 8, 3) - 有效性掩码
Step 1: 提取 GT
target_traj = sdc_planning[:, 0] # (B, 8, 3)
Step 2: 对每层 DiT 的输出计算损失（深度监督）
for poses_reg, poses_cls in zip(poses_reg_list, poses_cls_list):
cls_loss, reg_loss = loss_computer(
poses_reg, # (B, 20, 8, 3) - 20条预测轨迹
poses_cls, # (B, 20) - 20条轨迹的分数
target_traj, # (B, 8, 3) - GT
plan_anchor, # (B, 20, 8, 2) - 锚点（用于找最近的）
target_mask # (B, 8) - 掩码
)
cls_total += cls_loss
reg_total += reg_loss
Step 3: 返回
{"loss.diff_cls": cls_total, "loss.diff_reg": reg_total}
e2e_diffusiondrive.py (配置)
|
| planning_head=dict(type="DiffusionPlanningHead", num_poses=8, ...)
|
v
DiffusionPlanningHead.__init__(num_poses=8, ...)
|
| 创建子模块时传递参数
|
+---> GridSampleCrossBEVAttention(embed_dims=256, bev_range_x=51.2, ...)
+---> CustomTransformerDecoder(layer, num_layers=2)
| |
| +---> CustomTransformerDecoderLayer(num_poses=8, d_model=256, ...)
| |
| +---> GridSampleCrossBEVAttention(...)
| +---> nn.MultiheadAttention(256, 8) x2
| +---> ModulationLayer(256)
| +---> DiffMotionPlanningRefinementModule(256, 8, 20)
|
+---> LossComputer(cls_loss_weight=10.0, reg_loss_weight=8.0)
+---> DDIMScheduler(num_train_timesteps=1000, ...)
+---> plan_anchor: (20, 8, 2) from .npy file
+---> vocab_8192: (8192, 40, 3) from .npy file
| 方面 | 训练 (_forward_train) | 推理 (_forward_test) |
|---|---|---|
| 噪声来源 | 随机高斯噪声 | 截断噪声（t=8，很小） |
| 时间步 | 随机 t in [0, 50) | 固定 [8, 4]（2步DDIM） |
| DiT 解码器调用次数 | 1次 | 2次（每个DDIM步1次） |
| 输出用途 | 计算 loss | 直接作为预测结果 |
| poses_reg_list | 保存（给 loss 用） | 只用最后一层 |
| 梯度 | 有 | 无（torch.no_grad） |
训练时：
干净锚点轨迹 --加噪声--> 噪声轨迹 --DiT解码--> 预测轨迹
|
与 GT 比较算 loss
网络学习的是：给定一个带噪声的轨迹，如何把它"修正"成正确的轨迹。
推理时：
干净锚点轨迹 --加很小噪声--> 微噪声轨迹 --DDIM去噪2步--> 最终轨迹
因为训练时 t 最大只到 50（而不是 1000），所以推理时也只需要很少的去噪步骤。 这就是 DiffusionDrive 论文的核心创新：截断扩散，大幅减少推理时间。
传统扩散模型需要 1000 步去噪，DiffusionDrive 只需要 2 步，原因：
- 锚点初始化：不是从纯噪声开始，而是从 20 条有意义的锚点轨迹开始
- 截断扩散：训练时 t 最大只到 50（
train_timestep_max=50），噪声很小 - 推理起点：从 t=8 开始（
trunc_timesteps=8），噪声更小 - DDIM：比 DDPM 更高效的采样方法，可以跳步
这两个函数是兼容性适配器，让 DiffusionDrive 的输出格式匹配 NavFormer 的期望：
def _expand_to_40(self, traj_8):
# traj_8: (B, 8, 3) - DiffusionDrive 原生输出 8 步
# NavFormer 期望 40 步，评估时用 [4::5] 取回 8 步
# 所以我们把每步重复 5 次：[a,b,c,...] -> [a,a,a,a,a, b,b,b,b,b, ...]
return traj_8.unsqueeze(2).expand(-1,-1,5,-1).reshape(B, 40, 3)
def _snap_to_vocab(self, traj_8):
# traj_8: (B, 8, 3) - 预测的 8 步轨迹
# vocab_8192: (8192, 40, 3) - 预计算的 8192 条轨迹
# 找到 vocab 中与预测最接近的那条，返回其索引
vocab_8 = self.vocab_8192[:, 4::5, :2] # (8192, 8, 2)
dist = ||traj_8[:,:,:2] - vocab_8|| # (B, 8192)
return dist.argmin(dim=-1) # (B,)为什么需要 selected_indices？
- NavFormer 的评估流程会用这个索引去查 PDM 评分张量
- PDM 评分是预计算的：对 8192 条轨迹分别模拟，得到碰撞/舒适度等分数
- 虽然 DiffusionDrive 生成的轨迹不在 8192 中，但找最近的那条来查分数是合理的近似
| 原始 DiffusionDrive | 迁移后 DiffusionPlanningHead | 变化原因 |
|---|---|---|
| BEV 来自 ResNet34+LiDAR | BEV 来自 BEVFormer | WorldEngine 用纯视觉 |
| BEV 范围 32m | BEV 范围 51.2m | BEVFormer 覆盖更大 |
| BEV 形状 (B,C,H,W) | 入口 reshape (H*W,B,C)->(B,C,H,W) | NavFormer 的 BEV 格式不同 |
| config 对象传参 | 构造函数直接传参 | 去掉对 TransfuserConfig 的依赖 |
| 多文件分散 | 单文件 self-contained | 可删除原始目录 |
| 输出 8 步 | 输出 40 步 + vocab 索引 | 兼容 NavFormer 评估 |
| 有自己的 backbone/detector | 无（由 NavFormer 提供） | 只迁移规划头 |
| query 由外部模型提供 | query 在 head 内部准备 | 保持 NavFormer 调用接口不变 |
| 无 vocab 概念 | 添加 8192-vocab 快照 | 兼容 PDM 评估流程 |
| 张量名 | 形状 | 含义 |
|---|---|---|
| bev_embed (NavFormer输出) | (40000, B, 256) | 200x200 BEV 特征，序列优先格式 |
| bev_feature (head内部) | (B, 256, 200, 200) | 同上，reshape 为图像格式 |
| command | (B,) | 导航指令 [0,1,2,3] |
| sdc_planning_past | (B, 1, 4, 3) | 过去4步轨迹 |
| sdc_status | (B, 3) | 当前状态 [vx, vy, heading_rate] |
| status_token | (B, 1, 256) | 编码后的状态特征 |
| ego_query | (B, 1, 256) | 自车全局特征 |
| agents_query | (B, 30, 256) | 30个周围物体特征 |
| plan_anchor | (20, 8, 2) | 20条锚点轨迹的 (x,y) |
| traj_feature | (B, 20, 256) | 20条轨迹的特征表示 |
| noisy_traj_points | (B, 20, 8, 2) | 加噪后的轨迹坐标 |
| time_embed | (B, 1, 256) | 时间步编码 |
| poses_reg | (B, 20, 8, 3) | 预测的轨迹 (x,y,heading) |
| poses_cls | (B, 20) | 每条轨迹的置信度 |
| trajectory (输出) | (B, 40, 3) | 最终轨迹（40步格式） |
| selected_indices (输出) | (B,) | vocab 中最近邻的索引 |
| sdc_planning (GT) | (B, 1, 8, 3) | GT 未来轨迹 |
| vocab_8192 | (8192, 40, 3) | 预计算的 8192 条轨迹 |
# dense_heads/__init__.py
from .traj_scoring_head import TrajScoringHead
from .traj_scoring_head_RL import TrajScoringHeadRL
from .diffusion_planning_head import DiffusionPlanningHead这三行的作用：
- 当 Python 导入
dense_heads包时，自动执行这些 import - 每个 import 会触发对应文件中的
@HEADS.register_module()装饰器 - 装饰器把类注册到全局字典
HEADS - 之后配置文件中写
type="DiffusionPlanningHead"就能找到这个类
x_0 (干净数据) --加噪声--> x_1 --加噪声--> ... --加噪声--> x_T (纯噪声)
每一步加的噪声量由 beta schedule 控制。DiffusionDrive 用 "scaled_linear" schedule。
x_T (噪声) --网络预测--> x_{T-1} --网络预测--> ... --网络预测--> x_0 (干净数据)
网络学习的是：给定 x_t 和 t，预测 x_0（prediction_type="sample"）。
- DDPM：每步加随机噪声，需要很多步（1000步）
- DDIM：确定性采样，可以跳步（如 1000步 -> 2步）
- 不从纯噪声 x_1000 开始
- 而是从"锚点 + 少量噪声" x_50 开始
- 推理时从 x_8 开始，只需 2 步就能得到 x_0