huangfu/MAGAIL4AutoDrive
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

mixed_training

Browse Source
This commit is contained in:
huangfu
2026-01-17 14:50:20 +08:00
parent c94571ddaa
commit 4dbea5f0a6
30 changed files with 2075 additions and 2006 deletions
Download Patch File Download Diff File Expand all files Collapse all files

279
Env/ENVIRONMENT_INFO.md
View File

@@ -1,279 +0,0 @@
# 多智能体场景环境详细说明
## 1. 观测信息详解
### 观测向量结构（总维度：107维）
每个智能体的观测向量包含以下信息：
```python
观测向量 = [
# 1. 车辆状态信息 (5维)
position_x, # 车辆X坐标
position_y, # 车辆Y坐标
velocity_x, # X方向速度
velocity_y, # Y方向速度
heading_theta, # 朝向角度
# 2. 前向激光雷达 (80维)
lidar_1, # 第1个激光束的距离
lidar_2, # 第2个激光束的距离
...
lidar_80, # 第80个激光束的距离
# 范围：30米，用于前方障碍物检测
# 3. 侧向激光雷达 (10维)
side_lidar_1, # 第1个侧向激光束的距离
...
side_lidar_10, # 第10个侧向激光束的距离
# 范围：8米，用于侧方障碍物检测
# 4. 车道线检测 (10维)
lane_line_1, # 第1个车道线检测距离
...
lane_line_10, # 第10个车道线检测距离
# 范围：3米，用于车道线识别
# 5. 导航信息 (2维)
destination_x, # 目标点X坐标
destination_y, # 目标点Y坐标
]
```
### 观测信息说明
1. **车辆状态 (5维)**
- 位置：全局坐标系下的(x, y)
- 速度：车辆在全局坐标系下的速度分量
- 朝向：车辆的航向角（弧度）
2. **激光雷达 (100维)**
- 前向80束：覆盖前方视野，检测动态和静态障碍物
- 侧向10束：检测侧方物体，用于变道等操作
- 车道线10束：专门检测车道线位置
3. **导航信息 (2维)**
- 目标位置：从专家数据中提取的车辆最终位置
**总维度：5 + 80 + 10 + 10 + 2 = 107维**
---
## 2. 多场景加载逻辑
### 2.1 场景加载机制
MetaDrive的ScenarioEnv通过以下配置管理多场景：
```python
config = {
"data_directory": "path/to/dataset", # 包含dataset_mapping.pkl的目录
"num_scenarios": 3, # 场景总数（从mapping文件读取）
"start_scenario_index": 0, # 起始场景索引
"sequential_seed": True, # 是否顺序切换场景
}
```
### 2.2 场景切换逻辑
#### 场景索引管理
- **初始化**：环境启动时从`start_scenario_index`开始
- **顺序模式**（`sequential_seed=True`）：
- 每次`env.reset()`后自动切换到下一个场景
- 循环顺序：场景0 → 场景1 → 场景2 → 场景0 ...
#### 示例流程
```python
env = MultiAgentScenarioEnv(config={
"data_directory": "path/to/dataset", # 假设有3个场景
"start_scenario_index": 0,
"sequential_seed": True,
})
obs = env.reset(0) # 使用场景0
# ... 运行场景0 ...
obs = env.reset() # 自动切换到场景1
# ... 运行场景1 ...
obs = env.reset() # 自动切换到场景2
# ... 运行场景2 ...
obs = env.reset() # 循环回到场景0
```
### 2.3 当前场景信息获取
可以通过以下方式查看当前场景信息：
```python
# 获取当前场景索引
current_scenario = env.engine.current_seed
# 获取场景总数
total_scenarios = env.config["num_scenarios"]
# 查看场景ID
scenario_id = env.engine.traffic_manager.current_scenario_id
```
### 2.4 手动指定场景
如果需要固定使用某个场景：
```python
# 方法1：在reset时指定
obs = env.reset(seed=1) # 使用场景1
# 方法2：配置固定场景
config = {
"start_scenario_index": 2, # 从场景2开始
"sequential_seed": False, # 禁用自动切换
}
```
---
## 3. 车辆观测获取机制
### 3.1 观测获取流程
```
step() 被调用
↓
更新所有车辆物理状态
↓
_spawn_controlled_agents() # 生成新车辆（按时间步）
↓
_get_all_obs() # 获取所有车辆观测
↓
遍历 controlled_agents:
├─ 获取车辆状态 (position, velocity, heading)
├─ 调用激光雷达传感器
├─ 组装观测向量
└─ 添加到 obs_list
↓
返回 obs_list（包含所有车辆的观测）
```
### 3.2 观测返回格式
```python
obs = env.reset()
# obs 是一个列表，每个元素是一个车辆的观测向量
obs = [
[obs_vehicle_0], # 第1辆车的107维观测
[obs_vehicle_1], # 第2辆车的107维观测
...
]
# 在step中
obs, rewards, dones, infos = env.step(actions)
# obs格式相同，但只包含当前存活的车辆
```
### 3.3 观测一致性保证
1. **物理状态同步**
- 所有车辆在同一物理时间步获取观测
- 保证观测的时间一致性
2. **传感器独立性**
- 每个车辆有独立的传感器
- 激光雷达从各自位置发射
3. **动态车辆管理**
- 新车辆在生成时立即获取观测
- 观测列表动态更新
---
## 4. 常见问题解答
### Q1: 为什么观测维度是107而不是其他？
**A**:
- 车辆状态: 5维 (x, y, vx, vy, heading)
- 前向激光雷达: 80维
- 侧向激光雷达: 10维
- 车道线检测: 10维
- 目标位置: 2维
- **总计: 5 + 80 + 10 + 10 + 2 = 107维**
### Q2: 如何确保多场景都被使用？
**A**: 设置`sequential_seed=True`，环境会自动循环遍历所有场景。
### Q3: 车辆在不同时间步生成，如何获取观测？
**A**: 每次调用`step()`时：
1. 先检查是否有新车辆需要生成（`_spawn_controlled_agents`）
2. 为所有现存车辆（包括新生成的）获取观测
3. 返回的obs_list包含所有当前存活车辆的观测
### Q4: 如果场景中车辆数量不同怎么办？
**A**:
- 观测列表长度动态调整
- 使用`max_controlled_vehicles`可限制最大车辆数
- 使用`filter_offroad_vehicles`可过滤无效车辆
### Q5: 观测数据的坐标系是什么？
**A**:
- 位置/速度/目标：**全局坐标系**（世界坐标）
- 激光雷达：**车辆局部坐标系**（以车辆为中心）
---
## 5. 配置建议
### 5.1 训练配置
```python
config = {
"data_directory": "path/to/dataset",
"sequential_seed": True, # 循环使用所有场景
"filter_offroad_vehicles": True, # 过滤无效车辆
"max_controlled_vehicles": 20, # 限制车辆数防止过载
"inherit_expert_velocity": False, # 训练时不继承速度
"horizon": 300, # 每场景运行300步
}
```
### 5.2 评估配置
```python
config = {
"data_directory": "path/to/dataset",
"sequential_seed": False, # 固定场景
"start_scenario_index": 0, # 指定场景
"filter_offroad_vehicles": True,
"max_controlled_vehicles": None, # 不限制车辆数
"inherit_expert_velocity": True, # 评估时继承速度
"verbose_reset": True, # 输出详细信息
}
```
---
## 6. 调试技巧
### 6.1 查看当前场景信息
```python
# 在reset后查看
print(f"当前场景: {env.engine.current_seed}")
print(f"总场景数: {env.config['num_scenarios']}")
print(f"可控车辆数: {len(env.controlled_agents)}")
```
### 6.2 查看观测维度
```python
obs = env.reset()
print(f"车辆数量: {len(obs)}")
if len(obs) > 0:
print(f"单车辆观测维度: {len(obs[0])}")
```
### 6.3 启用详细日志
```python
config = {
"verbose_reset": True, # 重置时详细统计
"debug_lane_filter": True, # 车道过滤调试
}
```

413
Env/LOGGING_GUIDE.md
View File

@@ -1,413 +0,0 @@
# 日志记录功能使用指南
## 📋 概述
为所有运行脚本添加了日志记录功能，可以将终端输出同时保存到文本文件，方便后续分析和问题排查。
---
## 🎯 功能特点
1. **双向输出**：同时输出到终端和文件，不影响实时查看
2. **自动管理**：使用上下文管理器，自动处理文件开启/关闭
3. **灵活配置**：支持自定义文件名和日志目录
4. **时间戳命名**：默认使用时间戳生成唯一文件名
5. **无缝集成**：只需添加命令行参数，无需修改代码
---
## 🚀 快速使用
### 1. 基础用法
```bash
# 不启用日志（默认）
python Env/run_multiagent_env.py
# 启用日志记录
python Env/run_multiagent_env.py --log
# 或使用短选项
python Env/run_multiagent_env.py -l
```
### 2. 自定义文件名
```bash
# 使用自定义日志文件名
python Env/run_multiagent_env.py --log --log-file=my_test.log
# 测试脚本也支持
python Env/test_lane_filter.py --log --log-file=test_results.log
```
### 3. 组合使用调试和日志
```bash
# 测试脚本：调试模式 + 日志记录
python Env/test_lane_filter.py --debug --log
# 会生成类似：test_debug_20251021_123456.log
```
---
## 📁 日志文件位置
默认日志目录：`Env/logs/`
### 文件命名规则
| 脚本 | 默认文件名格式 | 示例 |
|------|---------------|------|
| `run_multiagent_env.py` | `run_YYYYMMDD_HHMMSS.log` | `run_20251021_143022.log` |
| `run_multiagent_env_fast.py` | `run_fast.log` | `run_fast.log` |
| `test_lane_filter.py` | `test_{mode}_YYYYMMDD_HHMMSS.log` | `test_debug_20251021_143500.log` |
**说明**：
- `YYYYMMDD_HHMMSS` 是时间戳（年月日_时分秒）
- `{mode}` 是测试模式（`standard` 或 `debug`）
---
## 📝 所有支持的脚本
### 1. run_multiagent_env.py（标准运行脚本）
```bash
# 不启用日志
python Env/run_multiagent_env.py
# 启用日志（自动生成时间戳文件名）
python Env/run_multiagent_env.py --log
# 自定义文件名
python Env/run_multiagent_env.py --log --log-file=run_test1.log
```
**日志位置**：`Env/logs/run_YYYYMMDD_HHMMSS.log`
---
### 2. run_multiagent_env_fast.py（高性能版本）
```bash
# 启用日志
python Env/run_multiagent_env_fast.py --log
# 自定义文件名
python Env/run_multiagent_env_fast.py --log --log-file=fast_test.log
```
**日志位置**：`Env/logs/run_fast.log`（默认）
---
### 3. test_lane_filter.py（测试脚本）
```bash
# 标准测试 + 日志
python Env/test_lane_filter.py --log
# 调试测试 + 日志
python Env/test_lane_filter.py --debug --log
# 自定义文件名
python Env/test_lane_filter.py --log --log-file=my_test.log
# 组合使用
python Env/test_lane_filter.py --debug --log --log-file=debug_run.log
```
**日志位置**：
- 标准模式：`Env/logs/test_standard_YYYYMMDD_HHMMSS.log`
- 调试模式：`Env/logs/test_debug_YYYYMMDD_HHMMSS.log`
---
## 💻 编程接口
如果您想在代码中直接使用日志功能：
```python
from logger_utils import setup_logger
# 方式1：使用上下文管理器（推荐）
with setup_logger(log_file="my_log.log", log_dir="logs"):
print("这条消息会同时输出到终端和文件")
# 运行您的代码
# ...
# 方式2：手动管理
from logger_utils import LoggerContext
logger = LoggerContext(log_file="custom.log", log_dir="output")
logger.__enter__() # 开启日志
print("输出消息")
logger.__exit__(None, None, None) # 关闭日志
```
---
## 📊 日志内容示例
### 标准运行
```
📝 日志记录已启用
📁 日志文件: Env/logs/run_20251021_143022.log
------------------------------------------------------------
💡 提示: 使用 --log 或 -l 参数启用日志记录
示例: python run_multiagent_env.py --log
自定义文件名: python run_multiagent_env.py --log --log-file=my_run.log
------------------------------------------------------------
[INFO] Environment: MultiAgentScenarioEnv
[INFO] MetaDrive version: 0.4.3
...
------------------------------------------------------------
✅ 日志已保存到: Env/logs/run_20251021_143022.log
```
### 调试模式
```
📝 日志记录已启用
📁 日志文件: Env/logs/test_debug_20251021_143500.log
------------------------------------------------------------
🐛 调试模式启用
============================================================
📍 场景信息统计:
- 总车道数: 123
- 有红绿灯的车道数: 0
⚠️ 场景中没有红绿灯！
🔍 开始车道过滤: 共 51 辆车待检测
...
------------------------------------------------------------
✅ 日志已保存到: Env/logs/test_debug_20251021_143500.log
```
---
## 🔧 高级配置
### 自定义日志目录
```python
from logger_utils import setup_logger
# 指定不同的日志目录
with setup_logger(log_file="test.log", log_dir="my_logs"):
print("日志会保存到 my_logs/test.log")
```
### 追加模式
```python
from logger_utils import setup_logger
# 追加到现有文件（而不是覆盖）
with setup_logger(log_file="test.log", mode='a'): # mode='a' 表示追加
print("这条消息会追加到文件末尾")
```
### 只重定向特定输出
```python
from logger_utils import LoggerContext
# 只重定向stdout，不重定向stderr
logger = LoggerContext(
log_file="test.log",
redirect_stdout=True, # 重定向标准输出
redirect_stderr=False # 不重定向错误输出
)
```
---
## 📋 命令行参数总结
| 参数 | 短选项 | 说明 | 示例 |
|------|--------|------|------|
| `--log` | `-l` | 启用日志记录 | `--log` |
| `--log-file=NAME` | 无 | 指定日志文件名 | `--log-file=test.log` |
| `--debug` | `-d` | 启用调试模式（test_lane_filter.py） | `--debug` |
### 参数组合
```bash
# 示例1：标准模式 + 日志
python Env/test_lane_filter.py --log
# 示例2：调试模式 + 日志
python Env/test_lane_filter.py --debug --log
# 示例3：调试 + 自定义文件名
python Env/test_lane_filter.py -d --log --log-file=my_debug.log
# 示例4：所有参数
python Env/test_lane_filter.py --debug --log --log-file=full_test.log
```
---
## 🛠️ 常见问题
### Q1: 日志文件在哪里？
**A**: 默认在 `Env/logs/` 目录下。如果目录不存在，会自动创建。
```bash
# 查看所有日志文件
ls -lh Env/logs/
# 查看最新的日志
ls -lt Env/logs/ | head -5
```
---
### Q2: 如何查看日志内容？
**A**: 使用任何文本编辑器或命令行工具：
```bash
# 方式1：使用cat
cat Env/logs/run_20251021_143022.log
# 方式2：使用less（可翻页）
less Env/logs/run_20251021_143022.log
# 方式3：查看末尾内容
tail -n 50 Env/logs/run_20251021_143022.log
# 方式4：实时监控（适合长时间运行）
tail -f Env/logs/run_20251021_143022.log
```
---
### Q3: 日志文件太多怎么办？
**A**: 可以定期清理旧日志：
```bash
# 删除7天前的日志
find Env/logs/ -name "*.log" -mtime +7 -delete
# 只保留最新的10个日志
cd Env/logs && ls -t *.log | tail -n +11 | xargs rm -f
```
---
### Q4: 日志会影响性能吗？
**A**: 影响很小，因为：
1. 文件I/O是异步的
2. 使用了缓冲区
3. 立即刷新确保数据不丢失
如果追求极致性能，建议训练时不启用日志，只在需要分析时启用。
---
### Q5: 可以同时记录多个脚本的日志吗？
**A**: 可以，每个脚本使用不同的日志文件：
```bash
# 终端1
python Env/run_multiagent_env.py --log --log-file=script1.log
# 终端2（同时运行）
python Env/test_lane_filter.py --log --log-file=script2.log
```
---
## 💡 最佳实践
### 1. 开发阶段
```bash
# 使用调试模式 + 日志，方便排查问题
python Env/test_lane_filter.py --debug --log
```
### 2. 长时间运行
```bash
# 启用日志，避免输出丢失
nohup python Env/run_multiagent_env.py --log > /dev/null 2>&1 &
# 查看实时输出
tail -f Env/logs/run_*.log
```
### 3. 批量实验
```bash
# 为每次实验使用不同的日志文件
for i in {1..5}; do
python Env/run_multiagent_env.py --log --log-file=exp_${i}.log
done
```
### 4. 性能测试
```bash
# 不启用日志，获得最佳性能
python Env/run_multiagent_env_fast.py
```
---
## 📖 相关文档
- `README.md` - 项目总览
- `DEBUG_GUIDE.md` - 调试功能使用指南
- `CHANGELOG.md` - 更新日志
---
## 🔍 技术细节
### 实现原理
1. **TeeLogger类**：实现同时写入终端和文件
2. **上下文管理器**：自动管理资源（文件打开/关闭）
3. **sys.stdout重定向**：拦截所有print输出
4. **即时刷新**：每次写入后立即刷新，确保数据不丢失
### 源代码
详见 `Env/logger_utils.py`
```python
# 简化示例
class TeeLogger:
def write(self, message):
self.terminal.write(message) # 输出到终端
self.log_file.write(message) # 写入文件
self.log_file.flush() # 立即刷新
```
---
## ✅ 总结
- ✅ 简单易用：只需添加 `--log` 参数
- ✅ 不影响输出：终端仍可实时查看
- ✅ 自动管理：文件自动开启/关闭
- ✅ 灵活配置：支持自定义文件名和目录
- ✅ 完整记录：包含所有调试信息
立即开始使用：
```bash
python Env/test_lane_filter.py --debug --log
```

