14 KiB
MAGAIL 训练方案架构文档
目录
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,用零填充
- 每个邻居:
- Ego状态 (5维):
动作空间 (Action Space)
- 维度: 2维
- 组成:
[steering, accel] - 范围: 归一化到
[-1, 1]
1.2 数据格式
专家数据文件结构 (.pkl 文件):
# 每个 .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
流程:
- 从 Waymo 数据 (
data/exp_filtered) 加载场景 - 使用
ExpertReplayEnv回放专家轨迹 - 通过逆动力学 (
Env/inverse_dynamics.py) 计算动作 - 构建45维观测(Ego + 10个最近邻居)
- 过滤无效轨迹(长度 < 10)
- 保存为
.pkl文件到data/training_data/
关键代码位置:
- 观测构建:
Env/expert_replay_env.py的_get_all_obs()方法 - 动作计算:
Env/inverse_dynamics.py的compute_action()方法
2. 多智能体训练机制
2.1 可变长度处理
问题: 不同场景中智能体数量不同,每个智能体的轨迹长度也不同。
解决方案:
- 数据层面 (
dataset/magail_dataset.py):- 将轨迹展平为独立的
(state, action)对 - 每个样本是独立的,不保留序列信息
- 这样所有轨迹可以统一处理,不受长度限制
- 将轨迹展平为独立的
# 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,)
- 训练环境层面 (
train_magail.py):- 每个 episode 动态处理不同数量的智能体
- 在 rollout 循环中,为每个活跃智能体独立收集数据
- 所有智能体的数据合并到一个
memory中
# 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])
...
- 观测维度固定:
- 通过
MAGAILScenarioEnv确保观测维度始终为45维 - 邻居数量不足时用零填充,保证维度一致
- 通过
2.2 多智能体交互
环境设置:
- 使用
MAGAILScenarioEnv(继承自MultiAgentScenarioEnv) - 自定义
_get_all_obs()方法,确保观测格式与专家数据一致 - 每个智能体独立选择动作,环境统一执行
关键点:
- 所有智能体共享同一个策略网络(参数共享)
- 每个智能体独立计算动作和奖励
- 数据收集时将所有智能体的经验合并
3. 完整训练流程
3.1 数据准备阶段
步骤 1: 生成专家数据
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 模型初始化
网络架构:
-
Actor (策略网络):
- 输入: 45维状态
- 输出: 2维动作(连续)
- 结构: MLP (45 → 256 → 256 → 2)
- 输出分布: 高斯分布(均值 + 可学习标准差)
-
Critic (价值网络):
- 输入: 45维状态
- 输出: 标量价值
- 结构: MLP (45 → 256 → 256 → 1)
-
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 更新:
# 优势函数 (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)有事件文件
可能原因:
-
TensorBoard 未安装:
- 代码中有 try-except 处理(
train_magail.py:269-274) - 如果 TensorBoard 未安装,
writer会被设置为None - 训练会继续,但不会写入日志
- 代码中有 try-except 处理(
-
日志写入失败:
- 即使
SummaryWriter创建成功,如果写入时出错,可能不会生成文件 - 需要检查是否有异常被静默捕获
- 即使
-
训练中断:
- 如果训练在写入第一个日志前中断,可能没有事件文件
- 但模型文件已保存,说明训练至少运行了一段时间
5.2 检查方法
步骤 1: 检查 TensorBoard 安装
python -c "import tensorboard; print(tensorboard.__version__)"
步骤 2: 检查训练脚本中的日志写入
查看 train_magail.py:493-496:
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: 检查日志目录权限
ls -la runs/magail_0112/
5.3 解决方案
方案 1: 确保 TensorBoard 已安装
pip install tensorboard
方案 2: 添加显式刷新
在训练循环结束后,显式调用 writer.flush():
if writer:
writer.flush() # 确保数据写入磁盘
方案 3: 添加日志验证 在训练开始时检查日志目录:
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 不可用,也可以写入文本日志:
import logging
logging.basicConfig(
filename=os.path.join(args.log_dir, 'training.log'),
level=logging.INFO
)
5.4 代码修复建议
在 train_magail.py 中添加以下改进:
- 确保 disc_loss 在 CPU 上:
# 第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)
- 添加显式刷新:
# 第496行后添加
if writer:
writer.flush() # 确保数据写入磁盘
- 添加初始化验证:
# 第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.")
- 在训练结束时确保关闭:
# 第505行后添加
if writer:
writer.flush() # 最后一次刷新
writer.close()
print(f"TensorBoard logs saved to {args.log_dir}")
5.5 验证修复
重新训练测试:
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
检查输出:
# 应该看到事件文件
ls runs/test_tensorboard/events.out.tfevents.*
# 启动 TensorBoard
tensorboard --logdir runs/test_tensorboard
对于 magail_0112 训练: 由于该训练已经完成且没有日志文件,建议:
- 检查训练时的控制台输出,确认是否有 "TensorBoard not installed" 消息
- 如果确实没有 TensorBoard,可以重新运行少量 episode 来验证修复
- 或者查看是否有其他日志文件(如
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 训练方案通过以下方式处理多智能体可变长度问题:
- 数据层面: 将轨迹展平为独立样本,统一处理
- 环境层面: 动态处理不同数量的智能体,合并经验
- 网络层面: 固定输入维度(45维),通过零填充处理邻居不足的情况
训练流程遵循标准的 GAIL 框架,使用 PPO 作为策略优化算法。当前主要问题集中在环境稳定性和日志记录方面,需要进一步优化。