代替点云

.gitignore

@@ -1,57 +0,0 @@
-# 日志文件
-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

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

Algorithm/bc.py

@@ -1,52 +0,0 @@
-"""
-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

@@ -1,136 +0,0 @@
-# 更新日志
-
-## 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

@@ -1,188 +0,0 @@
-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 all static background vehicles once at reset (no show_time filter; same as ExpertReplayEnv)."""
-        for sid, car in self.background_vehicles.items():
-            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.get("valid")
-            v.start_t = car.get("show_time")
-
-    def _update_background_vehicles(self):
-        # Static vehicles are spawned once at init and never removed.
-        pass
-
-    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()
-
-        # Reward shaping for evaluation/rollout monitoring (BC training itself doesn't use env reward).
-        speed_coef = float(self.config.get("reward_speed_coef", 0.05))
-        collision_distance = float(self.config.get("collision_distance", 6.0))
-        collision_penalty = float(self.config.get("collision_penalty", 100.0))
-
-        # Pre-collect all active vehicles (includes background vehicles).
-        active_agents = list(self.engine.agent_manager.active_agents.items())
-
-        rewards = {}
-        infos = {}
-        for aid, vehicle in self.controlled_agents.items():
-            # Speed reward
-            speed = getattr(vehicle, "speed", None)
-            if speed is None:
-                speed = float(np.linalg.norm(vehicle.velocity))
-            r_speed = speed_coef * float(speed)
-
-            # Near-collision penalty (distance-based, simulator-agnostic)
-            min_dist = float("inf")
-            for other_id, other_vehicle in active_agents:
-                if other_id == aid:
-                    continue
-                try:
-                    dist = float(np.linalg.norm(vehicle.position - other_vehicle.position))
-                except Exception:
-                    continue
-                if dist < min_dist:
-                    min_dist = dist
-
-            near_collision = bool(min_dist < collision_distance)
-            r_collision = -collision_penalty if near_collision else 0.0
-
-            rewards[aid] = float(r_speed + r_collision)
-            infos[aid] = {
-                "near_collision": near_collision,
-                "min_dist": (min_dist if np.isfinite(min_dist) else None),
-                "r_speed": float(r_speed),
-                "r_collision": float(r_collision),
-            }
-        dones = {aid: False for aid in self.controlled_agents}
-        dones["__all__"] = self.episode_step >= self.config["horizon"]
-        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

@@ -0,0 +1,109 @@
+#!/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

@@ -1,353 +0,0 @@
-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_all_background_vehicles_at_init()
-
-        # 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_all_background_vehicles_at_init(self):
-        """Spawn all static background vehicles once at reset (no show_time filter; no removal by valid)."""
-        for sid, car in self.background_vehicles.items():
-            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.get('valid')
-            v.start_t = car['show_time']
-
-    def _update_background_vehicles(self):
-        # Static vehicles are spawned once at init and never removed (no spawn/remove by show_time or valid).
-        pass
-
-    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

@@ -1,65 +0,0 @@
-import numpy as np
-import math
-
-class InverseDynamics:
-    def __init__(self, max_steering=0.7, max_acc=8.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

@@ -0,0 +1,62 @@
+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,41 +1,390 @@
+import argparse
 from scenario_env import MultiAgentScenarioEnv
-from Env.simple_idm_policy import ConstantVelocityPolicy
+from simple_idm_policy import ConstantVelocityPolicy
+from replay_policy import ReplayPolicy
 from metadrive.engine.asset_loader import AssetLoader
 
-WAYMO_DATA_DIR = r"/home/huangfukk/MAGAIL4AutoDrive/data"
+WAYMO_DATA_DIR = r"/home/huangfukk/mdsn"
 
-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(AssetLoader.asset_path, "waymo", unix_style=False),
-            "data_directory": AssetLoader.file_path(WAYMO_DATA_DIR, "exp_converted", unix_style=False),
+            "data_directory": AssetLoader.file_path(data_dir, "exp_filtered", unix_style=False),
+            "is_multi_agent": True,
+            "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": 300,
-            "use_render": True,
+            "horizon": horizon,
+            "use_render": render,  # 如果False会完全禁用渲染，避免LANE_FREEWAY错误
             "sequential_seed": True,
             "reactive_traffic": True,
-            "manual_control": True,
+            "manual_control": False,
+            "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)
     )
 
-    obs = env.reset(0
-                    )
-    for step in range(10000):
+    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)
+
+            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")
 
-        if dones["__all__"]:
-            break
+                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,6 +15,7 @@ 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
@@ -22,6 +23,9 @@ 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)
@@ -53,6 +57,15 @@ 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
 
