mixed_training

scripts/visualize_replay.py

@@ -0,0 +1,93 @@
+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)