代替点云

.gitignore

@@ -1,3 +0,0 @@
-# 日志文件
-Env/logs/
-*.log

.vscode/settings.json

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

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/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,527 +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 = {} # Vehicles that exist but are static/background
-        
-        # Helper function to check if a position is on a valid lane
-        def is_on_lane(pos, map_manager, threshold=2.0):
-            # Check if point is close to any lane in the road network
-            # This can be expensive if checked for every point, so we check sample points
-            # or rely on lane index if available.
-            # Waymo tracks don't have lane index, just positions.
-            # We can use map.road_network.get_closest_lane_index(pos)
-            if map_manager is None or map_manager.current_map is None:
-                return True # If no map, assume valid
-                
-            try:
-                # Use a larger search radius to catch slightly offset lanes
-                lane, lane_index = map_manager.current_map.road_network.get_closest_lane_index(pos, return_lane=True)
-                if lane is None:
-                    return False
-                
-                # Check lateral distance
-                long, lat = lane.local_coordinates(pos)
-                width = lane.width
-                # Allow being slightly off-lane (e.g. changing lanes)
-                # But parking lots are usually far from defined lanes in Waymo converted maps
-                if abs(lat) <= (width / 2 + threshold):
-                    return True
-                return False
-            except:
-                return False
-
-        # --- MODIFIED SECTION START ---
-        # Capture expert tracks before they are cleaned
-        self.expert_tracks = {}
-        # Capture SDC track for ego replay (MetaDrive default agent)
-        self.sdc_track = None
-        self.sdc_vehicle = 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)
-        _obj_to_clean_this_frame = []
-        self.car_birth_info_list = []
-        
-        # Pre-filter: Check tracks against map AND check for static vehicles
-        
-        for scenario_id, track in self.engine.traffic_manager.current_traffic_data.items():
-            if scenario_id == self.engine.traffic_manager.sdc_scenario_id:
-                continue
-            else:
-                if track["type"] == MetaDriveType.VEHICLE:
-                    _obj_to_clean_this_frame.append(scenario_id)
-                    
-                    valid = track['state']['valid']
-                    if not valid.any():
-                        continue
-                        
-                    first_show = np.argmax(valid)
-                    last_show = len(valid) - 1 - np.argmax(valid[::-1])
-                    mid_show = (first_show + last_show) // 2
-                    
-                    # 1. Lane check (existing logic)
-                    points_to_check = [first_show, mid_show, last_show]
-                    on_road_count = 0
-                    is_valid_track = True
-                    start_pos = track['state']['position'][first_show]
-                    if not is_on_lane(start_pos, self.engine.map_manager, threshold=5.0): # 5m tolerance
-                        mid_pos = track['state']['position'][mid_show]
-                        if not is_on_lane(mid_pos, self.engine.map_manager, threshold=5.0):
-                            is_valid_track = False
-                    
-                    # 2. Static check
-                    # Calculate total displacement and max speed
-                    positions = track['state']['position'][valid.astype(bool)]
-                    velocities = track['state']['velocity'][valid.astype(bool)]
-                    
-                    total_displacement = 0
-                    max_speed = 0
-                    if len(positions) > 1:
-                        total_displacement = np.linalg.norm(positions[-1] - positions[0])
-                        max_speed = np.max(np.linalg.norm(velocities, axis=1))
-                    
-                    is_static = False
-                    if total_displacement < 5.0 and max_speed < 1.0: # Relaxed threshold: <5m move and <1m/s
-                        is_static = True
-                        
-                    # Decision logic:
-                    # - If off-road AND static: Skip completely (don't even spawn as background)
-                    # - If off-road but moving: Maybe keep? Or skip? Usually off-road moving is weird, skip.
-                    # - If on-road but static: Spawn as BACKGROUND (visible but not controlled agent)
-                    # - If on-road and moving: Spawn as CONTROLLED agent
-                    
-                    if not is_valid_track:
-                        # Skip off-road vehicles entirely (both static and moving off-road)
-                        continue
-                        
-                    if is_static:
-                        # Add to background list, but NOT to car_birth_info_list (which is for controlled agents)
-                        # We need a way to spawn them. Let's add a separate list.
-                         self.background_vehicles[scenario_id] = {
-                            '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]),
-                            'scenario_id': scenario_id,
-                            'length': track['state']['length'][first_show],
-                            'width': track['state']['width'][first_show],
-                            'valid': valid # Need validity to know when to show/hide
-                        }
-                         continue # Do not add to controlled list
-
-                    # Store the full track for replay (only for controlled agents)
-                    self.expert_tracks[scenario_id] = track
-                    
-                    self.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]),
-                        'scenario_id': scenario_id, # Keep track of original ID to lookup tracks
-                        'length': track['state']['length'][first_show],
-                        'width': track['state']['width'][first_show]
-                    })
-
-        for scenario_id in _obj_to_clean_this_frame:
-            self.engine.traffic_manager.current_traffic_data.pop(scenario_id)
-        # --- MODIFIED SECTION END ---
-
-        self.engine.reset()
-        self.reset_sensors()
-        self.engine.taskMgr.step()
-
-        self.lanes = self.engine.map_manager.current_map.road_network.graph
-
-        if self.top_down_renderer is not None:
-            self.top_down_renderer.clear()
-            self.engine.top_down_renderer = None
-
-        self.dones = {}
-        self.episode_rewards = defaultdict(float)
-        self.episode_lengths = defaultdict(int)
-
-        self.controlled_agents.clear()
-        self.controlled_agent_ids.clear()
-        
-        # We skip calling super().reset() to avoid double reset
-        # But we need to ensure ScenarioEnv-specific setup is done if any.
-        # ScenarioEnv.reset() basically does engine.reset() and some cleanup.
-        # We covered most of it.
-        
-        self._spawn_controlled_agents()
-        self._spawn_background_vehicles() # Initial spawn for background
-
-        # Ensure SDC/ego is moved to the correct initial expert state.
-        if self.replay_sdc:
-            self.sdc_vehicle = self.engine.agent_manager.active_agents.get(self.sdc_agent_id, None)
-            if self.sdc_vehicle is not None and self.sdc_track is not None:
-                valid = self.sdc_track["state"]["valid"]
-                t0 = int(np.argmax(valid)) if valid.any() else 0
-                pos0 = self.sdc_track["state"]["position"][t0]
-                heading0 = self.sdc_track["state"]["heading"][t0]
-                vel0 = self.sdc_track["state"]["velocity"][t0]
-                self.sdc_vehicle.set_position(pos0)
-                self.sdc_vehicle.set_heading_theta(heading0)
-                self.sdc_vehicle.set_velocity(vel0)
-        
-        return self._get_all_obs()
-
-    def _spawn_background_vehicles(self):
-        # Spawn static/background vehicles
-        # Since they are static, we might just spawn them once if their show_time is 0
-        # But Waymo tracks have valid bits, they might appear/disappear.
-        # For optimization, if they are truly static (never move), we just spawn them when show_time matches.
-        
-        # We need to track spawned background vehicles to remove them if they become invalid?
-        # Since we defined them as "static", they probably stay put. 
-        # But validity might change (e.g. late spawn).
-        
-        # For simplicity in this step, let's just iterate and spawn if time matches
-        for sid, car in self.background_vehicles.items():
-            if car['show_time'] == self.round:
-                # Spawn as a Traffic Vehicle (not PolicyVehicle), or just a static object?
-                # Using DefaultVehicle is fine, but don't add to controlled_agents
-                
-                # Check duplication
-                bg_id = f"bg_{car['id']}"
-                # if bg_id in self.engine.obj_to_id: # obj_to_id might not be available in all versions
-                if bg_id in self.engine.agent_manager.active_agents:
-                    continue
-                    
-                vehicle_config = {}
-                if 'length' in car and 'width' in car:
-                     vehicle_config = {
-                         "length": car['length'],
-                         "width": car['width']
-                     }
-
-                v = self.engine.spawn_object(
-                    DefaultVehicle,
-                    name=bg_id,
-                    vehicle_config=vehicle_config,
-                    position=car['begin'],
-                    heading=car['heading']
-                )
-                
-                # Set color to grey/dark to indicate background
-                v.set_velocity([0, 0])
-                # Maybe set color? MetaDrive vehicles random color.
-                # v.set_color(...) if supported
-                
-                # Register as an active object but NOT controlled agent
-                # The engine manages it. 
-                # CRITICAL: We need it in self.engine.agent_manager.active_agents for Observation?
-                # If we want it to be seen by Lidar/Observation, it needs to be an "agent" or "traffic".
-                # DefaultVehicle spawned this way is just an object.
-                # We should add it to traffic manager? Or just leave it as object?
-                # MultiAgentScenarioEnv._get_all_obs iterates self.engine.agent_manager.active_agents
-                
-                # If we want it in observation, we must add it to active_agents OR iterate over all objects.
-                # Adding to active_agents is easier for compatibility.
-                self.engine.agent_manager.active_agents[bg_id] = v
-                
-                # Store valid mask to remove it later if needed?
-                v.valid_mask = car['valid']
-                v.start_t = car['show_time']
-
-    def _update_background_vehicles(self):
-        # Remove background vehicles if they become invalid
-        # Or spawn new ones
-        self._spawn_background_vehicles()
-        
-        # Check validity for existing
-        to_remove = []
-        for aid, v in self.engine.agent_manager.active_agents.items():
-            if aid.startswith("bg_"):
-                # Check validity
-                if hasattr(v, 'valid_mask'):
-                    curr_step = self.round
-                    if curr_step >= len(v.valid_mask) or not v.valid_mask[curr_step]:
-                        to_remove.append(aid)
-        
-        for aid in to_remove:
-             self.engine.agent_manager.active_agents.pop(aid, None)
-             # if aid in self.engine.obj_to_id:
-             #    self.engine.clear_objects([self.engine.obj_to_id[aid]])
-             # Instead, we should find the object by ID and clear it.
-             # Since we don't track obj directly, we can't easily clear it without obj ref.
-             # Wait, active_agents stores the vehicle object.
-             # So we can just clear that object.
-             pass 
-             
-        # Re-iterate to clear objects properly
-        for aid in to_remove:
-             # We need to find the vehicle object to clear it.
-             # But we popped it from active_agents.
-             # Wait, we should get it before pop.
-             pass
-             
-    def _update_background_vehicles(self):
-        # Remove background vehicles if they become invalid
-        # Or spawn new ones
-        self._spawn_background_vehicles()
-        
-        # Check validity for existing
-        to_remove = []
-        objects_to_clear = []
-        
-        for aid, v in self.engine.agent_manager.active_agents.items():
-            if aid.startswith("bg_"):
-                # Check validity
-                if hasattr(v, 'valid_mask'):
-                    curr_step = self.round
-                    if curr_step >= len(v.valid_mask) or not v.valid_mask[curr_step]:
-                        to_remove.append(aid)
-                        objects_to_clear.append(v)
-        
-        for aid in to_remove:
-             self.engine.agent_manager.active_agents.pop(aid, None)
-             
-        if objects_to_clear:
-            self.engine.clear_objects(objects_to_clear)
-
-    def _spawn_controlled_agents(self):
-        for car in self.car_birth_info_list:
-            if car['show_time'] == self.round:
-                agent_id = f"controlled_{car['id']}"
-                
-                # Check if we already have this agent (shouldn't happen with unique IDs but safety check)
-                if agent_id in self.controlled_agents:
-                    continue
-                
-                # Handling ID flickering / merging
-                # If this ID is new, check if there's an existing agent very close to its start position
-                # that just disappeared? (Not implemented here, complex logic)
-                # But we can check if there's an overlap with existing agents?
-                # For now, just spawn.
-                
-                # Read vehicle type/size if available
-                vehicle_config = {}
-                if 'length' in car and 'width' in car:
-                     vehicle_config = {
-                         "length": car['length'],
-                         "width": car['width']
-                     }
-
-                vehicle = self.engine.spawn_object(
-                    PolicyVehicle,
-                    name=agent_id,
-                    vehicle_config=vehicle_config,
-                    position=car['begin'],
-                    heading=car['heading']
-                )
-                vehicle.reset(position=car['begin'], heading=car['heading'])
-                
-                # We don't set policy or destination in the same way, or maybe we do for compatibility
-                vehicle.set_destination(car['end'])
-                
-                # Store extra info for replay
-                vehicle.expert_track = self.expert_tracks[car['scenario_id']]
-                vehicle.original_id = car['id']
-                
-                self.controlled_agents[agent_id] = vehicle
-                self.controlled_agent_ids.append(agent_id)
-                self.engine.agent_manager.active_agents[agent_id] = vehicle
-
-    def step(self, action_dict=None):
-        # We ignore input action_dict for the purpose of controlling agents
-        # Instead, we calculate what the action *should* be
-        
-        self.round += 1
-        expert_actions = {}
-        
-        # 1. Update state of all controlled agents to the current timestep (self.round)
-        #    and compute action from (self.round-1) to (self.round).
-        #    Wait, usually step() moves T -> T+1.
-        #    Current state is T. We want to move to T+1.
-        #    So we need state at T and T+1.
-        
-        # Identify agents that are done (valid=0 at T+1 or T+1 >= length)
-        agents_to_remove = []
-
-        # Update SDC/ego first (otherwise it will stay still with default action=0)
-        if self.replay_sdc and self.sdc_vehicle is not None and self.sdc_track is not None:
-            next_step = self.round
-            curr_step = self.round - 1
-            if next_step < len(self.sdc_track["state"]["position"]) and self.sdc_track["state"]["valid"][next_step]:
-                curr_state = {
-                    "position": self.sdc_track["state"]["position"][curr_step],
-                    "heading": self.sdc_track["state"]["heading"][curr_step],
-                    "velocity": self.sdc_track["state"]["velocity"][curr_step],
-                }
-                next_state = {
-                    "position": self.sdc_track["state"]["position"][next_step],
-                    "heading": self.sdc_track["state"]["heading"][next_step],
-                    "velocity": self.sdc_track["state"]["velocity"][next_step],
-                }
-                action, _ = self.inverse_dynamics.compute_action(curr_state, next_state, dt=0.1)
-                expert_actions[self.sdc_agent_id] = action
-                self.sdc_vehicle.set_position(next_state["position"])
-                self.sdc_vehicle.set_heading_theta(next_state["heading"])
-                self.sdc_vehicle.set_velocity(next_state["velocity"])
-                self.sdc_vehicle.last_expert_action = action
-
-        for agent_id, vehicle in self.controlled_agents.items():
-            track = vehicle.expert_track
-            # current_step = self.round - 1 # Since we incremented at start
-            # But vehicle is currently at state corresponding to self.round - 1.
-            # We want to move it to self.round.
-            
-            # Check bounds
-            next_step = self.round
-            curr_step = self.round - 1
-            
-            if next_step >= len(track['state']['position']):
-                agents_to_remove.append(agent_id)
-                continue
-                
-            valid = track['state']['valid'][next_step]
-            if not valid:
-                agents_to_remove.append(agent_id)
-                continue
-            
-            # Get states
-            curr_pos = track['state']['position'][curr_step]
-            next_pos = track['state']['position'][next_step]
-            curr_heading = track['state']['heading'][curr_step]
-            next_heading = track['state']['heading'][next_step]
-            curr_vel = track['state']['velocity'][curr_step]
-            next_vel = track['state']['velocity'][next_step]
-            
-            # Prepare state dicts for Inverse Dynamics
-            curr_state = {
-                'position': curr_pos,
-                'heading': curr_heading,
-                'velocity': curr_vel
-            }
-            next_state = {
-                'position': next_pos,
-                'heading': next_heading,
-                'velocity': next_vel
-            }
-            
-            # Calculate action
-            action, raw_info = self.inverse_dynamics.compute_action(curr_state, next_state, dt=0.1) # Waymo is 10Hz?
-            expert_actions[agent_id] = action
-            
-            # Force update vehicle state
-            vehicle.set_position(next_pos)
-            vehicle.set_heading_theta(next_heading)
-            vehicle.set_velocity(next_vel)
-            
-            # Also record this action in the vehicle for later retrieval if needed
-            vehicle.last_expert_action = action
-            
-        # Remove finished agents
-        for agent_id in agents_to_remove:
-            vehicle = self.controlled_agents[agent_id]
-            self.controlled_agents.pop(agent_id)
-            self.controlled_agent_ids.remove(agent_id)
-            self.engine.agent_manager.active_agents.pop(agent_id, None)
-            
-            self.engine.clear_objects([vehicle.id])
-
-        # Step physics world to update sensors/collision detection
-        # We don't need full integration, but we need to update the physics world state
-        self.engine.taskMgr.step()
-        
-        # Spawn new agents for this turn
-        self._spawn_controlled_agents()
-        self._update_background_vehicles()
-        
-        # Get observations
-        obs = self._get_all_obs()
-        
-        rewards = {aid: 0.0 for aid in self.controlled_agents}
-        dones = {aid: False for aid in self.controlled_agents}
-        dones["__all__"] = (self.round >= self.config["horizon"]) or (len(self.controlled_agents) == 0 and self.round > 190) # Waymo scenarios are usually ~198 steps (20s @ 10Hz) or 90 steps (9s)
-        
-        infos = {aid: {"expert_action": expert_actions.get(aid, np.zeros(2))} for aid in self.controlled_agents}
-        
-        return obs, rewards, dones, infos
-
-    def _get_all_obs(self):
-        # Implement custom observation: 30m range, 10 nearest vehicles
-        obs_dict = {}
-        
-        for agent_id, vehicle in self.controlled_agents.items():
-            # 1. Ego State
-            ego_state = [
-                vehicle.position[0], vehicle.position[1],
-                vehicle.velocity[0], vehicle.velocity[1],
-                vehicle.heading_theta
-            ]
-            
-            # 2. Neighbors
-            neighbors = []
-            # Iterate through all vehicles in the engine
-            candidates = []
-            for other_id, other_vehicle in self.engine.agent_manager.active_agents.items():
-                if other_id == agent_id:
-                    continue
-                
-                dist = np.linalg.norm(vehicle.position - other_vehicle.position)
-                if dist < 30.0:
-                    candidates.append((dist, other_vehicle))
-            
-            # Sort by distance
-            candidates.sort(key=lambda x: x[0])
-            
-            # Take top 10
-            top_10 = candidates[:10]
-            
-            neighbor_feats = []
-            for _, neighbor in top_10:
-                neighbor_feats.extend([
-                    neighbor.position[0] - vehicle.position[0], # Relative pos
-                    neighbor.position[1] - vehicle.position[1],
-                    neighbor.velocity[0], # Absolute vel? or Relative? Usually relative in MultiAgent
-                    neighbor.velocity[1]
-                ])
-                
-            # Pad if < 10
-            missing = 10 - len(top_10)
-            if missing > 0:
-                neighbor_feats.extend([0.0] * (4 * missing))
-                
-            # Flatten
-            obs = np.array(ego_state + neighbor_feats, dtype=np.float32)
-            obs_dict[agent_id] = obs
-            
-        return obs_dict

