import os from typing import Any, Callable, Optional, Union import gymnasium as gym from stable_baselines3.common.atari_wrappers import AtariWrapper from stable_baselines3.common.monitor import Monitor from stable_baselines3.common.vec_env import DummyVecEnv, SubprocVecEnv, VecEnv from stable_baselines3.common.vec_env.patch_gym import _patch_env def unwrap_wrapper(env: gym.Env, wrapper_class: type[gym.Wrapper]) -> Optional[gym.Wrapper]: """ Retrieve a ``VecEnvWrapper`` object by recursively searching. :param env: Environment to unwrap :param wrapper_class: Wrapper to look for :return: Environment unwrapped till ``wrapper_class`` if it has been wrapped with it """ env_tmp = env while isinstance(env_tmp, gym.Wrapper): if isinstance(env_tmp, wrapper_class): return env_tmp env_tmp = env_tmp.env return None def is_wrapped(env: gym.Env, wrapper_class: type[gym.Wrapper]) -> bool: """ Check if a given environment has been wrapped with a given wrapper. :param env: Environment to check :param wrapper_class: Wrapper class to look for :return: True if environment has been wrapped with ``wrapper_class``. """ return unwrap_wrapper(env, wrapper_class) is not None def make_vec_env( env_id: Union[str, Callable[..., gym.Env]], n_envs: int = 1, seed: Optional[int] = None, start_index: int = 0, monitor_dir: Optional[str] = None, wrapper_class: Optional[Callable[[gym.Env], gym.Env]] = None, env_kwargs: Optional[dict[str, Any]] = None, vec_env_cls: Optional[type[Union[DummyVecEnv, SubprocVecEnv]]] = None, vec_env_kwargs: Optional[dict[str, Any]] = None, monitor_kwargs: Optional[dict[str, Any]] = None, wrapper_kwargs: Optional[dict[str, Any]] = None, ) -> VecEnv: """ Create a wrapped, monitored ``VecEnv``. By default it uses a ``DummyVecEnv`` which is usually faster than a ``SubprocVecEnv``. :param env_id: either the env ID, the env class or a callable returning an env :param n_envs: the number of environments you wish to have in parallel :param seed: the initial seed for the random number generator :param start_index: start rank index :param monitor_dir: Path to a folder where the monitor files will be saved. If None, no file will be written, however, the env will still be wrapped in a Monitor wrapper to provide additional information about training. :param wrapper_class: Additional wrapper to use on the environment. This can also be a function with single argument that wraps the environment in many things. Note: the wrapper specified by this parameter will be applied after the ``Monitor`` wrapper. if some cases (e.g. with TimeLimit wrapper) this can lead to undesired behavior. See here for more details: https://github.com/DLR-RM/stable-baselines3/issues/894 :param env_kwargs: Optional keyword argument to pass to the env constructor :param vec_env_cls: A custom ``VecEnv`` class constructor. Default: None. :param vec_env_kwargs: Keyword arguments to pass to the ``VecEnv`` class constructor. :param monitor_kwargs: Keyword arguments to pass to the ``Monitor`` class constructor. :param wrapper_kwargs: Keyword arguments to pass to the ``Wrapper`` class constructor. :return: The wrapped environment """ env_kwargs = env_kwargs or {} vec_env_kwargs = vec_env_kwargs or {} monitor_kwargs = monitor_kwargs or {} wrapper_kwargs = wrapper_kwargs or {} assert vec_env_kwargs is not None # for mypy def make_env(rank: int) -> Callable[[], gym.Env]: def _init() -> gym.Env: # For type checker: assert monitor_kwargs is not None assert wrapper_kwargs is not None assert env_kwargs is not None if isinstance(env_id, str): # if the render mode was not specified, we set it to `rgb_array` as default. kwargs = {"render_mode": "rgb_array"} kwargs.update(env_kwargs) try: env = gym.make(env_id, **kwargs) # type: ignore[arg-type] except TypeError: env = gym.make(env_id, **env_kwargs) else: env = env_id(**env_kwargs) # Patch to support gym 0.21/0.26 and gymnasium env = _patch_env(env) if seed is not None: # Note: here we only seed the action space # We will seed the env at the next reset env.action_space.seed(seed + rank) # Wrap the env in a Monitor wrapper # to have additional training information monitor_path = os.path.join(monitor_dir, str(rank)) if monitor_dir is not None else None # Create the monitor folder if needed if monitor_path is not None and monitor_dir is not None: os.makedirs(monitor_dir, exist_ok=True) env = Monitor(env, filename=monitor_path, **monitor_kwargs) # Optionally, wrap the environment with the provided wrapper if wrapper_class is not None: env = wrapper_class(env, **wrapper_kwargs) return env return _init # No custom VecEnv is passed if vec_env_cls is None: # Default: use a DummyVecEnv vec_env_cls = DummyVecEnv vec_env = vec_env_cls([make_env(i + start_index) for i in range(n_envs)], **vec_env_kwargs) # Prepare the seeds for the first reset vec_env.seed(seed) return vec_env def make_atari_env( env_id: Union[str, Callable[..., gym.Env]], n_envs: int = 1, seed: Optional[int] = None, start_index: int = 0, monitor_dir: Optional[str] = None, wrapper_kwargs: Optional[dict[str, Any]] = None, env_kwargs: Optional[dict[str, Any]] = None, vec_env_cls: Optional[Union[type[DummyVecEnv], type[SubprocVecEnv]]] = None, vec_env_kwargs: Optional[dict[str, Any]] = None, monitor_kwargs: Optional[dict[str, Any]] = None, ) -> VecEnv: """ Create a wrapped, monitored VecEnv for Atari. It is a wrapper around ``make_vec_env`` that includes common preprocessing for Atari games. :param env_id: either the env ID, the env class or a callable returning an env :param n_envs: the number of environments you wish to have in parallel :param seed: the initial seed for the random number generator :param start_index: start rank index :param monitor_dir: Path to a folder where the monitor files will be saved. If None, no file will be written, however, the env will still be wrapped in a Monitor wrapper to provide additional information about training. :param wrapper_kwargs: Optional keyword argument to pass to the ``AtariWrapper`` :param env_kwargs: Optional keyword argument to pass to the env constructor :param vec_env_cls: A custom ``VecEnv`` class constructor. Default: None. :param vec_env_kwargs: Keyword arguments to pass to the ``VecEnv`` class constructor. :param monitor_kwargs: Keyword arguments to pass to the ``Monitor`` class constructor. :return: The wrapped environment """ return make_vec_env( env_id, n_envs=n_envs, seed=seed, start_index=start_index, monitor_dir=monitor_dir, wrapper_class=AtariWrapper, env_kwargs=env_kwargs, vec_env_cls=vec_env_cls, vec_env_kwargs=vec_env_kwargs, monitor_kwargs=monitor_kwargs, wrapper_kwargs=wrapper_kwargs, )