@@ -62,48 +75,180 @@ 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)
 
-    @property
-    def num_controlled_in_scenario(self) -> int:
-        """整个场景中受控车轨迹总数（car_birth_info_list 长度），会在不同 show_time 陆续 spawn。"""
-        return len(getattr(self, "car_birth_info_list", []))
-
     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)
 
-        self.lazy_init()
-        self._reset_global_seed(seed)
-        if self.engine is None:
-            raise ValueError("Broken MetaDrive instance.")
+        # ✅ 关键修复：在每次 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)
                 
-        # 注意：_build_birth_lists_from_traffic() 在 engine.reset() 之前执行，读的是当前 engine 的
-        # current_traffic_data 与 map_manager.current_map。若复用同一 env 连续 reset(0)、reset(1)，
-        # MetaDrive 可能已按 seed 更新了 traffic 为 scenario 1，但 map 仍为 scenario 0（在 engine.reset() 才切图），
-        # 导致 is_on_lane( scenario_1 车位, scenario_0 地图 ) 全为 False → 全部 off_lane → 0 受控车。
-        # 因此多场景时应“每个 scenario 单独建 env”（start_scenario_index=i, num_scenarios=1）再 reset(seed=i)。
-        self.background_vehicles = getattr(self, "background_vehicles", {})
-        self.car_birth_info_list, self.background_vehicles, _obj_to_clean = self._build_birth_lists_from_traffic()
-        for scenario_id in _obj_to_clean:
-            self.engine.traffic_manager.current_traffic_data.pop(scenario_id)
+                # 然后清理对象
+                for agent_id, vehicle in list(self.controlled_agents.items()):
+                    try:
+                        self.engine.clear_objects([vehicle.id])
+                    except:
+                        pass
                 
-        # 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()
         
+        self.lazy_init()
+        self._reset_global_seed(seed)
+
+        if self.engine is None:
+            raise ValueError("Broken MetaDrive instance.")
+
+        # 如果指定了场景ID，修改start_scenario_index
+        if self.config.get("specific_scenario_id") is not None:
+            scenario_id = self.config.get("specific_scenario_id")
+            self.config["start_scenario_index"] = scenario_id
+            if self.config.get("debug", False):
+                self.logger.info(f"Using specific scenario ID: {scenario_id}")
+
+        # ✅ 先初始化引擎和 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
@@ -111,64 +256,129 @@ 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 _build_birth_lists_from_traffic(self):
-        """Build car_birth_info_list and obj_to_clean from current_traffic_data. Override for filtered (lane/static) selection."""
-        _obj_to_clean_this_frame = []
-        car_birth_info_list = []
-        for scenario_id, track in self.engine.traffic_manager.current_traffic_data.items():
-            if scenario_id == self.engine.traffic_manager.sdc_scenario_id:
+    def _is_position_on_lane(self, position, tolerance=None):
+        if tolerance is None:
+            tolerance = self.config.get("lane_tolerance", 3.0)
+
+        # 确保 self.lanes 已初始化
+        if not hasattr(self, 'lanes') or self.lanes is None:
+            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:
-                _obj_to_clean_this_frame.append(scenario_id)
-                valid = track["state"]["valid"]
-                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
-                car_birth_info_list.append({
-                    "id": track["metadata"]["object_id"],
-                    "show_time": first_show,
-                    "begin": (track["state"]["position"][first_show, 0], track["state"]["position"][first_show, 1]),
-                    "heading": track["state"]["heading"][first_show],
-                    "end": (track["state"]["position"][last_show, 0], track["state"]["position"][last_show, 1]),
-                })
-        return car_birth_info_list, {}, _obj_to_clean_this_frame
+
+                lane_vec_normalized = lane_vec / lane_length
+                point_vec = position_2d - lane_start
+                projection = np.dot(point_vec, lane_vec_normalized)
+
+                if 0 <= projection <= lane_length:
+                    closest_point = lane_start + projection * lane_vec_normalized
+                    distance = np.linalg.norm(position_2d - closest_point)
+                    if distance <= tolerance:
+                        return True
+        except Exception as e:
+            if self.config.get("debug", False):
+                self.logger.warning(f"Lane check error: {e}")
+            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
+
             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):
@@ -183,8 +393,20 @@ class MultiAgentScenarioEnv(ScenarioEnv):
                             traffic_light = 0
                     break
 
