BC算法实现

Algorithm/bc.py

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

Env/bc_env.py

@@ -0,0 +1,64 @@
+from Env.scenario_env import MultiAgentScenarioEnv
+import numpy as np
+
+class BCScenarioEnv(MultiAgentScenarioEnv):
+    """
+    Environment for Behavior Cloning Evaluation.
+    Uses the same 45-dim observation as ExpertReplayEnv:
+    - Ego State (5): x, y, vx, vy, heading
+    - Neighbors (40): 10 nearest * (rel_x, rel_y, vx, vy)
+    """
+    def _get_all_obs(self):
+        # Implement custom observation: 30m range, 10 nearest vehicles
+        obs_dict = {}
+        
+        for agent_id, vehicle in self.controlled_agents.items():
+            # 1. Ego State
+            ego_state = [
+                vehicle.position[0], vehicle.position[1],
+                vehicle.velocity[0], vehicle.velocity[1],
+                vehicle.heading_theta
+            ]
+            
+            # 2. Neighbors
+            neighbors = []
+            # Iterate through all vehicles in the engine
+            candidates = []
+            # Use engine.agent_manager.active_agents to find neighbors
+            # Note: This includes background vehicles if they are in active_agents
+            for other_id, other_vehicle in self.engine.agent_manager.active_agents.items():
+                if other_id == agent_id:
+                    continue
+                
+                # Check if vehicle is valid/active
+                # (MetaDrive manages active_agents, so they should be active)
+                
+                dist = np.linalg.norm(vehicle.position - other_vehicle.position)
+                if dist < 30.0:
+                    candidates.append((dist, other_vehicle))
+            
+            # Sort by distance
+            candidates.sort(key=lambda x: x[0])
+            
+            # Take top 10
+            top_10 = candidates[:10]
+            
+            neighbor_feats = []
+            for _, neighbor in top_10:
+                neighbor_feats.extend([
+                    neighbor.position[0] - vehicle.position[0], # Relative pos
+                    neighbor.position[1] - vehicle.position[1],
+                    neighbor.velocity[0], # Absolute vel
+                    neighbor.velocity[1]
+                ])
+                
+            # Pad if < 10
+            missing = 10 - len(top_10)
+            if missing > 0:
+                neighbor_feats.extend([0.0] * (4 * missing))
+                
+            # Flatten
+            obs = np.array(ego_state + neighbor_feats, dtype=np.float32)
+            obs_dict[agent_id] = obs
+            
+        return obs_dict

Env/scenario_env.py

@@ -100,6 +100,33 @@ 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
+        # Instead of manually clearing, we just let the engine handle it, but we might need to 
+        # ensure no stale references in managers.
+        
+        # The error "KeyError" in clear_objects usually means we are trying to clear an object 
+        # that is already gone from _spawned_objects but still tracked by a manager.
+        
+        # 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()

README.md

@@ -1,98 +1,96 @@
 # MAGAIL4AutoDrive
 
-> 基于多智能体生成对抗模仿学习(MAGAIL)的自动驾驶训练系统 | MetaDrive + Waymo Open Motion Dataset
+基于 **MetaDrive** 仿真器和 **Waymo Open Motion Dataset** 的自动驾驶多智能体模仿学习（MAGAIL）与行为克隆（BC）训练系统。
 
-本项目利用 Waymo 真实驾驶数据，通过 MetaDrive 仿真环境构建专家回放系统，提取车辆状态与动作，用于训练多智能体模仿学习算法 (MAGAIL)。
+本项目旨在从真实的 Waymo 驾驶数据中提取专家轨迹，并通过模仿学习（Imitation Learning）训练能够适应复杂交互场景的自动驾驶策略。
 
-## 📁 核心模块
+## 目录结构
 
