MAGAIL4AutoDrive/MAGAIL算法应用指南.md

MAGAIL算法应用指南

目录

1. Algorithm模块概览
2. 如何应用到环境
3. 完整训练流程
4. 当前实现状态
5. 需要完善的部分

---

Algorithm模块概览

📁 模块文件说明

Algorithm/
├── bert.py           # BERT判别器/价值网络
├── disc.py           # GAIL判别器（继承BERT）
├── policy.py         # 策略网络（Actor）
├── ppo.py            # PPO算法基类
├── magail.py         # MAGAIL主算法（继承PPO）
├── buffer.py         # 经验回放缓冲区
└── utils.py          # 工具函数（标准化等）

🔗 模块依赖关系

MAGAIL (magail.py)
  ├─ 继承 PPO (ppo.py)
  │   ├─ 使用 RolloutBuffer (buffer.py)
  │   ├─ 使用 StateIndependentPolicy (policy.py)
  │   └─ 使用 Bert作为Critic (bert.py)
  │
  ├─ 使用 GAILDiscrim (disc.py)
  │   └─ 继承 Bert (bert.py)
  │
  └─ 使用 Normalizer (utils.py)

---

如何应用到环境

✅ 已完成的准备工作

我已经为您：

1. 修复了PPO代码bug：添加了缺失的action_shape参数
2. 创建了训练脚本：train_magail.py
3. 提供了完整框架：包含环境初始化、训练循环、模型保存等

🚀 快速开始

方法1：使用训练脚本（推荐）

# 基本训练（使用默认参数）
python train_magail.py

# 自定义参数
python train_magail.py \
    --data-dir /path/to/waymo/data \
    --episodes 1000 \
    --horizon 300 \
    --batch-size 256 \
    --lr-actor 3e-4 \
    --render  # 可视化

# 查看所有参数
python train_magail.py --help

方法2：在Jupyter Notebook中使用

import sys
sys.path.append('Algorithm')
sys.path.append('Env')

from Algorithm.magail import MAGAIL
from Env.scenario_env import MultiAgentScenarioEnv

# 初始化环境
env = MultiAgentScenarioEnv(config={...})

# 初始化MAGAIL
magail = MAGAIL(
    buffer_exp=expert_buffer,
    input_dim=(108,),  # 观测维度
    device="cuda"
)

# 训练循环
for episode in range(1000):
    obs = env.reset()
    for step in range(300):
        actions, log_pis = magail.explore(obs)
        next_obs, rewards, dones, infos = env.step(actions)
        # ... 更新逻辑

---

完整训练流程

📊 数据流程图

┌─────────────────────────────────────────────────────────────┐
│                    MAGAIL训练流程                            │
└─────────────────────────────────────────────────────────────┘

第1步: 初始化
├─ 加载Waymo专家数据 → ExpertBuffer
├─ 创建MAGAIL算法实例
│  ├─ Actor (policy.py)
│  ├─ Critic (bert.py)
│  ├─ Discriminator (disc.py)
│  └─ Buffers (buffer.py)
└─ 创建多智能体环境

