saymrwulf/stable-baselines3
1
0
Fork 0
mirror of https://github.com/saymrwulf/stable-baselines3.git synced 2026-09-15 22:10:25 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Implement HER (#120)

Browse source 
* 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 <antonin.raffin@dlr.de>
Co-authored-by: Antonin RAFFIN <antonin.raffin@ensta.org>
This commit is contained in:
Megan Klaiber 2020-10-22 11:56:43 +02:00 committed by GitHub
parent 15e94a6d14
commit dd6e361204
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
34 changed files with 1899 additions and 102 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
Makefile
View file

@ -5,7 +5,7 @@ pytest:
./scripts/run_tests.sh ./scripts/run_tests.sh
type: type:
pytype pytype -j auto
lint: lint:
# stop the build if there are Python syntax errors or undefined names # stop the build if there are Python syntax errors or undefined names

13
README.md
View file

@ -35,20 +35,9 @@ These algorithms will make it easier for the research community and industry to
| Type hints | :heavy_check_mark: | | Type hints | :heavy_check_mark: |
### Roadmap to V1.0
Please look at the issue for more details.
Planned features:
- [ ] HER
### Planned features (v1.1+) ### Planned features (v1.1+)
- [ ] DQN extensions (prioritized replay, double q-learning, ...) 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).
- [ ] Support for `Tuple` and `Dict` observation spaces
- [ ] Recurrent Policies
- [ ] TRPO
## Migration guide: from Stable-Baselines (SB2) to Stable-Baselines3 (SB3) ## Migration guide: from Stable-Baselines (SB2) to Stable-Baselines3 (SB3)

78
docs/guide/examples.rst
View file

@ -18,8 +18,7 @@ notebooks:
- `Atari Games`_ - `Atari Games`_
- `RL Baselines zoo`_ - `RL Baselines zoo`_
- `PyBullet`_ - `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 .. _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 .. _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 env.norm_reward = False
Hindsight Experience Replay (HER)
---------------------------------
For this example, we are using `Highway-Env <https://github.com/eleurent/highway-env>`_ by `@eleurent <https://github.com/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 Record a Video
-------------- --------------

8
docs/guide/migration.rst
View file

