完善项目目录结构

Env/scenario_env.py

@@ -100,33 +100,13 @@ class MultiAgentScenarioEnv(ScenarioEnv):
         for scenario_id in _obj_to_clean_this_frame:
             self.engine.traffic_manager.current_traffic_data.pop(scenario_id)
 
-        # Fix: Ensure all objects are cleared properly before reset
+        # Clear vehicles we spawned via engine.spawn_object() so _object_clean_check() passes
-        # Instead of manually clearing, we just let the engine handle it, but we might need to 
+        ids_to_clear = [v.id for v in self.controlled_agents.values()]
-        # ensure no stale references in managers.
+        if ids_to_clear:
-        
+            self.engine.clear_objects(ids_to_clear)
-        # The error "KeyError" in clear_objects usually means we are trying to clear an object 
+        self.controlled_agents.clear()
-        # that is already gone from _spawned_objects but still tracked by a manager.
+        self.controlled_agent_ids.clear()
-        
+
-        # Try to clear only objects that actually exist in the engine
-        # existing_objects = list(self.engine.get_objects().keys())
-        # if existing_objects:
-        #      self.engine.clear_objects(existing_objects)
-        
-        # Force clear agent manager's spawned objects to avoid stale references
-        if hasattr(self.engine, 'agent_manager') and self.engine.agent_manager:
-            # Check if it's ScenarioAgentManager or VehicleAgentManager
-            # ScenarioAgentManager might not have spawned_objects directly exposed or named differently
-            # But BaseAgentManager usually has it.
-            # If it's ScenarioAgentManager, it might be using a different structure.
-            
-            # Safe clear for BaseAgentManager subclasses
-            if hasattr(self.engine.agent_manager, 'spawned_objects'):
-                self.engine.agent_manager.spawned_objects.clear()
-            
-            # Also clear active_objects if present (VehicleAgentManager uses this)
-            if hasattr(self.engine.agent_manager, '_active_objects'):
-                self.engine.agent_manager._active_objects.clear()
-                
         self.engine.reset()
         self.reset_sensors()
         self.engine.taskMgr.step()
@@ -141,9 +121,6 @@ class MultiAgentScenarioEnv(ScenarioEnv):
         self.episode_rewards = defaultdict(float)
         self.episode_lengths = defaultdict(int)
 
-        self.controlled_agents.clear()
-        self.controlled_agent_ids.clear()
-
         super().reset(seed)  # 初始化场景
         self._spawn_controlled_agents()
 
@@ -217,6 +194,7 @@ class MultiAgentScenarioEnv(ScenarioEnv):
                 self.controlled_agents[agent_id].before_step(action)
 
         self.engine.step()
+        self.engine.after_step()
 
         for agent_id in action_dict:
             if agent_id in self.controlled_agents:

README.md

@@ -22,12 +22,11 @@ MAGAIL4AutoDrive/
 │   ├── simple_idm_policy.py  # ConstantVelocityPolicy 占位策略
 │   └── ...
 ├── dataset/                   # 数据集加载器
-│   ├── expert_dataset.py      # 通用专家数据加载类
+│   ├── loader.py              # 主流水线：load_expert_pkl、MAGAILExpertDataset
-│   └── magail_dataset.py      # MAGAIL 训练专用数据加载器
+│   └── expert_dataset.py      # 可选 107 维/5 维管线
 ├── scripts/                   # 工具脚本（数据、回放、可视化、分析）
 │   ├── generate_expert_data.py    # 从 Waymo 生成专家 (obs, act) pkl
-│   ├── visualize_replay.py       # 原始专家数据回放
+│   ├── visualize.py              # 可视化统一入口（replay / policy / trajectory）
-│   ├── visualize_trained_policy.py # BC/MAGAIL 策略可视化统一入口
 │   ├── analyze_expert_data.py    # 数据分布分析
 │   ├── launch_tensorboard.py     # 启动 TensorBoard
 │   ├── README.md                  # 脚本用法说明
@@ -44,7 +43,6 @@ MAGAIL4AutoDrive/
 │   └── magail/
 ├── train_bc.py                # [根目录] BC 训练
 ├── train_magail.py            # [根目录] MAGAIL 训练
