huangfu/MAGAIL4AutoDrive
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

BC算法实现

Browse Source
This commit is contained in:
huangfu
2026-02-02 01:18:18 +08:00
parent 265b0eade1
commit 21c046aef0
25 changed files with 632 additions and 322 deletions
Download Patch File Download Diff File Expand all files Collapse all files

82
scripts/README.md Normal file
View File

@@ -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`

113
scripts/README_visualize.md
View File

@@ -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`参数（在脚本中修改）
- 使用无头模式（需要修改脚本）

4
scripts/generate_expert_data.py
View File

@@ -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("--output_dir", type=str, default="/home/huangfukk/MAGAIL4AutoDrive/data/training", help="Output directory")
parser.add_argument("--data_dir", type=str, default="data/exp_filtered", help="Path to Waymo pickles (or filtered index)")
parser.add_argument("--output_dir", type=str, default="data/training_data", help="Output directory")
parser.add_argument("--start_index", type=int, default=0)
parser.add_argument("--num_scenarios", type=int, default=10)

186
scripts/visualize_trained_policy.py
View File

@@ -1,143 +1,189 @@
"""
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
data_path = os.path.abspath(args.data_dir)
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, # 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
actor = Actor(state_dim, action_dim).cuda()
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.")
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}...")
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={})
except:
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 = {}
# Inference
for agent_id, obs in obs_dict.items():
# Preprocess obs: (45,) -> (1, 45) tensor
obs_tensor = torch.FloatTensor(obs).unsqueeze(0).cuda()
with torch.no_grad():
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)
# 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()
action = torch.tanh(pre_tanh)
actions[agent_id] = action.cpu().numpy().flatten()
# Step
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)
# Render
episode_reward += sum(rewards.values())
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:
print("Interrupted.")
finally:
env.close()
if __name__ == "__main__":
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)")
parser.add_argument("--data_dir", type=str, default="data/exp_filtered")
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="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)