WJAD - 端到端自动驾驶模型

基于 Design.md 实现的端到端自动驾驶模型，用于 NVIDIA Cosmos-Drive-Dreams 数据集。

架构概览

• 视觉编码器：本地 DINOv3 ViT-B/16（dinov3-vitb16-pretrain-lvd1689m），SDPA 注意力。
• 时空压缩：2×2×2 Conv3D，将 8 帧 × 24 × 64 patch tokens 压缩为 1536 个视觉 token。
• 在线校准：dim=256，6 层 (1 GateCrossAttn + 2 GateSelfAttn) × 2，跨注意力 K/V 来自 DINOv3 patch；输入与残差均在 symlog 空间，输出 SE3 + 内外参修正量。
• 主干：18 层 GateSelfAttention（前 9 Dense + 后 9 MoE，每层独立 7 路由 + 1 共享专家，GAP 序列级 Sigmoid Top-3），dim=768，12 头 SDPA + PreNorm + SwiGLU。
• 位置编码：3D RoPE 仅作用于视觉 token 的 Q/K——头 0-3 编码自车系单位射线，头 4-7 编码 H/W/T，头 8-11 零频段（identity，统一代码路径）。其余 token（ego/det/ctrl/extra）使用一一对应的可学习 PE。
• 统一检测+预测头：1024 token 同时输出 cls + is_dynamic + box3d(μ,logσ) + 未来 24 帧轨迹(μ,logσ)。
• 控制头：24 token 输出自车未来轨迹与全局控制（均 NLL μ/logσ）。
• 多任务训练：GradNorm 自适应任务权重 + Stage2 启用 PCGrad 正交化梯度冲突。
• 训练阶段：Stage1 Dense + 路由锐化 + 中期运动学/内外参扰动；Stage2 切 Top-3 + DINOv3 低 LR 微调。

三步训练路径

# 1. 本地跑通（纯随机张量）
python -m scripts.smoke_test

# 2. HF Sandbox 微小训练
python -m scripts.push_to_sandbox

# 3. HF Jobs 全量训练
python -m scripts.push_to_jobs

数据准备

python -m scripts.download_data --odir ./data/cosmos --file_types synthetic,lidar,hdmap

项目结构

src/wjad/
├── modules/        # 公用算子：FFN/门控注意力/MoE/RoPE/可学习PE/symlog/...
├── encoders/       # DINOv3 包装 + 2x2x2 时空压缩
├── calibration/    # 在线校准网络
├── backbone/       # 18 层主干
├── heads/          # 检测+预测头、控制头
├── data/           # Cosmos-Drive-Dreams 加载器、f-theta、增广
├── losses/         # NLL/检测/轨迹/控制/MoE/校准正则
├── train/          # 多任务（GradNorm+PCGrad）、Trainer、调度
└── model.py        # 顶层 E2EAVModel

License

代码遵循仓库根目录指定的开源协议。DINOv3 权重遵循 Meta DINOv3 License；Cosmos-Drive-Dreams 数据集遵循 CC BY 4.0。