Add support for different vector autoreset modes

gymnasium/vector/async_vector_env.py

@@ -35,7 +35,7 @@
     read_from_shared_memory,
     write_to_shared_memory,
 )
-from gymnasium.vector.vector_env import ArrayType, VectorEnv
+from gymnasium.vector.vector_env import ArrayType, AutoresetMode, VectorEnv
 
 
 __all__ = ["AsyncVectorEnv", "AsyncState"]
@@ -101,6 +101,7 @@ def __init__(
             | None
         ) = None,
         observation_mode: str | Space = "same",
+        autoreset_mode: str | AutoresetMode = AutoresetMode.NEXT_STEP,
     ):
         """Vectorized environment that runs multiple environments in parallel.
 
@@ -120,6 +121,7 @@ def __init__(
                 'different' defines that there can be multiple observation spaces with different parameters though requires the same shape and dtype,
                 warning, may raise unexpected errors. Passing a ``Tuple[Space, Space]`` object allows defining a custom ``single_observation_space`` and
                 ``observation_space``, warning, may raise unexpected errors.
+            autoreset_mode: The Autoreset Mode used, see todo for more details.
 
         Warnings:
             worker is an advanced mode option. It provides a high degree of flexibility and a high chance
@@ -135,7 +137,15 @@ def __init__(
         self.env_fns = env_fns
         self.shared_memory = shared_memory
         self.copy = copy
+        self.context = context
+        self.daemon = daemon
+        self.worker = worker
         self.observation_mode = observation_mode
+        self.autoreset_mode = (
+            autoreset_mode
+            if isinstance(autoreset_mode, AutoresetMode)
+            else AutoresetMode(autoreset_mode)
+        )
 
         self.num_envs = len(env_fns)
 
@@ -145,6 +155,7 @@ def __init__(
 
         # As we support `make_vec(spec)` then we can't include a `spec = dummy_env.spec` as this doesn't guarantee we can actual recreate the vector env.
         self.metadata = dummy_env.metadata
+        self.metadata["autoreset_mode"] = self.autoreset_mode
         self.render_mode = dummy_env.render_mode
 
         self.single_action_space = dummy_env.action_space
@@ -211,6 +222,7 @@ def __init__(
                         parent_pipe,
                         _obs_buffer,
                         self.error_queue,
+                        self.autoreset_mode,
                     ),
                 )
 
@@ -287,9 +299,32 @@ def reset_async(
                 str(self._state.value),
             )
 
-        for pipe, env_seed in zip(self.parent_pipes, seed):
-            env_kwargs = {"seed": env_seed, "options": options}
-            pipe.send(("reset", env_kwargs))
+        if options is not None and "reset_mask" in options:
+            reset_mask = options.pop("reset_mask")
+            assert isinstance(
+                reset_mask, np.ndarray
+            ), f"`options['reset_mask': mask]` must be a numpy array, got {type(reset_mask)}"
+            assert reset_mask.shape == (
+                self.num_envs,
+            ), f"`options['reset_mask': mask]` must have shape `({self.num_envs},)`, got {reset_mask.shape}"
+            assert (
+                reset_mask.dtype == np.bool_
+            ), f"`options['reset_mask': mask]` must have `dtype=np.bool_`, got {reset_mask.dtype}"
+            assert np.any(
+                reset_mask
+            ), f"`options['reset_mask': mask]` must contain a boolean array, got reset_mask={reset_mask}"
+
+            for pipe, env_seed, env_reset in zip(self.parent_pipes, seed, reset_mask):
+                if env_reset:
+                    env_kwargs = {"seed": env_seed, "options": options}
+                    pipe.send(("reset", env_kwargs))
+                else:
+                    pipe.send(("reset-noop", None))
+        else:
+            for pipe, env_seed in zip(self.parent_pipes, seed):
+                env_kwargs = {"seed": env_seed, "options": options}
+                pipe.send(("reset", env_kwargs))
+
         self._state = AsyncState.WAITING_RESET
 
     def reset_wait(
@@ -688,11 +723,13 @@ def _async_worker(
     parent_pipe: Connection,
     shared_memory: multiprocessing.Array | dict[str, Any] | tuple[Any, ...],
     error_queue: Queue,
+    autoreset_mode: AutoresetMode,
 ):
     env = env_fn()
     observation_space = env.observation_space
     action_space = env.action_space
     autoreset = False
+    observation = None
 
     parent_pipe.close()
 
@@ -709,19 +746,51 @@ def _async_worker(
                     observation = None
                     autoreset = False
                 pipe.send(((observation, info), True))
+            elif command == "reset-noop":
+                pipe.send(((observation, {}), True))
             elif command == "step":
-                if autoreset:
-                    observation, info = env.reset()
-                    reward, terminated, truncated = 0, False, False
-                else:
+                if autoreset_mode == AutoresetMode.NEXT_STEP:
+                    if autoreset:
+                        observation, info = env.reset()
+                        reward, terminated, truncated = 0, False, False
+                    else:
+                        (
+                            observation,
+                            reward,
+                            terminated,
+                            truncated,
+                            info,
+                        ) = env.step(data)
+                    autoreset = terminated or truncated
+                elif autoreset_mode == AutoresetMode.SAME_STEP:
                     (
                         observation,
                         reward,
                         terminated,
                         truncated,
                         info,
                     ) = env.step(data)
-                autoreset = terminated or truncated
+
+                    if terminated or truncated:
+                        reset_observation, reset_info = env.reset()
+
+                        info = {
+                            "final_info": info,
+                            "final_obs": observation,
+                            **reset_info,
+                        }
+                        observation = reset_observation
+                elif autoreset_mode == AutoresetMode.DISABLED:
+                    assert autoreset is False
+                    (
+                        observation,
+                        reward,
+                        terminated,
+                        truncated,
+                        info,
+                    ) = env.step(data)
+                else:
+                    raise ValueError(f"Unexpected autoreset_mode: {autoreset_mode}")
 
                 if shared_memory:
                     write_to_shared_memory(

gymnasium/vector/sync_vector_env.py

@@ -17,7 +17,7 @@
     create_empty_array,
     iterate,
 )
-from gymnasium.vector.vector_env import ArrayType, VectorEnv
+from gymnasium.vector.vector_env import ArrayType, AutoresetMode, VectorEnv
 
 
 __all__ = ["SyncVectorEnv"]
@@ -65,6 +65,7 @@ def __init__(
         env_fns: Iterator[Callable[[], Env]] | Sequence[Callable[[], Env]],
         copy: bool = True,
         observation_mode: str | Space = "same",
+        autoreset_mode: str | AutoresetMode = AutoresetMode.NEXT_STEP,
     ):
         """Vectorized environment that serially runs multiple environments.
 
