环境代码优化

Signed-off-by: huangfu <3045324663@qq.com>

.gitignore

@@ -0,0 +1,57 @@
+# 日志文件
+Env/logs/
+*.log
+
+# Python
+__pycache__/
+*.py[cod]
+*$py.class
+*.so
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+
+# 虚拟环境
+venv/
+env/
+ENV/
+.venv
+
+# IDE
+.vscode/
+.idea/
+*.swp
+*.swo
+*~
+
+# 数据和模型文件
+data/
+runs/
+*.pkl
+*.h5
+*.ckpt
+*.pth
+*.pt
+checkpoints/
+models/
+
+# 第三方库（如果已安装）
+metadrive/
+scenarionet/
+
+# 系统文件
+.DS_Store
+Thumbs.db

.vscode/settings.json

@@ -1,6 +0,0 @@
-{
-    "cursorpyright.analysis.extraPaths": [
-        "/home/huangfukk/mdsn/metadrive",
-        "/home/huangfukk/mdsn/scenarionet"
-    ]
-}

Algorithm/bc.py

@@ -0,0 +1,52 @@
+"""
+Behavior Cloning (BC) 算法：仅包含损失与单 epoch 训练/评估逻辑。
+数据加载、环境评估、日志与保存由训练脚本 (train_bc.py) 负责。
+"""
+import torch
+
+
+def bc_loss(policy, states, actions):
+    """
+    BC 损失：负对数似然 -E[log pi(a|s)]。
+    states: (B, state_dim), actions: (B, action_dim), 均在 policy 所在 device 上。
+    """
+    log_pi = policy.evaluate_log_pi(states, actions)
+    return -log_pi.mean()
+
+
+def train_bc_epoch(policy, train_loader, optimizer, device):
+    """
+    训练一个 epoch，返回平均 train loss。
+    policy 与 optimizer 由调用方管理，本函数只做前向、损失、反向与 step。
+    """
+    policy.train()
+    total_loss = 0.0
+    n_batches = 0
+    for states, actions in train_loader:
+        states = states.to(device)
+        actions = actions.to(device)
+        loss = bc_loss(policy, states, actions)
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+        total_loss += loss.item()
+        n_batches += 1
+    return total_loss / n_batches if n_batches else 0.0
+
+
+def eval_bc_epoch(policy, val_loader, device):
+    """
+    在验证集上评估一个 epoch，返回平均 val loss（无梯度）。
+    """
+    policy.eval()
+    total_loss = 0.0
+    n_batches = 0
+    with torch.no_grad():
+        for states, actions in val_loader:
+            states = states.to(device)
+            actions = actions.to(device)
+            log_pi = policy.evaluate_log_pi(states, actions)
+            loss = -log_pi.mean().item()
+            total_loss += loss
+            n_batches += 1
+    return total_loss / n_batches if n_batches else 0.0

CHANGELOG.md

@@ -0,0 +1,136 @@
+# 更新日志
+
+## 2025-01-20 问题修复与优化
+
+### ✅ 已解决的问题
+
+#### 1. 车辆生成位置偏差问题
+**问题描述：** 部分车辆生成于草坪、停车场等非车道区域
+
+**解决方案：**
+- 实现 `_is_position_on_lane()` 方法：检测位置是否在有效车道上
+- 实现 `_filter_valid_spawn_positions()` 方法：自动过滤非车道区域车辆
+- 支持容差参数（默认3米）处理边界情况
+- 在 `reset()` 时自动执行过滤，并输出统计信息
+
+**配置参数：**
+```python
+"filter_offroad_vehicles": True,  # 启用/禁用过滤
+"lane_tolerance": 3.0,           # 容差范围（米）
+"max_controlled_vehicles": 10,   # 最大车辆数限制
+```
+
+#### 2. 红绿灯信息采集问题
+**问题描述：**
+- 部分红绿灯状态为 None
+- 车道分段时部分车辆无法获取红绿灯状态
+
+**解决方案：**
+- 实现 `_get_traffic_light_state()` 方法，采用双重检测策略
+- 方法1（优先）：从导航模块获取当前车道，直接查询（高效）
+- 方法2（兜底）：遍历所有车道匹配位置（处理特殊情况）
+- 完善异常处理，None 状态返回 0（无红绿灯）
+- 返回值：0=无/未知, 1=绿灯, 2=黄灯, 3=红灯
+
+#### 3. 性能优化问题
+**问题描述：** FPS只有15帧，CPU利用率不到20%
+
+**解决方案：**
+- 创建 `run_multiagent_env_fast.py`：激光雷达优化版（30-60 FPS）
+- 创建 `run_multiagent_env_parallel.py`：多进程并行版（300-600 steps/s）
+- 提供详细的性能优化文档
+
+### 📝 修改的文件
+
+1. **Env/scenario_env.py**
+   - 新增 `_is_position_on_lane()` 方法
+   - 新增 `_filter_valid_spawn_positions()` 方法
+   - 新增 `_get_traffic_light_state()` 方法
+   - 更新 `default_config()` 添加配置参数
+   - 更新 `reset()` 调用过滤逻辑
+   - 更新 `_get_all_obs()` 使用新的红绿灯检测方法
+
+2. **Env/run_multiagent_env.py**
+   - 添加车道过滤配置参数
+
+3. **Env/run_multiagent_env_fast.py**
+   - 添加车道过滤配置
+   - 性能优化配置
+
+4. **Env/run_multiagent_env_parallel.py**
+   - 添加车道过滤配置
+   - 多进程并行实现
+
+5. **README.md**
+   - 更新问题说明，添加解决方案
+   - 添加配置示例和测试方法
+   - 添加问题解决总结
+
+6. **新增文件**
+   - `Env/test_lane_filter.py`：功能测试脚本
+
+### 🧪 测试方法
+
+```bash
+# 测试车道过滤和红绿灯检测功能
+python Env/test_lane_filter.py
+
+# 运行标准版本（带过滤和可视化）
+python Env/run_multiagent_env.py
+
+# 运行高性能版本（适合训练）
+python Env/run_multiagent_env_fast.py
+
+# 运行多进程并行版本（最高吞吐量）
+python Env/run_multiagent_env_parallel.py
+```
+
+### 💡 使用建议
+
+1. **调试阶段**：使用 `run_multiagent_env.py`，启用渲染和车道过滤
+2. **训练阶段**：使用 `run_multiagent_env_fast.py`，关闭渲染，启用所有优化
+3. **大规模训练**：使用 `run_multiagent_env_parallel.py`，充分利用多核CPU
+
+### ⚙️ 配置说明
+
+所有配置参数都可以在创建环境时通过 `config` 字典传递：
+
+```python
+env = MultiAgentScenarioEnv(
+    config={
+        # 基础配置
+        "data_directory": "...",
+        "is_multi_agent": True,
+        "horizon": 300,
+        
+        # 车道过滤（新增）
+        "filter_offroad_vehicles": True,  # 启用车道过滤
+        "lane_tolerance": 3.0,           # 容差3米
+        "max_controlled_vehicles": 10,   # 最多10辆车
+        
+        # 性能优化
+        "use_render": False,
+        "decision_repeat": 5,
+        ...
+    },
+    agent2policy=your_policy
+)
+```
+
+### 🔍 技术细节
+
+**车道检测逻辑：**
+1. 使用 `lane.lane.point_on_lane()` 精确检测
+2. 使用 `lane.local_coordinates()` 计算横向距离
+3. 支持容差参数处理边界情况
+
+**红绿灯检测逻辑：**
+1. 优先从 `vehicle.navigation.current_lane` 获取
+2. 失败时遍历所有车道查找
+3. 所有异常均有保护，确保稳定性
+
+**性能优化原理：**
+- 减少激光束数量降低计算量
+- 多进程绕过Python GIL限制
+- 充分利用多核CPU
+

Env/bc_env.py

