[论文解读] Waymax: An Accelerated, Data-Driven Simulator for Large-Scale Autonomous Driving Research

Cole Gulino, Justin Fu+20 more|arXiv (Cornell University)|Oct 12, 2023

Autonomous Vehicle Technology and Safety被引用 22

一句话总结

Waymax 是一个可微分的、硬件加速的多智能体自动驾驶仿真器，基于现实世界数据（Waymo Open Motion Dataset）构建，旨在实现闭环与开环设置下的快速、图内训练与规划代理的基准测试。

ABSTRACT

Simulation is an essential tool to develop and benchmark autonomous vehicle planning software in a safe and cost-effective manner. However, realistic simulation requires accurate modeling of nuanced and complex multi-agent interactive behaviors. To address these challenges, we introduce Waymax, a new data-driven simulator for autonomous driving in multi-agent scenes, designed for large-scale simulation and testing. Waymax uses publicly-released, real-world driving data (e.g., the Waymo Open Motion Dataset) to initialize or play back a diverse set of multi-agent simulated scenarios. It runs entirely on hardware accelerators such as TPUs/GPUs and supports in-graph simulation for training, making it suitable for modern large-scale, distributed machine learning workflows. To support online training and evaluation, Waymax includes several learned and hard-coded behavior models that allow for realistic interaction within simulation. To supplement Waymax, we benchmark a suite of popular imitation and reinforcement learning algorithms with ablation studies on different design decisions, where we highlight the effectiveness of routes as guidance for planning agents and the ability of RL to overfit against simulated agents.

研究动机与目标

• 在最小化仿真到真实的差距的同时，促进自动驾驶规划研究的快速、真实感仿真。
• 提供一个可微分的、硬件加速的多智能体仿真器，能够在图内运行以在加速器上进行训练。
• 实现在线训练与评估，涵盖学习型和硬编码行为模型以及真实交互代理。
• 在多样化、以真实数据初始化的场景中对模仿学习与强化学习规划算法进行基准测试。

提出的方法

• 使用来自 Waymo Open Motion Dataset 的真实世界驾驶数据来初始化并回放多样化的多智能体场景。
• 实现一个用 JAX 构建、在 GPU/TPU 上运行并支持图内训练的可微分、硬件加速的仿真器。
• 提供多智能体和自车为中心的规划环境，包含路线/目标信息以及反应型仿真代理（如 IDM）。
• 通过 delta 或自行车动作空间定义动态对象动力学，并支持多种观测函数以及路径在路线内/外的概念。
• 提供在闭环中计算的一组评测指标（如路线进度、偏离路线、出线/越线、碰撞、运动学不可行性、位移误差）以及多种基线代理（专家、BC、DQN、Wayformer）。

实验结果

研究问题

• RQ1 一个可扩展、数据驱动的仿真器如何利用真实世界日志再现现实的多智能体行驶交互？
• RQ2 硬件加速的可微分仿真是否能促进自动驾驶规划策略的更快训练与评估？
• RQ3 路由条件化和交互式仿真代理对闭环设置中规划性能有何影响？
• RQ4 在 Waymax 中对比反应型仿真器和记录轨迹时，模仿学习与强化学习基线的表现如何？

主要发现

• Waymax 具有高吞吐量，在 GPU 上对批量大小为 16 及以上的单步执行，步长运行时间低于 3 ms，且单步执行频率超过 1000 Hz。
• 路由条件化在模仿学习中提升路线遵循并降低出线/越线和碰撞率。
• 采用离散自行车动作的模仿学习在与专家基线相比具有竞争力，且路由条件化的 BC 能实现强烈的路线遵循。
• 对 IDM 互动代理进行强化学习代理的训练可能不如对记录代理进行有效，暗示对交互式仿真器行为的过拟合。
• 与 IDM 仿真器的对比评估能减少碰撞，但由于更容易的协作性动力学，可能阻碍强化学习代理的学习。
• Waymax 使在加速器上的图内训练成为可能，从而实现端到端的可微分仿真工作流。