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

更新 .gitignore 和训练脚本，添加可视化脚本

Browse Source
This commit is contained in:
huangfu
2026-01-17 20:24:02 +08:00
parent 4dbea5f0a6
commit 265b0eade1
6 changed files with 453 additions and 1718 deletions
Download Patch File Download Diff File Expand all files Collapse all files

54
.gitignore vendored
View File

@@ -1,3 +1,57 @@
# 日志文件
Env/logs/
*.log
# Python
__pycache__/
*.py[cod]
*$py.class
*.so
.Python
build/
develop-eggs/
dist/
downloads/
eggs/
.eggs/
lib/
lib64/
parts/
sdist/
var/
wheels/
*.egg-info/
.installed.cfg
*.egg
# 虚拟环境
venv/
env/
ENV/
.venv
# IDE
.vscode/
.idea/
*.swp
*.swo
*~
# 数据和模型文件
data/
runs/
*.pkl
*.h5
*.ckpt
*.pth
*.pt
checkpoints/
models/
# 第三方库（如果已安装）
metadrive/
scenarionet/
# 系统文件
.DS_Store
Thumbs.db

113
scripts/README_visualize.md Normal file
View File

@@ -0,0 +1,113 @@
# 模型可视化脚本使用说明
## 功能
使用训练好的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`参数（在脚本中修改）
- 使用无头模式（需要修改脚本）

18
scripts/launch_tensorboard.py Normal file
View File

@@ -0,0 +1,18 @@
import sys
import types
import os
# Mock imghdr module for Python 3.13 compatibility
# TensorBoard depends on imghdr which was removed in Python 3.13
if sys.version_info >= (3, 13):
if 'imghdr' not in sys.modules:
imghdr_mock = types.ModuleType('imghdr')
imghdr_mock.what = lambda filename, h=None: None
# Mock tests list which tensorboard appends to
imghdr_mock.tests = []
sys.modules['imghdr'] = imghdr_mock
from tensorboard import main as tb_main
if __name__ == '__main__':
sys.exit(tb_main.run_main())

143
scripts/visualize_trained_policy.py Normal file
View File

@@ -0,0 +1,143 @@
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 metadrive.engine.engine_utils import close_engine
def visualize_model(args):
# 1. Load Environment
data_path = os.path.abspath(args.data_dir)
env_config = {
"data_directory": data_path,
"is_multi_agent": True,
"num_controlled_agents": 3,
"horizon": args.horizon,
"use_render": True, # Visualisation enabled
"sequential_seed": True,
"start_scenario_index": args.start_index,
"num_scenarios": args.num_scenarios,
"log_level": 40,
}
print("Initializing MAGAILScenarioEnv...")
try:
env = MAGAILScenarioEnv(config=env_config, agent2policy={})
except Exception as e:
print(f"Error init env: {e}. Trying to close lingering engine...")
try:
close_engine()
except:
pass
env = MAGAILScenarioEnv(config=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.")
print(f"Loading model from {model_path}...")
actor.load_state_dict(torch.load(model_path))
actor.eval()
# 3. Run Loop
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:
pass
continue
print(f"Scenario loaded. Controlled agents: {len(obs_dict)}")
step_count = 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():
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
else:
pre_tanh = dist.sample()
action = torch.tanh(pre_tanh)
actions[agent_id] = action.cpu().numpy().flatten()
# Step
obs_dict, rewards, dones, infos = env.step(actions)
# Render
env.render(
mode="top_down",
text={
"Scenario": i,
"Step": step_count,
"Agents": len(obs_dict)
}
)
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}")
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.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")
args = parser.parse_args()
visualize_model(args)

173
train_magail.py
View File

@@ -79,18 +79,30 @@ class PPO:
self.K_epochs = K_epochs
self.mse_loss = nn.MSELoss()
def _log_prob_from_dist(self, dist, pre_tanh_action):
# Tanh-squashed Gaussian log-prob with correction term.
log_prob = dist.log_prob(pre_tanh_action)
correction = torch.log(1 - torch.tanh(pre_tanh_action) ** 2 + 1e-6)
return (log_prob - correction).sum(dim=-1)
def select_action(self, state):
with torch.no_grad():
state = torch.FloatTensor(state).cuda()
dist = self.actor(state)
action = dist.sample()
action_logprob = dist.log_prob(action).sum(dim=-1)
return action.cpu().numpy(), action_logprob.cpu().numpy()
pre_tanh_action = dist.sample()
action = torch.tanh(pre_tanh_action)
action_logprob = self._log_prob_from_dist(dist, pre_tanh_action)
return (
action.cpu().numpy(),
action_logprob.cpu().numpy(),
pre_tanh_action.cpu().numpy()
)
def update(self, memory):
# Convert memory to tensors
states = torch.FloatTensor(np.array(memory['states'])).cuda()
actions = torch.FloatTensor(np.array(memory['actions'])).cuda()
pre_tanh_actions = torch.FloatTensor(np.array(memory['pre_tanh_actions'])).cuda()
logprobs = torch.FloatTensor(np.array(memory['logprobs'])).cuda()
rewards = torch.FloatTensor(np.array(memory['rewards'])).cuda()
next_states = torch.FloatTensor(np.array(memory['next_states'])).cuda()
@@ -124,7 +136,7 @@ class PPO:
for _ in range(self.K_epochs):
# Evaluating old actions and values :
dist = self.actor(states)
action_logprobs = dist.log_prob(actions).sum(dim=-1)
action_logprobs = self._log_prob_from_dist(dist, pre_tanh_actions)
dist_entropy = dist.entropy().sum(dim=-1)
state_values = self.critic(states).squeeze()
@@ -149,6 +161,49 @@ 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):
@@ -220,47 +275,47 @@ def train(args):
# We need to inject that same logic into the training env, OR
# subclass MultiAgentScenarioEnv in the training script to override observation.
class MAGAILScenarioEnv(MultiAgentScenarioEnv):
def _get_all_obs(self):
# Same logic as ExpertReplayEnv to ensure compatibility
obs_dict = {}
for agent_id, vehicle in self.controlled_agents.items():
# 1. Ego State
ego_state = [
vehicle.position[0], vehicle.position[1],
vehicle.velocity[0], vehicle.velocity[1],
vehicle.heading_theta
]
# 2. Neighbors
candidates = []
for other_id, other_vehicle in self.engine.agent_manager.active_agents.items():
if other_id == agent_id:
continue
dist = np.linalg.norm(vehicle.position - other_vehicle.position)
if dist < 30.0:
candidates.append((dist, other_vehicle))
candidates.sort(key=lambda x: x[0])
top_10 = candidates[:10]
neighbor_feats = []
for _, neighbor in top_10:
neighbor_feats.extend([
neighbor.position[0] - vehicle.position[0],
neighbor.position[1] - vehicle.position[1],
neighbor.velocity[0],
neighbor.velocity[1]
])
missing = 10 - len(top_10)
if missing > 0:
neighbor_feats.extend([0.0] * (4 * missing))
obs = np.array(ego_state + neighbor_feats, dtype=np.float32)
obs_dict[agent_id] = obs
return obs_dict
# class MAGAILScenarioEnv(MultiAgentScenarioEnv):
# def _get_all_obs(self):
# # Same logic as ExpertReplayEnv to ensure compatibility
# obs_dict = {}
# for agent_id, vehicle in self.controlled_agents.items():
# # 1. Ego State
# ego_state = [
# vehicle.position[0], vehicle.position[1],
# vehicle.velocity[0], vehicle.velocity[1],
# vehicle.heading_theta
# ]
#
# # 2. Neighbors
# candidates = []
# for other_id, other_vehicle in self.engine.agent_manager.active_agents.items():
# if other_id == agent_id:
# continue
# dist = np.linalg.norm(vehicle.position - other_vehicle.position)
# if dist < 30.0:
# candidates.append((dist, other_vehicle))
#
# candidates.sort(key=lambda x: x[0])
# top_10 = candidates[:10]
#
# neighbor_feats = []
# for _, neighbor in top_10:
# neighbor_feats.extend([
# neighbor.position[0] - vehicle.position[0],
# neighbor.position[1] - vehicle.position[1],
# neighbor.velocity[0],
# neighbor.velocity[1]
# ])
#
# missing = 10 - len(top_10)
# if missing > 0:
# neighbor_feats.extend([0.0] * (4 * missing))
#
# obs = np.array(ego_state + neighbor_feats, dtype=np.float32)
# obs_dict[agent_id] = obs
# return obs_dict
env = MAGAILScenarioEnv(config=env_config, agent2policy={}) # Pass empty dict if we control all externally
print("Starting training...")
@@ -277,7 +332,15 @@ def train(args):
for i_episode in range(args.max_episodes):
# --- 1. Collect Rollouts (Interaction) ---
memory = {'states': [], 'actions': [], 'logprobs': [], 'rewards': [], 'next_states': [], 'dones': []}
memory = {
'states': [],
'actions': [],
'pre_tanh_actions': [],
'logprobs': [],
'rewards': [],
'next_states': [],
'dones': []
}
# Prepare seed
available_scenarios = env.config["num_scenarios"]
@@ -349,6 +412,7 @@ def train(args):
# Select actions for all agents
actions = {}
action_logprobs = {}
pre_tanh_actions = {}
# obs_dict: {agent_id: obs}
# MultiAgentScenarioEnv usually returns a dict {agent_id: obs}
@@ -386,9 +450,10 @@ def train(args):
obs_dict = new_obs_dict
for agent_id, obs in obs_dict.items():
act, logprob = ppo_agent.select_action(obs) # Select action returns numpy
act, logprob, pre_tanh = ppo_agent.select_action(obs) # Select action returns numpy
actions[agent_id] = act.flatten() # (2,)
action_logprobs[agent_id] = logprob # scalar
pre_tanh_actions[agent_id] = pre_tanh.flatten()
# Step Env
next_obs_dict, rewards, dones, infos = env.step(actions)
@@ -398,6 +463,7 @@ def train(args):
if agent_id in actions:
memory['states'].append(obs)
memory['actions'].append(actions[agent_id])
memory['pre_tanh_actions'].append(pre_tanh_actions[agent_id])
memory['logprobs'].append(action_logprobs[agent_id])
# Store standard environmental reward for logging (not used for update in GAIL)
@@ -407,7 +473,7 @@ def train(args):
# Next state
if agent_id in next_obs_dict:
memory['next_states'].append(next_obs_dict[agent_id])
memory['dones'].append(False)
memory['dones'].append(dones.get("__all__", False))
else:
# Agent finished/vanished
# We need a dummy next state or handle done correctly
@@ -460,6 +526,10 @@ def train(args):
disc_loss = exp_loss + pol_loss
disc_loss.backward()
disc_optimizer.step()
with torch.no_grad():
disc_acc_exp = (exp_preds > 0.5).float().mean().item()
disc_acc_pol = (pol_preds < 0.5).float().mean().item()
# --- 3. Update Policy with GAIL Rewards ---
# Reward = -log(1 - D(s, a))
@@ -494,6 +564,13 @@ def train(args):
writer.add_scalar('Loss/Discriminator', disc_loss.item(), i_episode)
writer.add_scalar('Loss/Policy', ppo_loss, i_episode)
writer.add_scalar('Reward/Mean_GAIL', np.mean(all_gail_rewards), i_episode)
if batch_size > 0:
writer.add_scalar('Acc/Disc_Expert', disc_acc_exp, i_episode)
writer.add_scalar('Acc/Disc_Policy', disc_acc_pol, i_episode)
if len(memory['actions']) > 0:
action_arr = np.array(memory['actions'])
action_clip_ratio = (np.abs(action_arr) > 0.98).mean()
writer.add_scalar('Policy/ActionClipRatio', action_clip_ratio, i_episode)
print(f"Episode {i_episode}: Disc Loss {disc_loss.item():.4f} | PPO Loss {ppo_loss:.4f} | Mean Reward {np.mean(all_gail_rewards):.4f}")

1670
技术说明文档.md
View File

File diff suppressed because it is too large Load Diff