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
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229 | from __future__ import annotations
from pathlib import Path
from typing import TYPE_CHECKING, Any, cast
import warnings
import gymnasium as gym
import lightning as L # noqa: N812
import numpy as np
from lit_sb3.core.evaluation import evaluate_policy
from lit_sb3.vec_env import DummyVecEnv, VecEnv, sync_envs_normalization
from .base import BaseCallback, EventCallback
if TYPE_CHECKING:
pass
class EvalCallback(EventCallback):
"""Callback for evaluating an agent.
.. warning::
When using multiple environments, each call to ``env.step()``
will effectively correspond to ``n_envs`` steps.
To account for that, you can use ``eval_freq = max(eval_freq // n_envs, 1)``
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
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: Verbosity level: 0 for no output, 1 for indicating information about evaluation results
warn: Passed to ``evaluate_policy`` (warns if ``eval_env`` has not been
wrapped with a Monitor wrapper)
"""
def __init__(
self,
eval_env: gym.Env | VecEnv,
callback_on_new_best: BaseCallback | None = None,
callback_after_eval: BaseCallback | None = None,
n_eval_episodes: int = 5,
eval_freq: int = 10000,
log_path: str | Path | None = None,
best_model_save_path: str | Path | None = None,
deterministic: bool = True,
render: bool = False,
verbose: int = 1,
warn: bool = True,
) -> None:
super().__init__(callback_after_eval, verbose=verbose)
self.callback_on_new_best = callback_on_new_best
if self.callback_on_new_best is not None:
# Give access to the parent
self.callback_on_new_best.parent = self
self.n_eval_episodes = n_eval_episodes
self.eval_freq = eval_freq
self.best_mean_reward = -np.inf
self.last_mean_reward = -np.inf
self.deterministic = deterministic
self.render = render
self.warn = warn
# Convert to VecEnv for consistency
if not isinstance(eval_env, VecEnv):
eval_env = DummyVecEnv([lambda: eval_env])
self.eval_env = eval_env
self.best_model_save_path = best_model_save_path
# Logs will be written in ``evaluations.npz``
if log_path is not None:
log_path = str(Path(log_path) / "evaluations")
self.log_path = log_path
self.evaluations_results: list[list[float]] = []
self.evaluations_timesteps: list[int] = []
self.evaluations_length: list[list[int]] = []
# For computing success rate
self._is_success_buffer: list[bool] = []
self.evaluations_successes: list[list[bool]] = []
self._lightning_logger: Any = None
@staticmethod
def _resolve_lightning_logger(trainer: L.Trainer) -> Any:
loggers = getattr(trainer, "loggers", [])
return loggers[0] if loggers else None
def on_train_start(self, trainer: L.Trainer, pl_module: L.LightningModule) -> None:
self._lightning_logger = self._resolve_lightning_logger(trainer)
def _init_callback(self) -> None:
# Does not work in some corner cases, where the wrapper is not the same
if not isinstance(self.training_env, type(self.eval_env)):
warnings.warn(
f"Training and eval env are not of the same type{self.training_env} != {self.eval_env}", stacklevel=2
)
# Create folders if needed
if self.best_model_save_path is not None:
Path(self.best_model_save_path).mkdir(parents=True, exist_ok=True)
if self.log_path is not None:
Path(self.log_path).parent.mkdir(parents=True, exist_ok=True)
# Init callback called on new best model
if self.callback_on_new_best is not None:
self.callback_on_new_best.init_callback(self.model)
def _log_success_callback(self, locals_: dict[str, Any], globals_: dict[str, Any]) -> None:
"""Callback passed to the ``evaluate_policy`` function.
in order to log the success rate (when applicable),
for instance when using HER.
Args:
globals_: Global variables dict (from Python `globals()`).
locals_: Local variables.
"""
info = locals_["info"]
if locals_["done"]:
maybe_is_success = info.get("is_success")
if maybe_is_success is not None:
self._is_success_buffer.append(maybe_is_success)
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 = []
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,
)
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 >= 1:
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 >= 1:
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 >= 1:
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
def update_child_locals(self, locals_: dict[str, Any]) -> None:
"""Update the references to the local variables.
Args:
locals_: the local variables during rollout collection
"""
if self.callback:
self.callback.update_locals(locals_)
|