本文经自动驾驶之心公众号授权转载，转载请联系出处。
论文链接：https://arxiv.org/pdf/2410.15774
项目网页：https://tsinghua-mars-lab.github.io/StateTransformer/
代码开源：https://github.com/Tsinghua-MARS-Lab/StateTransformer
主要内容：
大型实际驾驶数据集推动了有关自动驾驶数据驱动运动规划器的各个方面的研究，包括数据增强、模型架构、奖励设计、训练策略和规划器架构。在处理复杂和少样本情况下，这些方法有较好的表现。但是由于设计过于复杂或训练范式的问题，这些方法在规划性能上的泛化能力有限。在本文中，我们回顾并比较了以前的方法，重点关注泛化能力。实验结果显示，随着模型的适当扩展，许多设计元素变得冗余。我们介绍了StateTransformer-2 (STR2)，这是一种可扩展的、仅使用解码器的运动规划器，它结合了Vision Transformer (ViT) 编码器和混合专家(MoE) 的Transformer架构。MoE骨干通过训练期间的专家路由解决了模态崩溃和奖励平衡问题。在NuPlan数据集上的大量实验表明，我们的方法在不同测试集和闭环模拟中比以前的方法具有更好的泛化能力。此外，我们评估了其在真实城市驾驶场景中的可扩展性，显示出随着数据和模型规模的增长其一致的准确性提升。
下面我们详细解释STR2的模型设计。我们选择的不是向量化的输入，而是栅格化图片的输入，可以方便进行规模化训练。输入还加入了聚类的轨迹作为引导线，结果可以看到在大曲率的场景表现会更合理一些。输入经过一个使用MoE架构增强的Transformer模型，自回归的方式输出引导线类别，关键点以及轨迹。
ViT编码器。我们采用仅解码的ViT图像编码器，以实现更好的可扩展性和性能，它由堆叠的12层Transformer组成。栅格化的图像被切分成16个小块。我们选择GeLU作为ViT编码器的激活函数。
Mixture-of-Expert。语言建模任务要求模型从复杂且通常具有统计争议的专家数据奖励中学习和实现平衡。受MoE模型在语言建模任务上泛化结果的启发，我们将GPT-2骨干网络替换为MoE骨干网络用于序列建模。MoE层通过专用内核和专家并行（EP）提供了更好的内存效率。我们还利用了Flash Attention2 和数据并行（DP）以提高训练效率。
自回归。在生成序列中我们添加了聚类轨迹作为嵌入特征用于模态分类，并使用交叉熵损失。我们使用K-Means聚类，从0.7百万个随机选择的动态可行轨迹中按其时空距离提取了512个候选轨迹。每条归一化的轨迹包括未来8秒的80个轨迹点（x, y和偏航角）。
liauto数据集上的规模化实验。我们采用了liauto数据集进行了scaling law的探索，liauto数据集是一个工业级的超大规模现实世界驾驶数据集。该数据集包括车道级导航地图和来自7个RGB摄像头、1个LiDAR和1个毫米波雷达（MMWR）的传感器设置的跟踪结果。我们选择了过去6个月内收集的城市驾驶场景，其中没有任何人工标注。我们筛选出错误的导航路线，因为这些路线与实际的未来驾驶轨迹不匹配。最终，我们将所有驾驶日志重新整理为长达10秒的训练和测试样本，其中包括过去的2秒和未来的8秒。最终的训练数据集拥有超过1b训练样本。实验结果（如图2）可以看到随着数据规模的增加以及模型参数的增加，test loss都有下降的趋势。均衡考虑训练成本和收益的关系，我们最后采用的为800m的模型。
从图3 nuplan闭环仿真的结果可以看到STR2取得了全面SOTA。结果来看专家轨迹的NR分数较高，R的分数却很低，说明专家轨迹没有一个适应环境变化的能力，仿真环境和实车有一定的区别。结果可以看出PDM-Hybrid相比于其他的方法的R得分明显高于NR，说明生成-评估的范式在模型泛化性能上表现优异。我们借鉴了PDM的生成-评估范式，基于STR2模型的输出结果进行了候选轨迹生成，再经过PDM打分器输出得分最高的轨迹。
图4可视化结果可以看出PDM-Hybrid轨迹因为基于当前车道中心线，没有主动变道和绕障的能力，我们的模型由于泛化性能强，所以很多场景下表现比PDM-Hybrid合理很多。本工作在nuplan testhard 数据集上的实验结果证明了生成-评估范式的合理性。对于生成-评估范式，在后续工作中我们还可以有更多尝试，例如如何在模型层面输出更合理的候选轨迹，如何将未来的不确定性考虑进打分器中，同时可以考虑将本文的方法作为真值标注的一种方式，帮助模型朝着更合理的方向迭代。为了感受模型的泛化性能我们将nuplan数据训练得到的模型直接应用于liauto数据集上推理，结果表现良好，对于动态障碍物的避让交互等都有合理的输出。
