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 | from collections.abc import Callable
from typing import Any, Optional, Union
import warnings
import gymnasium as gym
import numpy as np
from lit_sb3.core import types as type_aliases
from lit_sb3.utils.maskable import get_action_masks, is_masking_supported
from lit_sb3.vec_env import DummyVecEnv, VecEnv, VecMonitor, is_vecenv_wrapped
def evaluate_policy(
model: "type_aliases.PolicyPredictor",
env: gym.Env | VecEnv,
n_eval_episodes: int = 10,
deterministic: bool = True,
render: bool = False,
callback: Callable[[dict[str, Any], dict[str, Any]], None] | None = None,
reward_threshold: float | None = None,
return_episode_rewards: bool = False,
use_masking: bool = False,
) -> tuple[float, float] | tuple[list[float], list[int]]:
"""Runs the policy for ``n_eval_episodes`` episodes and outputs the average return.
If a vector env is passed in, this divides the episodes to evaluate onto the
different elements of the vector env. This static division of work is done to
remove bias. See https://github.com/DLR-RM/stable-baselines3/issues/402 for more
details and discussion.
.. note::
If environment has not been wrapped with ``Monitor`` wrapper, reward and
episode lengths are counted as it appears with ``env.step`` calls. If
the environment contains wrappers that modify rewards or episode lengths
(e.g. reward scaling, early episode reset), these will affect the evaluation
results as well. You can avoid this by wrapping environment with ``Monitor``
wrapper before anything else.
Args:
model: The RL agent you want to evaluate. This can be any object
that implements a ``predict`` method, such as an RL algorithm (``BaseAlgorithm``)
or policy (``BasePolicy``).
env: The gym environment or ``VecEnv`` environment.
n_eval_episodes: Number of episode to evaluate the agent
deterministic: Whether to use deterministic or stochastic actions
render: Whether to render the environment or not
callback: callback function to perform additional checks,
called ``n_envs`` times after each step.
Gets locals() and globals() passed as parameters.
See https://github.com/DLR-RM/stable-baselines3/issues/1912 for more details.
reward_threshold: Minimum expected reward per episode,
this will raise an error if the performance is not met
return_episode_rewards: If True, a list of rewards and episode lengths
per episode will be returned instead of the mean.
warn: If True (default), warns user about lack of a Monitor wrapper in the
evaluation environment.
use_masking: Whether or not to use invalid action masks during evaluation.
Requires the environment to support action masking.
Returns:
Mean return per episode (sum of rewards), std of reward per episode.
Returns (list[float], list[int]) when ``return_episode_rewards`` is True, first
list containing per-episode return and second containing per-episode lengths
(in number of steps).
"""
if use_masking and not is_masking_supported(env):
raise ValueError("Environment does not support action masking. Consider using ActionMasker wrapper")
is_monitor_wrapped = False
# Avoid circular import
from lit_sb3.wrappers.monitor import Monitor
if not isinstance(env, VecEnv):
env = DummyVecEnv([lambda: env])
is_monitor_wrapped = is_vecenv_wrapped(env, VecMonitor) or env.env_is_wrapped(Monitor)[0]
if not is_monitor_wrapped:
warnings.warn(
"Evaluation environment is not wrapped with a ``Monitor`` wrapper. "
"This may result in reporting modified episode lengths and rewards, if other wrappers happen to modify these. "
"Consider wrapping environment first with ``Monitor`` wrapper.",
UserWarning,
stacklevel=2,
)
n_envs = env.num_envs
episode_rewards = []
episode_lengths = []
episode_counts = np.zeros(n_envs, dtype="int")
# Divides episodes among different sub environments in the vector as evenly as possible
episode_count_targets = np.array([(n_eval_episodes + i) // n_envs for i in range(n_envs)], dtype="int")
current_rewards = np.zeros(n_envs)
current_lengths = np.zeros(n_envs, dtype="int")
observations = env.reset()
states = None
episode_starts = np.ones((env.num_envs,), dtype=bool)
while (episode_counts < episode_count_targets).any():
if use_masking:
action_masks = get_action_masks(env)
actions, states = model.predict(
observations, # type: ignore[arg-type]
state=states,
episode_start=episode_starts,
deterministic=deterministic,
action_masks=action_masks, # ty:ignore[unknown-argument]
)
else:
actions, states = model.predict(
observations, # type: ignore[arg-type]
state=states,
episode_start=episode_starts,
deterministic=deterministic,
)
new_observations, rewards, dones, infos = env.step(actions)
current_rewards += rewards
current_lengths += 1
for i in range(n_envs):
if episode_counts[i] < episode_count_targets[i]:
# unpack values so that the callback can access the local variables
reward = rewards[i]
done = dones[i]
info = infos[i]
episode_starts[i] = done
if callback is not None:
callback(locals(), globals())
if dones[i]:
if is_monitor_wrapped:
# Atari wrapper can send a "done" signal when
# the agent loses a life, but it does not correspond
# to the true end of episode
if "episode" in info:
# Do not trust "done" with episode endings.
# Monitor wrapper includes "episode" key in info if environment
# has been wrapped with it. Use those rewards instead.
episode_rewards.append(info["episode"]["r"])
episode_lengths.append(info["episode"]["l"])
# Only increment at the real end of an episode
episode_counts[i] += 1
else:
episode_rewards.append(current_rewards[i])
episode_lengths.append(current_lengths[i])
episode_counts[i] += 1
current_rewards[i] = 0
current_lengths[i] = 0
observations = new_observations
if render:
env.render()
mean_reward = np.mean(episode_rewards)
std_reward = np.std(episode_rewards)
if reward_threshold is not None:
assert mean_reward > reward_threshold, (
f"Mean reward below threshold: {mean_reward:.2f} < {reward_threshold:.2f}"
)
if return_episode_rewards:
return episode_rewards, episode_lengths
return mean_reward, std_reward
|