@@ -0,0 +1,164 @@
+from Env.scenario_env import MultiAgentScenarioEnv
+from Env.utils import filter_traffic_tracks_to_birth_lists
+from metadrive.component.vehicle.vehicle_type import DefaultVehicle
+import numpy as np
+
+
+class BCScenarioEnv(MultiAgentScenarioEnv):
+    """
+    Environment for Behavior Cloning Evaluation.
+    Uses the same 45-dim observation as ExpertReplayEnv:
+    - Ego State (5): x, y, vx, vy, heading
+    - Neighbors (40): 10 nearest * (rel_x, rel_y, vx, vy)
+
+    Spawns background (static) vehicles so that observation distribution matches expert data collection:
+    expert data is generated with ExpertReplayEnv which includes bg_* in active_agents, so the policy
+    was trained on obs that can include those neighbors. Demo should use the same scene for consistency.
+    """
+    def reset(self, seed=None):
+        # Clear background vehicles from previous episode so engine.reset() passes _object_clean_check
+        if getattr(self, "engine", None) is not None:
+            ids_bg = [
+                oid for oid, obj in self.engine.get_objects().items()
+                if (getattr(obj, "name", None) or getattr(obj, "id", None) or "").startswith("bg_")
+            ]
+            if ids_bg:
+                self.engine.clear_objects(ids_bg, force_destroy=True)
+            for aid in list(self.engine.agent_manager.active_agents.keys()):
+                if aid.startswith("bg_"):
+                    self.engine.agent_manager.active_agents.pop(aid, None)
+        obs = super().reset(seed=seed)
+        self._spawn_background_vehicles()
+        return self._get_all_obs()
+
+    def _build_birth_lists_from_traffic(self):
+        """Same lane/static filter as expert data; return background_vehicles so we spawn them (match training obs)."""
+        car_birth_info_list, background_vehicles, obj_to_clean, stats = filter_traffic_tracks_to_birth_lists(
+            self.engine.traffic_manager.current_traffic_data,
+            self.engine.traffic_manager.sdc_scenario_id,
+            self.engine.map_manager,
+            return_stats=True,
+        )
+        if stats["n_controlled"] == 0 and stats["n_total"] > 0:
+            print(
+                "[BCScenarioEnv] 0 controlled agents: total_vehicles={}, off_lane={}, static={}, no_valid={}.".format(
+                    stats["n_total"],
+                    stats["n_off_lane"],
+                    stats["n_static"],
+                    stats["n_no_valid"],
+                )
+            )
+        return car_birth_info_list, background_vehicles, obj_to_clean
+
+    def _spawn_background_vehicles(self):
+        """Spawn static background vehicles so they appear in active_agents and thus in obs (same as ExpertReplayEnv)."""
+        for sid, car in self.background_vehicles.items():
+            if car["show_time"] != self.round:
+                continue
+            bg_id = f"bg_{car['id']}"
+            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"],
+            )
+            v.set_velocity([0, 0])
+            self.engine.agent_manager.active_agents[bg_id] = v
+            v.valid_mask = car["valid"]
+            v.start_t = car["show_time"]
+
+    def _update_background_vehicles(self):
+        self._spawn_background_vehicles()
+        to_remove = []
+        objects_to_clear = []
+        for aid, v in self.engine.agent_manager.active_agents.items():
+            if not aid.startswith("bg_"):
+                continue
+            if hasattr(v, "valid_mask"):
+                if self.round >= len(v.valid_mask) or not v.valid_mask[self.round]:
+                    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([v.id for v in objects_to_clear])
+
+    def step(self, action_dict):
+        self.round += 1
+        for agent_id, action in action_dict.items():
+            if agent_id in self.controlled_agents:
+                self.controlled_agents[agent_id].before_step(action)
+        self.engine.step()
+        self.engine.after_step()
+        for agent_id in action_dict:
+            if agent_id in self.controlled_agents:
+                self.controlled_agents[agent_id].after_step()
+        self._spawn_controlled_agents()
+        self._update_background_vehicles()
+        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.episode_step >= self.config["horizon"]
+        infos = {aid: {} 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 = []
+            # Use engine.agent_manager.active_agents to find neighbors
+            # Note: This includes background vehicles if they are in active_agents
+            for other_id, other_vehicle in self.engine.agent_manager.active_agents.items():
+                if other_id == agent_id:
+                    continue
+                
+                # Check if vehicle is valid/active
+                # (MetaDrive manages active_agents, so they should be active)
+                
+                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
+                    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

Env/check_dataset.py

@@ -1,109 +0,0 @@
-#!/usr/bin/env python
-# -*- coding: utf-8 -*-
-"""
-数据集检查脚本：统计可用的场景数量
-"""
-
-import pickle
-import os
-from pathlib import Path
-from metadrive.engine.asset_loader import AssetLoader
-
-def check_dataset(data_dir, subfolder="exp_filtered"):
-    """
-    检查数据集中的场景数量
-    
-    Args:
-        data_dir: 数据根目录
-        subfolder: 数据子目录（exp_filtered 或 exp_converted）
-    """
-    print("=" * 80)
-    print("数据集检查工具")
-    print("=" * 80)
-    
-    # 获取完整路径
-    full_path = AssetLoader.file_path(data_dir, subfolder, unix_style=False)
-    print(f"\n数据集路径: {full_path}")
-    
-    # 检查文件结构
-    if not os.path.exists(full_path):
-        print(f"❌ 错误：路径不存在: {full_path}")
-        return
-    
-    print(f"✅ 路径存在")
-    
-    # 读取数据集映射
-    mapping_file = os.path.join(full_path, "dataset_mapping.pkl")
-    if os.path.exists(mapping_file):
-        print(f"\n读取 dataset_mapping.pkl...")
-        with open(mapping_file, 'rb') as f:
-            dataset_mapping = pickle.load(f)
-        
-        print(f"✅ 数据集映射文件存在")
-        print(f"   映射的场景数量: {len(dataset_mapping)}")
-        
-        # 统计各个子目录的分布
-        subdirs = {}
-        for filename, subdir in dataset_mapping.items():
-            if subdir not in subdirs:
-                subdirs[subdir] = []
-            subdirs[subdir].append(filename)
-        
-        print(f"\n场景分布：")
-        for subdir, files in subdirs.items():
-            print(f"  {subdir}: {len(files)} 个场景")
-    
-    # 读取数据集摘要
-    summary_file = os.path.join(full_path, "dataset_summary.pkl")
-    if os.path.exists(summary_file):
-        print(f"\n读取 dataset_summary.pkl...")
-        with open(summary_file, 'rb') as f:
-            dataset_summary = pickle.load(f)
-        
-        print(f"✅ 数据集摘要文件存在")
-        print(f"   摘要的场景数量: {len(dataset_summary)}")
-        
-        # 打印前几个场景的ID
-        print(f"\n前10个场景ID:")
-        for i, (scenario_id, info) in enumerate(list(dataset_summary.items())[:10]):
-            if isinstance(info, dict):
-                track_length = info.get('track_length', 'N/A')
-                num_objects = info.get('number_summary', {}).get('num_objects', 'N/A')
-                print(f"  {i}: {scenario_id[:16]}... (时长: {track_length}, 对象数: {num_objects})")
-    
-    # 检查实际文件
-    print(f"\n检查实际文件...")
-    pkl_files = list(Path(full_path).rglob("*.pkl"))
-    # 排除dataset_mapping和dataset_summary
-    scenario_files = [f for f in pkl_files if f.name not in ["dataset_mapping.pkl", "dataset_summary.pkl"]]
-    print(f"   实际场景文件数量: {len(scenario_files)}")
-    
-    # 检查子目录
-    print(f"\n子目录结构:")
-    for item in os.listdir(full_path):
-        item_path = os.path.join(full_path, item)
-        if os.path.isdir(item_path):
-            pkl_count = len([f for f in os.listdir(item_path) if f.endswith('.pkl')])
-            print(f"  {item}/: {pkl_count} 个pkl文件")
-    
-    print("\n" + "=" * 80)
-    print("检查完成")
-    print("=" * 80)
-    
-    # 返回场景数量
-    return len(dataset_mapping) if 'dataset_mapping' in locals() else 0
-
-
-if __name__ == "__main__":
-    WAYMO_DATA_DIR = r"/home/huangfukk/mdsn"
-    
-    print("\n【检查 exp_filtered 数据集】")
-    num_filtered = check_dataset(WAYMO_DATA_DIR, "exp_filtered")
-    
-    print("\n\n【检查 exp_converted 数据集】")
-    num_converted = check_dataset(WAYMO_DATA_DIR, "exp_converted")
-    
-    print("\n\n总结:")
-    print(f"  exp_filtered: {num_filtered} 个场景")
-    print(f"  exp_converted: {num_converted} 个场景")
-

Env/expert_replay_env.py

@@ -0,0 +1,407 @@
+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 = {}
+        self.expert_tracks = {}
+        self.sdc_track = None
+        self.sdc_vehicle = None
+
+        # 在加载新场景前，必须清除上一轮通过 spawn_object 生成的物体，否则 engine.reset() 内 _object_clean_check 会报错
+        # 从 engine 当前对象中按名称筛选（与 manager 无关的对象需在此清理），并强制销毁
+        ids_to_clear = []
+        for oid, obj in self.engine.get_objects().items():
+            name = getattr(obj, "name", None) or getattr(obj, "id", None)
+            if name and (str(name).startswith("controlled_") or str(name).startswith("bg_")):
+                ids_to_clear.append(oid)
+        if ids_to_clear:
+            self.engine.clear_objects(ids_to_clear, force_destroy=True)
+        self.controlled_agents.clear()
+        self.controlled_agent_ids.clear()
+        for aid in list(self.engine.agent_manager.active_agents.keys()):
+            if aid.startswith("bg_") or aid.startswith("controlled_"):
+                self.engine.agent_manager.active_agents.pop(aid, 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)
+
+        from Env.utils import filter_traffic_tracks_to_birth_lists
+        traffic_data = self.engine.traffic_manager.current_traffic_data
+        car_birth_info_list, self.background_vehicles, obj_to_clean = filter_traffic_tracks_to_birth_lists(
+            traffic_data,
+            self.engine.traffic_manager.sdc_scenario_id,
+            self.engine.map_manager,
+        )
+        for entry in car_birth_info_list:
+            sid = entry["scenario_id"]
+            if sid in traffic_data:
+                self.expert_tracks[sid] = traffic_data[sid]
+        self.car_birth_info_list = car_birth_info_list
+        for scenario_id in obj_to_clean:
+            self.engine.traffic_manager.current_traffic_data.pop(scenario_id)
+
+        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 = []
+        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([v.id for v in 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

Env/inverse_dynamics.py

@@ -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}

Env/replay_policy.py

@@ -1,62 +0,0 @@
-import numpy as np
-
-class ReplayPolicy:
-    """
-    严格回放策略：根据专家轨迹数据，逐帧回放车辆状态
-    """
-
-    def __init__(self, expert_trajectory, vehicle_id):
-        """
-        Args:
-            expert_trajectory: 专家轨迹字典，包含 positions, headings, velocities, valid
-            vehicle_id: 车辆ID（用于调试）
-        """
-        self.trajectory = expert_trajectory
-        self.vehicle_id = vehicle_id
-        self.current_step = 0
-
-    def act(self, observation=None):
-        """
-        返回动作：在回放模式下返回空动作
-        实际状态由环境直接设置
-        """
-        return [0.0, 0.0]
-
-    def get_target_state(self, step):
-        """
-        获取指定时间步的目标状态
-
-        Args:
-            step: 时间步
-
-        Returns:
-            dict: 包含 position, heading, velocity 的字典，如果无效则返回 None
-        """
-        if step >= len(self.trajectory['valid']):
-            return None
-
-        if not self.trajectory['valid'][step]:
-            return None
-
-        return {
-            'position': self.trajectory['positions'][step],
-            'heading': self.trajectory['headings'][step],
-            'velocity': self.trajectory['velocities'][step]
-        }
-
-    def is_finished(self, step):
-        """
-        判断轨迹是否已经播放完毕
-
-        Args:
-            step: 当前时间步
-
-        Returns:
-            bool: 如果轨迹已播放完或当前步无效，返回 True
-        """
-        # 超出轨迹长度
-        if step >= len(self.trajectory['valid']):
-            return True
-
-        # 当前步及之后都无效
-        return not any(self.trajectory['valid'][step:])

Env/run_multiagent_env.py

@@ -1,390 +1,41 @@
-import argparse
 from scenario_env import MultiAgentScenarioEnv
-from simple_idm_policy import ConstantVelocityPolicy
+from Env.simple_idm_policy import ConstantVelocityPolicy
-from replay_policy import ReplayPolicy
 from metadrive.engine.asset_loader import AssetLoader
 
-WAYMO_DATA_DIR = r"/home/huangfukk/mdsn"
+WAYMO_DATA_DIR = r"/home/huangfukk/MAGAIL4AutoDrive/data"
 
-
+def main():
-def run_replay_mode(data_dir, num_episodes=1, horizon=300, render=True, debug=False,
-                    scenario_id=None, use_scenario_duration=False,
-                    spawn_vehicles=True, spawn_pedestrians=True, spawn_cyclists=True):
-    """
-    回放模式：严格按照专家轨迹回放
-
-    Args:
-        data_dir: 数据目录
-        num_episodes: 回合数（如果指定scenario_id，则忽略）
-        horizon: 最大步数（如果use_scenario_duration=True，则自动设置）
-        render: 是否渲染
-        debug: 是否调试模式
-        scenario_id: 指定场景ID（可选）
-        use_scenario_duration: 是否使用场景原始时长
-        spawn_vehicles: 是否生成车辆（默认True）
-        spawn_pedestrians: 是否生成行人（默认True）
-        spawn_cyclists: 是否生成自行车（默认True）
-    """
-    print("=" * 50)
-    print("运行模式: 专家轨迹回放 (Replay Mode)")
-    if scenario_id is not None:
-        print(f"指定场景ID: {scenario_id}")
-    if use_scenario_duration:
-        print("使用场景原始时长")
-    print("=" * 50)
-
-    # 如果指定了场景ID，只运行1个回合
-    if scenario_id is not None:
-        num_episodes = 1
-
-    # ✅ 环境创建移到循环外面，避免重复创建
     env = MultiAgentScenarioEnv(
         config={
-            "data_directory": AssetLoader.file_path(data_dir, "exp_filtered", unix_style=False),
+            # "data_directory": AssetLoader.file_path(AssetLoader.asset_path, "waymo", unix_style=False),
-            "is_multi_agent": True,
+            "data_directory": AssetLoader.file_path(WAYMO_DATA_DIR, "exp_converted", unix_style=False),
-            "horizon": horizon,
-            "use_render": render,  # 如果False会完全禁用渲染，避免LANE_FREEWAY错误
-            "sequential_seed": True,
-            "reactive_traffic": False,  # 回放模式下不需要反应式交通
-            "manual_control": False,
-            "filter_offroad_vehicles": True,  # 启用车道过滤
-            "lane_tolerance": 3.0,
-            "replay_mode": True,  # 标记为回放模式
-            "debug": debug,
-            "specific_scenario_id": scenario_id,  # 指定场景ID
-            "use_scenario_duration": use_scenario_duration,  # 使用场景时长
-            # 对象类型过滤
-            "spawn_vehicles": spawn_vehicles,
-            "spawn_pedestrians": spawn_pedestrians,
-            "spawn_cyclists": spawn_cyclists,
-            # ✅ 关键：设置可用场景数量
-            #"num_scenarios": 19012,  # 从dataset_mapping.pkl中统计的实际场景数
-        },
-        agent2policy=None  # 回放模式不需要统一策略
-    )
-
-    try:
-        # 获取可用场景数量
-        num_scenarios = env.config.get("num_scenarios", 1)
-        print(f"可用场景数量: {num_scenarios}")
-        
-        for episode in range(num_episodes):
-            print(f"\n{'='*50}")
-            print(f"回合 {episode + 1}/{num_episodes}")
-            if scenario_id is not None:
-                print(f"场景ID: {scenario_id}")
-            else:
-                # 循环使用场景
-                scenario_idx = episode % num_scenarios
-                print(f"使用场景索引: {scenario_idx}")
-            print(f"{'='*50}")
-
-            # ✅ 如果不是指定场景，使用循环的场景索引
-            if scenario_id is not None:
-                seed = scenario_id
-            else:
-                seed = episode % num_scenarios
-            obs = env.reset(seed=seed)
-
-            # 为每个车辆分配 ReplayPolicy
-            replay_policies = {}
-            for agent_id, vehicle in env.controlled_agents.items():
-                vehicle_id = vehicle.expert_vehicle_id
-                if vehicle_id in env.expert_trajectories:
-                    replay_policy = ReplayPolicy(
-                        env.expert_trajectories[vehicle_id],
-                        vehicle_id
-                    )
-                    vehicle.set_policy(replay_policy)
-                    replay_policies[agent_id] = replay_policy
-
-            # 输出场景信息
-            actual_horizon = env.config["horizon"]
-            print(f"初始化完成:")
-            print(f"  可控车辆数: {len(env.controlled_agents)}")
-            print(f"  专家轨迹数: {len(env.expert_trajectories)}")
-            print(f"  场景时长: {env.scenario_max_duration} 步")
-            print(f"  实际Horizon: {actual_horizon} 步")
-
-            step_count = 0
-            active_vehicles_count = []
-
-            while True:
-                # 在回放模式下，直接使用专家轨迹设置车辆状态
-                for agent_id, vehicle in list(env.controlled_agents.items()):
-                    vehicle_id = vehicle.expert_vehicle_id
-                    if vehicle_id in env.expert_trajectories and agent_id in replay_policies:
-                        target_state = replay_policies[agent_id].get_target_state(env.round)
-                        if target_state is not None:
-                            # 直接设置车辆状态（绕过物理引擎）
-                            # 只使用xy坐标，保持车辆在地面上
-                            position_2d = target_state['position'][:2]
-                            vehicle.set_position(position_2d)
-                            vehicle.set_heading_theta(target_state['heading'])
-                            vehicle.set_velocity(target_state['velocity'][:2] if len(target_state['velocity']) > 2 else target_state['velocity'])
-
-                # 使用空动作进行步进
-                actions = {aid: [0.0, 0.0] for aid in env.controlled_agents}
-                obs, rewards, dones, infos = env.step(actions)
-
-                if render:
-                    env.render(mode="topdown")
-
-                step_count += 1
-                active_vehicles_count.append(len(env.controlled_agents))
-
-                # 每50步打印一次状态
-                if step_count % 50 == 0:
-                    print(f"Step {step_count}: {len(env.controlled_agents)} 辆活跃车辆")
-                    
-                    # 调试模式下打印车辆高度信息
-                    if debug and len(env.controlled_agents) > 0:
-                        sample_vehicle = list(env.controlled_agents.values())[0]
-                        z_pos = sample_vehicle.position[2] if len(sample_vehicle.position) > 2 else 0
-                        print(f"  [DEBUG] 示例车辆高度: z={z_pos:.3f}m")
-
-                if dones["__all__"]:
-                    print(f"\n回合结束统计:")
-                    print(f"  总步数: {step_count}")
-                    print(f"  最大同时车辆数: {max(active_vehicles_count) if active_vehicles_count else 0}")
-                    print(f"  平均车辆数: {sum(active_vehicles_count) / len(active_vehicles_count) if active_vehicles_count else 0:.1f}")
-                    if use_scenario_duration:
-                        print(f"  场景完整回放: {'是' if step_count >= env.scenario_max_duration else '否'}")
-                    break
-    finally:
-        # ✅ 确保环境被正确关闭
-        env.close()
-
-    print("\n" + "=" * 50)
-    print("回放完成！")
-    print("=" * 50)
-
-
-def run_simulation_mode(data_dir, num_episodes=1, horizon=300, render=True, debug=False,
-                        scenario_id=None, use_scenario_duration=False,
-                        spawn_vehicles=True, spawn_pedestrians=True, spawn_cyclists=True):
-    """
-    仿真模式：使用自定义策略控制车辆
-    车辆根据专家数据的初始位姿生成，然后由策略控制
-
-    Args:
-        data_dir: 数据目录
-        num_episodes: 回合数
-        horizon: 最大步数
-        render: 是否渲染
-        debug: 是否调试模式
-        scenario_id: 指定场景ID（可选）
-        use_scenario_duration: 是否使用场景原始时长
-        spawn_vehicles: 是否生成车辆（默认True）
-        spawn_pedestrians: 是否生成行人（默认True）
-        spawn_cyclists: 是否生成自行车（默认True）
-    """
-    print("=" * 50)
-    print("运行模式: 策略仿真 (Simulation Mode)")
-    if scenario_id is not None:
-        print(f"指定场景ID: {scenario_id}")
-    if use_scenario_duration:
-        print("使用场景原始时长")
-    print("=" * 50)
-
-    # 如果指定了场景ID，只运行1个回合
-    if scenario_id is not None:
-        num_episodes = 1
-
-    env = MultiAgentScenarioEnv(
-        config={
-            "data_directory": AssetLoader.file_path(data_dir, "exp_filtered", unix_style=False),
             "is_multi_agent": True,
             "num_controlled_agents": 3,
-            "horizon": horizon,
+            "horizon": 300,
-            "use_render": render,  # 如果False会完全禁用渲染，避免LANE_FREEWAY错误
+            "use_render": True,
             "sequential_seed": True,
             "reactive_traffic": True,
-            "manual_control": False,
+            "manual_control": True,
-            "filter_offroad_vehicles": True,  # 启用车道过滤
-            "lane_tolerance": 3.0,
-            "replay_mode": False,  # 仿真模式
-            "debug": debug,
-            "specific_scenario_id": scenario_id,  # 指定场景ID
-            "use_scenario_duration": use_scenario_duration,  # 使用场景时长
-            # 对象类型过滤
-            "spawn_vehicles": spawn_vehicles,
-            "spawn_pedestrians": spawn_pedestrians,
-            "spawn_cyclists": spawn_cyclists,
-            # ✅ 关键：设置可用场景数量
-            #"num_scenarios": 19012,  # 从dataset_mapping.pkl中统计的实际场景数
         },
         agent2policy=ConstantVelocityPolicy(target_speed=50)
     )
 
-    try:
+    obs = env.reset(0
-        # 获取可用场景数量
+                    )
-        num_scenarios = env.config.get("num_scenarios", 1)
+    for step in range(10000):
-        print(f"可用场景数量: {num_scenarios}")
-        
-        for episode in range(num_episodes):
-            print(f"\n{'='*50}")
-            print(f"回合 {episode + 1}/{num_episodes}")
-            if scenario_id is not None:
-                print(f"场景ID: {scenario_id}")
-            else:
-                # 循环使用场景
-                scenario_idx = episode % num_scenarios
-                print(f"使用场景索引: {scenario_idx}")
-            print(f"{'='*50}")
-
-            # ✅ 如果不是指定场景，使用循环的场景索引
-            if scenario_id is not None:
-                seed = scenario_id
-            else:
-                seed = episode % num_scenarios
-            obs = env.reset(seed=seed)
-
-            actual_horizon = env.config["horizon"]
-            print(f"初始化完成:")
-            print(f"  可控车辆数: {len(env.controlled_agents)}")
-            print(f"  场景时长: {env.scenario_max_duration} 步")
-            print(f"  实际Horizon: {actual_horizon} 步")
-
-            step_count = 0
-            total_reward = 0.0
-
-            while True:
-                # 使用策略生成动作
         actions = {
             aid: env.controlled_agents[aid].policy.act()
             for aid in env.controlled_agents
         }
 
         obs, rewards, dones, infos = env.step(actions)
-
-                if render:
         env.render(mode="topdown")
 
-                step_count += 1
-                total_reward += sum(rewards.values())
-
-                # 每50步打印一次状态
-                if step_count % 50 == 0:
-                    print(f"Step {step_count}: {len(env.controlled_agents)} 辆活跃车辆")
-
         if dones["__all__"]:
-                    print(f"\n回合结束统计:")
-                    print(f"  总步数: {step_count}")
-                    print(f"  总奖励: {total_reward:.2f}")
             break
-    finally:
+
     env.close()
 
-    print("\n" + "=" * 50)
-    print("仿真完成！")
-    print("=" * 50)
-
-
-def main():
-    parser = argparse.ArgumentParser(description="MetaDrive 多智能体环境运行脚本")
-
-    parser.add_argument(
-        "--mode",
-        type=str,
-        choices=["replay", "simulation"],
-        default="simulation",
-        help="运行模式: replay=专家轨迹回放, simulation=策略仿真 (默认: simulation)"
-    )
-
-    parser.add_argument(
-        "--data_dir",
-        type=str,
-        default=WAYMO_DATA_DIR,
-        help=f"数据目录路径 (默认: {WAYMO_DATA_DIR})"
-    )
-
-    parser.add_argument(
-        "--episodes",
-        type=int,
-        default=1,
-        help="运行回合数 (默认: 1)"
-    )
-
-    parser.add_argument(
-        "--horizon",
-        type=int,
-        default=300,
-        help="每回合最大步数 (默认: 300，如果启用 --use_scenario_duration 则自动设置)"
-    )
-
-    parser.add_argument(
-        "--no_render",
-        action="store_true",
-        help="禁用渲染（加速运行）"
-    )
-
-    parser.add_argument(
-        "--debug",
-        action="store_true",
-        help="启用调试模式（显示详细日志）"
-    )
-
-    parser.add_argument(
-        "--scenario_id",
-        type=int,
-        default=None,
-        help="指定场景ID（可选，如指定则只运行该场景）"
-    )
-
-    parser.add_argument(
-        "--use_scenario_duration",
-        action="store_true",
-        help="使用场景原始时长作为horizon（自动停止）"
-    )
-
-    parser.add_argument(
-        "--no_vehicles",
-        action="store_true",
-        help="禁止生成车辆"
-    )
-
-    parser.add_argument(
-        "--no_pedestrians",
-        action="store_true",
-        help="禁止生成行人"
-    )
-
-    parser.add_argument(
-        "--no_cyclists",
-        action="store_true",
-        help="禁止生成自行车"
-    )
-
-    args = parser.parse_args()
-
-    if args.mode == "replay":
-        run_replay_mode(
-            data_dir=args.data_dir,
-            num_episodes=args.episodes,
-            horizon=args.horizon,
-            render=not args.no_render,
-            debug=args.debug,
-            scenario_id=args.scenario_id,
-            use_scenario_duration=args.use_scenario_duration,
-            spawn_vehicles=not args.no_vehicles,
-            spawn_pedestrians=not args.no_pedestrians,
-            spawn_cyclists=not args.no_cyclists
-        )
-    else:
-        run_simulation_mode(
-            data_dir=args.data_dir,
-            num_episodes=args.episodes,
-            horizon=args.horizon,
-            render=not args.no_render,
-            debug=args.debug,
-            scenario_id=args.scenario_id,
-            use_scenario_duration=args.use_scenario_duration,
-            spawn_vehicles=not args.no_vehicles,
-            spawn_pedestrians=not args.no_pedestrians,
-            spawn_cyclists=not args.no_cyclists
-        )
-
 
 if __name__ == "__main__":
     main()

Env/scenario_env.py

@@ -15,7 +15,6 @@ class PolicyVehicle(DefaultVehicle):
         super().__init__(*args, **kwargs)
         self.policy = None
         self.destination = None
-        self.expert_vehicle_id = None  # 关联专家车辆ID
 
     def set_policy(self, policy):
         self.policy = policy
@@ -23,9 +22,6 @@ class PolicyVehicle(DefaultVehicle):
     def set_destination(self, des):
         self.destination = des
 
-    def set_expert_vehicle_id(self, vid):
-        self.expert_vehicle_id = vid
-
     def act(self, observation, policy=None):
         if self.policy is not None:
             return self.policy.act(observation)
@@ -57,15 +53,6 @@ class MultiAgentScenarioEnv(ScenarioEnv):
             data_directory=None,
             num_controlled_agents=3,
             horizon=1000,
-            filter_offroad_vehicles=True,  # 车道过滤开关
-            lane_tolerance=3.0,  # 车道检测容差（米）
-            replay_mode=False,  # 回放模式开关
-            specific_scenario_id=None,  # 新增：指定场景ID（仅回放模式）
-            use_scenario_duration=False,  # 新增：使用场景原始时长作为horizon
-            # 对象类型过滤选项
-            spawn_vehicles=True,  # 是否生成车辆
-            spawn_pedestrians=True,  # 是否生成行人
-            spawn_cyclists=True,  # 是否生成自行车
         ))
         return config
 
@@ -75,180 +62,38 @@ class MultiAgentScenarioEnv(ScenarioEnv):
         self.controlled_agent_ids = []
         self.obs_list = []
         self.round = 0
-        self.expert_trajectories = {}  # 存储完整专家轨迹
-        self.replay_mode = config.get("replay_mode", False)
-        self.scenario_max_duration = 0  # 场景实际最大时长
         super().__init__(config)
 
     def reset(self, seed: Union[None, int] = None):
         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)
 
-        # ✅ 关键修复：在每次 reset 前清理所有自定义生成的对象
-        if hasattr(self, 'engine') and self.engine is not None:
-            if hasattr(self, 'controlled_agents') and self.controlled_agents:
-                # 先从 agent_manager 中移除
-                if hasattr(self.engine, 'agent_manager'):
-                    for agent_id in list(self.controlled_agents.keys()):
-                        if agent_id in self.engine.agent_manager.active_agents:
-                            self.engine.agent_manager.active_agents.pop(agent_id)
-                
-                # 然后清理对象
-                for agent_id, vehicle in list(self.controlled_agents.items()):
-                    try:
-                        self.engine.clear_objects([vehicle.id])
-                    except:
-                        pass
-                
-                self.controlled_agents.clear()
-                self.controlled_agent_ids.clear()
-        
         self.lazy_init()
         self._reset_global_seed(seed)
-
         if self.engine is None:
             raise ValueError("Broken MetaDrive instance.")
 
-        # 如果指定了场景ID，修改start_scenario_index
+        self.background_vehicles = getattr(self, "background_vehicles", {})
-        if self.config.get("specific_scenario_id") is not None:
+        self.car_birth_info_list, self.background_vehicles, _obj_to_clean = self._build_birth_lists_from_traffic()
-            scenario_id = self.config.get("specific_scenario_id")
+        for scenario_id in _obj_to_clean:
-            self.config["start_scenario_index"] = scenario_id
+            self.engine.traffic_manager.current_traffic_data.pop(scenario_id)
-            if self.config.get("debug", False):
+
-                self.logger.info(f"Using specific scenario ID: {scenario_id}")
+        # Clear vehicles we spawned via engine.spawn_object() so _object_clean_check() passes
+        ids_to_clear = [v.id for v in self.controlled_agents.values()]
+        if ids_to_clear:
+            self.engine.clear_objects(ids_to_clear)
+        self.controlled_agents.clear()
+        self.controlled_agent_ids.clear()
 
