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 | from copy import deepcopy
from typing import Optional, TypeVar
from lit_sb3.vec_env.base_vec_env import CloudpickleWrapper, VecEnv, VecEnvWrapper
from lit_sb3.vec_env.dummy_vec_env import DummyVecEnv
from lit_sb3.vec_env.stacked_observations import StackedObservations
from lit_sb3.vec_env.subproc_vec_env import SubprocVecEnv
from lit_sb3.vec_env.vec_check_nan import VecCheckNan
from lit_sb3.vec_env.vec_extract_dict_obs import VecExtractDictObs
from lit_sb3.vec_env.vec_frame_stack import VecFrameStack
from lit_sb3.vec_env.vec_monitor import VecMonitor
from lit_sb3.vec_env.vec_normalize import VecNormalize
from lit_sb3.vec_env.vec_transpose import VecTransposeImage
from lit_sb3.vec_env.vec_video_recorder import VecVideoRecorder
VecEnvWrapperT = TypeVar("VecEnvWrapperT", bound=VecEnvWrapper)
def unwrap_vec_wrapper[VecEnvWrapperT: VecEnvWrapper](
env: VecEnv, vec_wrapper_class: type[VecEnvWrapperT]
) -> VecEnvWrapperT | None:
"""Retrieve a ``VecEnvWrapper`` object by recursively searching.
Args:
env: The ``VecEnv`` that is going to be unwrapped
vec_wrapper_class: The desired ``VecEnvWrapper`` class.
Returns:
The ``VecEnvWrapper`` object if the ``VecEnv`` is wrapped with the desired wrapper, None otherwise
"""
env_tmp = env
while isinstance(env_tmp, VecEnvWrapper):
if isinstance(env_tmp, vec_wrapper_class):
return env_tmp
env_tmp = env_tmp.venv
return None
def unwrap_vec_normalize(env: VecEnv) -> VecNormalize | None:
"""Retrieve a ``VecNormalize`` object by recursively searching.
Args:
env: The VecEnv that is going to be unwrapped
Returns:
The ``VecNormalize`` object if the ``VecEnv`` is wrapped with ``VecNormalize``, None otherwise
"""
return unwrap_vec_wrapper(env, VecNormalize)
def is_vecenv_wrapped(env: VecEnv, vec_wrapper_class: type[VecEnvWrapper]) -> bool:
"""Check if an environment is already wrapped in a given ``VecEnvWrapper``.
Args:
env: The VecEnv that is going to be checked
vec_wrapper_class: The desired ``VecEnvWrapper`` class.
Returns:
True if the ``VecEnv`` is wrapped with the desired wrapper, False otherwise
"""
return unwrap_vec_wrapper(env, vec_wrapper_class) is not None
def sync_envs_normalization(env: VecEnv, eval_env: VecEnv) -> None:
"""Synchronize the normalization statistics of an eval environment and train environment.
when they are both wrapped in a ``VecNormalize`` wrapper.
Args:
env: Training env
eval_env: Environment used for evaluation.
"""
env_tmp, eval_env_tmp = env, eval_env
while isinstance(env_tmp, VecEnvWrapper):
assert isinstance(eval_env_tmp, VecEnvWrapper), (
"Error while synchronizing normalization stats: expected the eval env to be "
f"a VecEnvWrapper but got {eval_env_tmp} instead. "
"This is probably due to the training env not being wrapped the same way as the evaluation env. "
f"Training env type: {env_tmp}."
)
if isinstance(env_tmp, VecNormalize):
assert isinstance(eval_env_tmp, VecNormalize), (
"Error while synchronizing normalization stats: expected the eval env to be "
f"a VecNormalize but got {eval_env_tmp} instead. "
"This is probably due to the training env not being wrapped the same way as the evaluation env. "
f"Training env type: {env_tmp}."
)
# Only synchronize if observation normalization exists
if hasattr(env_tmp, "obs_rms"):
eval_env_tmp.obs_rms = deepcopy(env_tmp.obs_rms)
eval_env_tmp.ret_rms = deepcopy(env_tmp.ret_rms)
env_tmp = env_tmp.venv
eval_env_tmp = eval_env_tmp.venv
__all__ = [
"CloudpickleWrapper",
"DummyVecEnv",
"StackedObservations",
"SubprocVecEnv",
"VecCheckNan",
"VecEnv",
"VecEnvWrapper",
"VecExtractDictObs",
"VecFrameStack",
"VecMonitor",
"VecNormalize",
"VecTransposeImage",
"VecVideoRecorder",
"is_vecenv_wrapped",
"sync_envs_normalization",
"unwrap_vec_normalize",
"unwrap_vec_wrapper",
]
|