Waymax 安装与使用教程（一）

Waymax 介绍

模拟是以安全和经济的方式开发和基准测试自主车辆规划软件的重要工具。然而，实际仿真需要准确建模微妙且复杂的多智能体互动行为。为了解决这些挑战，Waymax，一种新型的数据驱动自动驾驶多智能体场景模拟器，专为大规模仿真和测试而设计。Waymax使用公开发布的现实世界驾驶数据（例如Waymo Open Motion Dataset ）来初始化或回放各种多智能体仿真场景。它完全运行在TPUs/GPUs等硬件加速器上，并支持图内仿真进行训练，适用于现代大规模分布式机器学习工作流程。为了支持在线训练和评估，Waymax包含了几个学习得到或硬编码的行为模型，在仿真中实现了逼真互动。作为补充，Waymax对一套流行的模仿学习和强化学习算法进行了基线研究，并针对不同设计决策进行消融研究，在其中突出显示路线作为规划代理指导以及RL过度拟合于虚拟代理方面的有效性。

Waymax 是一个基于 JAX 的自动驾驶研究模拟器。它是一个轻量级、多代理系统，基于 Waymo 开放运动数据集。Waymax 旨在支持自动驾驶行为研究的各个方面，包括闭环规划和模拟代理研究以及开环行为预测。为了简化研究过程，Waymax 使用边界框来表示物体（如车辆和行人），而非原始传感器输出。此外，其整个架构均用 JAX 编写，便于在硬件加速器（如 GPU 和 TPU）上分发和部署​。

原文链接：

github链接：

对比其他模拟器：

相关驾驶模拟器的比较（按时间顺序）。多智能体：模拟多个智能体。支持的智能体数量在括号中（如果可用）。加速度：用于硬件（GPU/TPU）加速的图内编译。传感器模拟：传感器（例如相机、激光雷达和雷达）输入模拟。专家数据：通过专家策略收集的人类演示或轨迹回放数据。仿真代理：用于模拟对象（例如其他车辆）的代理模型。真实数据：真实世界驾驶数据。路线/目标：“-”表示没有提供路线或目标；“航点”表示从轨迹中采样得到的位置；“方向”表示包括左转、直行和右转在内的离散行驶方向；“目标点”表示目标位置。

主要指标：

这些指标在Waymax中是在闭环中计算的，意味着它们是通过在模拟中运行代理来计算的，而不是在开环中基于每个时间步骤单独计算。这种方法更能反映代理在模拟环境中的实际表现。

Software API：

1. 环境接口：Waymax主要以部分可观察的随机游戏的形式与用户交互。它遵循Brax的设计，仅定义了初始化和转换函数。这种无状态的设计允许通过JAX的JIT编译器和功能库进行高效优化，并容易实现需要回溯等控制算法。

2. 硬件加速和图内训练：Waymax支持在GPU和TPU上的硬件加速，以及将训练和模拟合并在同一个计算图（即图内训练）中，这使得训练和模拟可以完全在加速器上进行，无需通过主机进行通信。这些特性得益于Waymax完全使用JAX库编写，该库将操作转换为XLA，一个线性代数指令集和优化编译器。

3. 单代理和多代理模拟：Waymax支持多代理模拟，允许用户控制场景中的任意对象，也支持单代理工作流程，其中单个自动驾驶车辆代理被训练以控制其他车辆。

Waymax安装：

可以通过 Python 的 pip 包管理器进行。安装过程包括升级 pip，然后直接从其 GitHub 仓库安装 Waymax。需要运行的命令如下：

pip install --upgrade pip
pip install git+https://github.com/waymo-research/waymax.git@main#egg=waymo-waymax

如果需要，建议参考 JAX 文档，了解如何在 GPU/CUDA 支持下设置 ​​。

Jax安装（On Mac）：

1. Setup

python3 -m venv ~/jax-metal
source ~/jax-metal/bin/activate
python -m pip install -U pip
python -m pip install numpy wheel ml-dtypes==0.2.0

2. 安装

python -m pip install jax-metal

3. 测试

python -c 'import jax; print(jax.numpy.arange(10))'

在组件方面，Waymax 包括用于加载开放运动数据、可视化、计算常见指标、智能模拟代理以及适配常见强化学习接口（如 dm-env）的一系列库。这些库可以作为独立模块使用，也可以一起用于完整的闭环模拟体验​​。

后续将继续更新使用教程。