Files
MAGAIL4AutoDrive/README.md
2026-03-02 10:58:20 +08:00

7.6 KiB
Raw Blame History

MAGAIL4AutoDrive

基于 MetaDrive 仿真器和 Waymo Open Motion Dataset 的自动驾驶多智能体模仿学习MAGAIL与行为克隆BC训练系统。

本项目旨在从真实的 Waymo 驾驶数据中提取专家轨迹并通过模仿学习Imitation Learning训练能够适应复杂交互场景的自动驾驶策略。

目录结构

MAGAIL4AutoDrive/
├── Algorithm/                 # 强化学习与模仿学习算法实现
│   ├── policy.py              # 基础策略网络 (MLP 等)
│   ├── ppo.py                 # PPO 算法实现
│   ├── magail.py              # MAGAIL 算法核心逻辑
│   ├── disc.py                # 判别器 (Discriminator) 网络
│   └── ...
├── Env/                       # 仿真环境封装 (MetaDrive Wrapper)
│   ├── bc_env.py              # BCScenarioEnv45 维观测BC/MAGAIL 共用)
│   ├── bc_ego_replay_env.py   # BCEgoReplayEnv单智能体 BC 评估(仅 ego 受控)
│   ├── scenario_env.py       # 多智能体基础场景环境
│   ├── expert_replay_env.py  # 专家轨迹回放环境(数据生成与回放)
│   ├── inverse_dynamics.py   # 逆动力学模块 (轨迹 -> 动作)
│   ├── simple_idm_policy.py  # ConstantVelocityPolicy 占位策略
│   └── ...
├── dataset/                   # 数据集加载器
│   ├── loader.py              # 主流水线load_expert_pkl、MAGAILExpertDataset
│   └── expert_dataset.py      # 可选 107 维/5 维管线
├── scripts/                   # 工具脚本(数据、回放、可视化、分析)
│   ├── generate_expert_data.py    # 从 Waymo 生成专家 (obs, act) pkl
│   ├── visualize.py              # 可视化统一入口replay / policy / trajectory
│   ├── analyze_expert_data.py    # 数据分布分析
│   ├── launch_tensorboard.py     # 启动 TensorBoard
│   ├── README.md                  # 脚本用法说明
│   └── ...
├── data/                      # 数据目录(相对路径)
│   ├── exp_filtered/          # Waymo 场景数据
│   ├── training_data/        # 专家 pkl 输出generate_expert_data
│   └── trajectories/         # 其他轨迹 pkl如 expert_dataset 输出)
├── models/                    # 模型保存目录(相对路径)
│   ├── bc/                    # BC 模型 (.pt)
│   └── magail/                # MAGAIL 模型 (*_actor.pth, *_critic.pth)
├── logs/                      # 训练日志 (TensorBoard)
│   ├── bc/
│   └── magail/
├── train_bc.py                # [根目录] BC 训练
├── train_magail.py            # [根目录] MAGAIL 训练
└── README.md

路径约定(相对项目根)

  • 数据Waymo 场景 data/exp_filtered;专家 pkl data/training_data;其他轨迹 data/trajectories
  • 模型BC models/bc/MAGAIL models/magail/
  • 日志TensorBoard 写入 logs/bc/logs/magail/

所有默认路径均为相对项目根,便于在不同设备上复用。

数据处理流程

从 Waymo Motion 原始数据到本项目训练用专家 pkl依次为

1) 下载 Waymo MotionTFRecord
安装 gsutil 并登录 Google 账号后,例如只下载 training_20s

gsutil -m cp -r "gs://waymo_open_dataset_motion_v_1_2_0/uncompressed/scenario/training_20s" ./waymo/

2) ScenarioNet ConvertTFRecord → ScenarioNet 场景库)
需安装 ScenarioNet、MetaDrive 及 TensorFlow 2.11、protobuf 3.20;转换时不用 GPU。

python -m scenarionet.convert_waymo -d data/exp_converted --raw_data_path ./waymo/training_20s --num_workers 64