BIN
Env/__pycache__/expert_replay_env.cpython-313.pyc Normal file
View File

Binary file not shown.

BIN
Env/__pycache__/expert_replay_env.cpython-39.pyc Normal file
View File

Binary file not shown.

BIN
Env/__pycache__/inverse_dynamics.cpython-313.pyc Normal file
View File

Binary file not shown.

BIN
Env/__pycache__/inverse_dynamics.cpython-39.pyc Normal file
View File

Binary file not shown.

BIN
Env/__pycache__/logger_utils.cpython-39.pyc Normal file
View File

Binary file not shown.

BIN
Env/__pycache__/replay_policy.cpython-39.pyc Normal file
View File

Binary file not shown.

BIN
Env/__pycache__/scenario_env.cpython-313.pyc
View File

Binary file not shown.

BIN
Env/__pycache__/scenario_env.cpython-39.pyc Normal file
View File

Binary file not shown.

BIN
Env/__pycache__/simple_idm_policy.cpython-313.pyc
View File

Binary file not shown.

BIN
Env/__pycache__/simple_idm_policy.cpython-39.pyc Normal file
View File

Binary file not shown.

116
Env/example_with_logging.py
View File

@@ -1,116 +0,0 @@
"""
日志记录功能示例
演示如何在自定义脚本中使用日志功能
"""
from logger_utils import setup_logger
from datetime import datetime
import time
def example_without_logging():
"""示例1：不使用日志"""
print("=" * 60)
print("示例1：普通输出（不记录日志）")
print("=" * 60)
print("这是普通的print输出")
print("只会显示在终端")
print("不会保存到文件")
print()
def example_with_logging():
"""示例2：使用日志记录"""
print("=" * 60)
print("示例2：使用日志记录")
print("=" * 60)
# 使用with语句，自动管理日志文件
with setup_logger(log_file="example_demo.log", log_dir="logs"):
print("✅ 这条消息会同时输出到终端和文件")
print("✅ 运行一些计算...")
for i in range(5):
print(f" 步骤 {i+1}/5: 处理中...")
time.sleep(0.1)
print("✅ 计算完成！")
print("日志文件已关闭")
print()
def example_custom_filename():
"""示例3：使用时间戳命名"""
print("=" * 60)
print("示例3：自动生成时间戳文件名")
print("=" * 60)
# log_file=None 会自动生成时间戳文件名
with setup_logger(log_file=None, log_dir="logs"):
print("文件名会自动包含时间戳")
print("适合批量实验，避免覆盖")
print()
def example_append_mode():
"""示例4：追加模式"""
print("=" * 60)
print("示例4：追加到现有文件")
print("=" * 60)
# 第一次写入
with setup_logger(log_file="append_test.log", log_dir="logs", mode='w'):
print("第一次写入：这会覆盖文件")
# 第二次写入（追加）
with setup_logger(log_file="append_test.log", log_dir="logs", mode='a'):
print("第二次写入：这会追加到文件末尾")
print()
def example_complex_output():
"""示例5：复杂输出（包含颜色、格式）"""
print("=" * 60)
print("示例5：复杂输出格式")
print("=" * 60)
with setup_logger(log_file="complex_output.log", log_dir="logs"):
# 模拟多种输出格式
print("\n📊 实验统计：")
print(" - 实验名称：车道过滤测试")
print(" - 开始时间：", datetime.now().strftime("%Y-%m-%d %H:%M:%S"))
print(" - 车辆总数：51")
print(" - 过滤后：45")
print("\n🚦 红绿灯检测：")
print(" ✅ 方法1成功：3辆")
print(" ✅ 方法2成功：2辆")
print(" ⚠️ 未检测到：40辆")
print("\n" + "="*50)
print("实验完成！")
print()
def main():
"""运行所有示例"""
print("\n" + "🎯 " + "="*56)
print("日志记录功能完整示例")
print("="*60 + "\n")
example_without_logging()
example_with_logging()
example_custom_filename()
example_append_mode()
example_complex_output()
print("="*60)
print("✅ 所有示例运行完成！")
print("📁 查看日志文件：ls -lh logs/")
print("="*60)
if __name__ == "__main__":
main()

527
Env/expert_replay_env.py Normal file
View File

@@ -0,0 +1,527 @@
import logging
import numpy as np
from collections import defaultdict
from metadrive.component.vehicle.vehicle_type import DefaultVehicle
from metadrive.type import MetaDriveType
from Env.scenario_env import MultiAgentScenarioEnv, PolicyVehicle
from Env.inverse_dynamics import InverseDynamics
class ExpertReplayEnv(MultiAgentScenarioEnv):
def __init__(self, config=None):
# Allow passing config without agent2policy since we don't use policies for replay
if config is None:
config = {}
# Ensure we don't simulate physics for the controlled agents in the traditional sense
# but we still need the engine to run
super().__init__(config, agent2policy={})
self.inverse_dynamics = InverseDynamics()
self.expert_tracks = {}
# Replay SDC/ego ("default_agent" in MetaDrive) as well; otherwise it will keep default action=0 and look stuck.
self.replay_sdc = self.config.get("replay_sdc", True)
self.sdc_track = None
self.sdc_vehicle = None
self.sdc_agent_id = "default_agent"
def reset(self, seed=None):
self.round = 0
if self.logger is None:
from metadrive.engine.logger import get_logger, set_log_level
self.logger = get_logger()
log_level = self.config.get("log_level", logging.INFO)
set_log_level(log_level)
self.lazy_init()
self._reset_global_seed(seed)
if self.engine is None:
raise ValueError("Broken MetaDrive instance.")
self.background_vehicles = {} # Vehicles that exist but are static/background
# Helper function to check if a position is on a valid lane
def is_on_lane(pos, map_manager, threshold=2.0):
# Check if point is close to any lane in the road network
# This can be expensive if checked for every point, so we check sample points
# or rely on lane index if available.
# Waymo tracks don't have lane index, just positions.
# We can use map.road_network.get_closest_lane_index(pos)
if map_manager is None or map_manager.current_map is None:
return True # If no map, assume valid
try:
# Use a larger search radius to catch slightly offset lanes
lane, lane_index = map_manager.current_map.road_network.get_closest_lane_index(pos, return_lane=True)
if lane is None:
return False
# Check lateral distance
long, lat = lane.local_coordinates(pos)
width = lane.width
# Allow being slightly off-lane (e.g. changing lanes)
# But parking lots are usually far from defined lanes in Waymo converted maps
if abs(lat) <= (width / 2 + threshold):
return True
return False
except:
return False
# --- MODIFIED SECTION START ---
# Capture expert tracks before they are cleaned
self.expert_tracks = {}
# Capture SDC track for ego replay (MetaDrive default agent)
self.sdc_track = None
self.sdc_vehicle = None
if self.replay_sdc and hasattr(self.engine, "traffic_manager"):
sdc_sid = self.engine.traffic_manager.sdc_scenario_id
self.sdc_track = self.engine.traffic_manager.current_traffic_data.get(sdc_sid, None)
_obj_to_clean_this_frame = []
self.car_birth_info_list = []
# Pre-filter: Check tracks against map AND check for static vehicles
for scenario_id, track in self.engine.traffic_manager.current_traffic_data.items():
if scenario_id == self.engine.traffic_manager.sdc_scenario_id:
continue
else:
if track["type"] == MetaDriveType.VEHICLE:
_obj_to_clean_this_frame.append(scenario_id)
valid = track['state']['valid']
if not valid.any():
continue
first_show = np.argmax(valid)
last_show = len(valid) - 1 - np.argmax(valid[::-1])
mid_show = (first_show + last_show) // 2
# 1. Lane check (existing logic)
points_to_check = [first_show, mid_show, last_show]
on_road_count = 0
is_valid_track = True
start_pos = track['state']['position'][first_show]
if not is_on_lane(start_pos, self.engine.map_manager, threshold=5.0): # 5m tolerance
mid_pos = track['state']['position'][mid_show]
if not is_on_lane(mid_pos, self.engine.map_manager, threshold=5.0):
is_valid_track = False
# 2. Static check
# Calculate total displacement and max speed
positions = track['state']['position'][valid.astype(bool)]
velocities = track['state']['velocity'][valid.astype(bool)]
total_displacement = 0
max_speed = 0
if len(positions) > 1:
total_displacement = np.linalg.norm(positions[-1] - positions[0])
max_speed = np.max(np.linalg.norm(velocities, axis=1))
is_static = False
if total_displacement < 5.0 and max_speed < 1.0: # Relaxed threshold: <5m move and <1m/s
is_static = True
# Decision logic:
# - If off-road AND static: Skip completely (don't even spawn as background)
# - If off-road but moving: Maybe keep? Or skip? Usually off-road moving is weird, skip.
# - If on-road but static: Spawn as BACKGROUND (visible but not controlled agent)
# - If on-road and moving: Spawn as CONTROLLED agent
if not is_valid_track:
# Skip off-road vehicles entirely (both static and moving off-road)
continue
if is_static:
# Add to background list, but NOT to car_birth_info_list (which is for controlled agents)
# We need a way to spawn them. Let's add a separate list.
self.background_vehicles[scenario_id] = {
'id': track['metadata']['object_id'],
'show_time': first_show,
'begin': (track['state']['position'][first_show, 0], track['state']['position'][first_show, 1]),
'heading': track['state']['heading'][first_show],
'end': (track['state']['position'][last_show, 0], track['state']['position'][last_show, 1]),
'scenario_id': scenario_id,
'length': track['state']['length'][first_show],
'width': track['state']['width'][first_show],
'valid': valid # Need validity to know when to show/hide
}
continue # Do not add to controlled list
# Store the full track for replay (only for controlled agents)
self.expert_tracks[scenario_id] = track
self.car_birth_info_list.append({
'id': track['metadata']['object_id'],
'show_time': first_show,
'begin': (track['state']['position'][first_show, 0], track['state']['position'][first_show, 1]),
'heading': track['state']['heading'][first_show],
'end': (track['state']['position'][last_show, 0], track['state']['position'][last_show, 1]),
'scenario_id': scenario_id, # Keep track of original ID to lookup tracks
'length': track['state']['length'][first_show],
'width': track['state']['width'][first_show]
})
for scenario_id in _obj_to_clean_this_frame:
self.engine.traffic_manager.current_traffic_data.pop(scenario_id)
# --- MODIFIED SECTION END ---
self.engine.reset()
self.reset_sensors()
self.engine.taskMgr.step()
self.lanes = self.engine.map_manager.current_map.road_network.graph
if self.top_down_renderer is not None:
self.top_down_renderer.clear()
self.engine.top_down_renderer = None
self.dones = {}
self.episode_rewards = defaultdict(float)
self.episode_lengths = defaultdict(int)
self.controlled_agents.clear()
self.controlled_agent_ids.clear()
# We skip calling super().reset() to avoid double reset
# But we need to ensure ScenarioEnv-specific setup is done if any.
# ScenarioEnv.reset() basically does engine.reset() and some cleanup.
# We covered most of it.
self._spawn_controlled_agents()
self._spawn_background_vehicles() # Initial spawn for background
# Ensure SDC/ego is moved to the correct initial expert state.
if self.replay_sdc:
self.sdc_vehicle = self.engine.agent_manager.active_agents.get(self.sdc_agent_id, None)
if self.sdc_vehicle is not None and self.sdc_track is not None:
valid = self.sdc_track["state"]["valid"]
t0 = int(np.argmax(valid)) if valid.any() else 0
pos0 = self.sdc_track["state"]["position"][t0]
heading0 = self.sdc_track["state"]["heading"][t0]
vel0 = self.sdc_track["state"]["velocity"][t0]
self.sdc_vehicle.set_position(pos0)
self.sdc_vehicle.set_heading_theta(heading0)
self.sdc_vehicle.set_velocity(vel0)
return self._get_all_obs()
def _spawn_background_vehicles(self):
# Spawn static/background vehicles
# Since they are static, we might just spawn them once if their show_time is 0
# But Waymo tracks have valid bits, they might appear/disappear.
# For optimization, if they are truly static (never move), we just spawn them when show_time matches.
# We need to track spawned background vehicles to remove them if they become invalid?
# Since we defined them as "static", they probably stay put.
# But validity might change (e.g. late spawn).
# For simplicity in this step, let's just iterate and spawn if time matches
for sid, car in self.background_vehicles.items():
if car['show_time'] == self.round:
# Spawn as a Traffic Vehicle (not PolicyVehicle), or just a static object?
# Using DefaultVehicle is fine, but don't add to controlled_agents
# Check duplication
bg_id = f"bg_{car['id']}"
# if bg_id in self.engine.obj_to_id: # obj_to_id might not be available in all versions
if bg_id in self.engine.agent_manager.active_agents:
continue
vehicle_config = {}
if 'length' in car and 'width' in car:
vehicle_config = {
"length": car['length'],
"width": car['width']
}
v = self.engine.spawn_object(
DefaultVehicle,
name=bg_id,
vehicle_config=vehicle_config,
position=car['begin'],
heading=car['heading']
)
# Set color to grey/dark to indicate background
v.set_velocity([0, 0])
# Maybe set color? MetaDrive vehicles random color.
# v.set_color(...) if supported
# Register as an active object but NOT controlled agent
# The engine manages it.
# CRITICAL: We need it in self.engine.agent_manager.active_agents for Observation?
# If we want it to be seen by Lidar/Observation, it needs to be an "agent" or "traffic".
# DefaultVehicle spawned this way is just an object.
# We should add it to traffic manager? Or just leave it as object?
# MultiAgentScenarioEnv._get_all_obs iterates self.engine.agent_manager.active_agents
# If we want it in observation, we must add it to active_agents OR iterate over all objects.
# Adding to active_agents is easier for compatibility.
self.engine.agent_manager.active_agents[bg_id] = v
# Store valid mask to remove it later if needed?
v.valid_mask = car['valid']
v.start_t = car['show_time']
def _update_background_vehicles(self):
# Remove background vehicles if they become invalid
# Or spawn new ones
self._spawn_background_vehicles()
# Check validity for existing
to_remove = []
for aid, v in self.engine.agent_manager.active_agents.items():
if aid.startswith("bg_"):
# Check validity
if hasattr(v, 'valid_mask'):
curr_step = self.round
if curr_step >= len(v.valid_mask) or not v.valid_mask[curr_step]:
to_remove.append(aid)
for aid in to_remove:
self.engine.agent_manager.active_agents.pop(aid, None)
# if aid in self.engine.obj_to_id:
# self.engine.clear_objects([self.engine.obj_to_id[aid]])
# Instead, we should find the object by ID and clear it.
# Since we don't track obj directly, we can't easily clear it without obj ref.
# Wait, active_agents stores the vehicle object.
# So we can just clear that object.
pass
# Re-iterate to clear objects properly
for aid in to_remove:
# We need to find the vehicle object to clear it.
# But we popped it from active_agents.
# Wait, we should get it before pop.
pass
def _update_background_vehicles(self):
# Remove background vehicles if they become invalid
# Or spawn new ones
self._spawn_background_vehicles()
# Check validity for existing
to_remove = []
objects_to_clear = []
for aid, v in self.engine.agent_manager.active_agents.items():
if aid.startswith("bg_"):
# Check validity
if hasattr(v, 'valid_mask'):
curr_step = self.round
if curr_step >= len(v.valid_mask) or not v.valid_mask[curr_step]:
to_remove.append(aid)
objects_to_clear.append(v)
for aid in to_remove:
self.engine.agent_manager.active_agents.pop(aid, None)
if objects_to_clear:
self.engine.clear_objects(objects_to_clear)
def _spawn_controlled_agents(self):
for car in self.car_birth_info_list:
if car['show_time'] == self.round:
agent_id = f"controlled_{car['id']}"
# Check if we already have this agent (shouldn't happen with unique IDs but safety check)
if agent_id in self.controlled_agents:
continue
# Handling ID flickering / merging
# If this ID is new, check if there's an existing agent very close to its start position
# that just disappeared? (Not implemented here, complex logic)
# But we can check if there's an overlap with existing agents?
# For now, just spawn.
# Read vehicle type/size if available
vehicle_config = {}
if 'length' in car and 'width' in car:
vehicle_config = {
"length": car['length'],
"width": car['width']
}
vehicle = self.engine.spawn_object(
PolicyVehicle,
name=agent_id,
vehicle_config=vehicle_config,
position=car['begin'],
heading=car['heading']
)
vehicle.reset(position=car['begin'], heading=car['heading'])
# We don't set policy or destination in the same way, or maybe we do for compatibility
vehicle.set_destination(car['end'])
# Store extra info for replay
vehicle.expert_track = self.expert_tracks[car['scenario_id']]
vehicle.original_id = car['id']
self.controlled_agents[agent_id] = vehicle
self.controlled_agent_ids.append(agent_id)
self.engine.agent_manager.active_agents[agent_id] = vehicle
def step(self, action_dict=None):
# We ignore input action_dict for the purpose of controlling agents
# Instead, we calculate what the action *should* be
self.round += 1
expert_actions = {}
# 1. Update state of all controlled agents to the current timestep (self.round)
# and compute action from (self.round-1) to (self.round).
# Wait, usually step() moves T -> T+1.
# Current state is T. We want to move to T+1.
# So we need state at T and T+1.
# Identify agents that are done (valid=0 at T+1 or T+1 >= length)
agents_to_remove = []
# Update SDC/ego first (otherwise it will stay still with default action=0)
if self.replay_sdc and self.sdc_vehicle is not None and self.sdc_track is not None:
next_step = self.round
curr_step = self.round - 1
if next_step < len(self.sdc_track["state"]["position"]) and self.sdc_track["state"]["valid"][next_step]:
curr_state = {
"position": self.sdc_track["state"]["position"][curr_step],
"heading": self.sdc_track["state"]["heading"][curr_step],
"velocity": self.sdc_track["state"]["velocity"][curr_step],
}
next_state = {
"position": self.sdc_track["state"]["position"][next_step],
"heading": self.sdc_track["state"]["heading"][next_step],
"velocity": self.sdc_track["state"]["velocity"][next_step],
}
action, _ = self.inverse_dynamics.compute_action(curr_state, next_state, dt=0.1)
expert_actions[self.sdc_agent_id] = action
self.sdc_vehicle.set_position(next_state["position"])
self.sdc_vehicle.set_heading_theta(next_state["heading"])
self.sdc_vehicle.set_velocity(next_state["velocity"])
self.sdc_vehicle.last_expert_action = action
for agent_id, vehicle in self.controlled_agents.items():
track = vehicle.expert_track
# current_step = self.round - 1 # Since we incremented at start
# But vehicle is currently at state corresponding to self.round - 1.
# We want to move it to self.round.
# Check bounds
next_step = self.round
curr_step = self.round - 1
if next_step >= len(track['state']['position']):
agents_to_remove.append(agent_id)
continue
valid = track['state']['valid'][next_step]
if not valid:
agents_to_remove.append(agent_id)
continue
# Get states
curr_pos = track['state']['position'][curr_step]
next_pos = track['state']['position'][next_step]
curr_heading = track['state']['heading'][curr_step]
next_heading = track['state']['heading'][next_step]
curr_vel = track['state']['velocity'][curr_step]
next_vel = track['state']['velocity'][next_step]
# Prepare state dicts for Inverse Dynamics
curr_state = {
'position': curr_pos,
'heading': curr_heading,
'velocity': curr_vel
}
next_state = {
'position': next_pos,
'heading': next_heading,
'velocity': next_vel
}
# Calculate action
action, raw_info = self.inverse_dynamics.compute_action(curr_state, next_state, dt=0.1) # Waymo is 10Hz?
expert_actions[agent_id] = action
# Force update vehicle state
vehicle.set_position(next_pos)
vehicle.set_heading_theta(next_heading)
vehicle.set_velocity(next_vel)
# Also record this action in the vehicle for later retrieval if needed
vehicle.last_expert_action = action
# Remove finished agents
for agent_id in agents_to_remove:
vehicle = self.controlled_agents[agent_id]
self.controlled_agents.pop(agent_id)
self.controlled_agent_ids.remove(agent_id)
self.engine.agent_manager.active_agents.pop(agent_id, None)
self.engine.clear_objects([vehicle.id])
# Step physics world to update sensors/collision detection
# We don't need full integration, but we need to update the physics world state
self.engine.taskMgr.step()
# Spawn new agents for this turn
self._spawn_controlled_agents()
self._update_background_vehicles()
# Get observations
obs = self._get_all_obs()
rewards = {aid: 0.0 for aid in self.controlled_agents}
dones = {aid: False for aid in self.controlled_agents}
dones["__all__"] = (self.round >= self.config["horizon"]) or (len(self.controlled_agents) == 0 and self.round > 190) # Waymo scenarios are usually ~198 steps (20s @ 10Hz) or 90 steps (9s)
infos = {aid: {"expert_action": expert_actions.get(aid, np.zeros(2))} for aid in self.controlled_agents}
return obs, rewards, dones, infos
def _get_all_obs(self):
# Implement custom observation: 30m range, 10 nearest vehicles
obs_dict = {}
for agent_id, vehicle in self.controlled_agents.items():
# 1. Ego State
ego_state = [
vehicle.position[0], vehicle.position[1],
vehicle.velocity[0], vehicle.velocity[1],
vehicle.heading_theta
]
# 2. Neighbors
neighbors = []
# Iterate through all vehicles in the engine
candidates = []
for other_id, other_vehicle in self.engine.agent_manager.active_agents.items():
if other_id == agent_id:
continue
dist = np.linalg.norm(vehicle.position - other_vehicle.position)
if dist < 30.0:
candidates.append((dist, other_vehicle))
# Sort by distance
candidates.sort(key=lambda x: x[0])
# Take top 10
top_10 = candidates[:10]
neighbor_feats = []
for _, neighbor in top_10:
neighbor_feats.extend([
neighbor.position[0] - vehicle.position[0], # Relative pos
neighbor.position[1] - vehicle.position[1],
neighbor.velocity[0], # Absolute vel? or Relative? Usually relative in MultiAgent
neighbor.velocity[1]
])
# Pad if < 10
missing = 10 - len(top_10)
if missing > 0:
neighbor_feats.extend([0.0] * (4 * missing))
# Flatten
obs = np.array(ego_state + neighbor_feats, dtype=np.float32)
obs_dict[agent_id] = obs
return obs_dict

