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
230
231
232
233
234
235
236
237
238
239 | __all__ = ["Monitor", "ResultsWriter", "get_monitor_files", "load_results"]
import csv
import json
import os
from pathlib import Path
import time
from typing import Any, Optional, SupportsFloat, Union
import gymnasium as gym
from gymnasium.core import ActType, ObsType
import pandas
class Monitor(gym.Wrapper[ObsType, ActType, ObsType, ActType]):
"""A monitor wrapper for Gym environments, it is used to know the episode reward, length, time and other data.
Args:
env: The environment
filename: the location to save a log file, can be None for no log
allow_early_resets: allows the reset of the environment before it is done
reset_keywords: extra keywords for the reset call,
if extra parameters are needed at reset
info_keywords: extra information to log, from the information return of env.step()
override_existing: appends to file if ``filename`` exists, otherwise
override existing files (default)
"""
EXT = "monitor.csv"
def __init__(
self,
env: gym.Env,
filename: str | Path | None = None,
allow_early_resets: bool = True,
reset_keywords: tuple[str, ...] = (),
info_keywords: tuple[str, ...] = (),
override_existing: bool = True,
) -> None:
super().__init__(env=env)
self.t_start = time.time()
self.results_writer = None
if filename is not None:
env_id = env.spec.id if env.spec is not None else None
self.results_writer = ResultsWriter(
filename,
header={"t_start": self.t_start, "env_id": str(env_id)},
extra_keys=reset_keywords + info_keywords,
override_existing=override_existing,
)
self.reset_keywords = reset_keywords
self.info_keywords = info_keywords
self.allow_early_resets = allow_early_resets
self.rewards: list[float] = []
self.needs_reset = True
self.episode_returns: list[float] = []
self.episode_lengths: list[int] = []
self.episode_times: list[float] = []
self.total_steps = 0
# extra info about the current episode, that was passed in during reset()
self.current_reset_info: dict[str, Any] = {}
def reset(self, **kwargs) -> tuple[ObsType, dict[str, Any]]:
"""Calls the Gym environment reset. Can only be called if the environment is over, or if allow_early_resets is.
True.
Args:
kwargs: Extra keywords saved for the next episode. only if defined by reset_keywords
Returns:
the first observation of the environment
"""
if not self.allow_early_resets and not self.needs_reset:
raise RuntimeError(
"Tried to reset an environment before done. If you want to allow early resets, "
"wrap your env with Monitor(env, path, allow_early_resets=True)"
)
self.rewards = []
self.needs_reset = False
for key in self.reset_keywords:
value = kwargs.get(key)
if value is None:
raise ValueError(f"Expected you to pass keyword argument {key} into reset")
self.current_reset_info[key] = value
return self.env.reset(**kwargs)
def step(self, action: ActType) -> tuple[ObsType, SupportsFloat, bool, bool, dict[str, Any]]:
"""Step the environment with the given action.
Args:
action: the action
Returns:
observation, reward, terminated, truncated, information
"""
if self.needs_reset:
raise RuntimeError("Tried to step environment that needs reset")
observation, reward, terminated, truncated, info = self.env.step(action)
self.rewards.append(float(reward))
if terminated or truncated:
self.needs_reset = True
ep_rew = sum(self.rewards)
ep_len = len(self.rewards)
ep_info = {"r": round(ep_rew, 6), "l": ep_len, "t": round(time.time() - self.t_start, 6)}
for key in self.info_keywords:
ep_info[key] = info[key]
self.episode_returns.append(ep_rew)
self.episode_lengths.append(ep_len)
self.episode_times.append(time.time() - self.t_start)
ep_info.update(self.current_reset_info)
if self.results_writer:
self.results_writer.write_row(ep_info)
info["episode"] = ep_info
self.total_steps += 1
return observation, reward, terminated, truncated, info
def close(self) -> None:
"""Closes the environment."""
super().close()
if self.results_writer is not None:
self.results_writer.close()
def get_total_steps(self) -> int:
"""Returns the total number of timesteps."""
return self.total_steps
def get_episode_rewards(self) -> list[float]:
"""Returns the rewards of all the episodes."""
return self.episode_returns
def get_episode_lengths(self) -> list[int]:
"""Returns the number of timesteps of all the episodes."""
return self.episode_lengths
def get_episode_times(self) -> list[float]:
"""Returns the runtime in seconds of all the episodes."""
return self.episode_times
class LoadMonitorResultsError(Exception):
"""Raised when loading the monitor log fails."""
pass
class ResultsWriter:
"""A result writer that saves the data from the `Monitor` class.
Args:
filename: the location to save a log file. When it does not end in
the string ``"monitor.csv"``, this suffix will be appended to it
header: the header dictionary object of the saved csv
extra_keys: the extra information to log, typically is composed of
``reset_keywords`` and ``info_keywords``
override_existing: appends to file if ``filename`` exists, otherwise
override existing files (default)
"""
def __init__(
self,
filename: str | Path = "",
header: dict[str, float | str] | None = None,
extra_keys: tuple[str, ...] = (),
override_existing: bool = True,
) -> None:
if header is None:
header = {}
if not str(filename).endswith(Monitor.EXT):
filename = Path(filename) / Monitor.EXT if Path(filename).is_dir() else f"{filename}.{Monitor.EXT}"
filename = os.path.realpath(filename)
# Create (if any) missing filename directories
Path(filename).parent.mkdir(parents=True, exist_ok=True)
# Append mode when not overriding existing file
mode = "w" if override_existing else "a"
# Prevent newline issue on Windows, see GH issue #692
self.file_handler = Path(filename).open(f"{mode}t", newline="\n") # noqa: SIM115
self.logger = csv.DictWriter(self.file_handler, fieldnames=("r", "l", "t", *extra_keys))
if override_existing:
self.file_handler.write(f"#{json.dumps(header)}\n")
self.logger.writeheader()
self.file_handler.flush()
def write_row(self, epinfo: dict[str, float]) -> None:
"""Write row of monitor data to csv log file.
Args:
epinfo: the information on episodic return, length, and time
"""
if self.logger:
self.logger.writerow(epinfo)
self.file_handler.flush()
def close(self) -> None:
"""Close the file handler."""
self.file_handler.close()
def get_monitor_files(path: str | Path) -> list[str]:
"""Get all the monitor files in the given path.
Args:
path: the logging folder
Returns:
the log files
"""
return [str(p) for p in Path(path).glob("*" + Monitor.EXT)]
def load_results(path: str | Path) -> pandas.DataFrame:
"""Load all Monitor logs from a given directory path matching ``*monitor.csv``.
Args:
path: the directory path containing the log file(s)
Returns:
the logged data
"""
monitor_files = get_monitor_files(path)
if len(monitor_files) == 0:
raise LoadMonitorResultsError(f"No monitor files of the form *{Monitor.EXT} found in {path}")
data_frames, headers = [], []
for file_name in monitor_files:
with Path(file_name).open() as file_handler:
first_line = file_handler.readline()
assert first_line[0] == "#"
header = json.loads(first_line[1:])
data_frame = pandas.read_csv(file_handler, index_col=None)
headers.append(header)
data_frame["t"] += header["t_start"]
data_frames.append(data_frame)
data_frame = pandas.concat(data_frames)
data_frame.sort_values("t", inplace=True)
data_frame.reset_index(inplace=True)
data_frame["t"] -= min(header["t_start"] for header in headers)
return data_frame
|