-*   **`Env/expert_replay_env.py`**: 专家回放环境。核心类 `ExpertReplayEnv`，负责读取 Waymo 轨迹，计算逆动力学动作，并过滤非道路/静态车辆。
+```text
-*   **`Env/inverse_dynamics.py`**: 逆动力学模块。根据车辆位置和航向计算油门、刹车和转向动作。
+MAGAIL4AutoDrive/
-*   **`scripts/generate_expert_data.py`**: 数据收集脚本。批量运行场景并保存训练数据。
+├── Algorithm/                 # 强化学习与模仿学习算法实现
-*   **`scripts/visualize_replay.py`**: 可视化脚本。用于观察回放效果和数据质量。
+│   ├── policy.py              # 基础策略网络 (MLP 等)
-
+│   ├── ppo.py                 # PPO 算法实现
-***
+│   ├── magail.py              # MAGAIL 算法核心逻辑
-
+│   ├── disc.py                # 判别器 (Discriminator) 网络
-## 🚀 1. 数据收集
+│   └── ...
-
+├── Env/                       # 仿真环境封装 (MetaDrive Wrapper)
-### 生成专家数据
+│   ├── bc_env.py              # BCScenarioEnv，45 维观测（BC/MAGAIL 共用）
-使用 `generate_expert_data.py` 脚本从 Waymo 数据集中批量提取 (State, Action) 对。
+│   ├── scenario_env.py       # 多智能体基础场景环境
-
+│   ├── expert_replay_env.py  # 专家轨迹回放环境（数据生成与回放）
-```bash
+│   ├── inverse_dynamics.py   # 逆动力学模块 (轨迹 -> 动作)
-# 设置 Python 路径
+│   ├── simple_idm_policy.py  # ConstantVelocityPolicy 占位策略
-export PYTHONPATH=$PYTHONPATH:.:./metadrive
+│   └── ...
-
+├── dataset/                   # 数据集加载器
-# 运行生成脚本
+│   ├── expert_dataset.py      # 通用专家数据加载类
-# --data_dir: Waymo 数据路径 (建议使用 exp_filtered)
+│   └── magail_dataset.py      # MAGAIL 训练专用数据加载器
-# --output_dir: 结果保存路径
+├── scripts/                   # 工具脚本（数据、回放、可视化、分析）
-# --num_scenarios: 要处理的场景数量
+│   ├── generate_expert_data.py    # 从 Waymo 生成专家 (obs, act) pkl
-python scripts/generate_expert_data.py \
+│   ├── visualize_replay.py       # 原始专家数据回放
-    --data_dir data/exp_filtered \
+│   ├── visualize_trained_policy.py # BC/MAGAIL 策略可视化统一入口
-    --output_dir data/training_data \
+│   ├── analyze_expert_data.py    # 数据分布分析
-    --num_scenarios 100 \
+│   ├── launch_tensorboard.py     # 启动 TensorBoard
-    --start_index 0
+│   ├── README.md                  # 脚本用法说明
+│   └── ...
+├── data/                      # 数据目录（相对路径）
+│   ├── exp_filtered/          # Waymo 场景数据
+│   ├── training_data/        # 专家 pkl 输出（generate_expert_data）
+│   └── trajectories/         # 其他轨迹 pkl（如 expert_dataset 输出）
+├── models/                    # 模型保存目录（相对路径）
+│   ├── bc/                    # BC 模型 (.pt)
+│   └── magail/                # MAGAIL 模型 (*_actor.pth, *_critic.pth)
+├── logs/                      # 训练日志 (TensorBoard)
+│   ├── bc/
+│   └── magail/
+├── train_bc.py                # [根目录] BC 训练
+├── train_magail.py            # [根目录] MAGAIL 训练
+├── visualize_bc.py            # [根目录] BC 可视化薄包装 -> scripts/visualize_trained_policy.py
+└── README.md
 ```
 
-**生成的 `.pkl` 文件结构**：
+## 路径约定（相对项目根）
-包含一个列表，每个元素是一条车辆轨迹（Trajectory Dictionary）：
-*   `obs`: `(T, 45)` - 观测矩阵。包含 Ego 状态 (5维) + 10辆邻居车相对信息 (40维)。
-*   `acts`: `(T, 2)` - 动作矩阵。`[Steering, Accel]`，归一化到 `[-1, 1]`。
-*   `agent_id`: 车辆 ID。
-*   `scenario_id`: 所属场景 ID。
 
-**内置过滤器**：
+- **数据**：Waymo 场景 `data/exp_filtered`；专家 pkl `data/training_data`；其他轨迹 `data/trajectories`
-脚本会自动过滤掉以下无效车辆：
+- **模型**：BC `models/bc/`，MAGAIL `models/magail/`
-1.  **非道路车辆**：始终在停车场或路外行驶的车辆。
+- **日志**：TensorBoard 写入 `logs/bc/`、`logs/magail/`
-2.  **静态车辆**：全称移动距离小于 5米 且速度从未超过 1m/s 的车辆（作为背景流存在，不收集数据）。
 
----
+所有默认路径均为相对项目根，便于在不同设备上复用。
 
-## 🔍 2. 数据可视化与验证
+## 核心工作流
 
-### 回放可视化
+### 1. 数据准备
-使用 `visualize_replay.py` 直观地观察回放效果，确认车辆行为是否自然，以及过滤逻辑是否生效。
+使用 `scripts/generate_expert_data.py` 将 Waymo 数据转换为训练用 `.pkl`，输出到 `data/training_data/`。
 
 ```bash