@ -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`` 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 <her>` documentation).
New logger API New logger API
-------------- --------------

1
docs/index.rst
View file

@ -57,6 +57,7 @@ Main Features
modules/a2c modules/a2c
modules/ddpg modules/ddpg
modules/dqn modules/dqn
modules/her
modules/ppo modules/ppo
modules/sac modules/sac
modules/td3 modules/td3

10
docs/misc/changelog.rst
View file

@ -4,7 +4,7 @@ Changelog
========== ==========
Pre-Release 0.10.0a0 (WIP) Pre-Release 0.10.0a1 (WIP)
------------------------------ ------------------------------
Breaking Changes: Breaking Changes:
@ -14,11 +14,14 @@ Breaking Changes:
New Features: 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) - 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) - Support logging videos to Tensorboard (@SwamyDev)
Bug Fixes: Bug Fixes:
^^^^^^^^^^ ^^^^^^^^^^
- Fix GAE computation for on-policy algorithms (off-by one for the last value) (thanks @Wovchena) - 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) - 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) - 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 - Fix model creation initializing CUDA even when `device="cpu"` is provided
@ -37,6 +40,7 @@ Others:
Documentation: Documentation:
^^^^^^^^^^^^^^ ^^^^^^^^^^^^^^
- Added first draft of migration guide - Added first draft of migration guide
- Enabled doc for ``CnnPolicies``
Pre-Release 0.9.0 (2020-10-03) Pre-Release 0.9.0 (2020-10-03)
@ -68,6 +72,7 @@ New Features:
Bug Fixes: 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`` - Fixed a bug where the environment was reset twice when using ``evaluate_policy``
- Fix logging of ``clip_fraction`` in PPO (@diditforlulz273) - 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) - 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) - Fixed typo in custom policy doc (@RaphaelWag)
Pre-Release 0.7.0 (2020-06-10) 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 @MarvineGothic @jdossgollin @SyllogismRXS @rusu24edward @jbulow @Antymon @seheevic @justinkterry @edbeeching
@flodorner @KuKuXia @NeoExtended @PartiallyTyped @mmcenta @richardwu @kinalmehta @rolandgvc @tkelestemur @mloo3 @flodorner @KuKuXia @NeoExtended @PartiallyTyped @mmcenta @richardwu @kinalmehta @rolandgvc @tkelestemur @mloo3
@tirafesi @blurLake @koulakis @joeljosephjin @shwang @rk37 @andyshih12 @RaphaelWag @xicocaio @tirafesi @blurLake @koulakis @joeljosephjin @shwang @rk37 @andyshih12 @RaphaelWag @xicocaio
@diditforlulz273 @liorcohen5 @ManifoldFR @mloo3 @SwamyDev @wmmc88 @diditforlulz273 @liorcohen5 @ManifoldFR @mloo3 @SwamyDev @wmmc88 @megan-klaiber

19
docs/modules/a2c.rst
View file

@ -79,3 +79,22 @@ Parameters
.. autoclass:: A2C .. autoclass:: A2C
:members: :members:
:inherited-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:

8
docs/modules/ddpg.rst
View file

@ -98,7 +98,9 @@ DDPG Policies
:members: :members:
:inherited-members: :inherited-members:
.. autoclass:: stable_baselines3.td3.policies.TD3Policy
:members:
:noindex:
.. .. autoclass:: CnnPolicy .. autoclass:: CnnPolicy
.. :members: :members:
.. :inherited-members:

4
docs/modules/dqn.rst
View file

@ -90,5 +90,9 @@ DQN Policies
:members: :members:
:inherited-members: :inherited-members:
.. autoclass:: stable_baselines3.dqn.policies.DQNPolicy
:members:
:noindex:
.. autoclass:: CnnPolicy .. autoclass:: CnnPolicy
:members: :members:

121
docs/modules/her.rst Normal file
View file

@ -0,0 +1,121 @@
.. _her:
.. automodule:: stable_baselines3.her
HER
====
`Hindsight Experience Replay (HER) <https://arxiv.org/abs/1707.01495>`_
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 <https://github.com/openai/gym/blob/3394e245727c1ae6851b504a50ba77c73cd4c65b/gym/core.py#L160>`_
.. 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:

19
docs/modules/ppo.rst
View file

@ -80,3 +80,22 @@ Parameters
.. autoclass:: PPO .. autoclass:: PPO
:members: :members:
:inherited-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:

11
docs/modules/sac.rst
View file

@ -82,7 +82,7 @@ Example
obs = env.reset() obs = env.reset()
while True: while True:
action, _states = model.predict(obs) action, _states = model.predict(obs, deterministic=True)
obs, reward, done, info = env.step(action) obs, reward, done, info = env.step(action)
env.render() env.render()
if done: if done:
@ -104,6 +104,9 @@ SAC Policies
:members: :members:
:inherited-members: :inherited-members:
.. .. autoclass:: CnnPolicy .. autoclass:: stable_baselines3.sac.policies.SACPolicy
.. :members: :members:
.. :inherited-members: :noindex:
.. autoclass:: CnnPolicy
:members:

8
docs/modules/td3.rst
View file

@ -101,7 +101,9 @@ TD3 Policies
:members: :members:
:inherited-members: :inherited-members:
.. autoclass:: stable_baselines3.td3.policies.TD3Policy
:members:
:noindex:
.. .. autoclass:: CnnPolicy .. autoclass:: CnnPolicy
.. :members: :members:
.. :inherited-members:

1
setup.cfg
View file