Env/inverse_dynamics.py

@@ -1,65 +0,0 @@
-import numpy as np
-import math
-
-class InverseDynamics:
-    def __init__(self, max_steering=0.7, max_acc=15.0, length=4.5):
-        """
-        :param max_steering: Max steering angle in radians (approx 40 degrees)
-        :param max_acc: Max acceleration in m/s^2
-        :param length: Vehicle length in meters (Waymo default approx 4.5m)
-        """
-        self.max_steering = max_steering
-        self.max_acc = max_acc
-        self.wheelbase = 0.7 * length  # Approximation as per request
-
-    def compute_action(self, current_state, next_state, dt=0.1):
-        """
-        Compute action [steering, acceleration] from current and next state.
-        State format: dictionary or object with keys/attrs: position (x, y), heading, velocity (v_x, v_y)
-        or numpy array [x, y, vx, vy, heading]
-        
-        Using Bicycle Model:
-        delta = arctan(L * theta_dot / v)
-        acc = (v_next - v_curr) / dt
-        """
-        
-        # Extract state
-        # Assume state is dict-like for now, can adapt if needed
-        # We need: velocity (scalar), heading
-        
-        # Helper to get speed
-        def get_speed(vel):
-            return np.linalg.norm(vel)
-
-        v_curr = get_speed(current_state['velocity'])
-        v_next = get_speed(next_state['velocity'])
-        
-        # 1. Acceleration (longitudinal)
-        acc = (v_next - v_curr) / dt
-        
-        # 2. Steering (lateral)
-        # theta_dot = (theta_next - theta_curr) / dt
-        theta_curr = current_state['heading']
-        theta_next = next_state['heading']
-        
-        # Handle angle wrapping [-pi, pi]
-        diff_theta = theta_next - theta_curr
-        if diff_theta > np.pi:
-            diff_theta -= 2 * np.pi
-        elif diff_theta < -np.pi:
-            diff_theta += 2 * np.pi
-            
-        theta_dot = diff_theta / dt
-        
-        # Avoid division by zero for stationary vehicles
-        if v_curr < 0.1:
-            steering = 0.0
-        else:
-            # delta = arctan(L * theta_dot / v)
-            steering = np.arctan(self.wheelbase * theta_dot / v_curr)
-
-        # Normalize actions to [-1, 1]
-        norm_acc = np.clip(acc / self.max_acc, -1.0, 1.0)
-        norm_steering = np.clip(steering / self.max_steering, -1.0, 1.0)
-        
-        return np.array([norm_steering, norm_acc]), {'raw_acc': acc, 'raw_steering': steering}

