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

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)