Function make_vec_env

stable_baselines3/common/env_util.py:40–129

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

(
    env_id: str | Callable[..., gym.Env],
    n_envs: int = 1,
    seed: int | None = None,
    start_index: int = 0,
    monitor_dir: str | None = None,
    wrapper_class: Callable[[gym.Env], gym.Env] | None = None,
    env_kwargs: dict[str, Any] | None = None,
    vec_env_cls: type[DummyVecEnv | SubprocVecEnv] | None = None,
    vec_env_kwargs: dict[str, Any] | None = None,
    monitor_kwargs: dict[str, Any] | None = None,
    wrapper_kwargs: dict[str, Any] | None = None,
)

def make_vec_env(
    env_id: str | Callable[..., gym.Env],
    n_envs: int = 1,
    seed: int | None = None,
    start_index: int = 0,
    monitor_dir: str | None = None,
    wrapper_class: Callable[[gym.Env], gym.Env] | None = None,
    env_kwargs: dict[str, Any] | None = None,
    vec_env_cls: type[DummyVecEnv | SubprocVecEnv] | None = None,
    vec_env_kwargs: dict[str, Any] | None = None,
    monitor_kwargs: dict[str, Any] | None = None,
    wrapper_kwargs: dict[str, Any] | None = 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)