564 lines
21 KiB
Python
564 lines
21 KiB
Python
"""
|
||
MAGAIL训练脚本
|
||
|
||
将Algorithm模块中的MAGAIL算法应用到多智能体环境中进行训练
|
||
"""
|
||
|
||
import os
|
||
import sys
|
||
import torch
|
||
import numpy as np
|
||
import pickle
|
||
import time
|
||
from datetime import datetime
|
||
from torch.utils.tensorboard import SummaryWriter
|
||
|
||
# 添加路径
|
||
sys.path.append(os.path.join(os.path.dirname(__file__), 'Algorithm'))
|
||
sys.path.append(os.path.join(os.path.dirname(__file__), 'Env'))
|
||
|
||
from Algorithm.magail import MAGAIL
|
||
from Algorithm.buffer import RolloutBuffer
|
||
from Env.scenario_env import MultiAgentScenarioEnv
|
||
from metadrive.engine.asset_loader import AssetLoader
|
||
|
||
|
||
class ExpertBuffer:
|
||
"""
|
||
专家数据缓冲区
|
||
|
||
从Waymo数据集中加载专家轨迹，用于GAIL判别器训练
|
||
"""
|
||
|
||
def __init__(self, data_dir, device, max_samples=100000):
|
||
"""
|
||
初始化专家缓冲区
|
||
|
||
Args:
|
||
data_dir: 专家数据目录
|
||
device: 计算设备
|
||
max_samples: 最大样本数
|
||
"""
|
||
self.device = device
|
||
self.max_samples = max_samples
|
||
self.states = []
|
||
self.next_states = []
|
||
|
||
print(f"📚 加载专家数据从: {data_dir}")
|
||
self._load_expert_data(data_dir)
|
||
|
||
# 数据已经在_extract_trajectories中转换为tensor并放到设备上了
|
||
if len(self.states) > 0:
|
||
print(f"✅ 加载完成: {len(self.states)} 条专家轨迹")
|
||
else:
|
||
print(f"⚠️ 警告: 未找到专家数据")
|
||
|
||
def _load_expert_data(self, data_dir):
|
||
"""
|
||
从pkl文件加载专家数据
|
||
|
||
注意: 这里需要根据实际的数据格式进行调整
|
||
"""
|
||
# 查找所有pkl文件
|
||
pkl_files = []
|
||
for root, dirs, files in os.walk(data_dir):
|
||
for file in files:
|
||
if file.endswith('.pkl') and 'sd_waymo' in file:
|
||
pkl_files.append(os.path.join(root, file))
|
||
|
||
if not pkl_files:
|
||
print(f"⚠️ 在 {data_dir} 中未找到专家数据文件")
|
||
return
|
||
|
||
print(f" 找到 {len(pkl_files)} 个数据文件")
|
||
|
||
# 只加载第一个文件作为示例
|
||
# 实际使用时可以加载多个文件
|
||
for pkl_file in pkl_files[:1]: # 只加载第一个文件
|
||
try:
|
||
with open(pkl_file, 'rb') as f:
|
||
data = pickle.load(f)
|
||
print(f" 正在处理: {os.path.basename(pkl_file)}")
|
||
self._extract_trajectories(data)
|
||
|
||
if len(self.states) >= self.max_samples:
|
||
break
|
||
except Exception as e:
|
||
print(f" ⚠️ 加载 {pkl_file} 失败: {e}")
|
||
|
||
def _extract_trajectories(self, scenario_data):
|
||
"""
|
||
从MetaDrive场景数据中提取车辆轨迹
|
||
|
||
Args:
|
||
scenario_data: MetaDrive格式的场景数据字典
|
||
"""
|
||
try:
|
||
# 方法1: 如果是字典且有'tracks'键
|
||
if isinstance(scenario_data, dict) and 'tracks' in scenario_data:
|
||
for vehicle_id, track_data in scenario_data['tracks'].items():
|
||
if track_data.get('type') == 'VEHICLE':
|
||
states = track_data.get('state', {})
|
||
|
||
# 获取有效帧
|
||
valid = states.get('valid', [])
|
||
if not hasattr(valid, 'any') or not valid.any():
|
||
continue
|
||
|
||
# 提取位置、速度、朝向等
|
||
positions = states.get('position', [])
|
||
velocities = states.get('velocity', [])
|
||
headings = states.get('heading', [])
|
||
|
||
# 构建state序列
|
||
for t in range(len(positions) - 1):
|
||
if valid[t] and valid[t+1]:
|
||
# 当前状态
|
||
state = np.concatenate([
|
||
positions[t][:2], # x, y
|
||
velocities[t], # vx, vy
|
||
[headings[t]], # heading
|
||
# ... 其他观测维度（激光雷达等暂时用0填充）
|
||
])
|
||
|
||
# 下一状态
|
||
next_state = np.concatenate([
|
||
positions[t+1][:2],
|
||
velocities[t+1],
|
||
[headings[t+1]],
|
||
])
|
||
|
||
# 补齐到108维（匹配实际观测维度）
|
||
state = np.pad(state, (0, 108 - len(state)))
|
||
next_state = np.pad(next_state, (0, 108 - len(next_state)))
|
||
|
||
# 转换为tensor并移到指定设备
|
||
self.states.append(torch.tensor(state, dtype=torch.float32, device=self.device))
|
||
self.next_states.append(torch.tensor(next_state, dtype=torch.float32, device=self.device))
|
||
|
||
if len(self.states) >= self.max_samples:
|
||
return
|
||
|
||
# 方法2: 其他可能的格式
|
||
# ...
|
||
|
||
except Exception as e:
|
||
print(f" ⚠️ 提取轨迹失败: {e}")
|
||
import traceback
|
||
traceback.print_exc()
|
||
|
||
def sample(self, batch_size):
|
||
"""
|
||
随机采样一批专家数据
|
||
|
||
Returns:
|
||
(states, next_states)
|
||
"""
|
||
if len(self.states) == 0:
|
||
# 如果没有专家数据，返回零张量
|
||
return (torch.zeros(batch_size, 108, device=self.device),
|
||
torch.zeros(batch_size, 108, device=self.device))
|
||
|
||
# 使用numpy随机采样（避免索引问题）
|
||
indices = np.random.randint(0, len(self.states), size=batch_size)
|
||
# 将list中的tensor堆叠成batch
|
||
states_batch = torch.stack([self.states[i] for i in indices])
|
||
next_states_batch = torch.stack([self.next_states[i] for i in indices])
|
||
return states_batch, next_states_batch
|
||
|
||
|
||
class MAGAILPolicy:
|
||
"""
|
||
MAGAIL策略包装器
|
||
|
||
将MAGAIL算法包装成环境可用的策略接口
|
||
"""
|
||
|
||
def __init__(self, magail_agent, device):
|
||
self.magail = magail_agent
|
||
self.device = device
|
||
|
||
def act(self, observation=None):
|
||
"""
|
||
执行动作（与环境兼容的接口）
|
||
|
||
注意: 由于环境调用方式的限制，这里采用简化处理
|
||
实际训练时需要通过主循环统一调用
|
||
"""
|
||
# 这个方法在训练时不会被使用
|
||
# 训练时统一通过 magail.explore() 获取动作
|
||
return [0.0, 0.0]
|
||
|
||
|
||
def collect_observations(env):
|
||
"""
|
||
收集所有智能体的观测
|
||
|
||
Args:
|
||
env: 多智能体环境
|
||
|
||
Returns:
|
||
obs_array: numpy数组 (n_agents, obs_dim)
|
||
"""
|
||
obs_list = env.obs_list
|
||
if len(obs_list) == 0:
|
||
return np.array([])
|
||
return np.array(obs_list)
|
||
|
||
|
||
def train_magail(
|
||
data_dir,
|
||
output_dir="./outputs",
|
||
num_episodes=1000,
|
||
horizon=300,
|
||
rollout_length=512, # 改为512，更适合300步的episode
|
||
batch_size=128, # 减小batch_size
|
||
lr_actor=3e-4,
|
||
lr_critic=3e-4,
|
||
lr_disc=3e-4,
|
||
epoch_disc=5,
|
||
epoch_ppo=10,
|
||
render=False,
|
||
device="cuda",
|
||
):
|
||
"""
|
||
MAGAIL训练主函数
|
||
|
||
Args:
|
||
data_dir: Waymo数据目录
|
||
output_dir: 输出目录（模型、日志）
|
||
num_episodes: 训练轮数
|
||
horizon: 每轮最大步数
|
||
rollout_length: PPO更新间隔
|
||
batch_size: 批次大小
|
||
lr_actor: Actor学习率
|
||
lr_critic: Critic学习率
|
||
lr_disc: 判别器学习率
|
||
epoch_disc: 判别器更新轮数
|
||
epoch_ppo: PPO更新轮数
|
||
render: 是否渲染
|
||
device: 计算设备
|
||
"""
|
||
|
||
# 创建输出目录
|
||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||
run_dir = os.path.join(output_dir, f"magail_{timestamp}")
|
||
os.makedirs(run_dir, exist_ok=True)
|
||
os.makedirs(os.path.join(run_dir, "models"), exist_ok=True)
|
||
|
||
# TensorBoard
|
||
writer = SummaryWriter(os.path.join(run_dir, "logs"))
|
||
|
||
# 设备
|
||
device = torch.device(device if torch.cuda.is_available() else "cpu")
|
||
print(f"🖥️ 使用设备: {device}")
|
||
|
||
# 观测和动作维度
|
||
# 根据scenario_env.py的_get_all_obs():
|
||
# 位置(2) + 速度(2) + 朝向(1) + 激光雷达(80) + 侧向(10) + 车道线(10) + 红绿灯(1) + 目标点(2) = 108
|
||
obs_dim = 108
|
||
action_dim = 2 # [转向, 油门/刹车]
|
||
|
||
print(f"📊 观测维度: {obs_dim}, 动作维度: {action_dim}")
|
||
|
||
# 加载专家数据
|
||
expert_buffer = ExpertBuffer(
|
||
data_dir=os.path.join(data_dir, "exp_converted"),
|
||
device=device,
|
||
max_samples=50000
|
||
)
|
||
|
||
# 初始化MAGAIL
|
||
print(f"🤖 初始化MAGAIL算法...")
|
||
magail = MAGAIL(
|
||
buffer_exp=expert_buffer,
|
||
input_dim=(obs_dim,),
|
||
device=device,
|
||
action_shape=(action_dim,),
|
||
rollout_length=rollout_length,
|
||
disc_coef=20.0,
|
||
disc_grad_penalty=0.1,
|
||
disc_logit_reg=0.25,
|
||
disc_weight_decay=0.0005,
|
||
lr_disc=lr_disc,
|
||
lr_actor=lr_actor,
|
||
lr_critic=lr_critic,
|
||
epoch_disc=epoch_disc,
|
||
epoch_ppo=epoch_ppo,
|
||
batch_size=batch_size,
|
||
use_gail_norm=True,
|
||
gamma=0.995,
|
||
lambd=0.97,
|
||
)
|
||
|
||
# 创建策略包装器
|
||
policy = MAGAILPolicy(magail, device)
|
||
|
||
# 环境配置（稍后为每个episode创建）
|
||
env_config = {
|
||
"data_directory": AssetLoader.file_path(data_dir, "exp_converted", unix_style=False),
|
||
"is_multi_agent": True,
|
||
"num_controlled_agents": 5,
|
||
"horizon": horizon,
|
||
"use_render": render,
|
||
"sequential_seed": True,
|
||
"reactive_traffic": True,
|
||
"manual_control": False,
|
||
"filter_offroad_vehicles": True,
|
||
"lane_tolerance": 3.0,
|
||
"max_controlled_vehicles": 5,
|
||
"debug_lane_filter": False,
|
||
"debug_traffic_light": False,
|
||
}
|
||
|
||
env = None # 每个episode创建新环境
|
||
|
||
print(f"\n{'='*60}")
|
||
print(f"🚀 开始训练 MAGAIL")
|
||
print(f"{'='*60}")
|
||
print(f"训练轮数: {num_episodes}")
|
||
print(f"每轮步数: {horizon}")
|
||
print(f"更新间隔: {rollout_length}")
|
||
print(f"输出目录: {run_dir}")
|
||
print(f"{'='*60}\n")
|
||
|
||
# 训练循环
|
||
total_steps = 0
|
||
best_reward = -float('inf')
|
||
|
||
for episode in range(num_episodes):
|
||
# 为每个episode创建新环境（避免MetaDrive的对象清理问题）
|
||
if env is not None:
|
||
env.close()
|
||
|
||
print(f"🌍 初始化Episode {episode + 1}环境...")
|
||
env = MultiAgentScenarioEnv(config=env_config, agent2policy=policy)
|
||
|
||
# 重置环境（场景索引要在范围内，循环使用场景）
|
||
scenario_index = episode % 3 # 只有3个场景，循环使用
|
||
obs_list = env.reset(scenario_index)
|
||
episode_reward = 0
|
||
episode_length = 0
|
||
|
||
# 检查是否有车辆
|
||
if len(env.controlled_agents) == 0:
|
||
print(f"⚠️ Episode {episode}: 没有可控车辆，跳过")
|
||
continue
|
||
|
||
print(f"\n📍 Episode {episode + 1}/{num_episodes}")
|
||
print(f" 可控车辆数: {len(env.controlled_agents)}")
|
||
|
||
for step in range(horizon):
|
||
# 收集观测
|
||
obs_array = collect_observations(env)
|
||
|
||
if len(obs_array) == 0:
|
||
break
|
||
|
||
# 策略采样动作
|
||
actions, log_pis = magail.explore(obs_array)
|
||
|
||
# 调试：打印第一步的动作（查看动作范围）
|
||
if step == 0 and episode == 0:
|
||
print(f"\n🔍 调试信息 - 第一个动作:")
|
||
print(f" 动作数量: {len(actions)}")
|
||
if len(actions) > 0:
|
||
print(f" 第一个动作: {actions[0]}")
|
||
print(f" 动作范围: [{np.min(actions):.3f}, {np.max(actions):.3f}]")
|
||
# 检查车辆初始位置
|
||
first_vehicle = list(env.controlled_agents.values())[0]
|
||
print(f" 第一辆车初始位置: {first_vehicle.position}")
|
||
print(f" 第一辆车初始速度: {first_vehicle.speed:.2f} m/s")
|
||
|
||
# 每50步打印一次位置变化
|
||
if step % 50 == 0 and step > 0 and episode == 0:
|
||
if len(env.controlled_agents) > 0:
|
||
first_vehicle = list(env.controlled_agents.values())[0]
|
||
print(f" 步数{step}: 位置={first_vehicle.position}, 速度={first_vehicle.speed:.2f}m/s")
|
||
|
||
# 构建动作字典
|
||
action_dict = {}
|
||
for i, agent_id in enumerate(env.controlled_agents.keys()):
|
||
if i < len(actions):
|
||
action_dict[agent_id] = actions[i]
|
||
|
||
# 环境步进
|
||
next_obs_list, rewards, dones, infos = env.step(action_dict)
|
||
next_obs_array = collect_observations(env)
|
||
|
||
# 渲染
|
||
if render:
|
||
env.render(mode="topdown")
|
||
time.sleep(0.02) # 20ms延迟，让渲染更平滑，约50fps
|
||
|
||
# 存储经验到buffer（为每个智能体存储）
|
||
for i, agent_id in enumerate(env.controlled_agents.keys()):
|
||
if i < len(obs_array) and i < len(actions) and i < len(next_obs_array):
|
||
# 获取该智能体的数据
|
||
state = obs_array[i]
|
||
action = actions[i]
|
||
reward = rewards.get(agent_id, 0.0)
|
||
done = dones.get(agent_id, False)
|
||
tm_done = done # 暂时使用相同的done标志
|
||
log_pi = log_pis[i]
|
||
next_state = next_obs_array[i]
|
||
|
||
# 获取策略参数
|
||
mean = magail.actor.means[i].detach().cpu().numpy() if i < len(magail.actor.means) else np.zeros(action_dim)
|
||
std = magail.actor.log_stds.exp()[0].detach().cpu().numpy()
|
||
|
||
# 存储到buffer
|
||
try:
|
||
magail.buffer.append(
|
||
state=torch.tensor(state, dtype=torch.float32, device=device),
|
||
action=action,
|
||
reward=reward,
|
||
done=done,
|
||
tm_dones=tm_done, # 修正参数名
|
||
log_pi=log_pi,
|
||
next_state=next_state,
|
||
next_state_gail=next_state,
|
||
means=mean,
|
||
stds=std
|
||
)
|
||
# 调试：只在第一个episode打印一次
|
||
if episode == 0 and step == 0 and i == 0:
|
||
print(f" ✅ 成功存入第一条数据 (buffer._n={magail.buffer._n})")
|
||
except Exception as e:
|
||
# 打印错误信息
|
||
if episode == 0 and step == 0:
|
||
print(f" ❌ buffer存储失败: {e}")
|
||
import traceback
|
||
traceback.print_exc()
|
||
|
||
# 计算平均奖励
|
||
avg_reward = np.mean(list(rewards.values())) if rewards else 0.0
|
||
episode_reward += avg_reward
|
||
episode_length += 1
|
||
total_steps += 1
|
||
|
||
# 检查是否结束
|
||
if dones.get("__all__", False):
|
||
break
|
||
|
||
# 定期更新
|
||
if total_steps % rollout_length == 0 and total_steps > 0:
|
||
print(f"\n 🔄 步数 {total_steps}: 更新模型...")
|
||
print(f" Buffer状态: _n={magail.buffer._n}, _p={magail.buffer._p}, buffer_size={magail.buffer.buffer_size}")
|
||
|
||
# 检查buffer是否有足够数据
|
||
if magail.buffer._n < batch_size:
|
||
print(f" ⚠️ Buffer数据不足: _n={magail.buffer._n} < batch_size={batch_size}, 跳过本次更新")
|
||
continue
|
||
|
||
try:
|
||
# 调用MAGAIL更新
|
||
gail_reward = magail.update(writer, total_steps)
|
||
print(f" GAIL奖励: {gail_reward:.4f}")
|
||
|
||
writer.add_scalar('Training/GAILReward', gail_reward, total_steps)
|
||
except Exception as e:
|
||
print(f" ⚠️ 更新失败: {e}")
|
||
if episode < 5: # 只在前几个episode打印详细信息
|
||
import traceback
|
||
traceback.print_exc()
|
||
|
||
# 记录训练指标
|
||
writer.add_scalar('Training/EpisodeReward', episode_reward, episode)
|
||
writer.add_scalar('Training/EpisodeLength', episode_length, episode)
|
||
|
||
# Episode结束
|
||
avg_episode_reward = episode_reward / max(episode_length, 1)
|
||
print(f" ✅ Episode {episode + 1} 完成:")
|
||
print(f" 步数: {episode_length}")
|
||
print(f" 总奖励: {episode_reward:.2f}")
|
||
print(f" 平均奖励: {avg_episode_reward:.4f}")
|
||
print(f" 车辆数: {len(env.controlled_agents)}")
|
||
|
||
# 记录到TensorBoard
|
||
writer.add_scalar('Episode/Reward', episode_reward, episode)
|
||
writer.add_scalar('Episode/Length', episode_length, episode)
|
||
writer.add_scalar('Episode/AvgReward', avg_episode_reward, episode)
|
||
writer.add_scalar('Episode/NumVehicles', len(env.controlled_agents), episode)
|
||
writer.add_scalar('Training/TotalSteps', total_steps, episode)
|
||
|
||
# 保存最佳模型
|
||
if episode_reward > best_reward:
|
||
best_reward = episode_reward
|
||
save_path = os.path.join(run_dir, "models", "best_model")
|
||
magail.save_models(save_path)
|
||
print(f" 💾 保存最佳模型 (奖励: {best_reward:.2f})")
|
||
|
||
# 定期保存
|
||
if (episode + 1) % 50 == 0:
|
||
save_path = os.path.join(run_dir, "models", f"checkpoint_{episode + 1}")
|
||
magail.save_models(save_path)
|
||
print(f" 💾 保存检查点: {save_path}")
|
||
|
||
# 训练结束
|
||
print(f"\n{'='*60}")
|
||
print(f"✅ 训练完成!")
|
||
print(f"{'='*60}")
|
||
print(f"总步数: {total_steps}")
|
||
print(f"最佳奖励: {best_reward:.2f}")
|
||
print(f"模型保存位置: {run_dir}/models")
|
||
print(f"日志位置: {run_dir}/logs")
|
||
|
||
# 关闭
|
||
env.close()
|
||
writer.close()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
import argparse
|
||
|
||
parser = argparse.ArgumentParser(description="训练MAGAIL算法")
|
||
parser.add_argument("--data-dir", type=str,
|
||
default="/home/huangfukk/MAGAIL4AutoDrive/Env",
|
||
help="Waymo数据目录")
|
||
parser.add_argument("--output-dir", type=str,
|
||
default="./outputs",
|
||
help="输出目录")
|
||
parser.add_argument("--episodes", type=int, default=1000,
|
||
help="训练轮数")
|
||
parser.add_argument("--horizon", type=int, default=300,
|
||
help="每轮最大步数")
|
||
parser.add_argument("--rollout-length", type=int, default=2048,
|
||
help="PPO更新间隔")
|
||
parser.add_argument("--batch-size", type=int, default=256,
|
||
help="批次大小")
|
||
parser.add_argument("--lr-actor", type=float, default=3e-4,
|
||
help="Actor学习率")
|
||
parser.add_argument("--lr-critic", type=float, default=3e-4,
|
||
help="Critic学习率")
|
||
parser.add_argument("--lr-disc", type=float, default=3e-4,
|
||
help="判别器学习率")
|
||
parser.add_argument("--epoch-disc", type=int, default=5,
|
||
help="判别器更新轮数")
|
||
parser.add_argument("--epoch-ppo", type=int, default=10,
|
||
help="PPO更新轮数")
|
||
parser.add_argument("--render", action="store_true",
|
||
help="是否渲染")
|
||
parser.add_argument("--device", type=str, default="cuda",
|
||
choices=["cuda", "cpu"],
|
||
help="计算设备")
|
||
|
||
args = parser.parse_args()
|
||
|
||
train_magail(
|
||
data_dir=args.data_dir,
|
||
output_dir=args.output_dir,
|
||
num_episodes=args.episodes,
|
||
horizon=args.horizon,
|
||
rollout_length=args.rollout_length,
|
||
batch_size=args.batch_size,
|
||
lr_actor=args.lr_actor,
|
||
lr_critic=args.lr_critic,
|
||
lr_disc=args.lr_disc,
|
||
epoch_disc=args.epoch_disc,
|
||
epoch_ppo=args.epoch_ppo,
|
||
render=args.render,
|
||
device=args.device,
|
||
)
|
||
|