第2步: 训练循环
for episode in range(episodes):
    ├─ env.reset() → 重置环境，生成车辆
    │
    for step in range(horizon):
        ├─ obs = env._get_all_obs()  # 收集观测
        │
        ├─ actions = magail.explore(obs)  # 策略采样
        │
        ├─ next_obs, rewards, dones = env.step(actions)
        │
        ├─ buffer.append(obs, actions, rewards, ...)  # 存储经验
        │
        └─ if step % rollout_length == 0:
            ├─ 更新判别器
            │  ├─ 采样策略数据: buffer.sample()
            │  ├─ 采样专家数据: expert_buffer.sample()
            │  └─ update_disc(policy_data, expert_data)
            │
            ├─ 计算GAIL奖励
            │  └─ reward = -log(1 - D(s, s'))
            │
            └─ 更新PPO
               ├─ 计算GAE优势
               ├─ update_actor()
               └─ update_critic()

第3步: 评估与保存
└─ 保存模型、记录指标

🔑 关键代码段

1. 初始化MAGAIL

from Algorithm.magail import MAGAIL

magail = MAGAIL(
    buffer_exp=expert_buffer,      # 专家数据缓冲区
    input_dim=(obs_dim,),          # 观测维度 (108,)
    device=device,                 # cuda/cpu
    action_shape=(2,),             # 动作维度 [转向, 油门]
    
    # 判别器参数
    disc_coef=20.0,                # 判别器损失系数
    disc_grad_penalty=0.1,         # 梯度惩罚系数
    disc_logit_reg=0.25,           # Logit正则化
    disc_weight_decay=0.0005,      # 权重衰减
    lr_disc=3e-4,                  # 判别器学习率
    epoch_disc=5,                  # 判别器更新轮数
    
    # PPO参数
    rollout_length=2048,           # 更新间隔
    lr_actor=3e-4,                 # Actor学习率
    lr_critic=3e-4,                # Critic学习率
    epoch_ppo=10,                  # PPO更新轮数
    batch_size=256,                # 批次大小
    gamma=0.995,                   # 折扣因子
    lambd=0.97,                    # GAE lambda
    
    # 其他
    use_gail_norm=True,            # 使用数据标准化
)

2. 环境交互

# 重置环境
obs_list = env.reset(episode)

# 收集观测（所有车辆）
obs_array = np.array(env.obs_list)  # shape: (n_agents, 108)

# 策略采样
actions, log_pis = magail.explore(obs_array)
# actions: list of [转向, 油门] for each agent
# log_pis: list of log probabilities

# 构建动作字典
action_dict = {
    agent_id: actions[i] 
    for i, agent_id in enumerate(env.controlled_agents.keys())
}

# 环境步进
next_obs, rewards, dones, infos = env.step(action_dict)

3. 模型更新

from torch.utils.tensorboard import SummaryWriter

writer = SummaryWriter('logs')

# 更新判别器和策略
if total_steps % rollout_length == 0:
    # MAGAIL会自动：
    # 1. 从buffer采样策略数据
    # 2. 从expert_buffer采样专家数据
    # 3. 更新判别器
    # 4. 计算GAIL奖励
    # 5. 更新PPO（Actor + Critic）
    
    reward = magail.update(writer, total_steps)
    
    print(f"Step {total_steps}, Reward: {reward:.4f}")

4. 保存和加载模型

# 保存
magail.save_models('outputs/models/checkpoint_100')

# 加载
magail.load_models('outputs/models/checkpoint_100/model.pth')

---

当前实现状态

✅ 已实现功能

模块	状态	说明
BERT判别器	✅ 完整	支持动态车辆数量
GAIL判别器	✅ 完整	包含梯度惩罚、正则化
策略网络	✅ 完整	高斯策略，重参数化
PPO算法	✅ 完整	GAE、裁剪目标、自适应LR
MAGAIL	✅ 完整	判别器+PPO整合
缓冲区	✅ 完整	经验存储和采样
数据标准化	✅ 完整	运行时统计量
环境接口	✅ 完整	多智能体场景环境

⚠️ 需要注意的问题

1. 多智能体适配问题

当前状态： Algorithm模块设计为单智能体，但环境是多智能体

影响：

• buffer.append() 接受单个状态-动作对
• 但环境返回多个智能体的数据

解决方案A： 将所有智能体视为一个整体

# 拼接所有智能体的观测
all_obs = np.concatenate([obs for obs in obs_list])
all_actions = np.concatenate([actions for actions in action_list])

解决方案B： 为每个智能体独立存储

for i, agent_id in enumerate(env.controlled_agents):
    buffer.append(obs_list[i], actions[i], rewards[i], ...)

推荐： 解决方案B，因为MAGAIL的设计就是处理多智能体的

2. 专家数据加载

当前状态： ExpertBuffer 类只有框架，未实现实际加载

需要完善：

def _extract_trajectories(self, scenario_data):
    """
    需要根据Waymo数据格式实现
    
    示例结构：
    scenario_data = {
        'tracks': {
            'vehicle_id': {
                'states': [...],  # 状态序列
                'actions': [...], # 动作序列（如果有）
            }
        }
    }
    """
    # TODO: 提取state和next_state对
    for track_id, track_data in scenario_data['tracks'].items():
        states = track_data['states']
        for i in range(len(states) - 1):
            self.states.append(states[i])
            self.next_states.append(states[i+1])

3. 观测维度对齐

当前假设： 观测维度为108

• 位置(2) + 速度(2) + 朝向(1) + 激光雷达(80) + 侧向(10) + 车道线(10) + 红绿灯(1) + 目标点(2) = 108

需要验证： 实际运行时打印观测shape

obs = env.reset()
print(f"观测维度: {len(obs[0]) if obs else 0}")

---

需要完善的部分

🔨 短期TODO

1. 修复多智能体buffer问题

创建文件： Algorithm/multi_agent_buffer.py

class MultiAgentRolloutBuffer:
    """
    多智能体经验缓冲区
    
    支持动态数量的智能体
    """
    
    def __init__(self, buffer_size, state_shape, action_shape, device):
        self.buffer_size = buffer_size
        self.state_shape = state_shape
        self.action_shape = action_shape
        self.device = device
        
        # 使用列表存储，支持动态智能体数量
        self.episodes = []
        self.current_episode = {
            'states': [],
            'actions': [],
            'rewards': [],
            'dones': [],
            'log_pis': [],
            'next_states': [],
        }
    
    def append(self, state, action, reward, done, log_pi, next_state):
        """添加单步经验"""
        self.current_episode['states'].append(state)
        self.current_episode['actions'].append(action)
        self.current_episode['rewards'].append(reward)
        self.current_episode['dones'].append(done)
        self.current_episode['log_pis'].append(log_pi)
        self.current_episode['next_states'].append(next_state)
    
    def finish_episode(self):
        """完成一个episode"""
        self.episodes.append(self.current_episode)
        self.current_episode = {
            'states': [],
            'actions': [],
            'rewards': [],
            'dones': [],
            'log_pis': [],
            'next_states': [],
        }
    
    def sample(self, batch_size):
        """采样批次"""
        # 从所有episode中随机采样
        all_states = []
        all_next_states = []
        
        for episode in self.episodes:
            all_states.extend(episode['states'])
            all_next_states.extend(episode['next_states'])
        
        indices = np.random.choice(len(all_states), batch_size, replace=False)
        
        states = torch.tensor([all_states[i] for i in indices], device=self.device)
        next_states = torch.tensor([all_next_states[i] for i in indices], device=self.device)
        
        return states, next_states

2. 实现专家数据加载

需要了解： Waymo数据的实际格式

# 示例：读取一个pkl文件并打印结构
import pickle

with open('Env/exp_converted/exp_converted_0/sd_waymo_*.pkl', 'rb') as f:
    data = pickle.load(f)
    print(type(data))
    print(data.keys() if isinstance(data, dict) else len(data))
    # 根据实际结构调整加载代码

3. 完善训练循环

在 train_magail.py 中添加：

# 完整的buffer存储逻辑
for i, agent_id in enumerate(env.controlled_agents.keys()):
    if i < len(obs_array) and i < len(actions):
        magail.buffer.append(
            state=obs_array[i],
            action=actions[i],
            reward=rewards.get(agent_id, 0.0),
            done=dones.get(agent_id, False),
            tm_done=dones.get(agent_id, False),
            log_pi=log_pis[i],
            next_state=next_obs_array[i] if i < len(next_obs_array) else obs_array[i],
            next_state_gail=next_obs_array[i] if i < len(next_obs_array) else obs_array[i],
            means=magail.actor.means[i].detach().cpu().numpy(),
            stds=magail.actor.log_stds.exp()[0].detach().cpu().numpy()
        )

🎯 中期TODO

1. 实现多智能体BERT：当前BERT接受(batch, N, obs_dim)，需要确保正确处理
2. 奖励设计：当前环境奖励为0，需要设计合理的任务奖励
3. 评估脚本：创建评估脚本，可视化训练好的策略
4. 超参数调优：使用wandb或tensorboard进行超参数搜索

---

使用示例

示例1：简单训练

# 1. 确保环境正常
python Env/run_multiagent_env.py

# 2. 开始训练（不渲染，快速训练）
python train_magail.py \
    --episodes 100 \
    --horizon 200 \
    --rollout-length 1024 \
    --batch-size 128

# 3. 查看训练日志
tensorboard --logdir outputs/magail_*/logs

示例2：调试模式

# 少量episode，启用渲染
python train_magail.py \
    --episodes 5 \
    --horizon 100 \
    --render

示例3：在代码中使用

# test_algorithm.py
import sys
sys.path.append('Algorithm')

from Algorithm.magail import MAGAIL
import torch

# 创建虚拟数据测试
class DummyExpertBuffer:
    def __init__(self, device):
        self.device = device
    
    def sample(self, batch_size):
        states = torch.randn(batch_size, 108, device=self.device)
        next_states = torch.randn(batch_size, 108, device=self.device)
        return states, next_states

# 初始化
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
expert_buffer = DummyExpertBuffer(device)

magail = MAGAIL(
    buffer_exp=expert_buffer,
    input_dim=(108,),
    device=device,
    action_shape=(2,),
)

# 测试前向传播
test_obs = torch.randn(5, 108, device=device)  # 5个智能体
actions, log_pis = magail.explore(test_obs)

print(f"观测形状: {test_obs.shape}")
print(f"动作数量: {len(actions)}")
print(f"单个动作形状: {actions[0].shape}")
print(f"测试成功！✅")

---

总结

✅ 现在可以做什么

1. 运行环境测试：run_multiagent_env.py 已经可以正常运行
2. 测试算法模块：Algorithm中的所有模块都已实现
3. 开始初步训练：使用 train_magail.py（但需要完善buffer逻辑）

⚠️ 需要您完成的

1. 调试多智能体buffer：确保经验正确存储
2. 实现专家数据加载：根据实际数据格式调整
3. 验证观测维度：确认实际观测是否为108维
4. 调整训练参数：根据训练效果调优

🎯 最终目标

环境 (Env/) + 算法 (Algorithm/) = 完整的MAGAIL训练系统
                    ↓
            训练出能够模仿专家行为的
              多智能体自动驾驶策略

祝训练顺利！🚀