@@ -74,13 +75,22 @@ def __init__(
             observation_mode: Defines how environment observation spaces should be batched. 'same' defines that there should be ``n`` copies of identical spaces.
                 'different' defines that there can be multiple observation spaces with the same length but different high/low values batched together. Passing a ``Space`` object
                 allows the user to set some custom observation space mode not covered by 'same' or 'different.'
+            autoreset_mode: The Autoreset Mode used, see todo for more details.
+
         Raises:
             RuntimeError: If the observation space of some sub-environment does not match observation_space
                 (or, by default, the observation space of the first sub-environment).
         """
-        self.copy = copy
+        super().__init__()
+
         self.env_fns = env_fns
+        self.copy = copy
         self.observation_mode = observation_mode
+        self.autoreset_mode = (
+            autoreset_mode
+            if isinstance(autoreset_mode, AutoresetMode)
+            else AutoresetMode(autoreset_mode)
+        )
 
         # Initialise all sub-environments
         self.envs = [env_fn() for env_fn in env_fns]
@@ -89,6 +99,7 @@ def __init__(
         # As we support `make_vec(spec)` then we can't include a `spec = self.envs[0].spec` as this doesn't guarantee we can actual recreate the vector env.
         self.num_envs = len(self.envs)
         self.metadata = self.envs[0].metadata
+        self.metadata["autoreset_mode"] = self.autoreset_mode
         self.render_mode = self.envs[0].render_mode
 
         self.single_action_space = self.envs[0].action_space
@@ -130,6 +141,7 @@ def __init__(
             ), f"Sub-environment action space doesn't make the `single_action_space`, action_space={env.action_space}, single_action_space={self.single_action_space}"
 
         # Initialise attributes used in `step` and `reset`
+        self._env_obs = [None for _ in range(self.num_envs)]
         self._observations = create_empty_array(
             self.single_observation_space, n=self.num_envs, fn=np.zeros
         )
@@ -175,23 +187,52 @@ def reset(
             len(seed) == self.num_envs
         ), f"If seeds are passed as a list the length must match num_envs={self.num_envs} but got length={len(seed)}."
 
-        self._terminations = np.zeros((self.num_envs,), dtype=np.bool_)
-        self._truncations = np.zeros((self.num_envs,), dtype=np.bool_)
-
-        observations, infos = [], {}
-        for i, (env, single_seed) in enumerate(zip(self.envs, seed)):
-            env_obs, env_info = env.reset(seed=single_seed, options=options)
+        if options is not None and "reset_mask" in options:
+            reset_mask = options.pop("reset_mask")
+            assert isinstance(
+                reset_mask, np.ndarray
+            ), f"`options['reset_mask': mask]` must be a numpy array, got {type(reset_mask)}"
+            assert reset_mask.shape == (
+                self.num_envs,
+            ), f"`options['reset_mask': mask]` must have shape `({self.num_envs},)`, got {reset_mask.shape}"
+            assert (
+                reset_mask.dtype == np.bool_
+            ), f"`options['reset_mask': mask]` must have `dtype=np.bool_`, got {reset_mask.dtype}"
+            assert np.any(
+                reset_mask
+            ), f"`options['reset_mask': mask]` must contain a boolean array, got reset_mask={reset_mask}"
+
+            self._terminations[reset_mask] = False
+            self._truncations[reset_mask] = False
+            self._autoreset_envs[reset_mask] = False
+
+            infos = {}
+            for i, (env, single_seed, env_mask) in enumerate(
+                zip(self.envs, seed, reset_mask)
+            ):
+                if env_mask:
+                    self._env_obs[i], env_info = env.reset(
+                        seed=single_seed, options=options
+                    )
+
+                    infos = self._add_info(infos, env_info, i)
+        else:
+            self._terminations = np.zeros((self.num_envs,), dtype=np.bool_)
+            self._truncations = np.zeros((self.num_envs,), dtype=np.bool_)
+            self._autoreset_envs = np.zeros((self.num_envs,), dtype=np.bool_)
+
+            infos = {}
+            for i, (env, single_seed) in enumerate(zip(self.envs, seed)):
+                self._env_obs[i], env_info = env.reset(
+                    seed=single_seed, options=options
+                )
 
-            observations.append(env_obs)
-            infos = self._add_info(infos, env_info, i)
+                infos = self._add_info(infos, env_info, i)
 
         # Concatenate the observations
         self._observations = concatenate(
-            self.single_observation_space, observations, self._observations
+            self.single_observation_space, self._env_obs, self._observations
         )
-
-        self._autoreset_envs = np.zeros((self.num_envs,), dtype=np.bool_)
-
         return deepcopy(self._observations) if self.copy else self._observations, infos
 
     def step(
@@ -204,29 +245,58 @@ def step(
         """
         actions = iterate(self.action_space, actions)
 
-        observations, infos = [], {}
+        infos = {}
         for i, action in enumerate(actions):
-            if self._autoreset_envs[i]:
-                env_obs, env_info = self.envs[i].reset()
-
-                self._rewards[i] = 0.0
-                self._terminations[i] = False
-                self._truncations[i] = False
-            else:
+            if self.autoreset_mode == AutoresetMode.NEXT_STEP:
+                if self._autoreset_envs[i]:
+                    self._env_obs[i], env_info = self.envs[i].reset()
+
+                    self._rewards[i] = 0.0
+                    self._terminations[i] = False
+                    self._truncations[i] = False
+                else:
+                    (
+                        self._env_obs[i],
+                        self._rewards[i],
+                        self._terminations[i],
+                        self._truncations[i],
+                        env_info,
+                    ) = self.envs[i].step(action)
+            elif self.autoreset_mode == AutoresetMode.DISABLED:
+                # assumes that the user has correctly autoreset
+                assert not self._autoreset_envs[i], f"{self._autoreset_envs=}"
                 (
-                    env_obs,
+                    self._env_obs[i],
                     self._rewards[i],
                     self._terminations[i],
                     self._truncations[i],
                     env_info,
                 ) = self.envs[i].step(action)
+            elif self.autoreset_mode == AutoresetMode.SAME_STEP:
+                (
+                    self._env_obs[i],
+                    self._rewards[i],
+                    self._terminations[i],
+                    self._truncations[i],
+                    env_info,
+                ) = self.envs[i].step(action)
+
+                if self._terminations[i] or self._truncations[i]:
+                    infos = self._add_info(
+                        infos,
+                        {"final_obs": self._env_obs[i], "final_info": env_info},
+                        i,
+                    )
+
+                    self._env_obs[i], env_info = self.envs[i].reset()
+            else:
+                raise ValueError(f"Unexpected autoreset mode, {self.autoreset_mode}")
 
-            observations.append(env_obs)
             infos = self._add_info(infos, env_info, i)
 
         # Concatenate the observations
         self._observations = concatenate(
-            self.single_observation_space, observations, self._observations
+            self.single_observation_space, self._env_obs, self._observations
         )
         self._autoreset_envs = np.logical_or(self._terminations, self._truncations)