-        # ✅ 先初始化引擎和 lanes
         self.engine.reset()
         self.reset_sensors()
         self.engine.taskMgr.step()
+
         self.lanes = self.engine.map_manager.current_map.road_network.graph
 
-        # 记录专家数据（现在 self.lanes 已经初始化）
-        _obj_to_clean_this_frame = []
-        self.car_birth_info_list = []
-        self.expert_trajectories.clear()
-        total_vehicles = 0
-        total_pedestrians = 0
-        total_cyclists = 0
-        filtered_vehicles = 0
-        filtered_by_type = 0
-        self.scenario_max_duration = 0  # 重置场景时长
-
-        for scenario_id, track in self.engine.traffic_manager.current_traffic_data.items():
-            if scenario_id == self.engine.traffic_manager.sdc_scenario_id:
-                continue
-
-            # 对象类型过滤
-            obj_type = track["type"]
-            
-            # 统计对象类型
-            if obj_type == MetaDriveType.VEHICLE:
-                total_vehicles += 1
-            elif obj_type == MetaDriveType.PEDESTRIAN:
-                total_pedestrians += 1
-            elif obj_type == MetaDriveType.CYCLIST:
-                total_cyclists += 1
-            
-            # 根据配置过滤对象类型
-            if obj_type == MetaDriveType.VEHICLE and not self.config.get("spawn_vehicles", True):
-                _obj_to_clean_this_frame.append(scenario_id)
-                filtered_by_type += 1
-                if self.config.get("debug", False):
-                    self.logger.debug(f"Filtering VEHICLE {track['metadata']['object_id']} - spawn_vehicles=False")
-                continue
-            
-            if obj_type == MetaDriveType.PEDESTRIAN and not self.config.get("spawn_pedestrians", True):
-                _obj_to_clean_this_frame.append(scenario_id)
-                filtered_by_type += 1
-                if self.config.get("debug", False):
-                    self.logger.debug(f"Filtering PEDESTRIAN {track['metadata']['object_id']} - spawn_pedestrians=False")
-                continue
-            
-            if obj_type == MetaDriveType.CYCLIST and not self.config.get("spawn_cyclists", True):
-                _obj_to_clean_this_frame.append(scenario_id)
-                filtered_by_type += 1
-                if self.config.get("debug", False):
-                    self.logger.debug(f"Filtering CYCLIST {track['metadata']['object_id']} - spawn_cyclists=False")
-                continue
-
-            # 只处理车辆类型（行人和自行车暂时只做过滤）
-            if track["type"] == MetaDriveType.VEHICLE:
-                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
-
-                # 更新场景最大时长
-                self.scenario_max_duration = max(self.scenario_max_duration, last_show + 1)
-
-                # 获取车辆初始位置
-                initial_position = (
-                    track['state']['position'][first_show, 0],
-                    track['state']['position'][first_show, 1]
-                )
-
-                # 车道过滤
-                if self.config.get("filter_offroad_vehicles", True):
-                    if not self._is_position_on_lane(initial_position):
-                        filtered_vehicles += 1
-                        _obj_to_clean_this_frame.append(scenario_id)
-                        if self.config.get("debug", False):
-                            self.logger.debug(
-                                f"Filtering vehicle {track['metadata']['object_id']} - "
-                                f"not on lane at position {initial_position}"
-                            )
-                        continue
-
-                # 存储完整专家轨迹（只使用2D位置，避免高度问题）
-                object_id = track['metadata']['object_id']
-                positions_2d = track['state']['position'].copy()
-                positions_2d[:, 2] = 0  # 将z坐标设为0，让MetaDrive自动处理高度
-                
-                self.expert_trajectories[object_id] = {
-                    'positions': positions_2d,
-                    'headings': track['state']['heading'].copy(),
-                    'velocities': track['state']['velocity'].copy(),
-                    'valid': track['state']['valid'].copy(),
-                }
-
-                # 保存车辆生成信息
-                self.car_birth_info_list.append({
-                    'id': object_id,
-                    'show_time': first_show,
-                    'begin': initial_position,
-                    'heading': track['state']['heading'][first_show],
-                    'velocity': track['state']['velocity'][first_show] if self.config.get("inherit_expert_velocity", False) else None,
-                    'end': (track['state']['position'][last_show, 0], track['state']['position'][last_show, 1])
-                })
-
-                # 在回放和仿真模式下都清除原始专家车辆
-                _obj_to_clean_this_frame.append(scenario_id)
-
-        # 清除专家车辆和过滤的对象
-        for scenario_id in _obj_to_clean_this_frame:
-            self.engine.traffic_manager.current_traffic_data.pop(scenario_id)
-
-        # 输出统计信息
-        if self.config.get("debug", False):
-            self.logger.info(f"=== 对象统计 ===")
-            self.logger.info(f"车辆 (VEHICLE): 总数={total_vehicles}, 车道过滤={filtered_vehicles}, 保留={total_vehicles - filtered_vehicles}")
-            self.logger.info(f"行人 (PEDESTRIAN): 总数={total_pedestrians}")
-            self.logger.info(f"自行车 (CYCLIST): 总数={total_cyclists}")
-            self.logger.info(f"类型过滤: {filtered_by_type} 个对象")
-            self.logger.info(f"场景时长: {self.scenario_max_duration} 步")
-
-        # 如果启用场景时长控制，更新horizon
-        if self.config.get("use_scenario_duration", False) and self.scenario_max_duration > 0:
-            original_horizon = self.config["horizon"]
-            self.config["horizon"] = self.scenario_max_duration
-            if self.config.get("debug", False):
-                self.logger.info(f"Horizon updated from {original_horizon} to {self.scenario_max_duration} (scenario duration)")
-
         if self.top_down_renderer is not None:
             self.top_down_renderer.clear()
             self.engine.top_down_renderer = None
@@ -256,129 +101,64 @@ class MultiAgentScenarioEnv(ScenarioEnv):
         self.dones = {}
         self.episode_rewards = defaultdict(float)
         self.episode_lengths = defaultdict(int)
-        self.controlled_agents.clear()
-        self.controlled_agent_ids.clear()
 
         super().reset(seed)  # 初始化场景
         self._spawn_controlled_agents()
 
         return self._get_all_obs()
 
-    def _is_position_on_lane(self, position, tolerance=None):
+    def _build_birth_lists_from_traffic(self):
-        if tolerance is None:
+        """Build car_birth_info_list and obj_to_clean from current_traffic_data. Override for filtered (lane/static) selection."""
-            tolerance = self.config.get("lane_tolerance", 3.0)
+        _obj_to_clean_this_frame = []
-
+        car_birth_info_list = []
-        # 确保 self.lanes 已初始化
+        for scenario_id, track in self.engine.traffic_manager.current_traffic_data.items():
-        if not hasattr(self, 'lanes') or self.lanes is None:
+            if scenario_id == self.engine.traffic_manager.sdc_scenario_id:
-            if self.config.get("debug", False):
-                self.logger.warning("Lanes not initialized, skipping lane check")
-            return True
-
-        position_2d = np.array(position[:2]) if len(position) > 2 else np.array(position)
-
-        try:
-            for lane in self.lanes.values():
-                if lane.lane.point_on_lane(position_2d):
-                    return True
-
-                lane_start = np.array(lane.lane.start)[:2]
-                lane_end = np.array(lane.lane.end)[:2]
-                lane_vec = lane_end - lane_start
-                lane_length = np.linalg.norm(lane_vec)
-
-                if lane_length < 1e-6:
                 continue
-
+            if track["type"] == MetaDriveType.VEHICLE:
-                lane_vec_normalized = lane_vec / lane_length
+                _obj_to_clean_this_frame.append(scenario_id)
-                point_vec = position_2d - lane_start
+                valid = track["state"]["valid"]
-                projection = np.dot(point_vec, lane_vec_normalized)
+                first_show = int(np.argmax(valid)) if valid.any() else -1
-
+                last_show = len(valid) - 1 - int(np.argmax(valid[::-1])) if valid.any() else -1
-                if 0 <= projection <= lane_length:
+                car_birth_info_list.append({
-                    closest_point = lane_start + projection * lane_vec_normalized
+                    "id": track["metadata"]["object_id"],
-                    distance = np.linalg.norm(position_2d - closest_point)
+                    "show_time": first_show,
-                    if distance <= tolerance:
+                    "begin": (track["state"]["position"][first_show, 0], track["state"]["position"][first_show, 1]),
-                        return True
+                    "heading": track["state"]["heading"][first_show],
-        except Exception as e:
+                    "end": (track["state"]["position"][last_show, 0], track["state"]["position"][last_show, 1]),
-            if self.config.get("debug", False):
+                })
-                self.logger.warning(f"Lane check error: {e}")
+        return car_birth_info_list, {}, _obj_to_clean_this_frame
-            return False
-
-        return False
 
     def _spawn_controlled_agents(self):
-        """
+        # ego_vehicle = self.engine.agent_manager.active_agents.get("default_agent")
-        生成应该在当前或之前出现的车辆
+        # ego_position = ego_vehicle.position if ego_vehicle else np.array([0, 0])
-        如果round=0且所有车辆的show_time>0，则生成show_time最小的车辆（保证至少有车辆出现）
-        """
-        vehicles_to_spawn = []
-        
         for car in self.car_birth_info_list:
-            if car['show_time'] <= self.round:
+            if car['show_time'] == self.round:
-                vehicles_to_spawn.append(car)
-        
-        # 如果当前round没有车辆应该出现，但车辆列表不为空，则生成最早出现的车辆
-        # 这样可以确保在reset时至少有车辆出现
-        # if len(vehicles_to_spawn) == 0 and len(self.car_birth_info_list) > 0:
-        #     if self.config.get("debug", False):
-        #         self.logger.debug(
-        #             f"No vehicles to spawn at round {self.round}, "
-        #             f"spawning earliest vehicle instead"
-        #         )
-        #     # 找到show_time最小的车辆
-        #     earliest_car = min(self.car_birth_info_list, key=lambda x: x['show_time'])
-        #     vehicles_to_spawn.append(earliest_car)
-        
-        for car in vehicles_to_spawn:
                 agent_id = f"controlled_{car['id']}"
 
-            # 避免重复生成
-            if agent_id in self.controlled_agents:
-                continue
-                
-            vehicle_config = {}
                 vehicle = self.engine.spawn_object(
                     PolicyVehicle,
-                vehicle_config=vehicle_config,
+                    vehicle_config={},
                     position=car['begin'],
                     heading=car['heading']
                 )
+                vehicle.reset(position=car['begin'], heading=car['heading'])
 
-            # 重置车辆状态
-            reset_kwargs = {
-                'position': car['begin'],
-                'heading': car['heading']
-            }
-
-            # 如果启用速度继承，设置初始速度
-            if car.get('velocity') is not None:
-                reset_kwargs['velocity'] = car['velocity']
-
-            vehicle.reset(**reset_kwargs)
-
-            # 设置策略和目的地
                 vehicle.set_policy(self.policy)
                 vehicle.set_destination(car['end'])
-            vehicle.set_expert_vehicle_id(car['id'])
 
                 self.controlled_agents[agent_id] = vehicle
                 self.controlled_agent_ids.append(agent_id)
 
-            # 注册到引擎的 active_agents
+                # ✅ 关键：注册到引擎的 active_agents，才能参与物理更新
                 self.engine.agent_manager.active_agents[agent_id] = vehicle
 
-            if self.config.get("debug", False):
-                self.logger.debug(
-                    f"Spawned vehicle {agent_id} at round {self.round} "
-                    f"(show_time={car['show_time']}), position {car['begin']}"
-                )
-
     def _get_all_obs(self):
+        # position, velocity, heading, lidar, navigation, TODO: trafficlight -> list
         self.obs_list = []
-
         for agent_id, vehicle in self.controlled_agents.items():
             state = vehicle.get_state()
-            traffic_light = 0
 
+            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):
@@ -393,20 +173,8 @@ class MultiAgentScenarioEnv(ScenarioEnv):
                             traffic_light = 0
                     break
 
-            # 使用最近10辆车的相对位置与相对速度替代原80维LiDAR点云
+            lidar = self.engine.get_sensor("lidar").perceive(num_lasers=80, distance=30, base_vehicle=vehicle,
-            lidar_cloud_points, detected_objects = self.engine.get_sensor("lidar").perceive(
+                                                             physics_world=self.engine.physics_world.dynamic_world)
-                num_lasers=80,
-                distance=30,
-                base_vehicle=vehicle,
-                physics_world=self.engine.physics_world.dynamic_world
-            )
-            nearest_vehicle_info = self.engine.get_sensor("lidar").get_surrounding_vehicles_info(
-                vehicle,
-                detected_objects,
-                perceive_distance=30,
-                num_others=10,
-                add_others_navi=False
-            )
             side_lidar = self.engine.get_sensor("side_detector").perceive(num_lasers=10, distance=8,
                                                                           base_vehicle=vehicle,
                                                                           physics_world=self.engine.physics_world.static_world)
@@ -414,61 +182,30 @@ class MultiAgentScenarioEnv(ScenarioEnv):
                                                                                     base_vehicle=vehicle,
                                                                                     physics_world=self.engine.physics_world.static_world)
 
-            obs = (list(state['position'][:2]) + list(state['velocity']) + [state['heading_theta']]
+            obs = (state['position'][:2] + list(state['velocity']) + [state['heading_theta']]
-                   + nearest_vehicle_info + side_lidar[0] + lane_line_lidar[0] + [traffic_light]
+                   + 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]]):
         self.round += 1
 
-        # 应用动作
         for agent_id, action in action_dict.items():
             if agent_id in self.controlled_agents:
                 self.controlled_agents[agent_id].before_step(action)
 
-        # 物理引擎步进
         self.engine.step()
+        self.engine.after_step()
 
-        # 后处理
         for agent_id in action_dict:
             if agent_id in self.controlled_agents:
                 self.controlled_agents[agent_id].after_step()
 
-        # 生成新车辆
         self._spawn_controlled_agents()
-
-        # 获取观测
         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.episode_step >= self.config["horizon"]
-        # ✅ 修复：添加回放模式的完成检查
-        replay_finished = False
-        if self.replay_mode and self.config.get("use_scenario_duration", False):
-            # 检查是否所有专家轨迹都已播放完毕
-            if self.round >= self.scenario_max_duration:
-                replay_finished = True
-                if self.config.get("debug", False):
-                    self.logger.info(f"Replay finished at step {self.round}/{self.scenario_max_duration}")
-        
-        dones["__all__"] = self.episode_step >= self.config["horizon"] or replay_finished
-
         infos = {aid: {} for aid in self.controlled_agents}
-
         return obs, rewards, dones, infos
-
-    def close(self):
-        # ✅ 清理所有生成的车辆
-        if hasattr(self, 'controlled_agents') and self.controlled_agents:
-            for agent_id, vehicle in list(self.controlled_agents.items()):
-                if vehicle in self.engine.get_objects():
-                    self.engine.clear_objects([vehicle.id])
-            self.controlled_agents.clear()
-            self.controlled_agent_ids.clear()
-        
-        super().close()

Env/utils.py

@@ -2,6 +2,143 @@ import numpy as np
 import torch
 import random
 
+from metadrive.type import MetaDriveType
+
+
+def is_on_lane(pos, map_manager, threshold=2.0):
+    """Check if a position is on a valid lane (within lateral tolerance)."""
+    if map_manager is None or map_manager.current_map is None:
+        return True
+    try:
+        lane, _ = map_manager.current_map.road_network.get_closest_lane_index(pos, return_lane=True)
+        if lane is None:
+            return False
+        long, lat = lane.local_coordinates(pos)
+        width = lane.width
+        if abs(lat) <= (width / 2 + threshold):
+            return True
+        return False
+    except Exception:
+        return False
+
+
+def filter_traffic_tracks_to_birth_lists(
+    current_traffic_data,
+    sdc_scenario_id,
+    map_manager,
+    *,
+    lane_threshold=5.0,
+    static_displacement_threshold=5.0,
+    static_speed_threshold=1.0,
+    return_stats=False,
+):
+    """
+    Filter traffic tracks into controlled (car_birth_info_list) and background lists.
+
+    - controlled (car_birth_info_list): 非 SDC、类型 VEHICLE、至少一帧 valid、在车道内、且非静态
+      （位移/速度超过阈值）。用于策略控制或专家回放，spawn 时机为 show_time == round。
+    - background (background_vehicles): 同上但在车道内且判定为静态（位移 < 5m、速度 < 1 m/s）。
+      仅作场景占位与观测邻居，spawn 时机为 show_time == round，按 valid 在 step 中移除。
+
+    Returns (car_birth_info_list, background_vehicles, obj_to_clean) or, if return_stats=True,
+    (car_birth_info_list, background_vehicles, obj_to_clean, stats_dict).
+    stats_dict: n_total, n_no_valid, n_off_lane, n_static, n_controlled.
+    """
+    car_birth_info_list = []
+    background_vehicles = {}
+    obj_to_clean = []
+    n_total = 0
+    n_no_valid = 0
+    n_off_lane = 0
+    n_static = 0
+
+    for scenario_id, track in current_traffic_data.items():
+        if scenario_id == sdc_scenario_id:
+            continue
+        if track["type"] != MetaDriveType.VEHICLE:
+            continue
+
+        n_total += 1
+        obj_to_clean.append(scenario_id)
+        valid = track["state"]["valid"]
+        if not valid.any():
+            n_no_valid += 1
+            continue
+
+        first_show = int(np.argmax(valid))
+        last_show = len(valid) - 1 - int(np.argmax(valid[::-1]))
+        mid_show = (first_show + last_show) // 2
+
+        start_pos = track["state"]["position"][first_show]
+        is_valid_track = True
+        if not is_on_lane(start_pos, map_manager, threshold=lane_threshold):
+            mid_pos = track["state"]["position"][mid_show]
+            if not is_on_lane(mid_pos, map_manager, threshold=lane_threshold):
+                is_valid_track = False
+
+        if not is_valid_track:
+            n_off_lane += 1
+            continue
+
+        positions = track["state"]["position"][valid.astype(bool)]
+        velocities = track["state"]["velocity"][valid.astype(bool)]
+        total_displacement = 0.0
+        max_speed = 0.0
+        if len(positions) > 1:
+            total_displacement = float(np.linalg.norm(positions[-1] - positions[0]))
+            max_speed = float(np.max(np.linalg.norm(velocities, axis=1)))
+        is_static = total_displacement < static_displacement_threshold and max_speed < static_speed_threshold
+
+        if is_static:
+            n_static += 1
+            background_vehicles[scenario_id] = {
+                "id": track["metadata"]["object_id"],
+                "show_time": first_show,
+                "begin": (
+                    float(track["state"]["position"][first_show, 0]),
+                    float(track["state"]["position"][first_show, 1]),
+                ),
+                "heading": float(track["state"]["heading"][first_show]),
+                "end": (
+                    float(track["state"]["position"][last_show, 0]),
+                    float(track["state"]["position"][last_show, 1]),
+                ),
+                "scenario_id": scenario_id,
+                "length": track["state"]["length"][first_show],
+                "width": track["state"]["width"][first_show],
+                "valid": valid,
+            }
+            continue
+
+        car_birth_info_list.append({
+            "id": track["metadata"]["object_id"],
+            "show_time": first_show,
+            "begin": (
+                float(track["state"]["position"][first_show, 0]),
+                float(track["state"]["position"][first_show, 1]),
+            ),
+            "heading": float(track["state"]["heading"][first_show]),
+            "end": (
+                float(track["state"]["position"][last_show, 0]),
+                float(track["state"]["position"][last_show, 1]),
+            ),
+            "scenario_id": scenario_id,
+            "length": track["state"]["length"][first_show],
+            "width": track["state"]["width"][first_show],
+        })
+
+    if return_stats:
+        stats = {
+            "n_total": n_total,
+            "n_no_valid": n_no_valid,
+            "n_off_lane": n_off_lane,
+            "n_static": n_static,
+            "n_controlled": len(car_birth_info_list),
+        }
+        return car_birth_info_list, background_vehicles, obj_to_clean, stats
+    return car_birth_info_list, background_vehicles, obj_to_clean
+
+
 def set_seed(seed):
     if seed == -1:
         seed = np.random.randint(0, 10000)

