-
Notifications
You must be signed in to change notification settings - Fork 5
Expand file tree
/
Copy pathbase_env_manager.py
More file actions
executable file
·412 lines (366 loc) · 15.1 KB
/
Copy pathbase_env_manager.py
File metadata and controls
executable file
·412 lines (366 loc) · 15.1 KB
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
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
from types import MethodType
from typing import Union, Any, List, Callable, Dict, Optional
from functools import partial, wraps
from easydict import EasyDict
import copy
import platform
from collections import namedtuple
import numbers
import logging
import enum
import time
import traceback
from ding.utils import ENV_MANAGER_REGISTRY, import_module, one_time_warning
from ding.envs.env.base_env import BaseEnvTimestep
from ding.utils.time_helper import WatchDog
class EnvState(enum.IntEnum):
VOID = 0
INIT = 1
RUN = 2
RESET = 3
DONE = 4
ERROR = 5
def retry_wrapper(func: Callable = None, max_retry: int = 10, waiting_time: float = 0.1) -> Callable:
"""
Overview:
Retry the function until exceeding the maximum retry times.
"""
if func is None:
return partial(retry_wrapper, max_retry=max_retry)
if max_retry == 1:
return func
@wraps(func)
def wrapper(*args, **kwargs):
exceptions = []
for _ in range(max_retry):
try:
ret = func(*args, **kwargs)
return ret
except BaseException as e:
exceptions.append(e)
time.sleep(waiting_time)
logging.error("Function {} has exceeded max retries({})".format(func, max_retry))
runtime_error = RuntimeError(
"Function {} has exceeded max retries({}), and the latest exception is: {}".format(
func, max_retry, repr(exceptions[-1])
)
)
runtime_error.__traceback__ = exceptions[-1].__traceback__
raise runtime_error
return wrapper
def timeout_wrapper(func: Callable = None, timeout: int = 10) -> Callable:
"""
Overview:
Watch the function that must be finihsed within a period of time. If timeout, raise the captured error.
"""
if func is None:
return partial(timeout_wrapper, timeout=timeout)
windows_flag = platform.system().lower() == 'windows'
if windows_flag:
one_time_warning("Timeout wrapper is not implemented in windows platform, so ignore it default")
return func
@wraps(func)
def wrapper(*args, **kwargs):
watchdog = WatchDog(timeout)
try:
watchdog.start()
except ValueError as e:
# watchdog invalid case
return func(*args, **kwargs)
try:
return func(*args, **kwargs)
except BaseException as e:
raise e
finally:
watchdog.stop()
return wrapper
@ENV_MANAGER_REGISTRY.register('base')
class BaseEnvManager(object):
"""
Overview:
Create a BaseEnvManager to manage multiple environments.
Interfaces:
reset, step, seed, close, enable_save_replay, launch, env_info, default_config
Properties:
env_num, ready_obs, done, method_name_list，active_env
"""
@classmethod
def default_config(cls: type) -> EasyDict:
cfg = EasyDict(copy.deepcopy(cls.config))
cfg.cfg_type = cls.__name__ + 'Dict'
return cfg
config = dict(
episode_num=float("inf"),
max_retry=1,
step_timeout=60,
auto_reset=True,
reset_timeout=60,
retry_waiting_time=0.1,
)
def __init__(
self,
env_fn: List[Callable],
cfg: EasyDict = EasyDict({}),
) -> None:
"""
Overview:
Initialize the BaseEnvManager.
Arguments:
- env_fn (:obj:`List[Callable]`): The function to create environment
- cfg (:obj:`EasyDict`): Config
"""
self._cfg = cfg
self._env_fn = env_fn
self._env_num = len(self._env_fn)
self._closed = True
self._env_replay_path = None
# env_ref is used to acquire some common attributes of env, like obs_shape and act_shape
self._env_ref = self._env_fn[0]()
self._env_states = {i: EnvState.VOID for i in range(self._env_num)}
self._env_seed = {i: None for i in range(self._env_num)}
self._episode_num = self._cfg.episode_num
self._max_retry = self._cfg.max_retry
self._step_timeout = self._cfg.step_timeout
self._auto_reset = self._cfg.auto_reset
self._reset_timeout = self._cfg.reset_timeout
self._retry_waiting_time = self._cfg.retry_waiting_time
@property
def env_num(self) -> int:
return self._env_num
@property
def ready_obs(self) -> Dict[int, Any]:
"""
Overview:
Get the next observations(in ``np.ndarray`` type) and corresponding env id.
Return:
A dictionary with observations and their environment IDs.
Example:
>>> obs_dict = env_manager.ready_obs
>>> actions_dict = {env_id: model.forward(obs) for env_id, obs in obs_dict.items())}
"""
return {i: self._ready_obs[i] for i in range(self.env_num) if self._env_episode_count[i] < self._episode_num}
@property
def done(self) -> bool:
return all([s == EnvState.DONE for s in self._env_states.values()])
@property
def method_name_list(self) -> list:
return ['reset', 'step', 'seed', 'close', 'enable_save_replay']
@property
def active_env(self) -> List[int]:
return [i for i, s in self._env_states.items() if s == EnvState.RUN]
def __getattr__(self, key: str) -> Any:
"""
Note:
If a python object doesn't have the attribute whose name is `key`, it will call this method.
We suppose that all envs have the same attributes.
If you need different envs, please implement other env managers.
"""
if not hasattr(self._env_ref, key):
raise AttributeError("env `{}` doesn't have the attribute `{}`".format(type(self._env_ref), key))
if isinstance(getattr(self._env_ref, key), MethodType) and key not in self.method_name_list:
raise RuntimeError("env getattr doesn't support method({}), please override method_name_list".format(key))
self._check_closed()
return [getattr(env, key) if hasattr(env, key) else None for env in self._envs]
def _check_closed(self):
"""
Overview:
Check whether the env manager is closed. Will be called in ``__getattr__`` and ``step``.
"""
assert not self._closed, "env manager is closed, please use the alive env manager"
def launch(self, reset_param: Optional[Dict] = None) -> None:
"""
Overview:
Set up the environments and their parameters.
Arguments:
- reset_param (:obj:`Optional[Dict]`): Dict of reset parameters for each environment, key is the env_id, \
value is the cooresponding reset parameters.
"""
assert self._closed, "Please first close the env manager"
if reset_param is not None:
assert len(reset_param) == len(self._env_fn)
self._create_state()
self.reset(reset_param)
def _create_state(self) -> None:
self._env_episode_count = {i: 0 for i in range(self.env_num)}
self._ready_obs = {i: None for i in range(self.env_num)}
self._envs = [e() for e in self._env_fn]
# env_ref is used to acquire some common attributes of env, like obs_shape and act_shape
self._env_ref = self._envs[0]
assert len(self._envs) == self._env_num
self._reset_param = {i: {} for i in range(self.env_num)}
self._env_states = {i: EnvState.INIT for i in range(self.env_num)}
if self._env_replay_path is not None:
for e, s in zip(self._envs, self._env_replay_path):
e.enable_save_replay(s)
self._closed = False
def reset(self, reset_param: Optional[Dict] = None) -> None:
"""
Overview:
Reset the environments their parameters.
Arguments:
- reset_param (:obj:`List`): Dict of reset parameters for each environment, key is the env_id, \
value is the cooresponding reset parameters.
"""
self._check_closed()
# set seed if necessary
env_ids = list(range(self._env_num)) if reset_param is None else list(reset_param.keys())
for i, env_id in enumerate(env_ids): # loop-type is necessary
if self._env_seed[env_id] is not None:
if self._env_dynamic_seed is not None:
self._envs[env_id].seed(self._env_seed[env_id], self._env_dynamic_seed)
else:
self._envs[env_id].seed(self._env_seed[env_id])
self._env_seed[env_id] = None # seed only use once
# reset env
if reset_param is None:
env_range = range(self.env_num)
else:
for env_id in reset_param:
self._reset_param[env_id] = reset_param[env_id]
env_range = reset_param.keys()
for env_id in env_range:
if self._env_replay_path is not None and self._env_states[env_id] == EnvState.RUN:
logging.warning("please don't reset a unfinished env when you enable save replay, we just skip it")
continue
self._env_states[env_id] = EnvState.RESET
self._reset(env_id)
def _reset(self, env_id: int) -> None:
@retry_wrapper(max_retry=self._max_retry, waiting_time=self._retry_waiting_time)
@timeout_wrapper(timeout=self._reset_timeout)
def reset_fn():
# if self._reset_param[env_id] is None, just reset specific env, not pass reset param
if self._reset_param[env_id] is not None:
assert isinstance(self._reset_param[env_id], dict), type(self._reset_param[env_id])
return self._envs[env_id].reset(**self._reset_param[env_id])
else:
return self._envs[env_id].reset()
try:
obs = reset_fn()
except Exception as e:
self._env_states[env_id] = EnvState.ERROR
self.close()
raise e
self._ready_obs[env_id] = obs
self._env_states[env_id] = EnvState.RUN
def step(self, actions: Dict[int, Any]) -> Dict[int, namedtuple]:
"""
Overview:
Step all environments. Reset an env if done.
Arguments:
- actions (:obj:`Dict[int, Any]`): {env_id: action}
Returns:
- timesteps (:obj:`Dict[int, namedtuple]`): {env_id: timestep}. Timestep is a \
``BaseEnvTimestep`` tuple with observation, reward, done, env_info.
Example:
>>> actions_dict = {env_id: model.forward(obs) for env_id, obs in obs_dict.items())}
>>> timesteps = env_manager.step(actions_dict):
>>> for env_id, timestep in timesteps.items():
>>> pass
.. note:
- The env_id that appears in ``actions`` will also be returned in ``timesteps``.
- Once an environment is done, it is reset immediately.
"""
self._check_closed()
timesteps = {}
for env_id, act in actions.items():
timesteps[env_id] = self._step(env_id, act)
if timesteps[env_id].done:
self._env_episode_count[env_id] += 1
if self._env_episode_count[env_id] < self._episode_num and self._auto_reset:
self._env_states[env_id] = EnvState.RESET
self._reset(env_id)
else:
self._env_states[env_id] = EnvState.DONE
else:
self._ready_obs[env_id] = timesteps[env_id].obs
return timesteps
def _step(self, env_id: int, act: Any) -> namedtuple:
@retry_wrapper(max_retry=self._max_retry, waiting_time=self._retry_waiting_time)
@timeout_wrapper(timeout=self._step_timeout)
def step_fn():
return self._envs[env_id].step(act)
try:
return step_fn()
except Exception as e:
self._env_states[env_id] = EnvState.ERROR
raise e
def seed(self, seed: Union[Dict[int, int], List[int], int], dynamic_seed: bool = None) -> None:
"""
Overview:
Set the seed for each environment.
Arguments:
- seed (:obj:`Union[Dict[int, int], List[int], int]`): List of seeds for each environment; \
Or one seed for the first environment and other seeds are generated automatically.
"""
if isinstance(seed, numbers.Integral):
seed = [seed + i for i in range(self.env_num)]
self._env_seed = seed
elif isinstance(seed, list):
assert len(seed) == self._env_num, "len(seed) {:d} != env_num {:d}".format(len(seed), self._env_num)
self._env_seed = seed
elif isinstance(seed, dict):
if not hasattr(self, '_env_seed'):
raise RuntimeError("please indicate all the seed of each env in the beginning")
for env_id, s in seed.items():
self._env_seed[env_id] = s
else:
raise TypeError("invalid seed arguments type: {}".format(type(seed)))
self._env_dynamic_seed = dynamic_seed
def enable_save_replay(self, replay_path: Union[List[str], str]) -> None:
"""
Overview:
Set each env's replay save path.
Arguments:
- replay_path (:obj:`Union[List[str], str]`): List of paths for each environment; \
Or one path for all environments.
"""
if isinstance(replay_path, str):
replay_path = [replay_path] * self.env_num
self._env_replay_path = replay_path
def close(self) -> None:
"""
Overview:
Release the environment resources.
"""
if self._closed:
return
self._env_ref.close()
for env in self._envs:
env.close()
for i in range(self._env_num):
self._env_states[i] = EnvState.VOID
self._closed = True
def env_info(self) -> namedtuple:
"""
Overview:
Get one env's info, for example, action space, observation space, reward space, etc.
Returnns:
- info (:obj:`namedtuple`): Usually a namedtuple ``BaseEnvInfo``, each element is ``EnvElementInfo``.
"""
return self._env_ref.info()
def create_env_manager(manager_cfg: dict, env_fn: List[Callable]) -> BaseEnvManager:
r"""
Overview:
Create an env manager according to manager cfg and env function.
Arguments:
- manager_cfg (:obj:`EasyDict`): Env manager config.
- env_fn (:obj:` List[Callable]`): A list of envs' functions.
ArgumentsKeys:
- `manager_cfg`'s necessary: `type`
"""
manager_cfg = copy.deepcopy(manager_cfg)
if 'import_names' in manager_cfg:
import_module(manager_cfg.pop('import_names'))
manager_type = manager_cfg.pop('type')
return ENV_MANAGER_REGISTRY.build(manager_type, env_fn=env_fn, cfg=manager_cfg)
def get_env_manager_cls(cfg: EasyDict) -> type:
r"""
Overview:
Get an env manager class according to cfg.
Arguments:
- cfg (:obj:`EasyDict`): Env manager config.
ArgumentsKeys:
- necessary: `type`
"""
import_module(cfg.get('import_names', []))
return ENV_MANAGER_REGISTRY.get(cfg.type)