144 lines
5.2 KiB
Python
144 lines
5.2 KiB
Python
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)
|