diff --git a/README.md b/README.md index 2cce33e..62379db 100644 --- a/README.md +++ b/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: | -| DDPG | :x: | :heavy_check_mark: | :x: | :x: | :x: | :x: | -| DQN | :x: | :x: | :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: | -| SAC | :x: | :heavy_check_mark: | :x: | :x: | :x: | :x: | -| TD3 | :x: | :heavy_check_mark: | :x: | :x: | :x: | :x: | -| QR-DQN[1](#f1) | :x: | :x: | :heavy_check_mark: | :x: | :x: | :x: | -| TQC[1](#f1) | :x: | :heavy_check_mark: | :x: | :x: | :x: | :x: | +| 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: | :heavy_check_mark: | +| DQN | :x: | :x: | :heavy_check_mark: | :x: | :x: | :heavy_check_mark: | +| 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: | +| SAC | :x: | :heavy_check_mark: | :x: | :x: | :x: | :heavy_check_mark: | +| TD3 | :x: | :heavy_check_mark: | :x: | :x: | :x: | :heavy_check_mark: | +| QR-DQN[1](#f1) | :x: | :x: | :heavy_check_mark: | :x: | :x: | :heavy_check_mark: | +| TQC[1](#f1) | :x: | :heavy_check_mark: | :x: | :x: | :x: | :heavy_check_mark: | | Maskable PPO[1](#f1) | :x: | :x: | :heavy_check_mark: | :heavy_check_mark: | :heavy_check_mark: | :heavy_check_mark: | 1: Implemented in [SB3 Contrib](https://github.com/Stable-Baselines-Team/stable-baselines3-contrib) GitHub repository. diff --git a/docs/guide/algos.rst b/docs/guide/algos.rst index c9201e7..e362f8a 100644 --- a/docs/guide/algos.rst +++ b/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 ✔️ ❌ ❌ ❌ ❌ -DQN ❌ ✔️ ❌ ❌ ❌ -HER ✔️ ✔️ ❌ ❌ ❌ +DDPG ✔️ ❌ ❌ ❌ ✔️ +DQN ❌ ✔️ ❌ ❌ ✔️ +HER ✔️ ✔️ ❌ ❌ ❌ PPO ✔️ ✔️ ✔️ ✔️ ✔️ -SAC ✔️ ❌ ❌ ❌ ❌ -TD3 ✔️ ❌ ❌ ❌ ❌ -QR-DQN [#f1]_ ❌ ️ ✔️ ❌ ❌ ❌ -TQC [#f1]_ ✔️ ❌ ❌ ❌ ❌ +SAC ✔️ ❌ ❌ ❌ ✔️ +TD3 ✔️ ❌ ❌ ❌ ✔️ +QR-DQN [#f1]_ ❌ ️ ✔️ ❌ ❌ ✔️ +TQC [#f1]_ ✔️ ❌ ❌ ❌ ✔️ Maskable PPO [#f1]_ ❌ ✔️ ✔️ ✔️ ✔️ =================== =========== ============ ================= =============== ================ diff --git a/docs/guide/examples.rst b/docs/guide/examples.rst index a8f8de3..9317047 100644 --- a/docs/guide/examples.rst +++ b/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 `_ + + +.. 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 ----------------- diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index 7502c5b..8fd30a4 100644 --- a/docs/misc/changelog.rst +++ b/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``) diff --git a/docs/modules/ddpg.rst b/docs/modules/ddpg.rst index dd07c22..24d265f 100644 --- a/docs/modules/ddpg.rst +++ b/docs/modules/ddpg.rst @@ -39,7 +39,7 @@ Can I use? ---------- - Recurrent policies: ❌ -- Multi processing: ❌ +- Multi processing: ✔️ - Gym spaces: diff --git a/docs/modules/dqn.rst b/docs/modules/dqn.rst index 0c497ad..ce43855 100644 --- a/docs/modules/dqn.rst +++ b/docs/modules/dqn.rst @@ -34,7 +34,7 @@ Can I use? ---------- - Recurrent policies: ❌ -- Multi processing: ❌ +- Multi processing: ✔️ - Gym spaces: diff --git a/docs/modules/sac.rst b/docs/modules/sac.rst index 2f68047..a1156fd 100644 --- a/docs/modules/sac.rst +++ b/docs/modules/sac.rst @@ -46,7 +46,7 @@ Can I use? ---------- - Recurrent policies: ❌ -- Multi processing: ❌ +- Multi processing: ✔️ - Gym spaces: diff --git a/docs/modules/td3.rst b/docs/modules/td3.rst index 33cb38f..3bc93d7 100644 --- a/docs/modules/td3.rst +++ b/docs/modules/td3.rst @@ -39,7 +39,7 @@ Can I use? ---------- - Recurrent policies: ❌ -- Multi processing: ❌ +- Multi processing: ✔️ - Gym spaces: diff --git a/stable_baselines3/common/buffers.py b/stable_baselines3/common/buffers.py index 77d6c36..7e4f3f8 100644 --- a/stable_baselines3/common/buffers.py +++ b/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.actions[batch_inds, 0, :], + self._normalize_obs(self.observations[batch_inds, env_indices, :], env), + 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._normalize_reward(self.rewards[batch_inds], env), + (self.dones[batch_inds, env_indices] * (1 - self.timeouts[batch_inds, env_indices])).reshape(-1, 1), + 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.action_dim), dtype=action_space.dtype) + self.actions = np.zeros((self.buffer_size, self.n_envs, 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()}) - next_obs_ = self._normalize_obs({key: obs[batch_inds, 0, :] for key, obs in self.next_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, 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])), - rewards=self.to_torch(self._normalize_reward(self.rewards[batch_inds], env)), + dones=self.to_torch(self.dones[batch_inds, env_indices] * (1 - self.timeouts[batch_inds, env_indices])).reshape( + -1, 1 + ), + rewards=self.to_torch(self._normalize_reward(self.rewards[batch_inds, env_indices].reshape(-1, 1), env)), ) diff --git a/stable_baselines3/common/callbacks.py b/stable_baselines3/common/callbacks.py index 5f584da..cba6cb8 100644 --- a/stable_baselines3/common/callbacks.py +++ b/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: - # Some models use keyword 'done' (e.g.,: SAC, TD3, DQN, DDPG) - # While some models use keyword 'dones' (e.g.,: A2C, PPO) - 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() + # Check that the `dones` local variable is defined + assert "dones" in self.locals, "`dones` variable is not defined, please check your code next to `callback.on_step()`" + self.n_episodes += np.sum(self.locals["dones"]).item() continue_training = self.n_episodes < self._total_max_episodes diff --git a/stable_baselines3/common/off_policy_algorithm.py b/stable_baselines3/common/off_policy_algorithm.py index 8b3c667..b1528f5 100644 --- a/stable_baselines3/common/off_policy_algorithm.py +++ b/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: - next_obs = infos[0]["terminal_observation"] - # VecNormalize normalizes the terminal observation - if self._vec_normalize_env is not None: - next_obs = self._vec_normalize_env.unnormalize_obs(next_obs) - else: - next_obs = new_obs_ + for i, done in enumerate(dones): + if done and infos[i].get("terminal_observation") is not None: + if isinstance(next_obs, dict): + next_obs_ = infos[i]["terminal_observation"] + # VecNormalize normalizes the terminal observation + if self._vec_normalize_env is not None: + 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 - episode_reward, episode_timesteps = 0.0, 0 + if self.use_sde and self.sde_sample_freq > 0 and num_collected_steps % self.sde_sample_freq == 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: - # Sample a new noise matrix - self.actor.reset_noise() + # Rescale and perform action + new_obs, rewards, dones, infos = env.step(actions) - # Select action randomly or according to policy - action, buffer_action = self._sample_action(learning_starts, action_noise) + self.num_timesteps += env.num_envs + num_collected_steps += 1 - # Rescale and perform action - new_obs, reward, done, infos = env.step(action) + # Give access to local variables + 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 - episode_timesteps += 1 - num_collected_steps += 1 + # Retrieve reward and episode length if using Monitor wrapper + self._update_info_buffer(infos, dones) - # Give access to local variables - 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(0.0, num_collected_steps, num_collected_episodes, continue_training=False) + # Store data in replay buffer (normalized action and unnormalized observation) + self._store_transition(replay_buffer, buffer_actions, new_obs, rewards, dones, infos) - episode_reward += reward + self._update_current_progress_remaining(self.num_timesteps, self._total_timesteps) - # Retrieve reward and episode length if using Monitor wrapper - self._update_info_buffer(infos, done) + # For DQN, check if the target network should be updated + # 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) - self._store_transition(replay_buffer, buffer_action, new_obs, reward, done, infos) + for idx, done in enumerate(dones): + 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 - # and update the exploration schedule - # For SAC/TD3, the update is done as the same time as the gradient update - # 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 + # Log training infos + if log_interval is not None and self._episode_num % log_interval == 0: + self._dump_logs() 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) diff --git a/stable_baselines3/common/on_policy_algorithm.py b/stable_baselines3/common/on_policy_algorithm.py index 0aff9bb..a58d331 100644 --- a/stable_baselines3/common/on_policy_algorithm.py +++ b/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) ): diff --git a/stable_baselines3/common/type_aliases.py b/stable_baselines3/common/type_aliases.py index 45db9eb..7e69d39 100644 --- a/stable_baselines3/common/type_aliases.py +++ b/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 diff --git a/stable_baselines3/common/utils.py b/stable_baselines3/common/utils.py index 7548cdc..e0dfcb9 100644 --- a/stable_baselines3/common/utils.py +++ b/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"Tuple environment, please use {(obs.shape for obs in observation_space.spaces)} " + f"Error: Unexpected observation shape {observation[key].shape} for key {key}, " + + f"please use {observation_space.spaces[key]} " ) diff --git a/stable_baselines3/dqn/dqn.py b/stable_baselines3/dqn/dqn.py index 11e7ac7..668d729 100644 --- a/stable_baselines3/dqn/dqn.py +++ b/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) diff --git a/stable_baselines3/sac/sac.py b/stable_baselines3/sac/sac.py index e9502fd..f7059d7 100644 --- a/stable_baselines3/sac/sac.py +++ b/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 diff --git a/stable_baselines3/td3/td3.py b/stable_baselines3/td3/td3.py index e059761..eb257a6 100644 --- a/stable_baselines3/td3/td3.py +++ b/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 diff --git a/stable_baselines3/version.txt b/stable_baselines3/version.txt index 896c1f3..aca78c7 100644 --- a/stable_baselines3/version.txt +++ b/stable_baselines3/version.txt @@ -1 +1 @@ -1.3.1a3 +1.3.1a4 diff --git a/tests/test_dict_env.py b/tests/test_dict_env.py index f781999..93b13b4 100644 --- a/tests/test_dict_env.py +++ b/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) diff --git a/tests/test_run.py b/tests/test_run.py index c588a02..67b31c4 100644 --- a/tests/test_run.py +++ b/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, + ) diff --git a/tests/test_spaces.py b/tests/test_spaces.py index 1c66d04..deb09c4 100644 --- a/tests/test_spaces.py +++ b/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)