246 lines
7.2 KiB
Python
246 lines
7.2 KiB
Python
"""
|
|
Template script for evaluating trained RL agents with Stable Baselines3.
|
|
|
|
This template demonstrates:
|
|
- Loading trained models
|
|
- Evaluating performance with statistics
|
|
- Recording videos of agent behavior
|
|
- Visualizing agent performance
|
|
"""
|
|
|
|
import gymnasium as gym
|
|
import numpy as np
|
|
from stable_baselines3 import PPO
|
|
from stable_baselines3.common.evaluation import evaluate_policy
|
|
from stable_baselines3.common.vec_env import DummyVecEnv, VecVideoRecorder, VecNormalize
|
|
import os
|
|
|
|
|
|
def evaluate_agent(
|
|
model_path,
|
|
env_id="CartPole-v1",
|
|
n_eval_episodes=10,
|
|
deterministic=True,
|
|
render=False,
|
|
record_video=False,
|
|
video_folder="./videos/",
|
|
vec_normalize_path=None,
|
|
):
|
|
"""
|
|
Evaluate a trained RL agent.
|
|
|
|
Args:
|
|
model_path: Path to the saved model
|
|
env_id: Gymnasium environment ID
|
|
n_eval_episodes: Number of episodes to evaluate
|
|
deterministic: Use deterministic actions
|
|
render: Render the environment during evaluation
|
|
record_video: Record videos of the agent
|
|
video_folder: Folder to save videos
|
|
vec_normalize_path: Path to VecNormalize statistics (if used during training)
|
|
|
|
Returns:
|
|
mean_reward: Mean episode reward
|
|
std_reward: Standard deviation of episode rewards
|
|
"""
|
|
# Load the trained model
|
|
print(f"Loading model from {model_path}...")
|
|
model = PPO.load(model_path)
|
|
|
|
# Create evaluation environment
|
|
if render:
|
|
env = gym.make(env_id, render_mode="human")
|
|
else:
|
|
env = gym.make(env_id)
|
|
|
|
# Wrap in DummyVecEnv for consistency
|
|
env = DummyVecEnv([lambda: env])
|
|
|
|
# Load VecNormalize statistics if they were used during training
|
|
if vec_normalize_path and os.path.exists(vec_normalize_path):
|
|
print(f"Loading VecNormalize statistics from {vec_normalize_path}...")
|
|
env = VecNormalize.load(vec_normalize_path, env)
|
|
env.training = False # Don't update statistics during evaluation
|
|
env.norm_reward = False # Don't normalize rewards during evaluation
|
|
|
|
# Set up video recording if requested
|
|
if record_video:
|
|
os.makedirs(video_folder, exist_ok=True)
|
|
env = VecVideoRecorder(
|
|
env,
|
|
video_folder,
|
|
record_video_trigger=lambda x: x == 0, # Record all episodes
|
|
video_length=1000, # Max video length
|
|
name_prefix=f"eval-{env_id}",
|
|
)
|
|
print(f"Recording videos to {video_folder}...")
|
|
|
|
# Evaluate the agent
|
|
print(f"Evaluating for {n_eval_episodes} episodes...")
|
|
mean_reward, std_reward = evaluate_policy(
|
|
model,
|
|
env,
|
|
n_eval_episodes=n_eval_episodes,
|
|
deterministic=deterministic,
|
|
render=False, # VecEnv doesn't support render parameter
|
|
return_episode_rewards=False,
|
|
)
|
|
|
|
print(f"Mean reward: {mean_reward:.2f} +/- {std_reward:.2f}")
|
|
|
|
# Cleanup
|
|
env.close()
|
|
|
|
return mean_reward, std_reward
|
|
|
|
|
|
def watch_agent(
|
|
model_path,
|
|
env_id="CartPole-v1",
|
|
n_episodes=5,
|
|
deterministic=True,
|
|
vec_normalize_path=None,
|
|
):
|
|
"""
|
|
Watch a trained agent play (with rendering).
|
|
|
|
Args:
|
|
model_path: Path to the saved model
|
|
env_id: Gymnasium environment ID
|
|
n_episodes: Number of episodes to watch
|
|
deterministic: Use deterministic actions
|
|
vec_normalize_path: Path to VecNormalize statistics (if used during training)
|
|
"""
|
|
# Load the trained model
|
|
print(f"Loading model from {model_path}...")
|
|
model = PPO.load(model_path)
|
|
|
|
# Create environment with rendering
|
|
env = gym.make(env_id, render_mode="human")
|
|
|
|
# Load VecNormalize statistics if needed
|
|
obs_normalization = None
|
|
if vec_normalize_path and os.path.exists(vec_normalize_path):
|
|
print(f"Loading VecNormalize statistics from {vec_normalize_path}...")
|
|
# For rendering, we'll manually apply normalization
|
|
dummy_env = DummyVecEnv([lambda: gym.make(env_id)])
|
|
vec_env = VecNormalize.load(vec_normalize_path, dummy_env)
|
|
obs_normalization = vec_env
|
|
dummy_env.close()
|
|
|
|
# Run episodes
|
|
for episode in range(n_episodes):
|
|
obs, info = env.reset()
|
|
episode_reward = 0
|
|
done = False
|
|
step = 0
|
|
|
|
print(f"\nEpisode {episode + 1}/{n_episodes}")
|
|
|
|
while not done:
|
|
# Apply observation normalization if needed
|
|
if obs_normalization:
|
|
obs_normalized = obs_normalization.normalize_obs(obs)
|
|
else:
|
|
obs_normalized = obs
|
|
|
|
# Get action from model
|
|
action, _states = model.predict(obs_normalized, deterministic=deterministic)
|
|
|
|
# Take step in environment
|
|
obs, reward, terminated, truncated, info = env.step(action)
|
|
done = terminated or truncated
|
|
|
|
episode_reward += reward
|
|
step += 1
|
|
|
|
print(f"Episode reward: {episode_reward:.2f} ({step} steps)")
|
|
|
|
env.close()
|
|
|
|
|
|
def compare_models(
|
|
model_paths,
|
|
env_id="CartPole-v1",
|
|
n_eval_episodes=10,
|
|
deterministic=True,
|
|
):
|
|
"""
|
|
Compare performance of multiple trained models.
|
|
|
|
Args:
|
|
model_paths: List of paths to saved models
|
|
env_id: Gymnasium environment ID
|
|
n_eval_episodes: Number of episodes to evaluate each model
|
|
deterministic: Use deterministic actions
|
|
"""
|
|
results = {}
|
|
|
|
for model_path in model_paths:
|
|
print(f"\nEvaluating {model_path}...")
|
|
mean_reward, std_reward = evaluate_agent(
|
|
model_path,
|
|
env_id=env_id,
|
|
n_eval_episodes=n_eval_episodes,
|
|
deterministic=deterministic,
|
|
)
|
|
results[model_path] = {"mean": mean_reward, "std": std_reward}
|
|
|
|
# Print comparison
|
|
print("\n" + "=" * 60)
|
|
print("Model Comparison Results")
|
|
print("=" * 60)
|
|
for model_path, stats in results.items():
|
|
print(f"{model_path}: {stats['mean']:.2f} +/- {stats['std']:.2f}")
|
|
print("=" * 60)
|
|
|
|
return results
|
|
|
|
|
|
if __name__ == "__main__":
|
|
# Example 1: Evaluate a trained model
|
|
model_path = "./models/best_model/best_model.zip"
|
|
evaluate_agent(
|
|
model_path=model_path,
|
|
env_id="CartPole-v1",
|
|
n_eval_episodes=10,
|
|
deterministic=True,
|
|
)
|
|
|
|
# Example 2: Record videos of agent behavior
|
|
# evaluate_agent(
|
|
# model_path=model_path,
|
|
# env_id="CartPole-v1",
|
|
# n_eval_episodes=5,
|
|
# deterministic=True,
|
|
# record_video=True,
|
|
# video_folder="./videos/",
|
|
# )
|
|
|
|
# Example 3: Watch agent play with rendering
|
|
# watch_agent(
|
|
# model_path=model_path,
|
|
# env_id="CartPole-v1",
|
|
# n_episodes=3,
|
|
# deterministic=True,
|
|
# )
|
|
|
|
# Example 4: Compare multiple models
|
|
# compare_models(
|
|
# model_paths=[
|
|
# "./models/model_100k.zip",
|
|
# "./models/model_200k.zip",
|
|
# "./models/best_model/best_model.zip",
|
|
# ],
|
|
# env_id="CartPole-v1",
|
|
# n_eval_episodes=10,
|
|
# )
|
|
|
|
# Example 5: Evaluate with VecNormalize statistics
|
|
# evaluate_agent(
|
|
# model_path="./models/best_model/best_model.zip",
|
|
# env_id="Pendulum-v1",
|
|
# n_eval_episodes=10,
|
|
# vec_normalize_path="./models/vec_normalize.pkl",
|
|
# )
|