Env/verify_observation.py

@@ -1,184 +0,0 @@
-#!/usr/bin/env python
-# -*- coding: utf-8 -*-
-"""
-简单验证脚本：检查车辆是否正确获取观测空间
-用法：python verify_observations.py
-"""
-
-import argparse
-from scenario_env import MultiAgentScenarioEnv
-from replay_policy import ReplayPolicy
-from metadrive.engine.asset_loader import AssetLoader
-import numpy as np
-
-def verify_observations(data_dir, scenario_id=0):
-    """
-    验证观测空间是否正确获取
-
-    Args:
-        data_dir: 数据目录
-        scenario_id: 场景ID
-    """
-    print("=" * 60)
-    print("观测空间验证工具")
-    print("=" * 60)
-
-    # 创建环境
-    env = MultiAgentScenarioEnv(
-        config={
-            "data_directory": AssetLoader.file_path(data_dir, "exp_filtered", unix_style=False),
-            "is_multi_agent": True,
-            "horizon": 300,
-            "use_render": False,  # 不渲染，加速运行
-            "sequential_seed": True,
-            "reactive_traffic": False,
-            "manual_control": False,
-            "filter_offroad_vehicles": True,
-            "lane_tolerance": 3.0,
-            "replay_mode": True,
-            "debug": True,  # 启用调试以查看详细信息
-            "specific_scenario_id": scenario_id,
-            "use_scenario_duration": True,
-        },
-        agent2policy=None
-    )
-
-    # 重置环境
-    print(f"\n加载场景 {scenario_id}...")
-    obs = env.reset(seed=scenario_id)
-
-    # 输出基本信息
-    print(f"\n场景信息:")
-    print(f"  - 可控车辆数: {len(env.controlled_agents)}")
-    print(f"  - 观测数量: {len(obs)}")
-    print(f"  - 场景时长: {env.scenario_max_duration} 步")
-    print(f"  - 车辆生成列表长度: {len(env.car_birth_info_list)}")
-    print(f"  - 当前回合数 (round): {env.round}")
-    
-    # 检查车辆生成信息
-    if len(env.car_birth_info_list) > 0:
-        print(f"\n车辆生成信息分析:")
-        show_times = [car['show_time'] for car in env.car_birth_info_list]
-        print(f"  - show_time 分布: min={min(show_times)}, max={max(show_times)}")
-        print(f"  - show_time == 0 的车辆数: {sum(1 for st in show_times if st == 0)}")
-        print(f"  - 前5个车辆的 show_time: {show_times[:5]}")
-    else:
-        print(f"\n⚠️  警告: 车辆生成列表为空！可能原因:")
-        print(f"    1. 所有车辆都被车道过滤移除")
-        print(f"    2. 所有车辆都被类型过滤移除")
-        print(f"    3. 场景数据中没有有效车辆")
-
-    # 验证观测空间
-    print(f"\n" + "=" * 60)
-    print("观测空间验证")
-    print("=" * 60)
-
-    if len(obs) == 0:
-        print("❌ 错误：没有获取到任何观测！")
-        env.close()
-        return False
-
-    # 检查第一个观测
-    first_obs = obs[0]
-    print(f"\n第一个车辆的观测:")
-    print(f"  - 观测类型: {type(first_obs)}")
-    print(f"  - 观测维度: {len(first_obs)}")
-
-    # 详细解析观测
-    if isinstance(first_obs, (list, np.ndarray)):
-        obs_array = np.array(first_obs)
-        print(f"  - 观测形状: {obs_array.shape}")
-        print(f"  - 数据类型: {obs_array.dtype}")
-        print(f"\n观测内容分解:")
-        # 新观测划分（与 Env/scenario_env._get_all_obs 对齐）：
-        # [x, y] (2) + [vx, vy] (2) + heading (1)
-        # + nearest vehicles info (10 vehicles * 4 = 40)
-        # + side_lidar (10) + lane_line_lidar (10)
-        # + traffic_light (1) + destination (2)
-        lidar_len = 40
-        side_len = 10
-        lane_len = 10
-        pos_slice = slice(0, 2)
-        vel_slice = slice(2, 4)
-        heading_idx = 4
-        lidar_slice = slice(5, 5 + lidar_len)
-        side_slice = slice(lidar_slice.stop, lidar_slice.stop + side_len)
-        lane_slice = slice(side_slice.stop, side_slice.stop + lane_len)
-        tl_idx = lane_slice.stop
-        dest_slice = slice(tl_idx + 1, tl_idx + 3)
-
-        print(f"    位置 (x, y): {obs_array[pos_slice]}")
-        print(f"    速度 (vx, vy): {obs_array[vel_slice]}")
-        print(f"    航向角: {obs_array[heading_idx]:.3f} 弧度")
-        print(f"    最近车辆信息: {len(obs_array[lidar_slice])} 维 (10辆*4: 相对x/相对y/相对vx/相对vy)")
-        print(f"    侧向检测: {len(obs_array[side_slice])} 个点")
-        print(f"    车道线检测: {len(obs_array[lane_slice])} 个点")
-        print(f"    交通灯状态: {obs_array[tl_idx]}")
-        print(f"    目的地 (x, y): {obs_array[dest_slice]}")
-
-        # 检查数据有效性
-        print(f"\n数据有效性检查:")
-        has_nan = np.isnan(obs_array).any()
-        has_inf = np.isinf(obs_array).any()
-
-        if has_nan:
-            print(f"  ❌ 观测包含 NaN 值！")
-        else:
-            print(f"  ✅ 无 NaN 值")
-
-        if has_inf:
-            print(f"  ❌ 观测包含 Inf 值！")
-        else:
-            print(f"  ✅ 无 Inf 值")
-
-        # 检查最近车辆信息数据（40维）
-        lidar_data = obs_array[lidar_slice]
-        lidar_min = np.min(lidar_data)
-        lidar_max = np.max(lidar_data)
-        print(f"\n  最近车辆特征范围: [{lidar_min:.2f}, {lidar_max:.2f}]")
-
-        if lidar_max > 0:
-            print(f"  ✅ 存在非零最近车辆特征")
-        else:
-            print(f"  ⚠️  最近车辆特征全零（可能无邻车或距离过远）")
-
-    # 运行几步，验证观测持续有效
-    print(f"\n" + "=" * 60)
-    print("多步运行验证（前 5 步）")
-    print("=" * 60)
-
-    for step in range(5):
-        # 空动作
-        actions = {aid: [0.0, 0.0] for aid in env.controlled_agents}
-        obs, rewards, dones, infos = env.step(actions)
-
-        print(f"\nStep {step + 1}:")
-        print(f"  - 活跃车辆数: {len(env.controlled_agents)}")
-        print(f"  - 观测数量: {len(obs)}")
-
-        if len(obs) > 0:
-            sample_obs = np.array(obs[0])
-            print(f"  - 第一辆车位置: ({sample_obs[0]:.2f}, {sample_obs[1]:.2f})")
-            print(f"  - 数据有效: {'✅' if not (np.isnan(sample_obs).any() or np.isinf(sample_obs).any()) else '❌'}")
-
-        if dones["__all__"]:
-            print(f"  - 场景结束")
-            break
-
-    env.close()
-
-    print(f"\n" + "=" * 60)
-    print("✅ 验证完成！观测空间正常工作")
-    print("=" * 60)
-
-    return True
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="验证观测空间")
-    parser.add_argument("--data_dir", type=str, default="/home/huangfukk/mdsn", help="数据目录")
-    parser.add_argument("--scenario_id", type=int, default=0, help="场景ID")
-
-    args = parser.parse_args()
-
-    verify_observations(args.data_dir, args.scenario_id)

README.md

@@ -1,401 +1,121 @@
-# MAGAIL4AutoDrive - 多智能体自动驾驶环境
+# MAGAIL4AutoDrive
 
-基于 MetaDrive 的多智能体自动驾驶仿真与回放环境，支持 Waymo Open Dataset 的专家轨迹回放和自定义策略仿真。
+基于 **MetaDrive** 仿真器和 **Waymo Open Motion Dataset** 的自动驾驶多智能体模仿学习（MAGAIL）与行为克隆（BC）训练系统。
 
-## 📋 目录
+本项目旨在从真实的 Waymo 驾驶数据中提取专家轨迹，并通过模仿学习（Imitation Learning）训练能够适应复杂交互场景的自动驾驶策略。
 