-            lidar = self.engine.get_sensor("lidar").perceive(num_lasers=80, distance=30, base_vehicle=vehicle,
-                                                             physics_world=self.engine.physics_world.dynamic_world)
+            # 使用最近10辆车的相对位置与相对速度替代原80维LiDAR点云
+            lidar_cloud_points, detected_objects = self.engine.get_sensor("lidar").perceive(
+                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)
@@ -192,30 +414,61 @@ class MultiAgentScenarioEnv(ScenarioEnv):
                                                                                       base_vehicle=vehicle,
                                                                                       physics_world=self.engine.physics_world.static_world)
 
-            obs = (state['position'][:2] + list(state['velocity']) + [state['heading_theta']]
-                   + lidar[0] + side_lidar[0] + lane_line_lidar[0] + [traffic_light]
+            obs = (list(state['position'][:2]) + list(state['velocity']) + [state['heading_theta']]
+                   + nearest_vehicle_info + 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,227 +2,6 @@ import numpy as np
 import torch
 import random
 
-from metadrive.type import MetaDriveType
-
-
-def _static_obbs_overlap(car_a, car_b):
-    """
-    Check if two static vehicle OBBs overlap (2D SAT).
-    car_a, car_b: dicts with "begin" (x, y), "heading" (rad), "length", "width".
-    begin is center; half-extents are length/2, width/2.
-    """
-    def _get_corners(car):
-        cx, cy = car["begin"][0], car["begin"][1]
-        h = float(car["heading"])
-        L2 = float(car["length"]) / 2.0
-        W2 = float(car["width"]) / 2.0
-        ux, uy = np.cos(h), np.sin(h)
-        vx, vy = -np.sin(h), np.cos(h)
-        return np.array([
-            [cx + L2 * ux + W2 * vx, cy + L2 * uy + W2 * vy],
-            [cx + L2 * ux - W2 * vx, cy + L2 * uy - W2 * vy],
-            [cx - L2 * ux - W2 * vx, cy - L2 * uy - W2 * vy],
-            [cx - L2 * ux + W2 * vx, cy - L2 * uy + W2 * vy],
-        ])
-
-    def _get_axes(car):
-        h = float(car["heading"])
-        return [
-            np.array([np.cos(h), np.sin(h)]),
-            np.array([-np.sin(h), np.cos(h)]),
-        ]
-
-    corners_a = _get_corners(car_a)
-    corners_b = _get_corners(car_b)
-    axes = _get_axes(car_a) + _get_axes(car_b)
-
-    for axis in axes:
-        proj_a = corners_a @ axis
-        proj_b = corners_b @ axis
-        min_a, max_a = proj_a.min(), proj_a.max()
-        min_b, max_b = proj_b.min(), proj_b.max()
-        if max_a < min_b or max_b < min_a:
-            return False
-    return True
-
-
-def _deduplicate_background_by_collision(background_vehicles):
-    """
-    Merge static tracks that collide (same physical vehicle). Build collision graph,
-    find connected components, keep one representative per component (min show_time, then scenario_id).
-    """
-    if not background_vehicles:
-        return background_vehicles
-    items = list(background_vehicles.items())
-    n = len(items)
-    # Build adjacency by index
-    parent = list(range(n))
-
-    def find(i):
-        if parent[i] != i:
-            parent[i] = find(parent[i])
-        return parent[i]
-
-    def union(i, j):
-        pi, pj = find(i), find(j)
-        if pi != pj:
-            parent[pi] = pj
-
-    for i in range(n):
-        for j in range(i + 1, n):
-            if _static_obbs_overlap(items[i][1], items[j][1]):
-                union(i, j)
-
-    # Representative per component: index with min (show_time, scenario_id)
-    comp_rep = {}
-    for i in range(n):
-        r = find(i)
-        sid, car = items[i][0], items[i][1]
-        key = (car.get("show_time", 0), sid)
-        if r not in comp_rep or key < comp_rep[r][0]:
-            comp_rep[r] = (key, sid, car)
-
-    return {sid: car for (_, sid, car) in comp_rep.values()}
-
-
-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,
-    deduplicate_static_by_collision=True,
-):
-    """
-    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 deduplicate_static_by_collision and background_vehicles:
-        background_vehicles = _deduplicate_background_by_collision(background_vehicles)
-
-    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

@@ -0,0 +1,184 @@
+#!/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,125 +1,401 @@
-# MAGAIL4AutoDrive
+# MAGAIL4AutoDrive - 多智能体自动驾驶环境
 
-基于 **MetaDrive** 仿真器和 **Waymo Open Motion Dataset** 的自动驾驶多智能体模仿学习（MAGAIL）与行为克隆（BC）训练系统。
+基于 MetaDrive 的多智能体自动驾驶仿真与回放环境，支持 Waymo Open Dataset 的专家轨迹回放和自定义策略仿真。
 
-本项目旨在从真实的 Waymo 驾驶数据中提取专家轨迹，并通过模仿学习（Imitation Learning）训练能够适应复杂交互场景的自动驾驶策略。
+## 📋 目录
 
-## 目录结构
+- [项目简介](#项目简介)
+- [功能特性](#功能特性)
+- [环境要求](#环境要求)
+- [安装步骤](#安装步骤)
+- [快速开始](#快速开始)
+- [使用指南](#使用指南)
+- [项目结构](#项目结构)
+- [配置说明](#配置说明)
+- [常见问题](#常见问题)
 
-```text
-MAGAIL4AutoDrive/
-├── Algorithm/                 # 强化学习与模仿学习算法实现
-│   ├── policy.py              # 基础策略网络 (MLP 等)
-│   ├── ppo.py                 # PPO 算法实现
-│   ├── magail.py              # MAGAIL 算法核心逻辑
-│   ├── disc.py                # 判别器 (Discriminator) 网络
+## 项目简介
+
+MAGAIL4AutoDrive 是一个基于 MetaDrive 0.4.3 的多智能体自动驾驶环境，专为模仿学习（Imitation Learning）和强化学习（Reinforcement Learning）研究设计。项目支持从真实世界数据集（如 Waymo Open Dataset）中加载场景，并提供两种核心运行模式：
+
+- **回放模式（Replay Mode）**：严格按照专家轨迹回放，用于数据可视化和验证
+- **仿真模式（Simulation Mode）**：使用自定义策略控制车辆，用于算法训练和测试
+
+## 功能特性
+
+### 核心功能
+- ✅ **多智能体支持**：同时控制多辆车辆进行协同仿真
+- ✅ **专家轨迹回放**：精确回放 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/`
+### 回放模式（推荐先尝试）
+```
 
-所有默认路径均为相对项目根，便于在不同设备上复用。
 
-## 数据处理流程
+# 使用场景原始时长回放第一个场景
 
-从 Waymo Motion 原始数据到本项目训练用专家 pkl，依次为：
+python run_multiagent_env.py --mode replay --episodes 1 --use_scenario_duration
 
-**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
 ```
 
-**2) ScenarioNet Convert（TFRecord → ScenarioNet 场景库）**  
-需安装 ScenarioNet、MetaDrive 及 TensorFlow 2.11、protobuf 3.20；转换时不用 GPU。
+**重要提示**：回放模式建议始终启用 `--use_scenario_duration`，否则会出现场景播放完后继续运行的问题。
+
+### 仿真模式详解
+
+仿真模式使用自定义策略控制车辆，适合算法开发和测试：
 
 ```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）。
+### 自定义策略
 
-**4) 本项目：生成专家 pkl**  
-使用筛选后的场景目录，生成训练用 pkl 到 `data/training_data`：
+修改 `simple_idm_policy.py` 或创建新的策略类：
 
+```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 自动处理高度，确保车辆贴在地面上。
 
-### 1. 数据准备
-使用 `scripts/generate_expert_data.py` 将 Waymo 数据转换为训练用 `.pkl`，输出到 `data/training_data/`。
+### Q7: 如何导出观测数据？
+**A**: 在 `run_multiagent_env.py` 中添加数据保存逻辑：
+```python
+import pickle
 
-```bash
-python scripts/generate_expert_data.py --data_dir data/exp_filtered --output_dir data/training_data --num_scenarios 100
+obs_data = []
+while True:
+    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 train_bc.py --expert_data_path ./data/training/expert_data_0_50.pkl --epochs 100
-```
-- **可视化**：`python scripts/visualize.py policy --policy_type bc --model_path models/bc/policy_best.pt`
+## 更新日志
 
-### 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`
+### v1.2.0 (2025-10-26)
+- ✅ 修复车辆生成高度问题（车辆悬空）
+- ✅ 添加对象类型过滤功能（车辆/行人/自行车）
+- ✅ 新增命令行参数：`--no_vehicles`、`--no_pedestrians`、`--no_cyclists`
+- ✅ 改进调试信息输出，显示各类型对象统计
+- ✅ 优化位置处理逻辑，只使用 2D 坐标避免高度问题
 
-### 4. 可视化统一入口
-可视化统一使用 `scripts/visualize.py`，子命令：`replay`（场景回放）、`policy`（BC/MAGAIL 策略）、`trajectory`（专家轨迹 2D 动画）。详见 [scripts/README.md](scripts/README.md)。
+### v1.1.0 (2025-10-26)
+- ✅ 修复回放模式超出场景时长问题
+- ✅ 添加场景自动切换功能
+- ✅ 改进 `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/`
+## 贡献指南
 
-### 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 和 Pull Request！
 
-### Algorithm 模块
-- **Algorithm/policy.py**：`StateIndependentPolicy`，BC 使用的 MLP 策略
+### 提交 Issue
+- 请详细描述问题和复现步骤
+- 附上运行日志和错误信息
+- 说明运行环境（OS、Python 版本等）
 
-### scripts 目录
-工具脚本用途与用法见 [scripts/README.md](scripts/README.md)。
+### 提交 PR
+- Fork 本项目
+- 创建特性分支：`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

@@ -1,498 +0,0 @@
-# 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

@@ -1,305 +0,0 @@
-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

@@ -1,157 +0,0 @@
-"""
-统一数据加载：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, *, filter_terminal_last_step: bool = False):
-    """从目录或单个 pkl 加载专家 (obs, acts)，返回 concat 后的 obs_data, act_data。
-
-    Args:
-        expert_data_path: Directory containing pkl files or a single pkl file.
-        filter_terminal_last_step: If True, drop the last (obs, act) pair of each trajectory.
-            This approximates II's \"train only on non-terminal steps\" when the dataset doesn't
-            explicitly store dones.
-    """
-    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 = traj["obs"]
-                        acts = traj["acts"]
-                        if filter_terminal_last_step and len(obs) > 0 and len(acts) > 0:
-                            # Drop last step of each trajectory
-                            obs = obs[:-1]
-                            acts = acts[:-1]
-                        if len(obs) == 0 or len(acts) == 0:
-                            continue
-                        obs_data.append(obs)
-                        act_data.append(acts)
-            elif isinstance(data, dict):
-                if "observations" in data and "actions" in data:
-                    obs = data["observations"]
-                    acts = data["actions"]
-                    if filter_terminal_last_step and len(obs) > 0 and len(acts) > 0:
-                        obs = obs[:-1]
-                        acts = acts[:-1]
-                    if len(obs) == 0 or len(acts) == 0:
-                        continue
-                    obs_data.append(obs)
-                    act_data.append(acts)
-            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
-
-
-def get_expert_scenario_ids(expert_data_path, max_ids=10):
-    """
-    从专家 pkl 中收集出现过的 scenario_id（这些场景在采集时曾有受控车）。
-    用于 eval 时只在这些场景上评估，保证 eval 有受控车。
-    返回排序后的 list，最多 max_ids 个。
-    """
-    if os.path.isdir(expert_data_path):
-        pkl_files = glob.glob(os.path.join(expert_data_path, "*.pkl"))
-    elif os.path.exists(expert_data_path):
-        pkl_files = [expert_data_path]
-    else:
-        return []
-
-    seen = set()
-    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 "scenario_id" in traj:
-                        seen.add(traj["scenario_id"])
-            # dict 格式通常没有 per-trajectory scenario_id，跳过
-        except Exception:
-            continue
-    out = sorted(seen)[:max_ids]
-    return out
-
-
-class MAGAILExpertDataset(Dataset):
-    def __init__(self, data_dir, transform=None, *, filter_terminal_last_step: bool = False):
-        """
-        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)
-            max_i = len(obs)
-            if filter_terminal_last_step and max_i > 0:
-                max_i -= 1
-            for i in range(max_i):
-                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

@@ -1,76 +0,0 @@
-# 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

@@ -1,256 +0,0 @@
-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

@@ -1,47 +0,0 @@
-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

@@ -1,177 +0,0 @@
-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

@@ -1,162 +0,0 @@
-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

@@ -1,18 +0,0 @@
-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

@@ -1,400 +0,0 @@
-"""
-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)
-
-    print(f"Initializing ExpertReplayEnv with data from {data_path}...")
-
-    try:
-        for i in range(args.start_index, args.start_index + num_to_run):
-            print(f"\n--- Playing Scenario {i} ---")
-            # Each scenario uses a fresh env (start_scenario_index=i, num_scenarios=1) so the second
-            # scenario and beyond are fully cleaned and loaded like a single-scenario run.
-            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": i,
-                "num_scenarios": 1,
-                "log_level": 40,
-            }
-            env = ExpertReplayEnv(config=env_config)
-            try:
-                obs = env.reset(seed=i)
-            except Exception as e:
-                print(f"Error resetting scenario {i}: {e}")
-                env.close()
-                continue
-
-            print(f"Scenario loaded. Controlled agents (current): {len(env.controlled_agents)}, total in scenario: {env.num_controlled_in_scenario}")
-
-            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
-            env.close()
-    except KeyboardInterrupt:
-        print("Interrupted by user")
-    except Exception as e:
-        import traceback
-        traceback.print_exc()
-        print(f"Global error: {e}")
-    finally:
-        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 (current): {len(obs_dict)}, total in scenario: {env.num_controlled_in_scenario}")
-            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