Env/replay_policy.py

@@ -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,40 +1,389 @@
+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):
-        actions = {
-            aid: env.controlled_agents[aid].policy.act()
-            for aid in env.controlled_agents
-        }
+    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}")
 
-        obs, rewards, dones, infos = env.step(actions)
-        env.render(mode="topdown")
+            # ✅ 如果不是指定场景，使用循环的场景索引
+            if scenario_id is not None:
+                seed = scenario_id
+            else:
+                seed = episode % num_scenarios
+            obs = env.reset(seed=seed)
 
-        if dones["__all__"]:
-            break
+            actual_horizon = env.config["horizon"]
+            print(f"初始化完成:")
+            print(f"  可控车辆数: {len(env.controlled_agents)}")
+            print(f"  场景时长: {env.scenario_max_duration} 步")
+            print(f"  实际Horizon: {actual_horizon} 步")
 
-    env.close()
+            step_count = 0
+            total_reward = 0.0
+
+            while True:
+                # 使用策略生成动作
+                actions = {
+                    aid: env.controlled_agents[aid].policy.act()
+                    for aid in env.controlled_agents
+                }
+
+                obs, rewards, dones, infos = env.step(actions)
+
+                if render:
+                    env.render(mode="topdown")
+
+                step_count += 1
+                total_reward += sum(rewards.values())
+
+                # 每50步打印一次状态
+                if step_count % 50 == 0:
+                    print(f"Step {step_count}: {len(env.controlled_agents)} 辆活跃车辆")
+
+                if dones["__all__"]:
+                    print(f"\n回合结束统计:")
+                    print(f"  总步数: {step_count}")
+                    print(f"  总奖励: {total_reward:.2f}")
+                    break
+    finally:
+        env.close()
+
+    print("\n" + "=" * 50)
+    print("仿真完成！")
+    print("=" * 50)
+
+
+def main():
+    parser = argparse.ArgumentParser(description="MetaDrive 多智能体环境运行脚本")
+
+    parser.add_argument(
+        "--mode",
+        type=str,
+        choices=["replay", "simulation"],
+        default="simulation",
+        help="运行模式: replay=专家轨迹回放, simulation=策略仿真 (默认: simulation)"
+    )
+
+    parser.add_argument(
+        "--data_dir",
+        type=str,
+        default=WAYMO_DATA_DIR,
+        help=f"数据目录路径 (默认: {WAYMO_DATA_DIR})"
+    )
+
+    parser.add_argument(
+        "--episodes",
+        type=int,
+        default=1,
+        help="运行回合数 (默认: 1)"
+    )
+
+    parser.add_argument(
+        "--horizon",
+        type=int,
+        default=300,
+        help="每回合最大步数 (默认: 300，如果启用 --use_scenario_duration 则自动设置)"
+    )
+
+    parser.add_argument(
+        "--no_render",
+        action="store_true",
+        help="禁用渲染（加速运行）"
+    )
+
+    parser.add_argument(
+        "--debug",
+        action="store_true",
+        help="启用调试模式（显示详细日志）"
+    )
+
+    parser.add_argument(
+        "--scenario_id",
+        type=int,
+        default=None,
+        help="指定场景ID（可选，如指定则只运行该场景）"
+    )
+
+    parser.add_argument(
+        "--use_scenario_duration",
+        action="store_true",
+        help="使用场景原始时长作为horizon（自动停止）"
+    )
+
+    parser.add_argument(
+        "--no_vehicles",
+        action="store_true",
+        help="禁止生成车辆"
+    )
+
+    parser.add_argument(
+        "--no_pedestrians",
+        action="store_true",
+        help="禁止生成行人"
+    )
+
+    parser.add_argument(
+        "--no_cyclists",
+        action="store_true",
+        help="禁止生成自行车"
+    )
+
+    args = parser.parse_args()
+
+    if args.mode == "replay":
+        run_replay_mode(
+            data_dir=args.data_dir,
+            num_episodes=args.episodes,
+            horizon=args.horizon,
+            render=not args.no_render,
+            debug=args.debug,
+            scenario_id=args.scenario_id,
+            use_scenario_duration=args.use_scenario_duration,
+            spawn_vehicles=not args.no_vehicles,
+            spawn_pedestrians=not args.no_pedestrians,
+            spawn_cyclists=not args.no_cyclists
+        )
+    else:
+        run_simulation_mode(
+            data_dir=args.data_dir,
+            num_episodes=args.episodes,
+            horizon=args.horizon,
+            render=not args.no_render,
+            debug=args.debug,
+            scenario_id=args.scenario_id,
+            use_scenario_duration=args.use_scenario_duration,
+            spawn_vehicles=not args.no_vehicles,
+            spawn_pedestrians=not args.no_pedestrians,
+            spawn_cyclists=not args.no_cyclists
+        )
 
 
 if __name__ == "__main__":

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,50 +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)
 
     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)
