README.md

---
title: WJAD Sandbox
emoji: 🚗
colorFrom: blue
colorTo: indigo
sdk: docker
app_port: 7860
pinned: false
---

# WJAD - 端到端自动驾驶模型

基于 [Design.md](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 微调。

## 三步训练路径

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

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

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

## 数据准备

```bash
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。