From dd6e361204e24932184a4a6b05ce3b91d29671c7 Mon Sep 17 00:00:00 2001 From: Megan Klaiber Date: Thu, 22 Oct 2020 11:56:43 +0200 Subject: [PATCH] Implement HER (#120) * Added working her version, Online sampling is missing. * Updated test_her. * Added first version of online her sampling. Still problems with tensor dimensions. * Reformat * Fixed tests * Added some comments. * Updated changelog. * Add missing init file * Fixed some small bugs. * Reduced arguments for HER, small changes. * Added getattr. Fixed bug for online sampling. * Updated save/load funtions. Small changes. * Added her to init. * Updated save method. * Updated her ratio. * Move obs_wrapper * Added DQN test. * Fix potential bug * Offline and online her share same sample_goal function. * Changed lists into arrays. * Updated her test. * Fix online sampling * Fixed action bug. Updated time limit for episodes. * Updated convert_dict method to take keys as arguments. * Renamed obs dict wrapper. * Seed bit flipping env * Remove get_episode_dict * Add fast online sampling version * Added documentation. * Vectorized reward computation * Vectorized goal sampling * Update time limit for episodes in online her sampling. * Fix max episode length inference * Bug fix for Fetch envs * Fix for HER + gSDE * Reformat (new black version) * Added info dict to compute new reward. Check her_replay_buffer again. * Fix info buffer * Updated done flag. * Fixes for gSDE * Offline her version uses now HerReplayBuffer as episode storage. * Fix num_timesteps computation * Fix get torch params * Vectorized version for offline sampling. * Modified offline her sampling to use sample method of her_replay_buffer * Updated HER tests. * Updated documentation * Cleanup docstrings * Updated to review comments * Fix pytype * Update according to review comments. * Removed random goal strategy. Updated sample transitions. * Updated migration. Removed time signal removal. * Update doc * Fix potential load issue * Add VecNormalize support for dict obs * Updated saving/loading replay buffer for HER. * Fix test memory usage * Fixed save/load replay buffer. * Fixed save/load replay buffer * Fixed transition index after loading replay buffer in online sampling * Better error handling * Add tests for get_time_limit * More tests for VecNormalize with dict obs * Update doc * Improve HER description * Add test for sde support * Add comments * Add comments * Remove check that was always valid * Fix for terminal observation * Updated buffer size in offline version and reset of HER buffer * Reformat * Update doc * Remove np.empty + add doc * Fix loading * Updated loading replay buffer * Separate online and offline sampling + bug fixes * Update tensorboard log name * Version bump * Bug fix for special case Co-authored-by: Antonin Raffin Co-authored-by: Antonin RAFFIN --- Makefile | 2 +- README.md | 13 +- docs/guide/examples.rst | 78 ++- docs/guide/migration.rst | 8 + docs/index.rst | 1 + docs/misc/changelog.rst | 10 +- docs/modules/a2c.rst | 21 +- docs/modules/ddpg.rst | 8 +- docs/modules/dqn.rst | 4 + docs/modules/her.rst | 121 ++++ docs/modules/ppo.rst | 19 + docs/modules/sac.rst | 11 +- docs/modules/td3.rst | 8 +- setup.cfg | 1 + stable_baselines3/__init__.py | 1 + stable_baselines3/common/atari_wrappers.py | 2 +- stable_baselines3/common/base_class.py | 39 +- stable_baselines3/common/bit_flipping_env.py | 21 +- stable_baselines3/common/buffers.py | 8 +- .../common/off_policy_algorithm.py | 7 + stable_baselines3/common/policies.py | 6 +- stable_baselines3/common/utils.py | 11 +- stable_baselines3/common/vec_env/__init__.py | 11 + .../common/vec_env/obs_dict_wrapper.py | 68 +++ .../common/vec_env/vec_normalize.py | 80 ++- stable_baselines3/her/__init__.py | 4 + .../her/goal_selection_strategy.py | 26 + stable_baselines3/her/her.py | 567 ++++++++++++++++++ stable_baselines3/her/her_replay_buffer.py | 370 ++++++++++++ stable_baselines3/version.txt | 2 +- tests/test_her.py | 318 ++++++++++ tests/test_save_load.py | 8 +- tests/test_sde.py | 9 +- tests/test_vec_normalize.py | 138 ++++- 34 files changed, 1899 insertions(+), 102 deletions(-) create mode 100644 docs/modules/her.rst create mode 100644 stable_baselines3/common/vec_env/obs_dict_wrapper.py create mode 100644 stable_baselines3/her/__init__.py create mode 100644 stable_baselines3/her/goal_selection_strategy.py create mode 100644 stable_baselines3/her/her.py create mode 100644 stable_baselines3/her/her_replay_buffer.py create mode 100644 tests/test_her.py diff --git a/Makefile b/Makefile index 749bc02..9954c7d 100644 --- a/Makefile +++ b/Makefile @@ -5,7 +5,7 @@ pytest: ./scripts/run_tests.sh type: - pytype + pytype -j auto lint: # stop the build if there are Python syntax errors or undefined names diff --git a/README.md b/README.md index 94c005d..3415e4d 100644 --- a/README.md +++ b/README.md @@ -35,20 +35,9 @@ These algorithms will make it easier for the research community and industry to | Type hints | :heavy_check_mark: | -### Roadmap to V1.0 - -Please look at the issue for more details. -Planned features: - -- [ ] HER - ### Planned features (v1.1+) -- [ ] DQN extensions (prioritized replay, double q-learning, ...) -- [ ] Support for `Tuple` and `Dict` observation spaces -- [ ] Recurrent Policies -- [ ] TRPO - +Please take a look at the [Roadmap](https://github.com/DLR-RM/stable-baselines3/issues/1) and [Milestones](https://github.com/DLR-RM/stable-baselines3/milestones). ## Migration guide: from Stable-Baselines (SB2) to Stable-Baselines3 (SB3) diff --git a/docs/guide/examples.rst b/docs/guide/examples.rst index 6c44e38..4b702b7 100644 --- a/docs/guide/examples.rst +++ b/docs/guide/examples.rst @@ -18,8 +18,7 @@ notebooks: - `Atari Games`_ - `RL Baselines zoo`_ - `PyBullet`_ - -.. - `Hindsight Experience Replay`_ +- `Hindsight Experience Replay`_ .. _Getting Started: https://colab.research.google.com/github/Stable-Baselines-Team/rl-colab-notebooks/blob/sb3/stable_baselines_getting_started.ipynb .. _Training, Saving, Loading: https://colab.research.google.com/github/Stable-Baselines-Team/rl-colab-notebooks/blob/sb3/saving_loading_dqn.ipynb @@ -343,6 +342,81 @@ will compute a running average and standard deviation of input features (it can env.norm_reward = False +Hindsight Experience Replay (HER) +--------------------------------- + +For this example, we are using `Highway-Env `_ by `@eleurent `_. + + +.. image:: ../_static/img/colab-badge.svg + :target: https://colab.research.google.com/github/Stable-Baselines-Team/rl-colab-notebooks/blob/sb3/stable_baselines_her.ipynb + + +.. figure:: https://raw.githubusercontent.com/eleurent/highway-env/gh-media/docs/media/parking-env.gif + + The highway-parking-v0 environment. + +The parking env is a goal-conditioned continuous control task, in which the vehicle must park in a given space with the appropriate heading. + +.. note:: + + The hyperparameters in the following example were optimized for that environment. + + +.. code-block:: python + + import gym + import highway_env + import numpy as np + + from stable_baselines3 import HER, SAC, DDPG, TD3 + from stable_baselines3.common.noise import NormalActionNoise + + env = gym.make("parking-v0") + + # Create 4 artificial transitions per real transition + n_sampled_goal = 4 + + # SAC hyperparams: + model = HER( + "MlpPolicy", + env, + SAC, + 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, + verbose=1, + buffer_size=int(1e6), + learning_rate=1e-3, + gamma=0.95, + batch_size=256, + online_sampling=True, + policy_kwargs=dict(net_arch=[256, 256, 256]), + ) + + model.learn(int(2e5)) + model.save("her_sac_highway") + + # Load saved model + model = HER.load("her_sac_highway", env=env) + + obs = env.reset() + + # Evaluate the agent + episode_reward = 0 + for _ in range(100): + action, _ = model.predict(obs, deterministic=True) + obs, reward, done, info = env.step(action) + env.render() + episode_reward += reward + if done or info.get("is_success", False): + print("Reward:", episode_reward, "Success?", info.get("is_success", False)) + episode_reward = 0.0 + obs = env.reset() + + Record a Video -------------- diff --git a/docs/guide/migration.rst b/docs/guide/migration.rst index 53bc816..1f7f5ce 100644 --- a/docs/guide/migration.rst +++ b/docs/guide/migration.rst @@ -163,6 +163,14 @@ Despite this change, no change in performance should be expected. To match SB2 behavior, you need to explicitly pass ``deterministic=True`` +HER +^^^ + +The ``HER`` implementation now also 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/index.rst b/docs/index.rst index 76e4816..9b7edea 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -57,6 +57,7 @@ Main Features modules/a2c modules/ddpg modules/dqn + modules/her modules/ppo modules/sac modules/td3 diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index e896723..1474856 100644 --- a/docs/misc/changelog.rst +++ b/docs/misc/changelog.rst @@ -4,7 +4,7 @@ Changelog ========== -Pre-Release 0.10.0a0 (WIP) +Pre-Release 0.10.0a1 (WIP) ------------------------------ Breaking Changes: @@ -14,11 +14,14 @@ Breaking Changes: New Features: ^^^^^^^^^^^^^ - Allow custom actor/critic network architectures using ``net_arch=dict(qf=[400, 300], pi=[64, 64])`` for off-policy algorithms (SAC, TD3, DDPG) +- Added Hindsight Experience Replay ``HER``. (@megan-klaiber) +- ``VecNormalize`` now supports ``gym.spaces.Dict`` observation spaces - Support logging videos to Tensorboard (@SwamyDev) Bug Fixes: ^^^^^^^^^^ - Fix GAE computation for on-policy algorithms (off-by one for the last value) (thanks @Wovchena) +- Fixed potential issue when loading a different environment - Fix ignoring the exclude parameter when recording logs using json, csv or log as logging format (@SwamyDev) - Make ``make_vec_env`` support the ``env_kwargs`` argument when using an env ID str (@ManifoldFR) - Fix model creation initializing CUDA even when `device="cpu"` is provided @@ -37,6 +40,7 @@ Others: Documentation: ^^^^^^^^^^^^^^ - Added first draft of migration guide +- Enabled doc for ``CnnPolicies`` Pre-Release 0.9.0 (2020-10-03) @@ -68,6 +72,7 @@ New Features: Bug Fixes: ^^^^^^^^^^ +- Added ``unwrap_vec_wrapper()`` to ``common.vec_env`` to extract ``VecEnvWrapper`` if needed - Fixed a bug where the environment was reset twice when using ``evaluate_policy`` - Fix logging of ``clip_fraction`` in PPO (@diditforlulz273) - Fixed a bug where cuda support was wrongly checked when passing the GPU index, e.g., ``device="cuda:0"`` (@liorcohen5) @@ -160,7 +165,6 @@ Documentation: - Fixed typo in custom policy doc (@RaphaelWag) - Pre-Release 0.7.0 (2020-06-10) ------------------------------ @@ -461,4 +465,4 @@ And all the contributors: @MarvineGothic @jdossgollin @SyllogismRXS @rusu24edward @jbulow @Antymon @seheevic @justinkterry @edbeeching @flodorner @KuKuXia @NeoExtended @PartiallyTyped @mmcenta @richardwu @kinalmehta @rolandgvc @tkelestemur @mloo3 @tirafesi @blurLake @koulakis @joeljosephjin @shwang @rk37 @andyshih12 @RaphaelWag @xicocaio -@diditforlulz273 @liorcohen5 @ManifoldFR @mloo3 @SwamyDev @wmmc88 +@diditforlulz273 @liorcohen5 @ManifoldFR @mloo3 @SwamyDev @wmmc88 @megan-klaiber diff --git a/docs/modules/a2c.rst b/docs/modules/a2c.rst index bf223cb..90c5acc 100644 --- a/docs/modules/a2c.rst +++ b/docs/modules/a2c.rst @@ -11,7 +11,7 @@ It uses multiple workers to avoid the use of a replay buffer. .. warning:: - + If you find training unstable or want to match performance of stable-baselines A2C, consider using ``RMSpropTFLike`` optimizer from ``stable_baselines3.common.sb2_compat.rmsprop_tf_like``. You can change optimizer with ``A2C(policy_kwargs=dict(optimizer_class=RMSpropTFLike))``. @@ -79,3 +79,22 @@ Parameters .. autoclass:: A2C :members: :inherited-members: + + +A2C Policies +------------- + +.. autoclass:: MlpPolicy + :members: + :inherited-members: + +.. autoclass:: stable_baselines3.common.policies.ActorCriticPolicy + :members: + :noindex: + +.. autoclass:: CnnPolicy + :members: + +.. autoclass:: stable_baselines3.common.policies.ActorCriticCnnPolicy + :members: + :noindex: diff --git a/docs/modules/ddpg.rst b/docs/modules/ddpg.rst index dd74f3a..8add698 100644 --- a/docs/modules/ddpg.rst +++ b/docs/modules/ddpg.rst @@ -98,7 +98,9 @@ DDPG Policies :members: :inherited-members: +.. autoclass:: stable_baselines3.td3.policies.TD3Policy + :members: + :noindex: -.. .. autoclass:: CnnPolicy -.. :members: -.. :inherited-members: +.. autoclass:: CnnPolicy + :members: diff --git a/docs/modules/dqn.rst b/docs/modules/dqn.rst index 034e5b8..ca9ccca 100644 --- a/docs/modules/dqn.rst +++ b/docs/modules/dqn.rst @@ -90,5 +90,9 @@ DQN Policies :members: :inherited-members: +.. autoclass:: stable_baselines3.dqn.policies.DQNPolicy + :members: + :noindex: + .. autoclass:: CnnPolicy :members: diff --git a/docs/modules/her.rst b/docs/modules/her.rst new file mode 100644 index 0000000..355b36d --- /dev/null +++ b/docs/modules/her.rst @@ -0,0 +1,121 @@ +.. _her: + +.. automodule:: stable_baselines3.her + + +HER +==== + +`Hindsight Experience Replay (HER) `_ + +HER is an algorithm that works with off-policy methods (DQN, SAC, TD3 and DDPG for example). +HER uses the fact that even if a desired goal was not achieved, other goal may have been achieved during a rollout. +It creates "virtual" transitions by relabeling transitions (changing the desired goal) from past episodes. + + + +.. warning:: + + HER requires the environment to inherits from `gym.GoalEnv `_ + + +.. 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:: + + ``HER`` supports ``VecNormalize`` wrapper but only when ``online_sampling=True`` + + +Notes +----- + +- Original paper: https://arxiv.org/abs/1707.01495 +- OpenAI paper: `Plappert et al. (2018)`_ +- OpenAI blog post: https://openai.com/blog/ingredients-for-robotics-research/ + + +.. _Plappert et al. (2018): https://arxiv.org/abs/1802.09464 + +Can I use? +---------- + +Please refer to the used model (DQN, SAC, TD3 or DDPG) for that section. + +Example +------- + +.. code-block:: python + + from stable_baselines3 import HER, DDPG, DQN, SAC, TD3 + from stable_baselines3.her.goal_selection_strategy import GoalSelectionStrategy + from stable_baselines3.common.bit_flipping_env import BitFlippingEnv + from stable_baselines3.common.vec_env import DummyVecEnv + from stable_baselines3.common.vec_env.obs_dict_wrapper import ObsDictWrapper + + model_class = DQN # works also with SAC, DDPG and TD3 + N_BITS = 15 + + env = BitFlippingEnv(n_bits=N_BITS, continuous=model_class in [DDPG, SAC, TD3], max_steps=N_BITS) + + # 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 = HER('MlpPolicy', env, model_class, n_sampled_goal=4, goal_selection_strategy=goal_selection_strategy, online_sampling=online_sampling, + verbose=1, max_episode_length=max_episode_length) + # Train the model + model.learn(1000) + + model.save("./her_bit_env") + model = HER.load('./her_bit_env', env=env) + + obs = env.reset() + for _ in range(100): + action, _ = model.model.predict(obs, deterministic=True) + obs, reward, done, _ = env.step(action) + + if done: + obs = env.reset() + + +Parameters +---------- + +.. autoclass:: HER + :members: + +Goal Selection Strategies +------------------------- + +.. autoclass:: GoalSelectionStrategy + :members: + :inherited-members: + :undoc-members: + + +Obs Dict Wrapper +---------------- + +.. autoclass:: ObsDictWrapper + :members: + :inherited-members: + :undoc-members: + + +HER Replay Buffer +----------------- + +.. autoclass:: HerReplayBuffer + :members: + :inherited-members: diff --git a/docs/modules/ppo.rst b/docs/modules/ppo.rst index d54b970..fdcf121 100644 --- a/docs/modules/ppo.rst +++ b/docs/modules/ppo.rst @@ -80,3 +80,22 @@ Parameters .. autoclass:: PPO :members: :inherited-members: + + +PPO Policies +------------- + +.. autoclass:: MlpPolicy + :members: + :inherited-members: + +.. autoclass:: stable_baselines3.common.policies.ActorCriticPolicy + :members: + :noindex: + +.. autoclass:: CnnPolicy + :members: + +.. autoclass:: stable_baselines3.common.policies.ActorCriticCnnPolicy + :members: + :noindex: diff --git a/docs/modules/sac.rst b/docs/modules/sac.rst index 6d559d4..7b37974 100644 --- a/docs/modules/sac.rst +++ b/docs/modules/sac.rst @@ -82,7 +82,7 @@ Example obs = env.reset() while True: - action, _states = model.predict(obs) + action, _states = model.predict(obs, deterministic=True) obs, reward, done, info = env.step(action) env.render() if done: @@ -104,6 +104,9 @@ SAC Policies :members: :inherited-members: -.. .. autoclass:: CnnPolicy -.. :members: -.. :inherited-members: +.. autoclass:: stable_baselines3.sac.policies.SACPolicy + :members: + :noindex: + +.. autoclass:: CnnPolicy + :members: diff --git a/docs/modules/td3.rst b/docs/modules/td3.rst index 912fc1b..fbe6aab 100644 --- a/docs/modules/td3.rst +++ b/docs/modules/td3.rst @@ -101,7 +101,9 @@ TD3 Policies :members: :inherited-members: +.. autoclass:: stable_baselines3.td3.policies.TD3Policy + :members: + :noindex: -.. .. autoclass:: CnnPolicy -.. :members: -.. :inherited-members: +.. autoclass:: CnnPolicy + :members: diff --git a/setup.cfg b/setup.cfg index 011c3d9..4b5d439 100644 --- a/setup.cfg +++ b/setup.cfg @@ -29,6 +29,7 @@ per-file-ignores = ./stable_baselines3/a2c/__init__.py:F401 ./stable_baselines3/ddpg/__init__.py:F401 ./stable_baselines3/dqn/__init__.py:F401 + ./stable_baselines3/her/__init__.py:F401 ./stable_baselines3/ppo/__init__.py:F401 ./stable_baselines3/sac/__init__.py:F401 ./stable_baselines3/td3/__init__.py:F401 diff --git a/stable_baselines3/__init__.py b/stable_baselines3/__init__.py index b88ca5d..bcac479 100644 --- a/stable_baselines3/__init__.py +++ b/stable_baselines3/__init__.py @@ -3,6 +3,7 @@ import os from stable_baselines3.a2c import A2C from stable_baselines3.ddpg import DDPG from stable_baselines3.dqn import DQN +from stable_baselines3.her import HER from stable_baselines3.ppo import PPO from stable_baselines3.sac import SAC from stable_baselines3.td3 import TD3 diff --git a/stable_baselines3/common/atari_wrappers.py b/stable_baselines3/common/atari_wrappers.py index 7cc6836..b0c5295 100644 --- a/stable_baselines3/common/atari_wrappers.py +++ b/stable_baselines3/common/atari_wrappers.py @@ -34,7 +34,7 @@ class NoopResetEnv(gym.Wrapper): else: noops = self.unwrapped.np_random.randint(1, self.noop_max + 1) assert noops > 0 - obs = np.empty(0) + obs = np.zeros(0) for _ in range(noops): obs, _, done, _ = self.env.step(self.noop_action) if done: diff --git a/stable_baselines3/common/base_class.py b/stable_baselines3/common/base_class.py index c904925..f270c60 100644 --- a/stable_baselines3/common/base_class.py +++ b/stable_baselines3/common/base_class.py @@ -26,7 +26,15 @@ from stable_baselines3.common.utils import ( set_random_seed, update_learning_rate, ) -from stable_baselines3.common.vec_env import DummyVecEnv, VecEnv, VecNormalize, VecTransposeImage, unwrap_vec_normalize +from stable_baselines3.common.vec_env import ( + DummyVecEnv, + VecEnv, + VecNormalize, + VecTransposeImage, + is_wrapped, + unwrap_vec_normalize, +) +from stable_baselines3.common.vec_env.obs_dict_wrapper import ObsDictWrapper def maybe_make_env(env: Union[GymEnv, str, None], monitor_wrapper: bool, verbose: int) -> Optional[GymEnv]: @@ -146,7 +154,7 @@ class BaseAlgorithm(ABC): self.eval_env = maybe_make_env(env, monitor_wrapper, self.verbose) env = maybe_make_env(env, monitor_wrapper, self.verbose) - env = self._wrap_env(env) + env = self._wrap_env(env, self.verbose) self.observation_space = env.observation_space self.action_space = env.action_space @@ -158,19 +166,25 @@ class BaseAlgorithm(ABC): "Error: the model does not support multiple envs; it requires " "a single vectorized environment." ) - if self.use_sde and not isinstance(self.observation_space, gym.spaces.Box): + if self.use_sde and not isinstance(self.action_space, gym.spaces.Box): raise ValueError("generalized State-Dependent Exploration (gSDE) can only be used with continuous actions.") - def _wrap_env(self, env: GymEnv) -> VecEnv: + @staticmethod + def _wrap_env(env: GymEnv, verbose: int = 0) -> VecEnv: if not isinstance(env, VecEnv): - if self.verbose >= 1: + if verbose >= 1: print("Wrapping the env in a DummyVecEnv.") env = DummyVecEnv([lambda: env]) - if is_image_space(env.observation_space) and not isinstance(env, VecTransposeImage): - if self.verbose >= 1: + if is_image_space(env.observation_space) and not is_wrapped(env, VecTransposeImage): + if verbose >= 1: print("Wrapping the env in a VecTransposeImage.") env = VecTransposeImage(env) + + # check if wrapper for dict support is needed when using HER + if isinstance(env.observation_space, gym.spaces.dict.Dict): + env = ObsDictWrapper(env) + return env @abstractmethod @@ -188,7 +202,7 @@ class BaseAlgorithm(ABC): eval_env = self.eval_env if eval_env is not None: - eval_env = self._wrap_env(eval_env) + eval_env = self._wrap_env(eval_env, self.verbose) assert eval_env.num_envs == 1 return eval_env @@ -402,10 +416,11 @@ class BaseAlgorithm(ABC): :param env: The environment for learning a policy """ - check_for_correct_spaces(env, self.observation_space, self.action_space) - # it must be coherent now # if it is not a VecEnv, make it a VecEnv - env = self._wrap_env(env) + # and do other transformations (dict obs, image transpose) if needed + env = self._wrap_env(env, self.verbose) + # Check that the observation spaces match + check_for_correct_spaces(env, self.observation_space, self.action_space) self.n_envs = env.num_envs self.env = env @@ -576,6 +591,8 @@ class BaseAlgorithm(ABC): raise KeyError("The observation_space and action_space were not given, can't verify new environments") if env is not None: + # Wrap first if needed + cls._wrap_env(env, data["verbose"]) # Check if given env is valid check_for_correct_spaces(env, data["observation_space"], data["action_space"]) else: diff --git a/stable_baselines3/common/bit_flipping_env.py b/stable_baselines3/common/bit_flipping_env.py index 96e3cfd..62f0710 100644 --- a/stable_baselines3/common/bit_flipping_env.py +++ b/stable_baselines3/common/bit_flipping_env.py @@ -1,8 +1,9 @@ from collections import OrderedDict -from typing import Dict, Optional, Union +from typing import Any, Dict, Optional, Union import numpy as np from gym import GoalEnv, spaces +from gym.envs.registration import EnvSpec from stable_baselines3.common.type_aliases import GymStepReturn @@ -22,6 +23,8 @@ class BitFlippingEnv(GoalEnv): version or not, by default, it uses the MultiBinary one """ + spec = EnvSpec("BitFlippingEnv-v0") + def __init__( self, n_bits: int = 10, continuous: bool = False, max_steps: Optional[int] = None, discrete_obs_space: bool = False ): @@ -61,7 +64,9 @@ class BitFlippingEnv(GoalEnv): max_steps = n_bits self.max_steps = max_steps self.current_step = 0 - self.reset() + + def seed(self, seed: int) -> None: + self.obs_space.seed(seed) def convert_if_needed(self, state: np.ndarray) -> Union[int, np.ndarray]: """ @@ -101,7 +106,7 @@ class BitFlippingEnv(GoalEnv): else: self.state[action] = 1 - self.state[action] obs = self._get_obs() - reward = self.compute_reward(obs["achieved_goal"], obs["desired_goal"], None) + reward = float(self.compute_reward(obs["achieved_goal"], obs["desired_goal"], None)) done = reward == 0 self.current_step += 1 # Episode terminate when we reached the goal or the max number of steps @@ -109,11 +114,13 @@ class BitFlippingEnv(GoalEnv): done = done or self.current_step >= self.max_steps return obs, reward, done, info - def compute_reward(self, achieved_goal: np.ndarray, desired_goal: np.ndarray, _info) -> float: + def compute_reward( + self, achieved_goal: Union[int, np.ndarray], desired_goal: Union[int, np.ndarray], _info: Optional[Dict[str, Any]] + ) -> np.float32: # Deceptive reward: it is positive only when the goal is achieved - if self.discrete_obs_space: - return 0.0 if achieved_goal == desired_goal else -1.0 - return 0.0 if (achieved_goal == desired_goal).all() else -1.0 + # vectorized version + distance = np.linalg.norm(achieved_goal - desired_goal, axis=-1) + return -(distance > 0).astype(np.float32) def render(self, mode: str = "human") -> Optional[np.ndarray]: if mode == "rgb_array": diff --git a/stable_baselines3/common/buffers.py b/stable_baselines3/common/buffers.py index 563ddea..83e6554 100644 --- a/stable_baselines3/common/buffers.py +++ b/stable_baselines3/common/buffers.py @@ -1,6 +1,6 @@ import warnings from abc import ABC, abstractmethod -from typing import Generator, Optional, Union +from typing import Dict, Generator, Optional, Union import numpy as np import torch as th @@ -129,9 +129,11 @@ class BaseBuffer(ABC): return th.as_tensor(array).to(self.device) @staticmethod - def _normalize_obs(obs: np.ndarray, env: Optional[VecNormalize] = None) -> np.ndarray: + def _normalize_obs( + obs: Union[np.ndarray, Dict[str, np.ndarray]], env: Optional[VecNormalize] = None + ) -> Union[np.ndarray, Dict[str, np.ndarray]]: if env is not None: - return env.normalize_obs(obs).astype(np.float32) + return env.normalize_obs(obs) return obs @staticmethod diff --git a/stable_baselines3/common/off_policy_algorithm.py b/stable_baselines3/common/off_policy_algorithm.py index cf08b44..e3ffeb6 100644 --- a/stable_baselines3/common/off_policy_algorithm.py +++ b/stable_baselines3/common/off_policy_algorithm.py @@ -67,6 +67,8 @@ class OffPolicyAlgorithm(BaseAlgorithm): :param use_sde_at_warmup: Whether to use gSDE instead of uniform sampling during the warm up phase (before learning starts) :param sde_support: Whether the model support gSDE or not + :param remove_time_limit_termination: Remove terminations (dones) that are due to time limit. + See https://github.com/hill-a/stable-baselines/issues/863 """ def __init__( @@ -97,6 +99,7 @@ class OffPolicyAlgorithm(BaseAlgorithm): sde_sample_freq: int = -1, use_sde_at_warmup: bool = False, sde_support: bool = True, + remove_time_limit_termination: bool = False, ): super(OffPolicyAlgorithm, self).__init__( @@ -126,6 +129,10 @@ class OffPolicyAlgorithm(BaseAlgorithm): self.action_noise = action_noise self.optimize_memory_usage = optimize_memory_usage + # Remove terminations (dones) that are due to time limit + # see https://github.com/hill-a/stable-baselines/issues/863 + self.remove_time_limit_termination = remove_time_limit_termination + if train_freq > 0 and n_episodes_rollout > 0: warnings.warn( "You passed a positive value for `train_freq` and `n_episodes_rollout`." diff --git a/stable_baselines3/common/policies.py b/stable_baselines3/common/policies.py index fa21018..ef7ed05 100644 --- a/stable_baselines3/common/policies.py +++ b/stable_baselines3/common/policies.py @@ -23,6 +23,7 @@ from stable_baselines3.common.preprocessing import get_action_dim, is_image_spac from stable_baselines3.common.torch_layers import BaseFeaturesExtractor, FlattenExtractor, MlpExtractor, NatureCNN, create_mlp from stable_baselines3.common.utils import get_device, is_vectorized_observation from stable_baselines3.common.vec_env import VecTransposeImage +from stable_baselines3.common.vec_env.obs_dict_wrapper import ObsDictWrapper class BaseModel(nn.Module, ABC): @@ -234,7 +235,10 @@ class BasePolicy(BaseModel): # state = self.initial_state # if mask is None: # mask = [False for _ in range(self.n_envs)] - observation = np.array(observation) + if isinstance(observation, dict): + observation = ObsDictWrapper.convert_dict(observation) + else: + observation = np.array(observation) # Handle the different cases for images # as PyTorch use channel first format diff --git a/stable_baselines3/common/utils.py b/stable_baselines3/common/utils.py index 5787265..2f02ed7 100644 --- a/stable_baselines3/common/utils.py +++ b/stable_baselines3/common/utils.py @@ -16,9 +16,7 @@ except ImportError: SummaryWriter = None from stable_baselines3.common import logger -from stable_baselines3.common.preprocessing import is_image_space from stable_baselines3.common.type_aliases import GymEnv -from stable_baselines3.common.vec_env import VecTransposeImage def set_random_seed(seed: int, using_cuda: bool = False) -> None: @@ -204,14 +202,7 @@ def check_for_correct_spaces(env: GymEnv, observation_space: gym.spaces.Space, a :param observation_space: Observation space to check against :param action_space: Action space to check against """ - if ( - observation_space != env.observation_space - # Special cases for images that need to be transposed - and not ( - is_image_space(env.observation_space) - and observation_space == VecTransposeImage.transpose_space(env.observation_space) - ) - ): + if observation_space != env.observation_space: raise ValueError(f"Observation spaces do not match: {observation_space} != {env.observation_space}") if action_space != env.action_space: raise ValueError(f"Action spaces do not match: {action_space} != {env.action_space}") diff --git a/stable_baselines3/common/vec_env/__init__.py b/stable_baselines3/common/vec_env/__init__.py index 1940f20..0002788 100644 --- a/stable_baselines3/common/vec_env/__init__.py +++ b/stable_baselines3/common/vec_env/__init__.py @@ -41,6 +41,17 @@ def unwrap_vec_normalize(env: Union["GymEnv", VecEnv]) -> Optional[VecNormalize] return unwrap_vec_wrapper(env, VecNormalize) # pytype:disable=bad-return-type +def is_wrapped(env: Union["GymEnv", VecEnv], vec_wrapper_class: Type[VecEnvWrapper]) -> bool: + """ + Check if an environment is already wrapped by a given ``VecEnvWrapper``. + + :param env: + :param vec_wrapper_class: + :return: + """ + return unwrap_vec_wrapper(env, vec_wrapper_class) is not None + + # Define here to avoid circular import def sync_envs_normalization(env: "GymEnv", eval_env: "GymEnv") -> None: """ diff --git a/stable_baselines3/common/vec_env/obs_dict_wrapper.py b/stable_baselines3/common/vec_env/obs_dict_wrapper.py new file mode 100644 index 0000000..5b1dd1a --- /dev/null +++ b/stable_baselines3/common/vec_env/obs_dict_wrapper.py @@ -0,0 +1,68 @@ +from typing import Dict + +import numpy as np +from gym import spaces + +from stable_baselines3.common.vec_env import VecEnv, VecEnvWrapper + + +class ObsDictWrapper(VecEnvWrapper): + """ + Wrapper for a VecEnv which overrides the observation space for Hindsight Experience Replay to support dict observations. + + :param env: The vectorized environment to wrap. + """ + + def __init__(self, venv: VecEnv): + super(ObsDictWrapper, self).__init__(venv, venv.observation_space, venv.action_space) + + self.venv = venv + + self.spaces = list(venv.observation_space.spaces.values()) + + # get dimensions of observation and goal + if isinstance(self.spaces[0], spaces.Discrete): + self.obs_dim = 1 + self.goal_dim = 1 + else: + self.obs_dim = venv.observation_space.spaces["observation"].shape[0] + self.goal_dim = venv.observation_space.spaces["achieved_goal"].shape[0] + + # new observation space with concatenated observation and (desired) goal + # for the different types of spaces + if isinstance(self.spaces[0], spaces.Box): + low_values = np.concatenate( + [venv.observation_space.spaces["observation"].low, venv.observation_space.spaces["desired_goal"].low] + ) + high_values = np.concatenate( + [venv.observation_space.spaces["observation"].high, venv.observation_space.spaces["desired_goal"].high] + ) + self.observation_space = spaces.Box(low_values, high_values, dtype=np.float32) + elif isinstance(self.spaces[0], spaces.MultiBinary): + total_dim = self.obs_dim + self.goal_dim + self.observation_space = spaces.MultiBinary(total_dim) + elif isinstance(self.spaces[0], spaces.Discrete): + dimensions = [venv.observation_space.spaces["observation"].n, venv.observation_space.spaces["desired_goal"].n] + self.observation_space = spaces.MultiDiscrete(dimensions) + else: + raise NotImplementedError(f"{type(self.spaces[0])} space is not supported") + + def reset(self): + return self.venv.reset() + + def step_wait(self): + return self.venv.step_wait() + + @staticmethod + def convert_dict( + observation_dict: Dict[str, np.ndarray], observation_key: str = "observation", goal_key: str = "desired_goal" + ) -> np.ndarray: + """ + Concatenate observation and (desired) goal of observation dict. + + :param observation_dict: Dictionary with observation. + :param observation_key: Key of observation in dicitonary. + :param goal_key: Key of (desired) goal in dicitonary. + :return: Concatenated observation. + """ + return np.concatenate([observation_dict[observation_key], observation_dict[goal_key]], axis=-1) diff --git a/stable_baselines3/common/vec_env/vec_normalize.py b/stable_baselines3/common/vec_env/vec_normalize.py index 39a5d11..fcdefd8 100644 --- a/stable_baselines3/common/vec_env/vec_normalize.py +++ b/stable_baselines3/common/vec_env/vec_normalize.py @@ -1,8 +1,11 @@ import pickle -from typing import Any, Dict +from copy import deepcopy +from typing import Any, Dict, Union +import gym import numpy as np +from stable_baselines3.common import utils from stable_baselines3.common.running_mean_std import RunningMeanStd from stable_baselines3.common.vec_env.base_vec_env import VecEnv, VecEnvStepReturn, VecEnvWrapper @@ -34,7 +37,19 @@ class VecNormalize(VecEnvWrapper): epsilon: float = 1e-8, ): VecEnvWrapper.__init__(self, venv) - self.obs_rms = RunningMeanStd(shape=self.observation_space.shape) + + assert isinstance( + self.observation_space, (gym.spaces.Box, gym.spaces.Dict) + ), "VecNormalize only support `gym.spaces.Box` and `gym.spaces.Dict` observation spaces" + + if isinstance(self.observation_space, gym.spaces.Dict): + self.obs_keys = set(self.observation_space.spaces.keys()) + self.obs_spaces = self.observation_space.spaces + self.obs_rms = {key: RunningMeanStd(shape=space.shape) for key, space in self.obs_spaces.items()} + else: + self.obs_keys, self.obs_spaces = None, None + self.obs_rms = RunningMeanStd(shape=self.observation_space.shape) + self.ret_rms = RunningMeanStd(shape=()) self.clip_obs = clip_obs self.clip_reward = clip_reward @@ -83,8 +98,9 @@ class VecNormalize(VecEnvWrapper): if self.venv is not None: raise ValueError("Trying to set venv of already initialized VecNormalize wrapper.") VecEnvWrapper.__init__(self, venv) - if self.obs_rms.mean.shape != self.observation_space.shape: - raise ValueError("venv is incompatible with current statistics.") + + # Check only that the observation_space match + utils.check_for_correct_spaces(venv, self.observation_space, venv.action_space) self.ret = np.zeros(self.num_envs) def step_wait(self) -> VecEnvStepReturn: @@ -99,7 +115,12 @@ class VecNormalize(VecEnvWrapper): self.old_reward = rews if self.training: - self.obs_rms.update(obs) + if isinstance(obs, dict) and isinstance(self.obs_rms, dict): + for key in self.obs_rms.keys(): + self.obs_rms[key].update(obs[key]) + else: + self.obs_rms.update(obs) + obs = self.normalize_obs(obs) if self.training: @@ -114,14 +135,38 @@ class VecNormalize(VecEnvWrapper): self.ret = self.ret * self.gamma + reward self.ret_rms.update(self.ret) - def normalize_obs(self, obs: np.ndarray) -> np.ndarray: + def _normalize_obs(self, obs: np.ndarray, obs_rms: RunningMeanStd) -> np.ndarray: + """ + Helper to normalize observation. + :param obs: + :param obs_rms: associated statistics + :return: normalized observation + """ + return np.clip((obs - obs_rms.mean) / np.sqrt(obs_rms.var + self.epsilon), -self.clip_obs, self.clip_obs) + + def _unnormalize_obs(self, obs: np.ndarray, obs_rms: RunningMeanStd) -> np.ndarray: + """ + Helper to unnormalize observation. + :param obs: + :param obs_rms: associated statistics + :return: unnormalized observation + """ + return (obs * np.sqrt(obs_rms.var + self.epsilon)) + obs_rms.mean + + def normalize_obs(self, obs: Union[np.ndarray, Dict[str, np.ndarray]]) -> Union[np.ndarray, Dict[str, np.ndarray]]: """ Normalize observations using this VecNormalize's observations statistics. Calling this method does not update statistics. """ + # Avoid modifying by reference the original object + obs_ = deepcopy(obs) if self.norm_obs: - obs = np.clip((obs - self.obs_rms.mean) / np.sqrt(self.obs_rms.var + self.epsilon), -self.clip_obs, self.clip_obs) - return obs + if isinstance(obs, dict) and isinstance(self.obs_rms, dict): + for key in self.obs_rms.keys(): + obs_[key] = self._normalize_obs(obs[key], self.obs_rms[key]).astype(np.float32) + else: + obs_ = self._normalize_obs(obs, self.obs_rms).astype(np.float32) + return obs_ def normalize_reward(self, reward: np.ndarray) -> np.ndarray: """ @@ -132,22 +177,28 @@ class VecNormalize(VecEnvWrapper): reward = np.clip(reward / np.sqrt(self.ret_rms.var + self.epsilon), -self.clip_reward, self.clip_reward) return reward - def unnormalize_obs(self, obs: np.ndarray) -> np.ndarray: + def unnormalize_obs(self, obs: Union[np.ndarray, Dict[str, np.ndarray]]) -> Union[np.ndarray, Dict[str, np.ndarray]]: + # Avoid modifying by reference the original object + obs_ = deepcopy(obs) if self.norm_obs: - return (obs * np.sqrt(self.obs_rms.var + self.epsilon)) + self.obs_rms.mean - return obs + if isinstance(obs, dict) and isinstance(self.obs_rms, dict): + for key in self.obs_rms.keys(): + obs_[key] = self._unnormalize_obs(obs[key], self.obs_rms[key]) + else: + obs_ = self._unnormalize_obs(obs, self.obs_rms) + return obs_ def unnormalize_reward(self, reward: np.ndarray) -> np.ndarray: if self.norm_reward: return reward * np.sqrt(self.ret_rms.var + self.epsilon) return reward - def get_original_obs(self) -> np.ndarray: + def get_original_obs(self) -> Union[np.ndarray, Dict[str, np.ndarray]]: """ Returns an unnormalized version of the observations from the most recent step or reset. """ - return self.old_obs.copy() + return deepcopy(self.old_obs) def get_original_reward(self) -> np.ndarray: """ @@ -155,9 +206,10 @@ class VecNormalize(VecEnvWrapper): """ return self.old_reward.copy() - def reset(self) -> np.ndarray: + def reset(self) -> Union[np.ndarray, Dict[str, np.ndarray]]: """ Reset all environments + :return: first observation of the episode """ obs = self.venv.reset() self.old_obs = obs diff --git a/stable_baselines3/her/__init__.py b/stable_baselines3/her/__init__.py new file mode 100644 index 0000000..24f3473 --- /dev/null +++ b/stable_baselines3/her/__init__.py @@ -0,0 +1,4 @@ +from stable_baselines3.common.vec_env.obs_dict_wrapper import ObsDictWrapper +from stable_baselines3.her.goal_selection_strategy import GoalSelectionStrategy +from stable_baselines3.her.her import HER +from stable_baselines3.her.her_replay_buffer import HerReplayBuffer diff --git a/stable_baselines3/her/goal_selection_strategy.py b/stable_baselines3/her/goal_selection_strategy.py new file mode 100644 index 0000000..d4c6a93 --- /dev/null +++ b/stable_baselines3/her/goal_selection_strategy.py @@ -0,0 +1,26 @@ +from enum import Enum + + +class GoalSelectionStrategy(Enum): + """ + The strategies for selecting new goals when + creating artificial transitions. + """ + + # Select a goal that was achieved + # after the current step, in the same episode + FUTURE = 0 + # Select the goal that was achieved + # at the end of the episode + FINAL = 1 + # Select a goal that was achieved in the episode + EPISODE = 2 + + +# For convenience +# that way, we can use string to select a strategy +KEY_TO_GOAL_STRATEGY = { + "future": GoalSelectionStrategy.FUTURE, + "final": GoalSelectionStrategy.FINAL, + "episode": GoalSelectionStrategy.EPISODE, +} diff --git a/stable_baselines3/her/her.py b/stable_baselines3/her/her.py new file mode 100644 index 0000000..658abc6 --- /dev/null +++ b/stable_baselines3/her/her.py @@ -0,0 +1,567 @@ +import io +import pathlib +import warnings +from typing import Any, Iterable, List, Optional, Tuple, Type, Union + +import numpy as np +import torch as th + +from stable_baselines3.common.base_class import BaseAlgorithm +from stable_baselines3.common.callbacks import BaseCallback +from stable_baselines3.common.noise import ActionNoise +from stable_baselines3.common.off_policy_algorithm import OffPolicyAlgorithm +from stable_baselines3.common.policies import BasePolicy +from stable_baselines3.common.save_util import load_from_zip_file, recursive_setattr +from stable_baselines3.common.type_aliases import GymEnv, MaybeCallback, RolloutReturn +from stable_baselines3.common.utils import check_for_correct_spaces +from stable_baselines3.common.vec_env import VecEnv +from stable_baselines3.common.vec_env.obs_dict_wrapper import ObsDictWrapper +from stable_baselines3.her.goal_selection_strategy import KEY_TO_GOAL_STRATEGY, GoalSelectionStrategy +from stable_baselines3.her.her_replay_buffer import HerReplayBuffer + + +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: + 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" + ) + return current_max_episode_length + + +# TODO: rewrite HER class as soon as dict obs are supported +class HER(BaseAlgorithm): + """ + Hindsight Experience Replay (HER) + 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 model constructor + + + For additional offline algorithm specific arguments please have a look at the corresponding documentation. + + :param policy: The policy model to use. + :param env: The environment to learn from (if registered in Gym, can be str) + :param model_class: Off policy model which will be used with hindsight experience replay. (SAC, TD3, DDPG, DQN) + :param n_sampled_goal: Number of sampled goals for replay. (offline sampling) + :param goal_selection_strategy: Strategy for sampling goals for replay. + One of ['episode', 'final', 'future', 'random'] + :param online_sampling: Sample HER transitions online. + :param learning_rate: learning rate for the optimizer, + it can be a function of the current progress remaining (from 1 to 0) + :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. + """ + + def __init__( + self, + policy: Union[str, Type[BasePolicy]], + env: Union[GymEnv, str], + model_class: Type[OffPolicyAlgorithm], + n_sampled_goal: int = 4, + goal_selection_strategy: Union[GoalSelectionStrategy, str] = "future", + online_sampling: bool = False, + max_episode_length: Optional[int] = None, + *args, + **kwargs, + ): + + # we will use the policy and learning rate from the model + super(HER, self).__init__(policy=BasePolicy, env=env, policy_base=BasePolicy, learning_rate=0.0) + del self.policy, self.learning_rate + + if self.get_vec_normalize_env() is not None: + assert online_sampling, "You must pass `online_sampling=True` if you want to use `VecNormalize` with `HER`" + + _init_setup_model = kwargs.get("_init_setup_model", True) + if "_init_setup_model" in kwargs: + del kwargs["_init_setup_model"] + # model initialization + self.model_class = model_class + self.model = model_class( + policy=policy, + env=self.env, + _init_setup_model=False, # pytype: disable=wrong-keyword-args + *args, + **kwargs, # pytype: disable=wrong-keyword-args + ) + + self.verbose = self.model.verbose + self.tensorboard_log = self.model.tensorboard_log + + # convert goal_selection_strategy into GoalSelectionStrategy if string + if isinstance(goal_selection_strategy, str): + self.goal_selection_strategy = KEY_TO_GOAL_STRATEGY[goal_selection_strategy.lower()] + else: + self.goal_selection_strategy = goal_selection_strategy + + # check if goal_selection_strategy is valid + assert isinstance( + self.goal_selection_strategy, GoalSelectionStrategy + ), 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 + self.her_ratio = 1 - (1.0 / (self.n_sampled_goal + 1)) + # maximum steps in episode + self.max_episode_length = get_time_limit(self.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 = self.buffer_size if online_sampling else self.max_episode_length + self._episode_storage = HerReplayBuffer( + self.env, + her_buffer_size, + self.max_episode_length, + self.goal_selection_strategy, + self.env.observation_space, + self.env.action_space, + self.device, + self.n_envs, + self.her_ratio, # pytype: disable=wrong-arg-types + ) + + # counter for steps in episode + self.episode_steps = 0 + + if _init_setup_model: + self._setup_model() + + def _setup_model(self) -> None: + self.model._setup_model() + # assign episode storage to replay buffer when using online HER sampling + if self.online_sampling: + self.model.replay_buffer = self._episode_storage + + def predict( + self, + observation: np.ndarray, + state: Optional[np.ndarray] = None, + mask: Optional[np.ndarray] = None, + deterministic: bool = False, + ) -> Tuple[np.ndarray, Optional[np.ndarray]]: + + return self.model.predict(observation, state, mask, deterministic) + + def learn( + self, + total_timesteps: int, + callback: MaybeCallback = None, + log_interval: int = 4, + eval_env: Optional[GymEnv] = None, + eval_freq: int = -1, + n_eval_episodes: int = 5, + tb_log_name: str = "HER", + eval_log_path: Optional[str] = None, + reset_num_timesteps: bool = True, + ) -> BaseAlgorithm: + + total_timesteps, callback = self._setup_learn( + total_timesteps, eval_env, callback, eval_freq, n_eval_episodes, eval_log_path, reset_num_timesteps, tb_log_name + ) + self.model.start_time = self.start_time + self.model.ep_info_buffer = self.ep_info_buffer + self.model.ep_success_buffer = self.ep_success_buffer + self.model.num_timesteps = self.num_timesteps + self.model._episode_num = self._episode_num + self.model._last_obs = self._last_obs + self.model._total_timesteps = self._total_timesteps + + callback.on_training_start(locals(), globals()) + + while self.num_timesteps < total_timesteps: + + rollout = self.collect_rollouts( + self.env, + n_episodes=self.n_episodes_rollout, + n_steps=self.train_freq, + action_noise=self.action_noise, + callback=callback, + learning_starts=self.learning_starts, + log_interval=log_interval, + ) + + if rollout.continue_training is False: + break + + if self.num_timesteps > 0 and self.num_timesteps > self.learning_starts and self.replay_buffer.size() > 0: + # If no `gradient_steps` is specified, + # do as many gradients steps as steps performed during the rollout + gradient_steps = self.gradient_steps if self.gradient_steps > 0 else rollout.episode_timesteps + self.train(batch_size=self.batch_size, gradient_steps=gradient_steps) + + callback.on_training_end() + + return self + + def collect_rollouts( + self, + env: VecEnv, + callback: BaseCallback, + n_episodes: int = 1, + n_steps: int = -1, + action_noise: Optional[ActionNoise] = None, + learning_starts: int = 0, + log_interval: Optional[int] = None, + ) -> RolloutReturn: + """ + Collect experiences and store them into a ReplayBuffer. + + :param env: The training environment + :param callback: Callback that will be called at each step + (and at the beginning and end of the rollout) + :param n_episodes: Number of episodes to use to collect rollout data + You can also specify a ``n_steps`` instead + :param n_steps: Number of steps to use to collect rollout data + You can also specify a ``n_episodes`` instead. + :param action_noise: Action noise that will be used for exploration + 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 log_interval: Log data every ``log_interval`` episodes + :return: + """ + + episode_rewards, total_timesteps = [], [] + total_steps, total_episodes = 0, 0 + + assert isinstance(env, VecEnv), "You must pass a VecEnv" + assert env.num_envs == 1, "OffPolicyAlgorithm only support single environment" + + if self.model.use_sde: + self.actor.reset_noise() + + callback.on_rollout_start() + continue_training = True + + while total_steps < n_steps or total_episodes < n_episodes: + done = False + episode_reward, episode_timesteps = 0.0, 0 + + while not done: + # concatenate observation and (desired) goal + observation = self._last_obs + self._last_obs = ObsDictWrapper.convert_dict(observation) + + if self.model.use_sde and self.model.sde_sample_freq > 0 and total_steps % self.model.sde_sample_freq == 0: + # Sample a new noise matrix + self.actor.reset_noise() + + # Select action randomly or according to policy + self.model._last_obs = self._last_obs + action, buffer_action = self._sample_action(learning_starts, action_noise) + + # Perform action + new_obs, reward, done, infos = env.step(action) + + self.num_timesteps += 1 + self.model.num_timesteps = self.num_timesteps + episode_timesteps += 1 + total_steps += 1 + + # Only stop training if return value is False, not when it is None. + if callback.on_step() is False: + return RolloutReturn(0.0, total_steps, total_episodes, continue_training=False) + + episode_reward += reward + + # Retrieve reward and episode length if using Monitor wrapper + self._update_info_buffer(infos, done) + self.model.ep_info_buffer = self.ep_info_buffer + self.model.ep_success_buffer = self.ep_success_buffer + + # == Store transition in the replay buffer and/or in the episode storage == + + if self._vec_normalize_env is not None: + # Store only the unnormalized version + new_obs_ = self._vec_normalize_env.get_original_obs() + reward_ = self._vec_normalize_env.get_original_reward() + else: + # Avoid changing the original ones + self._last_original_obs, new_obs_, reward_ = observation, new_obs, reward + self.model._last_original_obs = self._last_original_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: + # The saved terminal_observation is not passed through other + # VecEnvWrapper, so no need to unnormalize + # NOTE: this may be an issue when using other wrappers + next_obs = infos[0]["terminal_observation"] + else: + next_obs = new_obs_ + + if self.online_sampling: + self.replay_buffer.add(self._last_original_obs, next_obs, buffer_action, reward_, done, infos) + else: + # concatenate observation with (desired) goal + flattened_obs = ObsDictWrapper.convert_dict(self._last_original_obs) + flattened_next_obs = ObsDictWrapper.convert_dict(next_obs) + # add to replay buffer + self.replay_buffer.add(flattened_obs, flattened_next_obs, buffer_action, reward_, done) + # add current transition to episode storage + self._episode_storage.add(self._last_original_obs, next_obs, buffer_action, reward_, done, infos) + + self._last_obs = new_obs + self.model._last_obs = self._last_obs + + # Save the unnormalized new observation + if self._vec_normalize_env is not None: + self._last_original_obs = new_obs_ + self.model._last_original_obs = self._last_original_obs + + self.model._update_current_progress_remaining(self.num_timesteps, self._total_timesteps) + + # 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.model._on_step() + + self.episode_steps += 1 + + if 0 < n_steps <= total_steps: + break + + if done or self.episode_steps >= self.max_episode_length: + if self.online_sampling: + self.replay_buffer.store_episode() + else: + self._episode_storage.store_episode() + # sample virtual transitions and store them in replay buffer + self._sample_her_transitions() + # clear storage for current episode + self._episode_storage.reset() + + total_episodes += 1 + self._episode_num += 1 + self.model._episode_num = self._episode_num + 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() + + self.episode_steps = 0 + + mean_reward = np.mean(episode_rewards) if total_episodes > 0 else 0.0 + + callback.on_rollout_end() + + return RolloutReturn(mean_reward, total_steps, total_episodes, continue_training) + + def _sample_her_transitions(self) -> None: + """ + Sample additional goals and store new transitions in replay buffer + when using offline sampling. + """ + + # Sample goals and get new observations + # maybe_vec_env=None as we should store unnormalized transitions, + # they will be normalized at sampling time + observations, next_observations, actions, rewards = self._episode_storage.sample_offline( + n_sampled_goal=self.n_sampled_goal + ) + + # store data in replay buffer + dones = np.zeros((len(observations)), dtype=bool) + self.replay_buffer.extend(observations, next_observations, actions, rewards, dones) + + def __getattr__(self, item: str) -> Any: + """ + Find attribute from model class if this class does not have it. + """ + if hasattr(self.model, item): + return getattr(self.model, item) + else: + raise AttributeError(f"{self} has no attribute {item}") + + def _get_torch_save_params(self) -> Tuple[List[str], List[str]]: + return self.model._get_torch_save_params() + + def save( + self, + path: Union[str, pathlib.Path, io.BufferedIOBase], + exclude: Optional[Iterable[str]] = None, + include: Optional[Iterable[str]] = None, + ) -> None: + """ + Save all the attributes of the object and the model parameters in a zip-file. + + :param path: path to the file where the rl agent should be saved + :param exclude: name of parameters that should be excluded in addition to the default one + :param include: name of parameters that might be excluded but should be included anyway + """ + + # add HER parameters to model + self.model.n_sampled_goal = self.n_sampled_goal + self.model.goal_selection_strategy = self.goal_selection_strategy + self.model.online_sampling = self.online_sampling + self.model.model_class = self.model_class + self.model.max_episode_length = self.max_episode_length + + self.model.save(path, exclude, include) + + @classmethod + def load( + cls, + path: Union[str, pathlib.Path, io.BufferedIOBase], + env: Optional[GymEnv] = None, + device: Union[th.device, str] = "auto", + **kwargs, + ) -> "BaseAlgorithm": + """ + Load the model from a zip-file + + :param path: path to the file (or a file-like) where to + load the agent from + :param env: the new environment to run the loaded model on + (can be None if you only need prediction from a trained model) has priority over any saved environment + :param device: Device on which the code should run. + :param kwargs: extra arguments to change the model when loading + """ + data, params, pytorch_variables = load_from_zip_file(path, device=device) + + # Remove stored device information and replace with ours + if "policy_kwargs" in data: + if "device" in data["policy_kwargs"]: + del data["policy_kwargs"]["device"] + + if "policy_kwargs" in kwargs and kwargs["policy_kwargs"] != data["policy_kwargs"]: + raise ValueError( + f"The specified policy kwargs do not equal the stored policy kwargs." + f"Stored kwargs: {data['policy_kwargs']}, specified kwargs: {kwargs['policy_kwargs']}" + ) + + # check if observation space and action space are part of the saved parameters + if "observation_space" not in data or "action_space" not in data: + raise KeyError("The observation_space and action_space were not given, can't verify new environments") + + # check if given env is valid + if env is not None: + # Wrap first if needed + env = cls._wrap_env(env, data["verbose"]) + # Check if given env is valid + check_for_correct_spaces(env, data["observation_space"], data["action_space"]) + else: + # Use stored env, if one exists. If not, continue as is (can be used for predict) + if "env" in data: + env = data["env"] + + if "use_sde" in data and data["use_sde"]: + kwargs["use_sde"] = True + + # Keys that cannot be changed + for key in {"model_class", "online_sampling", "max_episode_length"}: + if key in kwargs: + del kwargs[key] + + # Keys that can be changed + for key in {"n_sampled_goal", "goal_selection_strategy"}: + if key in kwargs: + data[key] = kwargs[key] # pytype: disable=unsupported-operands + del kwargs[key] + + # noinspection PyArgumentList + her_model = cls( + policy=data["policy_class"], + env=env, + model_class=data["model_class"], + n_sampled_goal=data["n_sampled_goal"], + goal_selection_strategy=data["goal_selection_strategy"], + online_sampling=data["online_sampling"], + max_episode_length=data["max_episode_length"], + policy_kwargs=data["policy_kwargs"], + _init_setup_model=False, # pytype: disable=not-instantiable,wrong-keyword-args + **kwargs, + ) + + # load parameters + her_model.model.__dict__.update(data) + her_model.model.__dict__.update(kwargs) + her_model._setup_model() + + her_model._total_timesteps = her_model.model._total_timesteps + her_model.num_timesteps = her_model.model.num_timesteps + her_model._episode_num = her_model.model._episode_num + + # put state_dicts back in place + her_model.model.set_parameters(params, exact_match=True, device=device) + + # put other pytorch variables back in place + if pytorch_variables is not None: + for name in pytorch_variables: + recursive_setattr(her_model.model, name, pytorch_variables[name]) + + # Sample gSDE exploration matrix, so it uses the right device + # see issue #44 + if her_model.model.use_sde: + her_model.model.policy.reset_noise() # pytype: disable=attribute-error + return her_model + + def load_replay_buffer( + self, path: Union[str, pathlib.Path, io.BufferedIOBase], truncate_last_trajectory: bool = True + ) -> None: + """ + Load a replay buffer from a pickle file and set environment for replay buffer (only online sampling). + + :param path: Path to the pickled replay buffer. + :param truncate_last_trajectory: Only for online sampling. + If set to ``True`` we assume that the last trajectory in the replay buffer was finished. + If it is set to ``False`` we assume that we continue the same trajectory (same episode). + """ + self.model.load_replay_buffer(path=path) + + if self.online_sampling: + # set environment + self.replay_buffer.set_env(self.env) + # If we are at the start of an episode, no need to truncate + current_idx = self.replay_buffer.current_idx + + # truncate interrupted episode + if truncate_last_trajectory and current_idx > 0: + 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.replay_buffer.pos + # set episode length for current episode + self.replay_buffer.episode_lengths[pos] = current_idx + # set done = True for current episode + # current_idx was already incremented + self.replay_buffer.buffer["done"][pos][current_idx - 1] = np.array([True], dtype=np.float32) + # reset current transition index + self.replay_buffer.current_idx = 0 + # increment episode counter + self.replay_buffer.pos = (self.replay_buffer.pos + 1) % self.replay_buffer.max_episode_stored + # update "full" indicator + self.replay_buffer.full = self.replay_buffer.full or self.replay_buffer.pos == 0 diff --git a/stable_baselines3/her/her_replay_buffer.py b/stable_baselines3/her/her_replay_buffer.py new file mode 100644 index 0000000..5bbf1b9 --- /dev/null +++ b/stable_baselines3/her/her_replay_buffer.py @@ -0,0 +1,370 @@ +from collections import deque +from typing import Any, Dict, List, Optional, Tuple, Union + +import numpy as np +import torch as th +from gym import spaces + +from stable_baselines3.common.buffers import ReplayBuffer +from stable_baselines3.common.type_aliases import ReplayBufferSamples, RolloutBufferSamples +from stable_baselines3.common.vec_env import VecNormalize +from stable_baselines3.common.vec_env.obs_dict_wrapper import ObsDictWrapper +from stable_baselines3.her.goal_selection_strategy import GoalSelectionStrategy + + +class HerReplayBuffer(ReplayBuffer): + """ + 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. + + :param env: The training environment + :param buffer_size: The size of the buffer measured in transitions. + :param max_episode_length: The length of an episode. (time horizon) + :param goal_selection_strategy: Strategy for sampling goals for replay. + One of ['episode', 'final', 'future'] + :param observation_space: Observation space + :param action_space: Action space + :param device: PyTorch device + :param n_envs: Number of parallel environments + :her_ratio: The ratio between HER transitions and regular transitions in percent + (between 0 and 1, for online sampling) + The default value ``her_ratio=0.8`` corresponds to 4 virtual transitions + for one real transition (4 / (4 + 1) = 0.8) + """ + + def __init__( + self, + env: ObsDictWrapper, + buffer_size: int, + max_episode_length: int, + goal_selection_strategy: GoalSelectionStrategy, + observation_space: spaces.Space, + action_space: spaces.Space, + device: Union[th.device, str] = "cpu", + n_envs: int = 1, + her_ratio: float = 0.8, + ): + + super(HerReplayBuffer, self).__init__(buffer_size, observation_space, action_space, device, n_envs) + + self.env = env + self.buffer_size = buffer_size + self.max_episode_length = max_episode_length + + # 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 + + # input dimensions for buffer initialization + input_shape = { + "observation": (self.env.num_envs, self.env.obs_dim), + "achieved_goal": (self.env.num_envs, self.env.goal_dim), + "desired_goal": (self.env.num_envs, self.env.goal_dim), + "action": (self.action_dim,), + "reward": (1,), + "next_obs": (self.env.num_envs, self.env.obs_dim), + "next_achieved_goal": (self.env.num_envs, self.env.goal_dim), + "next_desired_goal": (self.env.num_envs, self.env.goal_dim), + "done": (1,), + } + 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) + + self.goal_selection_strategy = goal_selection_strategy + # percentage of her indices + self.her_ratio = her_ratio + + def __getstate__(self) -> Dict[str, Any]: + """ + Gets state for pickling. + + Excludes self.env, as in general Env's may not be pickleable.""" + state = self.__dict__.copy() + # these attributes are not pickleable + del state["env"] + return state + + def __setstate__(self, state: Dict[str, Any]) -> None: + """ + Restores pickled state. + + User must call ``set_env()`` after unpickling before using. + + :param state: + """ + self.__dict__.update(state) + assert "env" not in state + self.env = None + + def set_env(self, env: ObsDictWrapper) -> None: + """ + Sets the environment. + + :param env: + """ + if self.env is not None: + raise ValueError("Trying to set env of already initialized environment.") + + self.env = env + + def _get_samples( + self, batch_inds: np.ndarray, env: Optional[VecNormalize] = None + ) -> Union[ReplayBufferSamples, RolloutBufferSamples]: + """ + Abstract method from base class. + """ + raise NotImplementedError() + + def sample( + self, + batch_size: int, + env: Optional[VecNormalize], + ) -> Union[ReplayBufferSamples, Tuple[np.ndarray, ...]]: + """ + 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. + """ + return self._sample_transitions(batch_size, maybe_vec_env=env, online_sampling=True) + + def sample_offline( + self, + n_sampled_goal: Optional[int] = None, + ) -> Union[ReplayBufferSamples, Tuple[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. + """ + # 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 + ) + + 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] + 1, 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["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[ReplayBufferSamples, Tuple[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" + 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] + + # Special case when using the "future" goal sampling strategy + # we cannot sample all transitions, we have to remove the last timestep + if self.goal_selection_strategy == GoalSelectionStrategy.FUTURE: + # restrict the sampling domain when ep_lengths > 1 + # otherwise filter out the indices + her_indices = her_indices[ep_lengths[her_indices] > 1] + ep_lengths[her_indices] -= 1 + + 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), 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) + ] + ) + + # 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 = ObsDictWrapper.convert_dict(self._normalize_obs(transitions, maybe_vec_env)) + # HACK to make normalize obs work with the next observation + transitions["observation"] = transitions["next_obs"] + next_observations = ObsDictWrapper.convert_dict(self._normalize_obs(transitions, maybe_vec_env)) + + if online_sampling: + data = ( + observations[:, 0], + transitions["action"], + next_observations[:, 0], + transitions["done"], + self._normalize_reward(transitions["reward"], maybe_vec_env), + ) + + return ReplayBufferSamples(*tuple(map(self.to_torch, data))) + else: + return observations, next_observations, transitions["action"], transitions["reward"] + + def add( + self, + obs: Dict[str, np.ndarray], + next_obs: Dict[str, np.ndarray], + action: np.ndarray, + reward: np.ndarray, + done: np.ndarray, + infos: List[dict], + ) -> None: + + if self.current_idx == 0 and self.full: + # Clear info buffer + self.info_buffer[self.pos] = deque(maxlen=self.max_episode_length) + + 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"] + + self.info_buffer[self.pos].append(infos) + + # update current pointer + self.current_idx += 1 + + def store_episode(self) -> None: + """ + Increment episode counter + and reset transition pointer. + """ + # add episode length to length storage + self.episode_lengths[self.pos] = self.current_idx + + # 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 + + @property + def n_episodes_stored(self) -> int: + if self.full: + return self.max_episode_stored + return self.pos + + def size(self) -> int: + """ + :return: The current number of transitions in the buffer. + """ + return int(np.sum(self.episode_lengths)) + + def reset(self) -> None: + """ + Reset the buffer. + """ + self.pos = 0 + self.current_idx = 0 + self.full = False + self.episode_lengths = np.zeros(self.max_episode_stored, dtype=np.int64) diff --git a/stable_baselines3/version.txt b/stable_baselines3/version.txt index 37f1777..8dabd1f 100644 --- a/stable_baselines3/version.txt +++ b/stable_baselines3/version.txt @@ -1 +1 @@ -0.10.0a0 +0.10.0a1 diff --git a/tests/test_her.py b/tests/test_her.py new file mode 100644 index 0000000..09d1a78 --- /dev/null +++ b/tests/test_her.py @@ -0,0 +1,318 @@ +import os +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, HER, SAC, TD3 +from stable_baselines3.common.bit_flipping_env import BitFlippingEnv +from stable_baselines3.common.vec_env import DummyVecEnv +from stable_baselines3.common.vec_env.obs_dict_wrapper import ObsDictWrapper +from stable_baselines3.her.goal_selection_strategy import GoalSelectionStrategy +from stable_baselines3.her.her import get_time_limit + + +@pytest.mark.parametrize("model_class", [SAC, TD3, DDPG, DQN]) +@pytest.mark.parametrize("online_sampling", [True, False]) +def test_her(model_class, online_sampling): + """ + Test Hindsight Experience Replay. + """ + n_bits = 4 + env = BitFlippingEnv(n_bits=n_bits, continuous=not (model_class == DQN)) + + model = HER( + "MlpPolicy", + env, + model_class, + goal_selection_strategy="future", + online_sampling=online_sampling, + gradient_steps=1, + train_freq=1, + n_episodes_rollout=-1, + max_episode_length=n_bits, + policy_kwargs=dict(net_arch=[64]), + learning_starts=100, + ) + + model.learn(total_timesteps=300) + + +@pytest.mark.parametrize( + "goal_selection_strategy", + [ + "final", + "episode", + "future", + GoalSelectionStrategy.FINAL, + GoalSelectionStrategy.EPISODE, + GoalSelectionStrategy.FUTURE, + ], +) +@pytest.mark.parametrize("online_sampling", [True, False]) +def test_goal_selection_strategy(goal_selection_strategy, online_sampling): + """ + Test different goal strategies. + """ + env = BitFlippingEnv(continuous=True) + + model = HER( + "MlpPolicy", + env, + SAC, + goal_selection_strategy=goal_selection_strategy, + online_sampling=online_sampling, + gradient_steps=1, + train_freq=1, + n_episodes_rollout=-1, + max_episode_length=10, + policy_kwargs=dict(net_arch=[64]), + learning_starts=100, + ) + model.learn(total_timesteps=300) + + +@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): + """ + 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_bits = 4 + env = BitFlippingEnv(n_bits=n_bits, continuous=not (model_class == DQN)) + + kwargs = dict(use_sde=True) if use_sde else {} + + # create model + model = HER( + "MlpPolicy", + env, + model_class, + n_sampled_goal=5, + goal_selection_strategy="future", + online_sampling=online_sampling, + verbose=0, + tau=0.05, + batch_size=128, + learning_rate=0.001, + policy_kwargs=dict(net_arch=[64]), + buffer_size=int(1e6), + gamma=0.98, + gradient_steps=1, + train_freq=4, + learning_starts=100, + n_episodes_rollout=-1, + max_episode_length=n_bits, + **kwargs + ) + + model.learn(total_timesteps=300) + + env.reset() + + observations_list = [] + for _ in range(10): + obs = env.step(env.action_space.sample())[0] + observation = ObsDictWrapper.convert_dict(obs) + observations_list.append(observation) + observations = np.array(observations_list) + + # Get dictionary of current parameters + params = deepcopy(model.policy.state_dict()) + + # Modify all parameters to be random values + random_params = dict((param_name, th.rand_like(param)) for param_name, param in params.items()) + + # Update model parameters with the new random values + model.policy.load_state_dict(random_params) + + new_params = model.policy.state_dict() + # Check that all params are different now + for k in params: + assert not th.allclose(params[k], new_params[k]), "Parameters did not change as expected." + + params = new_params + + # get selected actions + selected_actions, _ = model.predict(observations, deterministic=True) + + # Check + model.save(tmp_path / "test_save.zip") + del model + model = HER.load(str(tmp_path / "test_save.zip"), env=env) + + # check if params are still the same after load + new_params = model.policy.state_dict() + + # Check that all params are the same as before save load procedure now + for key in params: + assert th.allclose(params[key], new_params[key]), "Model parameters not the same after save and load." + + # check if model still selects the same actions + new_selected_actions, _ = model.predict(observations, deterministic=True) + assert np.allclose(selected_actions, new_selected_actions, 1e-4) + + # check if learn still works + model.learn(total_timesteps=300) + + # Test that the change of parameters works + model = HER.load(str(tmp_path / "test_save.zip"), env=env, verbose=3, learning_rate=2.0) + assert model.model.learning_rate == 2.0 + assert model.verbose == 3 + + # clear file from os + os.remove(tmp_path / "test_save.zip") + + +@pytest.mark.parametrize("online_sampling, truncate_last_trajectory", [(False, None), (True, True), (True, False)]) +def test_save_load_replay_buffer(tmp_path, recwarn, online_sampling, truncate_last_trajectory): + """ + Test if 'save_replay_buffer' and 'load_replay_buffer' works correctly + """ + # remove gym warnings + warnings.filterwarnings(action="ignore", category=DeprecationWarning) + warnings.filterwarnings(action="ignore", category=UserWarning, module="gym") + + path = pathlib.Path(tmp_path / "logs/replay_buffer.pkl") + path.parent.mkdir(exist_ok=True, parents=True) # to not raise a warning + env = BitFlippingEnv(n_bits=4, continuous=True) + model = HER( + "MlpPolicy", + env, + SAC, + goal_selection_strategy="future", + online_sampling=online_sampling, + gradient_steps=1, + train_freq=1, + n_episodes_rollout=-1, + max_episode_length=4, + buffer_size=int(2e4), + policy_kwargs=dict(net_arch=[64]), + seed=0, + ) + model.learn(200) + old_replay_buffer = deepcopy(model.replay_buffer) + model.save_replay_buffer(path) + del model.model.replay_buffer + + with pytest.raises(AttributeError): + model.replay_buffer + + # Check that there is no warning + assert len(recwarn) == 0 + + model.load_replay_buffer(path, truncate_last_trajectory) + + if truncate_last_trajectory: + assert len(recwarn) == 1 + warning = recwarn.pop(UserWarning) + assert "The last trajectory in the replay buffer will be truncated" in str(warning.message) + 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: + assert np.allclose(old_replay_buffer.observations, model.replay_buffer.observations) + assert np.allclose(old_replay_buffer.actions, model.replay_buffer.actions) + assert np.allclose(old_replay_buffer.rewards, model.replay_buffer.rewards) + assert np.allclose(old_replay_buffer.dones, model.replay_buffer.dones) + + # test if continuing training works properly + reset_num_timesteps = False if truncate_last_trajectory is False else True + model.learn(200, reset_num_timesteps=reset_num_timesteps) + + +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 + HER("MlpPolicy", dict_env, DQN, max_episode_length=5) + + with pytest.raises(ValueError): + HER("MlpPolicy", dict_env, DQN) + + # Wrapped in a timelimit, should be fine + # Note: it requires env.spec to be defined + env = DummyVecEnv([lambda: gym.wrappers.TimeLimit(BitFlippingEnv(), 10)]) + HER("MlpPolicy", env, DQN) + + +@pytest.mark.parametrize("online_sampling", [False, True]) +@pytest.mark.parametrize("n_bits", [10]) +def test_performance_her(online_sampling, 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) + + model = HER( + "MlpPolicy", + env, + DQN, + n_sampled_goal=5, + goal_selection_strategy="future", + online_sampling=online_sampling, + verbose=1, + learning_rate=5e-4, + max_episode_length=n_bits, + train_freq=1, + learning_starts=100, + exploration_final_eps=0.02, + target_update_interval=500, + seed=0, + batch_size=32, + ) + + model.learn(total_timesteps=5000, log_interval=50) + + # 90% training success + assert np.mean(model.ep_success_buffer) > 0.90 diff --git a/tests/test_save_load.py b/tests/test_save_load.py index 18fba37..e6230eb 100644 --- a/tests/test_save_load.py +++ b/tests/test_save_load.py @@ -231,8 +231,10 @@ def test_exclude_include_saved_params(tmp_path, model_class): def test_save_load_replay_buffer(tmp_path, model_class): path = pathlib.Path(tmp_path / "logs/replay_buffer.pkl") path.parent.mkdir(exist_ok=True, parents=True) # to not raise a warning - model = model_class("MlpPolicy", select_env(model_class), buffer_size=1000) - model.learn(500) + model = model_class( + "MlpPolicy", select_env(model_class), buffer_size=1000, policy_kwargs=dict(net_arch=[64]), learning_starts=200 + ) + model.learn(300) old_replay_buffer = deepcopy(model.replay_buffer) model.save_replay_buffer(path) model.replay_buffer = None @@ -305,7 +307,7 @@ def test_save_load_policy(tmp_path, model_class, policy_str): if policy_str == "MlpPolicy": env = select_env(model_class) else: - if model_class in [SAC, TD3, DQN]: + if model_class in [SAC, TD3, DQN, DDPG]: # Avoid memory error when using replay buffer # Reduce the size of the features kwargs = dict(buffer_size=250) diff --git a/tests/test_sde.py b/tests/test_sde.py index 3b010d3..74853a0 100644 --- a/tests/test_sde.py +++ b/tests/test_sde.py @@ -54,6 +54,11 @@ def test_state_dependent_exploration_grad(): assert sigma_hat.grad.allclose(grad) +def test_sde_check(): + with pytest.raises(ValueError): + PPO("MlpPolicy", "CartPole-v1", use_sde=True) + + @pytest.mark.parametrize("model_class", [SAC, A2C, PPO]) @pytest.mark.parametrize("sde_net_arch", [None, [32, 16], []]) @pytest.mark.parametrize("use_expln", [False, True]) @@ -65,9 +70,9 @@ def test_state_dependent_offpolicy_noise(model_class, sde_net_arch, use_expln): seed=None, create_eval_env=True, verbose=1, - policy_kwargs=dict(log_std_init=-2, sde_net_arch=sde_net_arch, use_expln=use_expln), + policy_kwargs=dict(log_std_init=-2, sde_net_arch=sde_net_arch, use_expln=use_expln, net_arch=[64]), ) - model.learn(total_timesteps=int(500), eval_freq=250) + model.learn(total_timesteps=int(300), eval_freq=250) model.policy.reset_noise() if model_class == SAC: model.policy.actor.get_std() diff --git a/tests/test_vec_normalize.py b/tests/test_vec_normalize.py index 311e3c9..a68e1b2 100644 --- a/tests/test_vec_normalize.py +++ b/tests/test_vec_normalize.py @@ -1,8 +1,9 @@ import gym import numpy as np import pytest +from gym import spaces -from stable_baselines3 import SAC, TD3 +from stable_baselines3 import HER, SAC, TD3 from stable_baselines3.common.running_mean_std import RunningMeanStd from stable_baselines3.common.vec_env import ( DummyVecEnv, @@ -15,14 +16,68 @@ from stable_baselines3.common.vec_env import ( ENV_ID = "Pendulum-v0" +class DummyDictEnv(gym.GoalEnv): + """ + Dummy gym goal env for testing purposes + """ + + def __init__(self): + super(DummyDictEnv, self).__init__() + self.observation_space = spaces.Dict( + { + "observation": spaces.Box(low=-20.0, high=20.0, shape=(4,), dtype=np.float32), + "achieved_goal": spaces.Box(low=-20.0, high=20.0, shape=(4,), dtype=np.float32), + "desired_goal": spaces.Box(low=-20.0, high=20.0, shape=(4,), dtype=np.float32), + } + ) + self.action_space = spaces.Box(low=-1, high=1, shape=(3,), dtype=np.float32) + + def reset(self): + return self.observation_space.sample() + + def step(self, action): + obs = self.observation_space.sample() + reward = self.compute_reward(obs["achieved_goal"], obs["desired_goal"], {}) + done = np.random.rand() > 0.8 + return obs, reward, done, {} + + def compute_reward(self, achieved_goal: np.ndarray, desired_goal: np.ndarray, _info) -> np.float32: + distance = np.linalg.norm(achieved_goal - desired_goal, axis=-1) + return -(distance > 0).astype(np.float32) + + +def allclose(obs_1, obs_2): + """ + Generalized np.allclose() to work with dict spaces. + """ + if isinstance(obs_1, dict): + all_close = True + for key in obs_1.keys(): + if not np.allclose(obs_1[key], obs_2[key]): + all_close = False + break + return all_close + return np.allclose(obs_1, obs_2) + + def make_env(): return gym.make(ENV_ID) +def make_dict_env(): + return DummyDictEnv() + + def check_rms_equal(rmsa, rmsb): - assert np.all(rmsa.mean == rmsb.mean) - assert np.all(rmsa.var == rmsb.var) - assert np.all(rmsa.count == rmsb.count) + if isinstance(rmsa, dict): + for key in rmsa.keys(): + assert np.all(rmsa[key].mean == rmsb[key].mean) + assert np.all(rmsa[key].var == rmsb[key].var) + assert np.all(rmsa[key].count == rmsb[key].count) + else: + assert np.all(rmsa.mean == rmsb.mean) + assert np.all(rmsa.var == rmsb.var) + assert np.all(rmsa.count == rmsb.count) def check_vec_norm_equal(norma, normb): @@ -56,6 +111,19 @@ def _make_warmstart_cartpole(): return venv +def _make_warmstart_dict_env(): + """Warm-start VecNormalize by stepping through BitFlippingEnv""" + venv = DummyVecEnv([make_dict_env]) + venv = VecNormalize(venv) + venv.reset() + venv.get_original_obs() + + for _ in range(100): + actions = [venv.action_space.sample()] + venv.step(actions) + return venv + + def test_runningmeanstd(): """Test RunningMeanStd object""" for (x_1, x_2, x_3) in [ @@ -74,7 +142,8 @@ def test_runningmeanstd(): assert np.allclose(moments_1, moments_2) -def test_vec_env(tmp_path): +@pytest.mark.parametrize("make_env", [make_env, make_dict_env]) +def test_vec_env(tmp_path, make_env): """Test VecNormalize Object""" clip_obs = 0.5 clip_reward = 5.0 @@ -85,7 +154,11 @@ def test_vec_env(tmp_path): while not done[0]: actions = [norm_venv.action_space.sample()] obs, rew, done, _ = norm_venv.step(actions) - assert np.max(np.abs(obs)) <= clip_obs + if isinstance(obs, dict): + for key in obs.keys(): + assert np.max(np.abs(obs[key])) <= clip_obs + else: + assert np.max(np.abs(obs)) <= clip_obs assert np.max(np.abs(rew)) <= clip_reward path = tmp_path / "vec_normalize" @@ -113,6 +186,26 @@ def test_get_original(): np.testing.assert_allclose(venv.normalize_reward(orig_rewards), rewards) +def test_get_original_dict(): + venv = _make_warmstart_dict_env() + for _ in range(3): + actions = [venv.action_space.sample()] + obs, rewards, _, _ = venv.step(actions) + # obs = obs[0] + orig_obs = venv.get_original_obs() + rewards = rewards[0] + orig_rewards = venv.get_original_reward()[0] + + for key in orig_obs.keys(): + assert orig_obs[key].shape == obs[key].shape + assert orig_rewards.dtype == rewards.dtype + + assert not allclose(orig_obs, obs) + assert not np.array_equal(orig_rewards, rewards) + assert allclose(venv.normalize_obs(orig_obs), obs) + np.testing.assert_allclose(venv.normalize_reward(orig_rewards), rewards) + + def test_normalize_external(): venv = _make_warmstart_cartpole() @@ -123,21 +216,24 @@ def test_normalize_external(): assert np.all(norm_rewards < 1) -@pytest.mark.parametrize("model_class", [SAC, TD3]) +@pytest.mark.parametrize("model_class", [SAC, TD3, HER]) def test_offpolicy_normalization(model_class): - env = DummyVecEnv([make_env]) + make_env_ = make_dict_env if model_class == HER else make_env + 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 = 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, policy_kwargs=dict(net_arch=[64])) - model.learn(total_timesteps=1000, eval_env=eval_env, eval_freq=500) + kwargs = dict(model_class=SAC, max_episode_length=200, online_sampling=True) if model_class == HER else {} + model = model_class("MlpPolicy", env, verbose=1, learning_starts=100, policy_kwargs=dict(net_arch=[64]), **kwargs) + model.learn(total_timesteps=500, eval_env=eval_env, eval_freq=250) # Check getter assert isinstance(model.get_vec_normalize_env(), VecNormalize) -def test_sync_vec_normalize(): +@pytest.mark.parametrize("make_env", [make_env, make_dict_env]) +def test_sync_vec_normalize(make_env): env = DummyVecEnv([make_env]) assert unwrap_vec_normalize(env) is None @@ -146,13 +242,15 @@ def test_sync_vec_normalize(): assert isinstance(unwrap_vec_normalize(env), VecNormalize) - env = VecFrameStack(env, 1) - - assert isinstance(unwrap_vec_normalize(env), VecNormalize) + if not isinstance(env.observation_space, spaces.Dict): + env = VecFrameStack(env, 1) + assert isinstance(unwrap_vec_normalize(env), VecNormalize) eval_env = DummyVecEnv([make_env]) eval_env = VecNormalize(eval_env, training=False, norm_obs=True, norm_reward=True, clip_obs=100.0, clip_reward=100.0) - eval_env = VecFrameStack(eval_env, 1) + + if not isinstance(env.observation_space, spaces.Dict): + eval_env = VecFrameStack(eval_env, 1) env.seed(0) env.action_space.seed(0) @@ -171,12 +269,12 @@ def test_sync_vec_normalize(): dummy_rewards = np.random.rand(10) original_obs = env.get_original_obs() # Check that unnormalization works - assert np.allclose(original_obs, env.unnormalize_obs(obs)) + assert allclose(original_obs, env.unnormalize_obs(obs)) # Normalization must be different (between different environments) - assert not np.allclose(obs, eval_env.normalize_obs(original_obs)) + assert not allclose(obs, eval_env.normalize_obs(original_obs)) # Test syncing of parameters sync_envs_normalization(env, eval_env) # Now they must be synced - assert np.allclose(obs, eval_env.normalize_obs(original_obs)) - assert np.allclose(env.normalize_reward(dummy_rewards), eval_env.normalize_reward(dummy_rewards)) + assert allclose(obs, eval_env.normalize_obs(original_obs)) + assert allclose(env.normalize_reward(dummy_rewards), eval_env.normalize_reward(dummy_rewards))