-├── visualize_bc.py            # [根目录] BC 可视化薄包装 -> scripts/visualize_trained_policy.py
 └── README.md
 ```
 
@@ -56,6 +54,34 @@ MAGAIL4AutoDrive/
 
 所有默认路径均为相对项目根，便于在不同设备上复用。
 
+## 数据处理流程
+
+从 Waymo Motion 原始数据到本项目训练用专家 pkl，依次为：
+
+**1) 下载 Waymo Motion（TFRecord）**  
+安装 `gsutil` 并登录 Google 账号后，例如只下载 training_20s：
+
+```bash
+gsutil -m cp -r "gs://waymo_open_dataset_motion_v_1_2_0/uncompressed/scenario/training_20s" ./waymo/
+```
+
+**2) ScenarioNet Convert（TFRecord → ScenarioNet 场景库）**  
+需安装 ScenarioNet、MetaDrive 及 TensorFlow 2.11、protobuf 3.20；转换时不用 GPU。
+
+```bash
+python -m scenarionet.convert_waymo -d data/exp_converted --raw_data_path ./waymo/training_20s --num_workers 64
+```
+
+**3) ScenarioNet Filter（按需筛选场景）**  
+从 convert 得到的场景库中筛掉含红绿灯、天桥等场景，输出到如 `data/exp_filtered`。具体命令以 ScenarioNet 文档为准（Operations → Filter）。
+
+**4) 本项目：生成专家 pkl**  
+使用筛选后的场景目录，生成训练用 pkl 到 `data/training_data`：
+
+```bash
+python scripts/generate_expert_data.py --data_dir data/exp_filtered --output_dir data/training_data --num_scenarios 100 --start_index 0
+```
+
 ## 核心工作流
 
 ### 1. 数据准备
@@ -67,21 +93,20 @@ python scripts/generate_expert_data.py --data_dir data/exp_filtered --output_dir
 
 ### 2. 行为克隆 (BC)
 - **训练**：`python train_bc.py`（模型保存到 `models/bc/`，日志到 `logs/bc/`）
-- **可视化**：`python visualize_bc.py` 或 `python scripts/visualize_trained_policy.py --policy_type bc --model_path models/bc/policy_best.pt`
+- **可视化**：`python scripts/visualize.py policy --policy_type bc --model_path models/bc/policy_best.pt`
 
 ### 3. 多智能体对抗模仿学习 (MAGAIL)
 - **训练**：`python train_magail.py`（模型保存到 `models/magail/`，日志到 `logs/magail/`）
-- **可视化**：`python scripts/visualize_trained_policy.py --policy_type magail --model_path models/magail/model_50_actor.pth`
+- **可视化**：`python scripts/visualize.py policy --policy_type magail --model_path models/magail/model_50_actor.pth`
 
-### 4. 策略可视化统一入口
+### 4. 可视化统一入口
-BC 与 MAGAIL 共用 `scripts/visualize_trained_policy.py`，通过 `--policy_type bc|magail`（或根据 `--model_path` 自动推断）选择模型类型。根目录 `visualize_bc.py` 为 BC 的薄包装。详见 [scripts/README_visualize.md](scripts/README_visualize.md) 与 [scripts/README.md](scripts/README.md)。
+可视化统一使用 `scripts/visualize.py`，子命令：`replay`（场景回放）、`policy`（BC/MAGAIL 策略）、`trajectory`（专家轨迹 2D 动画）。详见 [scripts/README.md](scripts/README.md)。
 
 ## 文件与模块职责
 
 ### 根目录脚本
-- **train_bc.py**：BC 训练，加载 `data/training_data` 下 pkl，模型与日志写入 `models/bc/`、`logs/bc/`
+- **train_bc.py**：BC 训练，从 `dataset.loader` 加载专家 pkl，模型与日志写入 `models/bc/`、`logs/bc/`
-- **train_magail.py**：MAGAIL 训练，环境使用 `BCScenarioEnv`（45 维），模型与日志写入 `models/magail/`、`logs/magail/`
+- **train_magail.py**：MAGAIL 训练，环境使用 `BCScenarioEnv`（45 维），从 `dataset.loader` 加载专家数据，模型与日志写入 `models/magail/`、`logs/magail/`
-- **visualize_bc.py**：薄包装，调用 `scripts/visualize_trained_policy.py --policy_type bc`
 
 ### Env 模块
 - **Env/bc_env.py**：`BCScenarioEnv`，45 维观测（Ego 5 维 + 10 邻居×4 维），BC 与 MAGAIL 训练/评估共用

dataset/loader.py

@@ -0,0 +1,103 @@
+"""
+统一数据加载：BC/MAGAIL 训练用专家 pkl 的加载函数与 Dataset。
+主训练流水线使用本模块；dataset/expert_dataset.py 为可选 107 维/5 维管线。
+"""
+import os
+import glob
+import pickle
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+
+
+def load_expert_pkl(expert_data_path):
+    """从目录或单个 pkl 加载专家 (obs, acts)，返回 concat 后的 obs_data, act_data。"""
+    if os.path.isdir(expert_data_path):
+        pkl_files = glob.glob(os.path.join(expert_data_path, "*.pkl"))
+        if not pkl_files:
+            raise FileNotFoundError(f"No .pkl files in {expert_data_path}")
+        print(f"Found {len(pkl_files)} pickle files in {expert_data_path}")
+    elif os.path.exists(expert_data_path):
+        pkl_files = [expert_data_path]
+    else:
+        raise FileNotFoundError(f"Expert data path not found: {expert_data_path}")
+
+    obs_data, act_data = [], []
+    for pkl_file in pkl_files:
+        try:
+            with open(pkl_file, "rb") as f:
+                data = pickle.load(f)
+            if isinstance(data, list):
+                for traj in data:
+                    if "obs" in traj and "acts" in traj:
+                        obs_data.append(traj["obs"])
+                        act_data.append(traj["acts"])
+            elif isinstance(data, dict):
+                if "observations" in data and "actions" in data:
+                    obs_data.append(data["observations"])
+                    act_data.append(data["actions"])
+            else:
+                print(f"Skipping {pkl_file}: Unknown data format {type(data)}")
+        except Exception as e:
+            print(f"Error loading {pkl_file}: {e}")
+
+    if len(obs_data) == 0:
+        raise ValueError("No valid data loaded from provided path.")
+    obs_data = np.concatenate(obs_data, axis=0)
+    act_data = np.concatenate(act_data, axis=0)
+    print(f"Total loaded samples: {len(obs_data)}")
+    return obs_data, act_data
+
+
+class MAGAILExpertDataset(Dataset):
+    def __init__(self, data_dir, transform=None):
+        """
+        Args:
+            data_dir (str): Directory containing .pkl files from generate_expert_data.py
+            transform (callable, optional): Optional transform to be applied on a sample.
+        """
+        self.data_dir = data_dir
+        self.transform = transform
+        self.trajectories = []
+        self.flat_data = []  # (obs, act) pairs
+
+        # Load all .pkl files
+        pkl_files = glob.glob(os.path.join(data_dir, "*.pkl"))
+        print(f"Loading data from {len(pkl_files)} files in {data_dir}...")
+
+        for pkl_file in pkl_files:
+            try:
+                with open(pkl_file, "rb") as f:
+                    data = pickle.load(f)
+                    # data is a list of dicts: {'obs': (T, 45), 'acts': (T, 2), ...}
+                    self.trajectories.extend(data)
+            except Exception as e:
+                print(f"Error loading {pkl_file}: {e}")
+
+        # Flatten for training Discriminator/BC
+        print(f"Processing {len(self.trajectories)} trajectories...")
+        for traj in self.trajectories:
+            obs = traj["obs"]
+            acts = traj["acts"]
+
+            # obs: (T, 45), acts: (T, 2)
+            for i in range(len(obs)):
+                self.flat_data.append((obs[i], acts[i]))
+
+        print(f"Total samples: {len(self.flat_data)}")
+
+    def __len__(self):
+        return len(self.flat_data)
+
+    def __getitem__(self, idx):
+        obs, act = self.flat_data[idx]
+
+        obs = torch.from_numpy(obs).float()
+        act = torch.from_numpy(act).float()
+
+        sample = {"state": obs, "action": act}
+
+        if self.transform:
+            sample = self.transform(sample)
+
+        return sample

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/README.md

@@ -22,36 +22,31 @@
 
 ---
 
-### 回放与可视化
+### 可视化（统一入口）
 
 | 脚本 | 用途 | 用法示例 |
 |------|------|----------|
-| [visualize_replay.py](visualize_replay.py) | 原始专家轨迹回放（ExpertReplayEnv） | `python scripts/visualize_replay.py --data_dir data/exp_filtered --num_scenarios 1 --horizon 200` |
+| [visualize.py](visualize.py) | **replay**：场景回放（ExpertReplayEnv）；**policy**：BC/MAGAIL 策略；**trajectory**：专家轨迹 2D 动画 | 见下方 |
-| [visualize_trained_policy.py](visualize_trained_policy.py) | **BC/MAGAIL 共用**：加载训练好的策略在 45 维场景中可视化 | 见下方「训练策略可视化」小节 |
 
-#### 训练策略可视化（visualize_trained_policy.py）
+**子命令**：
 
-使用训练好的 **BC** 或 **MAGAIL** 模型在 45 维场景环境中运行，并实时渲染俯瞰图（top-down view）。统一入口：`scripts/visualize_trained_policy.py`。
+- **replay**（原始专家轨迹回放）：
-
-**BC 模型**：
 ```bash
