完善项目目录结构
This commit is contained in:
Binary file not shown.
Binary file not shown.
@@ -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
|
||||
# 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()
|
||||
|
||||
# Clear vehicles we spawned via engine.spawn_object() so _object_clean_check() passes
|
||||
ids_to_clear = [v.id for v in self.controlled_agents.values()]
|
||||
if ids_to_clear:
|
||||
self.engine.clear_objects(ids_to_clear)
|
||||
self.controlled_agents.clear()
|
||||
self.controlled_agent_ids.clear()
|
||||
|
||||
self.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:
|
||||
|
||||
49
README.md
49
README.md
@@ -22,12 +22,11 @@ MAGAIL4AutoDrive/
|
||||
│ ├── simple_idm_policy.py # ConstantVelocityPolicy 占位策略
|
||||
│ └── ...
|
||||
├── dataset/ # 数据集加载器
|
||||
│ ├── expert_dataset.py # 通用专家数据加载类
|
||||
│ └── magail_dataset.py # MAGAIL 训练专用数据加载器
|
||||
│ ├── loader.py # 主流水线：load_expert_pkl、MAGAILExpertDataset
|
||||
│ └── expert_dataset.py # 可选 107 维/5 维管线
|
||||
├── scripts/ # 工具脚本（数据、回放、可视化、分析）
|
||||
│ ├── generate_expert_data.py # 从 Waymo 生成专家 (obs, act) pkl
|
||||
│ ├── visualize_replay.py # 原始专家数据回放
|
||||
│ ├── visualize_trained_policy.py # BC/MAGAIL 策略可视化统一入口
|
||||
│ ├── visualize.py # 可视化统一入口（replay / policy / trajectory）
|
||||
│ ├── 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. 策略可视化统一入口
|
||||
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)。
|
||||
### 4. 可视化统一入口
|
||||
可视化统一使用 `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_magail.py**：MAGAIL 训练，环境使用 `BCScenarioEnv`（45 维），模型与日志写入 `models/magail/`、`logs/magail/`
|
||||
- **visualize_bc.py**：薄包装，调用 `scripts/visualize_trained_policy.py --policy_type bc`
|
||||
- **train_bc.py**：BC 训练，从 `dataset.loader` 加载专家 pkl，模型与日志写入 `models/bc/`、`logs/bc/`
|
||||
- **train_magail.py**：MAGAIL 训练，环境使用 `BCScenarioEnv`（45 维），从 `dataset.loader` 加载专家数据，模型与日志写入 `models/magail/`、`logs/magail/`
|
||||
|
||||
### Env 模块
|
||||
- **Env/bc_env.py**：`BCScenarioEnv`，45 维观测（Ego 5 维 + 10 邻居×4 维），BC 与 MAGAIL 训练/评估共用
|
||||
|
||||
103
dataset/loader.py
Normal file
103
dataset/loader.py
Normal file
@@ -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
|
||||
@@ -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
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -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_trained_policy.py](visualize_trained_policy.py) | **BC/MAGAIL 共用**：加载训练好的策略在 45 维场景中可视化 | 见下方「训练策略可视化」小节 |
|
||||
| [visualize.py](visualize.py) | **replay**：场景回放（ExpertReplayEnv）；**policy**：BC/MAGAIL 策略；**trajectory**：专家轨迹 2D 动画 | 见下方 |
|
||||
|
||||
#### 训练策略可视化（visualize_trained_policy.py）
|
||||
**子命令**：
|
||||
|
||||
使用训练好的 **BC** 或 **MAGAIL** 模型在 45 维场景环境中运行，并实时渲染俯瞰图（top-down view）。统一入口：`scripts/visualize_trained_policy.py`。
|
||||
|
||||
**BC 模型**：
|
||||
- **replay**（原始专家轨迹回放）：
|
||||
```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_trained_policy.py --model_path models/magail/model_50_actor.pth
|
||||
python scripts/visualize.py trajectory --data_dir data/exp_filtered --scenario_idx 0
|
||||
```
|
||||
|
||||
**根目录 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`）。
|
||||
**公共参数**：`--data_dir`（默认 `data/exp_filtered`）、`--start_index`、`--num_scenarios`、`--horizon`。policy 模式另有 `--policy_type`（auto/bc/magail）、`--model_path`、`--deterministic`（仅 MAGAIL）。
|
||||
|
||||
---
|
||||
|
||||
@@ -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`
|
||||
|
||||
395
scripts/visualize.py
Normal file
395
scripts/visualize.py
Normal file
@@ -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()
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
44
train_bc.py
44
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
|
||||
|
||||
|
||||
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
|
||||
from dataset.loader import load_expert_pkl
|
||||
|
||||
|
||||
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)
|
||||
|
||||
@@ -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 ---
|
||||
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user