运动变换器 (MTR)：自动驾驶中多模态运动预测的强大基线

本仓库是 NeurIPS 2022 论文（口头报告）"具有全局意图定位与局部运动精炼的运动变换器" 的官方实现。

作者：ShaoShuai Shi、Li Jiang、Dengxin Dai、Bernt Schiele

[MTR (arXiv)]    [MTR++ (arXiv)]

最新消息

[2023-06] MTR++ 的正式论文已发布至 arXiv:2306.17770，该模型支持多智能体运动预测，并在 Waymo 开放数据集上取得了最先进性能。

[2023-05] MTR++ 荣获冠军，赢得了 2023 年 Waymo 开放数据集运动预测挑战赛的冠军，排行榜请见 此处。

[2022-06] MTR 荣获冠军，赢得了 2022 年 Waymo 开放数据集运动预测挑战赛的冠军，官方公告请见 此处。

摘要

预测交通参与者的多模态未来行为对于机器人车辆做出安全决策至关重要。现有方法要么直接基于潜在特征预测未来轨迹，要么利用密集的目标候选点来识别智能体的目的地。前者由于所有运动模式都源自同一特征，收敛速度较慢；后者则因性能高度依赖目标候选点的密度而存在效率问题。本文提出 运动变换器 (MTR) 框架，将运动预测建模为全局意图定位与局部运动精炼的联合优化过程。MTR 不使用目标候选点，而是通过采用一组可学习的运动查询对来引入空间意图先验。每对运动查询负责特定运动模式的轨迹预测与精炼，从而稳定训练过程并促进更优的多模态预测。实验表明，MTR 在边缘和联合运动预测挑战中均达到最先进水平，在 Waymo 开放运动数据集的排行榜上位居第一。

亮点

MTR 代码库

• 具有清晰代码结构的最先进性能，易于扩展
• 非常简单的上下文编码器，用于建模智能体与地图的关系
• 基于意图点的可学习查询的运动解码器
• 使用高斯混合模型的损失函数，用于多模态运动预测
• 清晰的 Waymo 开放运动数据集数据处理与组织
• 包含官方 Waymo 运动评估 API 的本地评估工具

方法

• 简单：纯基于 Transformer 的上下文编码器和运动解码器
• 高效：仅用少量可学习的意图查询即可建模多模态未来预测
• 准确：在 Waymo 运动预测排行榜上排名第一（最后更新：2023年2月）

快速入门

• 安装

• 数据集准备

• 训练与评估

主要结果

在 Waymo 开放运动数据集验证集上的性能

模型	训练集	minADE	minFDE	误检率	mAP
MTR	20%	0.6697	1.3712	0.1668	0.3437
MTR	100%	0.6046	1.2251	0.1366	0.4164
MTR-e2e	100%	0.5160	1.0404	0.1234	0.3245

在 Waymo 开放运动数据集测试集上的性能

模型	训练集	minADE	minFDE	误检率	mAP
MTR	100%	0.6050	1.2207	0.1351	0.4129
MTR-A (集成)	100%	0.5640	1.1344	0.1160	0.4492

引用

如果您在研究中发现本工作有所帮助，请考虑引用：

@article{shi2022motion,
  title={Motion transformer with global intention localization and local movement refinement},
  author={Shi, Shaoshuai and Jiang, Li and Dai, Dengxin and Schiele, Bernt},
  journal={Advances in Neural Information Processing Systems},
  year={2022}
}

@article{shi2023mtr,
  title={MTR++: Multi-Agent Motion Prediction with Symmetric Scene Modeling and Guided Intention Querying},
  author={Shi, Shaoshuai and Jiang, Li and Dai, Dengxin and Schiele, Bernt},
  journal={arXiv preprint arXiv:2306.17770},
  year={2023}
}

@article{shi2022mtra,
  title={MTR-A: 1st Place Solution for 2022 Waymo Open Dataset Challenge--Motion Prediction},
  author={Shi, Shaoshuai and Jiang, Li and Dai, Dengxin and Schiele, Bernt},
  journal={arXiv preprint arXiv:2209.10033},
  year={2022}
}

MTR (Motion Transformer) 快速上手指南

MTR 是 NeurIPS 2022 提出的自动驾驶多模态运动预测强基线模型，曾在 Waymo 开放数据集挑战赛中夺冠。本指南帮助开发者快速搭建环境并运行模型。

环境准备

在开始之前，请确保满足以下系统和依赖要求：

• 操作系统: Linux (推荐 Ubuntu 18.04/20.04)
• Python: 3.7 或更高版本
• GPU: NVIDIA GPU (支持 CUDA)，显存建议 16GB 以上以训练完整模型
• 核心依赖:
  • PyTorch >= 1.8
  • CUDA Toolkit (需与 PyTorch 版本匹配)
  • OpenCV
  • NumPy, SciPy, YAML

国内加速建议：安装 Python 依赖时，推荐使用清华或阿里镜像源以提升下载速度。

pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple

安装步骤

1. 克隆仓库

   git clone https://github.com/sshaoshuai/MTR.git
   cd MTR
   

2. 创建虚拟环境并安装依赖

   conda create -n mtr python=3.8 -y
   conda activate mtr
   
   # 安装 PyTorch (请根据实际 CUDA 版本调整，此处以 CUDA 11.1 为例)
   pip install torch==1.9.0+cu111 torchvision==0.10.0+cu111 -f https://download.pytorch.org/whl/torch_stable.html
   
   # 安装其他项目依赖
   pip install -r requirements.txt
   

3. 编译扩展模块 MTR 包含自定义的 CUDA 算子，需要编译安装：

   python setup.py develop
   

4. 数据集准备 下载 Waymo Open Motion Dataset，并按照 docs/DATASET_PREPARATION.md 中的说明进行预处理和目录组织。

   • 将处理后的数据放置在 data/waymo/ 目录下。

基本使用

以下命令展示如何加载预配置进行模型评估（假设已准备好验证集数据）。

单卡评估示例

使用官方提供的配置文件对验证集进行评估：

python test.py --cfg_file tools/cfgs/waymo/mtr+100_percent_data.yaml --batch_size 1 --eval_all

训练示例

启动模型训练（单卡）：

python train.py --cfg_file tools/cfgs/waymo/mtr+100_percent_data.yaml --batch_size 4

多卡分布式训练

使用 torch.distributed.launch 进行多卡加速训练（例如使用 8 张卡）：

python -m torch.distributed.launch --nproc_per_node=8 train.py --cfg_file tools/cfgs/waymo/mtr+100_percent_data.yaml --batch_size 4 --launcher pytorch

注意：具体的超参数、数据路径配置请查阅 tools/cfgs/ 目录下的 .yaml 文件。更多详细的高级用法请参考项目根目录下的 docs/TRAIN_EVAL.md 文档。