-# 运行可视化
+python scripts/generate_expert_data.py --data_dir data/exp_filtered --output_dir data/training_data --num_scenarios 100
-# --horizon: 回放的最大步数 (Waymo 场景通常为 90 或 198 步)
-python scripts/visualize_replay.py \
-    --data_dir data/exp_filtered \
-    --start_index 0 \
-    --num_scenarios 1 \
-    --horizon 200
 ```
 
-**观察要点**：
+### 2. 行为克隆 (BC)
-*   **受控车辆 (Controlled Agents)**：控制台会显示数量（如 `Controlled agents: 2`）。这些是真正产生数据的车辆。
+- **训练**：`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`
 
-### 数据分析
+### 3. 多智能体对抗模仿学习 (MAGAIL)
-使用 `analyze_expert_data.py` 查看生成数据的统计分布。
+- **训练**：`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`
 
-```bash
+### 4. 策略可视化统一入口
-python scripts/analyze_expert_data.py --data_path data/training_data/expert_data_0_100.pkl
+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)。
-```
 
----
+## 文件与模块职责
 
-## 🧠 3. 模型训练 (Next Steps)
+### 根目录脚本
+- **train_bc.py**：BC 训练，加载 `data/training_data` 下 pkl，模型与日志写入 `models/bc/`、`logs/bc/`
+- **train_magail.py**：MAGAIL 训练，环境使用 `BCScenarioEnv`（45 维），模型与日志写入 `models/magail/`、`logs/magail/`
+- **visualize_bc.py**：薄包装，调用 `scripts/visualize_trained_policy.py --policy_type bc`
 
-有了 `data/training_data/` 下的专家数据后，您可以开始训练 MAGAIL 模型。
+### Env 模块
+- **Env/bc_env.py**：`BCScenarioEnv`，45 维观测（Ego 5 维 + 10 邻居×4 维），BC 与 MAGAIL 训练/评估共用
+- **Env/scenario_env.py**：`MultiAgentScenarioEnv` 基类，Waymo 场景加载与步进
+- **Env/expert_replay_env.py**：专家轨迹回放与逆动力学动作，供 `generate_expert_data.py` 与回放可视化
+- **Env/inverse_dynamics.py**：轨迹 → 油门/转向动作
 
-### 训练流程
+### Algorithm 模块
-1.  **加载数据**：使用 `dataset/expert_dataset.py` 中的 `ExpertDataset` 类加载 `.pkl` 数据。
+- **Algorithm/policy.py**：`StateIndependentPolicy`，BC 使用的 MLP 策略
-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)。
 
-### 推荐配置
+### scripts 目录
-*   **Observation**: 45维 (Ego + 10 Neighbors)
+工具脚本用途与用法见 [scripts/README.md](scripts/README.md)。
-*   **Action**: 2维 Continuous (Steering, Accel)
-*   **Horizon**: 200 steps
-*   **Batch Size**: 1024+ (多智能体环境下数据量很大)

dataset/expert_dataset.py

@@ -244,6 +244,7 @@ class ExpertTrajectoryDataset(Dataset):
                 print(f"  观测维度: {obs_dim} (应为107)")
         
         if save_path:
+            os.makedirs(os.path.dirname(save_path), exist_ok=True)
             with open(save_path, "wb") as f:
                 pickle.dump({
                     "trajectories": all_trajectories,
@@ -282,7 +283,7 @@ if __name__ == "__main__":
         trajectories, observations = ExpertTrajectoryDataset.collect_with_full_obs(
             env_config,
             num_scenarios=10,
-            save_path="./expert_trajectories_full.pkl"
+            save_path="data/trajectories/expert_trajectories_full.pkl"
         )
         
         if len(trajectories) > 0:

scripts/README.md

@@ -0,0 +1,82 @@
+# scripts 工具脚本说明
+
+本目录包含数据生成、回放、可视化与分析等工具脚本。训练脚本（`train_bc.py`、`train_magail.py`）位于项目根目录。
+
+## 路径约定（相对项目根）
+
+- **数据**：`data/exp_filtered`（Waymo 场景）、`data/training_data`（专家 pkl 输出）
+- **模型**：`models/bc/`（BC）、`models/magail/`（MAGAIL）
+- **日志**：`logs/bc/`、`logs/magail/`（TensorBoard）
+
+---
+
+## 脚本列表与用法
+
+### 数据生成
+
+| 脚本 | 用途 | 用法示例 |
+|------|------|----------|
+| [generate_expert_data.py](generate_expert_data.py) | 从 Waymo 数据生成专家 (obs, act) 的 pkl | `python scripts/generate_expert_data.py --data_dir data/exp_filtered --output_dir data/training_data --num_scenarios 100` |
+
+**常用参数**：`--data_dir`（默认 `data/exp_filtered`）、`--output_dir`（默认 `data/training_data`）、`--start_index`、`--num_scenarios`。
+
+---
+
+### 回放与可视化
+
+| 脚本 | 用途 | 用法示例 |
+|------|------|----------|
+| [visualize_replay.py](visualize_replay.py) | 原始专家轨迹回放（ExpertReplayEnv） | `python scripts/visualize_replay.py --data_dir data/exp_filtered --num_scenarios 1 --horizon 200` |
+| [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`。
+
+**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
+```
+
+**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
+```
+
+**自动推断类型**（根据 `--model_path` 扩展名：`.pt` → BC，否则 → MAGAIL）：
+```bash
+python scripts/visualize_trained_policy.py --model_path models/bc/policy_best.pt
+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`
+
+**参数**：`--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`）。
+
+---
+
+### 数据分析与检查
+
+| 脚本 | 用途 | 用法示例 |
+|------|------|----------|
+| [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 接口可能不一致，可选使用 |
+
+---
+
+### 其他
+
+| 脚本 | 用途 | 用法示例 |
+|------|------|----------|
+| [launch_tensorboard.py](launch_tensorboard.py) | 启动 TensorBoard | `python scripts/launch_tensorboard.py --logdir logs`（或 `logs/bc` / `logs/magail`） |
+
+---
+
+## 与训练流程的对应关系
+
+1. **数据准备**：`generate_expert_data.py` → 输出到 `data/training_data/*.pkl`
+2. **BC 训练**：根目录 `train_bc.py` → 模型保存到 `models/bc/`，日志到 `logs/bc/`
+3. **MAGAIL 训练**：根目录 `train_magail.py` → 模型保存到 `models/magail/`，日志到 `logs/magail/`
+4. **可视化**：`visualize_trained_policy.py`（或根目录 `visualize_bc.py` 仅 BC）→ 从 `models/bc` 或 `models/magail` 加载模型，数据目录默认 `data/exp_filtered`

scripts/README_visualize.md

@@ -1,113 +0,0 @@
-# 模型可视化脚本使用说明
-
-## 功能
-使用训练好的MAGAIL模型在环境中运行，并生成俯瞰效果图（top-down view）。
-
-## 使用方法
-
-### 基本用法
-
-```bash
-python scripts/visualize_trained_model.py \
-    --model_dir runs/magail_0113 \
-    --episode 1250 \
-    --data_dir data/exp_filtered \
-    --num_scenarios 1 \
-    --output_dir visualizations
-```
-
-### 参数说明
-
-- `--model_dir`: 模型保存目录（例如：`runs/magail_0113`）
-- `--episode`: 要加载的episode编号（例如：`1250`）
-- `--data_dir`: Waymo数据目录（默认：`data/exp_filtered`）
-- `--start_index`: 起始场景索引（默认：`0`）
-- `--num_scenarios`: 要运行的场景数量（默认：`1`）
-- `--horizon`: 每个episode的最大步数（默认：`200`）
-- `--output_dir`: 输出图像保存目录（默认：`visualizations`）
-- `--save_all_frames`: 保存所有帧（否则按间隔保存）
-- `--save_interval`: 保存帧的间隔，当不使用`--save_all_frames`时生效（默认：`10`）
-- `--gif_duration`: GIF每帧持续时间（毫秒），默认50ms（20fps）。值越小，GIF播放越快
-
-### 示例
-
-#### 1. 查看最新训练的模型（episode 1250）
-```bash
-python scripts/visualize_trained_model.py \
-    --model_dir runs/magail_0113 \
-    --episode 1250 \
-    --num_scenarios 3 \
-    --output_dir visualizations/episode_1250
-```
-
-#### 2. 保存所有帧（用于制作视频）
-```bash
-python scripts/visualize_trained_model.py \
-    --model_dir runs/magail_0113 \
-    --episode 1250 \
-    --save_all_frames \
-    --output_dir visualizations/episode_1250_all_frames
-```
-
-#### 3. 每5步保存一帧
-```bash
-python scripts/visualize_trained_model.py \
-    --model_dir runs/magail_0113 \
-    --episode 1250 \
-    --save_interval 5 \
-    --output_dir visualizations/episode_1250_sparse
-```
-
-#### 4. 生成更快的GIF（30fps）
-```bash
-python scripts/visualize_trained_model.py \
-    --model_dir runs/magail_0113 \
-    --episode 1250 \
-    --gif_duration 33 \
-    --output_dir visualizations/episode_1250
-```
-
-## 输出
-
-脚本会在指定的输出目录中创建以下文件：
-- `scenario_{idx}.gif`: **场景动画GIF**（主要输出）
-- `scenario_{idx}_step_{step:04d}.png`: 每个保存步骤的俯瞰图（可选）
-- `scenario_{idx}_final.png`: 每个场景的最终状态图
-
-### GIF格式
-- 分辨率：1600x900
-- 格式：GIF动画
-- 包含完整的场景运行过程
-- 显示场景编号、步数、智能体数量和奖励信息
-- 默认帧率：20fps（可通过`--gif_duration`调整）
-
-### 图像格式
-- 分辨率：1600x900
-- 格式：PNG
-- 包含语义地图和车辆轨迹
-
-## 注意事项
-
-1. **GPU要求**: 脚本需要CUDA支持，如果没有GPU会自动使用CPU（速度较慢）
-2. **渲染模式**: 使用MetaDrive的top-down渲染模式，会弹出窗口显示实时渲染
-3. **内存占用**: 如果保存所有帧，会占用较多磁盘空间
-4. **场景数据**: 确保`--data_dir`指向正确的Waymo数据目录
-
-## 故障排除
-
-### 模型文件不存在
-```
-FileNotFoundError: 模型文件不存在: runs/magail_0113/model_1250_actor.pth
-```
-**解决**: 检查模型目录和episode编号是否正确
-
-### 场景数据不存在
-```
-ValueError: Data directory not found
-```
-**解决**: 确保`--data_dir`指向正确的数据目录
-
-### 渲染失败
-如果遇到渲染相关错误，可以尝试：
-- 降低`film_size`参数（在脚本中修改）
-- 使用无头模式（需要修改脚本）

scripts/generate_expert_data.py

@@ -153,8 +153,8 @@ def generate_data(args):
 
 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("--data_dir", type=str, default="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("--output_dir", type=str, default="data/training_data", help="Output directory")
     parser.add_argument("--start_index", type=int, default=0)
     parser.add_argument("--num_scenarios", type=int, default=10)
     

scripts/visualize_trained_policy.py

@@ -1,128 +1,159 @@
+"""
+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
-import time
 
 # 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 train_magail import Actor, MAGAILScenarioEnv
+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):
-    # 1. Load Environment
+    policy_type = (args.policy_type or "auto").lower()
-    data_path = os.path.abspath(args.data_dir)
+    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,  # Visualisation enabled
+        "use_render": True,
         "sequential_seed": True,
         "start_scenario_index": args.start_index,
         "num_scenarios": args.num_scenarios,
         "log_level": 40,
     }
 
-    print("Initializing MAGAILScenarioEnv...")
+    print(f"Initializing BCScenarioEnv (policy_type={policy_type})...")
     try:
-        env = MAGAILScenarioEnv(config=env_config, agent2policy={})
+        env = BCScenarioEnv(env_config, agent2policy={})
     except Exception as e:
         print(f"Error init env: {e}. Trying to close lingering engine...")
         try:
             close_engine()
-        except:
+        except Exception:
             pass
-        env = MAGAILScenarioEnv(config=env_config, agent2policy={})
+        env = BCScenarioEnv(env_config, agent2policy={})
 
-    # 2. Load Model
     state_dim = 45
     action_dim = 2
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-    actor = Actor(state_dim, action_dim).cuda()
+    model_path = _resolve_model_path(args.model_path, policy_type)
-    
-    model_path = args.model_path
-    if not os.path.exists(model_path):
-        # Try to find it in runs/
-        potential_path = os.path.join("runs", "magail_production", model_path)
-        if os.path.exists(potential_path):
-            model_path = potential_path
-        else:
-            # Try appending _actor.pth
-            potential_path = model_path + "_actor.pth"
-            if os.path.exists(potential_path):
-                model_path = potential_path
-            else:
-                raise ValueError(f"Model path {args.model_path} not found.")
-                
     print(f"Loading model from {model_path}...")
-    actor.load_state_dict(torch.load(model_path))
-    actor.eval()
 
-    # 3. Run Loop
+    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} ---")
-            
-            # Reset
             try:
-                # Use sequential seed logic or specific seed?
-                # ExpertReplayEnv/ScenarioEnv logic: seed matches scenario index if configured right
                 obs_dict = env.reset(seed=i)
             except Exception as e:
                 print(f"Error resetting {i}: {e}. Skipping.")
-                # Try soft reset
                 try:
                     close_engine()
-                    env = MAGAILScenarioEnv(config=env_config, agent2policy={})
+                    env = BCScenarioEnv(env_config, agent2policy={})
-                except:
+                except Exception:
                     pass
                 continue
 
             print(f"Scenario loaded. Controlled agents: {len(obs_dict)}")
-            
             step_count = 0
+            episode_reward = 0.0
+
             while True:
                 actions = {}
-                # Inference
+                agent_ids = list(obs_dict.keys())
-                for agent_id, obs in obs_dict.items():
+                obs_list = [obs_dict[aid] for aid in agent_ids]
-                    # Preprocess obs: (45,) -> (1, 45) tensor
+                obs_tensor = torch.FloatTensor(np.array(obs_list)).to(device)
-                    obs_tensor = torch.FloatTensor(obs).unsqueeze(0).cuda()
+
-                    with torch.no_grad():
+                with torch.no_grad():
+                    if policy_type == "bc":
+                        actions_np = policy(obs_tensor).cpu().numpy()
+                    else:
                         dist = actor(obs_tensor)
-                        # Deterministic action for viz? Or sample?
-                        # Usually deterministic (mean) is better for checking performance
-                        # But training uses sample.
                         if args.deterministic:
-                            action = torch.tanh(dist.mean) # Use mean of Gaussian
+                            actions_np = torch.tanh(dist.mean).cpu().numpy()
                         else:
-                            pre_tanh = dist.sample()
+                            actions_np = torch.tanh(dist.sample()).cpu().numpy()
-                            action = torch.tanh(pre_tanh)
 
-                    actions[agent_id] = action.cpu().numpy().flatten()
+                for idx, aid in enumerate(agent_ids):
+                    actions[aid] = actions_np[idx].flatten()
 
-                # Step
                 obs_dict, rewards, dones, infos = env.step(actions)
+                episode_reward += sum(rewards.values())
 
-                # Render
                 env.render(
                     mode="top_down",
                     text={
                         "Scenario": i,
                         "Step": step_count,
-                        "Agents": len(obs_dict)
+                        "Agents": len(obs_dict),
-                    }
+                        "Total Reward": f"{episode_reward:.2f}",
+                    },
                 )
-                
                 step_count += 1
-                # time.sleep(0.02) # Slow down if needed
 
                 if dones["__all__"] or step_count >= args.horizon:
-                    print(f"Scenario finished at step {step_count}")
+                    print(f"Scenario finished at step {step_count}, reward {episode_reward:.2f}")
                     break
 
     except KeyboardInterrupt:
@@ -130,14 +161,29 @@ def visualize_model(args):
     finally:
         env.close()
 
+
 if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
+    parser = argparse.ArgumentParser(
-    parser.add_argument("--model_path", type=str, required=True, help="Path to actor model pth (e.g. runs/magail_production/model_50_actor.pth)")
+        description="Visualize BC or MAGAIL trained policy in 45-dim scenario env."
-    parser.add_argument("--data_dir", type=str, default="data/exp_filtered")
+    )
+    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="Use mean action instead of sampling")
+    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

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

train_magail.py

@@ -10,6 +10,7 @@ import signal
 import sys
 from torch.utils.data import DataLoader
 from dataset.magail_dataset import MAGAILExpertDataset
+from Env.bc_env import BCScenarioEnv
 
 # --- Networks ---
 
@@ -161,58 +162,10 @@ class PPO:
         torch.save(self.actor.state_dict(), checkpoint_path + "_actor.pth")
         torch.save(self.critic.state_dict(), checkpoint_path + "_critic.pth")
 
-from Env.scenario_env import MultiAgentScenarioEnv
-
-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
-
 # --- Training Loop ---
 
 def train(args):
-    # 1. Setup Environment (Dummy for now, usually you run simulation here)
+    # 1. Setup Environment (45-dim obs via BCScenarioEnv)
-    # 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,
@@ -255,12 +208,7 @@ def train(args):
                 yield batch
     expert_iter = cycle(expert_loader)
     
-    # 4. Initialize Env
+    # 4. Initialize Env (BCScenarioEnv provides 45-dim obs)
-    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? 
@@ -316,7 +264,7 @@ def train(args):
     #             obs_dict[agent_id] = obs
     #         return obs_dict
     
-    env = MAGAILScenarioEnv(config=env_config, agent2policy={}) # Pass empty dict if we control all externally
+    env = BCScenarioEnv(env_config, agent2policy={})  # 45-dim obs
     
     print("Starting training...")
     
@@ -401,7 +349,7 @@ def train(args):
             import gc
             gc.collect()
             
-            env = MAGAILScenarioEnv(config=env_config, agent2policy={})
+            env = BCScenarioEnv(env_config, agent2policy={})
             obs_dict = env.reset(seed=seed)
         
         episode_reward = 0
@@ -575,7 +523,7 @@ def train(args):
         print(f"Episode {i_episode}: Disc Loss {disc_loss.item():.4f} | PPO Loss {ppo_loss:.4f} | Mean Reward {np.mean(all_gail_rewards):.4f}")
         
         if i_episode % 50 == 0:
-            ppo_agent.save(os.path.join(args.log_dir, f"model_{i_episode}"))
+            ppo_agent.save(os.path.join(args.save_dir, f"model_{i_episode}"))
 
     env.close()
     if writer:
@@ -588,11 +536,13 @@ if __name__ == '__main__':
     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")
+    parser.add_argument("--log_dir", type=str, default="logs/magail", help="TensorBoard log directory")
+    parser.add_argument("--save_dir", type=str, default="models/magail", help="Directory to save model checkpoints")
     
     args = parser.parse_args()
     
-    # Create log dir
+    # Create log dir and save dir
     os.makedirs(args.log_dir, exist_ok=True)
+    os.makedirs(args.save_dir, exist_ok=True)
     
     train(args)

visualize_bc.py

@@ -0,0 +1,17 @@
+"""
+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()