@ -29,6 +29,7 @@ per-file-ignores =
./stable_baselines3/a2c/__init__.py:F401 ./stable_baselines3/a2c/__init__.py:F401
./stable_baselines3/ddpg/__init__.py:F401 ./stable_baselines3/ddpg/__init__.py:F401
./stable_baselines3/dqn/__init__.py:F401 ./stable_baselines3/dqn/__init__.py:F401
./stable_baselines3/her/__init__.py:F401
./stable_baselines3/ppo/__init__.py:F401 ./stable_baselines3/ppo/__init__.py:F401
./stable_baselines3/sac/__init__.py:F401 ./stable_baselines3/sac/__init__.py:F401
./stable_baselines3/td3/__init__.py:F401 ./stable_baselines3/td3/__init__.py:F401

1
stable_baselines3/__init__.py
View file

@ -3,6 +3,7 @@ import os
from stable_baselines3.a2c import A2C from stable_baselines3.a2c import A2C
from stable_baselines3.ddpg import DDPG from stable_baselines3.ddpg import DDPG
from stable_baselines3.dqn import DQN from stable_baselines3.dqn import DQN
from stable_baselines3.her import HER
from stable_baselines3.ppo import PPO from stable_baselines3.ppo import PPO
from stable_baselines3.sac import SAC from stable_baselines3.sac import SAC
from stable_baselines3.td3 import TD3 from stable_baselines3.td3 import TD3

2
stable_baselines3/common/atari_wrappers.py
View file

@ -34,7 +34,7 @@ class NoopResetEnv(gym.Wrapper):
else: else:
noops = self.unwrapped.np_random.randint(1, self.noop_max + 1) noops = self.unwrapped.np_random.randint(1, self.noop_max + 1)
assert noops > 0 assert noops > 0
obs = np.empty(0) obs = np.zeros(0)
for _ in range(noops): for _ in range(noops):
obs, _, done, _ = self.env.step(self.noop_action) obs, _, done, _ = self.env.step(self.noop_action)
if done: if done:

39
stable_baselines3/common/base_class.py
View file

@ -26,7 +26,15 @@ from stable_baselines3.common.utils import (
set_random_seed, set_random_seed,
update_learning_rate, 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]: 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) self.eval_env = maybe_make_env(env, monitor_wrapper, self.verbose)
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.observation_space = env.observation_space
self.action_space = env.action_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." "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.") 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 not isinstance(env, VecEnv):
if self.verbose >= 1: if verbose >= 1:
print("Wrapping the env in a DummyVecEnv.") print("Wrapping the env in a DummyVecEnv.")
env = DummyVecEnv([lambda: env]) env = DummyVecEnv([lambda: env])
if is_image_space(env.observation_space) and not isinstance(env, VecTransposeImage): if is_image_space(env.observation_space) and not is_wrapped(env, VecTransposeImage):
if self.verbose >= 1: if verbose >= 1:
print("Wrapping the env in a VecTransposeImage.") print("Wrapping the env in a VecTransposeImage.")
env = VecTransposeImage(env) 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 return env
@abstractmethod @abstractmethod
@ -188,7 +202,7 @@ class BaseAlgorithm(ABC):
eval_env = self.eval_env eval_env = self.eval_env
if eval_env is not None: 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 assert eval_env.num_envs == 1
return eval_env return eval_env
@ -402,10 +416,11 @@ class BaseAlgorithm(ABC):
:param env: The environment for learning a policy :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 # 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.n_envs = env.num_envs
self.env = env 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") raise KeyError("The observation_space and action_space were not given, can't verify new environments")
if env is not None: if env is not None:
# Wrap first if needed
cls._wrap_env(env, data["verbose"])
# Check if given env is valid # Check if given env is valid
check_for_correct_spaces(env, data["observation_space"], data["action_space"]) check_for_correct_spaces(env, data["observation_space"], data["action_space"])
else: else:

21
stable_baselines3/common/bit_flipping_env.py
View file