+        log_level = self.config.get("log_level", logging.DEBUG if self.config.get("debug", False) else logging.INFO)
+        set_log_level(log_level)
 
+        # ✅ 关键修复：在每次 reset 前清理所有自定义生成的对象
+        if hasattr(self, 'engine') and self.engine is not None:
+            if hasattr(self, 'controlled_agents') and self.controlled_agents:
+                # 先从 agent_manager 中移除
+                if hasattr(self.engine, 'agent_manager'):
+                    for agent_id in list(self.controlled_agents.keys()):
+                        if agent_id in self.engine.agent_manager.active_agents:
+                            self.engine.agent_manager.active_agents.pop(agent_id)
+                
+                # 然后清理对象
+                for agent_id, vehicle in list(self.controlled_agents.items()):
+                    try:
+                        self.engine.clear_objects([vehicle.id])
+                    except:
+                        pass
+                
+                self.controlled_agents.clear()
+                self.controlled_agent_ids.clear()
+        
         self.lazy_init()
         self._reset_global_seed(seed)
+
         if self.engine is None:
             raise ValueError("Broken MetaDrive instance.")
 
-        # 记录专家数据中每辆车的位置，接着全部清除，只保留位置等信息，用于后续生成
-        _obj_to_clean_this_frame = []
-        self.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:
-                continue
-            else:
-                if track["type"] == MetaDriveType.VEHICLE:
-                    _obj_to_clean_this_frame.append(scenario_id)
-                    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
-                    # id，出现时间，出生点坐标，出生朝向，目的地
-                    self.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])
-                    })
-
-        for scenario_id in _obj_to_clean_this_frame:
-            self.engine.traffic_manager.current_traffic_data.pop(scenario_id)
+        # 如果指定了场景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
@@ -113,7 +256,6 @@ class MultiAgentScenarioEnv(ScenarioEnv):
         self.dones = {}
         self.episode_rewards = defaultdict(float)
         self.episode_lengths = defaultdict(int)
