Skills › Scientific Agent Skills › stable-baselines3 › evaluate_agent.py
evaluate_agent.py
92 lines · skills/stable-baselines3/scripts/evaluate_agent.py at 5f2bc03
1"""Evaluate trusted SB3 checkpoints with reproducible seeds and original rewards.
3Recreate the training environment and observation wrappers exactly. The helper
4supports core algorithms and RecurrentPPO; MaskablePPO needs contrib's specialized
5evaluation helper and is deliberately rejected here.
6"""
8from pathlib import Path
10from stable_baselines3 import PPO
11from stable_baselines3.common.env_util import make_vec_env
12from stable_baselines3.common.evaluation import evaluate_policy
13from stable_baselines3.common.vec_env import VecNormalize, VecVideoRecorder
16def evaluate_agent(
17 model_path, env_id="CartPole-v1", n_eval_episodes=10, deterministic=True,
18 render=False, record_video=False, video_folder="./videos/", vec_normalize_path=None,
19 *, algorithm=PPO, seed=20000, device="cpu", video_length=1000, env_kwargs=None,
20):
21 """Return mean and population SD of whole-episode undiscounted returns.
23 Pass the class used for training via algorithm, and the normalization file
24 from that exact checkpoint. Missing requested statistics fail immediately.
25 Video records one bounded clip starting at reset, not every evaluation episode.
26 Rendering and recording are separate modes. Files must come from a trusted
27 source because SB3/VecNormalize persistence includes serialized Python objects.
28 """
29 if not isinstance(n_eval_episodes, int) or n_eval_episodes <= 0:
30 raise ValueError("n_eval_episodes must be a positive integer")
31 if render and record_video:
32 raise ValueError("choose render or record_video for one environment")
33 if record_video and (not isinstance(video_length, int) or video_length < 1):
34 raise ValueError("video_length must be a positive integer")
35 if algorithm.__name__ == "MaskablePPO":
36 raise ValueError("MaskablePPO requires sb3_contrib.common.maskable.evaluation.evaluate_policy")
37 if vec_normalize_path is not None and not Path(vec_normalize_path).is_file():
38 raise FileNotFoundError(vec_normalize_path)
39 kwargs = dict(env_kwargs or {})
40 kwargs["render_mode"] = "rgb_array" if record_video else "human" if render else None
41 env = make_vec_env(env_id, n_envs=1, seed=seed, env_kwargs=kwargs)
42 try:
43 if vec_normalize_path is not None:
44 env = VecNormalize.load(str(vec_normalize_path), env)
45 env.training = False
46 env.norm_reward = False
47 model = algorithm.load(model_path, env=env, device=device)
48 if record_video:
49 env = VecVideoRecorder(
50 env, video_folder, record_video_trigger=lambda step: step == 0,
51 video_length=video_length, name_prefix="evaluation",
52 )
53 mean, std = evaluate_policy(
54 model, env, n_eval_episodes=n_eval_episodes,
55 deterministic=deterministic, render=render,
56 )
57 print(f"Mean reward: {mean:.2f} +/- {std:.2f} (episode SD)")
58 return float(mean), float(std)
59 finally:
60 env.close()
63def watch_agent(model_path, env_id="CartPole-v1", n_episodes=5,
64 deterministic=True, vec_normalize_path=None, **kwargs):
65 """Human rendering through the same normalization and recurrent-state path."""
66 return evaluate_agent(
67 model_path, env_id, n_episodes, deterministic, render=True,
68 vec_normalize_path=vec_normalize_path, **kwargs,
69 )
72def compare_models(model_paths, env_id="CartPole-v1", n_eval_episodes=10,
73 deterministic=True, *, normalization_paths=None, **kwargs):
74 """Compare checkpoints using the same reset seed; SD is not a confidence interval.
76 normalization_paths optionally maps each checkpoint path to its matching stats.
77 Reusing evaluation seeds aids comparison; use independent training runs to
78 estimate training variability and new held-out seeds after model selection.
79 """
80 results = {}
81 for path in model_paths:
82 stats_path = None if normalization_paths is None else normalization_paths[path]
83 mean, std = evaluate_agent(
84 path, env_id, n_eval_episodes, deterministic,
85 vec_normalize_path=stats_path, **kwargs,
86 )
87 results[path] = {"mean": mean, "std": std}
88 return results
91if __name__ == "__main__":
92 evaluate_agent("./models/best_model/best_model.zip")