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

.gitignore

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

scripts/README_visualize.md

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

scripts/launch_tensorboard.py

@@ -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())

scripts/visualize_trained_policy.py

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

train_magail.py

@@ -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()
+            pre_tanh_action = dist.sample()
-            action_logprob = dist.log_prob(action).sum(dim=-1)
+            action = torch.tanh(pre_tanh_action)
-        return action.cpu().numpy(), action_logprob.cpu().numpy()
+            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):
+    # class MAGAILScenarioEnv(MultiAgentScenarioEnv):
-        def _get_all_obs(self):
+    #     def _get_all_obs(self):
-            # Same logic as ExpertReplayEnv to ensure compatibility
+    #         # Same logic as ExpertReplayEnv to ensure compatibility
-            obs_dict = {}
+    #         obs_dict = {}
-            for agent_id, vehicle in self.controlled_agents.items():
+    #         for agent_id, vehicle in self.controlled_agents.items():
-                # 1. Ego State
+    #             # 1. Ego State
-                ego_state = [
+    #             ego_state = [
-                    vehicle.position[0], vehicle.position[1],
+    #                 vehicle.position[0], vehicle.position[1],
-                    vehicle.velocity[0], vehicle.velocity[1],
+    #                 vehicle.velocity[0], vehicle.velocity[1],
-                    vehicle.heading_theta
+    #                 vehicle.heading_theta
-                ]
+    #             ]
-                
+    #             
-                # 2. Neighbors
+    #             # 2. Neighbors
-                candidates = []
+    #             candidates = []
-                for other_id, other_vehicle in self.engine.agent_manager.active_agents.items():
+    #             for other_id, other_vehicle in self.engine.agent_manager.active_agents.items():
-                    if other_id == agent_id:
+    #                 if other_id == agent_id:
-                        continue
+    #                     continue
-                    dist = np.linalg.norm(vehicle.position - other_vehicle.position)
+    #                 dist = np.linalg.norm(vehicle.position - other_vehicle.position)
-                    if dist < 30.0:
+    #                 if dist < 30.0:
-                        candidates.append((dist, other_vehicle))
+    #                     candidates.append((dist, other_vehicle))
-                
+    #             
-                candidates.sort(key=lambda x: x[0])
+    #             candidates.sort(key=lambda x: x[0])
-                top_10 = candidates[:10]
+    #             top_10 = candidates[:10]
-                
+    #             
-                neighbor_feats = []
+    #             neighbor_feats = []
-                for _, neighbor in top_10:
+    #             for _, neighbor in top_10:
-                    neighbor_feats.extend([
+    #                 neighbor_feats.extend([
-                        neighbor.position[0] - vehicle.position[0], 
+    #                     neighbor.position[0] - vehicle.position[0], 
-                        neighbor.position[1] - vehicle.position[1],
+    #                     neighbor.position[1] - vehicle.position[1],
-                        neighbor.velocity[0], 
+    #                     neighbor.velocity[0], 
-                        neighbor.velocity[1]
+    #                     neighbor.velocity[1]
-                    ])
+    #                 ])
-                    
+    #                 
-                missing = 10 - len(top_10)
+    #             missing = 10 - len(top_10)
-                if missing > 0:
+    #             if missing > 0:
-                    neighbor_feats.extend([0.0] * (4 * missing))
+    #                 neighbor_feats.extend([0.0] * (4 * missing))
-                    
+    #                 
-                obs = np.array(ego_state + neighbor_feats, dtype=np.float32)
+    #             obs = np.array(ego_state + neighbor_feats, dtype=np.float32)
-                obs_dict[agent_id] = obs
+    #             obs_dict[agent_id] = obs
-            return obs_dict
+    #         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}")