-- [项目简介](#项目简介)
+## 目录结构
-- [功能特性](#功能特性)
-- [环境要求](#环境要求)
-- [安装步骤](#安装步骤)
-- [快速开始](#快速开始)
-- [使用指南](#使用指南)
-- [项目结构](#项目结构)
-- [配置说明](#配置说明)
-- [常见问题](#常见问题)
 
-## 项目简介
+```text
-
+MAGAIL4AutoDrive/
-MAGAIL4AutoDrive 是一个基于 MetaDrive 0.4.3 的多智能体自动驾驶环境，专为模仿学习（Imitation Learning）和强化学习（Reinforcement Learning）研究设计。项目支持从真实世界数据集（如 Waymo Open Dataset）中加载场景，并提供两种核心运行模式：
+├── Algorithm/                 # 强化学习与模仿学习算法实现
-
+│   ├── policy.py              # 基础策略网络 (MLP 等)
-- **回放模式（Replay Mode）**：严格按照专家轨迹回放，用于数据可视化和验证
+│   ├── ppo.py                 # PPO 算法实现
-- **仿真模式（Simulation Mode）**：使用自定义策略控制车辆，用于算法训练和测试
+│   ├── magail.py              # MAGAIL 算法核心逻辑
-
+│   ├── disc.py                # 判别器 (Discriminator) 网络
-## 功能特性
-
-### 核心功能
-- ✅ **多智能体支持**：同时控制多辆车辆进行协同仿真
-- ✅ **专家轨迹回放**：精确回放 Waymo 数据集中的专家驾驶行为
-- ✅ **自定义策略接口**：灵活接入各种控制策略（IDM、RL 等）
-- ✅ **智能车道过滤**：自动过滤不在车道上的异常车辆
-- ✅ **场景时长控制**：支持使用数据集原始场景时长或自定义 horizon
-- ✅ **丰富的传感器**：LiDAR、侧向检测器、车道线检测器、相机、仪表盘
-
-### 高级特性
-- 🎯 指定场景 ID 运行
-- 🔄 自动场景切换（修复版）
-- 📊 详细的调试日志输出
-- 🚗 车辆动态生成与管理
-- 🎮 支持可视化渲染和无头运行
-
-## 环境要求
-
-### 系统要求
-- **操作系统**：Ubuntu 18.04+ / macOS 10.14+ / Windows 10+
-- **Python 版本**：3.8 - 3.10
-- **GPU**：可选，但推荐使用（用于加速渲染）
-
-### 依赖库
-```
-
-metadrive-simulator==0.4.3
-numpy>=1.19.0
-pygame>=2.0.0
-
-```
-
-## 安装步骤
-
-### 1. 创建 Conda 环境
-```
-
-conda create -n metadrive python=3.10
-conda activate metadrive
-
-```
-
-### 2. 安装 MetaDrive
-```
-
-pip install metadrive-simulator==0.4.3
-
-```
-
-### 3. 克隆项目
-```
-
-git clone https://github.com/your-username/MAGAIL4AutoDrive.git
-cd MAGAIL4AutoDrive/Env
-
-```
-
-### 4. 准备数据集
-将 Waymo 数据集转换为 MetaDrive 格式并放置在项目目录下：
-```
-
-MAGAIL4AutoDrive/Env/
-├── exp_converted/
-│   ├── scenario_0/
-│   ├── scenario_1/
 │   └── ...
-
+├── Env/                       # 仿真环境封装 (MetaDrive Wrapper)
+│   ├── bc_env.py              # BCScenarioEnv，45 维观测（BC/MAGAIL 共用）
+│   ├── scenario_env.py       # 多智能体基础场景环境
+│   ├── expert_replay_env.py  # 专家轨迹回放环境（数据生成与回放）
+│   ├── inverse_dynamics.py   # 逆动力学模块 (轨迹 -> 动作)
+│   ├── simple_idm_policy.py  # ConstantVelocityPolicy 占位策略
+│   └── ...
+├── dataset/                   # 数据集加载器
+│   ├── loader.py              # 主流水线：load_expert_pkl、MAGAILExpertDataset
+│   └── expert_dataset.py      # 可选 107 维/5 维管线
+├── scripts/                   # 工具脚本（数据、回放、可视化、分析）
+│   ├── generate_expert_data.py    # 从 Waymo 生成专家 (obs, act) pkl
+│   ├── visualize.py              # 可视化统一入口（replay / policy / trajectory）
+│   ├── analyze_expert_data.py    # 数据分布分析
+│   ├── launch_tensorboard.py     # 启动 TensorBoard
+│   ├── README.md                  # 脚本用法说明
+│   └── ...
+├── data/                      # 数据目录（相对路径）
+│   ├── exp_filtered/          # Waymo 场景数据
+│   ├── training_data/        # 专家 pkl 输出（generate_expert_data）
+│   └── trajectories/         # 其他轨迹 pkl（如 expert_dataset 输出）
+├── models/                    # 模型保存目录（相对路径）
+│   ├── bc/                    # BC 模型 (.pt)
+│   └── magail/                # MAGAIL 模型 (*_actor.pth, *_critic.pth)
+├── logs/                      # 训练日志 (TensorBoard)
+│   ├── bc/
+│   └── magail/
+├── train_bc.py                # [根目录] BC 训练
+├── train_magail.py            # [根目录] MAGAIL 训练
+└── README.md
 ```
 
-## 快速开始
+## 路径约定（相对项目根）
 
-### 回放模式（推荐先尝试）
+- **数据**：Waymo 场景 `data/exp_filtered`；专家 pkl `data/training_data`；其他轨迹 `data/trajectories`
-```
+- **模型**：BC `models/bc/`，MAGAIL `models/magail/`
+- **日志**：TensorBoard 写入 `logs/bc/`、`logs/magail/`
 
+所有默认路径均为相对项目根，便于在不同设备上复用。
 
-# 使用场景原始时长回放第一个场景
+## 数据处理流程
 
-python run_multiagent_env.py --mode replay --episodes 1 --use_scenario_duration
+从 Waymo Motion 原始数据到本项目训练用专家 pkl，依次为：
 
-# 回放指定场景
+**1) 下载 Waymo Motion（TFRecord）**  
-
+安装 `gsutil` 并登录 Google 账号后，例如只下载 training_20s：
-python run_multiagent_env.py --mode replay --scenario_id 0 --use_scenario_duration
-
-# 回放多个场景
-
-python run_multiagent_env.py --mode replay --episodes 3 --use_scenario_duration
-
-```
-
-### 仿真模式
-```
-
-
-# 使用默认策略运行仿真
-
-python run_multiagent_env.py --mode simulation --episodes 1
-
-# 无渲染运行（加速训练）
-
-python run_multiagent_env.py --mode simulation --episodes 5 --no_render
-
-```
-
-## 使用指南
-
-### 命令行参数
-
-| 参数 | 类型 | 默认值 | 说明 |
-|------|------|--------|------|
-| `--mode` | str | simulation | 运行模式：`replay` 或 `simulation` |
-| `--data_dir` | str | 当前目录 | Waymo 数据目录路径 |
-| `--episodes` | int | 1 | 运行回合数 |
-| `--horizon` | int | 300 | 每回合最大步数 |
-| `--no_render` | flag | False | 禁用渲染（加速运行） |
-| `--debug` | flag | False | 启用调试模式 |
-| `--scenario_id` | int | None | 指定场景 ID |
-| `--use_scenario_duration` | flag | False | 使用场景原始时长 |
-| `--no_vehicles` | flag | False | 禁止生成车辆 |
-| `--no_pedestrians` | flag | False | 禁止生成行人 |
-| `--no_cyclists` | flag | False | 禁止生成自行车 |
-
-### 回放模式详解
-
-回放模式严格按照专家轨迹回放车辆状态，不涉及物理引擎控制。主要用途：
-- 数据集可视化
-- 验证数据质量
-- 生成演示视频
 
 ```bash
-# 完整参数示例
+gsutil -m cp -r "gs://waymo_open_dataset_motion_v_1_2_0/uncompressed/scenario/training_20s" ./waymo/
-python run_multiagent_env.py \
-  --mode replay \
-  --episodes 1 \
-  --use_scenario_duration \
-  --debug
-
-# 仅回放车辆，禁止行人和自行车
-python run_multiagent_env.py \
-  --mode replay \
-  --use_scenario_duration \
-  --no_pedestrians \
-  --no_cyclists
 ```
 
-**重要提示**：回放模式建议始终启用 `--use_scenario_duration`，否则会出现场景播放完后继续运行的问题。
+**2) ScenarioNet Convert（TFRecord → ScenarioNet 场景库）**  
-
+需安装 ScenarioNet、MetaDrive 及 TensorFlow 2.11、protobuf 3.20；转换时不用 GPU。
-### 仿真模式详解
-
-仿真模式使用自定义策略控制车辆，适合算法开发和测试：
 
 ```bash
-# 基础仿真
+python -m scenarionet.convert_waymo -d data/exp_converted --raw_data_path ./waymo/training_20s --num_workers 64
-python run_multiagent_env.py --mode simulation
-
-# 长时间训练（无渲染）
-python run_multiagent_env.py \
-  --mode simulation \
-  --episodes 100 \
-  --horizon 500 \
-  --no_render
-
-# 仅车辆仿真（用于专注车车交互场景）
-python run_multiagent_env.py \
-  --mode simulation \
-  --no_pedestrians \
-  --no_cyclists
 ```
 
-### 自定义策略
+**3) ScenarioNet Filter（按需筛选场景）**  
+从 convert 得到的场景库中筛掉含红绿灯、天桥等场景，输出到如 `data/exp_filtered`。具体命令以 ScenarioNet 文档为准（Operations → Filter）。
 
-修改 `simple_idm_policy.py` 或创建新的策略类：
+**4) 本项目：生成专家 pkl**  
+使用筛选后的场景目录，生成训练用 pkl 到 `data/training_data`：
 
-```python
-class CustomPolicy:
-def __init__(self, **kwargs):
-# 初始化策略参数
-pass
-
-    def act(self, observation=None):
-        # 返回动作 [steering, acceleration]
-        # steering: [-1, 1]
-        # acceleration: [-1, 1]
-        return [0.0, 0.5]
-```
-
-在 `run_multiagent_env.py` 中使用：
-```
-
-from custom_policy import CustomPolicy
-
-env = MultiAgentScenarioEnv(
-config={...},
-agent2policy=CustomPolicy()
-)
-
-```
-
-## 项目结构
-
-```
-
-MAGAIL4AutoDrive/Env/
-├── run_multiagent_env.py      \# 主运行脚本
-├── scenario_env.py             \# 多智能体场景环境
-├── replay_policy.py            \# 专家轨迹回放策略
-├── simple_idm_policy.py        \# IDM 策略实现
-├── utils.py                    \# 工具函数
-├── ENHANCED_USAGE_GUIDE.md     \# 详细使用指南
-├── README.md                   \# 本文档
-└── exp_converted/              \# Waymo 数据集（需自行准备）
-├── scenario_0/
-├── scenario_1/
-└── ...
-
-```
-
-### 核心文件说明
-
-**run_multiagent_env.py**
-- 主入口脚本
-- 处理命令行参数
-- 管理回放和仿真两种模式的运行逻辑
-
-**scenario_env.py**
-- 自定义多智能体环境类
-- 车辆生成与管理
-- 车道过滤逻辑
-- 观测空间定义
-
-**replay_policy.py**
-- 专家轨迹回放策略
-- 逐帧状态查询
-- 轨迹完成判断
-
-**simple_idm_policy.py**
-- 简单的恒速策略示例
-- 可作为自定义策略的模板
-
-## 配置说明
-
-### 环境配置参数
-
-在 `scenario_env.py` 的 `default_config()` 中可修改：
-
-```python
-config.update(dict(
-    data_directory=None,              # 数据目录
-    num_controlled_agents=3,          # 可控车辆数量（仅仿真模式）
-    horizon=1000,                     # 最大步数
-    filter_offroad_vehicles=True,     # 是否过滤车道外车辆
-    lane_tolerance=3.0,               # 车道容差（米）
-    replay_mode=False,                # 是否为回放模式
-    specific_scenario_id=None,        # 指定场景 ID
-    use_scenario_duration=False,      # 使用场景原始时长
-    # 对象类型过滤选项
-    spawn_vehicles=True,              # 是否生成车辆
-    spawn_pedestrians=True,           # 是否生成行人
-    spawn_cyclists=True,              # 是否生成自行车
-))
-```
-
-### 传感器配置
-
-默认启用的传感器（可在环境初始化时修改）：
-- **LiDAR**：80 条激光，探测距离 30 米
-- **侧向检测器**：10 条激光，探测距离 8 米
-- **车道线检测器**：10 条激光，探测距离 3 米
-- **主相机**：分辨率 1200x900
-- **仪表盘**：车辆状态信息
-
-## 常见问题
-
-### Q1: 回放模式为什么超出数据集的最大帧数还在继续？
-**A**: 需要添加 `--use_scenario_duration` 参数。修复版本已在 `scenario_env.py` 中添加了自动检测机制。
-
-### Q2: 如何切换不同的场景？
-**A**: 
-- 方法一：使用 `--scenario_id` 指定场景
-- 方法二：使用 `--episodes N` 自动遍历 N 个场景
-
-### Q3: 为什么有些车辆没有出现？
-**A**: 启用了车道过滤功能（`filter_offroad_vehicles=True`），不在车道上的车辆会被过滤。可以通过设置 `lane_tolerance` 调整容差或关闭此功能。
-
-### Q4: 如何提高运行速度？
-**A**: 
-- 使用 `--no_render` 禁用可视化
-- 减少 `num_controlled_agents` 数量
-- 使用 GPU 加速
-
-### Q5: 如何控制场景中的对象类型？
-**A**: 使用对象过滤参数：
 ```bash
-# 仅车辆，无行人和自行车
+python scripts/generate_expert_data.py --data_dir data/exp_filtered --output_dir data/training_data --num_scenarios 100 --start_index 0
-python run_multiagent_env.py --mode replay --no_pedestrians --no_cyclists
-
-# 仅行人和自行车，无车辆（特殊场景）
-python run_multiagent_env.py --mode replay --no_vehicles
-
-# 调试模式查看过滤统计
-python run_multiagent_env.py --mode replay --debug --no_pedestrians
 ```
 
-### Q6: 为什么有些车辆生成在空中？
+## 核心工作流
-**A**: 已在 v1.2.0 中修复。现在所有车辆位置都只使用 2D 坐标（x, y），z 坐标设为 0，让 MetaDrive 自动处理高度，确保车辆贴在地面上。
 
-### Q7: 如何导出观测数据？
+### 1. 数据准备
-**A**: 在 `run_multiagent_env.py` 中添加数据保存逻辑：
+使用 `scripts/generate_expert_data.py` 将 Waymo 数据转换为训练用 `.pkl`，输出到 `data/training_data/`。
-```python
-import pickle
 
-obs_data = []
+```bash
-while True:
+python scripts/generate_expert_data.py --data_dir data/exp_filtered --output_dir data/training_data --num_scenarios 100
-    obs, rewards, dones, infos = env.step(actions)
-    obs_data.append(obs)
-    if dones["__all__"]:
-        break
-
-with open('observations.pkl', 'wb') as f:
-    pickle.dump(obs_data, f)
 ```
 
-## 更新日志
+### 2. 行为克隆 (BC)
+- **训练**：`python train_bc.py`（模型保存到 `models/bc/`，日志到 `logs/bc/`）
+- **可视化**：`python scripts/visualize.py policy --policy_type bc --model_path models/bc/policy_best.pt`
 
-### v1.2.0 (2025-10-26)
+### 3. 多智能体对抗模仿学习 (MAGAIL)
-- ✅ 修复车辆生成高度问题（车辆悬空）
+- **训练**：`python train_magail.py`（模型保存到 `models/magail/`，日志到 `logs/magail/`）
-- ✅ 添加对象类型过滤功能（车辆/行人/自行车）
+- **可视化**：`python scripts/visualize.py policy --policy_type magail --model_path models/magail/model_50_actor.pth`
-- ✅ 新增命令行参数：`--no_vehicles`、`--no_pedestrians`、`--no_cyclists`
-- ✅ 改进调试信息输出，显示各类型对象统计
-- ✅ 优化位置处理逻辑，只使用 2D 坐标避免高度问题
 
-### v1.1.0 (2025-10-26)
+### 4. 可视化统一入口
-- ✅ 修复回放模式超出场景时长问题
+可视化统一使用 `scripts/visualize.py`，子命令：`replay`（场景回放）、`policy`（BC/MAGAIL 策略）、`trajectory`（专家轨迹 2D 动画）。详见 [scripts/README.md](scripts/README.md)。
-- ✅ 添加场景自动切换功能
-- ✅ 改进 `replay_policy.py`，新增 `is_finished()` 方法
-- ✅ 优化 `scenario_env.py` 的 done 判断逻辑
-- ✅ 修复多回合运行时的对象清理问题
 
-### v1.0.0 (初始版本)
+## 文件与模块职责
-- 基础多智能体环境实现
-- 回放和仿真两种模式
-- 车道过滤功能
-- Waymo 数据集支持
 
-## 贡献指南
+### 根目录脚本
+- **train_bc.py**：BC 训练，从 `dataset.loader` 加载专家 pkl，模型与日志写入 `models/bc/`、`logs/bc/`
+- **train_magail.py**：MAGAIL 训练，环境使用 `BCScenarioEnv`（45 维），从 `dataset.loader` 加载专家数据，模型与日志写入 `models/magail/`、`logs/magail/`
 
-欢迎提交 Issue 和 Pull Request！
+### Env 模块
+- **Env/bc_env.py**：`BCScenarioEnv`，45 维观测（Ego 5 维 + 10 邻居×4 维），BC 与 MAGAIL 训练/评估共用
+- **Env/scenario_env.py**：`MultiAgentScenarioEnv` 基类，Waymo 场景加载与步进
+- **Env/expert_replay_env.py**：专家轨迹回放与逆动力学动作，供 `generate_expert_data.py` 与回放可视化
+- **Env/inverse_dynamics.py**：轨迹 → 油门/转向动作
 
-### 提交 Issue
+### Algorithm 模块
-- 请详细描述问题和复现步骤
+- **Algorithm/policy.py**：`StateIndependentPolicy`，BC 使用的 MLP 策略
-- 附上运行日志和错误信息
-- 说明运行环境（OS、Python 版本等）
 
-### 提交 PR
+### scripts 目录
-- Fork 本项目
+工具脚本用途与用法见 [scripts/README.md](scripts/README.md)。
-- 创建特性分支：`git checkout -b feature/your-feature`
-- 提交更改：`git commit -m 'Add some feature'`
-- 推送分支：`git push origin feature/your-feature`
-- 提交 Pull Request
-
-## 许可证
-
-本项目基于 MIT 许可证开源。
-
-## 致谢
-
-- [MetaDrive](https://github.com/metadriverse/metadrive) - 优秀的驾驶仿真平台
-- [Waymo Open Dataset](https://waymo.com/open/) - 高质量的自动驾驶数据集
-
-## 联系方式
-
-如有问题或建议，请通过以下方式联系：
-- GitHub Issues: [项目 Issues 页面]
-- Email: huangfukk@xxx.com
-
----
-
-**Happy Driving! 🚗💨**

TRAINING_ARCHITECTURE.md

@@ -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 作为策略优化算法。当前主要问题集中在环境稳定性和日志记录方面，需要进一步优化。

dataset/expert_dataset.py

@@ -0,0 +1,305 @@
+import sys
+import os
+current_dir = os.path.dirname(os.path.abspath(__file__))
+project_root = os.path.dirname(current_dir)
+sys.path.insert(0, os.path.join(project_root, "Env"))
+
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+import pickle
+from scenario_env import MultiAgentScenarioEnv
+from metadrive.engine.asset_loader import AssetLoader
+
+class DummyPolicy:
+    def act(self, *args, **kwargs):
+        return np.array([0.0, 0.0])
+
+class ExpertTrajectoryDataset(Dataset):
+    """
+    完整107维观测的专家轨迹数据集
+    """
+    
+    def __init__(self, 
+                 trajectory_data: dict,
+                 observation_data: dict = None,  # 可选的完整观测
+                 sequence_length: int = 1,
+                 extract_actions: bool = True):
+        """
+        Args:
+            trajectory_data: 专家轨迹数据
+            observation_data: 完整107维观测数据(可选)
+            sequence_length: 序列长度
+            extract_actions: 是否提取动作
+        """
+        self.trajectory_data = trajectory_data
+        self.observation_data = observation_data if observation_data else {}
+        self.sequence_length = sequence_length
+        self.extract_actions = extract_actions
+        
+        # 构建索引
+        self.indices = []
+        for traj_id, traj in trajectory_data.items():
+            traj_len = traj["length"]
+            for start_idx in range(traj_len - sequence_length):
+                self.indices.append((traj_id, start_idx))
+        
+        obs_dim = 107 if len(self.observation_data) > 0 else 5
+        print(f"专家数据集: {len(trajectory_data)} 条轨迹, "
+              f"{len(self.indices)} 个训练样本, 观测维度: {obs_dim}")
+    
+    def __len__(self):
+        return len(self.indices)
+    
+    def __getitem__(self, idx):
+        traj_id, start_idx = self.indices[idx]
+        traj = self.trajectory_data[traj_id]
+        
+        end_idx = start_idx + self.sequence_length
+        
+        # 如果有完整观测,使用完整观测(107维)
+        if traj_id in self.observation_data and len(self.observation_data[traj_id]) > 0:
+            obs_sequence = self.observation_data[traj_id]
+            states = obs_sequence[start_idx:end_idx]  # (seq_len, 107)
+        else:
+            # 否则使用简化观测(5维)
+            positions = traj["positions"][start_idx:end_idx+1]
+            headings = traj["headings"][start_idx:end_idx+1]
+            velocities = traj["velocities"][start_idx:end_idx]
+            
+            states = []
+            for i in range(self.sequence_length):
+                state = np.concatenate([
+                    positions[i, :2],    # x, y
+                    velocities[i],       # vx, vy
+                    [headings[i]],       # heading
+                ])
+                states.append(state)
+            states = np.array(states)
+        
+        if self.extract_actions:
+            positions = traj["positions"][start_idx:end_idx+1]
+            headings = traj["headings"][start_idx:end_idx+1]
+            velocities = traj["velocities"][start_idx:end_idx]
+            
+            actions = self._extract_actions_from_states(
+                positions[:-1], positions[1:], 
+                headings[:-1], headings[1:],
+                velocities
+            )
+            return torch.FloatTensor(states), torch.FloatTensor(actions)
+        else:
+            next_states = states[1:]
+            return torch.FloatTensor(states[:-1]), torch.FloatTensor(next_states)
+    
+    def _extract_actions_from_states(self, pos_t, pos_t1, head_t, head_t1, vel_t):
+        """从状态序列反推动作"""
+        actions = []
+        dt = 0.1
+        
+        for i in range(len(pos_t)):
+            current_speed = np.linalg.norm(vel_t[i])
+            displacement = np.linalg.norm(pos_t1[i, :2] - pos_t[i, :2])
+            next_speed = displacement / dt
+            
+            speed_change = (next_speed - current_speed) / dt
+            if speed_change >= 0:
+                throttle = np.clip(speed_change / 5.0, 0.0, 1.0)
+            else:
+                throttle = np.clip(speed_change / 8.0, -1.0, 0.0)
+            
+            heading_change = head_t1[i] - head_t[i]
+            heading_change = np.arctan2(np.sin(heading_change), np.cos(heading_change))
+            steering = np.clip(heading_change / 0.2, -1.0, 1.0)
+            
+            actions.append([throttle, steering])
+        
+        return np.array(actions)
+    
+    @staticmethod
+    def collect_with_full_obs(env_config, num_scenarios=10, save_path=None):
+        """
+        ✅ 使用env._get_all_obs()收集完整107维观测
+        
+        这是正确的方法!直接利用环境已有的观测函数
+        """
+        all_trajectories = {}
+        all_observations = {}
+        
+        # 检查数据库
+        data_dir = env_config["config"]["data_directory"]
+        summary_path = os.path.join(data_dir, "dataset_summary.pkl")
+        
+        with open(summary_path, 'rb') as f:
+            summary = pickle.load(f)
+        
+        total_scenarios = len(summary)
+        print(f"数据库总场景数: {total_scenarios}")
+        
+        if num_scenarios is None:
+            num_scenarios = total_scenarios
+        else:
+            num_scenarios = min(num_scenarios, total_scenarios)
+        
+        print(f"计划收集(完整107维观测): {num_scenarios} 个场景")
+        
+        for i in range(num_scenarios):
+            try:
+                # 创建环境
+                env = MultiAgentScenarioEnv(
+                    config={
+                        **env_config["config"],
+                        "start_scenario_index": i,
+                        "num_scenarios": 1,
+                    },
+                    agent2policy=env_config["agent2policy"]
+                )
+                
+                # 重置环境
+                env.reset()
+                
+                if not hasattr(env, 'expert_trajectories'):
+                    print(f"⚠️ 场景 {i}: 缺少expert_trajectories")
+                    env.close()
+                    continue
+                
+                expert_trajs = env.expert_trajectories
+                
+                if len(expert_trajs) == 0:
+                    print(f"⚠️ 场景 {i}: 无专家轨迹")
+                    env.close()
+                    continue
+                
+                # 存储轨迹
+                scenario_id = env.engine.current_seed
+                for obj_id, traj in expert_trajs.items():
+                    unique_id = f"scenario{i}_{obj_id}"
+                    all_trajectories[unique_id] = traj
+                
+                # ✅ 关键: 使用_get_all_obs()获取完整观测
+                # 创建agent_id到unique_id的映射
+                agent_to_unique = {}
+                for agent_id in env.controlled_agents.keys():
+                    # 尝试匹配agent_id到expert_trajectories的obj_id
+                    for obj_id in expert_trajs.keys():
+                        if str(agent_id) in str(obj_id) or str(obj_id) in str(agent_id):
+                            unique_id = f"scenario{i}_{obj_id}"
+                            agent_to_unique[agent_id] = unique_id
+                            all_observations[unique_id] = []
+                            break
+                
+                # 遍历场景的每一步,收集完整观测
+                max_steps = min([traj["length"] for traj in expert_trajs.values()])
+                
+                for step in range(max_steps):
+                    # ✅ 直接调用_get_all_obs()获取107维观测!
+                    obs_list = env._get_all_obs()
+                    
+                    # 存储每个agent的观测
+                    for agent_idx, agent_id in enumerate(env.controlled_agents.keys()):
+                        if agent_id in agent_to_unique:
+                            unique_id = agent_to_unique[agent_id]
+                            if agent_idx < len(obs_list):
+                                # obs_list[agent_idx]已经是107维向量!
+                                all_observations[unique_id].append(np.array(obs_list[agent_idx]))
+                    
+                    # 执行零动作(保持场景状态)
+                    actions = {aid: np.array([0.0, 0.0]) 
+                              for aid in env.controlled_agents.keys()}
+                    env.step(actions)
+                
+                # 转换为numpy数组
+                for unique_id in list(all_observations.keys()):
+                    if len(all_observations[unique_id]) > 0:
+                        all_observations[unique_id] = np.array(all_observations[unique_id])
+                    else:
+                        del all_observations[unique_id]
+                
+                env.close()
+                
+                if (i + 1) % 5 == 0:
+                    print(f"✓ 已收集 {i+1}/{num_scenarios}, "
+                          f"轨迹: {len(all_trajectories)}, "
+                          f"观测: {len(all_observations)}")
+                          
+            except Exception as e:
+                print(f"✗ 场景 {i} 收集失败: {e}")
+                import traceback
+                traceback.print_exc()
+                try:
+                    env.close()
+                except:
+                    pass
+                continue
+        
+        print(f"\n收集完成!")
+        print(f"  轨迹数: {len(all_trajectories)}")
+        print(f"  完整观测数: {len(all_observations)}")
+        
+        # 验证观测维度
+        if len(all_observations) > 0:
+            sample_obs = list(all_observations.values())[0]
+            if len(sample_obs) > 0:
+                obs_dim = len(sample_obs[0])
+                print(f"  观测维度: {obs_dim} (应为107)")
+        
+        if save_path:
+            os.makedirs(os.path.dirname(save_path), exist_ok=True)
+            with open(save_path, "wb") as f:
+                pickle.dump({
+                    "trajectories": all_trajectories,
+                    "observations": all_observations
+                }, f)
+            print(f"数据已保存到: {save_path}")
+        
+        return all_trajectories, all_observations
+
+
+if __name__ == "__main__":
+    WAYMO_DATA_DIR = r"/home/huangfukk/mdsn"
+    data_dir = AssetLoader.file_path(WAYMO_DATA_DIR, "exp_filtered", unix_style=False)
+    
+    env_config = {
+        "config": {
+            "data_directory": data_dir,
+            "is_multi_agent": True,
+            "num_controlled_agents": 3,
+            "use_render": False,
+            "sequential_seed": True,
+        },
+        "agent2policy": DummyPolicy()
+    }
+    
+    print("=" * 60)
+    print("选择收集模式:")
+    print("1. 简化观测(5维) - 快速,已验证 ✅")
+    print("2. 完整观测(107维) - 使用_get_all_obs() ⭐")
+    print("=" * 60)
+    
+    mode = input("请选择模式(1或2,默认1): ").strip() or "1"
+    
+    if mode == "2":
+        print("\n开始收集完整107维观测...")
+        trajectories, observations = ExpertTrajectoryDataset.collect_with_full_obs(
+            env_config,
+            num_scenarios=10,
+            save_path="data/trajectories/expert_trajectories_full.pkl"
+        )
+        
+        if len(trajectories) > 0:
+            dataset = ExpertTrajectoryDataset(
+                trajectories, 
+                observations,
+                sequence_length=1
+            )
+            state, action = dataset[0]
+            print(f"\n数据集测试:")
+            print(f"  总轨迹数: {len(trajectories)}")
+            print(f"  总观测数: {len(observations)}")
+            print(f"  训练样本数: {len(dataset)}")
+            print(f"  状态维度: {state.shape}")
+            print(f"  动作维度: {action.shape}")
+    else:
+        print("\n开始收集简化5维观测...")
+        # 保持原有的简化版本代码...
+        print("(使用之前已成功的方法)")

dataset/loader.py

@@ -0,0 +1,103 @@
+"""
+统一数据加载：BC/MAGAIL 训练用专家 pkl 的加载函数与 Dataset。
+主训练流水线使用本模块；dataset/expert_dataset.py 为可选 107 维/5 维管线。
+"""
+import os
+import glob
+import pickle
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+
+
+def load_expert_pkl(expert_data_path):
+    """从目录或单个 pkl 加载专家 (obs, acts)，返回 concat 后的 obs_data, act_data。"""
+    if os.path.isdir(expert_data_path):
+        pkl_files = glob.glob(os.path.join(expert_data_path, "*.pkl"))
+        if not pkl_files:
+            raise FileNotFoundError(f"No .pkl files in {expert_data_path}")
+        print(f"Found {len(pkl_files)} pickle files in {expert_data_path}")
+    elif os.path.exists(expert_data_path):
+        pkl_files = [expert_data_path]
+    else:
+        raise FileNotFoundError(f"Expert data path not found: {expert_data_path}")
+
+    obs_data, act_data = [], []
+    for pkl_file in pkl_files:
+        try:
+            with open(pkl_file, "rb") as f:
+                data = pickle.load(f)
+            if isinstance(data, list):
+                for traj in data:
+                    if "obs" in traj and "acts" in traj:
+                        obs_data.append(traj["obs"])
+                        act_data.append(traj["acts"])
+            elif isinstance(data, dict):
+                if "observations" in data and "actions" in data:
+                    obs_data.append(data["observations"])
+                    act_data.append(data["actions"])
+            else:
+                print(f"Skipping {pkl_file}: Unknown data format {type(data)}")
+        except Exception as e:
+            print(f"Error loading {pkl_file}: {e}")
+
+    if len(obs_data) == 0:
+        raise ValueError("No valid data loaded from provided path.")
+    obs_data = np.concatenate(obs_data, axis=0)
+    act_data = np.concatenate(act_data, axis=0)
+    print(f"Total loaded samples: {len(obs_data)}")
+    return obs_data, act_data
+
+
+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)
+            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]
+
+        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

scripts/README.md

@@ -0,0 +1,76 @@
+# scripts 工具脚本说明
+
+本目录包含数据生成、回放、可视化与分析等工具脚本。训练脚本（`train_bc.py`、`train_magail.py`）位于项目根目录。
+
+## 路径约定（相对项目根）
+
+- **数据**：`data/exp_filtered`（Waymo 场景）、`data/training_data`（专家 pkl 输出）
+- **模型**：`models/bc/`（BC）、`models/magail/`（MAGAIL）
+- **日志**：`logs/bc/`、`logs/magail/`（TensorBoard）
+
+---
+
+## 脚本列表与用法
+
+### 数据生成
+
+| 脚本 | 用途 | 用法示例 |
+|------|------|----------|
+| [generate_expert_data.py](generate_expert_data.py) | 从 Waymo 数据生成专家 (obs, act) 的 pkl | `python scripts/generate_expert_data.py --data_dir data/exp_filtered --output_dir data/training_data --num_scenarios 100` |
+
+**常用参数**：`--data_dir`（默认 `data/exp_filtered`）、`--output_dir`（默认 `data/training_data`）、`--start_index`、`--num_scenarios`。
+
+---
+
+### 可视化（统一入口）
+
+| 脚本 | 用途 | 用法示例 |
+|------|------|----------|
+| [visualize.py](visualize.py) | **replay**：场景回放（ExpertReplayEnv）；**policy**：BC/MAGAIL 策略；**trajectory**：专家轨迹 2D 动画 | 见下方 |
+
+**子命令**：
+
+- **replay**（原始专家轨迹回放）：
+```bash
+python scripts/visualize.py replay --data_dir data/exp_filtered --num_scenarios 1 --horizon 500
+```
+
+- **policy**（BC 或 MAGAIL 训练策略）：与专家数据生成/回放一致——同一套车道+静态筛选、且会生成背景车（bg_*），使观测分布与训练集一致，便于在训练集上公平演示。
+```bash
+python scripts/visualize.py policy --policy_type bc --model_path models/bc/policy_best.pt --data_dir data/exp_filtered --num_scenarios 1
+python scripts/visualize.py policy --policy_type magail --model_path models/magail/model_50_actor.pth --num_scenarios 1 --deterministic
+```
+
+- **trajectory**（专家轨迹 matplotlib 俯视图动画）：
+```bash
+python scripts/visualize.py trajectory --data_dir data/exp_filtered --scenario_idx 0
+```
+
+**公共参数**：`--data_dir`（默认 `data/exp_filtered`）、`--start_index`、`--num_scenarios`、`--horizon`。policy 模式另有 `--policy_type`（auto/bc/magail）、`--model_path`、`--deterministic`（仅 MAGAIL）。
+
+---
+
+### 数据分析与检查
+
+| 脚本 | 用途 | 用法示例 |
+|------|------|----------|
+| [analyze_expert_data.py](analyze_expert_data.py) | 分析专家数据分布与统计 | 见脚本内 `__main__`（依赖 env 与数据目录配置） |
+| [check_track_fields.py](check_track_fields.py) | 检查 Waymo 轨迹字段 | 见脚本内 `__main__` |
+| [check_database_info.py](check_database_info.py) | 检查数据库/场景信息 | 见脚本内 `__main__`（含硬编码路径，可按需改为 `data/exp_filtered`） |
+
+---
+
+### 其他
+
+| 脚本 | 用途 | 用法示例 |
+|------|------|----------|
+| [launch_tensorboard.py](launch_tensorboard.py) | 启动 TensorBoard | `python scripts/launch_tensorboard.py --logdir logs`（或 `logs/bc` / `logs/magail`） |
+
+---
+
+## 与训练流程的对应关系
+
+1. **数据准备**：`generate_expert_data.py` → 输出到 `data/training_data/*.pkl`
+2. **BC 训练**：根目录 `train_bc.py` → 模型保存到 `models/bc/`，日志到 `logs/bc/`
+3. **MAGAIL 训练**：根目录 `train_magail.py` → 模型保存到 `models/magail/`，日志到 `logs/magail/`
+4. **可视化**：`scripts/visualize.py`（子命令 replay / policy / trajectory）→ 数据目录默认 `data/exp_filtered`

scripts/analyze_expert_data.py

@@ -0,0 +1,256 @@
+import sys
+import os
+
+# 添加路径
+current_dir = os.path.dirname(os.path.abspath(__file__))
+project_root = os.path.dirname(current_dir)
+env_dir = os.path.join(project_root, "Env")
+sys.path.insert(0, project_root)
+sys.path.insert(0, env_dir)
+
+import numpy as np
+import matplotlib.pyplot as plt
+from collections import defaultdict
+from scenario_env import MultiAgentScenarioEnv
+from metadrive.engine.asset_loader import AssetLoader
+import pickle
+import os
+
+class DummyPolicy:
+    """占位策略"""
+    def act(self, *args, **kwargs):
+        return np.array([0.0, 0.0])
+
+class ExpertDataAnalyzer:
+    def __init__(self, data_directory):
+        self.data_directory = data_directory
+        self.env = MultiAgentScenarioEnv(
+            config={
+                "data_directory": data_directory,
+                "is_multi_agent": True,
+                "num_controlled_agents": 3,
+                "use_render": False,
+                "sequential_seed": True,
+            },
+            agent2policy=DummyPolicy()  # 添加必需参数
+        )
+        
+        self.statistics = {
+            "num_scenarios": 0,
+            "num_trajectories": 0,
+            "trajectory_lengths": [],
+            "velocities": [],
+            "speeds": [],  # 速度大小
+            "accelerations": [],
+            "heading_changes": [],
+            "inter_vehicle_distances": [],
+            "num_vehicles_per_scenario": [],
+            "static_vehicles": 0,  # 统计静止车辆
+        }
+    
+    def analyze_all_scenarios(self, num_scenarios=None):
+        """遍历所有场景并收集统计信息"""
+        scenario_count = 0
+        
+        while True:
+            try:
+                obs = self.env.reset()
+                
+                if not hasattr(self.env, 'expert_trajectories'):
+                    print("⚠️ 环境缺少expert_trajectories属性")
+                    break
+                
+                expert_trajs = self.env.expert_trajectories
+                
+                if len(expert_trajs) == 0:
+                    continue
+                
+                scenario_count += 1
+                self.statistics["num_scenarios"] += 1
+                self.statistics["num_vehicles_per_scenario"].append(len(expert_trajs))
+                
+                # 分析每条轨迹
+                for obj_id, traj in expert_trajs.items():
+                    self.analyze_single_trajectory(traj)
+                
+                # 分析车辆间交互
+                self.analyze_vehicle_interactions(expert_trajs)
+                
+                print(f"已分析场景 {scenario_count}/{num_scenarios}, 车辆数: {len(expert_trajs)}")
+                
+                if num_scenarios and scenario_count >= num_scenarios:
+                    break
+                    
+            except Exception as e:
+                print(f"场景 {scenario_count} 处理失败: {e}")
+                break
+        
+        self.env.close()
+    
+    def analyze_single_trajectory(self, traj):
+        """分析单条轨迹"""
+        self.statistics["num_trajectories"] += 1
+        
+        length = traj["length"]
+        self.statistics["trajectory_lengths"].append(length)
+        
+        # 速度分析
+        velocities = traj["velocities"]
+        speeds = np.linalg.norm(velocities, axis=1)
+        self.statistics["velocities"].extend(velocities.tolist())
+        self.statistics["speeds"].extend(speeds.tolist())
+        
+        # 检查是否为静止车辆
+        if np.max(speeds) < 0.5:  # 最大速度小于0.5m/s视为静止
+            self.statistics["static_vehicles"] += 1
+        
+        # 加速度分析
+        if length > 1:
+            accelerations = np.diff(speeds) * 10  # 10Hz数据
+            self.statistics["accelerations"].extend(accelerations.tolist())
+        
+        # 航向角变化
+        headings = traj["headings"]
+        if length > 1:
+            heading_changes = np.diff(headings)
+            heading_changes = np.arctan2(np.sin(heading_changes), np.cos(heading_changes))
+            self.statistics["heading_changes"].extend(heading_changes.tolist())
+    
+    def analyze_vehicle_interactions(self, expert_trajs):
+        """分析车辆间的距离"""
+        if len(expert_trajs) < 2:
+            return
+        
+        traj_list = list(expert_trajs.values())
+        
+        for i in range(len(traj_list)):
+            for j in range(i+1, len(traj_list)):
+                traj_i = traj_list[i]
+                traj_j = traj_list[j]
+                
+                start_time = max(traj_i["start_timestep"], traj_j["start_timestep"])
+                end_time = min(traj_i["end_timestep"], traj_j["end_timestep"])
+                
+                if start_time >= end_time:
+                    continue
+                
+                idx_i_start = start_time - traj_i["start_timestep"]
+                idx_i_end = end_time - traj_i["start_timestep"]
+                idx_j_start = start_time - traj_j["start_timestep"]
+                idx_j_end = end_time - traj_j["start_timestep"]
+                
+                pos_i = traj_i["positions"][idx_i_start:idx_i_end, :2]
+                pos_j = traj_j["positions"][idx_j_start:idx_j_end, :2]
+                
+                distances = np.linalg.norm(pos_i - pos_j, axis=1)
+                self.statistics["inter_vehicle_distances"].extend(distances.tolist())
+    
+    def generate_report(self, save_dir="./analysis_results"):
+        """生成统计报告"""
+        os.makedirs(save_dir, exist_ok=True)
+        
+        stats = self.statistics
+        
+        print("\n" + "="*60)
+        print("专家数据集统计报告")
+        print("="*60)
+        print(f"总场景数: {stats['num_scenarios']}")
+        print(f"总轨迹数: {stats['num_trajectories']}")
+        print(f"静止车辆数: {stats['static_vehicles']} ({stats['static_vehicles']/stats['num_trajectories']*100:.1f}%)")
+        print(f"平均每场景车辆数: {np.mean(stats['num_vehicles_per_scenario']):.2f} ± {np.std(stats['num_vehicles_per_scenario']):.2f}")
+        
+        print(f"\n轨迹长度统计 (帧数 @ 10Hz):")
+        print(f"  平均: {np.mean(stats['trajectory_lengths']):.2f} 帧 ({np.mean(stats['trajectory_lengths'])*0.1:.2f}秒)")
+        print(f"  中位数: {np.median(stats['trajectory_lengths']):.2f} 帧")
+        print(f"  最小/最大: {np.min(stats['trajectory_lengths'])} / {np.max(stats['trajectory_lengths'])} 帧")
+        
+        print(f"\n速度统计 (m/s):")
+        speeds = np.array(stats['speeds'])
+        print(f"  平均: {np.mean(speeds):.2f} ± {np.std(speeds):.2f}")
+        print(f"  中位数: {np.median(speeds):.2f}")
+        print(f"  最小/最大: {np.min(speeds):.2f} / {np.max(speeds):.2f}")
+        print(f"  静止帧(<0.5m/s): {np.sum(speeds < 0.5)} ({np.sum(speeds < 0.5)/len(speeds)*100:.1f}%)")
+        
+        print(f"\n加速度统计 (m/s²):")
+        accs = np.array(stats['accelerations'])
+        print(f"  平均: {np.mean(accs):.4f} ± {np.std(accs):.2f}")
+        print(f"  最小/最大: {np.min(accs):.2f} / {np.max(accs):.2f}")
+        
+        if len(stats['inter_vehicle_distances']) > 0:
+            dists = np.array(stats['inter_vehicle_distances'])
+            print(f"\n车辆间距离统计 (m):")
+            print(f"  平均: {np.mean(dists):.2f} ± {np.std(dists):.2f}")
+            print(f"  最小: {np.min(dists):.2f}")
+            print(f"  近距离交互(<5m): {np.sum(dists < 5.0)} ({np.sum(dists < 5.0)/len(dists)*100:.2f}%)")
+        
+        # 保存数据
+        with open(os.path.join(save_dir, "statistics.pkl"), "wb") as f:
+            pickle.dump(stats, f)
+        
+        # 绘制可视化
+        self.plot_distributions(save_dir)
+        
+        print(f"\n✓ 报告已保存到: {save_dir}")
+    
+    def plot_distributions(self, save_dir):
+        """绘制分布图"""
+        stats = self.statistics
+        
+        fig, axes = plt.subplots(2, 3, figsize=(15, 10))
+        
+        # 1. 轨迹长度分布
+        axes[0, 0].hist(stats['trajectory_lengths'], bins=50, edgecolor='black')
+        axes[0, 0].set_xlabel('Trajectory Length (frames @ 10Hz)')
+        axes[0, 0].set_ylabel('Frequency')
+        axes[0, 0].set_title('Trajectory Length Distribution')
+        axes[0, 0].axvline(np.mean(stats['trajectory_lengths']), color='red', 
+                           linestyle='--', label=f'Mean: {np.mean(stats["trajectory_lengths"]):.1f}')
+        axes[0, 0].legend()
+        
+        # 2. 速度分布
+        axes[0, 1].hist(stats['speeds'], bins=50, edgecolor='black')
+        axes[0, 1].set_xlabel('Speed (m/s)')
+        axes[0, 1].set_ylabel('Frequency')
+        axes[0, 1].set_title('Speed Distribution')
+        axes[0, 1].axvline(np.mean(stats['speeds']), color='red', 
+                           linestyle='--', label=f'Mean: {np.mean(stats["speeds"]):.2f}')
+        axes[0, 1].legend()
+        
+        # 3. 加速度分布
+        axes[0, 2].hist(stats['accelerations'], bins=50, edgecolor='black')
+        axes[0, 2].set_xlabel('Acceleration (m/s²)')
+        axes[0, 2].set_ylabel('Frequency')
+        axes[0, 2].set_title('Acceleration Distribution')
+        
+        # 4. 每场景车辆数
+        axes[1, 0].hist(stats['num_vehicles_per_scenario'], bins=30, edgecolor='black')
+        axes[1, 0].set_xlabel('Vehicles per Scenario')
+        axes[1, 0].set_ylabel('Frequency')
+        axes[1, 0].set_title('Vehicles per Scenario')
+        
+        # 5. 航向角变化
+        axes[1, 1].hist(stats['heading_changes'], bins=50, edgecolor='black')
+        axes[1, 1].set_xlabel('Heading Change (rad)')
+        axes[1, 1].set_ylabel('Frequency')
+        axes[1, 1].set_title('Heading Change Distribution')
+        
+        # 6. 车辆间距离
+        if len(stats['inter_vehicle_distances']) > 0:
+            axes[1, 2].hist(stats['inter_vehicle_distances'], bins=50, 
+                           range=(0, 50), edgecolor='black')
+            axes[1, 2].set_xlabel('Inter-vehicle Distance (m)')
+            axes[1, 2].set_ylabel('Frequency')
+            axes[1, 2].set_title('Distance Distribution')
+        
+        plt.tight_layout()
+        plt.savefig(os.path.join(save_dir, "distributions.png"), dpi=300)
+        print(f"  ✓ 分布图已保存")
+
+if __name__ == "__main__":
+    WAYMO_DATA_DIR = r"/home/huangfukk/MAGAIL4AutoDrive/data"
+    data_dir = AssetLoader.file_path(WAYMO_DATA_DIR, "exp_filtered", unix_style=False)
+    
+    print("开始分析专家数据...")
+    analyzer = ExpertDataAnalyzer(data_dir)
+    analyzer.analyze_all_scenarios(num_scenarios=100)  # 分析100个场景
+    analyzer.generate_report()

scripts/check_database_info.py

@@ -0,0 +1,47 @@
+import pickle
+import os
+
+# 检查过滤后的数据库
+filtered_db = "/home/huangfukk/mdsn/exp_filtered"
+
+print("="*60)
+print("过滤后数据库信息")
+print("="*60)
+
+# 读取summary
+summary_path = os.path.join(filtered_db, "dataset_summary.pkl")
+with open(summary_path, 'rb') as f:
+    summary = pickle.load(f)
+
+print(f"\n总场景数: {len(summary)}")
+print(f"场景ID列表(前10个): {list(summary.keys())[:10]}")
+
+# 读取mapping
+mapping_path = os.path.join(filtered_db, "dataset_mapping.pkl")
+with open(mapping_path, 'rb') as f:
+    mapping = pickle.load(f)
+
+print(f"\n映射关系数量: {len(mapping)}")
+
+# 检查第一个场景的详细信息
+first_scenario_id = list(summary.keys())[0]
+first_scenario_info = summary[first_scenario_id]
+print(f"\n第一个场景详细信息:")
+print(f"  场景ID: {first_scenario_id}")
+print(f"  元数据: {first_scenario_info}")
+
+# 检查映射的文件路径
+first_scenario_path = mapping[first_scenario_id]
+print(f"  场景文件路径(相对): {first_scenario_path}")
+
+# 检查文件是否存在
+abs_path = os.path.join(filtered_db, first_scenario_path)
+print(f"  场景文件路径(绝对): {abs_path}")
+print(f"  文件存在: {os.path.exists(abs_path)}")
+
+# 统计源数据库的场景文件
+converted_db = "/home/huangfukk/mdsn/exp_converted"
+converted_files = [f for f in os.listdir(converted_db) if f.endswith('.pkl') and f.startswith('sd_')]
+print(f"\n源数据库 exp_converted:")
+print(f"  场景文件数量: {len(converted_files)}")
+print(f"  示例文件: {converted_files[:5]}")

scripts/check_track_fields.py

@@ -0,0 +1,177 @@
+import sys
+import os
+
+# 添加路径
+current_dir = os.path.dirname(os.path.abspath(__file__))
+project_root = os.path.dirname(current_dir)
+env_dir = os.path.join(project_root, "Env")
+
+sys.path.insert(0, project_root)
+sys.path.insert(0, env_dir)
+
+from scenario_env import MultiAgentScenarioEnv
+from metadrive.engine.asset_loader import AssetLoader
+import numpy as np
+
+class DummyPolicy:
+    """
+    占位策略,用于数据检查时初始化环境
+    不需要实际执行动作,只是为了满足环境初始化要求
+    """
+    def act(self, *args, **kwargs):
+        # 返回零动作 [throttle, steering]
+        return np.array([0.0, 0.0])
+
+def check_available_fields():
+    """
+    检查Waymo转MetaDrive数据中实际可用的字段
+    """
+    WAYMO_DATA_DIR = r"/home/huangfukk/mdsn"
+    data_dir = AssetLoader.file_path(WAYMO_DATA_DIR, "exp_filtered", unix_style=False)
+    
+    # 创建占位策略
+    dummy_policy = DummyPolicy()
+    
+    # 初始化环境,传入必需的agent2policy参数
+    env = MultiAgentScenarioEnv(
+        config={
+            "data_directory": data_dir,
+            "is_multi_agent": True,
+            "num_controlled_agents": 3,
+            "use_render": False,
+            "sequential_seed": True,
+        },
+        agent2policy=dummy_policy  # 添加这个必需参数
+    )
+    
+    print("✓ 环境初始化成功")
+    
+    # 重置环境以加载数据
+    print("正在加载场景数据...")
+    env.reset()
+    
+    # 检查是否有expert_trajectories属性
+    if hasattr(env, 'expert_trajectories'):
+        print(f"✓ expert_trajectories属性存在,包含 {len(env.expert_trajectories)} 条轨迹")
+    else:
+        print("⚠️ expert_trajectories属性不存在,请先修改scenario_env.py添加轨迹存储功能")
+    
+    # 获取一个track样本
+    sample_track = None
+    for scenario_id, track in env.engine.traffic_manager.current_traffic_data.items():
+        if track["type"] == "VEHICLE":
+            sample_track = track
+            print(f"\n找到样本车辆: scenario_id = {scenario_id}")
+            break
+    
+    if sample_track is None:
+        print("未找到车辆轨迹数据")
+        env.close()
+        return
+    
+    print("="*60)
+    print("Track数据结构分析")
+    print("="*60)
+    
+    # 1. 顶层字段
+    print("\n1. Track顶层字段:")
+    for key in sample_track.keys():
+        print(f"   - {key}: {type(sample_track[key])}")
+    
+    # 2. metadata字段
+    print("\n2. track['metadata']字段:")
+    if "metadata" in sample_track:
+        for key, value in sample_track["metadata"].items():
+            if isinstance(value, (str, int, float, bool)):
+                print(f"   - {key}: {type(value).__name__} = {value}")
+            else:
+                print(f"   - {key}: {type(value).__name__}")
+    
+    # 3. state字段
+    print("\n3. track['state']字段:")
+    if "state" in sample_track:
+        for key, value in sample_track["state"].items():
+            if isinstance(value, np.ndarray):
+                print(f"   - {key}: shape={value.shape}, dtype={value.dtype}")
+                # 打印第一个有效值
+                if "valid" in sample_track["state"]:
+                    valid_idx = np.argmax(sample_track["state"]["valid"])
+                    if valid_idx >= 0 and valid_idx < len(value):
+                        print(f"     示例值 (index {valid_idx}): {value[valid_idx]}")
+            else:
+                print(f"   - {key}: {type(value)} = {value}")
+    
+    print("\n" + "="*60)
+    print("建议存储的字段:")
+    print("="*60)
+    
+    # 检查必需字段
+    required_fields = ["position", "heading", "velocity", "valid"]
+    print("\n必需字段:")
+    all_required_exist = True
+    for field in required_fields:
+        if "state" in sample_track and field in sample_track["state"]:
+            print(f"  ✓ {field} (存在)")
+        else:
+            print(f"  ✗ {field} (缺失)")
+            all_required_exist = False
+    
+    # 检查可选字段
+    optional_fields = ["length", "width", "height", "bbox"]
+    print("\n可选字段:")
+    available_optional = []
+    for field in optional_fields:
+        if "state" in sample_track and field in sample_track["state"]:
+            print(f"  + {field} (在state中)")
+            available_optional.append(field)
+        elif "metadata" in sample_track and field in sample_track["metadata"]:
+            print(f"  + {field} (在metadata中)")
+            available_optional.append(field)
+        else:
+            print(f"  - {field} (不存在)")
+    
+    print("\n" + "="*60)
+    print("推荐的trajectory_data结构:")
+    print("="*60)
+    
+    if all_required_exist:
+        print("""
+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(),
+    "headings": track["state"]["heading"][first_show:last_show+1].copy(),
+    "velocities": track["state"]["velocity"][first_show:last_show+1].copy(),
+    "timesteps": np.arange(first_show, last_show+1),
+    "start_timestep": first_show,
+    "end_timestep": last_show,
+    "length": last_show - first_show + 1
+}
+""")
+        
+        if available_optional:
+            print("如果需要车辆尺寸,可选添加:")
+            for field in available_optional:
+                if field in ["length", "width", "height"]:
+                    print(f'  trajectory_data["vehicle_{field}"] = track["state" or "metadata"]["{field}"][first_show]')
+    else:
+        print("⚠️ 缺少必需字段,请检查数据转换流程")
+    
+    # 如果有expert_trajectories,展示一个样本
+    if hasattr(env, 'expert_trajectories') and len(env.expert_trajectories) > 0:
+        print("\n" + "="*60)
+        print("expert_trajectories样本:")
+        print("="*60)
+        sample_traj = list(env.expert_trajectories.values())[0]
+        for key, value in sample_traj.items():
+            if isinstance(value, np.ndarray):
+                print(f"  {key}: shape={value.shape}, dtype={value.dtype}")
+            else:
+                print(f"  {key}: {type(value).__name__} = {value}")
+    
+    env.close()
+    print("\n✓ 分析完成")
+
+if __name__ == "__main__":
+    check_available_fields()

scripts/generate_expert_data.py

@@ -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="data/exp_filtered", help="Path to Waymo pickles (or filtered index)")
+    parser.add_argument("--output_dir", type=str, default="data/training_data", 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)

scripts/launch_tensorboard.py

@@ -0,0 +1,18 @@
+import sys
+import types
+import os
+
+# Mock imghdr module for Python 3.13 compatibility
+# TensorBoard depends on imghdr which was removed in Python 3.13
+if sys.version_info >= (3, 13):
+    if 'imghdr' not in sys.modules:
+        imghdr_mock = types.ModuleType('imghdr')
+        imghdr_mock.what = lambda filename, h=None: None
+        # Mock tests list which tensorboard appends to
+        imghdr_mock.tests = []
+        sys.modules['imghdr'] = imghdr_mock
+
+from tensorboard import main as tb_main
+
+if __name__ == '__main__':
+    sys.exit(tb_main.run_main())

scripts/visualize.py

@@ -0,0 +1,398 @@
+"""
+Unified visualization: replay (scenario replay), policy (BC/MAGAIL), trajectory (2D expert trajectory animation).
+Usage: python scripts/visualize.py <replay|policy|trajectory> [args...]
+"""
+import argparse
+import os
+import sys
+import time
+import numpy as np
+import torch
+
+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)
+
+# --- Replay ---
+def _run_replay(args):
+    from Env.expert_replay_env import ExpertReplayEnv
+
+    data_path = os.path.abspath(args.data_dir)
+    if not os.path.exists(data_path):
+        raise ValueError(f"Data directory {data_path} not found")
+
+    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,
+        "sequential_seed": True,
+        "reactive_traffic": False,
+        "start_scenario_index": args.start_index,
+        "num_scenarios": -1,
+        "log_level": 40,
+    }
+
+    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):
+                obs, rewards, dones, infos = env.step(None)
+                env.render(
+                    mode="top_down",
+                    text={"Step": step, "Agents": len(env.controlled_agents), "Scenario": i},
+                )
+                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.")
+
+
+# --- Policy (BC / MAGAIL) ---
+def _resolve_data_dir(data_dir_arg):
+    if data_dir_arg:
+        data_dir = data_dir_arg
+    else:
+        data_dir = os.path.join(project_root, "data", "exp_filtered")
+        if not os.path.exists(data_dir):
+            data_dir = os.path.join(project_root, "data", "exp_converted")
+    if not os.path.exists(data_dir):
+        raise FileNotFoundError(f"Data directory not found at {data_dir}. Please specify --data_dir.")
+    return data_dir
+
+
+def _resolve_model_path(model_path, policy_type):
+    if os.path.exists(model_path):
+        return model_path
+    if policy_type == "bc":
+        candidate = os.path.join(project_root, "models", "bc", os.path.basename(model_path))
+    else:
+        candidate = os.path.join(project_root, "models", "magail", os.path.basename(model_path))
+    if os.path.exists(candidate):
+        return candidate
+    if policy_type == "magail" and not model_path.endswith("_actor.pth"):
+        candidate = os.path.join(project_root, "models", "magail", os.path.basename(model_path) + "_actor.pth")
+        if os.path.exists(candidate):
+            return candidate
+    raise FileNotFoundError(f"Model path {model_path} not found.")
+
+
+def _run_policy(args):
+    from Env.bc_env import BCScenarioEnv
+    from metadrive.engine.engine_utils import close_engine
+
+    policy_type = (args.policy_type or "auto").lower()
+    if policy_type == "auto":
+        policy_type = "bc" if args.model_path.endswith(".pt") else "magail"
+
+    data_dir = _resolve_data_dir(args.data_dir)
+    data_path = os.path.abspath(data_dir)
+    env_config = {
+        "data_directory": data_path,
+        "is_multi_agent": True,
+        "num_controlled_agents": 3,
+        "horizon": args.horizon,
+        "use_render": True,
+        "sequential_seed": True,
+        "start_scenario_index": args.start_index,
+        "num_scenarios": args.num_scenarios,
+        "log_level": 40,
+    }
+
+    print(f"Initializing BCScenarioEnv (policy_type={policy_type})...")
+    try:
+        env = BCScenarioEnv(env_config, agent2policy={})
+    except Exception as e:
+        print(f"Error init env: {e}. Trying to close lingering engine...")
+        try:
+            close_engine()
+        except Exception:
+            pass
+        env = BCScenarioEnv(env_config, agent2policy={})
+
+    state_dim = 45
+    action_dim = 2
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model_path = _resolve_model_path(args.model_path, policy_type)
+    print(f"Loading model from {model_path}...")
+
+    if policy_type == "bc":
+        from Algorithm.policy import StateIndependentPolicy
+        policy = StateIndependentPolicy(
+            state_shape=(state_dim,),
+            action_shape=(action_dim,),
+            hidden_units=(256, 256),
+            hidden_activation=torch.nn.Tanh(),
+        ).to(device)
+        policy.load_state_dict(torch.load(model_path, map_location=device))
+        policy.eval()
+    else:
+        from train_magail import Actor
+        actor = Actor(state_dim, action_dim).to(device)
+        actor.load_state_dict(torch.load(model_path, map_location=device))
+        actor.eval()
+
+    try:
+        for i in range(args.start_index, args.start_index + args.num_scenarios):
+            print(f"\n--- Playing Scenario {i} ---")
+            try:
+                obs_dict = env.reset(seed=i)
+            except Exception as e:
+                print(f"Error resetting {i}: {e}. Skipping.")
+                try:
+                    close_engine()
+                    env = BCScenarioEnv(env_config, agent2policy={})
+                except Exception:
+                    pass
+                continue
+
+            print(f"Scenario loaded. Controlled agents: {len(obs_dict)}")
+            if len(obs_dict) == 0:
+                print(f"Scenario {i} has no controlled agents (all filtered out). Skipping.")
+                continue
+            step_count = 0
+            episode_reward = 0.0
+
+            while True:
+                agent_ids = list(obs_dict.keys())
+                obs_list = [obs_dict[aid] for aid in agent_ids]
+                obs_tensor = torch.FloatTensor(np.array(obs_list)).to(device)
+
+                with torch.no_grad():
+                    if policy_type == "bc":
+                        actions_np = policy(obs_tensor).cpu().numpy()
+                    else:
+                        dist = actor(obs_tensor)
+                        if args.deterministic:
+                            actions_np = torch.tanh(dist.mean).cpu().numpy()
+                        else:
+                            actions_np = torch.tanh(dist.sample()).cpu().numpy()
+
+                actions = {aid: actions_np[idx].flatten() for idx, aid in enumerate(agent_ids)}
+                obs_dict, rewards, dones, infos = env.step(actions)
+                episode_reward += sum(rewards.values())
+
+                env.render(
+                    mode="top_down",
+                    text={
+                        "Scenario": i,
+                        "Step": step_count,
+                        "Agents": len(obs_dict),
+                        "Total Reward": f"{episode_reward:.2f}",
+                    },
+                )
+                step_count += 1
+
+                if dones["__all__"] or step_count >= args.horizon:
+                    print(f"Scenario finished at step {step_count}, reward {episode_reward:.2f}")
+                    break
+    except KeyboardInterrupt:
+        print("Interrupted.")
+    finally:
+        env.close()
+
+
+# --- Trajectory (matplotlib 2D animation) ---
+def _build_expert_trajectories_from_env(env):
+    """Build expert_trajectories dict from env (ExpertReplayEnv has traffic_manager.current_traffic_data)."""
+    if hasattr(env, "expert_trajectories") and env.expert_trajectories:
+        return env.expert_trajectories
+    if not hasattr(env, "engine") or not hasattr(env.engine, "traffic_manager"):
+        return {}
+    from metadrive.type import MetaDriveType
+    data = getattr(env.engine.traffic_manager, "current_traffic_data", None)
+    if not data:
+        return {}
+    expert_trajs = {}
+    for scenario_id, track in data.items():
+        if track.get("type") != MetaDriveType.VEHICLE or "state" not in track:
+            continue
+        state = track["state"]
+        positions = state.get("position")
+        if positions is None:
+            continue
+        valid = state.get("valid", np.ones(len(positions), dtype=bool))
+        valid = np.asarray(valid).flatten()
+        if valid.size != len(positions):
+            valid = np.ones(len(positions), dtype=bool)
+        first_show = int(np.argmax(valid)) if valid.any() else 0
+        last_show = len(valid) - 1 - int(np.argmax(valid[::-1])) if valid.any() else len(positions) - 1
+        obj_id = track.get("metadata", {}).get("object_id", str(scenario_id))
+        expert_trajs[obj_id] = {
+            "positions": np.asarray(positions),
+            "start_timestep": first_show,
+            "end_timestep": last_show,
+        }
+    return expert_trajs
+
+
+def _run_trajectory_animation(expert_trajs, scenario_idx):
+    import matplotlib.pyplot as plt
+    from matplotlib.animation import FuncAnimation
+
+    if len(expert_trajs) == 0:
+        print("No expert trajectories to visualize.")
+        return
+
+    fig, ax = plt.subplots(figsize=(12, 12))
+    max_timestep = max(t["end_timestep"] for t in expert_trajs.values())
+    min_timestep = min(t["start_timestep"] for t in expert_trajs.values())
+
+    colors = plt.cm.tab10(np.linspace(0, 1, len(expert_trajs)))
+    for idx, (obj_id, traj) in enumerate(expert_trajs.items()):
+        positions = np.asarray(traj["positions"])
+        if positions.ndim >= 2:
+            positions = positions[:, :2]
+        else:
+            continue
+        ax.plot(
+            positions[:, 0], positions[:, 1],
+            color=colors[idx], alpha=0.3, linewidth=1,
+            label=f"Vehicle {str(obj_id)[:6]}",
+        )
+
+    scatter = ax.scatter([], [], s=200, c="red", marker="o", edgecolors="black", linewidths=2)
+    time_text = ax.text(0.02, 0.95, "", transform=ax.transAxes, fontsize=14)
+    ax.set_xlabel("X (m)")
+    ax.set_ylabel("Y (m)")
+    ax.set_title(f"Expert Trajectory Visualization - Scenario {scenario_idx}")
+    ax.legend(loc="upper right", fontsize=8)
+    ax.grid(True, alpha=0.3)
+    ax.axis("equal")
+
+    def update(frame):
+        current_time = min_timestep + frame
+        current_positions = []
+        for traj in expert_trajs.values():
+            st, et = traj["start_timestep"], traj["end_timestep"]
+            if st <= current_time <= et:
+                pos = np.asarray(traj["positions"])
+                if pos.ndim >= 2:
+                    pos = pos[current_time - st, :2]
+                else:
+                    continue
+                current_positions.append(pos)
+        if current_positions:
+            scatter.set_offsets(np.array(current_positions))
+        time_text.set_text(f"Time: {frame * 0.1:.1f}s (Frame {frame})")
+        return scatter, time_text
+
+    anim = FuncAnimation(
+        fig, update, frames=max_timestep - min_timestep + 1,
+        interval=100, blit=True, repeat=True,
+    )
+    plt.tight_layout()
+    plt.show()
+    return anim
+
+
+def _run_trajectory(args):
+    from Env.expert_replay_env import ExpertReplayEnv
+
+    data_dir = _resolve_data_dir(args.data_dir)
+    data_path = os.path.abspath(data_dir)
+    env_config = {
+        "data_directory": data_path,
+        "is_multi_agent": True,
+        "num_controlled_agents": 100,
+        "horizon": 500,
+        "use_render": False,
+        "sequential_seed": True,
+        "reactive_traffic": False,
+        "start_scenario_index": args.scenario_idx,
+        "num_scenarios": 1,
+        "log_level": 40,
+    }
+
+    env = ExpertReplayEnv(config=env_config)
+    try:
+        env.reset(seed=args.scenario_idx)
+        expert_trajs = _build_expert_trajectories_from_env(env)
+        _run_trajectory_animation(expert_trajs, args.scenario_idx)
+    finally:
+        env.close()
+
+
+# --- Main ---
+def main():
+    parser = argparse.ArgumentParser(
+        description="Unified visualization: replay, policy (BC/MAGAIL), trajectory.",
+    )
+    subparsers = parser.add_subparsers(dest="mode", required=True, help="replay | policy | trajectory")
+
+    # Common args for data_dir (used by all)
+    def add_common_data_args(p):
+        p.add_argument("--data_dir", type=str, default="data/exp_filtered", help="Waymo scenario directory")
+        p.add_argument("--start_index", type=int, default=0)
+        p.add_argument("--num_scenarios", type=int, default=1)
+        p.add_argument("--horizon", type=int, default=200)
+
+    # replay
+    pr = subparsers.add_parser("replay", help="Replay scenario with ExpertReplayEnv (no policy)")
+    add_common_data_args(pr)
+    pr.set_defaults(horizon=500)
+
+    # policy
+    pp = subparsers.add_parser("policy", help="Visualize BC or MAGAIL trained policy")
+    add_common_data_args(pp)
+    pp.add_argument("--policy_type", type=str, default="auto", choices=["auto", "bc", "magail"])
+    pp.add_argument("--model_path", type=str, default="models/bc/policy_best.pt")
+    pp.add_argument("--deterministic", action="store_true", help="MAGAIL: use mean action")
+
+    # trajectory
+    pt = subparsers.add_parser("trajectory", help="2D matplotlib animation of expert trajectories")
+    pt.add_argument("--data_dir", type=str, default="data/exp_filtered")
+    pt.add_argument("--scenario_idx", type=int, default=0)
+
+    args = parser.parse_args()
+
+    # Resolve data_dir relative to project root when default
+    if args.mode != "trajectory":
+        if args.data_dir in ("data/exp_filtered", "data/exp_converted"):
+            args.data_dir = os.path.join(project_root, args.data_dir)
+    else:
+        if args.data_dir in ("data/exp_filtered", "data/exp_converted"):
+            args.data_dir = os.path.join(project_root, args.data_dir)
+
+    if args.mode == "replay":
+        _run_replay(args)
+    elif args.mode == "policy":
+        _run_policy(args)
+    elif args.mode == "trajectory":
+        _run_trajectory(args)
+    else:
+        parser.error(f"Unknown mode: {args.mode}")
+
+
+if __name__ == "__main__":
+    main()

train_bc.py

@@ -0,0 +1,146 @@
+"""
+BC 训练脚本：负责数据加载、环境评估、日志与保存；BC 算法由 Algorithm.bc 提供。
+使用方式不变：python train_bc.py [--expert_data_path data/training_data] [--save_dir models/bc] ...
+"""
+import os
+import numpy as np
+import torch
+import argparse
+from torch.utils.data import DataLoader, TensorDataset
+from torch.optim import Adam
+from torch.optim.lr_scheduler import ExponentialLR
+from datetime import datetime
+from torch.utils.tensorboard import SummaryWriter
+
+from Algorithm.policy import StateIndependentPolicy
+from Algorithm.bc import train_bc_epoch, eval_bc_epoch
+from Env.bc_env import BCScenarioEnv
+from dataset.loader import load_expert_pkl
+
+
+def evaluate_policy(policy, args, device):
+    """在 BCScenarioEnv 中评估策略，跑若干 episode，返回平均 reward。"""
+    waymo_data_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "data")
+    data_dir = os.path.join(waymo_data_dir, "exp_filtered")
+    if not os.path.exists(data_dir):
+        data_dir = os.path.join(waymo_data_dir, "exp_converted")
+    if not os.path.exists(data_dir):
+        print(f"[ERROR] Could not find scenario data in {waymo_data_dir}. Evaluation skipped.")
+        return 0.0
+
+    env_config = {
+        "data_directory": data_dir,
+        "is_multi_agent": True,
+        "num_controlled_agents": 3,
+        "use_render": False,
+        "sequential_seed": True,
+        "horizon": 200,
+    }
+    env = BCScenarioEnv(env_config, agent2policy=None)
+    total_rewards = []
+
+    try:
+        for i in range(3):
+            obs_dict = env.reset(seed=i)
+            episode_reward = 0
+            dones = {"__all__": False}
+            step_count = 0
+            horizon = 200
+            while not dones["__all__"]:
+                step_count += 1
+                if step_count >= horizon:
+                    break
+                if not obs_dict:
+                    obs_dict, _, dones, _ = env.step({})
+                    continue
+                agent_ids = list(obs_dict.keys())
+                obs_list = [obs_dict[aid] for aid in agent_ids]
+                obs_tensor = torch.FloatTensor(np.array(obs_list)).to(device)
+                with torch.no_grad():
+                    actions, _ = policy.sample(obs_tensor)
+                    actions = actions.cpu().numpy()
+                action_dict = {aid: act for aid, act in zip(agent_ids, actions)}
+                obs_dict, rewards, dones, _ = env.step(action_dict)
+                episode_reward += sum(rewards.values())
+            total_rewards.append(episode_reward)
+            print(f"  Eval Episode {i}: Total Reward {episode_reward:.2f}")
+        avg_reward = float(np.mean(total_rewards))
+        print(f"  Average Evaluation Reward: {avg_reward:.2f}")
+        return avg_reward
+    except Exception as e:
+        print(f"Evaluation failed: {e}")
+        import traceback
+        traceback.print_exc()
+        return 0.0
+    finally:
+        env.close()
+
+
+def main(args):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Using device: {device}")
+
+    os.makedirs("logs/bc", exist_ok=True)
+    log_dir = os.path.join("logs", "bc", datetime.now().strftime("%Y%m%d-%H%M%S"))
+    writer = SummaryWriter(log_dir)
+    print(f"TensorBoard logging to: {log_dir}")
+    os.makedirs(args.save_dir, exist_ok=True)
+
+    obs_data, act_data = load_expert_pkl(args.expert_data_path)
+    obs_tensor = torch.FloatTensor(obs_data)
+    act_tensor = torch.FloatTensor(act_data)
+    dataset = TensorDataset(obs_tensor, act_tensor)
+    train_size = int(0.8 * len(dataset))
+    val_size = len(dataset) - train_size
+    train_dataset, val_dataset = torch.utils.data.random_split(dataset, [train_size, val_size])
+    train_loader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True)
+    val_loader = DataLoader(val_dataset, batch_size=args.batch_size, shuffle=False)
+    print(f"Dataset loaded. Train size: {len(train_dataset)}, Val size: {len(val_dataset)}")
+
+    state_dim = obs_data.shape[1]
+    action_dim = act_data.shape[1]
+    print(f"State Dim: {state_dim}, Action Dim: {action_dim}")
+
+    policy = StateIndependentPolicy(
+        state_shape=(state_dim,),
+        action_shape=(action_dim,),
+        hidden_units=(256, 256),
+        hidden_activation=torch.nn.Tanh(),
+    ).to(device)
+    optimizer = Adam(policy.parameters(), lr=args.lr)
+    scheduler = ExponentialLR(optimizer, gamma=0.99)
+
+    best_val_loss = float("inf")
+    for epoch in range(args.epochs):
+        avg_train_loss = train_bc_epoch(policy, train_loader, optimizer, device)
+        scheduler.step()
+        avg_val_loss = eval_bc_epoch(policy, val_loader, device)
+
+        print(f"Epoch {epoch+1}/{args.epochs} | Train Loss: {avg_train_loss:.4f} | Val Loss: {avg_val_loss:.4f}")
+        writer.add_scalar("Loss/train", avg_train_loss, epoch)
+        writer.add_scalar("Loss/val", avg_val_loss, epoch)
+        writer.add_scalar("Learning_rate", scheduler.get_last_lr()[0], epoch)
+
+        if avg_val_loss < best_val_loss:
+            best_val_loss = avg_val_loss
+            torch.save(policy.state_dict(), os.path.join(args.save_dir, "policy_best.pt"))
+
+        if (epoch + 1) % args.eval_freq == 0:
+            eval_reward = evaluate_policy(policy, args, device)
+            writer.add_scalar("Reward/eval", eval_reward, epoch)
+
+    torch.save(policy.state_dict(), os.path.join(args.save_dir, "policy_final.pt"))
+    writer.close()
+    print("Training finished.")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--expert_data_path", type=str, default="data/training_data", help="Path to expert data pickle or directory")
+    parser.add_argument("--save_dir", type=str, default="models/bc", help="Directory to save models")
+    parser.add_argument("--epochs", type=int, default=100)
+    parser.add_argument("--batch_size", type=int, default=64)
+    parser.add_argument("--lr", type=float, default=3e-4)
+    parser.add_argument("--eval_freq", type=int, default=10)
+    args = parser.parse_args()
+    main(args)

train_magail.py

@@ -1,115 +1,548 @@
 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
-from torch.utils.tensorboard import SummaryWriter
+import os
-from Env.scenario_env import MultiAgentScenarioEnv
+import argparse
-from Algorithm.magail import MAGAIL
+import signal
-from Algorithm.buffer import RolloutBuffer # 假设 Buffer 在这里
+import sys
-# 假设你有加载专家数据的工具
+from torch.utils.data import DataLoader
-# from utils import load_expert_buffer 
+from dataset.loader import MAGAILExpertDataset
+from Env.bc_env import BCScenarioEnv
 
-# --- 配置 ---
+# --- Networks ---
-CONFIG = {
-    "data_dir": "/home/huangfukk/mdsn",
-    # ... 其他配置 ...
-    "state_dim": 30, # 你的 Observation 维度
-    "action_dim": 2, # [steering, throttle]
-    "rollout_length": 2048, # Buffer 大小
-}
 
-def train():
+class Actor(nn.Module):
-    # 1. 环境
+    def __init__(self, state_dim, action_dim, hidden_dim=256):
-    env = MultiAgentScenarioEnv(config={...}, agent2policy=None)
+        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))
 