@ -1,8 +1,9 @@
from collections import OrderedDict from collections import OrderedDict
from typing import Dict, Optional, Union from typing import Any, Dict, Optional, Union
import numpy as np import numpy as np
from gym import GoalEnv, spaces from gym import GoalEnv, spaces
from gym.envs.registration import EnvSpec
from stable_baselines3.common.type_aliases import GymStepReturn 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 version or not, by default, it uses the MultiBinary one
""" """
spec = EnvSpec("BitFlippingEnv-v0")
def __init__( def __init__(
self, n_bits: int = 10, continuous: bool = False, max_steps: Optional[int] = None, discrete_obs_space: bool = False 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 max_steps = n_bits
self.max_steps = max_steps self.max_steps = max_steps
self.current_step = 0 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]: def convert_if_needed(self, state: np.ndarray) -> Union[int, np.ndarray]:
""" """
@ -101,7 +106,7 @@ class BitFlippingEnv(GoalEnv):
else: else:
self.state[action] = 1 - self.state[action] self.state[action] = 1 - self.state[action]
obs = self._get_obs() 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 done = reward == 0
self.current_step += 1 self.current_step += 1
# Episode terminate when we reached the goal or the max number of steps # 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 done = done or self.current_step >= self.max_steps
return obs, reward, done, info 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 # Deceptive reward: it is positive only when the goal is achieved
if self.discrete_obs_space: # vectorized version
return 0.0 if achieved_goal == desired_goal else -1.0 distance = np.linalg.norm(achieved_goal - desired_goal, axis=-1)
return 0.0 if (achieved_goal == desired_goal).all() else -1.0 return -(distance > 0).astype(np.float32)
def render(self, mode: str = "human") -> Optional[np.ndarray]: def render(self, mode: str = "human") -> Optional[np.ndarray]:
if mode == "rgb_array": if mode == "rgb_array":

8
stable_baselines3/common/buffers.py
View file

@ -1,6 +1,6 @@
import warnings import warnings
from abc import ABC, abstractmethod from abc import ABC, abstractmethod
from typing import Generator, Optional, Union from typing import Dict, Generator, Optional, Union
import numpy as np import numpy as np
import torch as th import torch as th
@ -129,9 +129,11 @@ class BaseBuffer(ABC):
return th.as_tensor(array).to(self.device) return th.as_tensor(array).to(self.device)
@staticmethod @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: if env is not None:
return env.normalize_obs(obs).astype(np.float32) return env.normalize_obs(obs)
return obs return obs
@staticmethod @staticmethod

7
stable_baselines3/common/off_policy_algorithm.py
View file

@ -67,6 +67,8 @@ class OffPolicyAlgorithm(BaseAlgorithm):
:param use_sde_at_warmup: Whether to use gSDE instead of uniform sampling :param use_sde_at_warmup: Whether to use gSDE instead of uniform sampling
during the warm up phase (before learning starts) during the warm up phase (before learning starts)
:param sde_support: Whether the model support gSDE or not :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__( def __init__(
@ -97,6 +99,7 @@ class OffPolicyAlgorithm(BaseAlgorithm):
sde_sample_freq: int = -1, sde_sample_freq: int = -1,
use_sde_at_warmup: bool = False, use_sde_at_warmup: bool = False,
sde_support: bool = True, sde_support: bool = True,
remove_time_limit_termination: bool = False,
): ):
super(OffPolicyAlgorithm, self).__init__( super(OffPolicyAlgorithm, self).__init__(
@ -126,6 +129,10 @@ class OffPolicyAlgorithm(BaseAlgorithm):
self.action_noise = action_noise self.action_noise = action_noise
self.optimize_memory_usage = optimize_memory_usage 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: if train_freq > 0 and n_episodes_rollout > 0:
warnings.warn( warnings.warn(
"You passed a positive value for `train_freq` and `n_episodes_rollout`." "You passed a positive value for `train_freq` and `n_episodes_rollout`."

6
stable_baselines3/common/policies.py
View file

@ -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.torch_layers import BaseFeaturesExtractor, FlattenExtractor, MlpExtractor, NatureCNN, create_mlp
from stable_baselines3.common.utils import get_device, is_vectorized_observation 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 import VecTransposeImage
from stable_baselines3.common.vec_env.obs_dict_wrapper import ObsDictWrapper
class BaseModel(nn.Module, ABC): class BaseModel(nn.Module, ABC):
@ -234,7 +235,10 @@ class BasePolicy(BaseModel):
# state = self.initial_state # state = self.initial_state
# if mask is None: # if mask is None:
# mask = [False for _ in range(self.n_envs)] # 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 # Handle the different cases for images
# as PyTorch use channel first format # as PyTorch use channel first format

11
stable_baselines3/common/utils.py
View file

@ -16,9 +16,7 @@ except ImportError:
SummaryWriter = None SummaryWriter = None
from stable_baselines3.common import logger 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.type_aliases import GymEnv
from stable_baselines3.common.vec_env import VecTransposeImage
def set_random_seed(seed: int, using_cuda: bool = False) -> None: 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 observation_space: Observation space to check against
:param action_space: Action space to check against :param action_space: Action space to check against
""" """
if ( if observation_space != env.observation_space:
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)
)
):
raise ValueError(f"Observation spaces do not match: {observation_space} != {env.observation_space}") raise ValueError(f"Observation spaces do not match: {observation_space} != {env.observation_space}")
if action_space != env.action_space: if action_space != env.action_space:
raise ValueError(f"Action spaces do not match: {action_space} != {env.action_space}") raise ValueError(f"Action spaces do not match: {action_space} != {env.action_space}")

11
stable_baselines3/common/vec_env/__init__.py
View file

@ -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 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 # Define here to avoid circular import
def sync_envs_normalization(env: "GymEnv", eval_env: "GymEnv") -> None: def sync_envs_normalization(env: "GymEnv", eval_env: "GymEnv") -> None:
""" """

68
stable_baselines3/common/vec_env/obs_dict_wrapper.py Normal file
View file

@ -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)

80
stable_baselines3/common/vec_env/vec_normalize.py
View file

@ -1,8 +1,11 @@
import pickle import pickle
from typing import Any, Dict from copy import deepcopy
from typing import Any, Dict, Union
import gym
import numpy as np import numpy as np
from stable_baselines3.common import utils
from stable_baselines3.common.running_mean_std import RunningMeanStd from stable_baselines3.common.running_mean_std import RunningMeanStd
from stable_baselines3.common.vec_env.base_vec_env import VecEnv, VecEnvStepReturn, VecEnvWrapper from stable_baselines3.common.vec_env.base_vec_env import VecEnv, VecEnvStepReturn, VecEnvWrapper
@ -34,7 +37,19 @@ class VecNormalize(VecEnvWrapper):
epsilon: float = 1e-8, epsilon: float = 1e-8,
): ):
VecEnvWrapper.__init__(self, venv) 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.ret_rms = RunningMeanStd(shape=())
self.clip_obs = clip_obs self.clip_obs = clip_obs
self.clip_reward = clip_reward self.clip_reward = clip_reward
@ -83,8 +98,9 @@ class VecNormalize(VecEnvWrapper):
if self.venv is not None: if self.venv is not None:
raise ValueError("Trying to set venv of already initialized VecNormalize wrapper.") raise ValueError("Trying to set venv of already initialized VecNormalize wrapper.")
VecEnvWrapper.__init__(self, venv) 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) self.ret = np.zeros(self.num_envs)
def step_wait(self) -> VecEnvStepReturn: def step_wait(self) -> VecEnvStepReturn:
@ -99,7 +115,12 @@ class VecNormalize(VecEnvWrapper):
self.old_reward = rews self.old_reward = rews
if self.training: 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) obs = self.normalize_obs(obs)
if self.training: if self.training:
@ -114,14 +135,38 @@ class VecNormalize(VecEnvWrapper):
self.ret = self.ret * self.gamma + reward self.ret = self.ret * self.gamma + reward
self.ret_rms.update(self.ret) 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. Normalize observations using this VecNormalize's observations statistics.
Calling this method does not update statistics. Calling this method does not update statistics.
""" """
# Avoid modifying by reference the original object
obs_ = deepcopy(obs)
if self.norm_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) if isinstance(obs, dict) and isinstance(self.obs_rms, dict):
return obs 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: 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) reward = np.clip(reward / np.sqrt(self.ret_rms.var + self.epsilon), -self.clip_reward, self.clip_reward)
return 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: if self.norm_obs:
return (obs * np.sqrt(self.obs_rms.var + self.epsilon)) + self.obs_rms.mean if isinstance(obs, dict) and isinstance(self.obs_rms, dict):
return obs 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: def unnormalize_reward(self, reward: np.ndarray) -> np.ndarray:
if self.norm_reward: if self.norm_reward:
return reward * np.sqrt(self.ret_rms.var + self.epsilon) return reward * np.sqrt(self.ret_rms.var + self.epsilon)
return reward 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 Returns an unnormalized version of the observations from the most recent
step or reset. step or reset.
""" """
return self.old_obs.copy() return deepcopy(self.old_obs)
def get_original_reward(self) -> np.ndarray: def get_original_reward(self) -> np.ndarray:
""" """
@ -155,9 +206,10 @@ class VecNormalize(VecEnvWrapper):
""" """
return self.old_reward.copy() return self.old_reward.copy()
def reset(self) -> np.ndarray: def reset(self) -> Union[np.ndarray, Dict[str, np.ndarray]]:
""" """
Reset all environments Reset all environments
:return: first observation of the episode
""" """
obs = self.venv.reset() obs = self.venv.reset()
self.old_obs = obs self.old_obs = obs

