-
-
Notifications
You must be signed in to change notification settings - Fork 38
/
minari_dataset.py
328 lines (264 loc) · 11.9 KB
/
minari_dataset.py
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
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
from __future__ import annotations
import importlib.metadata
import os
import re
from dataclasses import dataclass, field
from typing import Callable, Iterable, Iterator, List, Optional, Union
import gymnasium as gym
import numpy as np
from gymnasium import error, logger
from gymnasium.envs.registration import EnvSpec
from packaging.specifiers import InvalidSpecifier, SpecifierSet
from packaging.version import Version
from minari.data_collector.episode_buffer import EpisodeBuffer
from minari.dataset.episode_data import EpisodeData
from minari.dataset.minari_storage import MinariStorage, PathLike
# Use importlib due to circular import when: "from minari import __version__"
__version__ = importlib.metadata.version("minari")
DATASET_ID_RE = re.compile(
r"(?:(?P<environment>[\w]+?))?(?:-(?P<dataset>[\w:.-]+?))(?:-v(?P<version>\d+))?$"
)
def parse_dataset_id(dataset_id: str) -> tuple[str | None, str, int]:
"""Parse dataset ID string format - ``(env_name-)(dataset_name)(-v(version))``.
Args:
dataset_id: The dataset id to parse
Returns:
A tuple of environment name, dataset name and version number
Raises:
Error: If the dataset id is not valid dataset regex
"""
match = DATASET_ID_RE.fullmatch(dataset_id)
if not match:
raise error.Error(
f"Malformed dataset ID: {dataset_id}. (Currently all IDs must be of the form (env_name-)(dataset_name)-v(version). (namespace is optional))"
)
env_name, dataset_name, version = match.group("environment", "dataset", "version")
version = int(version)
return env_name, dataset_name, version
@dataclass
class MinariDatasetSpec:
env_spec: Optional[EnvSpec]
total_episodes: int
total_steps: int
dataset_id: str
combined_datasets: List[str]
observation_space: gym.Space
action_space: gym.Space
data_path: str
minari_version: str
# post-init attributes
env_name: str | None = field(init=False)
dataset_name: str = field(init=False)
version: int | None = field(init=False)
def __post_init__(self):
"""Calls after the spec is created to extract the environment name, dataset name and version from the dataset id."""
self.env_name, self.dataset_name, self.version = parse_dataset_id(
self.dataset_id
)
class MinariDataset:
"""Main Minari dataset class to sample data and get metadata information from a dataset."""
def __init__(
self,
data: Union[MinariStorage, PathLike],
episode_indices: Optional[np.ndarray] = None,
):
"""Initialize properties of the Minari Dataset.
Args:
data (Union[MinariStorage, PathLike]): source of data.
episode_indices (Optional[np.ndarray]): slice of episode indices this dataset is pointing to.
"""
if isinstance(data, MinariStorage):
self._data = data
elif isinstance(data, (str, os.PathLike)):
self._data = MinariStorage.read(data)
else:
raise ValueError(f"Unrecognized type {type(data)} for data")
if episode_indices is None:
episode_indices = np.arange(self._data.total_episodes)
self._episode_indices: np.ndarray = episode_indices
self._total_steps = None
metadata = self._data.metadata
env_spec = metadata.get("env_spec")
if env_spec is not None:
assert isinstance(env_spec, str)
env_spec = EnvSpec.from_json(env_spec)
self._env_spec = env_spec
eval_env_spec = metadata.get("eval_env_spec")
if eval_env_spec is not None:
assert isinstance(eval_env_spec, str)
eval_env_spec = EnvSpec.from_json(eval_env_spec)
self._eval_env_spec = eval_env_spec
dataset_id = metadata["dataset_id"]
assert isinstance(dataset_id, str)
self._dataset_id = dataset_id
minari_version = metadata["minari_version"]
assert isinstance(minari_version, str)
# Check that the dataset is compatible with the current version of Minari
try:
assert Version(__version__) in SpecifierSet(
minari_version
), f"The installed Minari version {__version__} is not contained in the dataset version specifier {minari_version}."
self._minari_version = minari_version
except InvalidSpecifier:
print(f"{minari_version} is not a version specifier.")
self._combined_datasets = metadata.get("combined_datasets", [])
self._observation_space = metadata["observation_space"]
self._action_space = metadata["action_space"]
assert isinstance(self._observation_space, gym.spaces.Space)
assert isinstance(self._action_space, gym.spaces.Space)
self._generator = np.random.default_rng()
def recover_environment(self, eval_env: bool = False, **kwargs) -> gym.Env:
"""Recover the Gymnasium environment used to create the dataset.
Args:
eval_env (bool): if True the returned Gymnasium environment will be that intended to be used for evaluation. If no eval_env was specified when creating the dataset, the returned environment will be the same as the one used for creating the dataset. Default False.
**kwargs: any other parameter that you want to pass to the `gym.make` function.
Returns:
environment: Gymnasium environment
"""
if eval_env:
if self._eval_env_spec is not None:
return gym.make(self._eval_env_spec, **kwargs)
logger.info(
f"`eval_env` has been set to True but the dataset {self._dataset_id} doesn't provide an evaluation environment. Instead, the environment used for collecting the data will be returned: {self._env_spec}"
)
if self.env_spec is None:
raise ValueError("Environment cannot be recovered when env_spec is None")
return gym.make(self.env_spec, **kwargs)
def set_seed(self, seed: int):
"""Set seed for random episode sampling generator."""
self._generator = np.random.default_rng(seed)
def filter_episodes(
self, condition: Callable[[EpisodeData], bool]
) -> MinariDataset:
"""Filter the dataset episodes with a condition.
The condition must be a callable which takes an `EpisodeData` instance and returns a bool.
The callable must return a `bool` True if the condition is met and False otherwise.
i.e filtering for episodes that terminate:
```
dataset.filter(condition=lambda x: x['terminations'][-1] )
```
Args:
condition (Callable[[EpisodeData], bool]): function that gets in input an EpisodeData object and returns True if certain condition is met.
"""
def dict_to_episode_data_condition(episode: dict) -> bool:
return condition(EpisodeData(**episode))
mask = self.storage.apply(
dict_to_episode_data_condition, episode_indices=self.episode_indices
)
assert self.episode_indices is not None
filtered_indices = self.episode_indices[list(mask)]
return MinariDataset(self.storage, episode_indices=filtered_indices)
def sample_episodes(self, n_episodes: int) -> Iterable[EpisodeData]:
"""Sample n number of episodes from the dataset.
Args:
n_episodes (Optional[int], optional): number of episodes to sample.
"""
indices = self._generator.choice(
self.episode_indices, size=n_episodes, replace=False
)
episodes = self.storage.get_episodes(indices)
return list(map(lambda data: EpisodeData(**data), episodes))
def iterate_episodes(
self, episode_indices: Optional[List[int]] = None
) -> Iterator[EpisodeData]:
"""Iterate over episodes from the dataset.
Args:
episode_indices (Optional[List[int]], optional): episode indices to iterate over.
"""
if episode_indices is None:
assert self.episode_indices is not None
assert self.episode_indices.ndim == 1
episode_indices = self.episode_indices.tolist()
assert episode_indices is not None
for episode_index in episode_indices:
data = self.storage.get_episodes([episode_index])[0]
yield EpisodeData(**data)
def update_dataset_from_buffer(self, buffer: List[EpisodeBuffer]):
"""Additional data can be added to the Minari Dataset from a list of episode dictionary buffers.
Args:
buffer (list[EpisodeBuffer]): list of episode dictionary buffers to add to dataset
"""
first_id = self.storage.total_episodes
self.storage.update_episodes(buffer)
self.episode_indices = np.append(
self.episode_indices, first_id + np.arange(len(buffer))
)
def __iter__(self):
return self.iterate_episodes()
def __getitem__(self, idx: int) -> EpisodeData:
episodes_data = self.storage.get_episodes([self.episode_indices[idx]])
assert len(episodes_data) == 1
return EpisodeData(**episodes_data[0])
def __len__(self) -> int:
return self.total_episodes
@property
def total_episodes(self) -> int:
"""Total number of episodes in the Minari dataset."""
return len(self.episode_indices)
@property
def total_steps(self) -> int:
"""Total episodes steps in the Minari dataset."""
if self._total_steps is None:
if self.episode_indices is None:
self._total_steps = self.storage.total_steps
else:
self._total_steps = sum(
self.storage.apply(
lambda episode: episode["total_steps"],
episode_indices=self.episode_indices,
)
)
return int(self._total_steps)
@property
def episode_indices(self) -> np.ndarray:
"""Indices of the available episodes to sample within the Minari dataset."""
return self._episode_indices
@episode_indices.setter
def episode_indices(self, new_value: np.ndarray):
self._total_steps = None # invalidate cache
self._episode_indices = new_value
@property
def observation_space(self):
"""Original observation space of the environment before flatteining (if this is the case)."""
return self._observation_space
@property
def action_space(self):
"""Original action space of the environment before flatteining (if this is the case)."""
return self._action_space
@property
def env_spec(self):
"""Envspec of the environment that has generated the dataset."""
return self._env_spec
@property
def combined_datasets(self) -> List[str]:
"""If this Minari dataset is a combination of other subdatasets, return a list with the subdataset names."""
if self._combined_datasets is None:
return []
else:
return self._combined_datasets
@property
def id(self) -> str:
"""Name of the Minari dataset."""
return self._dataset_id
@property
def minari_version(self) -> str:
"""Version of Minari the dataset is compatible with."""
return self._minari_version
@property
def storage(self) -> MinariStorage:
"""Minari storage managing access to disk."""
return self._data
@property
def spec(self) -> MinariDatasetSpec:
"""Minari dataset specifier."""
return MinariDatasetSpec(
env_spec=self.env_spec,
total_episodes=self._episode_indices.size,
total_steps=self.total_steps,
dataset_id=self.id,
combined_datasets=self.combined_datasets,
observation_space=self.observation_space,
action_space=self.action_space,
data_path=str(self.storage.data_path),
minari_version=str(self.minari_version),
)