-    # 2. 专家 Buffer (伪代码)
+    def forward(self, state):
-    # expert_buffer = RolloutBuffer(...)
+        x = self.net(state)
-    # expert_buffer.load(...) 
+        mu = torch.tanh(self.mu_head(x)) # Action range [-1, 1]
-    # 必须保证 expert_buffer.sample() 返回 (state, next_state) 供 Discriminator 训练
+        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
 
-    # 3. 初始化 MAGAIL
+class Critic(nn.Module):
-    magail = MAGAIL(
+    def __init__(self, state_dim, hidden_dim=256):
-        buffer_exp=expert_buffer, # 传入专家 buffer
+        super(Critic, self).__init__()
-        input_dim=(CONFIG["state_dim"],),
+        self.net = nn.Sequential(
-        action_shape=(CONFIG["action_dim"],), # PPO 初始化可能需要 action_shape，原代码好像没传？
+            nn.Linear(state_dim, hidden_dim),
-        device=torch.device("cuda"),
+            nn.Tanh(),
-        rollout_length=CONFIG["rollout_length"]
+            nn.Linear(hidden_dim, hidden_dim),
+            nn.Tanh(),
+            nn.Linear(hidden_dim, 1)
         )
 
-    writer = SummaryWriter("./logs")
+    def forward(self, state):
+        return self.net(state)
 
-    # 4. 训练循环
+class Discriminator(nn.Module):
-    total_steps = 0
+    def __init__(self, state_dim, action_dim, hidden_dim=256):
-    obs_dict = env.reset()
+        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() 
+        )
 
