From c5adad82b2c1733bd8add7e2eeb49a895b635856 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= <45557362+qgallouedec@users.noreply.github.com> Date: Mon, 20 Mar 2023 12:03:57 +0100 Subject: [PATCH] Multiprocessing support for HerReplayBuffer (#704) * IM compat. modif from old fork * mp her working, without offline sampling * update readme and doc * fix discrete action/obs space case * handle offline sampling * fix pos to be consistent with the old version * improve typing and docstring * fix discrete obs special case * new her, using episode uid * deal with full buffer * offline not implemented * info storage; compute_reward as arg; offline sampling error * offline sampling; timeout_termination; fix last_trans detection * rm max_episode_length from tests * fix loading and loading test * Fix episode sampling strategy * Episode interrupted not valid * Typo * Fix infos sampling, next_obs desired goals, offline sampling * update tests for multienvs * speed up code * handle timeout sampling when samping * give up ep_uid for ep_start and ep_lenght * speed up sampling * Improve docstring * Typos and renaming * Fix typing * Fix linter warnings * Renaming + add note * fix reward type * Fix future sampling strategy * Fix future goal selection strategy * env_fn as lambda * Re-fix linter warnings * Formatting * Fix offline sampling * restore the initial performance budget * Remove max_episode_length for HerReplayBuffer kwargs * SubprcVecEnv compat test * Dedicated SubrocVecEnv test rm n_envs from parametrization * Back to using the env arg instead of compute_reward * Up VecEnv import * fix lint warnings * fix docstring * Fix device issue * actor_loss_modifier in SAV and TD3 * Merge RewardModifier and ActorLossModifier into Surgeon * update surgeon for rnd * fix uninteded merge * fix uninteded merge * fix unintended merge * Rm unintended merge * Fix KeyError * Remove useless `all_inds` * Minor docstring format * Fix hint * speedup! * Speedup again * speedup * np.nonzero * fix env normalization * flat sampling for speedup * typo * drop online * format * remove observation from env_cheker (see #1335) * update changelog * default device to "auto" * add comment for info storage * add comment for ep_start and ep_length attributes * a[b][c] to a[b, c] * comment flatnonzero and unravel_index * update _sample_goals docstring * Fix future gaol sampling for split episode * add informative error message for learning_starts too small * use keyword arg for env * try fix pytye * Update stable_baselines3/common/off_policy_algorithm.py Co-authored-by: Antonin RAFFIN * Add `copy_info_dict` option * Ignore pytype * Update changelog * Rename variables and improve documentation * Ignore new bug bear rule * Add note about future strategy * Add deprecation warning * Fix bug trying to pickle buffer kwargs --------- Co-authored-by: Antonin RAFFIN --- README.md | 2 +- docs/guide/algos.rst | 2 +- docs/guide/examples.rst | 4 - docs/guide/migration.rst | 4 +- docs/misc/changelog.rst | 8 +- docs/modules/her.rst | 21 +- pyproject.toml | 3 +- stable_baselines3/common/env_checker.py | 2 +- .../common/off_policy_algorithm.py | 58 +- stable_baselines3/her/her_replay_buffer.py | 686 +++++++----------- stable_baselines3/version.txt | 2 +- tests/test_callbacks.py | 1 - tests/test_her.py | 200 +++-- tests/test_vec_normalize.py | 60 +- 14 files changed, 426 insertions(+), 627 deletions(-) diff --git a/README.md b/README.md index a77dad4..0e0b38d 100644 --- a/README.md +++ b/README.md @@ -173,7 +173,7 @@ All the following examples can be executed online using Google Colab notebooks: | 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: | +| HER | :x: | :heavy_check_mark: | :heavy_check_mark: | :x: | :x: | :heavy_check_mark: | | PPO | :x: | :heavy_check_mark: | :heavy_check_mark: | :heavy_check_mark: | :heavy_check_mark: | :heavy_check_mark: | | QR-DQN[1](#f1) | :x: | :x: | :heavy_check_mark: | :x: | :x: | :heavy_check_mark: | | RecurrentPPO[1](#f1) | :heavy_check_mark: | :heavy_check_mark: | :heavy_check_mark: | :heavy_check_mark: | :heavy_check_mark: | :heavy_check_mark: | diff --git a/docs/guide/algos.rst b/docs/guide/algos.rst index 55aba35..33ac3ba 100644 --- a/docs/guide/algos.rst +++ b/docs/guide/algos.rst @@ -12,7 +12,7 @@ ARS [#f1]_ ✔️ ✔️ ❌ ❌ A2C ✔️ ✔️ ✔️ ✔️ ✔️ DDPG ✔️ ❌ ❌ ❌ ✔️ DQN ❌ ✔️ ❌ ❌ ✔️ -HER ✔️ ✔️ ❌ ❌ ❌ +HER ✔️ ✔️ ❌ ❌ ✔️ PPO ✔️ ✔️ ✔️ ✔️ ✔️ QR-DQN [#f1]_ ❌ ️ ✔️ ❌ ❌ ✔️ RecurrentPPO [#f1]_ ✔️ ✔️ ✔️ ✔️ ✔️ diff --git a/docs/guide/examples.rst b/docs/guide/examples.rst index 61b1fd1..a3f1dc6 100644 --- a/docs/guide/examples.rst +++ b/docs/guide/examples.rst @@ -450,10 +450,6 @@ The parking env is a goal-conditioned continuous control task, in which the vehi replay_buffer_kwargs=dict( n_sampled_goal=n_sampled_goal, goal_selection_strategy="future", - # IMPORTANT: because the env is not wrapped with a TimeLimit wrapper - # we have to manually specify the max number of steps per episode - max_episode_length=100, - online_sampling=True, ), verbose=1, buffer_size=int(1e6), diff --git a/docs/guide/migration.rst b/docs/guide/migration.rst index 571cf45..967a6ac 100644 --- a/docs/guide/migration.rst +++ b/docs/guide/migration.rst @@ -177,10 +177,8 @@ Despite this change, no change in performance should be expected. HER ^^^ -The ``HER`` implementation now also supports online sampling of the new goals. This is done in a vectorized version. +The ``HER`` implementation now only supports online sampling of the new goals. This is done in a vectorized version. The goal selection strategy ``RANDOM`` is no longer supported. -``HER`` now supports ``VecNormalize`` wrapper but only when ``online_sampling=True``. -For performance reasons, the maximum number of steps per episodes must be specified (see :ref:`HER ` documentation). New logger API diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index d129204..840ab60 100644 --- a/docs/misc/changelog.rst +++ b/docs/misc/changelog.rst @@ -4,7 +4,7 @@ Changelog ========== -Release 1.8.0a9 (WIP) +Release 1.8.0a10 (WIP) -------------------------- .. warning:: @@ -20,12 +20,18 @@ Breaking Changes: - Removed shared layers in ``mlp_extractor`` (@AlexPasqua) - Refactored ``StackedObservations`` (it now handles dict obs, ``StackedDictObservations`` was removed) - You must now explicitely pass a ``features_extractor`` parameter when calling ``extract_features()`` +- Dropped offline sampling for ``HerReplayBuffer`` +- As ``HerReplayBuffer`` was refactored to support multiprocessing, previous replay buffer are incompatible with this new version +- ``HerReplayBuffer`` doesn't require a ``max_episode_length`` anymore New Features: ^^^^^^^^^^^^^ - Added ``repeat_action_probability`` argument in ``AtariWrapper``. - Only use ``NoopResetEnv`` and ``MaxAndSkipEnv`` when needed in ``AtariWrapper`` - Added support for dict/tuple observations spaces for ``VecCheckNan``, the check is now active in the ``env_checker()`` (@DavyMorgan) +- Added multiprocessing support for ``HerReplayBuffer`` +- ``HerReplayBuffer`` now supports all datatypes supported by ``ReplayBuffer`` + `SB3-Contrib`_ ^^^^^^^^^^^^^^ diff --git a/docs/modules/her.rst b/docs/modules/her.rst index 817a991..f23c76c 100644 --- a/docs/modules/her.rst +++ b/docs/modules/her.rst @@ -27,14 +27,6 @@ It creates "virtual" transitions by relabeling transitions (changing the desired - a dictionary observation space with three keys: ``observation``, ``achieved_goal`` and ``desired_goal`` -.. warning:: - - For performance reasons, the maximum number of steps per episodes must be specified. - In most cases, it will be inferred if you specify ``max_episode_steps`` when registering the environment - or if you use a ``gym.wrappers.TimeLimit`` (and ``env.spec`` is not None). - Otherwise, you can directly pass ``max_episode_length`` to the model constructor - - .. warning:: Because it needs access to ``env.compute_reward()`` @@ -42,6 +34,12 @@ It creates "virtual" transitions by relabeling transitions (changing the desired without instantiating the environment, we recommend saving the policy only. +.. note:: + + Compared to other implementations, the ``future`` goal sampling strategy is inclusive: + the current transition can be used when re-sampling. + + Notes ----- @@ -77,11 +75,6 @@ This example is only to demonstrate the use of the library and its functions, an # Available strategies (cf paper): future, final, episode goal_selection_strategy = "future" # equivalent to GoalSelectionStrategy.FUTURE - # If True the HER transitions will get sampled online - online_sampling = True - # Time limit for the episodes - max_episode_length = N_BITS - # Initialize the model model = model_class( "MultiInputPolicy", @@ -91,8 +84,6 @@ This example is only to demonstrate the use of the library and its functions, an replay_buffer_kwargs=dict( n_sampled_goal=4, goal_selection_strategy=goal_selection_strategy, - online_sampling=online_sampling, - max_episode_length=max_episode_length, ), verbose=1, ) diff --git a/pyproject.toml b/pyproject.toml index 67a7edc..7941679 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -5,7 +5,8 @@ line-length = 127 target-version = "py37" # See https://beta.ruff.rs/docs/rules/ select = ["E", "F", "B", "UP", "C90", "RUF"] -ignore = [] +# Ignore explicit stacklevel` +ignore = ["B028"] [tool.ruff.per-file-ignores] # Default implementation in abstract methods diff --git a/stable_baselines3/common/env_checker.py b/stable_baselines3/common/env_checker.py index ce01f2e..950cea6 100644 --- a/stable_baselines3/common/env_checker.py +++ b/stable_baselines3/common/env_checker.py @@ -117,7 +117,7 @@ def _check_goal_env_obs(obs: dict, observation_space: spaces.Dict, method_name: f"The current observation contains {len(observation_space.spaces)} keys: {list(observation_space.spaces.keys())}" ) - for key in ["observation", "achieved_goal", "desired_goal"]: + for key in ["achieved_goal", "desired_goal"]: if key not in observation_space.spaces: raise AssertionError( f"The observation returned by the `{method_name}()` method of a goal-conditioned env requires the '{key}' " diff --git a/stable_baselines3/common/off_policy_algorithm.py b/stable_baselines3/common/off_policy_algorithm.py index c1ab215..943f2ab 100644 --- a/stable_baselines3/common/off_policy_algorithm.py +++ b/stable_baselines3/common/off_policy_algorithm.py @@ -125,10 +125,14 @@ class OffPolicyAlgorithm(BaseAlgorithm): self.gradient_steps = gradient_steps self.action_noise = action_noise self.optimize_memory_usage = optimize_memory_usage - self.replay_buffer_class = replay_buffer_class - if replay_buffer_kwargs is None: - replay_buffer_kwargs = {} - self.replay_buffer_kwargs = replay_buffer_kwargs + if replay_buffer_class is None: + if isinstance(self.observation_space, spaces.Dict): + self.replay_buffer_class = DictReplayBuffer + else: + self.replay_buffer_class = ReplayBuffer + else: + self.replay_buffer_class = replay_buffer_class + self.replay_buffer_kwargs = replay_buffer_kwargs or {} self._episode_storage = None # Save train freq parameter, will be converted later to TrainFreq object @@ -170,37 +174,13 @@ class OffPolicyAlgorithm(BaseAlgorithm): self._setup_lr_schedule() self.set_random_seed(self.seed) - # Use DictReplayBuffer if needed - if self.replay_buffer_class is None: - if isinstance(self.observation_space, spaces.Dict): - self.replay_buffer_class = DictReplayBuffer - else: - self.replay_buffer_class = ReplayBuffer - - elif self.replay_buffer_class == HerReplayBuffer: - assert self.env is not None, "You must pass an environment when using `HerReplayBuffer`" - - # If using offline sampling, we need a classic replay buffer too - if self.replay_buffer_kwargs.get("online_sampling", True): - replay_buffer = None - else: - replay_buffer = DictReplayBuffer( - self.buffer_size, - self.observation_space, - self.action_space, - device=self.device, - optimize_memory_usage=self.optimize_memory_usage, - ) - - self.replay_buffer = HerReplayBuffer( - self.env, - self.buffer_size, - device=self.device, - replay_buffer=replay_buffer, - **self.replay_buffer_kwargs, - ) - if self.replay_buffer is None: + # Make a local copy as we should not pickle + # the environment when using HerReplayBuffer + replay_buffer_kwargs = self.replay_buffer_kwargs.copy() + if issubclass(self.replay_buffer_class, HerReplayBuffer): + assert self.env is not None, "You must pass an environment when using `HerReplayBuffer`" + replay_buffer_kwargs["env"] = self.env self.replay_buffer = self.replay_buffer_class( self.buffer_size, self.observation_space, @@ -208,7 +188,7 @@ class OffPolicyAlgorithm(BaseAlgorithm): device=self.device, n_envs=self.n_envs, optimize_memory_usage=self.optimize_memory_usage, - **self.replay_buffer_kwargs, + **replay_buffer_kwargs, # pytype:disable=wrong-keyword-args ) self.policy = self.policy_class( # pytype:disable=not-instantiable @@ -276,12 +256,7 @@ class OffPolicyAlgorithm(BaseAlgorithm): # when using memory efficient replay buffer # see https://github.com/DLR-RM/stable-baselines3/issues/46 - # Special case when using HerReplayBuffer, - # the classic replay buffer is inside it when using offline sampling - if isinstance(self.replay_buffer, HerReplayBuffer): - replay_buffer = self.replay_buffer.replay_buffer - else: - replay_buffer = self.replay_buffer + replay_buffer = self.replay_buffer truncate_last_traj = ( self.optimize_memory_usage @@ -552,7 +527,6 @@ class OffPolicyAlgorithm(BaseAlgorithm): callback.on_rollout_start() continue_training = True - while should_collect_more_steps(train_freq, num_collected_steps, num_collected_episodes): if self.use_sde and self.sde_sample_freq > 0 and num_collected_steps % self.sde_sample_freq == 0: # Sample a new noise matrix diff --git a/stable_baselines3/her/her_replay_buffer.py b/stable_baselines3/her/her_replay_buffer.py index d4d1a43..5a438b4 100644 --- a/stable_baselines3/her/her_replay_buffer.py +++ b/stable_baselines3/her/her_replay_buffer.py @@ -1,87 +1,84 @@ +import copy import warnings -from collections import deque -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any, Dict, List, Optional, Union import numpy as np import torch as th +from gym import spaces from stable_baselines3.common.buffers import DictReplayBuffer -from stable_baselines3.common.preprocessing import get_obs_shape -from stable_baselines3.common.type_aliases import DictReplayBufferSamples +from stable_baselines3.common.type_aliases import DictReplayBufferSamples, TensorDict from stable_baselines3.common.vec_env import VecEnv, VecNormalize from stable_baselines3.her.goal_selection_strategy import KEY_TO_GOAL_STRATEGY, GoalSelectionStrategy -def get_time_limit(env: VecEnv, current_max_episode_length: Optional[int]) -> int: - """ - Get time limit from environment. - - :param env: Environment from which we want to get the time limit. - :param current_max_episode_length: Current value for max_episode_length. - :return: max episode length - """ - # try to get the attribute from environment - if current_max_episode_length is None: - try: - current_max_episode_length = env.get_attr("spec")[0].max_episode_steps - # Raise the error because the attribute is present but is None - if current_max_episode_length is None: - raise AttributeError - # if not available check if a valid value was passed as an argument - except AttributeError as e: - raise ValueError( - "The max episode length could not be inferred.\n" - "You must specify a `max_episode_steps` when registering the environment,\n" - "use a `gym.wrappers.TimeLimit` wrapper " - "or pass `max_episode_length` to the model constructor" - ) from e - return current_max_episode_length - - class HerReplayBuffer(DictReplayBuffer): """ Hindsight Experience Replay (HER) buffer. Paper: https://arxiv.org/abs/1707.01495 - .. warning:: - - For performance reasons, the maximum number of steps per episodes must be specified. - In most cases, it will be inferred if you specify ``max_episode_steps`` when registering the environment - or if you use a ``gym.wrappers.TimeLimit`` (and ``env.spec`` is not None). - Otherwise, you can directly pass ``max_episode_length`` to the replay buffer constructor. - - Replay buffer for sampling HER (Hindsight Experience Replay) transitions. - In the online sampling case, these new transitions will not be saved in the replay buffer - and will only be created at sampling time. + .. note:: + + Compared to other implementations, the ``future`` goal sampling strategy is inclusive: + the current transition can be used when re-sampling. + + :param buffer_size: Max number of element in the buffer + :param observation_space: Observation space + :param action_space: Action space :param env: The training environment - :param buffer_size: The size of the buffer measured in transitions. - :param max_episode_length: The maximum length of an episode. If not specified, - it will be automatically inferred if the environment uses a ``gym.wrappers.TimeLimit`` wrapper. - :param goal_selection_strategy: Strategy for sampling goals for replay. - One of ['episode', 'final', 'future'] :param device: PyTorch device - :param n_sampled_goal: Number of virtual transitions to create per real transition, - by sampling new goals. + :param n_envs: Number of parallel environments + :param optimize_memory_usage: Enable a memory efficient variant + Disabled for now (see https://github.com/DLR-RM/stable-baselines3/pull/243#discussion_r531535702) :param handle_timeout_termination: Handle timeout termination (due to timelimit) separately and treat the task as infinite horizon task. https://github.com/DLR-RM/stable-baselines3/issues/284 + :param n_sampled_goal: Number of virtual transitions to create per real transition, + by sampling new goals. + :param goal_selection_strategy: Strategy for sampling goals for replay. + One of ['episode', 'final', 'future'] + :param copy_info_dict: Whether to copy the info dictionary and pass it to + ``compute_reward()`` method. + Please note that the copy may cause a slowdown. + False by default. """ def __init__( self, - env: VecEnv, buffer_size: int, + observation_space: spaces.Space, + action_space: spaces.Space, + env: VecEnv, device: Union[th.device, str] = "auto", - replay_buffer: Optional[DictReplayBuffer] = None, - max_episode_length: Optional[int] = None, + n_envs: int = 1, + optimize_memory_usage: bool = False, + handle_timeout_termination: bool = True, n_sampled_goal: int = 4, goal_selection_strategy: Union[GoalSelectionStrategy, str] = "future", - online_sampling: bool = True, - handle_timeout_termination: bool = True, + copy_info_dict: bool = False, + online_sampling: Optional[bool] = None, ): - super().__init__(buffer_size, env.observation_space, env.action_space, device, env.num_envs) + super().__init__( + buffer_size, + observation_space, + action_space, + device=device, + n_envs=n_envs, + optimize_memory_usage=optimize_memory_usage, + handle_timeout_termination=handle_timeout_termination, + ) + self.env = env + self.copy_info_dict = copy_info_dict + + if online_sampling is not None: + assert online_sampling is True, "Since v1.8.0, SB3 only supports online sampling with HerReplayBuffer." + warnings.warn( + "Since v1.8.0, the `online_sampling` argument is deprecated " + "as SB3 only supports online sampling with HerReplayBuffer. It will be removed in v2.0", + stacklevel=1, + ) # convert goal_selection_strategy into GoalSelectionStrategy if string if isinstance(goal_selection_strategy, str): @@ -95,67 +92,24 @@ class HerReplayBuffer(DictReplayBuffer): ), f"Invalid goal selection strategy, please use one of {list(GoalSelectionStrategy)}" self.n_sampled_goal = n_sampled_goal - # if we sample her transitions online use custom replay buffer - self.online_sampling = online_sampling - # compute ratio between HER replays and regular replays in percent for online HER sampling + + # Compute ratio between HER replays and regular replays in percent self.her_ratio = 1 - (1.0 / (self.n_sampled_goal + 1)) - # maximum steps in episode - self.max_episode_length = get_time_limit(env, max_episode_length) - # storage for transitions of current episode for offline sampling - # for online sampling, it replaces the "classic" replay buffer completely - her_buffer_size = buffer_size if online_sampling else self.max_episode_length - - self.env = env - self.buffer_size = her_buffer_size - - if online_sampling: - replay_buffer = None - self.replay_buffer = replay_buffer - self.online_sampling = online_sampling - - # Handle timeouts termination properly if needed - # see https://github.com/DLR-RM/stable-baselines3/issues/284 - self.handle_timeout_termination = handle_timeout_termination - - # buffer with episodes - # number of episodes which can be stored until buffer size is reached - self.max_episode_stored = self.buffer_size // self.max_episode_length - self.current_idx = 0 - # Counter to prevent overflow - self.episode_steps = 0 - - # Get shape of observation and goal (usually the same) - self.obs_shape = get_obs_shape(self.env.observation_space.spaces["observation"]) - self.goal_shape = get_obs_shape(self.env.observation_space.spaces["achieved_goal"]) - - # input dimensions for buffer initialization - input_shape = { - "observation": (self.env.num_envs, *self.obs_shape), - "achieved_goal": (self.env.num_envs, *self.goal_shape), - "desired_goal": (self.env.num_envs, *self.goal_shape), - "action": (self.action_dim,), - "reward": (1,), - "next_obs": (self.env.num_envs, *self.obs_shape), - "next_achieved_goal": (self.env.num_envs, *self.goal_shape), - "next_desired_goal": (self.env.num_envs, *self.goal_shape), - "done": (1,), - } - self._observation_keys = ["observation", "achieved_goal", "desired_goal"] - self._buffer = { - key: np.zeros((self.max_episode_stored, self.max_episode_length, *dim), dtype=np.float32) - for key, dim in input_shape.items() - } - # Store info dicts are it can be used to compute the reward (e.g. continuity cost) - self.info_buffer = [deque(maxlen=self.max_episode_length) for _ in range(self.max_episode_stored)] - # episode length storage, needed for episodes which has less steps than the maximum length - self.episode_lengths = np.zeros(self.max_episode_stored, dtype=np.int64) + # In some environments, the info dict is used to compute the reward. Then, we need to store it. + self.infos = np.array([[{} for _ in range(self.n_envs)] for _ in range(self.buffer_size)]) + # To create virtual transitions, we need to know for each transition + # when an episode starts and ends. + # We use the following arrays to store the indices, + # and update them when an episode ends. + self.ep_start = np.zeros((self.buffer_size, self.n_envs), dtype=np.int64) + self.ep_length = np.zeros((self.buffer_size, self.n_envs), dtype=np.int64) + self._current_ep_start = np.zeros(self.n_envs, dtype=np.int64) def __getstate__(self) -> Dict[str, Any]: """ Gets state for pickling. Excludes self.env, as in general Env's may not be pickleable. - Note: when using offline sampling, this will also save the offline replay buffer. """ state = self.__dict__.copy() # these attributes are not pickleable @@ -185,347 +139,257 @@ class HerReplayBuffer(DictReplayBuffer): self.env = env - def _get_samples(self, batch_inds: np.ndarray, env: Optional[VecNormalize] = None) -> DictReplayBufferSamples: - """ - Abstract method from base class. - """ - raise NotImplementedError() - - def sample(self, batch_size: int, env: Optional[VecNormalize] = None) -> DictReplayBufferSamples: - """ - Sample function for online sampling of HER transition, - this replaces the "regular" replay buffer ``sample()`` - method in the ``train()`` function. - - :param batch_size: Number of element to sample - :param env: Associated gym VecEnv - to normalize the observations/rewards when sampling - :return: Samples. - """ - if self.replay_buffer is not None: - return self.replay_buffer.sample(batch_size, env) - return self._sample_transitions(batch_size, maybe_vec_env=env, online_sampling=True) # pytype: disable=bad-return-type - - def _sample_offline( - self, - n_sampled_goal: Optional[int] = None, - ) -> Tuple[Dict[str, np.ndarray], Dict[str, np.ndarray], np.ndarray, np.ndarray]: - """ - Sample function for offline sampling of HER transition, - in that case, only one episode is used and transitions - are added to the regular replay buffer. - - :param n_sampled_goal: Number of sampled goals for replay - :return: at most(n_sampled_goal * episode_length) HER transitions. - """ - # `maybe_vec_env=None` as we should store unnormalized transitions, - # they will be normalized at sampling time - return self._sample_transitions( - batch_size=None, - maybe_vec_env=None, - online_sampling=False, - n_sampled_goal=n_sampled_goal, - ) # pytype: disable=bad-return-type - - def sample_goals( - self, - episode_indices: np.ndarray, - her_indices: np.ndarray, - transitions_indices: np.ndarray, - ) -> np.ndarray: - """ - Sample goals based on goal_selection_strategy. - This is a vectorized (fast) version. - - :param episode_indices: Episode indices to use. - :param her_indices: HER indices. - :param transitions_indices: Transition indices to use. - :return: Return sampled goals. - """ - her_episode_indices = episode_indices[her_indices] - - if self.goal_selection_strategy == GoalSelectionStrategy.FINAL: - # replay with final state of current episode - transitions_indices = self.episode_lengths[her_episode_indices] - 1 - - elif self.goal_selection_strategy == GoalSelectionStrategy.FUTURE: - # replay with random state which comes from the same episode and was observed after current transition - transitions_indices = np.random.randint( - transitions_indices[her_indices], self.episode_lengths[her_episode_indices] - ) - - elif self.goal_selection_strategy == GoalSelectionStrategy.EPISODE: - # replay with random state which comes from the same episode as current transition - transitions_indices = np.random.randint(self.episode_lengths[her_episode_indices]) - - else: - raise ValueError(f"Strategy {self.goal_selection_strategy} for sampling goals not supported!") - - return self._buffer["next_achieved_goal"][her_episode_indices, transitions_indices] - - def _sample_transitions( - self, - batch_size: Optional[int], - maybe_vec_env: Optional[VecNormalize], - online_sampling: bool, - n_sampled_goal: Optional[int] = None, - ) -> Union[DictReplayBufferSamples, Tuple[Dict[str, np.ndarray], Dict[str, np.ndarray], np.ndarray, np.ndarray]]: - """ - :param batch_size: Number of element to sample (only used for online sampling) - :param env: associated gym VecEnv to normalize the observations/rewards - Only valid when using online sampling - :param online_sampling: Using online_sampling for HER or not. - :param n_sampled_goal: Number of sampled goals for replay. (offline sampling) - :return: Samples. - """ - # Select which episodes to use - if online_sampling: - assert batch_size is not None, "No batch_size specified for online sampling of HER transitions" - # Do not sample the episode with index `self.pos` as the episode is invalid - if self.full: - episode_indices = ( - np.random.randint(1, self.n_episodes_stored, batch_size) + self.pos - ) % self.n_episodes_stored - else: - episode_indices = np.random.randint(0, self.n_episodes_stored, batch_size) - # A subset of the transitions will be relabeled using HER algorithm - her_indices = np.arange(batch_size)[: int(self.her_ratio * batch_size)] - else: - assert maybe_vec_env is None, "Transitions must be stored unnormalized in the replay buffer" - assert n_sampled_goal is not None, "No n_sampled_goal specified for offline sampling of HER transitions" - # Offline sampling: there is only one episode stored - episode_length = self.episode_lengths[0] - # we sample n_sampled_goal per timestep in the episode (only one is stored). - episode_indices = np.tile(0, (episode_length * n_sampled_goal)) - # we only sample virtual transitions - # as real transitions are already stored in the replay buffer - her_indices = np.arange(len(episode_indices)) - - ep_lengths = self.episode_lengths[episode_indices] - - if online_sampling: - # Select which transitions to use - transitions_indices = np.random.randint(ep_lengths) - else: - if her_indices.size == 0: - # Episode of one timestep, not enough for using the "future" strategy - # no virtual transitions are created in that case - return {}, {}, np.zeros(0), np.zeros(0) - else: - # Repeat every transition index n_sampled_goals times - # to sample n_sampled_goal per timestep in the episode (only one is stored). - # Now with the corrected episode length when using "future" strategy - transitions_indices = np.tile(np.arange(ep_lengths[0]), n_sampled_goal) - episode_indices = episode_indices[transitions_indices] - her_indices = np.arange(len(episode_indices)) - - # get selected transitions - transitions = {key: self._buffer[key][episode_indices, transitions_indices].copy() for key in self._buffer.keys()} - - # sample new desired goals and relabel the transitions - new_goals = self.sample_goals(episode_indices, her_indices, transitions_indices) - transitions["desired_goal"][her_indices] = new_goals - - # Convert info buffer to numpy array - transitions["info"] = np.array( - [ - self.info_buffer[episode_idx][transition_idx] - for episode_idx, transition_idx in zip(episode_indices, transitions_indices) - ] - ) - - # Edge case: episode of one timesteps with the future strategy - # no virtual transition can be created - if len(her_indices) > 0: - # Vectorized computation of the new reward - transitions["reward"][her_indices, 0] = self.env.env_method( - "compute_reward", - # the new state depends on the previous state and action - # s_{t+1} = f(s_t, a_t) - # so the next_achieved_goal depends also on the previous state and action - # because we are in a GoalEnv: - # r_t = reward(s_t, a_t) = reward(next_achieved_goal, desired_goal) - # therefore we have to use "next_achieved_goal" and not "achieved_goal" - transitions["next_achieved_goal"][her_indices, 0], - # here we use the new desired goal - transitions["desired_goal"][her_indices, 0], - transitions["info"][her_indices, 0], - ) - - # concatenate observation with (desired) goal - observations = self._normalize_obs(transitions, maybe_vec_env) - - # HACK to make normalize obs and `add()` work with the next observation - next_observations = { - "observation": transitions["next_obs"], - "achieved_goal": transitions["next_achieved_goal"], - # The desired goal for the next observation must be the same as the previous one - "desired_goal": transitions["desired_goal"], - } - next_observations = self._normalize_obs(next_observations, maybe_vec_env) - - if online_sampling: - next_obs = {key: self.to_torch(next_observations[key][:, 0, :]) for key in self._observation_keys} - - normalized_obs = {key: self.to_torch(observations[key][:, 0, :]) for key in self._observation_keys} - - return DictReplayBufferSamples( - observations=normalized_obs, - actions=self.to_torch(transitions["action"]), - next_observations=next_obs, - dones=self.to_torch(transitions["done"]), - rewards=self.to_torch(self._normalize_reward(transitions["reward"], maybe_vec_env)), - ) - else: - return observations, next_observations, transitions["action"], transitions["reward"] - def add( self, - obs: Dict[str, np.ndarray], - next_obs: Dict[str, np.ndarray], + obs: TensorDict, + next_obs: TensorDict, action: np.ndarray, reward: np.ndarray, done: np.ndarray, infos: List[Dict[str, Any]], ) -> None: - if self.current_idx == 0 and self.full: - # Clear info buffer - self.info_buffer[self.pos] = deque(maxlen=self.max_episode_length) + # When the buffer is full, we rewrite on old episodes. When we start to + # rewrite on an old episodes, we want the whole old episode to be deleted + # (and not only the transition on which we rewrite). To do this, we set + # the length of the old episode to 0, so it can't be sampled anymore. + for env_idx in range(self.n_envs): + episode_start = self.ep_start[self.pos, env_idx] + episode_length = self.ep_length[self.pos, env_idx] + if episode_length > 0: + episode_end = episode_start + episode_length + episode_indices = np.arange(self.pos, episode_end) % self.buffer_size + self.ep_length[episode_indices, env_idx] = 0 - # Remove termination signals due to timeout - if self.handle_timeout_termination: - done_ = done * (1 - np.array([info.get("TimeLimit.truncated", False) for info in infos])) - else: - done_ = done + # Update episode start + self.ep_start[self.pos] = self._current_ep_start.copy() - self._buffer["observation"][self.pos][self.current_idx] = obs["observation"] - self._buffer["achieved_goal"][self.pos][self.current_idx] = obs["achieved_goal"] - self._buffer["desired_goal"][self.pos][self.current_idx] = obs["desired_goal"] - self._buffer["action"][self.pos][self.current_idx] = action - self._buffer["done"][self.pos][self.current_idx] = done_ - self._buffer["reward"][self.pos][self.current_idx] = reward - self._buffer["next_obs"][self.pos][self.current_idx] = next_obs["observation"] - self._buffer["next_achieved_goal"][self.pos][self.current_idx] = next_obs["achieved_goal"] - self._buffer["next_desired_goal"][self.pos][self.current_idx] = next_obs["desired_goal"] + if self.copy_info_dict: + self.infos[self.pos] = infos + # Store the transition + super().add(obs, next_obs, action, reward, done, infos) - # When doing offline sampling - # Add real transition to normal replay buffer - if self.replay_buffer is not None: - self.replay_buffer.add( - obs, - next_obs, - action, - reward, - done, - infos, + # When episode ends, compute and store the episode length + for env_idx in range(self.n_envs): + if done[env_idx]: + episode_start = self._current_ep_start[env_idx] + episode_end = self.pos + if episode_end < episode_start: + # Occurs when the buffer becomes full, the storage resumes at the + # beginning of the buffer. This can happen in the middle of an episode. + episode_end += self.buffer_size + episode_indices = np.arange(episode_start, episode_end) % self.buffer_size + self.ep_length[episode_indices, env_idx] = episode_end - episode_start + # Update the current episode start + self._current_ep_start[env_idx] = self.pos + + def sample(self, batch_size: int, env: Optional[VecNormalize] = None) -> DictReplayBufferSamples: + """ + Sample elements from the replay buffer. + + :param batch_size: Number of element to sample + :param env: Associated VecEnv to normalize the observations/rewards when sampling + :return: Samples + """ + # When the buffer is full, we rewrite on old episodes. We don't want to + # sample incomplete episode transitions, so we have to eliminate some indexes. + is_valid = self.ep_length > 0 + if not np.any(is_valid): + raise RuntimeError( + "Unable to sample before the end of the first episode. We recommend choosing a value " + "for learning_starts that is greater than the maximum number of timesteps in the environment." ) + # Get the indices of valid transitions + # Example: + # if is_valid = [[True, False, False], [True, False, True]], + # is_valid has shape (buffer_size=2, n_envs=3) + # then valid_indices = [0, 3, 5] + # they correspond to is_valid[0, 0], is_valid[1, 0] and is_valid[1, 2] + # or in numpy format ([rows], [columns]): (array([0, 1, 1]), array([0, 0, 2])) + # Those indices are obtained back using np.unravel_index(valid_indices, is_valid.shape) + valid_indices = np.flatnonzero(is_valid) + # Sample valid transitions that will constitute the minibatch of size batch_size + sampled_indices = np.random.choice(valid_indices, size=batch_size, replace=True) + # Unravel the indexes, i.e. recover the batch and env indices. + # Example: if sampled_indices = [0, 3, 5], then batch_indices = [0, 1, 1] and env_indices = [0, 0, 2] + batch_indices, env_indices = np.unravel_index(sampled_indices, is_valid.shape) - self.info_buffer[self.pos].append(infos) + # Split the indexes between real and virtual transitions. + nb_virtual = int(self.her_ratio * batch_size) + virtual_batch_indices, real_batch_indices = np.split(batch_indices, [nb_virtual]) + virtual_env_indices, real_env_indices = np.split(env_indices, [nb_virtual]) - # update current pointer - self.current_idx += 1 + # Get real and virtual data + real_data = self._get_real_samples(real_batch_indices, real_env_indices, env) + # Create virtual transitions by sampling new desired goals and computing new rewards + virtual_data = self._get_virtual_samples(virtual_batch_indices, virtual_env_indices, env) - self.episode_steps += 1 + # Concatenate real and virtual data + observations = { + key: th.cat((real_data.observations[key], virtual_data.observations[key])) + for key in virtual_data.observations.keys() + } + actions = th.cat((real_data.actions, virtual_data.actions)) + next_observations = { + key: th.cat((real_data.next_observations[key], virtual_data.next_observations[key])) + for key in virtual_data.next_observations.keys() + } + dones = th.cat((real_data.dones, virtual_data.dones)) + rewards = th.cat((real_data.rewards, virtual_data.rewards)) - if done or self.episode_steps >= self.max_episode_length: - self.store_episode() - if not self.online_sampling: - # sample virtual transitions and store them in replay buffer - self._sample_her_transitions() - # clear storage for current episode - self.reset() + return DictReplayBufferSamples( + observations=observations, + actions=actions, + next_observations=next_observations, + dones=dones, + rewards=rewards, + ) - self.episode_steps = 0 - - def store_episode(self) -> None: + def _get_real_samples( + self, + batch_indices: np.ndarray, + env_indices: np.ndarray, + env: Optional[VecNormalize] = None, + ) -> DictReplayBufferSamples: """ - Increment episode counter - and reset transition pointer. - """ - # add episode length to length storage - self.episode_lengths[self.pos] = self.current_idx + Get the samples corresponding to the batch and environment indices. - # update current episode pointer - # Note: in the OpenAI implementation - # when the buffer is full, the episode replaced - # is randomly chosen - self.pos += 1 - if self.pos == self.max_episode_stored: - self.full = True - self.pos = 0 - # reset transition pointer - self.current_idx = 0 - - def _sample_her_transitions(self) -> None: - """ - Sample additional goals and store new transitions in replay buffer - when using offline sampling. + :param batch_indices: Indices of the transitions + :param env_indices: Indices of the envrionments + :param env: associated gym VecEnv to normalize the + observations/rewards when sampling, defaults to None + :return: Samples """ + # Normalize if needed and remove extra dimension (we are using only one env for now) + obs_ = self._normalize_obs({key: obs[batch_indices, env_indices, :] for key, obs in self.observations.items()}, env) + next_obs_ = self._normalize_obs( + {key: obs[batch_indices, env_indices, :] for key, obs in self.next_observations.items()}, env + ) - # Sample goals to create virtual transitions for the last episode. - observations, next_observations, actions, rewards = self._sample_offline(n_sampled_goal=self.n_sampled_goal) + # Convert to torch tensor + observations = {key: self.to_torch(obs) for key, obs in obs_.items()} + next_observations = {key: self.to_torch(obs) for key, obs in next_obs_.items()} - # Store virtual transitions in the replay buffer, if available - if len(observations) > 0: - for i in range(len(observations["observation"])): - self.replay_buffer.add( - {key: obs[i] for key, obs in observations.items()}, - {key: next_obs[i] for key, next_obs in next_observations.items()}, - actions[i], - rewards[i], - # We consider the transition as non-terminal - done=[False], - infos=[{}], - ) + return DictReplayBufferSamples( + observations=observations, + actions=self.to_torch(self.actions[batch_indices, 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_indices, env_indices] * (1 - self.timeouts[batch_indices, env_indices]) + ).reshape(-1, 1), + rewards=self.to_torch(self._normalize_reward(self.rewards[batch_indices, env_indices].reshape(-1, 1), env)), + ) - @property - def n_episodes_stored(self) -> int: - if self.full: - return self.max_episode_stored - return self.pos + def _get_virtual_samples( + self, + batch_indices: np.ndarray, + env_indices: np.ndarray, + env: Optional[VecNormalize] = None, + ) -> DictReplayBufferSamples: + """ + Get the samples, sample new desired goals and compute new rewards. - def size(self) -> int: + :param batch_indices: Indices of the transitions + :param env_indices: Indices of the envrionments + :param env: associated gym VecEnv to normalize the + observations/rewards when sampling, defaults to None + :return: Samples, with new desired goals and new rewards """ - :return: The current number of transitions in the buffer. - """ - return int(np.sum(self.episode_lengths)) + # Get infos and obs + obs = {key: obs[batch_indices, env_indices, :] for key, obs in self.observations.items()} + next_obs = {key: obs[batch_indices, env_indices, :] for key, obs in self.next_observations.items()} + if self.copy_info_dict: + # The copy may cause a slow down + infos = copy.deepcopy(self.infos[batch_indices, env_indices]) + else: + infos = [{} for _ in range(len(batch_indices))] + # Sample and set new goals + new_goals = self._sample_goals(batch_indices, env_indices) + obs["desired_goal"] = new_goals + # The desired goal for the next observation must be the same as the previous one + next_obs["desired_goal"] = new_goals - def reset(self) -> None: + # Compute new reward + rewards = self.env.env_method( + "compute_reward", + # here we use the new desired goal + obs["desired_goal"], + # the new state depends on the previous state and action + # s_{t+1} = f(s_t, a_t) + # so the next achieved_goal depends also on the previous state and action + # because we are in a GoalEnv: + # r_t = reward(s_t, a_t) = reward(next_achieved_goal, desired_goal) + # therefore we have to use next_obs["achieved_goal"] and not obs["achieved_goal"] + next_obs["achieved_goal"], + infos, + # we use the method of the first environment assuming that all environments are identical. + indices=[0], + ) + rewards = rewards[0].astype(np.float32) # env_method returns a list containing one element + obs = self._normalize_obs(obs, env) + next_obs = self._normalize_obs(next_obs, env) + + # Convert to torch tensor + observations = {key: self.to_torch(obs) for key, obs in obs.items()} + next_observations = {key: self.to_torch(obs) for key, obs in next_obs.items()} + + return DictReplayBufferSamples( + observations=observations, + actions=self.to_torch(self.actions[batch_indices, 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_indices, env_indices] * (1 - self.timeouts[batch_indices, env_indices]) + ).reshape(-1, 1), + rewards=self.to_torch(self._normalize_reward(rewards.reshape(-1, 1), env)), + ) + + def _sample_goals(self, batch_indices: np.ndarray, env_indices: np.ndarray) -> np.ndarray: """ - Reset the buffer. + Sample goals based on goal_selection_strategy. + + :param batch_indices: Indices of the transitions + :param env_indices: Indices of the envrionments + :return: Sampled goals """ - self.pos = 0 - self.current_idx = 0 - self.full = False - self.episode_lengths = np.zeros(self.max_episode_stored, dtype=np.int64) + batch_ep_start = self.ep_start[batch_indices, env_indices] + batch_ep_length = self.ep_length[batch_indices, env_indices] + + if self.goal_selection_strategy == GoalSelectionStrategy.FINAL: + # Replay with final state of current episode + transition_indices_in_episode = batch_ep_length - 1 + + elif self.goal_selection_strategy == GoalSelectionStrategy.FUTURE: + # Replay with random state which comes from the same episode and was observed after current transition + # Note: our implementation is inclusive: current transition can be sampled + current_indices_in_episode = (batch_indices - batch_ep_start) % self.buffer_size + transition_indices_in_episode = np.random.randint(current_indices_in_episode, batch_ep_length) + + elif self.goal_selection_strategy == GoalSelectionStrategy.EPISODE: + # Replay with random state which comes from the same episode as current transition + transition_indices_in_episode = np.random.randint(0, batch_ep_length) + + else: + raise ValueError(f"Strategy {self.goal_selection_strategy} for sampling goals not supported!") + + transition_indices = (transition_indices_in_episode + batch_ep_start) % self.buffer_size + return self.next_observations["achieved_goal"][transition_indices, env_indices] def truncate_last_trajectory(self) -> None: """ - Only for online sampling, called when loading the replay buffer. If called, we assume that the last trajectory in the replay buffer was finished (and truncate it). If not called, we assume that we continue the same trajectory (same episode). """ # If we are at the start of an episode, no need to truncate - current_idx = self.current_idx - - # truncate interrupted episode - if current_idx > 0: + if (self.ep_start[self.pos] != self.pos).any(): warnings.warn( "The last trajectory in the replay buffer will be truncated.\n" "If you are in the same episode as when the replay buffer was saved,\n" "you should use `truncate_last_trajectory=False` to avoid that issue." ) - # get current episode and transition index - pos = self.pos - # set episode length for current episode - self.episode_lengths[pos] = current_idx - # set done = True for current episode - # current_idx was already incremented - self._buffer["done"][pos][current_idx - 1] = np.array([True], dtype=np.float32) - # reset current transition index - self.current_idx = 0 - # increment episode counter - self.pos = (self.pos + 1) % self.max_episode_stored - # update "full" indicator - self.full = self.full or self.pos == 0 + self.ep_start[-1] = self.pos + # set done = True for current episodes + self.dones[self.pos - 1] = True diff --git a/stable_baselines3/version.txt b/stable_baselines3/version.txt index 13ef2a8..cba76ac 100644 --- a/stable_baselines3/version.txt +++ b/stable_baselines3/version.txt @@ -1 +1 @@ -1.8.0a9 +1.8.0a10 diff --git a/tests/test_callbacks.py b/tests/test_callbacks.py index 420a16a..a9bdc43 100644 --- a/tests/test_callbacks.py +++ b/tests/test_callbacks.py @@ -140,7 +140,6 @@ def test_eval_success_logging(tmp_path): replay_buffer_class=HerReplayBuffer, learning_starts=100, seed=0, - replay_buffer_kwargs=dict(max_episode_length=n_bits), ) model.learn(500, callback=eval_callback) assert len(eval_callback._is_success_buffer) > 0 diff --git a/tests/test_her.py b/tests/test_her.py index 888d36a..f9794d5 100644 --- a/tests/test_her.py +++ b/tests/test_her.py @@ -3,19 +3,18 @@ import pathlib import warnings from copy import deepcopy -import gym import numpy as np import pytest import torch as th from stable_baselines3 import DDPG, DQN, SAC, TD3, HerReplayBuffer +from stable_baselines3.common.env_util import make_vec_env from stable_baselines3.common.envs import BitFlippingEnv from stable_baselines3.common.evaluation import evaluate_policy from stable_baselines3.common.monitor import Monitor from stable_baselines3.common.noise import NormalActionNoise -from stable_baselines3.common.vec_env import DummyVecEnv +from stable_baselines3.common.vec_env import SubprocVecEnv from stable_baselines3.her.goal_selection_strategy import GoalSelectionStrategy -from stable_baselines3.her.her_replay_buffer import get_time_limit def test_import_error(): @@ -27,18 +26,22 @@ def test_import_error(): @pytest.mark.parametrize("model_class", [SAC, TD3, DDPG, DQN]) -@pytest.mark.parametrize("online_sampling", [True, False]) @pytest.mark.parametrize("image_obs_space", [True, False]) -def test_her(model_class, online_sampling, image_obs_space): +def test_her(model_class, image_obs_space): """ Test Hindsight Experience Replay. """ + n_envs = 1 n_bits = 4 - env = BitFlippingEnv( - n_bits=n_bits, - continuous=not (model_class == DQN), - image_obs_space=image_obs_space, - ) + + def env_fn(): + return BitFlippingEnv( + n_bits=n_bits, + continuous=not (model_class == DQN), + image_obs_space=image_obs_space, + ) + + env = make_vec_env(env_fn, n_envs) model = model_class( "MultiInputPolicy", @@ -47,18 +50,28 @@ def test_her(model_class, online_sampling, image_obs_space): replay_buffer_kwargs=dict( n_sampled_goal=2, goal_selection_strategy="future", - online_sampling=online_sampling, - max_episode_length=n_bits, + copy_info_dict=True, ), train_freq=4, - gradient_steps=1, + gradient_steps=n_envs, policy_kwargs=dict(net_arch=[64]), learning_starts=100, buffer_size=int(2e4), ) model.learn(total_timesteps=150) - evaluate_policy(model, Monitor(env)) + evaluate_policy(model, Monitor(env_fn())) + + +@pytest.mark.parametrize("model_class", [TD3, DQN]) +@pytest.mark.parametrize("image_obs_space", [True, False]) +def test_multiprocessing(model_class, image_obs_space): + def env_fn(): + return BitFlippingEnv(n_bits=4, continuous=not (model_class == DQN), image_obs_space=image_obs_space) + + env = make_vec_env(env_fn, n_envs=2, vec_env_cls=SubprocVecEnv) + model = model_class("MultiInputPolicy", env, replay_buffer_class=HerReplayBuffer, buffer_size=int(2e4), train_freq=4) + model.learn(total_timesteps=150) @pytest.mark.parametrize( @@ -72,12 +85,16 @@ def test_her(model_class, online_sampling, image_obs_space): GoalSelectionStrategy.FUTURE, ], ) -@pytest.mark.parametrize("online_sampling", [True, False]) -def test_goal_selection_strategy(goal_selection_strategy, online_sampling): +def test_goal_selection_strategy(goal_selection_strategy): """ Test different goal strategies. """ - env = BitFlippingEnv(continuous=True) + n_envs = 2 + + def env_fn(): + return BitFlippingEnv(continuous=True) + + env = make_vec_env(env_fn, n_envs) normal_action_noise = NormalActionNoise(np.zeros(1), 0.1 * np.ones(1)) @@ -87,12 +104,10 @@ def test_goal_selection_strategy(goal_selection_strategy, online_sampling): replay_buffer_class=HerReplayBuffer, replay_buffer_kwargs=dict( goal_selection_strategy=goal_selection_strategy, - online_sampling=online_sampling, - max_episode_length=10, n_sampled_goal=2, ), train_freq=4, - gradient_steps=1, + gradient_steps=n_envs, policy_kwargs=dict(net_arch=[64]), learning_starts=100, buffer_size=int(1e5), @@ -104,16 +119,20 @@ def test_goal_selection_strategy(goal_selection_strategy, online_sampling): @pytest.mark.parametrize("model_class", [SAC, TD3, DDPG, DQN]) @pytest.mark.parametrize("use_sde", [False, True]) -@pytest.mark.parametrize("online_sampling", [False, True]) -def test_save_load(tmp_path, model_class, use_sde, online_sampling): +def test_save_load(tmp_path, model_class, use_sde): """ Test if 'save' and 'load' saves and loads model correctly """ if use_sde and model_class != SAC: pytest.skip("Only SAC has gSDE support") + n_envs = 2 n_bits = 4 - env = BitFlippingEnv(n_bits=n_bits, continuous=not (model_class == DQN)) + + def env_fn(): + return BitFlippingEnv(n_bits=n_bits, continuous=not (model_class == DQN)) + + env = make_vec_env(env_fn, n_envs) kwargs = dict(use_sde=True) if use_sde else {} @@ -125,8 +144,6 @@ def test_save_load(tmp_path, model_class, use_sde, online_sampling): replay_buffer_kwargs=dict( n_sampled_goal=2, goal_selection_strategy="future", - online_sampling=online_sampling, - max_episode_length=n_bits, ), verbose=0, tau=0.05, @@ -135,7 +152,7 @@ def test_save_load(tmp_path, model_class, use_sde, online_sampling): policy_kwargs=dict(net_arch=[64]), buffer_size=int(1e5), gamma=0.98, - gradient_steps=1, + gradient_steps=n_envs, train_freq=4, learning_starts=100, **kwargs @@ -143,14 +160,9 @@ def test_save_load(tmp_path, model_class, use_sde, online_sampling): model.learn(total_timesteps=150) - obs = env.reset() - - observations = {key: [] for key in obs.keys()} - for _ in range(10): - obs = env.step(env.action_space.sample())[0] - for key in obs.keys(): - observations[key].append(obs[key]) - observations = {key: np.array(obs) for key, obs in observations.items()} + env.reset() + action = np.array([env.action_space.sample() for _ in range(n_envs)]) + observations = env.step(action)[0] # Get dictionary of current parameters params = deepcopy(model.policy.state_dict()) @@ -210,9 +222,9 @@ def test_save_load(tmp_path, model_class, use_sde, online_sampling): os.remove(tmp_path / "test_save.zip") -@pytest.mark.parametrize("online_sampling", [False, True]) +@pytest.mark.parametrize("n_envs", [1, 2]) @pytest.mark.parametrize("truncate_last_trajectory", [False, True]) -def test_save_load_replay_buffer(tmp_path, recwarn, online_sampling, truncate_last_trajectory): +def test_save_load_replay_buffer(n_envs, tmp_path, recwarn, truncate_last_trajectory): """ Test if 'save_replay_buffer' and 'load_replay_buffer' works correctly """ @@ -222,7 +234,11 @@ def test_save_load_replay_buffer(tmp_path, recwarn, online_sampling, truncate_la path = pathlib.Path(tmp_path / "replay_buffer.pkl") path.parent.mkdir(exist_ok=True, parents=True) # to not raise a warning - env = BitFlippingEnv(n_bits=4, continuous=True) + + def env_fn(): + return BitFlippingEnv(n_bits=4, continuous=True) + + env = make_vec_env(env_fn, n_envs) model = SAC( "MultiInputPolicy", env, @@ -230,20 +246,16 @@ def test_save_load_replay_buffer(tmp_path, recwarn, online_sampling, truncate_la replay_buffer_kwargs=dict( n_sampled_goal=2, goal_selection_strategy="future", - online_sampling=online_sampling, - max_episode_length=4, ), - gradient_steps=1, + gradient_steps=n_envs, train_freq=4, buffer_size=int(2e4), policy_kwargs=dict(net_arch=[64]), seed=1, ) model.learn(200) - if online_sampling: - old_replay_buffer = deepcopy(model.replay_buffer) - else: - old_replay_buffer = deepcopy(model.replay_buffer.replay_buffer) + old_replay_buffer = deepcopy(model.replay_buffer) + model.save_replay_buffer(path) del model.replay_buffer @@ -262,36 +274,15 @@ def test_save_load_replay_buffer(tmp_path, recwarn, online_sampling, truncate_la else: assert len(recwarn) == 0 - if online_sampling: - n_episodes_stored = model.replay_buffer.n_episodes_stored - assert np.allclose( - old_replay_buffer._buffer["observation"][:n_episodes_stored], - model.replay_buffer._buffer["observation"][:n_episodes_stored], - ) - assert np.allclose( - old_replay_buffer._buffer["next_obs"][:n_episodes_stored], - model.replay_buffer._buffer["next_obs"][:n_episodes_stored], - ) - assert np.allclose( - old_replay_buffer._buffer["action"][:n_episodes_stored], - model.replay_buffer._buffer["action"][:n_episodes_stored], - ) - assert np.allclose( - old_replay_buffer._buffer["reward"][:n_episodes_stored], - model.replay_buffer._buffer["reward"][:n_episodes_stored], - ) - # we might change the last done of the last trajectory so we don't compare it - assert np.allclose( - old_replay_buffer._buffer["done"][: n_episodes_stored - 1], - model.replay_buffer._buffer["done"][: n_episodes_stored - 1], - ) - else: - replay_buffer = model.replay_buffer.replay_buffer - assert np.allclose(old_replay_buffer.observations["observation"], replay_buffer.observations["observation"]) - assert np.allclose(old_replay_buffer.observations["desired_goal"], replay_buffer.observations["desired_goal"]) - assert np.allclose(old_replay_buffer.actions, replay_buffer.actions) - assert np.allclose(old_replay_buffer.rewards, replay_buffer.rewards) - assert np.allclose(old_replay_buffer.dones, replay_buffer.dones) + replay_buffer = model.replay_buffer + pos = replay_buffer.pos + for key in ["observation", "desired_goal", "achieved_goal"]: + assert np.allclose(old_replay_buffer.observations[key][:pos], replay_buffer.observations[key][:pos]) + assert np.allclose(old_replay_buffer.next_observations[key][:pos], replay_buffer.next_observations[key][:pos]) + assert np.allclose(old_replay_buffer.actions[:pos], replay_buffer.actions[:pos]) + assert np.allclose(old_replay_buffer.rewards[:pos], replay_buffer.rewards[:pos]) + # we might change the last done of the last trajectory so we don't compare it + assert np.allclose(old_replay_buffer.dones[: pos - 1], replay_buffer.dones[: pos - 1]) # test if continuing training works properly reset_num_timesteps = False if truncate_last_trajectory is False else True @@ -304,7 +295,12 @@ def test_full_replay_buffer(): It should not sample the current episode which is not finished. """ n_bits = 4 - env = BitFlippingEnv(n_bits=n_bits, continuous=True) + n_envs = 2 + + def env_fn(): + return BitFlippingEnv(n_bits=n_bits, continuous=True) + + env = make_vec_env(env_fn, n_envs) # use small buffer size to get the buffer full model = SAC( @@ -314,14 +310,12 @@ def test_full_replay_buffer(): replay_buffer_kwargs=dict( n_sampled_goal=2, goal_selection_strategy="future", - online_sampling=True, - max_episode_length=n_bits, ), gradient_steps=1, train_freq=4, policy_kwargs=dict(net_arch=[64]), - learning_starts=1, - buffer_size=20, + learning_starts=n_bits * n_envs, + buffer_size=20 * n_envs, verbose=1, seed=757, ) @@ -329,49 +323,18 @@ def test_full_replay_buffer(): model.learn(total_timesteps=100) -def test_get_max_episode_length(): - dict_env = DummyVecEnv([lambda: BitFlippingEnv()]) - - # Cannot infer max epsiode length - with pytest.raises(ValueError): - get_time_limit(dict_env, current_max_episode_length=None) - - default_length = 10 - assert get_time_limit(dict_env, current_max_episode_length=default_length) == default_length - - env = gym.make("CartPole-v1") - vec_env = DummyVecEnv([lambda: env]) - - assert get_time_limit(vec_env, current_max_episode_length=None) == 500 - # Overwrite max_episode_steps - assert get_time_limit(vec_env, current_max_episode_length=default_length) == default_length - - # Set max_episode_steps to None - env.spec.max_episode_steps = None - vec_env = DummyVecEnv([lambda: env]) - with pytest.raises(ValueError): - get_time_limit(vec_env, current_max_episode_length=None) - - # Initialize HER and specify max_episode_length, should not raise an issue - DQN("MultiInputPolicy", dict_env, replay_buffer_class=HerReplayBuffer, replay_buffer_kwargs=dict(max_episode_length=5)) - - with pytest.raises(ValueError): - DQN("MultiInputPolicy", dict_env, replay_buffer_class=HerReplayBuffer) - - # Wrapped in a timelimit, should be fine - # Note: it requires env.spec to be defined - env = DummyVecEnv([lambda: gym.wrappers.TimeLimit(BitFlippingEnv(), 10)]) - DQN("MultiInputPolicy", env, replay_buffer_class=HerReplayBuffer, replay_buffer_kwargs=dict(max_episode_length=5)) - - -@pytest.mark.parametrize("online_sampling", [False, True]) @pytest.mark.parametrize("n_bits", [10]) -def test_performance_her(online_sampling, n_bits): +def test_performance_her(n_bits): """ That DQN+HER can solve BitFlippingEnv. It should not work when n_sampled_goal=0 (DQN alone). """ - env = BitFlippingEnv(n_bits=n_bits, continuous=False) + n_envs = 2 + + def env_fn(): + return BitFlippingEnv(n_bits=n_bits, continuous=False) + + env = make_vec_env(env_fn, n_envs) model = DQN( "MultiInputPolicy", @@ -380,12 +343,11 @@ def test_performance_her(online_sampling, n_bits): replay_buffer_kwargs=dict( n_sampled_goal=5, goal_selection_strategy="future", - online_sampling=online_sampling, - max_episode_length=n_bits, ), verbose=1, learning_rate=5e-4, train_freq=1, + gradient_steps=n_envs, learning_starts=100, exploration_final_eps=0.02, target_update_interval=500, diff --git a/tests/test_vec_normalize.py b/tests/test_vec_normalize.py index aca113d..27bba9a 100644 --- a/tests/test_vec_normalize.py +++ b/tests/test_vec_normalize.py @@ -338,36 +338,44 @@ def test_normalize_dict_selected_keys(): np.testing.assert_array_equal(obs["achieved_goal"], orig_obs["achieved_goal"]) -@pytest.mark.parametrize("model_class", [SAC, TD3, HerReplayBuffer]) -@pytest.mark.parametrize("online_sampling", [False, True]) -def test_offpolicy_normalization(model_class, online_sampling): - if online_sampling and model_class != HerReplayBuffer: - pytest.skip() - - make_env_ = make_dict_env if model_class == HerReplayBuffer else make_env - env = DummyVecEnv([make_env_]) +def test_her_normalization(): + env = DummyVecEnv([make_dict_env]) env = VecNormalize(env, norm_obs=True, norm_reward=True, clip_obs=10.0, clip_reward=10.0) - eval_env = DummyVecEnv([make_env_]) + eval_env = DummyVecEnv([make_dict_env]) eval_env = VecNormalize(eval_env, training=False, norm_obs=True, norm_reward=False, clip_obs=10.0, clip_reward=10.0) - if model_class == HerReplayBuffer: - model = SAC( - "MultiInputPolicy", - env, - verbose=1, - learning_starts=100, - policy_kwargs=dict(net_arch=[64]), - replay_buffer_kwargs=dict( - max_episode_length=100, - online_sampling=online_sampling, - n_sampled_goal=2, - ), - replay_buffer_class=HerReplayBuffer, - seed=2, - ) - else: - model = model_class("MlpPolicy", env, verbose=1, learning_starts=100, policy_kwargs=dict(net_arch=[64])) + model = SAC( + "MultiInputPolicy", + env, + verbose=1, + learning_starts=100, + policy_kwargs=dict(net_arch=[64]), + replay_buffer_kwargs=dict(n_sampled_goal=2), + replay_buffer_class=HerReplayBuffer, + seed=2, + ) + + # Check that VecNormalize object is correctly updated + assert model.get_vec_normalize_env() is env + model.set_env(eval_env) + assert model.get_vec_normalize_env() is eval_env + model.learn(total_timesteps=10) + model.set_env(env) + model.learn(total_timesteps=150) + # Check getter + assert isinstance(model.get_vec_normalize_env(), VecNormalize) + + +@pytest.mark.parametrize("model_class", [SAC, TD3]) +def test_offpolicy_normalization(model_class): + env = DummyVecEnv([make_env]) + env = VecNormalize(env, norm_obs=True, norm_reward=True, clip_obs=10.0, clip_reward=10.0) + + eval_env = DummyVecEnv([make_env]) + eval_env = VecNormalize(eval_env, training=False, norm_obs=True, norm_reward=False, clip_obs=10.0, clip_reward=10.0) + + model = model_class("MlpPolicy", env, verbose=1, learning_starts=100, policy_kwargs=dict(net_arch=[64])) # Check that VecNormalize object is correctly updated assert model.get_vec_normalize_env() is env