Fix reset options/seeds ignored on auto-reset in VecEnv

stable_baselines3/common/vec_env/dummy_vec_env.py

@@ -68,7 +68,13 @@ def step_wait(self) -> VecEnvStepReturn:
             if self.buf_dones[env_idx]:
                 # save final observation where user can get it, then reset
                 self.buf_infos[env_idx]["terminal_observation"] = obs
-                obs, self.reset_infos[env_idx] = self.envs[env_idx].reset()
+                maybe_options = {"options": self._options[env_idx]} if self._options[env_idx] else {}
+                obs, self.reset_infos[env_idx] = self.envs[env_idx].reset(
+                    seed=self._seeds[env_idx], **maybe_options
+                )
+                # Seeds and options are only used once per environment
+                self._seeds[env_idx] = None
+                self._options[env_idx] = None
             self._save_obs(env_idx, obs)
         return (self._obs_from_buf(), np.copy(self.buf_rews), np.copy(self.buf_dones), deepcopy(self.buf_infos))
 

stable_baselines3/common/vec_env/subproc_vec_env.py

@@ -32,14 +32,16 @@ def _worker(  # noqa: C901
         try:
             cmd, data = remote.recv()
             if cmd == "step":
-                observation, reward, terminated, truncated, info = env.step(data)
+                action, seed, options = data
+                observation, reward, terminated, truncated, info = env.step(action)
                 # convert to SB3 VecEnv api
                 done = terminated or truncated
                 info["TimeLimit.truncated"] = truncated and not terminated
                 if done:
                     # save final observation where user can get it, then reset
                     info["terminal_observation"] = observation
-                    observation, reset_info = env.reset()
+                    maybe_options = {"options": options} if options else {}
+                    observation, reset_info = env.reset(seed=seed, **maybe_options)
                 remote.send((observation, reward, done, info, reset_info))
             elif cmd == "reset":
                 maybe_options = {"options": data[1]} if data[1] else {}
@@ -85,13 +87,13 @@ class SubprocVecEnv(VecEnv):
     number of logical cores on your CPU.
 
     .. warning::
-
         Only 'forkserver' and 'spawn' start methods are thread-safe,
         which is important when TensorFlow sessions or other non thread-safe
         libraries are used in the parent (see issue #217). However, compared to
         'fork' they incur a small start-up cost and have restrictions on
         global variables. With those methods, users must wrap the code in an
         ``if __name__ == "__main__":`` block.
+
         For more information, see the multiprocessing documentation.
 
     :param env_fns: Environments to run in subprocesses
@@ -104,7 +106,6 @@ def __init__(self, env_fns: list[Callable[[], gym.Env]], start_method: str | Non
         self.waiting = False
         self.closed = False
         n_envs = len(env_fns)
-
         if start_method is None:
             # Fork is not a thread safe method (see issue #217)
             # but is more user friendly (does not require to wrap the code in
@@ -125,18 +126,22 @@ def __init__(self, env_fns: list[Callable[[], gym.Env]], start_method: str | Non
 
         self.remotes[0].send(("get_spaces", None))
         observation_space, action_space = self.remotes[0].recv()
-
         super().__init__(len(env_fns), observation_space, action_space)
 
     def step_async(self, actions: np.ndarray) -> None:
-        for remote, action in zip(self.remotes, actions, strict=True):
-            remote.send(("step", action))
+        for env_idx, (remote, action) in enumerate(zip(self.remotes, actions, strict=True)):
+            remote.send(("step", (action, self._seeds[env_idx], self._options[env_idx])))
         self.waiting = True
 
     def step_wait(self) -> VecEnvStepReturn:
         results = [remote.recv() for remote in self.remotes]
         self.waiting = False
         obs, rews, dones, infos, self.reset_infos = zip(*results, strict=True)  # type: ignore[assignment]
+        # Seeds and options are only used once per environment
+        for env_idx in range(self.num_envs):
+            if dones[env_idx]:
+                self._seeds[env_idx] = None
+                self._options[env_idx] = None
         return _stack_obs(obs, self.observation_space), np.stack(rews), np.stack(dones), infos  # type: ignore[return-value]
 
     def reset(self) -> VecEnvObs: