1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125 | import os
from pathlib import Path
from typing import Any
import numpy as np
from lit_sb3.callbacks import EvalCallback
from lit_sb3.core.evaluation import evaluate_policy
from lit_sb3.vec_env import sync_envs_normalization
class MaskableEvalCallback(EvalCallback):
"""Callback for evaluating an agent. Supports invalid action masking.
Args:
eval_env: The environment used for initialization
callback_on_new_best: Callback to trigger
when there is a new best model according to the ``mean_reward``
callback_after_eval: Callback to trigger after every evaluation
when there is a new best model according to the ``mean_reward``
n_eval_episodes: The number of episodes to test the agent
eval_freq: Evaluate the agent every eval_freq call of the callback.
log_path: Path to a folder where the evaluations (``evaluations.npz``)
will be saved. It will be updated at each evaluation.
best_model_save_path: Path to a folder where the best model
according to performance on the eval env will be saved.
deterministic: Whether the evaluation should
use a stochastic or deterministic actions.
render: Whether to render or not the environment during evaluation
verbose:
warn: Passed to ``evaluate_policy`` (warns if ``eval_env`` has not been
wrapped with a Monitor wrapper)
use_masking: Whether to use invalid action masks during evaluation
"""
def __init__(self, *args, use_masking: bool = True, **kwargs) -> None:
super().__init__(*args, **kwargs)
self.use_masking = use_masking
def _on_step(self) -> bool:
continue_training = True
if self.eval_freq > 0 and self.n_calls % self.eval_freq == 0:
# Sync training and eval env if there is VecNormalize
if self.model.get_vec_normalize_env() is not None:
try:
sync_envs_normalization(self.training_env, self.eval_env)
except AttributeError as e:
raise AssertionError(
"Training and eval env are not wrapped the same way, "
"see https://stable-baselines3.readthedocs.io/en/master/guide/callbacks.html#evalcallback "
"and warning above."
) from e
# Reset success rate buffer
self._is_success_buffer = []
# Note that evaluate_policy() has been patched to support masking
episode_rewards, episode_lengths = evaluate_policy(
self.model,
self.eval_env,
n_eval_episodes=self.n_eval_episodes,
render=self.render,
deterministic=self.deterministic,
return_episode_rewards=True,
warn=self.warn,
callback=self._log_success_callback,
use_masking=self.use_masking,
)
if self.log_path is not None:
assert isinstance(episode_rewards, list)
assert isinstance(episode_lengths, list)
self.evaluations_timesteps.append(self.num_timesteps)
self.evaluations_results.append(episode_rewards)
self.evaluations_length.append(episode_lengths)
kwargs = {}
# Save success log if present
if len(self._is_success_buffer) > 0:
self.evaluations_successes.append(self._is_success_buffer)
kwargs = {"successes": self.evaluations_successes}
np.savez(
self.log_path,
timesteps=self.evaluations_timesteps,
results=self.evaluations_results,
ep_lengths=self.evaluations_length,
**kwargs, # type: ignore[arg-type]
)
mean_reward, std_reward = np.mean(episode_rewards), np.std(episode_rewards)
mean_ep_length, std_ep_length = np.mean(episode_lengths), np.std(episode_lengths)
self.last_mean_reward = float(mean_reward)
if self.verbose > 0:
print(f"Eval num_timesteps={self.num_timesteps}, episode_reward={mean_reward:.2f} +/- {std_reward:.2f}")
print(f"Episode length: {mean_ep_length:.2f} +/- {std_ep_length:.2f}")
eval_metrics: dict[str, Any] = {
"eval/mean_reward": float(mean_reward),
"eval/mean_ep_length": mean_ep_length,
}
if len(self._is_success_buffer) > 0:
success_rate = np.mean(self._is_success_buffer)
if self.verbose > 0:
print(f"Success rate: {100 * success_rate:.2f}%")
eval_metrics["eval/success_rate"] = float(success_rate)
if self._lightning_logger is not None:
self._lightning_logger.log_metrics(eval_metrics, step=self.num_timesteps)
if mean_reward > self.best_mean_reward:
if self.verbose > 0:
print("New best mean reward!")
if self.best_model_save_path is not None:
self.model.save(Path(self.best_model_save_path) / "best_model")
self.best_mean_reward = float(mean_reward)
# Trigger callback on new best model, if needed
if self.callback_on_new_best is not None:
continue_training = self.callback_on_new_best.on_step()
# Trigger callback after every evaluation, if needed
if self.callback is not None:
continue_training = continue_training and self._on_event()
return continue_training
|