手把手教你用nuPlan数据集和PyTorch框架训练你的第一个自动驾驶规划模型
手把手教你用nuPlan数据集和PyTorch框架训练你的第一个自动驾驶规划模型
自动驾驶技术的核心挑战之一是如何让车辆在复杂环境中做出安全、舒适的行驶决策。nuPlan数据集的出现为这一领域的研究者提供了前所未有的实验平台——它不仅包含1300小时的真实驾驶数据,更配备了完整的训练框架、闭环仿真器和可视化工具。本文将带您从零开始,用PyTorch构建一个基础的规划模型,并通过nuBoard分析其表现。
1. 环境配置与数据准备
在开始建模前,需要搭建支持GPU加速的Python环境。推荐使用conda创建独立环境:
conda create -n nuplan python=3.8 conda install pytorch torchvision cudatoolkit=11.3 -c pytorch pip install nuplan-devkit==1.1.0nuPlan数据集下载后需解压到指定目录,其结构包含以下关键部分:
sensor_blobs/:原始传感器数据maps/:高精地图文件db/:SQLite格式的场景数据库features/:预提取的特征数据
提示:首次使用时建议先下载最小的
mini版本(约15GB)进行功能验证,再根据需要下载完整数据集。
2. 数据加载与特征工程
nuPlan的NuPlanScenarioBuilder提供了标准化的数据接口。以下代码展示了如何加载波士顿地区的训练数据:
from nuplan.planning.scenario_builder.nuplan_db.nuplan_scenario_builder import NuPlanScenarioBuilder from nuplan.planning.scenario_builder.scenario_filter import ScenarioFilter scenario_builder = NuPlanScenarioBuilder( data_root="~/nuplan/dataset", map_root="~/nuplan/maps", db_files=["~/nuplan/dataset/nuplan_v1.0/mini/boston"], ) scenario_filter = ScenarioFilter( scenario_types=["lane_following", "lane_change"], map_names=["us-ma-boston"] ) scenarios = scenario_builder.get_scenarios(scenario_filter)典型输入特征包括:
- 自车状态:速度、加速度、航向角
- 周围物体:位置、速度、类型(车辆/行人)
- 地图信息:车道拓扑、交通信号、可行驶区域
# 特征提取示例 def extract_features(scenario): ego_states = scenario.get_ego_past_trajectory() detections = scenario.get_tracked_objects() roadgraph = scenario.get_roadgraph_features() return { 'ego': ego_states, 'objects': detections, 'map': roadgraph }3. 模型架构设计
我们采用基于Transformer的编码器-解码器结构处理时序规划问题。编码器负责理解环境上下文,解码器生成未来轨迹分布。
import torch import torch.nn as nn from torch.nn import Transformer class PlanningTransformer(nn.Module): def __init__(self, d_model=256, nhead=8, num_layers=6): super().__init__() self.encoder = TransformerEncoderLayer(d_model, nhead) self.decoder = TransformerDecoderLayer(d_model, nhead) self.traj_head = nn.Linear(d_model, 2*5) # 预测5个时间点的(x,y)坐标 def forward(self, src, tgt): memory = self.encoder(src) output = self.decoder(tgt, memory) return self.traj_head(output)关键设计考量:
- 多模态融合:使用注意力机制统一处理车辆状态、环境物体和地图特征
- 课程学习:先训练短期(3秒)预测,逐步扩展到长期(8秒)规划
- 不确定性建模:输出高斯混合分布而非单一轨迹
4. 训练流程与技巧
训练循环需要特别处理nuPlan的大规模数据。以下是优化后的训练步骤:
数据批处理:
from nuplan.planning.training.data_loader.datamodule import NuPlanDataModule datamodule = NuPlanDataModule( scenario_builder=scenario_builder, batch_size=32, num_workers=8 )损失函数设计:
def planning_loss(pred, target): # 轨迹端点误差 endpoint_error = F.mse_loss(pred[:, -1], target[:, -1]) # 整体形状相似性 chamfer_dist = compute_chamfer_distance(pred, target) return 0.7*endpoint_error + 0.3*chamfer_dist关键训练参数:
参数 推荐值 说明 学习率 3e-4 使用OneCycle调度 批大小 32 需根据GPU显存调整 训练轮次 50 早停法监控验证损失
注意:使用
nuplan-devkit内置的MetricAggregator可实时计算超过20种规划质量指标。
5. 结果分析与调试
训练完成后,通过nuBoard进行可视化分析:
python -m nuplan.planning.script.run_nuboard \ --scenario_dir ~/nuplan/dataset \ --simulation_dir ./output/simulations常见问题诊断表:
| 现象 | 可能原因 | 解决方案 |
|---|---|---|
| 轨迹抖动 | 损失函数未考虑运动平滑性 | 添加加速度惩罚项 |
| 偏离车道 | 地图特征编码不足 | 增强车道拓扑表示 |
| 碰撞率高 | 物体交互建模薄弱 | 引入GNN处理关系 |
可视化分析时应特别关注:
- 关键帧对比:选择变道、路口等复杂场景
- 指标相关性:检查舒适度与轨迹曲率的关系
- 失败案例:分析前10%最差表现的共性特征
6. 进阶优化方向
当基础模型跑通后,可尝试以下提升方案:
多任务学习框架:
class MultiTaskHead(nn.Module): def __init__(self, d_model): super().__init__() self.traj = nn.Linear(d_model, 10) # 轨迹预测 self.intent = nn.Linear(d_model, 3) # 行为意图分类 self.risk = nn.Linear(d_model, 1) # 碰撞风险估计 def forward(self, x): return { 'trajectory': self.traj(x), 'intention': self.intent(x), 'risk': self.risk(x) }实时性优化技巧:
- 使用TensorRT加速模型推理
- 实现自定义CUDA核函数处理注意力计算
- 采用滑动窗口缓存机制减少重复计算
在实际项目中,我们发现将规划频率从10Hz降到5Hz可节省40%计算资源,而对规划质量影响有限。这种权衡需要根据具体应用场景进行评估。