-    # 将 Dict Obs 转为 Array: (N_agents, Obs_dim)
+    def forward(self, state, action):
-    # 假设所有 Agent 的 Obs 维度相同
+        x = torch.cat([state, action], dim=-1)
-    agents = list(obs_dict.keys())
+        return self.net(x)
-    current_obs = np.stack([obs_dict[a] for a in agents]) 
 
-    # 如果需要 state_gail，假设它就是 obs
+# --- PPO Algorithm ---
-    current_state_gail = current_obs.copy()
 
-    while total_steps < 1e7:
+class PPO:
-        # --- 收集数据 (Rollout) ---
+    def __init__(self, state_dim, action_dim, lr=3e-4, gamma=0.99, eps_clip=0.2, K_epochs=10):
-        # PPO 的 buffer 长度是 rollout_length
+        self.actor = Actor(state_dim, action_dim).cuda()
-        # 我们可以在这里调用 magail.step 来自动处理 explore 和 buffer append
+        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)
         
-        # 注意：magail.step 内部调用了 env.step，这可能不适用于 Multi-Agent 环境返回 Dict 的情况
+        self.gamma = gamma
-        # 原 PPO.step 似乎是为 Single-Agent 或 VectorEnv 设计的
+        self.eps_clip = eps_clip
-        # 这里我们需要手动写 Rollout 循环或者修改 PPO.step 以适配 Dict 返回值
+        self.K_epochs = K_epochs
+        self.mse_loss = nn.MSELoss()
 