-python scripts/visualize_trained_policy.py --policy_type bc --model_path models/bc/policy_best.pt --data_dir data/exp_filtered --num_scenarios 1
+python scripts/visualize.py replay --data_dir data/exp_filtered --num_scenarios 1 --horizon 500
 ```
 
-**MAGAIL 模型**：
+- **policy**（BC 或 MAGAIL 训练策略）：
 ```bash
-python scripts/visualize_trained_policy.py --policy_type magail --model_path models/magail/model_50_actor.pth --data_dir data/exp_filtered --num_scenarios 1 --deterministic
+python scripts/visualize.py policy --policy_type bc --model_path models/bc/policy_best.pt --data_dir data/exp_filtered --num_scenarios 1
+python scripts/visualize.py policy --policy_type magail --model_path models/magail/model_50_actor.pth --num_scenarios 1 --deterministic
 ```
 
-**自动推断类型**（根据 `--model_path` 扩展名：`.pt` → BC，否则 → MAGAIL）：
+- **trajectory**（专家轨迹 matplotlib 俯视图动画）：
 ```bash
-python scripts/visualize_trained_policy.py --model_path models/bc/policy_best.pt
+python scripts/visualize.py trajectory --data_dir data/exp_filtered --scenario_idx 0
-python scripts/visualize_trained_policy.py --model_path models/magail/model_50_actor.pth
 ```
 
-**根目录 BC 薄包装**：`python visualize_bc.py --model_path models/bc/policy_best.pt`
+**公共参数**：`--data_dir`（默认 `data/exp_filtered`）、`--start_index`、`--num_scenarios`、`--horizon`。policy 模式另有 `--policy_type`（auto/bc/magail）、`--model_path`、`--deterministic`（仅 MAGAIL）。
-
-**参数**：`--policy_type`（`auto`|`bc`|`magail`）、`--model_path`（默认 `models/bc/policy_best.pt`）、`--data_dir`、`--start_index`、`--num_scenarios`、`--horizon`、`--deterministic`（仅 MAGAIL）。环境统一为 45 维 `BCScenarioEnv`，渲染为 MetaDrive top_down。数据目录未指定时默认 `data/exp_filtered`（不存在则 `data/exp_converted`）。
 
 ---
 
@@ -62,7 +57,6 @@ python scripts/visualize_trained_policy.py --model_path models/magail/model_50_a
 | [analyze_expert_data.py](analyze_expert_data.py) | 分析专家数据分布与统计 | 见脚本内 `__main__`（依赖 env 与数据目录配置） |
 | [check_track_fields.py](check_track_fields.py) | 检查 Waymo 轨迹字段 | 见脚本内 `__main__` |
 | [check_database_info.py](check_database_info.py) | 检查数据库/场景信息 | 见脚本内 `__main__`（含硬编码路径，可按需改为 `data/exp_filtered`） |
-| [visualize_expert_trajectory.py](visualize_expert_trajectory.py) | 用 matplotlib 画专家轨迹动画 | 依赖 `env.expert_trajectories`，与当前 env 接口可能不一致，可选使用 |
 
 ---
 
@@ -79,4 +73,4 @@ python scripts/visualize_trained_policy.py --model_path models/magail/model_50_a
 1. **数据准备**：`generate_expert_data.py` → 输出到 `data/training_data/*.pkl`
 2. **BC 训练**：根目录 `train_bc.py` → 模型保存到 `models/bc/`，日志到 `logs/bc/`
 3. **MAGAIL 训练**：根目录 `train_magail.py` → 模型保存到 `models/magail/`，日志到 `logs/magail/`
-4. **可视化**：`visualize_trained_policy.py`（或根目录 `visualize_bc.py` 仅 BC）→ 从 `models/bc` 或 `models/magail` 加载模型，数据目录默认 `data/exp_filtered`
+4. **可视化**：`scripts/visualize.py`（子命令 replay / policy / trajectory）→ 数据目录默认 `data/exp_filtered`

scripts/visualize.py

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

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)

scripts/visualize_trained_policy.py

@@ -1,189 +0,0 @@
-"""
-Unified visualization for BC and MAGAIL trained policies.
-Use --policy_type bc or magail (or auto-detect from --model_path: .pt -> bc, else magail).
-"""
-import argparse
-import os
-import sys
-import torch
-import numpy as np
-
-# Add project root to Python path
-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.bc_env import BCScenarioEnv
-from metadrive.engine.engine_utils import close_engine
-
-
-def _resolve_data_dir(args):
-    """Resolve data directory: explicit or auto-detect under project data/."""
-    if args.data_dir:
-        data_dir = args.data_dir
-    else:
-        current_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
-        data_dir = os.path.join(current_dir, "data", "exp_filtered")
-        if not os.path.exists(data_dir):
-            data_dir = os.path.join(current_dir, "data", "exp_converted")
-    if not os.path.exists(data_dir):
-        raise FileNotFoundError(f"Data directory not found at {data_dir}. Please specify --data_dir.")
-    return data_dir
-
-
-def _resolve_model_path(model_path, policy_type):
-    """Resolve model path: if not found, try models/bc or models/magail."""
-    if os.path.exists(model_path):
-        return model_path
-    if policy_type == "bc":
-        candidate = os.path.join("models", "bc", model_path)
-    else:
-        candidate = os.path.join("models", "magail", model_path)
-    if os.path.exists(candidate):
-        return candidate
-    if policy_type == "magail" and not model_path.endswith("_actor.pth"):
-        candidate = model_path + "_actor.pth"
-        if os.path.exists(candidate):
-            return candidate
-    raise FileNotFoundError(f"Model path {model_path} not found (tried {candidate}).")
-
-
-def visualize_model(args):
-    policy_type = (args.policy_type or "auto").lower()
-    if policy_type == "auto":
-        policy_type = "bc" if args.model_path.endswith(".pt") else "magail"
-
-    data_dir = _resolve_data_dir(args)
-    data_path = os.path.abspath(data_dir)
-    env_config = {
-        "data_directory": data_path,
-        "is_multi_agent": True,
-        "num_controlled_agents": 3,
-        "horizon": args.horizon,
-        "use_render": True,
-        "sequential_seed": True,
-        "start_scenario_index": args.start_index,
-        "num_scenarios": args.num_scenarios,
-        "log_level": 40,
-    }
-
-    print(f"Initializing BCScenarioEnv (policy_type={policy_type})...")
-    try:
-        env = BCScenarioEnv(env_config, agent2policy={})
-    except Exception as e:
-        print(f"Error init env: {e}. Trying to close lingering engine...")
-        try:
-            close_engine()
-        except Exception:
-            pass
-        env = BCScenarioEnv(env_config, agent2policy={})
-
-    state_dim = 45
-    action_dim = 2
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-    model_path = _resolve_model_path(args.model_path, policy_type)
-    print(f"Loading model from {model_path}...")
-
-    if policy_type == "bc":
-        from Algorithm.policy import StateIndependentPolicy
-        policy = StateIndependentPolicy(
-            state_shape=(state_dim,),
-            action_shape=(action_dim,),
-            hidden_units=(256, 256),
-            hidden_activation=torch.nn.Tanh(),
-        ).to(device)
-        policy.load_state_dict(torch.load(model_path, map_location=device))
-        policy.eval()
-    else:
-        from train_magail import Actor
-        actor = Actor(state_dim, action_dim).to(device)
-        actor.load_state_dict(torch.load(model_path, map_location=device))
-        actor.eval()
-
-    try:
-        for i in range(args.start_index, args.start_index + args.num_scenarios):
-            print(f"\n--- Playing Scenario {i} ---")
-            try:
-                obs_dict = env.reset(seed=i)
-            except Exception as e:
-                print(f"Error resetting {i}: {e}. Skipping.")
-                try:
-                    close_engine()
-                    env = BCScenarioEnv(env_config, agent2policy={})
-                except Exception:
-                    pass
-                continue
-
-            print(f"Scenario loaded. Controlled agents: {len(obs_dict)}")
-            step_count = 0
-            episode_reward = 0.0
-
-            while True:
-                actions = {}
-                agent_ids = list(obs_dict.keys())
-                obs_list = [obs_dict[aid] for aid in agent_ids]
-                obs_tensor = torch.FloatTensor(np.array(obs_list)).to(device)
-
-                with torch.no_grad():
-                    if policy_type == "bc":
-                        actions_np = policy(obs_tensor).cpu().numpy()
-                    else:
-                        dist = actor(obs_tensor)
-                        if args.deterministic:
-                            actions_np = torch.tanh(dist.mean).cpu().numpy()
-                        else:
-                            actions_np = torch.tanh(dist.sample()).cpu().numpy()
-
-                for idx, aid in enumerate(agent_ids):
-                    actions[aid] = actions_np[idx].flatten()
-
-                obs_dict, rewards, dones, infos = env.step(actions)
-                episode_reward += sum(rewards.values())
-
-                env.render(
-                    mode="top_down",
-                    text={
-                        "Scenario": i,
-                        "Step": step_count,
-                        "Agents": len(obs_dict),
-                        "Total Reward": f"{episode_reward:.2f}",
-                    },
-                )
-                step_count += 1
-
-                if dones["__all__"] or step_count >= args.horizon:
-                    print(f"Scenario finished at step {step_count}, reward {episode_reward:.2f}")
-                    break
-
-    except KeyboardInterrupt:
-        print("Interrupted.")
-    finally:
-        env.close()
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser(
-        description="Visualize BC or MAGAIL trained policy in 45-dim scenario env."
-    )
-    parser.add_argument(
-        "--policy_type",
-        type=str,
-        default="auto",
-        choices=["auto", "bc", "magail"],
-        help="Policy type: bc (StateIndependentPolicy .pt) or magail (Actor _actor.pth). auto = infer from model_path.",
-    )
-    parser.add_argument(
-        "--model_path",
-        type=str,
-        default="models/bc/policy_best.pt",
-        help="Path to model: BC .pt (e.g. models/bc/policy_best.pt) or MAGAIL _actor.pth (e.g. models/magail/model_50_actor.pth)",
-    )
-    parser.add_argument("--data_dir", type=str, default=None, help="Waymo data directory (default: data/exp_filtered)")
-    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=200)
-    parser.add_argument("--deterministic", action="store_true", help="For MAGAIL: use mean action instead of sampling")
-
-    args = parser.parse_args()
-    visualize_model(args)

train_bc.py

@@ -3,8 +3,6 @@ BC 训练脚本：负责数据加载、环境评估、日志与保存；BC 算
 使用方式不变：python train_bc.py [--expert_data_path data/training_data] [--save_dir models/bc] ...
 """
 import os
-import glob
-import pickle
 import numpy as np
 import torch
 import argparse
@@ -17,45 +15,7 @@ from torch.utils.tensorboard import SummaryWriter
 from Algorithm.policy import StateIndependentPolicy
 from Algorithm.bc import train_bc_epoch, eval_bc_epoch
 from Env.bc_env import BCScenarioEnv
-
+from dataset.loader import load_expert_pkl
-
-def load_expert_data(expert_data_path):
-    """从目录或单个 pkl 加载专家 (obs, acts)，返回 concat 后的 obs_data, act_data."""
-    if os.path.isdir(expert_data_path):
-        pkl_files = glob.glob(os.path.join(expert_data_path, "*.pkl"))
-        if not pkl_files:
-            raise FileNotFoundError(f"No .pkl files in {expert_data_path}")
-        print(f"Found {len(pkl_files)} pickle files in {expert_data_path}")
-    elif os.path.exists(expert_data_path):
-        pkl_files = [expert_data_path]
-    else:
-        raise FileNotFoundError(f"Expert data path not found: {expert_data_path}")
-
-    obs_data, act_data = [], []
-    for pkl_file in pkl_files:
-        try:
-            with open(pkl_file, "rb") as f:
-                data = pickle.load(f)
-            if isinstance(data, list):
-                for traj in data:
-                    if "obs" in traj and "acts" in traj:
-                        obs_data.append(traj["obs"])
-                        act_data.append(traj["acts"])
-            elif isinstance(data, dict):
-                if "observations" in data and "actions" in data:
-                    obs_data.append(data["observations"])
-                    act_data.append(data["actions"])
-            else:
-                print(f"Skipping {pkl_file}: Unknown data format {type(data)}")
-        except Exception as e:
-            print(f"Error loading {pkl_file}: {e}")
-
-    if len(obs_data) == 0:
-        raise ValueError("No valid data loaded from provided path.")
-    obs_data = np.concatenate(obs_data, axis=0)
-    act_data = np.concatenate(act_data, axis=0)
-    print(f"Total loaded samples: {len(obs_data)}")
-    return obs_data, act_data
 
 
 def evaluate_policy(policy, args, device):
@@ -126,7 +86,7 @@ def main(args):
     print(f"TensorBoard logging to: {log_dir}")
     os.makedirs(args.save_dir, exist_ok=True)
 
-    obs_data, act_data = load_expert_data(args.expert_data_path)
+    obs_data, act_data = load_expert_pkl(args.expert_data_path)
     obs_tensor = torch.FloatTensor(obs_data)
     act_tensor = torch.FloatTensor(act_data)
     dataset = TensorDataset(obs_tensor, act_tensor)

train_magail.py

@@ -9,7 +9,7 @@ import argparse
 import signal
 import sys
 from torch.utils.data import DataLoader
-from dataset.magail_dataset import MAGAILExpertDataset
+from dataset.loader import MAGAILExpertDataset
 from Env.bc_env import BCScenarioEnv
 
 # --- Networks ---

visualize_bc.py

@@ -1,17 +0,0 @@
-"""
-Thin wrapper: forwards to scripts/visualize_trained_policy.py --policy_type bc.
-Use: python visualize_bc.py [--model_path models/bc/policy_best.pt] [other args...]
-Or call directly: python scripts/visualize_trained_policy.py --policy_type bc --model_path models/bc/policy_best.pt
-"""
-import subprocess
-import sys
-import os
-
-def main():
-    script_dir = os.path.dirname(os.path.abspath(__file__))
-    script = os.path.join(script_dir, "scripts", "visualize_trained_policy.py")
-    cmd = [sys.executable, script, "--policy_type", "bc"] + sys.argv[1:]
-    sys.exit(subprocess.run(cmd).returncode)
-
-if __name__ == "__main__":
-    main()