65
Env/inverse_dynamics.py Normal file
View File

@@ -0,0 +1,65 @@
import numpy as np
import math
class InverseDynamics:
def __init__(self, max_steering=0.7, max_acc=15.0, length=4.5):
"""
:param max_steering: Max steering angle in radians (approx 40 degrees)
:param max_acc: Max acceleration in m/s^2
:param length: Vehicle length in meters (Waymo default approx 4.5m)
"""
self.max_steering = max_steering
self.max_acc = max_acc
self.wheelbase = 0.7 * length # Approximation as per request
def compute_action(self, current_state, next_state, dt=0.1):
"""
Compute action [steering, acceleration] from current and next state.
State format: dictionary or object with keys/attrs: position (x, y), heading, velocity (v_x, v_y)
or numpy array [x, y, vx, vy, heading]
Using Bicycle Model:
delta = arctan(L * theta_dot / v)
acc = (v_next - v_curr) / dt
"""
# Extract state
# Assume state is dict-like for now, can adapt if needed
# We need: velocity (scalar), heading
# Helper to get speed
def get_speed(vel):
return np.linalg.norm(vel)
v_curr = get_speed(current_state['velocity'])
v_next = get_speed(next_state['velocity'])
# 1. Acceleration (longitudinal)
acc = (v_next - v_curr) / dt
# 2. Steering (lateral)
# theta_dot = (theta_next - theta_curr) / dt
theta_curr = current_state['heading']
theta_next = next_state['heading']
# Handle angle wrapping [-pi, pi]
diff_theta = theta_next - theta_curr
if diff_theta > np.pi:
diff_theta -= 2 * np.pi
elif diff_theta < -np.pi:
diff_theta += 2 * np.pi
theta_dot = diff_theta / dt
# Avoid division by zero for stationary vehicles
if v_curr < 0.1:
steering = 0.0
else:
# delta = arctan(L * theta_dot / v)
steering = np.arctan(self.wheelbase * theta_dot / v_curr)
# Normalize actions to [-1, 1]
norm_acc = np.clip(acc / self.max_acc, -1.0, 1.0)
norm_steering = np.clip(steering / self.max_steering, -1.0, 1.0)
return np.array([norm_steering, norm_acc]), {'raw_acc': acc, 'raw_steering': steering}

170
Env/logger_utils.py
View File

@@ -1,170 +0,0 @@
"""
日志工具模块
提供将终端输出同时保存到文件的功能
"""
import sys
import os
from datetime import datetime
class TeeLogger:
"""
双向输出类：同时输出到终端和文件
"""
def __init__(self, filename, mode='w', terminal=None):
"""
Args:
filename: 日志文件路径
mode: 文件打开模式 ('w'=覆盖, 'a'=追加)
terminal: 原始输出流（通常是sys.stdout或sys.stderr）
"""
self.terminal = terminal or sys.stdout
self.log_file = open(filename, mode, encoding='utf-8')
def write(self, message):
"""写入消息到终端和文件"""
self.terminal.write(message)
self.log_file.write(message)
self.log_file.flush() # 立即写入磁盘
def flush(self):
"""刷新缓冲区"""
self.terminal.flush()
self.log_file.flush()
def close(self):
"""关闭日志文件"""
if self.log_file:
self.log_file.close()
class LoggerContext:
"""
日志上下文管理器
使用with语句自动管理日志的开启和关闭
"""
def __init__(self, log_file=None, log_dir="logs", mode='w',
redirect_stdout=True, redirect_stderr=True):
"""
Args:
log_file: 日志文件名（None则自动生成时间戳文件名）
log_dir: 日志目录
mode: 文件打开模式 ('w'=覆盖, 'a'=追加)
redirect_stdout: 是否重定向标准输出
redirect_stderr: 是否重定向标准错误
"""
self.log_dir = log_dir
self.mode = mode
self.redirect_stdout = redirect_stdout
self.redirect_stderr = redirect_stderr
# 创建日志目录
os.makedirs(log_dir, exist_ok=True)
# 生成日志文件名
if log_file is None:
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
log_file = f"run_{timestamp}.log"
self.log_path = os.path.join(log_dir, log_file)
# 保存原始的stdout和stderr
self.original_stdout = sys.stdout
self.original_stderr = sys.stderr
# 日志对象
self.stdout_logger = None
self.stderr_logger = None
def __enter__(self):
"""进入上下文：开启日志"""
print(f"📝 日志记录已启用")
print(f"📁 日志文件: {self.log_path}")
print("-" * 60)
# 创建TeeLogger对象
if self.redirect_stdout:
self.stdout_logger = TeeLogger(
self.log_path,
mode=self.mode,
terminal=self.original_stdout
)
sys.stdout = self.stdout_logger
if self.redirect_stderr:
self.stderr_logger = TeeLogger(
self.log_path,
mode='a', # stderr总是追加模式
terminal=self.original_stderr
)
sys.stderr = self.stderr_logger
return self
def __exit__(self, exc_type, exc_val, exc_tb):
"""退出上下文：关闭日志"""
# 恢复原始输出
sys.stdout = self.original_stdout
sys.stderr = self.original_stderr
# 关闭日志文件
if self.stdout_logger:
self.stdout_logger.close()
if self.stderr_logger:
self.stderr_logger.close()
print("-" * 60)
print(f"✅ 日志已保存到: {self.log_path}")
# 返回False表示不抑制异常
return False
def setup_logger(log_file=None, log_dir="logs", mode='w'):
"""
快速设置日志记录
Args:
log_file: 日志文件名（None则自动生成）
log_dir: 日志目录
mode: 文件模式 ('w'=覆盖, 'a'=追加)
Returns:
LoggerContext对象
Example:
with setup_logger("my_test.log"):
print("这条消息会同时输出到终端和文件")
"""
return LoggerContext(log_file=log_file, log_dir=log_dir, mode=mode)
def get_default_log_filename(prefix="run"):
"""
生成默认的日志文件名（带时间戳）
Args:
prefix: 文件名前缀
Returns:
str: 格式为 "prefix_YYYYMMDD_HHMMSS.log"
"""
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
return f"{prefix}_{timestamp}.log"
if __name__ == "__main__":
# 测试代码
print("测试1: 使用默认配置")
with setup_logger():
print("这是测试消息1")
print("这是测试消息2")
print("日志记录已结束\n")
print("测试2: 使用自定义文件名")
with setup_logger(log_file="test_custom.log"):
print("自定义文件名测试")
for i in range(3):
print(f" 消息 {i+1}")
print("完成")

54
Env/run_multiagent_env.py
View File