-
         self.controlled_agents.clear()
         self.controlled_agent_ids.clear()
 
@@ -122,37 +264,121 @@ class MultiAgentScenarioEnv(ScenarioEnv):
 
         return self._get_all_obs()
 
+    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
+
+                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:
-                agent_id = f"controlled_{car['id']}"
+            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,
+                position=car['begin'],
+                heading=car['heading']
+            )
 
-                vehicle = self.engine.spawn_object(
-                    PolicyVehicle,
-                    vehicle_config={},
-                    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
+            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']}"
                 )
-                vehicle.reset(position=car['begin'], heading=car['heading'])
-
-                vehicle.set_policy(self.policy)
-                vehicle.set_destination(car['end'])
-
-                self.controlled_agents[agent_id] = vehicle
-                self.controlled_agent_ids.append(agent_id)
-
-                # ✅ 关键：注册到引擎的 active_agents，才能参与物理更新
-                self.engine.agent_manager.active_agents[agent_id] = vehicle
 
     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):
@@ -167,38 +393,82 @@ 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)
+                                                                           base_vehicle=vehicle,
+                                                                           physics_world=self.engine.physics_world.static_world)
             lane_line_lidar = self.engine.get_sensor("lane_line_detector").perceive(num_lasers=10, distance=3,
-                                                                                    base_vehicle=vehicle,
-                                                                                    physics_world=self.engine.physics_world.static_world)
+                                                                                      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()
 
+        # 后处理
         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
+
+        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/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,98 +1,401 @@
-# MAGAIL4AutoDrive
+# MAGAIL4AutoDrive - 多智能体自动驾驶环境
 
-> 基于多智能体生成对抗模仿学习(MAGAIL)的自动驾驶训练系统 | MetaDrive + Waymo Open Motion Dataset
+基于 MetaDrive 的多智能体自动驾驶仿真与回放环境，支持 Waymo Open Dataset 的专家轨迹回放和自定义策略仿真。
 
-本项目利用 Waymo 真实驾驶数据，通过 MetaDrive 仿真环境构建专家回放系统，提取车辆状态与动作，用于训练多智能体模仿学习算法 (MAGAIL)。
+## 📋 目录
 