@@ -1,199 +0,0 @@
-"""
-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, get_expert_scenario_ids
-
-
-def evaluate_policy(policy, args, device):
-    """在 BCScenarioEnv 中评估策略：仅使用专家数据中出现过的 scenario_id，保证 eval 有受控车。
-    输出与 replay 对齐：agents (current)=reset 时受控车数，total in scenario=该场景受控轨迹总数（car_birth_info_list 长度）。"""
-    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
-
-    scenario_ids = get_expert_scenario_ids(args.expert_data_path, max_ids=5)
-    if not scenario_ids:
-        print("[WARN] No scenario_id in expert pkl, falling back to scenarios [0,1,2]. Eval may have 0 controlled agents.")
-        scenario_ids = [0, 1, 2]
-
-    total_rewards = []
-    total_steps = []
-    collision_episodes = 0
-    horizon = 200
-
-    for idx, scenario_id in enumerate(scenario_ids):
-        env_config = {
-            "data_directory": data_dir,
-            "is_multi_agent": True,
-            "num_controlled_agents": 100,
-            "use_render": False,
-            "sequential_seed": True,
-            "horizon": horizon,
-            "start_scenario_index": scenario_id,
-            "num_scenarios": 1,
-        }
-        env = BCScenarioEnv(env_config, agent2policy=None)
-        try:
-            obs_dict = env.reset(seed=scenario_id)
-        except Exception as e:
-            print(f"  Eval Episode {idx} (scenario {scenario_id}): reset failed: {e}")
-            env.close()
-            continue
-
-        n_controlled = len(env.controlled_agents)
-        n_total_in_scenario = getattr(env, "num_controlled_in_scenario", n_controlled)
-        if n_controlled == 0:
-            print(
-                f"  Eval Episode {idx} (scenario {scenario_id}): 0 controlled agents (total in scenario: {n_total_in_scenario}), skip."
-            )
-            env.close()
-            continue
-
-        episode_reward = 0.0
-        step_count = 0
-        had_near_collision = False
-        dones = {"__all__": False}
-        while not dones["__all__"] and step_count < horizon:
-            step_count += 1
-            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, infos = env.step(action_dict)
-            episode_reward += sum(rewards.values())
-            if infos:
-                for _aid, info in infos.items():
-                    if isinstance(info, dict) and info.get("near_collision", False):
-                        had_near_collision = True
-                        break
-
-        total_rewards.append(episode_reward)
-        total_steps.append(step_count)
-        if had_near_collision:
-            collision_episodes += 1
-        print(
-            f"  Eval Episode {idx} (scenario {scenario_id}): Total Reward {episode_reward:.2f}, steps {step_count}, "
-            f"agents (current): {n_controlled}, total in scenario: {n_total_in_scenario}"
-        )
-        env.close()
-
-    if not total_rewards:
-        print("  No valid eval episodes (all skipped or failed).")
-        return 0.0, 0.0, 0.0
-    avg_reward = float(np.mean(total_rewards))
-    avg_steps = float(np.mean(total_steps)) if total_steps else 0.0
-    collision_rate = float(collision_episodes / max(1, len(total_rewards)))
-    print(
-        f"  Average Evaluation Reward: {avg_reward:.2f} | Mean Episode Length: {avg_steps:.1f} | "
-        f"Collision Rate (near): {collision_rate:.3f}"
-    )
-    return avg_reward, collision_rate, avg_steps
-
-
-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,
-        filter_terminal_last_step=args.filter_terminal_last_step,
-    )
-    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"))
-
-        # Periodic checkpointing (II-style)
-        if args.checkpoint_freq > 0 and (epoch + 1) % args.checkpoint_freq == 0:
-            torch.save(policy.state_dict(), os.path.join(args.save_dir, f"policy_epoch{epoch+1}.pt"))
-
-        if (epoch + 1) % args.eval_freq == 0:
-            eval_reward, eval_collision_rate, eval_mean_steps = evaluate_policy(policy, args, device)
-            writer.add_scalar("Reward/eval", eval_reward, epoch)
-            writer.add_scalar("Eval/collision_rate_near", eval_collision_rate, epoch)
-            writer.add_scalar("Eval/mean_episode_length", eval_mean_steps, 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)
-    parser.add_argument("--checkpoint_freq", type=int, default=50, help="Save policy_epochN.pt every N epochs. Set <=0 to disable.")
-    parser.add_argument(
-        "--filter_terminal_last_step",
-        action="store_true",
-        help="Drop the last (obs, act) pair of each trajectory to approximate training on non-terminal steps (II-style).",
-    )
-    args = parser.parse_args()
-    main(args)

train_magail.py

@@ -1,548 +1,115 @@
 import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import torch.optim as optim
-from torch.distributions import Normal
 import numpy as np
-import os
-import argparse
-import signal
-import sys
-from torch.utils.data import DataLoader
-from dataset.loader import MAGAILExpertDataset
-from Env.bc_env import BCScenarioEnv
+from torch.utils.tensorboard import SummaryWriter
+from Env.scenario_env import MultiAgentScenarioEnv
+from Algorithm.magail import MAGAIL
+from Algorithm.buffer import RolloutBuffer # 假设 Buffer 在这里
+# 假设你有加载专家数据的工具
+# from utils import load_expert_buffer 
 
-# --- Networks ---
+# --- 配置 ---
+CONFIG = {
+    "data_dir": "/home/huangfukk/mdsn",
+    # ... 其他配置 ...
+    "state_dim": 30, # 你的 Observation 维度
+    "action_dim": 2, # [steering, throttle]
+    "rollout_length": 2048, # Buffer 大小
+}
 
-class Actor(nn.Module):
-    def __init__(self, state_dim, action_dim, hidden_dim=256):
-        super(Actor, self).__init__()
-        self.net = nn.Sequential(
-            nn.Linear(state_dim, hidden_dim),
-            nn.Tanh(),
-            nn.Linear(hidden_dim, hidden_dim),
-            nn.Tanh(),
-        )
-        self.mu_head = nn.Linear(hidden_dim, action_dim)
-        self.log_std_head = nn.Parameter(torch.zeros(1, action_dim))
+def train():
+    # 1. 环境
+    env = MultiAgentScenarioEnv(config={...}, agent2policy=None)
     
-    def forward(self, state):
-        x = self.net(state)
-        mu = torch.tanh(self.mu_head(x)) # Action range [-1, 1]
-        if mu.dim() == 1:
-            mu = mu.unsqueeze(0) # Handle single sample
-        log_std = self.log_std_head.expand_as(mu)
-        std = torch.exp(log_std)
-        dist = Normal(mu, std)
-        return dist
+    # 2. 专家 Buffer (伪代码)
+    # expert_buffer = RolloutBuffer(...)
+    # expert_buffer.load(...) 
+    # 必须保证 expert_buffer.sample() 返回 (state, next_state) 供 Discriminator 训练
     
-class Critic(nn.Module):
-    def __init__(self, state_dim, hidden_dim=256):
-        super(Critic, self).__init__()
-        self.net = nn.Sequential(
-            nn.Linear(state_dim, hidden_dim),
-            nn.Tanh(),
-            nn.Linear(hidden_dim, hidden_dim),
-            nn.Tanh(),
-            nn.Linear(hidden_dim, 1)
+    # 3. 初始化 MAGAIL
+    magail = MAGAIL(
+        buffer_exp=expert_buffer, # 传入专家 buffer
+        input_dim=(CONFIG["state_dim"],),
+        action_shape=(CONFIG["action_dim"],), # PPO 初始化可能需要 action_shape，原代码好像没传？
+        device=torch.device("cuda"),
+        rollout_length=CONFIG["rollout_length"]
     )
     
-    def forward(self, state):
-        return self.net(state)
+    writer = SummaryWriter("./logs")
     
-class Discriminator(nn.Module):
-    def __init__(self, state_dim, action_dim, hidden_dim=256):
-        super(Discriminator, self).__init__()
-        self.net = nn.Sequential(
-            nn.Linear(state_dim + action_dim, hidden_dim),
-            nn.Tanh(),
-            nn.Linear(hidden_dim, hidden_dim),
-            nn.Tanh(),
-            nn.Linear(hidden_dim, 1),
-            nn.Sigmoid() 
-        )
+    # 4. 训练循环
+    total_steps = 0
+    obs_dict = env.reset()
     
-    def forward(self, state, action):
-        x = torch.cat([state, action], dim=-1)
-        return self.net(x)
+    # 将 Dict Obs 转为 Array: (N_agents, Obs_dim)
+    # 假设所有 Agent 的 Obs 维度相同
+    agents = list(obs_dict.keys())
+    current_obs = np.stack([obs_dict[a] for a in agents]) 
     
-# --- PPO Algorithm ---
+    # 如果需要 state_gail，假设它就是 obs
+    current_state_gail = current_obs.copy()
     
-class PPO:
-    def __init__(self, state_dim, action_dim, lr=3e-4, gamma=0.99, eps_clip=0.2, K_epochs=10):
-        self.actor = Actor(state_dim, action_dim).cuda()
-        self.critic = Critic(state_dim).cuda()
-        self.optimizer_actor = optim.Adam(self.actor.parameters(), lr=lr)
-        self.optimizer_critic = optim.Adam(self.critic.parameters(), lr=lr)
+    while total_steps < 1e7:
+        # --- 收集数据 (Rollout) ---
+        # PPO 的 buffer 长度是 rollout_length
+        # 我们可以在这里调用 magail.step 来自动处理 explore 和 buffer append
         
-        self.gamma = gamma
-        self.eps_clip = eps_clip
-        self.K_epochs = K_epochs
-        self.mse_loss = nn.MSELoss()
+        # 注意：magail.step 内部调用了 env.step，这可能不适用于 Multi-Agent 环境返回 Dict 的情况
+        # 原 PPO.step 似乎是为 Single-Agent 或 VectorEnv 设计的
+        # 这里我们需要手动写 Rollout 循环或者修改 PPO.step 以适配 Dict 返回值
         
-    def _log_prob_from_dist(self, dist, pre_tanh_action):
-        # Tanh-squashed Gaussian log-prob with correction term.
-        log_prob = dist.log_prob(pre_tanh_action)
-        correction = torch.log(1 - torch.tanh(pre_tanh_action) ** 2 + 1e-6)
-        return (log_prob - correction).sum(dim=-1)
+        # --- 方案：手动 Rollout 适配 Multi-Agent ---
+        for _ in range(CONFIG["rollout_length"]):
+            # 1. 决策
+            actions_list, log_pis_list = magail.explore(current_obs)
             
-    def select_action(self, state):
+            # 2. 拼装 Action Dict
+            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():
-            state = torch.FloatTensor(state).cuda()
-            dist = self.actor(state)
-            pre_tanh_action = dist.sample()
-            action = torch.tanh(pre_tanh_action)
-            action_logprob = self._log_prob_from_dist(dist, pre_tanh_action)
-        return (
-            action.cpu().numpy(),
-            action_logprob.cpu().numpy(),
-            pre_tanh_action.cpu().numpy()
+                # 重新计算一遍或者在 explore 里返回
+                # 这里 magail.actor(state) 返回的是 mean (deterministic action)
+                means = magail.actor(torch.tensor(current_obs, device=magail.device, dtype=torch.float)).cpu().numpy()
+                stds = magail.actor.log_stds.exp().detach().cpu().numpy()
+                # 如果是 StateIndependentPolicy，std 可能是共享的参数，维度需要广播
+                if stds.shape[0] != len(agents):
+                     stds = np.repeat(stds, len(agents), axis=0)
+
+            magail.buffer.append(
+                current_obs, current_state_gail, actions_list, 
+                rewards, dones, terminated, log_pis_list,
+                next_obs, next_obs, # next_state_gail = next_obs
+                means, stds
             )
             
-    def update(self, memory):
-        # Convert memory to tensors
-        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()
+            current_obs = next_obs
+            current_state_gail = next_obs # 更新 gail state
+            total_steps += len(agents) # 步数增加 N
             
-        # Monte Carlo estimate of state rewards (or GAE if implemented, simplistic here)
-        # Usually for PPO we use GAE. Let's do a simple discounted return for now or bootstrapping.
-        # Let's use bootstrapping from critic for returns.
+            # 处理 Done
+            if all(dones):
+                obs_dict = env.reset()
+                current_obs = np.stack([obs_dict[a] for a in agents])
+                current_state_gail = current_obs
         
-        returns = []
-        discounted_reward = 0
-        # This simple loop assumes full episode or consistent batch. 
-        # For multi-agent disjoint steps, bootstrapping is better.
-        # But let's calculate advantage using GAE for stability.
+        # --- 更新 ---
+        # 当 Buffer 满时 (rollout_length)，调用 update
+        magail.update(writer, total_steps)
         
-        values = self.critic(states).detach()
-        next_values = self.critic(next_states).detach()
+        # 保存
+        if total_steps % 10000 == 0:
+            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)