4
stable_baselines3/her/__init__.py Normal file
View file

@ -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

26
stable_baselines3/her/goal_selection_strategy.py Normal file
View file

@ -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,
}

567
stable_baselines3/her/her.py Normal file
View file

@ -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

370
stable_baselines3/her/her_replay_buffer.py Normal file
View file

@ -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)

2
stable_baselines3/version.txt
View file

@ -1 +1 @@
0.10.0a0 0.10.0a1

318
tests/test_her.py Normal file
View file

@ -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

8
tests/test_save_load.py
View file

@ -231,8 +231,10 @@ def test_exclude_include_saved_params(tmp_path, model_class):
def test_save_load_replay_buffer(tmp_path, model_class): def test_save_load_replay_buffer(tmp_path, model_class):
path = pathlib.Path(tmp_path / "logs/replay_buffer.pkl") path = pathlib.Path(tmp_path / "logs/replay_buffer.pkl")
path.parent.mkdir(exist_ok=True, parents=True) # to not raise a warning 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 = model_class(
model.learn(500) "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) old_replay_buffer = deepcopy(model.replay_buffer)
model.save_replay_buffer(path) model.save_replay_buffer(path)
model.replay_buffer = None model.replay_buffer = None
@ -305,7 +307,7 @@ def test_save_load_policy(tmp_path, model_class, policy_str):
if policy_str == "MlpPolicy": if policy_str == "MlpPolicy":
env = select_env(model_class) env = select_env(model_class)
else: else:
if model_class in [SAC, TD3, DQN]: if model_class in [SAC, TD3, DQN, DDPG]:
# Avoid memory error when using replay buffer # Avoid memory error when using replay buffer
# Reduce the size of the features # Reduce the size of the features
kwargs = dict(buffer_size=250) kwargs = dict(buffer_size=250)