-        # --- 方案：手动 Rollout 适配 Multi-Agent ---
+    def _log_prob_from_dist(self, dist, pre_tanh_action):
-        for _ in range(CONFIG["rollout_length"]):
+        # Tanh-squashed Gaussian log-prob with correction term.
-            # 1. 决策
+        log_prob = dist.log_prob(pre_tanh_action)
-            actions_list, log_pis_list = magail.explore(current_obs)
+        correction = torch.log(1 - torch.tanh(pre_tanh_action) ** 2 + 1e-6)
+        return (log_prob - correction).sum(dim=-1)
 
-            # 2. 拼装 Action Dict
+    def select_action(self, state):
-            action_dict = {agent_id: action for agent_id, action in zip(agents, actions_list)}
-            
-            # 3. 环境步进
-            next_obs_dict, rewards_dict, dones_dict, infos_dict = env.step(action_dict)
-            
-            # 4. 处理返回值
-            # 需要处理 Agent 死亡/重置的情况。这里简化假设 Agent 数量不变
-            next_obs = np.stack([next_obs_dict[a] for a in agents])
-            rewards = np.array([rewards_dict[a] for a in agents])
-            dones = np.array([dones_dict[a] for a in agents])
-            # truncated 通常在 infos 里或者 dones 里隐含，需根据 gym 版本确认
-            terminated = dones # 简化
-            truncated = [False] * len(agents) # 简化
-            
-            # 5. 存入 Buffer
-            # 获取当前 Actor 的均值方差用于 PPO 更新
         with torch.no_grad():
-                # 重新计算一遍或者在 explore 里返回
+            state = torch.FloatTensor(state).cuda()
-                # 这里 magail.actor(state) 返回的是 mean (deterministic action)
+            dist = self.actor(state)
-                means = magail.actor(torch.tensor(current_obs, device=magail.device, dtype=torch.float)).cpu().numpy()
+            pre_tanh_action = dist.sample()
-                stds = magail.actor.log_stds.exp().detach().cpu().numpy()
+            action = torch.tanh(pre_tanh_action)
-                # 如果是 StateIndependentPolicy，std 可能是共享的参数，维度需要广播
+            action_logprob = self._log_prob_from_dist(dist, pre_tanh_action)
-                if stds.shape[0] != len(agents):
+        return (
-                     stds = np.repeat(stds, len(agents), axis=0)
+            action.cpu().numpy(),
-
+            action_logprob.cpu().numpy(),
-            magail.buffer.append(
+            pre_tanh_action.cpu().numpy()
-                current_obs, current_state_gail, actions_list, 
-                rewards, dones, terminated, log_pis_list,
-                next_obs, next_obs, # next_state_gail = next_obs
-                means, stds
         )
 
-            current_obs = next_obs
+    def update(self, memory):
-            current_state_gail = next_obs # 更新 gail state
+        # Convert memory to tensors
-            total_steps += len(agents) # 步数增加 N
+        states = torch.FloatTensor(np.array(memory['states'])).cuda()
+        actions = torch.FloatTensor(np.array(memory['actions'])).cuda()
+        pre_tanh_actions = torch.FloatTensor(np.array(memory['pre_tanh_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()
         
-            # 处理 Done
+        # Monte Carlo estimate of state rewards (or GAE if implemented, simplistic here)
-            if all(dones):
+        # Usually for PPO we use GAE. Let's do a simple discounted return for now or bootstrapping.
-                obs_dict = env.reset()
+        # Let's use bootstrapping from critic for returns.
-                current_obs = np.stack([obs_dict[a] for a in agents])
-                current_state_gail = current_obs
         
-        # --- 更新 ---
+        returns = []
-        # 当 Buffer 满时 (rollout_length)，调用 update
+        discounted_reward = 0
-        magail.update(writer, total_steps)
+        # 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()
-        if total_steps % 10000 == 0:
+        next_values = self.critic(next_states).detach()
-            magail.save_models("./models")
         
+        # 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 = self._log_prob_from_dist(dist, pre_tanh_actions)
+            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 (45-dim obs via BCScenarioEnv)
+    # 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 (BCScenarioEnv provides 45-dim obs)
+    # 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 = BCScenarioEnv(env_config, agent2policy={})  # 45-dim obs
+    
+    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': [],
+            'pre_tanh_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 = BCScenarioEnv(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 = {}
+            pre_tanh_actions = {}
+            
+            # 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, pre_tanh = ppo_agent.select_action(obs) # Select action returns numpy
+                actions[agent_id] = act.flatten() # (2,)
+                action_logprobs[agent_id] = logprob # scalar
+                pre_tanh_actions[agent_id] = pre_tanh.flatten()
+                
+            # 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['pre_tanh_actions'].append(pre_tanh_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(dones.get("__all__", 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()
+
+            with torch.no_grad():
+                disc_acc_exp = (exp_preds > 0.5).float().mean().item()
+                disc_acc_pol = (pol_preds < 0.5).float().mean().item()
+            
+            # --- 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)
+            if batch_size > 0:
+                writer.add_scalar('Acc/Disc_Expert', disc_acc_exp, i_episode)
+                writer.add_scalar('Acc/Disc_Policy', disc_acc_pol, i_episode)
+            if len(memory['actions']) > 0:
+                action_arr = np.array(memory['actions'])
+                action_clip_ratio = (np.abs(action_arr) > 0.98).mean()
+                writer.add_scalar('Policy/ActionClipRatio', action_clip_ratio, 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.save_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="logs/magail", help="TensorBoard log directory")
+    parser.add_argument("--save_dir", type=str, default="models/magail", help="Directory to save model checkpoints")
+    
+    args = parser.parse_args()
+    
+    # Create log dir and save dir
+    os.makedirs(args.log_dir, exist_ok=True)
+    os.makedirs(args.save_dir, exist_ok=True)
+    
+    train(args)