-## 📁 核心模块
+- [项目简介](#项目简介)
+- [功能特性](#功能特性)
+- [环境要求](#环境要求)
+- [安装步骤](#安装步骤)
+- [快速开始](#快速开始)
+- [使用指南](#使用指南)
+- [项目结构](#项目结构)
+- [配置说明](#配置说明)
+- [常见问题](#常见问题)
 
-*   **`Env/expert_replay_env.py`**: 专家回放环境。核心类 `ExpertReplayEnv`，负责读取 Waymo 轨迹，计算逆动力学动作，并过滤非道路/静态车辆。
-*   **`Env/inverse_dynamics.py`**: 逆动力学模块。根据车辆位置和航向计算油门、刹车和转向动作。
-*   **`scripts/generate_expert_data.py`**: 数据收集脚本。批量运行场景并保存训练数据。
-*   **`scripts/visualize_replay.py`**: 可视化脚本。用于观察回放效果和数据质量。
+## 项目简介
 
-***
+MAGAIL4AutoDrive 是一个基于 MetaDrive 0.4.3 的多智能体自动驾驶环境，专为模仿学习（Imitation Learning）和强化学习（Reinforcement Learning）研究设计。项目支持从真实世界数据集（如 Waymo Open Dataset）中加载场景，并提供两种核心运行模式：
 
-## 🚀 1. 数据收集
+- **回放模式（Replay Mode）**：严格按照专家轨迹回放，用于数据可视化和验证
+- **仿真模式（Simulation Mode）**：使用自定义策略控制车辆，用于算法训练和测试
 
-### 生成专家数据
-使用 `generate_expert_data.py` 脚本从 Waymo 数据集中批量提取 (State, Action) 对。
+## 功能特性
 
-```bash
-# 设置 Python 路径
-export PYTHONPATH=$PYTHONPATH:.:./metadrive
+### 核心功能
+- ✅ **多智能体支持**：同时控制多辆车辆进行协同仿真
+- ✅ **专家轨迹回放**：精确回放 Waymo 数据集中的专家驾驶行为
+- ✅ **自定义策略接口**：灵活接入各种控制策略（IDM、RL 等）
+- ✅ **智能车道过滤**：自动过滤不在车道上的异常车辆
+- ✅ **场景时长控制**：支持使用数据集原始场景时长或自定义 horizon
+- ✅ **丰富的传感器**：LiDAR、侧向检测器、车道线检测器、相机、仪表盘
 
-# 运行生成脚本
-# --data_dir: Waymo 数据路径 (建议使用 exp_filtered)
-# --output_dir: 结果保存路径
-# --num_scenarios: 要处理的场景数量
-python scripts/generate_expert_data.py \
-    --data_dir data/exp_filtered \
-    --output_dir data/training_data \
-    --num_scenarios 100 \
-    --start_index 0
+### 高级特性
+- 🎯 指定场景 ID 运行
+- 🔄 自动场景切换（修复版）
+- 📊 详细的调试日志输出
+- 🚗 车辆动态生成与管理
+- 🎮 支持可视化渲染和无头运行
+
+## 环境要求
+
+### 系统要求
+- **操作系统**：Ubuntu 18.04+ / macOS 10.14+ / Windows 10+
+- **Python 版本**：3.8 - 3.10
+- **GPU**：可选，但推荐使用（用于加速渲染）
+
+### 依赖库
 ```
 
-**生成的 `.pkl` 文件结构**：
-包含一个列表，每个元素是一条车辆轨迹（Trajectory Dictionary）：
-*   `obs`: `(T, 45)` - 观测矩阵。包含 Ego 状态 (5维) + 10辆邻居车相对信息 (40维)。
-*   `acts`: `(T, 2)` - 动作矩阵。`[Steering, Accel]`，归一化到 `[-1, 1]`。
-*   `agent_id`: 车辆 ID。
-*   `scenario_id`: 所属场景 ID。
+metadrive-simulator==0.4.3
+numpy>=1.19.0
+pygame>=2.0.0
 
-**内置过滤器**：
-脚本会自动过滤掉以下无效车辆：
-1.  **非道路车辆**：始终在停车场或路外行驶的车辆。
-2.  **静态车辆**：全称移动距离小于 5米 且速度从未超过 1m/s 的车辆（作为背景流存在，不收集数据）。
+```
+
+## 安装步骤
+
+### 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/
+│   └── ...
+
+```
+
+## 快速开始
+
+### 回放模式（推荐先尝试）
+```
+
+
+# 使用场景原始时长回放第一个场景
+
+python run_multiagent_env.py --mode replay --episodes 1 --use_scenario_duration
+
+# 回放指定场景
+
+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
+# 完整参数示例
+python run_multiagent_env.py \
+  --mode replay \
+  --episodes 1 \
+  --use_scenario_duration \
+  --debug
+
+# 仅回放车辆，禁止行人和自行车
+python run_multiagent_env.py \
+  --mode replay \
+  --use_scenario_duration \
+  --no_pedestrians \
+  --no_cyclists
+```
+
+**重要提示**：回放模式建议始终启用 `--use_scenario_duration`，否则会出现场景播放完后继续运行的问题。
+
+### 仿真模式详解
+
+仿真模式使用自定义策略控制车辆，适合算法开发和测试：
+
+```bash
+# 基础仿真
+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
+```
+
+### 自定义策略
+
+修改 `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 run_multiagent_env.py --mode replay --no_pedestrians --no_cyclists
+
+# 仅行人和自行车，无车辆（特殊场景）
+python run_multiagent_env.py --mode replay --no_vehicles
+
+# 调试模式查看过滤统计
+python run_multiagent_env.py --mode replay --debug --no_pedestrians
+```
+
+### Q6: 为什么有些车辆生成在空中？
+**A**: 已在 v1.2.0 中修复。现在所有车辆位置都只使用 2D 坐标（x, y），z 坐标设为 0，让 MetaDrive 自动处理高度，确保车辆贴在地面上。
+
+### Q7: 如何导出观测数据？
+**A**: 在 `run_multiagent_env.py` 中添加数据保存逻辑：
+```python
+import pickle
+
+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)
+```
+
+## 更新日志
+
+### v1.2.0 (2025-10-26)
+- ✅ 修复车辆生成高度问题（车辆悬空）
+- ✅ 添加对象类型过滤功能（车辆/行人/自行车）
+- ✅ 新增命令行参数：`--no_vehicles`、`--no_pedestrians`、`--no_cyclists`
+- ✅ 改进调试信息输出，显示各类型对象统计
+- ✅ 优化位置处理逻辑，只使用 2D 坐标避免高度问题
+
+### v1.1.0 (2025-10-26)
+- ✅ 修复回放模式超出场景时长问题
+- ✅ 添加场景自动切换功能
+- ✅ 改进 `replay_policy.py`，新增 `is_finished()` 方法
+- ✅ 优化 `scenario_env.py` 的 done 判断逻辑
+- ✅ 修复多回合运行时的对象清理问题
+
+### v1.0.0 (初始版本)
+- 基础多智能体环境实现
+- 回放和仿真两种模式
+- 车道过滤功能
+- Waymo 数据集支持
+
+## 贡献指南
+
+欢迎提交 Issue 和 Pull Request！
+
+### 提交 Issue
+- 请详细描述问题和复现步骤
+- 附上运行日志和错误信息
+- 说明运行环境（OS、Python 版本等）
+
+### 提交 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
 
 ---
 
-## 🔍 2. 数据可视化与验证
-
-### 回放可视化
-使用 `visualize_replay.py` 直观地观察回放效果，确认车辆行为是否自然，以及过滤逻辑是否生效。
-
-```bash
-# 运行可视化
-# --horizon: 回放的最大步数 (Waymo 场景通常为 90 或 198 步)
-python scripts/visualize_replay.py \
-    --data_dir data/exp_filtered \
-    --start_index 0 \
-    --num_scenarios 1 \
-    --horizon 200
-```
-
-**观察要点**：
-*   **受控车辆 (Controlled Agents)**：控制台会显示数量（如 `Controlled agents: 2`）。这些是真正产生数据的车辆。
-*   **背景车辆**：如果在渲染图中看到其他车（通常是路边停放的），但受控数量很少，说明静态过滤生效了。
-
-### 数据分析
-使用 `analyze_expert_data.py` 查看生成数据的统计分布。
-
-```bash
-python scripts/analyze_expert_data.py --data_path data/training_data/expert_data_0_100.pkl
-```
-
----
-
-## 🧠 3. 模型训练 (Next Steps)
-
-有了 `data/training_data/` 下的专家数据后，您可以开始训练 MAGAIL 模型。
-
-### 训练流程
-1.  **加载数据**：使用 `dataset/expert_dataset.py` 中的 `ExpertDataset` 类加载 `.pkl` 数据。
-2.  **初始化 MAGAIL**：
-    *   **Generator (Policy)**: 接收观测 `(B, 45)`，输出动作 `(B, 2)`。
-    *   **Discriminator**: 接收状态-动作对 `(s, a)`，判断是专家还是生成器。
-3.  **交互采样**：
-    *   在 `MultiAgentScenarioEnv`（非回放模式）中运行 Policy。
-    *   收集 Policy 生成的轨迹。
-4.  **对抗更新**：
-    *   利用专家数据和 Policy 数据训练 Discriminator。
-    *   利用 Discriminator 的输出作为 Reward (GAIL Reward) 训练 Policy (PPO/TRPO)。
-
-### 推荐配置
-*   **Observation**: 45维 (Ego + 10 Neighbors)
-*   **Action**: 2维 Continuous (Steering, Accel)
-*   **Horizon**: 200 steps
-*   **Batch Size**: 1024+ (多智能体环境下数据量很大)
+**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,304 +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:
-            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="./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/magail_dataset.py

@@ -1,61 +0,0 @@
-import torch
-from torch.utils.data import Dataset
-import pickle
-import numpy as np
-import os
-import glob
-
-class MAGAILExpertDataset(Dataset):
-    def __init__(self, data_dir, transform=None):
-        """
-        Args:
-            data_dir (str): Directory containing .pkl files from generate_expert_data.py
-            transform (callable, optional): Optional transform to be applied on a sample.
-        """
-        self.data_dir = data_dir
-        self.transform = transform
-        self.trajectories = []
-        self.flat_data = [] # (obs, act) pairs
-        
-        # Load all .pkl files
-        pkl_files = glob.glob(os.path.join(data_dir, "*.pkl"))
-        print(f"Loading data from {len(pkl_files)} files in {data_dir}...")
-        
-        for pkl_file in pkl_files:
-            try:
-                with open(pkl_file, 'rb') as f:
-                    data = pickle.load(f)
-                    # data is a list of dicts: {'obs': (T, 45), 'acts': (T, 2), ...}
-                    self.trajectories.extend(data)
-            except Exception as e:
-                print(f"Error loading {pkl_file}: {e}")
-                
-        # Flatten for training Discriminator/BC
-        print(f"Processing {len(self.trajectories)} trajectories...")
-        for traj in self.trajectories:
-            obs = traj['obs']
-            acts = traj['acts']
-            
-            # obs: (T, 45), acts: (T, 2)
-            # We pair them up
-            for i in range(len(obs)):
-                self.flat_data.append((obs[i], acts[i]))
-                
-        print(f"Total samples: {len(self.flat_data)}")
-        
-    def __len__(self):
-        return len(self.flat_data)
-
-    def __getitem__(self, idx):
-        obs, act = self.flat_data[idx]
-        
-        # Convert to tensor
-        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/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="/home/huangfukk/MAGAIL4AutoDrive/data/exp_filtered", help="Path to Waymo pickles (or filtered index)")
-    parser.add_argument("--output_dir", type=str, default="/home/huangfukk/MAGAIL4AutoDrive/data/training", 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/visualize_expert_trajectory.py

@@ -1,105 +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
-import matplotlib.pyplot as plt
-from matplotlib.animation import FuncAnimation
-
-class DummyPolicy:
-    """
-    占位策略,用于数据检查时初始化环境
-    不需要实际执行动作,只是为了满足环境初始化要求
-    """
-    def act(self, *args, **kwargs):
-        # 返回零动作 [throttle, steering]
-        return np.array([0.0, 0.0])
-
-def visualize_expert_trajectory(env, scenario_idx=0):
-    """
-    可视化专家轨迹的俯视图动画
-    """
-    env.reset()
-    expert_trajs = env.expert_trajectories
-    
-    if len(expert_trajs) == 0:
-        print("当前场景无专家轨迹")
-        return
-    
-    # 设置绘图
-    fig, ax = plt.subplots(figsize=(12, 12))
-    
-    # 获取所有轨迹的最大时间长度
-    max_timestep = max(traj["end_timestep"] for traj in expert_trajs.values())
-    min_timestep = min(traj["start_timestep"] for traj 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 = traj["positions"][:, :2]
-        ax.plot(positions[:, 0], positions[:, 1], 
-                color=colors[idx], alpha=0.3, linewidth=1,
-                label=f'Vehicle {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():
-            if traj["start_timestep"] <= current_time <= traj["end_timestep"]:
-                idx = current_time - traj["start_timestep"]
-                pos = traj["positions"][idx, :2]
-                current_positions.append(pos)
-        
-        if len(current_positions) > 0:
-            current_positions = np.array(current_positions)
-            scatter.set_offsets(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
-
-if __name__ == "__main__":
-    WAYMO_DATA_DIR = r"/home/huangfukk/mdsn"
-    data_dir = AssetLoader.file_path(WAYMO_DATA_DIR, "exp_filtered", unix_style=False)
-    
-    env = MultiAgentScenarioEnv(
-        config={
-            "data_directory": data_dir,
-            "is_multi_agent": True,
-            "num_controlled_agents": 3,
-            "use_render": False,
-        },
-        agent2policy=DummyPolicy()
-    )
-    
-    # 可视化第一个场景
-    anim = visualize_expert_trajectory(env, scenario_idx=0)

scripts/visualize_replay.py

@@ -1,93 +0,0 @@
-import argparse
-import os
-import sys
-import time
-
-# 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 Env.expert_replay_env import ExpertReplayEnv
-
-def visualize_replay(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")
-
-    # Same as data generation: avoid MetaDrive assertion when requested num_scenarios > available.
-    from metadrive.scenario.utils import read_dataset_summary
-    _, summary_lookup, _ = read_dataset_summary(data_path)
-    if args.start_index >= len(summary_lookup):
-        raise ValueError(
-            f"start_index={args.start_index} out of range. Dataset has {len(summary_lookup)} scenarios."
-        )
-    max_available = len(summary_lookup) - args.start_index
-    num_to_run = min(args.num_scenarios, max_available)
-
-    env_config = {
-        "data_directory": data_path,
-        "is_multi_agent": True,
-        "num_controlled_agents": 100,
-        "horizon": args.horizon,
-        "use_render": True,  # Enable rendering
-        "sequential_seed": True,
-        "reactive_traffic": False,
-        "start_scenario_index": args.start_index,
-        "num_scenarios": -1,
-        "log_level": 40, # ERROR
-        # "pstats": True, # For performance debugging
-    }
-
-    print(f"Initializing ExpertReplayEnv with data from {data_path}...")
-    env = ExpertReplayEnv(config=env_config)
-
-    try:
-        for i in range(args.start_index, args.start_index + num_to_run):
-            print(f"\n--- Playing Scenario {i} ---")
-            try:
-                obs = env.reset(seed=i)
-            except Exception as e:
-                print(f"Error resetting scenario {i}: {e}")
-                continue
-            
-            print(f"Scenario loaded. Controlled agents: {len(env.controlled_agents)}")
-            
-            for step in range(args.horizon):
-                # Step
-                obs, rewards, dones, infos = env.step(None)
-                
-                # Render
-                env.render(mode="top_down", 
-                           text={
-                               "Step": step, 
-                               "Agents": len(env.controlled_agents),
-                               "Scenario": i
-                           })
-                
-                # Sleep to control playback speed
-                time.sleep(0.05)
-                
-                if dones["__all__"]:
-                    print(f"Scenario {i} finished at step {step}")
-                    break
-                    
-    except KeyboardInterrupt:
-        print("Interrupted by user")
-    except Exception as e:
-        import traceback
-        traceback.print_exc()
-        print(f"Global error: {e}")
-    finally:
-        env.close()
-        print("Environment closed.")
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--data_dir", type=str, default="/home/huangfukk/MAGAIL4AutoDrive/data/exp_filtered", help="Path to Waymo data")
-    parser.add_argument("--start_index", type=int, default=0)
-    parser.add_argument("--num_scenarios", type=int, default=1)
-    parser.add_argument("--horizon", type=int, default=500)
-    
-    args = parser.parse_args()
-    visualize_replay(args)

train_magail.py

@@ -1,521 +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.magail_dataset import MAGAILExpertDataset
+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 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
-
-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)
-        )
-
-    def forward(self, state):
-        return self.net(state)
-
-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() 
-        )
-
-    def forward(self, state, action):
-        x = torch.cat([state, action], dim=-1)
-        return self.net(x)
-
-# --- PPO Algorithm ---
-
-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)
+def train():
+    # 1. 环境
+    env = MultiAgentScenarioEnv(config={...}, agent2policy=None)
+    
+    # 2. 专家 Buffer (伪代码)
+    # expert_buffer = RolloutBuffer(...)
+    # expert_buffer.load(...) 
+    # 必须保证 expert_buffer.sample() 返回 (state, next_state) 供 Discriminator 训练
+    
+    # 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"]
+    )
+    
+    writer = SummaryWriter("./logs")
+    
+    # 4. 训练循环
+    total_steps = 0
+    obs_dict = env.reset()
+    
+    # 将 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]) 
+    
+    # 如果需要 state_gail，假设它就是 obs
+    current_state_gail = current_obs.copy()
+    
+    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()
-
-    def select_action(self, state):
-        with torch.no_grad():
-            state = torch.FloatTensor(state).cuda()
-            dist = self.actor(state)
-            action = dist.sample()
-            action_logprob = dist.log_prob(action).sum(dim=-1)
-        return action.cpu().numpy(), action_logprob.cpu().numpy()
-
-    def update(self, memory):
-        # Convert memory to tensors
-        states = torch.FloatTensor(np.array(memory['states'])).cuda()
-        actions = torch.FloatTensor(np.array(memory['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()
+        # 注意：magail.step 内部调用了 env.step，这可能不适用于 Multi-Agent 环境返回 Dict 的情况
+        # 原 PPO.step 似乎是为 Single-Agent 或 VectorEnv 设计的
+        # 这里我们需要手动写 Rollout 循环或者修改 PPO.step 以适配 Dict 返回值
         
-        # 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.
-        
-        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.
-        
-        values = self.critic(states).detach()
-        next_values = self.critic(next_states).detach()
-        
-        # 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)
+        # --- 方案：手动 Rollout 适配 Multi-Agent ---
+        for _ in range(CONFIG["rollout_length"]):
+            # 1. 决策
+            actions_list, log_pis_list = magail.explore(current_obs)
             
-        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 = dist.log_prob(actions).sum(dim=-1)
-            dist_entropy = dist.entropy().sum(dim=-1)
-            state_values = self.critic(states).squeeze()
+            # 2. 拼装 Action Dict
+            action_dict = {agent_id: action for agent_id, action in zip(agents, actions_list)}
             
-            # 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
+            # 3. 环境步进
+            next_obs_dict, rewards_dict, dones_dict, infos_dict = env.step(action_dict)
             
-            # 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 (Dummy for now, usually you run simulation here)
-    # But for MAGAIL we need to collect generated trajectories.
-    # We need the Env class to be importable.
-    from Env.scenario_env import MultiAgentScenarioEnv
-    from Env.simple_idm_policy import ConstantVelocityPolicy # Just for init
-    
-    # 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
-    from Env.expert_replay_env import ExpertReplayEnv # Using ReplayEnv for config, but we need ScenarioEnv for simulation?
-    # Actually we need MultiAgentScenarioEnv for interactive training, not Replay.
-    from Env.scenario_env import MultiAgentScenarioEnv
-    from Env.simple_idm_policy import ConstantVelocityPolicy # Placeholder policy for init
-    
-    # 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 = MAGAILScenarioEnv(config=env_config, agent2policy={}) # Pass empty dict if we control all externally
-    
-    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': [], '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 = MAGAILScenarioEnv(config=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 = {}
-            
-            # 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 = ppo_agent.select_action(obs) # Select action returns numpy
-                actions[agent_id] = act.flatten() # (2,)
-                action_logprobs[agent_id] = logprob # scalar
-                
-            # 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['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(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()
-            
-            # --- 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()
+            # 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():
-                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()
+                # 重新计算一遍或者在 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
+            )
             
-            # Update PPO
-            ppo_loss = ppo_agent.update(memory)
+            current_obs = next_obs
+            current_state_gail = next_obs # 更新 gail state
+            total_steps += len(agents) # 步数增加 N
             
-            # 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)
+            # 处理 Done
+            if all(dones):
+                obs_dict = env.reset()
+                current_obs = np.stack([obs_dict[a] for a in agents])
+                current_state_gail = current_obs
         
-        print(f"Episode {i_episode}: Disc Loss {disc_loss.item():.4f} | PPO Loss {ppo_loss:.4f} | Mean Reward {np.mean(all_gail_rewards):.4f}")
+        # --- 更新 ---
+        # 当 Buffer 满时 (rollout_length)，调用 update
+        magail.update(writer, total_steps)
         
-        if i_episode % 50 == 0:
-            ppo_agent.save(os.path.join(args.log_dir, f"model_{i_episode}"))
+        # 保存
+        if total_steps % 10000 == 0:
+            magail.save_models("./models")
 
-    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="runs/magail_exp")
-    
-    args = parser.parse_args()
-    
-    # Create log dir
-    os.makedirs(args.log_dir, exist_ok=True)
-    
-    train(args)