9
tests/test_sde.py
View file

@ -54,6 +54,11 @@ def test_state_dependent_exploration_grad():
assert sigma_hat.grad.allclose(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("model_class", [SAC, A2C, PPO])
@pytest.mark.parametrize("sde_net_arch", [None, [32, 16], []]) @pytest.mark.parametrize("sde_net_arch", [None, [32, 16], []])
@pytest.mark.parametrize("use_expln", [False, True]) @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, seed=None,
create_eval_env=True, create_eval_env=True,
verbose=1, 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() model.policy.reset_noise()
if model_class == SAC: if model_class == SAC:
model.policy.actor.get_std() model.policy.actor.get_std()

138
tests/test_vec_normalize.py
View file

@ -1,8 +1,9 @@
import gym import gym
import numpy as np import numpy as np
import pytest 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.running_mean_std import RunningMeanStd
from stable_baselines3.common.vec_env import ( from stable_baselines3.common.vec_env import (
DummyVecEnv, DummyVecEnv,
@ -15,14 +16,68 @@ from stable_baselines3.common.vec_env import (
ENV_ID = "Pendulum-v0" 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(): def make_env():
return gym.make(ENV_ID) return gym.make(ENV_ID)
def make_dict_env():
return DummyDictEnv()
def check_rms_equal(rmsa, rmsb): def check_rms_equal(rmsa, rmsb):
assert np.all(rmsa.mean == rmsb.mean) if isinstance(rmsa, dict):
assert np.all(rmsa.var == rmsb.var) for key in rmsa.keys():
assert np.all(rmsa.count == rmsb.count) 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): def check_vec_norm_equal(norma, normb):
@ -56,6 +111,19 @@ def _make_warmstart_cartpole():
return venv 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(): def test_runningmeanstd():
"""Test RunningMeanStd object""" """Test RunningMeanStd object"""
for (x_1, x_2, x_3) in [ for (x_1, x_2, x_3) in [
@ -74,7 +142,8 @@ def test_runningmeanstd():
assert np.allclose(moments_1, moments_2) 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""" """Test VecNormalize Object"""
clip_obs = 0.5 clip_obs = 0.5
clip_reward = 5.0 clip_reward = 5.0
@ -85,7 +154,11 @@ def test_vec_env(tmp_path):
while not done[0]: while not done[0]:
actions = [norm_venv.action_space.sample()] actions = [norm_venv.action_space.sample()]
obs, rew, done, _ = norm_venv.step(actions) 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 assert np.max(np.abs(rew)) <= clip_reward
path = tmp_path / "vec_normalize" path = tmp_path / "vec_normalize"
@ -113,6 +186,26 @@ def test_get_original():
np.testing.assert_allclose(venv.normalize_reward(orig_rewards), rewards) 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(): def test_normalize_external():
venv = _make_warmstart_cartpole() venv = _make_warmstart_cartpole()
@ -123,21 +216,24 @@ def test_normalize_external():
assert np.all(norm_rewards < 1) 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): 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) 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) 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])) kwargs = dict(model_class=SAC, max_episode_length=200, online_sampling=True) if model_class == HER else {}
model.learn(total_timesteps=1000, eval_env=eval_env, eval_freq=500) 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 # Check getter
assert isinstance(model.get_vec_normalize_env(), VecNormalize) 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]) env = DummyVecEnv([make_env])
assert unwrap_vec_normalize(env) is None assert unwrap_vec_normalize(env) is None
@ -146,13 +242,15 @@ def test_sync_vec_normalize():
assert isinstance(unwrap_vec_normalize(env), VecNormalize) assert isinstance(unwrap_vec_normalize(env), VecNormalize)
env = VecFrameStack(env, 1) if not isinstance(env.observation_space, spaces.Dict):
env = VecFrameStack(env, 1)
assert isinstance(unwrap_vec_normalize(env), VecNormalize) assert isinstance(unwrap_vec_normalize(env), VecNormalize)
eval_env = DummyVecEnv([make_env]) 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 = 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.seed(0)
env.action_space.seed(0) env.action_space.seed(0)
@ -171,12 +269,12 @@ def test_sync_vec_normalize():
dummy_rewards = np.random.rand(10) dummy_rewards = np.random.rand(10)
original_obs = env.get_original_obs() original_obs = env.get_original_obs()
# Check that unnormalization works # 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) # 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 # Test syncing of parameters
sync_envs_normalization(env, eval_env) sync_envs_normalization(env, eval_env)
# Now they must be synced # Now they must be synced
assert np.allclose(obs, eval_env.normalize_obs(original_obs)) assert allclose(obs, eval_env.normalize_obs(original_obs))
assert np.allclose(env.normalize_reward(dummy_rewards), eval_env.normalize_reward(dummy_rewards)) assert allclose(env.normalize_reward(dummy_rewards), eval_env.normalize_reward(dummy_rewards))