3) ScenarioNet Filter按需筛选场景
从 convert 得到的场景库中筛掉含红绿灯、天桥等场景,输出到如 data/exp_filtered。具体命令以 ScenarioNet 文档为准Operations → Filter

4) 本项目:生成专家 pkl
使用筛选后的场景目录,生成训练用 pkl 到 data/training_data

  • 多智能体(所有受控车轨迹,输出 expert_data_{start_index}_{num_scenarios}.pkl
python scripts/generate_expert_data.py --data_dir data/exp_filtered --output_dir data/training_data --num_scenarios 100 --start_index 0
  • 单智能体(仅 ego 车轨迹,输出 expert_data_ego_{start_index}_{num_scenarios}.pkl,用于单智能体 BC
python scripts/generate_expert_data.py --data_dir data/exp_filtered --output_dir data/training_data --num_scenarios 100 --start_index 0 --ego_only

核心工作流

1. 数据准备

使用 scripts/generate_expert_data.py 将 Waymo 数据转换为训练用 .pkl,输出到 data/training_data/

  • 多智能体
python scripts/generate_expert_data.py --data_dir data/exp_filtered --output_dir data/training_data --num_scenarios 100 --start_index 0
  • 单智能体(仅 ego
python scripts/generate_expert_data.py --data_dir data/exp_filtered --output_dir data/training_data --num_scenarios 100 --start_index 0 --ego_only

2. 行为克隆 (BC)

BC 支持两种模式:多智能体(默认,所有受控车共用同一策略)与 单智能体(仅 ego 车,评估时其他车按专家轨迹回放)。

  • 多智能体训练(模型保存到 models/bc/,日志到 logs/bc/
python train_bc.py --expert_data_path data/training_data/expert_data_0_50.pkl --epochs 100
  • 单智能体训练(使用 ego-only 数据,评估时仅 ego 受策略控制,其他车专家回放):
python train_bc.py --expert_data_path data/training_data/expert_data_ego_0_50.pkl --epochs 100 --single_agent
  • 可视化python scripts/visualize.py policy --policy_type bc --model_path models/bc/policy_best.pt
    仅自车用策略、其他车回放(单智能体可视化):加 --ego_only,例如
    python scripts/visualize.py policy --policy_type bc --model_path models/bc/policy_best.pt --ego_only --num_scenarios 1

3. 多智能体对抗模仿学习 (MAGAIL)

  • 训练python train_magail.py(模型保存到 models/magail/,日志到 logs/magail/
  • 可视化python scripts/visualize.py policy --policy_type magail --model_path models/magail/model_50_actor.pth

4. 可视化统一入口

可视化统一使用 scripts/visualize.py,子命令:replay(场景回放)、policyBC/MAGAIL 策略)、trajectory(专家轨迹 2D 动画)。详见 scripts/README.md

文件与模块职责

根目录脚本

  • train_bc.pyBC 训练,从 dataset.loader 加载专家 pkl模型与日志写入 models/bc/logs/bc/
  • train_magail.pyMAGAIL 训练,环境使用 BCScenarioEnv45 维),从 dataset.loader 加载专家数据,模型与日志写入 models/magail/logs/magail/

Env 模块

  • Env/bc_env.pyBCScenarioEnv45 维观测Ego 5 维 + 10 邻居×4 维BC 与 MAGAIL 训练/评估共用
  • Env/bc_ego_replay_env.pyBCEgoReplayEnv,单智能体 BC 评估环境,仅 ego 受策略控制,其他车按专家轨迹回放
  • Env/scenario_env.pyMultiAgentScenarioEnv 基类Waymo 场景加载与步进
  • Env/expert_replay_env.py:专家轨迹回放与逆动力学动作,供 generate_expert_data.py 与回放可视化
  • Env/inverse_dynamics.py:轨迹 → 油门/转向动作

Algorithm 模块

  • Algorithm/policy.pyStateIndependentPolicyBC 使用的 MLP 策略

scripts 目录

工具脚本用途与用法见 scripts/README.md