Multiprocessing support for off policy algorithms (#439)

* Add multi-env training support for SAC

* Fix for dict obs

* Pytype fixes

* Fix assert on number of envs

* Remove for loop

* Add support for Dict obs

* Start cleanup

* Update doc and bug fix

* Add support for vectorized action noise
and add multi env example for off-policy

* Update version

* Bug fix with VecNormalize

* Update README table

* Update variable names

* Update changelog and version

* Update doc and fix for `gradient_steps=-1`

* Add test for `gradient_steps=-1`

* Disable pytype pyi errors

* Fix for DQN

* Update comment on deepcopy

* Remove episode_reward field

* Fix RolloutReturn

* Avoid modification by reference

* Fix error message

Co-authored-by: Anssi <kaneran21@hotmail.com>

README.md

@@ -160,15 +160,15 @@ All the following examples can be executed online using Google colab notebooks:
 
 | **Name**         | **Recurrent**      | `Box`          | `Discrete`     | `MultiDiscrete` | `MultiBinary`  | **Multi Processing**              |
 | ------------------- | ------------------ | ------------------ | ------------------ | ------------------- | ------------------ | --------------------------------- |
-| A2C   | :x: | :heavy_check_mark: | :heavy_check_mark: | :heavy_check_mark: | :heavy_check_mark: | :heavy_check_mark:                |
+| A2C   | :x: | :heavy_check_mark: | :heavy_check_mark: | :heavy_check_mark: | :heavy_check_mark: | :heavy_check_mark: |
-| DDPG  | :x: | :heavy_check_mark: | :x:                | :x:                 | :x:                | :x:                               |
+| DDPG  | :x: | :heavy_check_mark: | :x:                | :x:                 | :x:                | :heavy_check_mark: |
-| DQN   | :x: | :x: | :heavy_check_mark: | :x:                 | :x:                | :x:                               |
+| DQN   | :x: | :x: | :heavy_check_mark: | :x:                 | :x:                | :heavy_check_mark: |
-| HER   | :x: | :heavy_check_mark: | :heavy_check_mark: | :x:                 | :x:                | :x:                               |
+| HER   | :x: | :heavy_check_mark: | :heavy_check_mark: | :x:                 | :x:                | :x: |
-| PPO   | :x: | :heavy_check_mark: | :heavy_check_mark: | :heavy_check_mark:  | :heavy_check_mark: | :heavy_check_mark:                |
+| PPO   | :x: | :heavy_check_mark: | :heavy_check_mark: | :heavy_check_mark:  | :heavy_check_mark: | :heavy_check_mark: |
-| SAC   | :x: | :heavy_check_mark: | :x:                | :x:                 | :x:                | :x:                               |
+| SAC   | :x: | :heavy_check_mark: | :x:                | :x:                 | :x:                | :heavy_check_mark: |
-| TD3   | :x: | :heavy_check_mark: | :x:                | :x:                 | :x:                | :x:                               |
+| TD3   | :x: | :heavy_check_mark: | :x:                | :x:                 | :x:                | :heavy_check_mark: |
-| QR-DQN<sup>[1](#f1)</sup>  | :x: | :x: | :heavy_check_mark: | :x:                 | :x:                | :x:                               |
+| QR-DQN<sup>[1](#f1)</sup>  | :x: | :x: | :heavy_check_mark: | :x:                 | :x:                | :heavy_check_mark: |
-| TQC<sup>[1](#f1)</sup>   | :x: | :heavy_check_mark: | :x:                | :x:                 | :x:                | :x:                               |
+| TQC<sup>[1](#f1)</sup>   | :x: | :heavy_check_mark: | :x:                | :x:                 | :x: | :heavy_check_mark: |
 | Maskable PPO<sup>[1](#f1)</sup>   | :x: | :x: | :heavy_check_mark: | :heavy_check_mark:  | :heavy_check_mark: | :heavy_check_mark:  |
 
 <b id="f1">1</b>: Implemented in [SB3 Contrib](https://github.com/Stable-Baselines-Team/stable-baselines3-contrib) GitHub repository.

docs/guide/algos.rst

@@ -9,14 +9,14 @@ along with some useful characteristics: support for discrete/continuous actions,
 Name                 ``Box``     ``Discrete`` ``MultiDiscrete`` ``MultiBinary`` Multi Processing
 ===================  =========== ============ ================= =============== ================
 A2C                  ✔️           ✔️            ✔️                 ✔️               ✔️
-DDPG                 ✔️           ❌            ❌                ❌              ❌
+DDPG                 ✔️           ❌            ❌                ❌               ✔️
-DQN                  ❌           ✔️            ❌                ❌              ❌
+DQN                  ❌           ✔️            ❌                ❌               ✔️
-HER                  ✔️           ✔️            ❌                ❌              ❌
+HER                  ✔️           ✔️            ❌                ❌               ❌
 PPO                  ✔️           ✔️            ✔️                 ✔️               ✔️
-SAC                  ✔️           ❌            ❌                ❌              ❌
+SAC                  ✔️           ❌            ❌                ❌               ✔️
-TD3                  ✔️           ❌            ❌                ❌              ❌
+TD3                  ✔️           ❌            ❌                ❌               ✔️
-QR-DQN [#f1]_        ❌          ️ ✔️            ❌                ❌              ❌
+QR-DQN [#f1]_        ❌          ️ ✔️            ❌                ❌               ✔️
-TQC [#f1]_           ✔️           ❌            ❌                ❌              ❌
+TQC [#f1]_           ✔️           ❌            ❌                ❌               ✔️
 Maskable PPO [#f1]_  ❌           ✔️            ✔️                 ✔️               ✔️
 ===================  =========== ============ ================= =============== ================
 

docs/guide/examples.rst

@@ -158,6 +158,33 @@ Multiprocessing: Unleashing the Power of Vectorized Environments
           env.render()
 
 
+Multiprocessing with off-policy algorithms
+------------------------------------------
+
+.. warning::
+  
+  When using multiple environments with off-policy algorithms, you should update the ``gradient_steps``
+  parameter too. Set it to ``gradient_steps=-1`` to perform as many gradient steps as transitions collected.
+  There is usually a compromise between wall-clock time and sample efficiency,
+  see this `example in PR #439 <https://github.com/DLR-RM/stable-baselines3/pull/439#issuecomment-961796799>`_
+
+
+.. code-block:: python
+
+  import gym
+
+  from stable_baselines3 import SAC
+  from stable_baselines3.common.env_util import make_vec_env
+
+  env = make_vec_env("Pendulum-v0", n_envs=4, seed=0)
+
+  # We collect 4 transitions per call to `ènv.step()`
+  # and performs 2 gradient steps per call to `ènv.step()`
+  # if gradient_steps=-1, then we would do 4 gradients steps per call to `ènv.step()`
+  model = SAC('MlpPolicy', env, train_freq=1, gradient_steps=2, verbose=1)
+  model.learn(total_timesteps=10_000)
+
+
 Dict Observations
 -----------------
 

docs/misc/changelog.rst

@@ -4,16 +4,20 @@ Changelog
 ==========
 
 
-Release 1.3.1a3 (WIP)
+Release 1.3.1a4 (WIP)
 ---------------------------
 
 Breaking Changes:
 ^^^^^^^^^^^^^^^^^
 - Renamed ``mask`` argument of the ``predict()`` method to ``episode_start`` (used with RNN policies only)
+- local variables ``action``, ``done`` and ``reward`` were renamed to their plural form for offpolicy algorithms (``actions``, ``dones``, ``rewards``),
+  this may affect custom callbacks.
+- Removed ``episode_reward`` field from ``RolloutReturn()`` type
 
 New Features:
 ^^^^^^^^^^^^^
 - Added ``norm_obs_keys`` param for ``VecNormalize`` wrapper to configure which observation keys to normalize (@kachayev)
+- Added experimental support to train off-policy algorithms with multiple envs (note: ``HerReplayBuffer`` currently not supported)
 - Handle timeout termination properly for on-policy algorithms (when using ``TimeLimit``)
 
 

docs/modules/ddpg.rst

@@ -39,7 +39,7 @@ Can I use?
 ----------
 
 -  Recurrent policies: ❌
--  Multi processing: ❌
+-  Multi processing: ✔️
 -  Gym spaces:
 
 

docs/modules/dqn.rst

@@ -34,7 +34,7 @@ Can I use?
 ----------
 
 -  Recurrent policies: ❌
--  Multi processing: ❌
+-  Multi processing: ✔️
 -  Gym spaces:
 
 

docs/modules/sac.rst

@@ -46,7 +46,7 @@ Can I use?
 ----------
 
 -  Recurrent policies: ❌
--  Multi processing: ❌
+-  Multi processing: ✔️
 -  Gym spaces:
 
 

docs/modules/td3.rst

@@ -39,7 +39,7 @@ Can I use?
 ----------
 
 -  Recurrent policies: ❌
--  Multi processing: ❌
+-  Multi processing: ✔️
 -  Gym spaces:
 
 

stable_baselines3/common/buffers.py

@@ -181,7 +181,8 @@ class ReplayBuffer(BaseBuffer):
     ):
         super(ReplayBuffer, self).__init__(buffer_size, observation_space, action_space, device, n_envs=n_envs)
 
-        assert n_envs == 1, "Replay buffer only support single environment for now"
+        # Adjust buffer size
+        self.buffer_size = max(buffer_size // n_envs, 1)
 
         # Check that the replay buffer can fit into the memory
         if psutil is not None:
@@ -230,6 +231,17 @@ class ReplayBuffer(BaseBuffer):
         done: np.ndarray,
         infos: List[Dict[str, Any]],
     ) -> None:
+
+        # Reshape needed when using multiple envs with discrete observations
+        # as numpy cannot broadcast (n_discrete,) to (n_discrete, 1)
+        if isinstance(self.observation_space, spaces.Discrete):
+            obs = obs.reshape((self.n_envs,) + self.obs_shape)
+            next_obs = next_obs.reshape((self.n_envs,) + self.obs_shape)
+
+        # Same, for actions
+        if isinstance(self.action_space, spaces.Discrete):
+            action = action.reshape((self.n_envs, self.action_dim))
+
         # Copy to avoid modification by reference
         self.observations[self.pos] = np.array(obs).copy()
 
@@ -273,20 +285,22 @@ class ReplayBuffer(BaseBuffer):
         return self._get_samples(batch_inds, env=env)
 
     def _get_samples(self, batch_inds: np.ndarray, env: Optional[VecNormalize] = None) -> ReplayBufferSamples:
+        # Sample randomly the env idx
+        env_indices = np.random.randint(0, high=self.n_envs, size=(len(batch_inds),))
 
         if self.optimize_memory_usage:
-            next_obs = self._normalize_obs(self.observations[(batch_inds + 1) % self.buffer_size, 0, :], env)
+            next_obs = self._normalize_obs(self.observations[(batch_inds + 1) % self.buffer_size, env_indices, :], env)
         else:
-            next_obs = self._normalize_obs(self.next_observations[batch_inds, 0, :], env)
+            next_obs = self._normalize_obs(self.next_observations[batch_inds, env_indices, :], env)
 
         data = (
-            self._normalize_obs(self.observations[batch_inds, 0, :], env),
+            self._normalize_obs(self.observations[batch_inds, env_indices, :], env),
-            self.actions[batch_inds, 0, :],
+            self.actions[batch_inds, env_indices, :],
             next_obs,
             # Only use dones that are not due to timeouts
             # deactivated by default (timeouts is initialized as an array of False)
-            self.dones[batch_inds] * (1 - self.timeouts[batch_inds]),
+            (self.dones[batch_inds, env_indices] * (1 - self.timeouts[batch_inds, env_indices])).reshape(-1, 1),
-            self._normalize_reward(self.rewards[batch_inds], env),
+            self._normalize_reward(self.rewards[batch_inds, env_indices].reshape(-1, 1), env),
         )
         return ReplayBufferSamples(*tuple(map(self.to_torch, data)))
 
@@ -491,7 +505,7 @@ class DictReplayBuffer(ReplayBuffer):
         super(ReplayBuffer, self).__init__(buffer_size, observation_space, action_space, device, n_envs=n_envs)
 
         assert isinstance(self.obs_shape, dict), "DictReplayBuffer must be used with Dict obs space only"
-        assert n_envs == 1, "Replay buffer only support single environment for now"
+        self.buffer_size = max(buffer_size // n_envs, 1)
 
         # Check that the replay buffer can fit into the memory
         if psutil is not None:
@@ -511,8 +525,7 @@ class DictReplayBuffer(ReplayBuffer):
             for key, _obs_shape in self.obs_shape.items()
         }
 
-        # only 1 env is supported
+        self.actions = np.zeros((self.buffer_size, self.n_envs, self.action_dim), dtype=action_space.dtype)
-        self.actions = np.zeros((self.buffer_size, self.action_dim), dtype=action_space.dtype)
         self.rewards = np.zeros((self.buffer_size, self.n_envs), dtype=np.float32)
         self.dones = np.zeros((self.buffer_size, self.n_envs), dtype=np.float32)
 
@@ -553,11 +566,21 @@ class DictReplayBuffer(ReplayBuffer):
     ) -> None:
         # Copy to avoid modification by reference
         for key in self.observations.keys():
-            self.observations[key][self.pos] = np.array(obs[key]).copy()
+            # Reshape needed when using multiple envs with discrete observations
+            # as numpy cannot broadcast (n_discrete,) to (n_discrete, 1)
+            if isinstance(self.observation_space.spaces[key], spaces.Discrete):
+                obs[key] = obs[key].reshape((self.n_envs,) + self.obs_shape[key])
+            self.observations[key][self.pos] = np.array(obs[key])
 
         for key in self.next_observations.keys():
+            if isinstance(self.observation_space.spaces[key], spaces.Discrete):
+                next_obs[key] = next_obs[key].reshape((self.n_envs,) + self.obs_shape[key])
             self.next_observations[key][self.pos] = np.array(next_obs[key]).copy()
 
+        # Same reshape, for actions
+        if isinstance(self.action_space, spaces.Discrete):
+            action = action.reshape((self.n_envs, self.action_dim))
+
         self.actions[self.pos] = np.array(action).copy()
         self.rewards[self.pos] = np.array(reward).copy()
         self.dones[self.pos] = np.array(done).copy()
@@ -582,10 +605,12 @@ class DictReplayBuffer(ReplayBuffer):
         return super(ReplayBuffer, self).sample(batch_size=batch_size, env=env)
 
     def _get_samples(self, batch_inds: np.ndarray, env: Optional[VecNormalize] = None) -> DictReplayBufferSamples:
+        # Sample randomly the env idx
+        env_indices = np.random.randint(0, high=self.n_envs, size=(len(batch_inds),))
 
         # Normalize if needed and remove extra dimension (we are using only one env for now)
-        obs_ = self._normalize_obs({key: obs[batch_inds, 0, :] for key, obs in self.observations.items()})
+        obs_ = self._normalize_obs({key: obs[batch_inds, env_indices, :] for key, obs in self.observations.items()})
-        next_obs_ = self._normalize_obs({key: obs[batch_inds, 0, :] for key, obs in self.next_observations.items()})
+        next_obs_ = self._normalize_obs({key: obs[batch_inds, env_indices, :] for key, obs in self.next_observations.items()})
 
         # Convert to torch tensor
         observations = {key: self.to_torch(obs) for key, obs in obs_.items()}
@@ -593,12 +618,14 @@ class DictReplayBuffer(ReplayBuffer):
 
         return DictReplayBufferSamples(
             observations=observations,
-            actions=self.to_torch(self.actions[batch_inds]),
+            actions=self.to_torch(self.actions[batch_inds, env_indices]),
             next_observations=next_observations,
             # Only use dones that are not due to timeouts
             # deactivated by default (timeouts is initialized as an array of False)
-            dones=self.to_torch(self.dones[batch_inds] * (1 - self.timeouts[batch_inds])),
+            dones=self.to_torch(self.dones[batch_inds, env_indices] * (1 - self.timeouts[batch_inds, env_indices])).reshape(
-            rewards=self.to_torch(self._normalize_reward(self.rewards[batch_inds], env)),
+                -1, 1
+            ),
+            rewards=self.to_torch(self._normalize_reward(self.rewards[batch_inds, env_indices].reshape(-1, 1), env)),
         )
 
 

stable_baselines3/common/callbacks.py

@@ -519,11 +519,9 @@ class StopTrainingOnMaxEpisodes(BaseCallback):
         self._total_max_episodes = self.max_episodes * self.training_env.num_envs
 
     def _on_step(self) -> bool:
-        # Checking for both 'done' and 'dones' keywords because:
+        # Check that the `dones` local variable is defined
-        # Some models use keyword 'done' (e.g.,: SAC, TD3, DQN, DDPG)
+        assert "dones" in self.locals, "`dones` variable is not defined, please check your code next to `callback.on_step()`"
-        # While some models use keyword 'dones' (e.g.,: A2C, PPO)
+        self.n_episodes += np.sum(self.locals["dones"]).item()
-        done_array = np.array(self.locals.get("done") if self.locals.get("done") is not None else self.locals.get("dones"))
-        self.n_episodes += np.sum(done_array).item()
 
         continue_training = self.n_episodes < self._total_max_episodes
 

stable_baselines3/common/off_policy_algorithm.py

@@ -2,6 +2,7 @@ import io
 import pathlib
 import time
 import warnings
+from copy import deepcopy
 from typing import Any, Dict, List, Optional, Tuple, Type, Union
 
 import gym
@@ -11,7 +12,7 @@ import torch as th
 from stable_baselines3.common.base_class import BaseAlgorithm
 from stable_baselines3.common.buffers import DictReplayBuffer, ReplayBuffer
 from stable_baselines3.common.callbacks import BaseCallback
-from stable_baselines3.common.noise import ActionNoise
+from stable_baselines3.common.noise import ActionNoise, VectorizedActionNoise
 from stable_baselines3.common.policies import BasePolicy
 from stable_baselines3.common.save_util import load_from_pkl, save_to_pkl
 from stable_baselines3.common.type_aliases import GymEnv, MaybeCallback, RolloutReturn, Schedule, TrainFreq, TrainFrequencyUnit
@@ -214,6 +215,7 @@ class OffPolicyAlgorithm(BaseAlgorithm):
                 self.observation_space,
                 self.action_space,
                 self.device,
+                n_envs=self.n_envs,
                 optimize_memory_usage=self.optimize_memory_usage,
                 **self.replay_buffer_kwargs,
             )
@@ -382,7 +384,10 @@ class OffPolicyAlgorithm(BaseAlgorithm):
         raise NotImplementedError()
 
     def _sample_action(
-        self, learning_starts: int, action_noise: Optional[ActionNoise] = None
+        self,
+        learning_starts: int,
+        action_noise: Optional[ActionNoise] = None,
+        n_envs: int = 1,
     ) -> Tuple[np.ndarray, np.ndarray]:
         """
         Sample an action according to the exploration policy.
@@ -394,6 +399,7 @@ class OffPolicyAlgorithm(BaseAlgorithm):
             Required for deterministic policy (e.g. TD3). This can also be used
             in addition to the stochastic policy for SAC.
         :param learning_starts: Number of steps before learning for the warm-up phase.
+        :param n_envs:
         :return: action to take in the environment
             and scaled action that will be stored in the replay buffer.
             The two differs when the action space is not normalized (bounds are not [-1, 1]).
@@ -401,7 +407,7 @@ class OffPolicyAlgorithm(BaseAlgorithm):
         # Select action randomly or according to policy
         if self.num_timesteps < learning_starts and not (self.use_sde and self.use_sde_at_warmup):
             # Warmup phase
-            unscaled_action = np.array([self.action_space.sample()])
+            unscaled_action = np.array([self.action_space.sample() for _ in range(n_envs)])
         else:
             # Note: when using continuous actions,
             # we assume that the policy uses tanh to scale the action
@@ -458,9 +464,9 @@ class OffPolicyAlgorithm(BaseAlgorithm):
         self,
         replay_buffer: ReplayBuffer,
         buffer_action: np.ndarray,
-        new_obs: np.ndarray,
+        new_obs: Union[np.ndarray, Dict[str, np.ndarray]],
         reward: np.ndarray,
-        done: np.ndarray,
+        dones: np.ndarray,
         infos: List[Dict[str, Any]],
     ) -> None:
         """
@@ -471,9 +477,9 @@ class OffPolicyAlgorithm(BaseAlgorithm):
         :param replay_buffer: Replay buffer object where to store the transition.
         :param buffer_action: normalized action
         :param new_obs: next observation in the current episode
-            or first observation of the episode (when done is True)
+            or first observation of the episode (when dones is True)
         :param reward: reward for the current transition
-        :param done: Termination signal
+        :param dones: Termination signal
         :param infos: List of additional information about the transition.
             It may contain the terminal observations and information about timeout.
         """
@@ -485,22 +491,32 @@ class OffPolicyAlgorithm(BaseAlgorithm):
             # Avoid changing the original ones
             self._last_original_obs, new_obs_, reward_ = self._last_obs, new_obs, reward
 
+        # Avoid modification by reference
+        next_obs = deepcopy(new_obs_)
         # As the VecEnv resets automatically, new_obs is already the
         # first observation of the next episode
-        if done and infos[0].get("terminal_observation") is not None:
+        for i, done in enumerate(dones):
-            next_obs = infos[0]["terminal_observation"]
+            if done and infos[i].get("terminal_observation") is not None:
-            # VecNormalize normalizes the terminal observation
+                if isinstance(next_obs, dict):
-            if self._vec_normalize_env is not None:
+                    next_obs_ = infos[i]["terminal_observation"]
-                next_obs = self._vec_normalize_env.unnormalize_obs(next_obs)
+                    # VecNormalize normalizes the terminal observation
-        else:
+                    if self._vec_normalize_env is not None:
-            next_obs = new_obs_
+                        next_obs_ = self._vec_normalize_env.unnormalize_obs(next_obs_)
+                    # Replace next obs for the correct envs
+                    for key in next_obs.keys():
+                        next_obs[key][i] = next_obs_[key]
+                else:
+                    next_obs[i] = infos[i]["terminal_observation"]
+                    # VecNormalize normalizes the terminal observation
+                    if self._vec_normalize_env is not None:
+                        next_obs[i] = self._vec_normalize_env.unnormalize_obs(next_obs[i, :])
 
         replay_buffer.add(
             self._last_original_obs,
             next_obs,
             buffer_action,
             reward_,
-            done,
+            dones,
             infos,
         )
 
@@ -541,79 +557,72 @@ class OffPolicyAlgorithm(BaseAlgorithm):
         # Switch to eval mode (this affects batch norm / dropout)
         self.policy.set_training_mode(False)
 
-        episode_rewards, total_timesteps = [], []
         num_collected_steps, num_collected_episodes = 0, 0
 
         assert isinstance(env, VecEnv), "You must pass a VecEnv"
-        assert env.num_envs == 1, "OffPolicyAlgorithm only support single environment"
         assert train_freq.frequency > 0, "Should at least collect one step or episode."
 
+        if env.num_envs > 1:
+            assert train_freq.unit == TrainFrequencyUnit.STEP, "You must use only one env when doing episodic training."
+
+        # Vectorize action noise if needed
+        if action_noise is not None and env.num_envs > 1 and not isinstance(action_noise, VectorizedActionNoise):
+            action_noise = VectorizedActionNoise(action_noise, env.num_envs)
+
         if self.use_sde:
-            self.actor.reset_noise()
+            self.actor.reset_noise(env.num_envs)
 
         callback.on_rollout_start()
         continue_training = True
 
         while should_collect_more_steps(train_freq, num_collected_steps, num_collected_episodes):
-            done = False
+            if self.use_sde and self.sde_sample_freq > 0 and num_collected_steps % self.sde_sample_freq == 0:
-            episode_reward, episode_timesteps = 0.0, 0
+                # Sample a new noise matrix
+                self.actor.reset_noise(env.num_envs)
 
-            while not done:
+            # Select action randomly or according to policy
+            actions, buffer_actions = self._sample_action(learning_starts, action_noise, env.num_envs)
 
-                if self.use_sde and self.sde_sample_freq > 0 and num_collected_steps % self.sde_sample_freq == 0:
+            # Rescale and perform action
-                    # Sample a new noise matrix
+            new_obs, rewards, dones, infos = env.step(actions)
-                    self.actor.reset_noise()
 
-                # Select action randomly or according to policy
+            self.num_timesteps += env.num_envs
-                action, buffer_action = self._sample_action(learning_starts, action_noise)
+            num_collected_steps += 1
 
-                # Rescale and perform action
+            # Give access to local variables
-                new_obs, reward, done, infos = env.step(action)
+            callback.update_locals(locals())
+            # Only stop training if return value is False, not when it is None.
+            if callback.on_step() is False:
+                return RolloutReturn(num_collected_steps * env.num_envs, num_collected_episodes, continue_training=False)
 
-                self.num_timesteps += 1
+            # Retrieve reward and episode length if using Monitor wrapper
-                episode_timesteps += 1
+            self._update_info_buffer(infos, dones)
-                num_collected_steps += 1
 
-                # Give access to local variables
+            # Store data in replay buffer (normalized action and unnormalized observation)
-                callback.update_locals(locals())
+            self._store_transition(replay_buffer, buffer_actions, new_obs, rewards, dones, infos)
-                # Only stop training if return value is False, not when it is None.
-                if callback.on_step() is False:
-                    return RolloutReturn(0.0, num_collected_steps, num_collected_episodes, continue_training=False)
 
-                episode_reward += reward
+            self._update_current_progress_remaining(self.num_timesteps, self._total_timesteps)
 
-                # Retrieve reward and episode length if using Monitor wrapper
+            # For DQN, check if the target network should be updated
-                self._update_info_buffer(infos, done)
+            # and update the exploration schedule
+            # For SAC/TD3, the update is dones as the same time as the gradient update
+            # see https://github.com/hill-a/stable-baselines/issues/900
+            self._on_step()
 
-                # Store data in replay buffer (normalized action and unnormalized observation)
+            for idx, done in enumerate(dones):
-                self._store_transition(replay_buffer, buffer_action, new_obs, reward, done, infos)
+                if done:
+                    # Update stats
+                    num_collected_episodes += 1
+                    self._episode_num += 1
 
-                self._update_current_progress_remaining(self.num_timesteps, self._total_timesteps)
+                    if action_noise is not None:
+                        kwargs = dict(indices=[idx]) if env.num_envs > 1 else {}
+                        action_noise.reset(**kwargs)
 
-                # For DQN, check if the target network should be updated
+                    # Log training infos
-                # and update the exploration schedule
+                    if log_interval is not None and self._episode_num % log_interval == 0:
-                # For SAC/TD3, the update is done as the same time as the gradient update
+                        self._dump_logs()
-                # see https://github.com/hill-a/stable-baselines/issues/900
-                self._on_step()
-
-                if not should_collect_more_steps(train_freq, num_collected_steps, num_collected_episodes):
-                    break
-
-            if done:
-                num_collected_episodes += 1
-                self._episode_num += 1
-                episode_rewards.append(episode_reward)
-                total_timesteps.append(episode_timesteps)
-
-                if action_noise is not None:
-                    action_noise.reset()
-
-                # Log training infos
-                if log_interval is not None and self._episode_num % log_interval == 0:
-                    self._dump_logs()
-
-        mean_reward = np.mean(episode_rewards) if num_collected_episodes > 0 else 0.0
 
         callback.on_rollout_end()
 
-        return RolloutReturn(mean_reward, num_collected_steps, num_collected_episodes, continue_training)
+        return RolloutReturn(num_collected_steps * env.num_envs, num_collected_episodes, continue_training)

stable_baselines3/common/on_policy_algorithm.py

@@ -193,9 +193,9 @@ class OnPolicyAlgorithm(BaseAlgorithm):
 
             # Handle timeout by bootstraping with value function
             # see GitHub issue #633
-            for idx, done_ in enumerate(dones):
+            for idx, done in enumerate(dones):
                 if (
-                    done_
+                    done
                     and infos[idx].get("terminal_observation") is not None
                     and infos[idx].get("TimeLimit.truncated", False)
                 ):

stable_baselines3/common/type_aliases.py

@@ -56,7 +56,6 @@ class DictReplayBufferSamples(ReplayBufferSamples):
 
 
 class RolloutReturn(NamedTuple):
-    episode_reward: float
     episode_timesteps: int
     n_episodes: int
     continue_training: bool

stable_baselines3/common/utils.py

@@ -332,8 +332,8 @@ def is_vectorized_dict_observation(observation: np.ndarray, observation_space: g
         return True
     else:
         raise ValueError(
-            f"Error: Unexpected observation shape {observation.shape} for "
+            f"Error: Unexpected observation shape {observation[key].shape} for key {key}, "
-            + f"Tuple environment, please use {(obs.shape for obs in observation_space.spaces)} "
+            + f"please use {observation_space.spaces[key]} "
         )
 
 

stable_baselines3/dqn/dqn.py

@@ -1,3 +1,4 @@
+import warnings
 from typing import Any, Dict, List, Optional, Tuple, Type, Union
 
 import gym
@@ -111,12 +112,15 @@ class DQN(OffPolicyAlgorithm):
             sde_support=False,
             optimize_memory_usage=optimize_memory_usage,
             supported_action_spaces=(gym.spaces.Discrete,),
+            support_multi_env=True,
         )
 
         self.exploration_initial_eps = exploration_initial_eps
         self.exploration_final_eps = exploration_final_eps
         self.exploration_fraction = exploration_fraction
         self.target_update_interval = target_update_interval
+        # For updating the target network with multiple envs:
+        self._n_calls = 0
         self.max_grad_norm = max_grad_norm
         # "epsilon" for the epsilon-greedy exploration
         self.exploration_rate = 0.0
@@ -135,6 +139,18 @@ class DQN(OffPolicyAlgorithm):
             self.exploration_final_eps,
             self.exploration_fraction,
         )
+        # Account for multiple environments
+        # each call to step() corresponds to n_envs transitions
+        if self.n_envs > 1:
+            if self.n_envs > self.target_update_interval:
+                warnings.warn(
+                    "The number of environments used is greater than the target network "
+                    f"update interval ({self.n_envs} > {self.target_update_interval}), "
+                    "therefore the target network will be updated after each call to env.step() "
+                    f"which corresponds to {self.n_envs} steps."
+                )
+
+            self.target_update_interval = max(self.target_update_interval // self.n_envs, 1)
 
     def _create_aliases(self) -> None:
         self.q_net = self.policy.q_net
@@ -145,7 +161,8 @@ class DQN(OffPolicyAlgorithm):
         Update the exploration rate and target network if needed.
         This method is called in ``collect_rollouts()`` after each step in the environment.
         """
-        if self.num_timesteps % self.target_update_interval == 0:
+        self._n_calls += 1
+        if self._n_calls % self.target_update_interval == 0:
             polyak_update(self.q_net.parameters(), self.q_net_target.parameters(), self.tau)
 
         self.exploration_rate = self.exploration_schedule(self._current_progress_remaining)

stable_baselines3/sac/sac.py

@@ -129,6 +129,7 @@ class SAC(OffPolicyAlgorithm):
             use_sde_at_warmup=use_sde_at_warmup,
             optimize_memory_usage=optimize_memory_usage,
             supported_action_spaces=(gym.spaces.Box),
+            support_multi_env=True,
         )
 
         self.target_entropy = target_entropy

stable_baselines3/td3/td3.py

@@ -112,6 +112,7 @@ class TD3(OffPolicyAlgorithm):
             sde_support=False,
             optimize_memory_usage=optimize_memory_usage,
             supported_action_spaces=(gym.spaces.Box),
+            support_multi_env=True,
         )
 
         self.policy_delay = policy_delay

stable_baselines3/version.txt

@@ -1 +1 @@
-1.3.1a3
+1.3.1a4

tests/test_dict_env.py

@@ -184,19 +184,19 @@ def test_dict_spaces(model_class, channel_last):
     evaluate_policy(model, env, n_eval_episodes=5, warn=False)
 
 
-@pytest.mark.parametrize("model_class", [PPO, A2C])
+@pytest.mark.parametrize("model_class", [PPO, A2C, SAC, DQN])
 def test_multiprocessing(model_class):
     use_discrete_actions = model_class not in [SAC, TD3, DDPG]
 
     def make_env():
         env = DummyDictEnv(use_discrete_actions=use_discrete_actions, channel_last=False)
-        env = gym.wrappers.TimeLimit(env, 100)
+        env = gym.wrappers.TimeLimit(env, 50)
         return env
 
     env = make_vec_env(make_env, n_envs=2, vec_env_cls=SubprocVecEnv)
 
     kwargs = {}
-    n_steps = 256
+    n_steps = 128
 
     if model_class in {A2C, PPO}:
         kwargs = dict(
@@ -206,6 +206,15 @@ def test_multiprocessing(model_class):
                 features_extractor_kwargs=dict(cnn_output_dim=32),
             ),
         )
+    elif model_class in {SAC, TD3, DQN}:
+        kwargs = dict(
+            buffer_size=1000,
+            policy_kwargs=dict(
+                net_arch=[32],
+                features_extractor_kwargs=dict(cnn_output_dim=16),
+            ),
+            train_freq=5,
+        )
 
     model = model_class("MultiInputPolicy", env, gamma=0.5, seed=1, **kwargs)
 

tests/test_run.py

@@ -1,7 +1,9 @@
+import gym
 import numpy as np
 import pytest
 
 from stable_baselines3 import A2C, DDPG, DQN, PPO, SAC, TD3
+from stable_baselines3.common.env_util import make_vec_env
 from stable_baselines3.common.noise import NormalActionNoise, OrnsteinUhlenbeckActionNoise
 
 normal_action_noise = NormalActionNoise(np.zeros(1), 0.1 * np.ones(1))
@@ -139,3 +141,65 @@ def test_train_freq_fail(train_freq):
             train_freq=train_freq,
         )
         model.learn(total_timesteps=250)
+
+
+@pytest.mark.parametrize("model_class", [SAC, TD3, DDPG, DQN])
+def test_offpolicy_multi_env(model_class):
+    kwargs = {}
+    if model_class in [SAC, TD3, DDPG]:
+        env_id = "Pendulum-v0"
+        policy_kwargs = dict(net_arch=[64], n_critics=1)
+        # Check auto-conversion to VectorizedActionNoise
+        kwargs = dict(action_noise=NormalActionNoise(np.zeros(1), 0.1 * np.ones(1)))
+        if model_class == SAC:
+            kwargs["use_sde"] = True
+            kwargs["sde_sample_freq"] = 4
+    else:
+        env_id = "CartPole-v1"
+        policy_kwargs = dict(net_arch=[64])
+
+    def make_env():
+        env = gym.make(env_id)
+        # to check that the code handling timeouts runs
+        env = gym.wrappers.TimeLimit(env, 50)
+        return env
+
+    env = make_vec_env(make_env, n_envs=2)
+    model = model_class(
+        "MlpPolicy",
+        env,
+        policy_kwargs=policy_kwargs,
+        learning_starts=100,
+        buffer_size=10000,
+        verbose=0,
+        train_freq=5,
+        **kwargs,
+    )
+    model.learn(total_timesteps=150)
+
+    # Check that gradient_steps=-1 works as expected:
+    # perform as many gradient_steps as transitions collected
+    train_freq = 3
+    model = model_class(
+        "MlpPolicy",
+        env,
+        policy_kwargs=policy_kwargs,
+        learning_starts=0,
+        buffer_size=10000,
+        verbose=0,
+        train_freq=train_freq,
+        gradient_steps=-1,
+        **kwargs,
+    )
+    model.learn(total_timesteps=train_freq)
+    assert model.logger.name_to_value["train/n_updates"] == train_freq * env.num_envs
+
+
+def test_warn_dqn_multi_env():
+    with pytest.warns(UserWarning, match="The number of environments used is greater"):
+        DQN(
+            "MlpPolicy",
+            make_vec_env("CartPole-v1", n_envs=2),
+            buffer_size=100,
+            target_update_interval=1,
+        )

tests/test_spaces.py

@@ -69,8 +69,13 @@ def test_action_spaces(model_class, env):
             model_class("MlpPolicy", env)
 
 
-@pytest.mark.parametrize("model_class", [A2C, PPO])
+@pytest.mark.parametrize("model_class", [A2C, PPO, DQN])
 @pytest.mark.parametrize("env", ["Taxi-v3"])
 def test_discrete_obs_space(model_class, env):
     env = make_vec_env(env, n_envs=2, seed=0)
-    model_class("MlpPolicy", env, n_steps=256).learn(500)
+    kwargs = {}
+    if model_class == DQN:
+        kwargs = dict(buffer_size=1000, learning_starts=100)
+    else:
+        kwargs = dict(n_steps=256)
+    model_class("MlpPolicy", env, **kwargs).learn(256)