@@ -1,24 +1,14 @@
from scenario_env import MultiAgentScenarioEnv
from simple_idm_policy import ConstantVelocityPolicy
from Env.simple_idm_policy import ConstantVelocityPolicy
from metadrive.engine.asset_loader import AssetLoader
from logger_utils import setup_logger
import sys
import os
WAYMO_DATA_DIR = r"/home/huangfukk/mdsn"
WAYMO_DATA_DIR = r"/home/huangfukk/MAGAIL4AutoDrive/data"
def main(enable_logging=False, log_file=None):
"""
主函数
Args:
enable_logging: 是否启用日志记录到文件
log_file: 日志文件名（None则自动生成时间戳文件名）
"""
def main():
env = MultiAgentScenarioEnv(
config={
# "data_directory": AssetLoader.file_path(AssetLoader.asset_path, "waymo", unix_style=False),
"data_directory": AssetLoader.file_path(WAYMO_DATA_DIR, "exp_filtered", unix_style=False),
"data_directory": AssetLoader.file_path(WAYMO_DATA_DIR, "exp_converted", unix_style=False),
"is_multi_agent": True,
"num_controlled_agents": 3,
"horizon": 300,
@@ -26,21 +16,12 @@ def main(enable_logging=False, log_file=None):
"sequential_seed": True,
"reactive_traffic": True,
"manual_control": True,
# 车道检测与过滤配置
"filter_offroad_vehicles": True, # 启用车道区域过滤，过滤草坪等非车道区域的车辆
"lane_tolerance": 3.0, # 车道检测容差（米），可根据需要调整
"max_controlled_vehicles": None, # 限制最大车辆数（可选，None表示不限制）
# 调试配置（可选）
# "debug_lane_filter": True, # 启用车道过滤详细调试
# "verbose_reset": True, # 启用重置详细统计
# "inherit_expert_velocity": True, # 继承专家速度
},
agent2policy=ConstantVelocityPolicy(target_speed=50)
)
obs = env.reset(0)
obs = env.reset(0
)
for step in range(10000):
actions = {
aid: env.controlled_agents[aid].policy.act()
@@ -57,25 +38,4 @@ def main(enable_logging=False, log_file=None):
if __name__ == "__main__":
# 解析命令行参数
enable_logging = "--log" in sys.argv or "-l" in sys.argv
# 提取自定义日志文件名
log_file = None
for arg in sys.argv:
if arg.startswith("--log-file="):
log_file = arg.split("=")[1]
break
if enable_logging:
# 使用日志记录
log_dir = os.path.join(os.path.dirname(__file__), "logs")
with setup_logger(log_file=log_file, log_dir=log_dir):
main(enable_logging=True, log_file=log_file)
else:
# 普通运行（只输出到终端）
print("💡 提示: 使用 --log 或 -l 参数启用日志记录")
print(" 示例: python run_multiagent_env.py --log")
print(" 自定义文件名: python run_multiagent_env.py --log --log-file=my_run.log")
print("-" * 60)
main(enable_logging=False)
main()

699
Env/scenario_env.py
View File

@@ -1,91 +1,3 @@
"""
多智能体场景环境 (MultiAgentScenarioEnv)
==================================
配置参数说明 (写在最前面)
==================================
基础配置：
data_directory (str): 专家数据目录路径
num_controlled_agents (int): 默认可控智能体数量，默认3
horizon (int): 每个回合的最大步数，默认1000
车道检测与过滤配置：
filter_offroad_vehicles (bool): 是否过滤非车道区域的车辆，默认True
- True: 过滤掉在草坪、停车场等非车道区域生成的车辆
- False: 保留所有车辆
lane_tolerance (float): 车道检测容差（米），默认3.0
- 用于放宽车道检测的边界条件
max_controlled_vehicles (int|None): 最大可控车辆数限制，默认None
- None: 不限制车辆数量
- int: 限制最多生成的车辆数
场景对象配置：
no_traffic_lights (bool): 是否禁用红绿灯渲染和逻辑，默认False
- True: 完全移除场景中的红绿灯
- False: 保留红绿灯（按数据集原样生成）
专家数据继承配置：
inherit_expert_velocity (bool): 是否继承专家数据中车辆的初始速度，默认False
- True: 车辆生成时使用专家数据中的速度
- False: 车辆生成时速度为0（由策略控制）
调试模式配置：
debug_lane_filter (bool): 车道过滤详细调试输出，默认False
- 输出每个车辆位置的车道检测详细过程
verbose_reset (bool): 重置时输出详细统计信息，默认False
- 输出场景统计、过滤详情等
使用示例：
env = MultiAgentScenarioEnv(
config={
"data_directory": "path/to/data",
"max_controlled_vehicles": 10,
"inherit_expert_velocity": True, # 继承专家速度
"no_traffic_lights": True, # 禁用红绿灯
"verbose_reset": True, # 详细输出
},
agent2policy=your_policy
)
==================================
整体逻辑和处理流程
==================================
1. 初始化阶段：
- 继承MetaDrive的ScenarioEnv基类
- 配置多智能体参数（车辆数量、调试模式等）
- 接收策略映射字典(agent2policy)
2. 环境重置阶段 (reset方法)：
- 解析专家数据，提取车辆生成信息(car_birth_info_list)
- 清理原始交通数据，只保留车辆位置、朝向、目的地、速度（可选）
- 禁用红绿灯（如果配置）
- 初始化地图和车道信息
- 执行车道过滤(_filter_valid_spawn_positions)，移除非车道区域的车辆
- 限制最大车辆数量
- 生成可控智能体(_spawn_controlled_agents)
3. 观测获取阶段 (_get_all_obs方法)：
- 遍历所有可控车辆
- 获取车辆状态（位置、速度、朝向）
- 检测红绿灯状态(_get_traffic_light_state)
- 获取激光雷达数据（前向、侧向、车道线检测）
- 组装完整观测向量
4. 环境步进阶段 (step方法)：
- 执行所有智能体的动作
- 更新物理引擎状态
- 生成新的智能体（按时间步）
- 返回新的观测、奖励、完成状态
核心功能模块：
- PolicyVehicle: 可控制策略的车辆类
- 车道检测与过滤: 确保车辆只在有效车道上生成
- 多智能体管理: 动态生成和管理可控车辆
- 专家速度继承: 可选地继承专家数据中的初始速度
"""
import numpy as np
from metadrive.component.navigation_module.node_network_navigation import NodeNetworkNavigation
from metadrive.envs.scenario_env import ScenarioEnv
@@ -99,693 +11,194 @@ from metadrive.type import MetaDriveType
class PolicyVehicle(DefaultVehicle):
"""
可控制策略的车辆类
继承自MetaDrive的DefaultVehicle，增加了策略控制和目标设置功能。
用于多智能体环境中的可控车辆，支持自定义策略和目的地。
"""
def __init__(self, *args, **kwargs):
"""
初始化策略车辆
Args:
*args: 传递给父类的位置参数
**kwargs: 传递给父类的关键字参数
"""
super().__init__(*args, **kwargs)
self.policy = None # 车辆的控制策略
self.destination = None # 车辆的目标目的地
self.policy = None
self.destination = None
def set_policy(self, policy):
"""
设置车辆的控制策略
Args:
policy: 控制策略对象，必须实现act(observation)方法
"""
self.policy = policy
def set_destination(self, des):
"""
设置车辆的目标目的地
Args:
des: 目标位置坐标 (x, y)
"""
self.destination = des
def act(self, observation, policy=None):
"""
根据观测获取动作
Args:
observation: 环境观测数据
policy: 可选的外部策略，如果提供则优先使用
Returns:
动作向量，如果无策略则返回随机动作
"""
if self.policy is not None:
return self.policy.act(observation)
else:
return self.action_space.sample()
def before_step(self, action):
"""
执行动作前的状态记录
在每步执行前记录当前状态，用于后续的状态追踪和分析。
Args:
action: 即将执行的动作
"""
self.last_position = self.position # 记录当前位置 (2D向量)
self.last_velocity = self.velocity # 记录当前速度 (2D向量)
self.last_speed = self.speed # 记录当前速度大小 (标量)
self.last_heading_dir = self.heading # 记录当前朝向
self.last_position = self.position # 2D vector
self.last_velocity = self.velocity # 2D vector
self.last_speed = self.speed # Scalar
self.last_heading_dir = self.heading
if action is not None:
self.last_current_action.append(action) # 记录动作历史
self._set_action(action) # 设置动作到车辆
self.last_current_action.append(action)
self._set_action(action)
def is_done(self):
"""
检查车辆是否完成任务
目前为空实现，可根据需要添加到达目的地或碰撞检测逻辑
Returns:
bool: True表示任务完成，False表示继续执行
"""
# arrive or crash
pass
# 将PolicyVehicle注册为默认车辆类型
vehicle_class_to_type[PolicyVehicle] = "default"
class MultiAgentScenarioEnv(ScenarioEnv):
"""
多智能体场景环境
基于MetaDrive的ScenarioEnv扩展，支持多智能体强化学习训练。
主要功能包括：
1. 从专家数据中提取车辆信息并生成可控智能体
2. 车道检测与过滤，确保车辆在有效区域生成
3. 红绿灯状态检测，为智能体提供交通信号信息
4. 多智能体观测、动作和奖励管理
"""
@classmethod
def default_config(cls):
"""
获取环境的默认配置
继承父类配置并添加多智能体相关的配置参数
配置参数说明：
- data_directory: 专家数据目录路径
- num_controlled_agents: 默认可控智能体数量
- horizon: 每个回合的最大步数
车道检测与过滤配置：
- filter_offroad_vehicles: 是否过滤非车道区域的车辆
- lane_tolerance: 车道检测容差（米）
- max_controlled_vehicles: 最大可控车辆数限制（None表示不限制）
场景对象配置：
- no_traffic_lights: 禁用红绿灯渲染和逻辑
专家数据继承配置：
- inherit_expert_velocity: 是否继承专家数据中车辆的初始速度（默认False）
调试模式配置：
- debug_traffic_light: 红绿灯检测详细调试输出
- debug_lane_filter: 车道过滤详细调试输出
- verbose_reset: 重置时输出详细统计信息
Returns:
dict: 包含所有配置参数的字典
"""
config = super().default_config()
config.update(dict(
# 基础配置
data_directory=None,
num_controlled_agents=3,
horizon=1000,
# 车道检测与过滤配置
filter_offroad_vehicles=True,
lane_tolerance=3.0,
max_controlled_vehicles=None,
# 场景对象配置
no_traffic_lights=False,
# 专家数据继承配置
inherit_expert_velocity=False,
# 调试模式配置
debug_lane_filter=False,
verbose_reset=False,
))
return config
def __init__(self, config, agent2policy):
"""
初始化多智能体场景环境
Args:
config: 环境配置字典，包含各种参数设置
agent2policy: 智能体ID到策略的映射字典
"""
self.policy = agent2policy # 智能体策略映射
self.controlled_agents = {} # 可控智能体字典 {agent_id: vehicle}
self.controlled_agent_ids = [] # 可控智能体ID列表
self.obs_list = [] # 观测数据列表
self.round = 0 # 当前时间步
# 调试模式配置
self.debug_lane_filter = config.get("debug_lane_filter", False)
# 调用父类初始化
self.policy = agent2policy
self.controlled_agents = {}
self.controlled_agent_ids = []
self.obs_list = []
self.round = 0
super().__init__(config)
def reset(self, seed: Union[None, int] = None):
"""
重置环境到初始状态
这是环境的核心重置方法，执行以下步骤：
1. 初始化日志系统
2. 解析专家数据，提取车辆生成信息
3. 清理原始交通数据
4. 初始化地图和车道信息
5. 执行车道过滤和车辆数量限制
6. 生成可控智能体
7. 返回初始观测
Args:
seed: 随机种子，用于环境重置时的随机性控制
Returns:
list: 所有智能体的初始观测数据
"""
self.round = 0 # 重置时间步计数器
# 初始化日志系统
self.round = 0
if self.logger is None:
self.logger = get_logger()
log_level = self.config.get("log_level", logging.DEBUG if self.config.get("debug", False) else logging.INFO)
set_log_level(log_level)
# 延迟初始化MetaDrive引擎
self.lazy_init()
self._reset_global_seed(seed)
if self.engine is None:
raise ValueError("Broken MetaDrive instance.")
# 在场景加载前禁用红绿灯（如果配置中启用了no_traffic_lights选项）
if self.config.get("no_traffic_lights", False):
# 重写红绿灯管理器的方法，阻止创建和使用红绿灯
if hasattr(self.engine, 'light_manager') and self.engine.light_manager is not None:
self.engine.light_manager.before_reset = lambda *args, **kwargs: None
self.engine.light_manager.after_reset = lambda *args, **kwargs: None
self.engine.light_manager.before_step = lambda *args, **kwargs: None
self.engine.light_manager.get_traffic_light = lambda *args, **kwargs: None
self.logger.info("已禁用红绿灯管理器")
# 步骤1：解析专家数据，提取车辆生成信息
# 记录专家数据中每辆车的位置，接着全部清除，只保留位置等信息，用于后续生成
_obj_to_clean_this_frame = [] # 需要清理的对象ID列表
self.car_birth_info_list = [] # 车辆生成信息列表
self.expert_trajectories = {} # 专家数据轨迹字典
_obj_to_clean_this_frame = []
self.car_birth_info_list = []
for scenario_id, track in self.engine.traffic_manager.current_traffic_data.items():
# 跳过自车（SDC - Self Driving Car）
if scenario_id == self.engine.traffic_manager.sdc_scenario_id:
continue
# 只处理车辆类型的对象
else:
if track["type"] == MetaDriveType.VEHICLE:
_obj_to_clean_this_frame.append(scenario_id)
valid = track['state']['valid'] # 车辆有效性标记
# 找到车辆首次出现和最后出现的时间步
valid = track['state']['valid']
first_show = np.argmax(valid) if valid.any() else -1
last_show = len(valid) - 1 - np.argmax(valid[::-1]) if valid.any() else -1
if first_show == -1 or last_show == -1:
continue
object_id = track["metadata"]["object_id"]
# 提取完整轨迹数据(只使用确认存在的字段)
trajectory_data = {
"object_id": object_id,
"scenario_id": scenario_id,
"valid_mask": valid[first_show:last_show+1].copy(), # 有效性掩码
"positions": track["state"]["position"][first_show:last_show+1].copy(), # (T, 3)
"headings": track["state"]["heading"][first_show:last_show+1].copy(), # (T,)
"velocities": track["state"]["velocity"][first_show:last_show+1].copy(), # (T, 2)
"timesteps": np.arange(first_show, last_show+1), # 时间戳
"start_timestep": first_show,
"end_timestep": last_show,
"length": last_show - first_show + 1
}
# 可选:如果数据中有车辆尺寸信息,则添加
# 方法1: 尝试从state中获取
if "length" in track["state"]:
trajectory_data["vehicle_length"] = track["state"]["length"][first_show]
if "width" in track["state"]:
trajectory_data["vehicle_width"] = track["state"]["width"][first_show]
if "height" in track["state"]:
trajectory_data["vehicle_height"] = track["state"]["height"][first_show]
# 方法2: 尝试从metadata中获取
if "vehicle_length" not in trajectory_data and "length" in track.get("metadata", {}):
trajectory_data["vehicle_length"] = track["metadata"]["length"]
if "vehicle_width" not in trajectory_data and "width" in track.get("metadata", {}):
trajectory_data["vehicle_width"] = track["metadata"]["width"]
if "vehicle_height" not in trajectory_data and "height" in track.get("metadata", {}):
trajectory_data["vehicle_height"] = track["metadata"]["height"]
# 方法3: 使用默认值(如果以上都没有)
if "vehicle_length" not in trajectory_data:
trajectory_data["vehicle_length"] = 4.5 # MetaDrive默认车长
if "vehicle_width" not in trajectory_data:
trajectory_data["vehicle_width"] = 2.0 # MetaDrive默认车宽
if "vehicle_height" not in trajectory_data:
trajectory_data["vehicle_height"] = 1.5 # MetaDrive默认车高
# 存储到专家轨迹字典
self.expert_trajectories[object_id] = trajectory_data
# 提取车辆关键信息
car_info = {
# id，出现时间，出生点坐标，出生朝向，目的地
self.car_birth_info_list.append({
'id': track['metadata']['object_id'],
'show_time': first_show,
'begin': (track['state']['position'][first_show, 0], track['state']['position'][first_show, 1]),
'heading': track['state']['heading'][first_show],
'end': (track['state']['position'][last_show, 0], track['state']['position'][last_show, 1])
}
})
# 如果配置要求继承专家速度，则提取初始速度
if self.config.get("inherit_expert_velocity", False):
velocity = track['state']['velocity'][first_show]
car_info['velocity'] = (velocity[0], velocity[1])
self.car_birth_info_list.append(car_info)
# 非车辆对象（如红绿灯、行人等）保留，不清理
# 清理车辆原始交通数据，释放内存（保留红绿灯等其他对象）
for scenario_id in _obj_to_clean_this_frame:
self.engine.traffic_manager.current_traffic_data.pop(scenario_id)
# 步骤2：重置MetaDrive引擎和传感器
self.engine.reset()
self.reset_sensors()
self.engine.taskMgr.step()
# 步骤3：获取地图车道信息
self.lanes = self.engine.map_manager.current_map.road_network.graph
# 调试：场景信息统计（仅在verbose_reset模式下输出）
if self.config.get("verbose_reset", False):
print(f"\n📍 场景信息统计:")
print(f" - 总车道数: {len(self.lanes)}")
# 统计红绿灯数量（如果未禁用红绿灯）
if not self.config.get("no_traffic_lights", False):
traffic_light_lanes = []
for lane in self.lanes.values():
if self.engine.light_manager.has_traffic_light(lane.lane.index):
traffic_light_lanes.append(lane.lane.index)
print(f" - 有红绿灯的车道数: {len(traffic_light_lanes)}")
if len(traffic_light_lanes) > 5:
print(f" 车道索引示例: {traffic_light_lanes[:5]} ...")
else:
print(f" - 红绿灯: 已禁用")
# 步骤4：执行车道区域过滤
total_cars_before = len(self.car_birth_info_list)
valid_count, filtered_count, filtered_list = self._filter_valid_spawn_positions()
# 输出过滤信息（仅在有过滤时输出）
if filtered_count > 0:
if self.config.get("verbose_reset", False):
self.logger.warning(f"车辆生成位置过滤: 原始 {total_cars_before} 辆, "
f"有效 {valid_count} 辆, 过滤 {filtered_count} 辆")
for filtered_car in filtered_list[:3]:
self.logger.debug(f" 过滤车辆 ID={filtered_car['id']}, "
f"位置={filtered_car['position']}, "
f"原因={filtered_car['reason']}")
if filtered_count > 3:
self.logger.debug(f" ... 还有 {filtered_count - 3} 辆车被过滤")
else:
self.logger.info(f"车辆过滤: {total_cars_before} 辆 -> {valid_count} 辆 (过滤 {filtered_count} 辆)")
# 步骤5：限制最大车辆数（在过滤后应用）
max_vehicles = self.config.get("max_controlled_vehicles", None)
if max_vehicles is not None and len(self.car_birth_info_list) > max_vehicles:
original_count = len(self.car_birth_info_list)
self.car_birth_info_list = self.car_birth_info_list[:max_vehicles]
if self.config.get("verbose_reset", False):
self.logger.info(f"限制最大车辆数: {original_count} 辆 -> {max_vehicles} 辆")
# 最终统计
if self.config.get("verbose_reset", False):
self.logger.info(f"✓ 最终生成 {len(self.car_birth_info_list)} 辆可控车辆")
# 清理渲染器
if self.top_down_renderer is not None:
self.top_down_renderer.clear()
self.engine.top_down_renderer = None
# 初始化回合相关变量
self.dones = {} # 智能体完成状态
self.episode_rewards = defaultdict(float) # 回合奖励累积
self.episode_lengths = defaultdict(int) # 回合长度累积
self.dones = {}
self.episode_rewards = defaultdict(float)
self.episode_lengths = defaultdict(int)
# 清空可控智能体
self.controlled_agents.clear()
self.controlled_agent_ids.clear()
# 步骤6：调用父类重置并生成可控智能体
super().reset(seed) # 初始化场景
self._spawn_controlled_agents()
# 步骤7：返回初始观测
return self._get_all_obs()
def _is_position_on_lane(self, position, tolerance=None):
"""
检测给定位置是否在有效车道范围内
这个函数用于验证车辆生成位置是否在合法的车道上，避免在草坪、停车场等
非车道区域生成车辆。支持两种检测方法：
1. 严格检测：直接使用MetaDrive的point_on_lane方法
2. 容差检测：考虑车道边缘的容差范围（当前已禁用）
Args:
position: (x, y) 车辆位置坐标
tolerance: 容差范围（米），用于放宽检测条件。None时使用配置中的默认值
Returns:
bool: True表示在车道上，False表示在非车道区域（如草坪、停车场等）
"""
# 检查车道信息是否已初始化
if not hasattr(self, 'lanes') or self.lanes is None:
if self.debug_lane_filter:
print(f" ⚠️ 车道信息未初始化，默认允许")
return True # 如果车道信息未初始化，默认允许生成
# 设置容差参数
if tolerance is None:
tolerance = self.config.get("lane_tolerance", 3.0)
position_2d = (position[0], position[1])
if self.debug_lane_filter:
print(f" 🔍 检测位置 ({position_2d[0]:.2f}, {position_2d[1]:.2f}), 容差={tolerance}m")
# 方法1：直接检测是否在任一车道上
# 遍历所有车道，使用MetaDrive的point_on_lane方法进行精确检测
checked_lanes = 0
for lane in self.lanes.values():
try:
checked_lanes += 1
if lane.lane.point_on_lane(position_2d):
if self.debug_lane_filter:
print(f" ✅ 在车道上 (车道{lane.lane.index}, 检查了{checked_lanes}条)")
return True
except:
# 如果检测过程中出现异常，继续检查下一条车道
continue
if self.debug_lane_filter:
print(f" ❌ 不在任何车道上 (检查了{checked_lanes}条车道)")
# 方法2：如果严格检测失败，使用容差范围检测（考虑车道边缘）
# 注释：此方法已被禁用，如需启用请取消注释
# 该方法通过计算点到车道中心线的横向距离来判断是否在容差范围内
# if tolerance > 0:
# for lane in self.lanes.values():
# try:
# # 计算点到车道中心线的距离
# lane_obj = lane.lane
# # 获取车道长度并检测最近点
# s, lateral = lane_obj.local_coordinates(position_2d)
# # 如果横向距离在容差范围内，认为是有效的
# if abs(lateral) <= tolerance and 0 <= s <= lane_obj.length:
# return True
# except:
# continue
return False
def _filter_valid_spawn_positions(self):
"""
过滤掉生成位置不在有效车道上的车辆信息
这个函数是车道过滤的核心实现，用于确保所有生成的车辆都在合法的车道上。
它会遍历所有车辆生成信息，使用_is_position_on_lane方法检测每个位置，
过滤掉在草坪、停车场等非车道区域的车辆。
过滤过程包括：
1. 检查配置是否启用过滤
2. 遍历所有车辆生成信息
3. 对每个车辆位置进行车道检测
4. 分离有效和无效的车辆
5. 更新车辆生成列表
Returns:
tuple: (有效车辆数量, 被过滤车辆数量, 被过滤车辆详细信息列表)
"""
# 检查配置是否启用车道过滤
if not self.config.get("filter_offroad_vehicles", True):
if self.debug_lane_filter:
print(f"🚫 车道过滤已禁用")
return len(self.car_birth_info_list), 0, []
if self.debug_lane_filter:
print(f"\n🔍 开始车道过滤: 共 {len(self.car_birth_info_list)} 辆车待检测")
# 初始化过滤结果
valid_cars = [] # 有效车辆列表
filtered_cars = [] # 被过滤车辆列表
tolerance = self.config.get("lane_tolerance", 3.0)
# 遍历所有车辆生成信息进行检测
for idx, car in enumerate(self.car_birth_info_list):
if self.debug_lane_filter:
print(f"\n车辆 {idx+1}/{len(self.car_birth_info_list)}: ID={car['id']}")
# 检测车辆生成位置是否在有效车道上
if self._is_position_on_lane(car['begin'], tolerance=tolerance):
valid_cars.append(car)
if self.debug_lane_filter:
print(f" ✅ 保留")
else:
# 记录被过滤的车辆信息
filtered_cars.append({
'id': car['id'],
'position': car['begin'],
'reason': '生成位置不在有效车道上（可能在草坪/停车场等区域）'
})
if self.debug_lane_filter:
print(f" ❌ 过滤 (原因: 不在车道上)")
# 更新车辆生成列表为过滤后的结果
self.car_birth_info_list = valid_cars
if self.debug_lane_filter:
print(f"\n📊 过滤结果: 保留 {len(valid_cars)} 辆, 过滤 {len(filtered_cars)} 辆")
return len(valid_cars), len(filtered_cars), filtered_cars
def _spawn_controlled_agents(self):
"""
生成可控智能体车辆
根据当前时间步和车辆生成信息，动态生成需要出现的可控车辆。
每个车辆都会被分配策略和目标目的地，并注册到MetaDrive引擎中
参与物理仿真。
生成过程：
1. 遍历所有车辆生成信息
2. 检查车辆是否应该在当前时间步出现
3. 创建PolicyVehicle实例
4. 设置车辆策略和目的地
5. 注册到环境管理和引擎中
"""
# 注释：可以获取自车位置用于相对位置计算（当前未使用）
# ego_vehicle = self.engine.agent_manager.active_agents.get("default_agent")
# ego_position = ego_vehicle.position if ego_vehicle else np.array([0, 0])
# 遍历所有车辆生成信息
for car in self.car_birth_info_list:
# 检查车辆是否应该在当前时间步出现
if car['show_time'] == self.round:
# 生成智能体ID
agent_id = f"controlled_{car['id']}"
# 在MetaDrive引擎中生成车辆对象
vehicle = self.engine.spawn_object(
PolicyVehicle, # 使用自定义的策略车辆类
vehicle_config={}, # 车辆配置（使用默认）
position=car['begin'], # 车辆生成位置
heading=car['heading'] # 车辆生成朝向
PolicyVehicle,
vehicle_config={},
position=car['begin'],
heading=car['heading']
)
# 重置车辆状态到指定位置和朝向
vehicle.reset(position=car['begin'], heading=car['heading'])
# 如果配置要求继承专家速度，则设置初始速度
if 'velocity' in car and self.config.get("inherit_expert_velocity", False):
vehicle.set_velocity(car['velocity'])
vehicle.set_policy(self.policy)
vehicle.set_destination(car['end'])
# 设置车辆的控制策略和目标
vehicle.set_policy(self.policy) # 设置策略
vehicle.set_destination(car['end']) # 设置目的地
# 注册到环境管理
self.controlled_agents[agent_id] = vehicle
self.controlled_agent_ids.append(agent_id)
# ✅ 关键：注册到引擎的 active_agents，才能参与物理更新
# 这是MetaDrive引擎识别和管理智能体的关键步骤
self.engine.agent_manager.active_agents[agent_id] = vehicle
def _get_all_obs(self):
"""
获取所有可控智能体的观测数据
这是环境的核心观测函数，为每个可控智能体组装完整的观测向量。
观测数据包括：
1. 车辆状态信息：位置、速度、朝向
2. 传感器数据：激光雷达（前向、侧向、车道线检测）
3. 导航信息：目标目的地
观测向量结构：
- position[2]: 车辆位置 (x, y)
- velocity[2]: 车辆速度 (vx, vy)
- heading[1]: 车辆朝向角度
- lidar[80]: 前向激光雷达数据 (80个激光束，30米范围)
- side_lidar[10]: 侧向激光雷达数据 (10个激光束，8米范围)
- lane_line_lidar[10]: 车道线检测数据 (10个激光束，3米范围)
- destination[2]: 目标目的地 (x, y)
Returns:
list: 所有智能体的观测数据列表
"""
self.obs_list = [] # 清空观测列表
# 遍历所有可控智能体
# position, velocity, heading, lidar, navigation, TODO: trafficlight -> list
self.obs_list = []
for agent_id, vehicle in self.controlled_agents.items():
# 获取车辆基本状态信息
state = vehicle.get_state()
# 获取激光雷达传感器数据
# 前向激光雷达：80个激光束，30米检测距离，用于障碍物检测
lidar = self.engine.get_sensor("lidar").perceive(
num_lasers=80,
distance=30,
traffic_light = 0
for lane in self.lanes.values():
if lane.lane.point_on_lane(state['position'][:2]):
if self.engine.light_manager.has_traffic_light(lane.lane.index):
traffic_light = self.engine.light_manager._lane_index_to_obj[lane.lane.index].status
if traffic_light == 'TRAFFIC_LIGHT_GREEN':
traffic_light = 1
elif traffic_light == 'TRAFFIC_LIGHT_YELLOW':
traffic_light = 2
elif traffic_light == 'TRAFFIC_LIGHT_RED':
traffic_light = 3
else:
traffic_light = 0
break
lidar = self.engine.get_sensor("lidar").perceive(num_lasers=80, distance=30, base_vehicle=vehicle,
physics_world=self.engine.physics_world.dynamic_world)
side_lidar = self.engine.get_sensor("side_detector").perceive(num_lasers=10, distance=8,
base_vehicle=vehicle,
physics_world=self.engine.physics_world.dynamic_world
)
# 侧向激光雷达：10个激光束，8米检测距离，用于侧向障碍物检测
side_lidar = self.engine.get_sensor("side_detector").perceive(
num_lasers=10,
distance=8,
physics_world=self.engine.physics_world.static_world)
lane_line_lidar = self.engine.get_sensor("lane_line_detector").perceive(num_lasers=10, distance=3,
base_vehicle=vehicle,
physics_world=self.engine.physics_world.static_world
)
# 车道线检测激光雷达：10个激光束，3米检测距离，用于车道线识别
lane_line_lidar = self.engine.get_sensor("lane_line_detector").perceive(
num_lasers=10,
distance=3,
base_vehicle=vehicle,
physics_world=self.engine.physics_world.static_world
)
# 组装完整的观测向量
obs = (state['position'][:2] + # 位置 (x, y)
list(state['velocity']) + # 速度 (vx, vy)
[state['heading_theta']] + # 朝向角度
lidar[0] + # 前向激光雷达数据
side_lidar[0] + # 侧向激光雷达数据
lane_line_lidar[0] + # 车道线检测数据
list(vehicle.destination)) # 目标目的地 (x, y)
physics_world=self.engine.physics_world.static_world)
obs = (state['position'][:2] + list(state['velocity']) + [state['heading_theta']]
+ lidar[0] + side_lidar[0] + lane_line_lidar[0] + [traffic_light]
+ list(vehicle.destination))
self.obs_list.append(obs)
return self.obs_list
def step(self, action_dict: Dict[AnyStr, Union[list, np.ndarray]]):
"""
执行环境的一个时间步
这是环境的核心步进函数，执行以下操作序列：
1. 更新时间步计数器
2. 执行所有智能体的动作（before_step）
3. 更新MetaDrive物理引擎状态
4. 执行智能体动作后的处理（after_step）
5. 生成新的智能体（按时间步）
6. 获取新的观测数据
7. 计算奖励和完成状态
8. 返回环境状态
Args:
action_dict: 智能体动作字典 {agent_id: action}
Returns:
tuple: (观测数据, 奖励字典, 完成状态字典, 信息字典)
- obs: 所有智能体的观测数据列表
- rewards: 每个智能体的奖励 {agent_id: reward}
- dones: 每个智能体的完成状态 {agent_id: done, "__all__": episode_done}
- infos: 每个智能体的额外信息 {agent_id: info}
"""
# 步骤1：更新时间步计数器
self.round += 1
# 步骤2：执行所有智能体的动作（动作执行前处理）
for agent_id, action in action_dict.items():
if agent_id in self.controlled_agents:
# 记录车辆状态并设置动作
self.controlled_agents[agent_id].before_step(action)
# 步骤3：更新MetaDrive物理引擎状态
# 这是核心的物理仿真步骤，所有车辆状态都会根据动作更新
self.engine.step()
# 步骤4：执行智能体动作后的处理
for agent_id in action_dict:
if agent_id in self.controlled_agents:
# 执行动作后的状态更新（如果有after_step方法）
self.controlled_agents[agent_id].after_step()
# 步骤5：生成新的智能体（按时间步动态生成）
self._spawn_controlled_agents()
# 步骤6：获取新的观测数据
obs = self._get_all_obs()
# 步骤7：计算奖励和完成状态
# 初始化所有智能体的奖励为0（可根据需要实现奖励计算逻辑）
rewards = {aid: 0.0 for aid in self.controlled_agents}
# 初始化所有智能体的完成状态为False（可根据需要实现完成条件）
dones = {aid: False for aid in self.controlled_agents}
# 检查整个回合是否结束（达到最大步数）
dones["__all__"] = self.episode_step >= self.config["horizon"]
# 初始化所有智能体的额外信息为空字典
infos = {aid: {} for aid in self.controlled_agents}
# 步骤8：返回环境状态
return obs, rewards, dones, infos

76
Env/test_no_traffic_lights.py
View File

@@ -1,76 +0,0 @@
"""测试禁用红绿灯功能"""
from scenario_env import MultiAgentScenarioEnv
from simple_idm_policy import ConstantVelocityPolicy
from metadrive.engine.asset_loader import AssetLoader
WAYMO_DATA_DIR = r"/home/huangfukk/MAGAIL4AutoDrive/Env"
def test_no_traffic_lights():
"""测试禁用红绿灯"""
print("=" * 60)
print("测试：禁用红绿灯功能")
print("=" * 60)
env = MultiAgentScenarioEnv(
config={
"data_directory": AssetLoader.file_path(WAYMO_DATA_DIR, "exp_converted", unix_style=False),
"is_multi_agent": True,
"num_controlled_agents": 3,
"horizon": 300,
"use_render": True,
"sequential_seed": True,
"reactive_traffic": True,
"manual_control": True,
# 车道检测与过滤配置
"filter_offroad_vehicles": True,
"lane_tolerance": 3.0,
"max_controlled_vehicles": 2,
# 禁用红绿灯
"no_traffic_lights": True, # 关键配置
# 调试模式
"debug_lane_filter": False,
"verbose_reset": False,
},
agent2policy=ConstantVelocityPolicy(target_speed=50)
)
print("\n重置环境...")
obs = env.reset(0)
# 检查红绿灯管理器状态
if hasattr(env.engine, 'light_manager') and env.engine.light_manager is not None:
num_lights = len(env.engine.light_manager._lane_index_to_obj)
print(f"✓ 红绿灯管理器中的红绿灯数量: {num_lights}")
if num_lights == 0:
print("✅ 成功：所有红绿灯已被移除！")
else:
print(f"⚠️ 警告：仍有 {num_lights} 个红绿灯")
print("\n运行几步测试...")
for step in range(100):
actions = {
aid: env.controlled_agents[aid].policy.act()
for aid in env.controlled_agents
}
obs, rewards, dones, infos = env.step(actions)
env.render(mode="topdown")
if step == 0:
print(f"步骤 {step}: 环境运行正常")
if dones["__all__"]:
break
print(f"\n测试完成，共运行 {step+1} 步")
print("请检查渲染窗口中是否还有红绿灯显示")
env.close()
print("=" * 60)
if __name__ == "__main__":
test_no_traffic_lights()

321
README.md
View File

@@ -2,274 +2,97 @@
> 基于多智能体生成对抗模仿学习(MAGAIL)的自动驾驶训练系统 | MetaDrive + Waymo Open Motion Dataset
[![MetaDrive](https://img.shields [![Python](https://img.shields.io/io/badge/Dataset目实现了适配多智能体场景的GAIL算法训练系统,核心创新在于**改进判别器架构支持动态车辆数量**,利用Transformer处理1-100+辆车的交互场景。
本项目利用 Waymo 真实驾驶数据，通过 MetaDrive 仿真环境构建专家回放系统，提取车辆状态与动作，用于训练多智能体模仿学习算法 (MAGAIL)。
**核心特性:**
- ✅ 完整的Waymo数据处理pipeline(12,201个场景)
- ✅ 车道过滤和红绿灯检测优化
- ✅ 支持5维简化/107维完整观测空间
- ✅ 专家轨迹数据集(52K+训练样本)
- 🚧 MAGAIL算法实现(判别器+策略网络)
## 📁 核心模块
* **`Env/expert_replay_env.py`**: 专家回放环境。核心类 `ExpertReplayEnv`，负责读取 Waymo 轨迹，计算逆动力学动作，并过滤非道路/静态车辆。
* **`Env/inverse_dynamics.py`**: 逆动力学模块。根据车辆位置和航向计算油门、刹车和转向动作。
* **`scripts/generate_expert_data.py`**: 数据收集脚本。批量运行场景并保存训练数据。
* **`scripts/visualize_replay.py`**: 可视化脚本。用于观察回放效果和数据质量。
***
## 🚀 快速开始
## 🚀 1. 数据收集
### 环境安装
### 生成专家数据
使用 `generate_expert_data.py` 脚本从 Waymo 数据集中批量提取 (State, Action) 对。
```bash
# 克隆项目
git clone <repository_url>
cd MAGAIL4AutoDrive
# 设置 Python 路径
export PYTHONPATH=$PYTHONPATH:.:./metadrive
# 安装依赖
pip install metadrive-simulator==0.4.3 torch numpy matplotlib scenarionet
# 创建必需目录
mkdir -p analysis_results
touch scripts/__init__.py dataset/__init__.py Algorithm/__init__.py
# 运行生成脚本
# --data_dir: Waymo 数据路径 (建议使用 exp_filtered)
# --output_dir: 结果保存路径
# --num_scenarios: 要处理的场景数量
python scripts/generate_expert_data.py \
--data_dir data/exp_filtered \
--output_dir data/training_data \
--num_scenarios 100 \
--start_index 0
```
### 数据准备
**生成的 `.pkl` 文件结构**：
包含一个列表，每个元素是一条车辆轨迹（Trajectory Dictionary）：
* `obs`: `(T, 45)` - 观测矩阵。包含 Ego 状态 (5维) + 10辆邻居车相对信息 (40维)。
* `acts`: `(T, 2)` - 动作矩阵。`[Steering, Accel]`，归一化到 `[-1, 1]`。
* `agent_id`: 车辆 ID。
* `scenario_id`: 所属场景 ID。
**内置过滤器**：
脚本会自动过滤掉以下无效车辆：
1. **非道路车辆**：始终在停车场或路外行驶的车辆。
2. **静态车辆**：全称移动距离小于 5米 且速度从未超过 1m/s 的车辆（作为背景流存在，不收集数据）。
---
## 🔍 2. 数据可视化与验证
### 回放可视化
使用 `visualize_replay.py` 直观地观察回放效果，确认车辆行为是否自然，以及过滤逻辑是否生效。
```bash
# 1. 转换Waymo数据
python -m scenarionet.convert_waymo -d ~/mdsn/exp_converted --raw_data_path /path/to/waymo --num_files=150
# 2. 筛选场景(无红绿灯)
python -m scenarionet.filter --database_path ~/mdsn/exp_filtered --from ~/mdsn/exp_converted --no_traffic_light
# 3. 验证数据集
python scripts/check_database_info.py
# 运行可视化
# --horizon: 回放的最大步数 (Waymo 场景通常为 90 或 198 步)
python scripts/visualize_replay.py \
--data_dir data/exp_filtered \
--start_index 0 \
--num_scenarios 1 \
--horizon 200
```
### 运行环境
**观察要点**：
* **受控车辆 (Controlled Agents)**：控制台会显示数量（如 `Controlled agents: 2`）。这些是真正产生数据的车辆。
* **背景车辆**：如果在渲染图中看到其他车（通常是路边停放的），但受控数量很少，说明静态过滤生效了。
### 数据分析
使用 `analyze_expert_data.py` 查看生成数据的统计分布。
```bash
# 测试多智能体环境
python Env/run_multiagent_env.py
# 收集专家数据(10个场景测试)
python dataset/expert_dataset.py
python scripts/analyze_expert_data.py --data_path data/training_data/expert_data_0_100.pkl
```
***
---
## 📁 项目结构
## 🧠 3. 模型训练 (Next Steps)
```
MAGAIL4AutoDrive/
├── Env/ # 仿真环境模块
│ ├── scenario_env.py # 多智能体场景环境(含轨迹存储)
│ ├── run_multiagent_env.py# 环境运行脚本
│ └── simple_idm_policy.py # 测试策略
│
├── dataset/ # 数据集模块
│ └── expert_dataset.py # PyTorch Dataset(5维观测)
│
├── scripts/ # 工具脚本
│ ├── check_track_fields.py # 数据字段验证
│ ├── check_database_info.py # 数据库信息检查
│ ├── analyze_expert_data.py # 统计分析
│ └── visualize_expert_trajectory.py # 轨迹可视化
│
├── Algorithm/ # MAGAIL算法(待完善)
│ ├── bert.py # Transformer判别器
│ ├── disc.py # 判别器网络
│ ├── policy.py # 策略网络
│ ├── ppo.py # PPO优化器
│ └── magail.py # MAGAIL训练循环
│
└── analysis_results/ # 分析输出
├── statistics.pkl # 数据统计
└── distributions.png # 可视化图表
```
有了 `data/training_data/` 下的专家数据后，您可以开始训练 MAGAIL 模型。
***
### 训练流程
1. **加载数据**：使用 `dataset/expert_dataset.py` 中的 `ExpertDataset` 类加载 `.pkl` 数据。
2. **初始化 MAGAIL**：
* **Generator (Policy)**: 接收观测 `(B, 45)`，输出动作 `(B, 2)`。
* **Discriminator**: 接收状态-动作对 `(s, a)`，判断是专家还是生成器。
3. **交互采样**：
* 在 `MultiAgentScenarioEnv`（非回放模式）中运行 Policy。
* 收集 Policy 生成的轨迹。
4. **对抗更新**：
* 利用专家数据和 Policy 数据训练 Discriminator。
* 利用 Discriminator 的输出作为 Reward (GAIL Reward) 训练 Policy (PPO/TRPO)。
## 🎯 核心功能
### 1. 环境与数据处理
**scenario_env.py** - 多智能体场景环境
- 专家轨迹完整存储(位置、速度、航向角、车辆尺寸)
- 车道区域过滤(自动移除非车道车辆)
- 红绿灯状态检测(双重保障机制)
- 107维完整观测空间(激光雷达+车道线)
**expert_dataset.py** - 专家数据集
- 状态-动作对提取(逆动力学)
- 批量采样和序列化
- 支持PyTorch DataLoader
### 2. 数据分析工具
| 脚本 | 功能 | 输出 |
|------|------|------|
| `check_database_info.py` | 验证数据库完整性 | 场景总数、映射关系 |
| `check_track_fields.py` | 检查可用字段 | 必需/可选字段列表 |
| `analyze_expert_data.py` | 统计分析 | 轨迹长度、速度、交互频率 |
| `visualize_expert_trajectory.py` | 轨迹可视化 | 动画展示车辆运动 |
### 3. MAGAIL算法
**判别器** (Algorithm/bert.py + disc.py)
- Transformer编码器处理动态车辆数量
- CLS标记或均值池化聚合特征
- 支持集中式/去中心化/零和模式
**策略网络** (Algorithm/policy.py + ppo.py)
- Actor-Critic架构
- 参数共享机制(所有车辆共享模型)
- PPO/TRPO优化器
***
## ⚙️ 配置说明
```python
# 环境配置
config = {
# 数据路径
"data_directory": "~/mdsn/exp_filtered",
# 多智能体设置
"num_controlled_agents": 3, # 初始车辆数
"max_controlled_vehicles": 10, # 最大车辆数限制
# 车道过滤
"filter_offroad_vehicles": True, # 启用车道过滤
"lane_tolerance": 3.0, # 容差(米)
# 场景加载
"sequential_seed": True, # 顺序加载场景
"horizon": 1000, # 最大步数
}
```
***
## 📊 数据集统计
**当前数据规模**(基于exp_filtered):
- 场景总数: **12,201**
- 已收集场景: 10个测试场景
- 轨迹数: 900条
- 训练样本: **52,065**个(s,a)对
- 观测维度: 5维(简化) / 107维(完整)
- 动作维度: 2维(油门/刹车, 转向)
**数据质量**:
- 静止车辆占比: 54.8%(正常,包含停车场和路边停车)
- 平均轨迹长度: 67帧(6.7秒 @ 10Hz)
- 平均速度: 1.46 m/s
- 近距离交互(<5m): 1.92%
***
## 🛠️ 使用示例
### 收集专家数据
```python
# dataset/expert_dataset.py
from expert_dataset import ExpertTrajectoryDataset
# 收集1000个场景
trajectories = ExpertTrajectoryDataset.collect_from_env(
env_config,
num_scenarios=1000,
save_path="./expert_trajectories.pkl"
)
# 创建数据集
dataset = ExpertTrajectoryDataset(trajectories, sequence_length=1)
```
### 环境测试
```python
from scenario_env import MultiAgentScenarioEnv
env = MultiAgentScenarioEnv(
config=config,
agent2policy=your_policy
)
obs = env.reset()
for step in range(1000):
actions = {aid: policy(obs[aid]) for aid in env.controlled_agents}
obs, rewards, dones, infos = env.step(actions)
```
***
## ❓ 常见问题
### Q1: KeyError: 'bbox'
**原因**: Waymo转换数据不含bbox字段
**解决**: 使用length/width/height,代码已添加条件检查
### Q2: ModuleNotFoundError: scenario_env
**原因**: Python路径问题
**解决**: 脚本开头添加:
```python
import sys, os
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "../Env"))
```
### Q3: 多次reset失败(clear_objects错误)
**原因**: MetaDrive对象管理bug
**解决**: 每次收集数据都重新创建环境(已实现)
### Q4: 静止车辆占比过高
**原因**: Waymo真实场景包含停车场等静止车辆
**解决**: 可在数据收集时过滤平均速度<2m/s的轨迹
***
## 📈 开发路线图
### ✅ 已完成(Phase 1)
- [x] 数据转换与筛选
- [x] 完整轨迹存储
- [x] 数据质量分析
- [x] PyTorch Dataset构建
### 🚧 进行中(Phase 2)
- [ ] 107维完整观测空间
- [ ] 数据质量过滤
- [ ] 轨迹可视化工具
### 📅 计划中(Phase 3-4)
- [ ] 判别器网络实现
- [ ] Actor-Critic策略网络
- [ ] MAGAIL训练循环
- [ ] TensorBoard监控
- [ ] 实验与评估
***
## 📚 参考资料
- [MetaDrive Documentation](https://metadrive-simulator.readthedocs.io/)
- [Waymo Open Dataset](https://waymo.com/open/)
- [MAGAIL Paper](https://arxiv.org/abs/1807.09936)
- [ScenarioNet](https://github.com/metadriverse/scenarionet)
## 📄 License
MIT License
***
**💡 提示**: 项目处于活跃开发中,欢迎提Issue或PR贡献代码!
[1](https://blog.csdn.net/BxuqBlockchain/article/details/133606934)
[2](https://blog.csdn.net/sinat_28461591/article/details/148351123)
[3](https://www.reddit.com/r/Python/comments/13kpoti/readmeai_autogenerate_readmemd_files/)
[4](https://www.reddit.com/r/learnprogramming/comments/1298ix8/what_does_a_good_readme_look_like_for_personal/)
[5](https://juejin.cn/post/7195763127883169853)
[6](https://jimmysong.io/trans/spec-driven-development-using-markdown/)
[7](https://www.showapi.com/news/article/66b602964ddd79f11a001e3c)
[8](https://learn.microsoft.com/zh-cn/nuget/nuget-org/package-readme-on-nuget-org)
### 推荐配置
* **Observation**: 45维 (Ego + 10 Neighbors)
* **Action**: 2维 Continuous (Steering, Accel)
* **Horizon**: 200 steps
* **Batch Size**: 1024+ (多智能体环境下数据量很大)

498
TRAINING_ARCHITECTURE.md Normal file
View File

@@ -0,0 +1,498 @@
# MAGAIL 训练方案架构文档
## 目录
1. [训练数据结构](#1-训练数据结构)
2. [多智能体训练机制](#2-多智能体训练机制)
3. [完整训练流程](#3-完整训练流程)
4. [当前项目问题](#4-当前项目问题)
5. [TensorBoard 日志问题](#5-tensorboard-日志问题)
---
## 1. 训练数据结构
### 1.1 数据维度
**观测空间 (Observation Space)**
- **维度**: 45维
- **组成**:
- **Ego状态** (5维): `[position_x, position_y, velocity_x, velocity_y, heading_theta]`
- **邻居信息** (40维): 最多10个邻居，每个邻居4维特征
- 每个邻居: `[relative_x, relative_y, velocity_x, velocity_y]`
- 如果邻居数量 < 10，用零填充
**动作空间 (Action Space)**
- **维度**: 2维
- **组成**: `[steering, accel]`
- **范围**: 归一化到 `[-1, 1]`
### 1.2 数据格式
**专家数据文件结构** (`.pkl` 文件):
```python
# 每个 .pkl 文件包含一个列表，每个元素是一条车辆轨迹
trajectories = [
{
'obs': np.array, # Shape: (T, 45) - T为轨迹长度（可变）
'acts': np.array, # Shape: (T, 2) - 对应的动作序列
'agent_id': str, # 车辆ID
'scenario_id': int # 场景ID
},
...
]
```
**数据特点**:
- 轨迹长度 `T` 是**可变的**，取决于车辆在场景中的存活时间
- 最小轨迹长度过滤: 只保留长度 > 10 的轨迹
- 数据已通过静态车辆过滤（移动距离 < 5m 且最大速度 < 1m/s 的车辆被过滤）
### 1.3 数据生成流程
**脚本**: `scripts/generate_expert_data.py`
**流程**:
1. 从 Waymo 数据 (`data/exp_filtered`) 加载场景
2. 使用 `ExpertReplayEnv` 回放专家轨迹
3. 通过逆动力学 (`Env/inverse_dynamics.py`) 计算动作
4. 构建45维观测（Ego + 10个最近邻居）
5. 过滤无效轨迹（长度 < 10）
6. 保存为 `.pkl` 文件到 `data/training_data/`
**关键代码位置**:
- 观测构建: `Env/expert_replay_env.py` 的 `_get_all_obs()` 方法
- 动作计算: `Env/inverse_dynamics.py` 的 `compute_action()` 方法
---
## 2. 多智能体训练机制
### 2.1 可变长度处理
**问题**: 不同场景中智能体数量不同，每个智能体的轨迹长度也不同。
**解决方案**:
1. **数据层面** (`dataset/magail_dataset.py`):
- 将轨迹**展平**为独立的 `(state, action)` 对
- 每个样本是独立的，不保留序列信息
- 这样所有轨迹可以统一处理，不受长度限制
```python
# MAGAILExpertDataset 的处理方式
for traj in self.trajectories:
obs = traj['obs'] # (T, 45)
acts = traj['acts'] # (T, 2)
# 展平为独立样本
for i in range(len(obs)):
self.flat_data.append((obs[i], acts[i])) # 每个样本: (45,), (2,)
```
2. **训练环境层面** (`train_magail.py`):
- 每个 episode 动态处理不同数量的智能体
- 在 rollout 循环中，为每个活跃智能体独立收集数据
- 所有智能体的数据合并到一个 `memory` 中
```python
# Rollout 循环
for agent_id, obs in obs_dict.items():
act, logprob = ppo_agent.select_action(obs)
actions[agent_id] = act
# 所有智能体的数据都存入同一个 memory
memory['states'].append(obs)
memory['actions'].append(actions[agent_id])
...
```
3. **观测维度固定**:
- 通过 `MAGAILScenarioEnv` 确保观测维度始终为45维
- 邻居数量不足时用零填充，保证维度一致
### 2.2 多智能体交互
**环境设置**:
- 使用 `MAGAILScenarioEnv` (继承自 `MultiAgentScenarioEnv`)
- 自定义 `_get_all_obs()` 方法，确保观测格式与专家数据一致
- 每个智能体独立选择动作，环境统一执行
**关键点**:
- 所有智能体共享同一个策略网络（参数共享）
- 每个智能体独立计算动作和奖励
- 数据收集时将所有智能体的经验合并
---
## 3. 完整训练流程
### 3.1 数据准备阶段
**步骤 1: 生成专家数据**
```bash
python scripts/generate_expert_data.py \
--data_dir data/exp_filtered \
--output_dir data/training_data \
--num_scenarios 100 \
--start_index 0
```
**输出**: `data/training_data/expert_data_*.pkl`
### 3.2 模型初始化
**网络架构**:
1. **Actor (策略网络)**:
- 输入: 45维状态
- 输出: 2维动作（连续）
- 结构: MLP (45 → 256 → 256 → 2)
- 输出分布: 高斯分布（均值 + 可学习标准差）
2. **Critic (价值网络)**:
- 输入: 45维状态
- 输出: 标量价值
- 结构: MLP (45 → 256 → 256 → 1)
3. **Discriminator (鉴别器)**:
- 输入: 45维状态 + 2维动作 = 47维
- 输出: 标量（0-1之间，表示专家概率）
- 结构: MLP (47 → 256 → 256 → 1) + Sigmoid
### 3.3 训练循环
**主循环** (`train_magail.py` 的 `train()` 函数):
```
For each episode:
1. 收集 Rollout
- 重置环境（随机选择场景）
- 运行策略收集轨迹
- 存储 (state, action, logprob, next_state, done)
2. 训练 Discriminator
- 采样专家批次
- 采样策略批次
- 更新鉴别器:
- Expert loss: BCE(D(s_e, a_e), 1)
- Policy loss: BCE(D(s_p, a_p), 0)
- Total: L_d = L_expert + L_policy
3. 计算 GAIL 奖励
- 对所有策略状态-动作对:
reward = -log(1 - D(s, a) + ε)
- 替换环境奖励
4. 更新策略 (PPO)
- 计算 GAE (Generalized Advantage Estimation)
- PPO 更新 (K epochs):
- 计算优势函数
- 计算策略损失（带clip）
- 计算价值损失
- 更新 Actor 和 Critic
```
### 3.4 训练目标
**Discriminator 目标**:
```
L_D = E_{(s,a)~π_E}[-log(D(s,a))] + E_{(s,a)~π_θ}[-log(1-D(s,a))]
```
- 最大化区分专家数据和策略数据的能力
**Policy (Generator) 目标**:
```
L_π = E_{(s,a)~π_θ}[-log(D(s,a))] - λ_H(π_θ)
```
- 通过 PPO 优化，使用 GAIL 奖励作为信号
- 最大化鉴别器给出的"专家概率"
- 同时保持策略熵（探索）
**PPO 更新**:
```python
# 优势函数 (GAE)
advantages = compute_gae(rewards, values, next_values, dones, gamma, lambda)
# 策略损失
ratios = exp(log_probs - old_log_probs)
surr1 = ratios * advantages
surr2 = clip(ratios, 1-ε, 1+ε) * advantages
policy_loss = -min(surr1, surr2) + 0.01 * entropy
# 价值损失
value_loss = MSE(critic(states), returns)
# 总损失
total_loss = policy_loss + 0.5 * value_loss
```
### 3.5 关键代码位置
- **训练主循环**: `train_magail.py:278-505`
- **PPO 更新**: `train_magail.py:90-146`
- **Discriminator 更新**: `train_magail.py:429-462`
- **GAIL 奖励计算**: `train_magail.py:472-477`
---
## 4. 当前项目问题
### 4.1 环境重置问题
**问题描述**:
- MetaDrive 环境在快速重置时可能出现对象清理不完整的问题
- 错误信息: "You should clear all generated objects..."
**当前处理**:
- 代码中已有异常处理机制（`train_magail.py:288-342`）
- 重置失败时会尝试关闭并重新创建环境
- 但可能导致训练不稳定
**建议修复**:
- 在每次重置前显式清理所有对象
- 增加重置间隔，避免过于频繁的重置
- 考虑使用环境池（Environment Pool）复用环境实例
### 4.2 观测维度对齐
**问题描述**:
- 原始 `MultiAgentScenarioEnv` 返回108维观测（包含Lidar）
- 专家数据使用45维观测
- 维度不匹配会导致训练失败
**当前解决方案**:
- 通过 `MAGAILScenarioEnv` 重写 `_get_all_obs()` 方法
- 确保训练环境与专家数据使用相同的观测格式
**代码位置**: `train_magail.py:223-262`
### 4.3 数据收集效率
**问题描述**:
- 每个 episode 都需要完整运行环境收集数据
- 可变长度轨迹导致 batch 大小不一致
- 可能影响训练稳定性
**当前处理**:
- 使用展平的数据集，每个样本独立
- 在 rollout 时收集所有智能体的数据，合并处理
**潜在改进**:
- 考虑使用经验回放缓冲区
- 实现轨迹级别的采样（保留序列信息）
### 4.4 内存管理
**问题描述**:
- 长时间训练可能导致内存泄漏
- 环境对象可能没有完全释放
**当前处理**:
- 代码中有显式的 `gc.collect()` 和 `torch.cuda.empty_cache()`
- 但可能不够彻底
**建议**:
- 定期检查内存使用
- 考虑限制 rollout 长度
- 使用更激进的清理策略
### 4.5 训练稳定性
**问题描述**:
- Discriminator 可能过早收敛，导致策略无法学习
- GAIL 奖励可能不稳定
**当前处理**:
- 使用标准的 GAIL 奖励公式: `-log(1 - D(s,a) + ε)`
- PPO 的 clip 机制提供稳定性
**潜在改进**:
- 考虑使用 WGAN-GP 或 LSGAN 损失
- 实现 Discriminator 的预训练
- 添加奖励归一化
---
## 5. TensorBoard 日志问题
### 5.1 问题分析
**现象**:
- `runs/magail_0112/` 目录下只有模型文件（`.pth`），没有 TensorBoard 事件文件（`events.out.tfevents.*`）
- 其他目录（`magail_full`, `magail_production`）有事件文件
**可能原因**:
1. **TensorBoard 未安装**:
- 代码中有 try-except 处理（`train_magail.py:269-274`）
- 如果 TensorBoard 未安装，`writer` 会被设置为 `None`
- 训练会继续，但不会写入日志
2. **日志写入失败**:
- 即使 `SummaryWriter` 创建成功，如果写入时出错，可能不会生成文件
- 需要检查是否有异常被静默捕获
3. **训练中断**:
- 如果训练在写入第一个日志前中断，可能没有事件文件
- 但模型文件已保存，说明训练至少运行了一段时间
### 5.2 检查方法
**步骤 1: 检查 TensorBoard 安装**
```bash
python -c "import tensorboard; print(tensorboard.__version__)"
```
**步骤 2: 检查训练脚本中的日志写入**
查看 `train_magail.py:493-496`:
```python
if writer:
writer.add_scalar('Loss/Discriminator', disc_loss.item(), i_episode)
writer.add_scalar('Loss/Policy', ppo_loss, i_episode)
writer.add_scalar('Reward/Mean_GAIL', np.mean(all_gail_rewards), i_episode)
```
**步骤 3: 检查日志目录权限**
```bash
ls -la runs/magail_0112/
```
### 5.3 解决方案
**方案 1: 确保 TensorBoard 已安装**
```bash
pip install tensorboard
```
**方案 2: 添加显式刷新**
在训练循环结束后，显式调用 `writer.flush()`:
```python
if writer:
writer.flush() # 确保数据写入磁盘
```
**方案 3: 添加日志验证**
在训练开始时检查日志目录:
```python
if writer:
# 测试写入
writer.add_scalar('Test/Initialization', 0.0, 0)
writer.flush()
print(f"TensorBoard logging enabled. Log dir: {args.log_dir}")
else:
print("WARNING: TensorBoard not available. Logging disabled.")
```
**方案 4: 使用文件日志作为备份**
即使 TensorBoard 不可用，也可以写入文本日志:
```python
import logging
logging.basicConfig(
filename=os.path.join(args.log_dir, 'training.log'),
level=logging.INFO
)
```
### 5.4 代码修复建议
**在 `train_magail.py` 中添加以下改进**:
1. **确保 disc_loss 在 CPU 上**:
```python
# 第425行附近
disc_loss = torch.tensor(0.0).cuda() # 改为 .cuda() 或保持 CPU
# 或者在使用时转换
if writer:
disc_loss_value = disc_loss.item() if isinstance(disc_loss, torch.Tensor) else disc_loss
writer.add_scalar('Loss/Discriminator', disc_loss_value, i_episode)
```
2. **添加显式刷新**:
```python
# 第496行后添加
if writer:
writer.flush() # 确保数据写入磁盘
```
3. **添加初始化验证**:
```python
# 第271行后添加
if writer:
# 测试写入
writer.add_scalar('Test/Initialization', 0.0, 0)
writer.flush()
print(f"✓ TensorBoard logging enabled. Log dir: {args.log_dir}")
# 检查文件是否创建
import glob
event_files = glob.glob(os.path.join(args.log_dir, "events.out.tfevents.*"))
if event_files:
print(f"✓ TensorBoard event file created: {event_files[0]}")
else:
print("⚠ WARNING: TensorBoard not available. Logging disabled.")
```
4. **在训练结束时确保关闭**:
```python
# 第505行后添加
if writer:
writer.flush() # 最后一次刷新
writer.close()
print(f"TensorBoard logs saved to {args.log_dir}")
```
### 5.5 验证修复
**重新训练测试**:
```bash
python train_magail.py \
--expert_data_dir data/training_data \
--data_dir data/exp_filtered \
--batch_size 1024 \
--max_episodes 10 \
--log_dir runs/test_tensorboard
```
**检查输出**:
```bash
# 应该看到事件文件
ls runs/test_tensorboard/events.out.tfevents.*
# 启动 TensorBoard
tensorboard --logdir runs/test_tensorboard
```
**对于 magail_0112 训练**:
由于该训练已经完成且没有日志文件，建议：
1. 检查训练时的控制台输出，确认是否有 "TensorBoard not installed" 消息
2. 如果确实没有 TensorBoard，可以重新运行少量 episode 来验证修复
3. 或者查看是否有其他日志文件（如 `training.log`）
---
## 附录: 关键文件清单
### 核心训练文件
- `train_magail.py`: 主训练脚本
- `dataset/magail_dataset.py`: 专家数据集加载
- `Env/expert_replay_env.py`: 专家回放环境
- `Env/scenario_env.py`: 多智能体场景环境
- `Env/inverse_dynamics.py`: 逆动力学计算
### 数据生成文件
- `scripts/generate_expert_data.py`: 专家数据生成
- `scripts/visualize_replay.py`: 数据可视化
- `scripts/analyze_expert_data.py`: 数据分析
### 配置文件
- `README.md`: 项目说明
- `TRAINING_ARCHITECTURE.md`: 本文档
---
## 总结
本项目的 MAGAIL 训练方案通过以下方式处理多智能体可变长度问题：
1. **数据层面**: 将轨迹展平为独立样本，统一处理
2. **环境层面**: 动态处理不同数量的智能体，合并经验
3. **网络层面**: 固定输入维度（45维），通过零填充处理邻居不足的情况
训练流程遵循标准的 GAIL 框架，使用 PPO 作为策略优化算法。当前主要问题集中在环境稳定性和日志记录方面，需要进一步优化。

BIN
dataset/__pycache__/__init__.cpython-313.pyc Normal file
View File

Binary file not shown.

BIN
dataset/__pycache__/__init__.cpython-39.pyc Normal file
View File

Binary file not shown.

BIN
dataset/__pycache__/magail_dataset.cpython-313.pyc Normal file
View File

Binary file not shown.

BIN
dataset/__pycache__/magail_dataset.cpython-39.pyc Normal file
View File

Binary file not shown.

61
dataset/magail_dataset.py Normal file
View File

@@ -0,0 +1,61 @@
import torch
from torch.utils.data import Dataset
import pickle
import numpy as np
import os
import glob
class MAGAILExpertDataset(Dataset):
def __init__(self, data_dir, transform=None):
"""
Args:
data_dir (str): Directory containing .pkl files from generate_expert_data.py
transform (callable, optional): Optional transform to be applied on a sample.
"""
self.data_dir = data_dir
self.transform = transform
self.trajectories = []
self.flat_data = [] # (obs, act) pairs
# Load all .pkl files
pkl_files = glob.glob(os.path.join(data_dir, "*.pkl"))
print(f"Loading data from {len(pkl_files)} files in {data_dir}...")
for pkl_file in pkl_files:
try:
with open(pkl_file, 'rb') as f:
data = pickle.load(f)
# data is a list of dicts: {'obs': (T, 45), 'acts': (T, 2), ...}
self.trajectories.extend(data)
except Exception as e:
print(f"Error loading {pkl_file}: {e}")
# Flatten for training Discriminator/BC
print(f"Processing {len(self.trajectories)} trajectories...")
for traj in self.trajectories:
obs = traj['obs']
acts = traj['acts']
# obs: (T, 45), acts: (T, 2)
# We pair them up
for i in range(len(obs)):
self.flat_data.append((obs[i], acts[i]))
print(f"Total samples: {len(self.flat_data)}")
def __len__(self):
return len(self.flat_data)
def __getitem__(self, idx):
obs, act = self.flat_data[idx]
# Convert to tensor
obs = torch.from_numpy(obs).float()
act = torch.from_numpy(act).float()
sample = {'state': obs, 'action': act}
if self.transform:
sample = self.transform(sample)
return sample

2
scripts/analyze_expert_data.py
View File

@@ -247,7 +247,7 @@ class ExpertDataAnalyzer:
print(f" ✓ 分布图已保存")
if __name__ == "__main__":
WAYMO_DATA_DIR = r"/home/huangfukk/mdsn"
WAYMO_DATA_DIR = r"/home/huangfukk/MAGAIL4AutoDrive/data"
data_dir = AssetLoader.file_path(WAYMO_DATA_DIR, "exp_filtered", unix_style=False)
print("开始分析专家数据...")

162
scripts/generate_expert_data.py Normal file
View File

@@ -0,0 +1,162 @@
import argparse
import os
import sys
import pickle
import numpy as np
from tqdm import tqdm
# Add project root to Python path so we can import Env module
project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
if project_root not in sys.path:
sys.path.insert(0, project_root)
from metadrive.engine.asset_loader import AssetLoader
from Env.expert_replay_env import ExpertReplayEnv
def generate_data(args):
data_path = os.path.abspath(args.data_dir)
if not os.path.exists(data_path):
raise ValueError(f"Data directory {data_path} not found")
# MetaDrive's ScenarioDataManager asserts if config["num_scenarios"] > available scenarios in data_directory.
# So we always set it to -1 (load all available) and clamp the loop range by reading dataset summary.
from metadrive.scenario.utils import read_dataset_summary
_, summary_lookup, _ = read_dataset_summary(data_path)
if args.start_index >= len(summary_lookup):
raise ValueError(
f"start_index={args.start_index} out of range. Dataset has {len(summary_lookup)} scenarios."
)
max_available = len(summary_lookup) - args.start_index
num_to_run = min(args.num_scenarios, max_available)
env_config = {
"data_directory": data_path,
"is_multi_agent": True,
"num_controlled_agents": 100, # Set high to catch all vehicles in scenario
"horizon": 1000,
"use_render": False,
"sequential_seed": True,
"reactive_traffic": False, # Important: we replay, not react
"start_scenario_index": args.start_index,
# Load all scenarios available in the directory to avoid assertion failure.
# We will still only iterate `num_to_run` scenarios below.
"num_scenarios": -1,
"log_level": 50 # ERROR to reduce noise
}
expert_trajectories = []
try:
# Loop through scenarios
for i in tqdm(range(args.start_index, args.start_index + num_to_run), desc="Scenarios"):
env = ExpertReplayEnv(config=env_config)
try:
obs_dict = env.reset(seed=i)
except Exception as e:
print(f"Error resetting scenario {i}: {e}")
try:
env.close()
except Exception:
pass
continue
# Storage for current episode
# dict of lists: {agent_id: {'obs': [], 'acts': []}}
episode_data = {}
# Map agent_id to original ID if possible, but agent_id is unique enough
for step in range(env.config["horizon"]):
# Step with dummy actions
obs, rewards, dones, infos = env.step(None)
# 'obs' is next observation (t+1)
# 'infos' contains 'expert_action' which took (t -> t+1)
# Wait, usually (obs_t, act_t) -> obs_{t+1}
# expert_replay_env.step():
# calc action (t -> t+1)
# move agents to t+1
# return obs_{t+1}
# So we have obs_dict (from reset or prev step) which is at 't'
# And we have 'infos' which has action at 't'.
current_agents = list(obs_dict.keys())
for agent_id in current_agents:
if agent_id not in episode_data:
episode_data[agent_id] = {'obs': [], 'acts': []}
# Check if we have action for this agent
if agent_id in infos and 'expert_action' in infos[agent_id]:
action = infos[agent_id]['expert_action']
observation = obs_dict[agent_id]
episode_data[agent_id]['obs'].append(observation)
episode_data[agent_id]['acts'].append(action)
# Update obs_dict for next step
obs_dict = obs
if dones["__all__"]:
break
# Post-process episode data
for agent_id, data in episode_data.items():
if len(data['obs']) > 10: # Minimum length filter
expert_trajectories.append({
'obs': np.array(data['obs']),
'acts': np.array(data['acts']),
'agent_id': agent_id,
'scenario_id': i
})
env.close()
except Exception as e:
import traceback
traceback.print_exc()
print(f"Global error: {e}")
finally:
# env is closed per-scenario above (more robust for MetaDrive object lifecycle)
pass
# Save data
output_file = os.path.join(args.output_dir, f"expert_data_{args.start_index}_{args.num_scenarios}.pkl")
os.makedirs(args.output_dir, exist_ok=True)
print(f"Saving {len(expert_trajectories)} trajectories to {output_file}")
with open(output_file, 'wb') as f:
pickle.dump(expert_trajectories, f)
# Verification stats
if len(expert_trajectories) > 0:
all_acts = np.concatenate([t['acts'] for t in expert_trajectories])
print("Action Stats:")
print(f" Steering: min={all_acts[:,0].min():.3f}, max={all_acts[:,0].max():.3f}, mean={all_acts[:,0].mean():.3f}")
print(f" Accel: min={all_acts[:,1].min():.3f}, max={all_acts[:,1].max():.3f}, mean={all_acts[:,1].mean():.3f}")
# Clipping ratio diagnostics (actions are normalized to [-1, 1])
# If this ratio is high, it usually indicates max_acc/max_steering too small or noisy finite-difference.
eps = 1e-6
steer = all_acts[:, 0]
accel = all_acts[:, 1]
steer_clipped = np.isclose(np.abs(steer), 1.0, atol=eps)
accel_clipped = np.isclose(np.abs(accel), 1.0, atol=eps)
print("Clipping Stats:")
print(
f" Steering clipped (|a|==1): {steer_clipped.mean()*100:.2f}% "
f"({steer_clipped.sum()}/{len(steer_clipped)})"
)
print(
f" Accel clipped (|a|==1): {accel_clipped.mean()*100:.2f}% "
f"({accel_clipped.sum()}/{len(accel_clipped)})"
)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--data_dir", type=str, default="/home/huangfukk/MAGAIL4AutoDrive/data/exp_filtered", help="Path to Waymo pickles (or filtered index)")
parser.add_argument("--output_dir", type=str, default="/home/huangfukk/MAGAIL4AutoDrive/data/training", help="Output directory")
parser.add_argument("--start_index", type=int, default=0)
parser.add_argument("--num_scenarios", type=int, default=10)
args = parser.parse_args()
generate_data(args)

93
scripts/visualize_replay.py Normal file
View File

@@ -0,0 +1,93 @@
import argparse
import os
import sys
import time
# Add project root to Python path so we can import Env module
project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
if project_root not in sys.path:
sys.path.insert(0, project_root)
from Env.expert_replay_env import ExpertReplayEnv
def visualize_replay(args):
data_path = os.path.abspath(args.data_dir)
if not os.path.exists(data_path):
raise ValueError(f"Data directory {data_path} not found")
# Same as data generation: avoid MetaDrive assertion when requested num_scenarios > available.
from metadrive.scenario.utils import read_dataset_summary
_, summary_lookup, _ = read_dataset_summary(data_path)
if args.start_index >= len(summary_lookup):
raise ValueError(
f"start_index={args.start_index} out of range. Dataset has {len(summary_lookup)} scenarios."
)
max_available = len(summary_lookup) - args.start_index
num_to_run = min(args.num_scenarios, max_available)
env_config = {
"data_directory": data_path,
"is_multi_agent": True,
"num_controlled_agents": 100,
"horizon": args.horizon,
"use_render": True, # Enable rendering
"sequential_seed": True,
"reactive_traffic": False,
"start_scenario_index": args.start_index,
"num_scenarios": -1,
"log_level": 40, # ERROR
# "pstats": True, # For performance debugging
}
print(f"Initializing ExpertReplayEnv with data from {data_path}...")
env = ExpertReplayEnv(config=env_config)
try:
for i in range(args.start_index, args.start_index + num_to_run):
print(f"\n--- Playing Scenario {i} ---")
try:
obs = env.reset(seed=i)
except Exception as e:
print(f"Error resetting scenario {i}: {e}")
continue
print(f"Scenario loaded. Controlled agents: {len(env.controlled_agents)}")
for step in range(args.horizon):
# Step
obs, rewards, dones, infos = env.step(None)
# Render
env.render(mode="top_down",
text={
"Step": step,
"Agents": len(env.controlled_agents),
"Scenario": i
})
# Sleep to control playback speed
time.sleep(0.05)
if dones["__all__"]:
print(f"Scenario {i} finished at step {step}")
break
except KeyboardInterrupt:
print("Interrupted by user")
except Exception as e:
import traceback
traceback.print_exc()
print(f"Global error: {e}")
finally:
env.close()
print("Environment closed.")
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--data_dir", type=str, default="/home/huangfukk/MAGAIL4AutoDrive/data/exp_filtered", help="Path to Waymo data")
parser.add_argument("--start_index", type=int, default=0)
parser.add_argument("--num_scenarios", type=int, default=1)
parser.add_argument("--horizon", type=int, default=500)
args = parser.parse_args()
visualize_replay(args)

521
train_magail.py Normal file
View File

@@ -0,0 +1,521 @@
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.distributions import Normal
import numpy as np
import os
import argparse
import signal
import sys
from torch.utils.data import DataLoader
from dataset.magail_dataset import MAGAILExpertDataset
# --- Networks ---
class Actor(nn.Module):
def __init__(self, state_dim, action_dim, hidden_dim=256):
super(Actor, self).__init__()
self.net = nn.Sequential(
nn.Linear(state_dim, hidden_dim),
nn.Tanh(),
nn.Linear(hidden_dim, hidden_dim),
nn.Tanh(),
)
self.mu_head = nn.Linear(hidden_dim, action_dim)
self.log_std_head = nn.Parameter(torch.zeros(1, action_dim))
def forward(self, state):
x = self.net(state)
mu = torch.tanh(self.mu_head(x)) # Action range [-1, 1]
if mu.dim() == 1:
mu = mu.unsqueeze(0) # Handle single sample
log_std = self.log_std_head.expand_as(mu)
std = torch.exp(log_std)
dist = Normal(mu, std)
return dist
class Critic(nn.Module):
def __init__(self, state_dim, hidden_dim=256):
super(Critic, self).__init__()
self.net = nn.Sequential(
nn.Linear(state_dim, hidden_dim),
nn.Tanh(),
nn.Linear(hidden_dim, hidden_dim),
nn.Tanh(),
nn.Linear(hidden_dim, 1)
)
def forward(self, state):
return self.net(state)
class Discriminator(nn.Module):
def __init__(self, state_dim, action_dim, hidden_dim=256):
super(Discriminator, self).__init__()
self.net = nn.Sequential(
nn.Linear(state_dim + action_dim, hidden_dim),
nn.Tanh(),
nn.Linear(hidden_dim, hidden_dim),
nn.Tanh(),
nn.Linear(hidden_dim, 1),
nn.Sigmoid()
)
def forward(self, state, action):
x = torch.cat([state, action], dim=-1)
return self.net(x)
# --- PPO Algorithm ---
class PPO:
def __init__(self, state_dim, action_dim, lr=3e-4, gamma=0.99, eps_clip=0.2, K_epochs=10):
self.actor = Actor(state_dim, action_dim).cuda()
self.critic = Critic(state_dim).cuda()
self.optimizer_actor = optim.Adam(self.actor.parameters(), lr=lr)
self.optimizer_critic = optim.Adam(self.critic.parameters(), lr=lr)
self.gamma = gamma
self.eps_clip = eps_clip
self.K_epochs = K_epochs
self.mse_loss = nn.MSELoss()
def select_action(self, state):
with torch.no_grad():
state = torch.FloatTensor(state).cuda()
dist = self.actor(state)
action = dist.sample()
action_logprob = dist.log_prob(action).sum(dim=-1)
return action.cpu().numpy(), action_logprob.cpu().numpy()
def update(self, memory):
# Convert memory to tensors
states = torch.FloatTensor(np.array(memory['states'])).cuda()
actions = torch.FloatTensor(np.array(memory['actions'])).cuda()
logprobs = torch.FloatTensor(np.array(memory['logprobs'])).cuda()
rewards = torch.FloatTensor(np.array(memory['rewards'])).cuda()
next_states = torch.FloatTensor(np.array(memory['next_states'])).cuda()
dones = torch.FloatTensor(np.array(memory['dones'])).cuda()
# Monte Carlo estimate of state rewards (or GAE if implemented, simplistic here)
# Usually for PPO we use GAE. Let's do a simple discounted return for now or bootstrapping.
# Let's use bootstrapping from critic for returns.
returns = []
discounted_reward = 0
# This simple loop assumes full episode or consistent batch.
# For multi-agent disjoint steps, bootstrapping is better.
# But let's calculate advantage using GAE for stability.
values = self.critic(states).detach()
next_values = self.critic(next_states).detach()
# GAE
advantages = []
gae = 0
for i in reversed(range(len(rewards))):
delta = rewards[i] + self.gamma * next_values[i] * (1 - dones[i]) - values[i]
gae = delta + self.gamma * 0.95 * (1 - dones[i]) * gae
advantages.insert(0, gae)
advantages = torch.FloatTensor(advantages).cuda()
returns = advantages + values.squeeze()
# Optimize policy for K epochs:
for _ in range(self.K_epochs):
# Evaluating old actions and values :
dist = self.actor(states)
action_logprobs = dist.log_prob(actions).sum(dim=-1)
dist_entropy = dist.entropy().sum(dim=-1)
state_values = self.critic(states).squeeze()
# Finding the ratio (pi_theta / pi_theta__old):
ratios = torch.exp(action_logprobs - logprobs)
# Finding Surrogate Loss:
surr1 = ratios * advantages
surr2 = torch.clamp(ratios, 1-self.eps_clip, 1+self.eps_clip) * advantages
loss = -torch.min(surr1, surr2) + 0.5*self.mse_loss(state_values, returns) - 0.01*dist_entropy
# take gradient step
self.optimizer_actor.zero_grad()
self.optimizer_critic.zero_grad()
loss.mean().backward()
self.optimizer_actor.step()
self.optimizer_critic.step()
return loss.mean().item()
def save(self, checkpoint_path):
torch.save(self.actor.state_dict(), checkpoint_path + "_actor.pth")
torch.save(self.critic.state_dict(), checkpoint_path + "_critic.pth")
# --- Training Loop ---
def train(args):
# 1. Setup Environment (Dummy for now, usually you run simulation here)
# But for MAGAIL we need to collect generated trajectories.
# We need the Env class to be importable.
from Env.scenario_env import MultiAgentScenarioEnv
from Env.simple_idm_policy import ConstantVelocityPolicy # Just for init
# Config for Env
env_config = {
"data_directory": args.data_dir,
"is_multi_agent": True,
"num_controlled_agents": 3, # Dynamic
"horizon": 200,
"use_render": False,
"sequential_seed": True,
"start_scenario_index": 0,
"num_scenarios": args.num_scenarios # Use argument
}
# Ideally we use a wrapper for RL
# env = MultiAgentScenarioEnv(config=env_config) # This requires Waymo data loader setup
# 2. Setup Models
state_dim = 45
action_dim = 2
ppo_agent = PPO(state_dim, action_dim)
discriminator = Discriminator(state_dim, action_dim).cuda()
disc_optimizer = optim.Adam(discriminator.parameters(), lr=3e-4)
disc_criterion = nn.BCELoss()
# 3. Load Expert Data
expert_dataset = MAGAILExpertDataset(args.expert_data_dir)
# Ensure batch_size is not larger than dataset
if len(expert_dataset) < args.batch_size:
print(f"Warning: Expert dataset size {len(expert_dataset)} < batch_size {args.batch_size}. Adjusting batch_size.")
args.batch_size = len(expert_dataset)
if args.batch_size == 0:
raise ValueError("Expert dataset is empty!")
expert_loader = DataLoader(expert_dataset, batch_size=args.batch_size, shuffle=True, drop_last=True)
# Create an infinite iterator
def cycle(loader):
while True:
for batch in loader:
yield batch
expert_iter = cycle(expert_loader)
# 4. Initialize Env
from Env.expert_replay_env import ExpertReplayEnv # Using ReplayEnv for config, but we need ScenarioEnv for simulation?
# Actually we need MultiAgentScenarioEnv for interactive training, not Replay.
from Env.scenario_env import MultiAgentScenarioEnv
from Env.simple_idm_policy import ConstantVelocityPolicy # Placeholder policy for init
# 2. Setup Models
# Determine state dim from environment if possible, or use fixed
# Expert data has 45 dim?
# But Env might return something else if we are using default ScenarioEnv settings.
# ScenarioEnv returns list of obs.
# The error says: "mat1 and mat2 shapes cannot be multiplied (1x108 and 45x256)"
# This means the Env is returning 108-dim observation (MetaDrive default + Lidar),
# but our Actor expects 45 (which is what we saved in expert data).
# We must align the environment observation space with our expert data format.
# Our ExpertReplayEnv used a custom _get_all_obs.
# We need to inject that same logic into the training env, OR
# subclass MultiAgentScenarioEnv in the training script to override observation.
class MAGAILScenarioEnv(MultiAgentScenarioEnv):
def _get_all_obs(self):
# Same logic as ExpertReplayEnv to ensure compatibility
obs_dict = {}
for agent_id, vehicle in self.controlled_agents.items():
# 1. Ego State
ego_state = [
vehicle.position[0], vehicle.position[1],
vehicle.velocity[0], vehicle.velocity[1],
vehicle.heading_theta
]
# 2. Neighbors
candidates = []
for other_id, other_vehicle in self.engine.agent_manager.active_agents.items():
if other_id == agent_id:
continue
dist = np.linalg.norm(vehicle.position - other_vehicle.position)
if dist < 30.0:
candidates.append((dist, other_vehicle))
candidates.sort(key=lambda x: x[0])
top_10 = candidates[:10]
neighbor_feats = []
for _, neighbor in top_10:
neighbor_feats.extend([
neighbor.position[0] - vehicle.position[0],
neighbor.position[1] - vehicle.position[1],
neighbor.velocity[0],
neighbor.velocity[1]
])
missing = 10 - len(top_10)
if missing > 0:
neighbor_feats.extend([0.0] * (4 * missing))
obs = np.array(ego_state + neighbor_feats, dtype=np.float32)
obs_dict[agent_id] = obs
return obs_dict
env = MAGAILScenarioEnv(config=env_config, agent2policy={}) # Pass empty dict if we control all externally
print("Starting training...")
# Tensorboard
try:
from torch.utils.tensorboard import SummaryWriter
writer = SummaryWriter(log_dir=args.log_dir)
except ImportError:
print("TensorBoard not installed. Logging to console only.")
writer = None
global_step = 0
for i_episode in range(args.max_episodes):
# --- 1. Collect Rollouts (Interaction) ---
memory = {'states': [], 'actions': [], 'logprobs': [], 'rewards': [], 'next_states': [], 'dones': []}
# Prepare seed
available_scenarios = env.config["num_scenarios"]
start_index = env.config["start_scenario_index"]
seed = np.random.randint(start_index, start_index + available_scenarios)
# Reset Env
try:
# MetaDrive sometimes complains about uncleared objects if reset happens too fast or with lingering objs
# We can try to force clear before reset or handle exception
# But standard reset should handle it.
# The error "You should clear all generated objects..." means some manager didn't clear its objects.
# This is likely due to TrafficManager or AgentManager holding refs.
# Re-creating env is safer but slower.
# Let's try closing and re-creating if reset fails frequently.
# Or just ignore this error and try reset again? No, reset failing is fatal usually.
# Hack: Manually clear objects if we can access engine
if env.engine is not None:
env.engine.clear_objects(list(env.engine.get_objects().keys()))
obs_dict = env.reset(seed=seed)
except Exception as e:
# print(f"Env reset failed: {e}. Recreating environment...")
try:
env.close()
except:
pass
# Ensure engine is closed properly
from metadrive.engine.engine_utils import close_engine
try:
close_engine()
except Exception as e2:
# Force cleanup of singleton if close failed
from metadrive.engine.base_engine import BaseEngine
if BaseEngine.singleton is not None:
BaseEngine.singleton = None
# Also need to clear ShowBase
try:
from direct.showbase.ShowBase import ShowBase
if hasattr(base, 'destroy'):
base.destroy()
except:
pass
# Brutal force: delete base from builtins if it exists
import builtins
if hasattr(builtins, 'base'):
del builtins.base
# print(f"Error closing engine: {e2}")
# Explicitly delete old env object to free memory
del env
import gc
gc.collect()
env = MAGAILScenarioEnv(config=env_config, agent2policy={})
obs_dict = env.reset(seed=seed)
episode_reward = 0
steps = 0
# Rollout loop
while True:
# Select actions for all agents
actions = {}
action_logprobs = {}
# obs_dict: {agent_id: obs}
# MultiAgentScenarioEnv usually returns a dict {agent_id: obs}
# BUT wait, check scenario_env.py implementation
if isinstance(obs_dict, list):
# This happens if the environment returns a list instead of a dict
# MultiAgentScenarioEnv._get_all_obs returns a list in original implementation?
# Let's check scenario_env.py
# If it returns list, we need to map it to agent ids or just iterate
pass
# Temporary fix if it returns list (which means my previous edit to Env/expert_replay_env.py
# changed it there, but maybe not in Env/scenario_env.py which we are using here!)
if isinstance(obs_dict, list):
# We need agent IDs to step
# In MultiAgentScenarioEnv, controlled_agents is a dict.
# If obs is a list, it probably corresponds to controlled_agents.values() order?
# This is risky.
# Let's assume obs_dict is actually just observations.
# We need to keys to create action dict.
current_agent_ids = list(env.controlled_agents.keys())
# Ensure length matches
if len(obs_dict) != len(current_agent_ids):
# print(f"Warning: Obs list len {len(obs_dict)} != agents {len(current_agent_ids)}")
pass
# Reconstruct dict
new_obs_dict = {}
for i, agent_id in enumerate(current_agent_ids):
if i < len(obs_dict):
new_obs_dict[agent_id] = obs_dict[i]
obs_dict = new_obs_dict
for agent_id, obs in obs_dict.items():
act, logprob = ppo_agent.select_action(obs) # Select action returns numpy
actions[agent_id] = act.flatten() # (2,)
action_logprobs[agent_id] = logprob # scalar
# Step Env
next_obs_dict, rewards, dones, infos = env.step(actions)
# Store in memory
for agent_id, obs in obs_dict.items():
if agent_id in actions:
memory['states'].append(obs)
memory['actions'].append(actions[agent_id])
memory['logprobs'].append(action_logprobs[agent_id])
# Store standard environmental reward for logging (not used for update in GAIL)
# For GAIL update we use Discriminator reward later
memory['rewards'].append(0) # Placeholder
# Next state
if agent_id in next_obs_dict:
memory['next_states'].append(next_obs_dict[agent_id])
memory['dones'].append(False)
else:
# Agent finished/vanished
# We need a dummy next state or handle done correctly
# Just duplicate current state and mark done?
memory['next_states'].append(obs)
memory['dones'].append(True)
obs_dict = next_obs_dict
steps += 1
if dones["__all__"] or steps >= 200: # Limit horizon
break
# Initialize losses to 0/None before potential loop skip
disc_loss = torch.tensor(0.0)
ppo_loss = 0.0
all_gail_rewards = [0.0]
# --- 2. Train Discriminator ---
# Convert policy memory to tensors
policy_states = torch.FloatTensor(np.array(memory['states'])).cuda()
policy_actions = torch.FloatTensor(np.array(memory['actions'])).cuda()
# Sample expert batch
expert_batch = next(expert_iter)
expert_states = expert_batch['state'].cuda()
expert_actions = expert_batch['action'].cuda()
# Minibatch size matching
batch_size = min(policy_states.size(0), expert_states.size(0))
if batch_size > 0: # Only train if we have data
policy_states = policy_states[:batch_size]
policy_actions = policy_actions[:batch_size]
expert_states = expert_states[:batch_size]
expert_actions = expert_actions[:batch_size]
# Update Discriminator
# Label 1 for Expert, 0 for Policy
# Train Expert
disc_optimizer.zero_grad()
exp_preds = discriminator(expert_states, expert_actions)
exp_loss = disc_criterion(exp_preds, torch.ones_like(exp_preds))
pol_preds = discriminator(policy_states.detach(), policy_actions.detach()) # Detach policy data
pol_loss = disc_criterion(pol_preds, torch.zeros_like(pol_preds))
disc_loss = exp_loss + pol_loss
disc_loss.backward()
disc_optimizer.step()
# --- 3. Update Policy with GAIL Rewards ---
# Reward = -log(1 - D(s, a))
# Or more stable: log(D(s, a)) ? Original GAIL uses -log(1-D) which is log(D) roughly.
# Let's use -log(1 - D(s, a) + eps)
# Actually PPO needs the full trajectory for GAE.
# So we should compute rewards for ALL policy samples in memory.
all_policy_states = torch.FloatTensor(np.array(memory['states'])).cuda()
all_policy_actions = torch.FloatTensor(np.array(memory['actions'])).cuda()
with torch.no_grad():
all_d_val = discriminator(all_policy_states, all_policy_actions)
all_gail_rewards = -torch.log(1 - all_d_val + 1e-8).cpu().numpy().flatten()
# Replace placeholders
memory['rewards'] = all_gail_rewards.tolist()
# Update PPO
ppo_loss = ppo_agent.update(memory)
# Clean up memory
del policy_states, policy_actions, expert_states, expert_actions, exp_preds, exp_loss, pol_preds, pol_loss
del all_policy_states, all_policy_actions, all_d_val
torch.cuda.empty_cache()
else:
print(f"Episode {i_episode}: No data collected (Env might have crashed or no agents). Skipping update.")
# --- 4. Logging ---
if writer:
writer.add_scalar('Loss/Discriminator', disc_loss.item(), i_episode)
writer.add_scalar('Loss/Policy', ppo_loss, i_episode)
writer.add_scalar('Reward/Mean_GAIL', np.mean(all_gail_rewards), i_episode)
print(f"Episode {i_episode}: Disc Loss {disc_loss.item():.4f} | PPO Loss {ppo_loss:.4f} | Mean Reward {np.mean(all_gail_rewards):.4f}")
if i_episode % 50 == 0:
ppo_agent.save(os.path.join(args.log_dir, f"model_{i_episode}"))
env.close()
if writer:
writer.close()
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument("--expert_data_dir", type=str, default="data/training_data", help="Directory with .pkl expert data")
parser.add_argument("--data_dir", type=str, default="data/exp_filtered", help="Waymo data dir for Env")
parser.add_argument("--batch_size", type=int, default=1024)
parser.add_argument("--max_episodes", type=int, default=1000)
parser.add_argument("--num_scenarios", type=int, default=100)
parser.add_argument("--log_dir", type=str, default="runs/magail_exp")
args = parser.parse_args()
# Create log dir
os.makedirs(args.log_dir, exist_ok=True)
train(args)