mirror of
https://github.com/saymrwulf/stable-baselines3.git
synced 2026-09-15 22:10:25 +00:00
Implement HER (#120)
* Added working her version, Online sampling is missing. * Updated test_her. * Added first version of online her sampling. Still problems with tensor dimensions. * Reformat * Fixed tests * Added some comments. * Updated changelog. * Add missing init file * Fixed some small bugs. * Reduced arguments for HER, small changes. * Added getattr. Fixed bug for online sampling. * Updated save/load funtions. Small changes. * Added her to init. * Updated save method. * Updated her ratio. * Move obs_wrapper * Added DQN test. * Fix potential bug * Offline and online her share same sample_goal function. * Changed lists into arrays. * Updated her test. * Fix online sampling * Fixed action bug. Updated time limit for episodes. * Updated convert_dict method to take keys as arguments. * Renamed obs dict wrapper. * Seed bit flipping env * Remove get_episode_dict * Add fast online sampling version * Added documentation. * Vectorized reward computation * Vectorized goal sampling * Update time limit for episodes in online her sampling. * Fix max episode length inference * Bug fix for Fetch envs * Fix for HER + gSDE * Reformat (new black version) * Added info dict to compute new reward. Check her_replay_buffer again. * Fix info buffer * Updated done flag. * Fixes for gSDE * Offline her version uses now HerReplayBuffer as episode storage. * Fix num_timesteps computation * Fix get torch params * Vectorized version for offline sampling. * Modified offline her sampling to use sample method of her_replay_buffer * Updated HER tests. * Updated documentation * Cleanup docstrings * Updated to review comments * Fix pytype * Update according to review comments. * Removed random goal strategy. Updated sample transitions. * Updated migration. Removed time signal removal. * Update doc * Fix potential load issue * Add VecNormalize support for dict obs * Updated saving/loading replay buffer for HER. * Fix test memory usage * Fixed save/load replay buffer. * Fixed save/load replay buffer * Fixed transition index after loading replay buffer in online sampling * Better error handling * Add tests for get_time_limit * More tests for VecNormalize with dict obs * Update doc * Improve HER description * Add test for sde support * Add comments * Add comments * Remove check that was always valid * Fix for terminal observation * Updated buffer size in offline version and reset of HER buffer * Reformat * Update doc * Remove np.empty + add doc * Fix loading * Updated loading replay buffer * Separate online and offline sampling + bug fixes * Update tensorboard log name * Version bump * Bug fix for special case Co-authored-by: Antonin Raffin <antonin.raffin@dlr.de> Co-authored-by: Antonin RAFFIN <antonin.raffin@ensta.org>
This commit is contained in:
Megan Klaiber
GitHub
parent
15e94a6d14
commit
dd6e361204
34 changed files with 1899 additions and 102 deletions
2
Makefile
2
Makefile
|
|
@ -5,7 +5,7 @@ pytest:
|
|||
./scripts/run_tests.sh
|
||||
|
||||
type:
|
||||
pytype
|
||||
pytype -j auto
|
||||
|
||||
lint:
|
||||
# stop the build if there are Python syntax errors or undefined names
|
||||
|
|
|
|||
13
README.md
13
README.md
|
|
@ -35,20 +35,9 @@ These algorithms will make it easier for the research community and industry to
|
|||
| Type hints | :heavy_check_mark: |
|
||||
|
||||
|
||||
### Roadmap to V1.0
|
||||
|
||||
Please look at the issue for more details.
|
||||
Planned features:
|
||||
|
||||
- [ ] HER
|
||||
|
||||
### Planned features (v1.1+)
|
||||
|
||||
- [ ] DQN extensions (prioritized replay, double q-learning, ...)
|
||||
- [ ] Support for `Tuple` and `Dict` observation spaces
|
||||
- [ ] Recurrent Policies
|
||||
- [ ] TRPO
|
||||
|
||||
Please take a look at the [Roadmap](https://github.com/DLR-RM/stable-baselines3/issues/1) and [Milestones](https://github.com/DLR-RM/stable-baselines3/milestones).
|
||||
|
||||
## Migration guide: from Stable-Baselines (SB2) to Stable-Baselines3 (SB3)
|
||||
|
||||
|
|
|
|||
|
|
@ -18,8 +18,7 @@ notebooks:
|
|||
- `Atari Games`_
|
||||
- `RL Baselines zoo`_
|
||||
- `PyBullet`_
|
||||
|
||||
.. - `Hindsight Experience Replay`_
|
||||
- `Hindsight Experience Replay`_
|
||||
|
||||
.. _Getting Started: https://colab.research.google.com/github/Stable-Baselines-Team/rl-colab-notebooks/blob/sb3/stable_baselines_getting_started.ipynb
|
||||
.. _Training, Saving, Loading: https://colab.research.google.com/github/Stable-Baselines-Team/rl-colab-notebooks/blob/sb3/saving_loading_dqn.ipynb
|
||||
|
|
@ -343,6 +342,81 @@ will compute a running average and standard deviation of input features (it can
|
|||
env.norm_reward = False
|
||||
|
||||
|
||||
Hindsight Experience Replay (HER)
|
||||
---------------------------------
|
||||
|
||||
For this example, we are using `Highway-Env <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
|
||||
--------------
|
||||
|
||||
|
|
|
|||
|
|
@ -163,6 +163,14 @@ Despite this change, no change in performance should be expected.
|
|||
To match SB2 behavior, you need to explicitly pass ``deterministic=True``
|
||||
|
||||
|
||||
HER
|
||||
^^^
|
||||
|
||||
The ``HER`` implementation now also supports online sampling of the new goals. This is done in a vectorized version.
|
||||
The goal selection strategy ``RANDOM`` is no longer supported.
|
||||
``HER`` now supports ``VecNormalize`` wrapper but only when ``online_sampling=True``.
|
||||
For performance reasons, the maximum number of steps per episodes must be specified (see :ref:`HER <her>` documentation).
|
||||
|
||||
|
||||
New logger API
|
||||
--------------
|
||||
|
|
|
|||
|
|
@ -57,6 +57,7 @@ Main Features
|
|||
modules/a2c
|
||||
modules/ddpg
|
||||
modules/dqn
|
||||
modules/her
|
||||
modules/ppo
|
||||
modules/sac
|
||||
modules/td3
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ Changelog
|
|||
==========
|
||||
|
||||
|
||||
Pre-Release 0.10.0a0 (WIP)
|
||||
Pre-Release 0.10.0a1 (WIP)
|
||||
------------------------------
|
||||
|
||||
Breaking Changes:
|
||||
|
|
@ -14,11 +14,14 @@ Breaking Changes:
|
|||
New Features:
|
||||
^^^^^^^^^^^^^
|
||||
- Allow custom actor/critic network architectures using ``net_arch=dict(qf=[400, 300], pi=[64, 64])`` for off-policy algorithms (SAC, TD3, DDPG)
|
||||
- Added Hindsight Experience Replay ``HER``. (@megan-klaiber)
|
||||
- ``VecNormalize`` now supports ``gym.spaces.Dict`` observation spaces
|
||||
- Support logging videos to Tensorboard (@SwamyDev)
|
||||
|
||||
Bug Fixes:
|
||||
^^^^^^^^^^
|
||||
- Fix GAE computation for on-policy algorithms (off-by one for the last value) (thanks @Wovchena)
|
||||
- Fixed potential issue when loading a different environment
|
||||
- Fix ignoring the exclude parameter when recording logs using json, csv or log as logging format (@SwamyDev)
|
||||
- Make ``make_vec_env`` support the ``env_kwargs`` argument when using an env ID str (@ManifoldFR)
|
||||
- Fix model creation initializing CUDA even when `device="cpu"` is provided
|
||||
|
|
@ -37,6 +40,7 @@ Others:
|
|||
Documentation:
|
||||
^^^^^^^^^^^^^^
|
||||
- Added first draft of migration guide
|
||||
- Enabled doc for ``CnnPolicies``
|
||||
|
||||
|
||||
Pre-Release 0.9.0 (2020-10-03)
|
||||
|
|
@ -68,6 +72,7 @@ New Features:
|
|||
|
||||
Bug Fixes:
|
||||
^^^^^^^^^^
|
||||
- Added ``unwrap_vec_wrapper()`` to ``common.vec_env`` to extract ``VecEnvWrapper`` if needed
|
||||
- Fixed a bug where the environment was reset twice when using ``evaluate_policy``
|
||||
- Fix logging of ``clip_fraction`` in PPO (@diditforlulz273)
|
||||
- Fixed a bug where cuda support was wrongly checked when passing the GPU index, e.g., ``device="cuda:0"`` (@liorcohen5)
|
||||
|
|
@ -160,7 +165,6 @@ Documentation:
|
|||
- Fixed typo in custom policy doc (@RaphaelWag)
|
||||
|
||||
|
||||
|
||||
Pre-Release 0.7.0 (2020-06-10)
|
||||
------------------------------
|
||||
|
||||
|
|
@ -461,4 +465,4 @@ And all the contributors:
|
|||
@MarvineGothic @jdossgollin @SyllogismRXS @rusu24edward @jbulow @Antymon @seheevic @justinkterry @edbeeching
|
||||
@flodorner @KuKuXia @NeoExtended @PartiallyTyped @mmcenta @richardwu @kinalmehta @rolandgvc @tkelestemur @mloo3
|
||||
@tirafesi @blurLake @koulakis @joeljosephjin @shwang @rk37 @andyshih12 @RaphaelWag @xicocaio
|
||||
@diditforlulz273 @liorcohen5 @ManifoldFR @mloo3 @SwamyDev @wmmc88
|
||||
@diditforlulz273 @liorcohen5 @ManifoldFR @mloo3 @SwamyDev @wmmc88 @megan-klaiber
|
||||
|
|
|
|||
|
|
@ -79,3 +79,22 @@ Parameters
|
|||
.. autoclass:: A2C
|
||||
:members:
|
||||
:inherited-members:
|
||||
|
||||
|
||||
A2C Policies
|
||||
-------------
|
||||
|
||||
.. autoclass:: MlpPolicy
|
||||
:members:
|
||||
:inherited-members:
|
||||
|
||||
.. autoclass:: stable_baselines3.common.policies.ActorCriticPolicy
|
||||
:members:
|
||||
:noindex:
|
||||
|
||||
.. autoclass:: CnnPolicy
|
||||
:members:
|
||||
|
||||
.. autoclass:: stable_baselines3.common.policies.ActorCriticCnnPolicy
|
||||
:members:
|
||||
:noindex:
|
||||
|
|
|
|||
|
|
@ -98,7 +98,9 @@ DDPG Policies
|
|||
:members:
|
||||
:inherited-members:
|
||||
|
||||
.. autoclass:: stable_baselines3.td3.policies.TD3Policy
|
||||
:members:
|
||||
:noindex:
|
||||
|
||||
.. .. autoclass:: CnnPolicy
|
||||
.. :members:
|
||||
.. :inherited-members:
|
||||
.. autoclass:: CnnPolicy
|
||||
:members:
|
||||
|
|
|
|||
|
|
@ -90,5 +90,9 @@ DQN Policies
|
|||
:members:
|
||||
:inherited-members:
|
||||
|
||||
.. autoclass:: stable_baselines3.dqn.policies.DQNPolicy
|
||||
:members:
|
||||
:noindex:
|
||||
|
||||
.. autoclass:: CnnPolicy
|
||||
:members:
|
||||
|
|
|
|||
121
docs/modules/her.rst
Normal file
121
docs/modules/her.rst
Normal 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:
|
||||
|
|
@ -80,3 +80,22 @@ Parameters
|
|||
.. autoclass:: PPO
|
||||
:members:
|
||||
:inherited-members:
|
||||
|
||||
|
||||
PPO Policies
|
||||
-------------
|
||||
|
||||
.. autoclass:: MlpPolicy
|
||||
:members:
|
||||
:inherited-members:
|
||||
|
||||
.. autoclass:: stable_baselines3.common.policies.ActorCriticPolicy
|
||||
:members:
|
||||
:noindex:
|
||||
|
||||
.. autoclass:: CnnPolicy
|
||||
:members:
|
||||
|
||||
.. autoclass:: stable_baselines3.common.policies.ActorCriticCnnPolicy
|
||||
:members:
|
||||
:noindex:
|
||||
|
|
|
|||
|
|
@ -82,7 +82,7 @@ Example
|
|||
|
||||
obs = env.reset()
|
||||
while True:
|
||||
action, _states = model.predict(obs)
|
||||
action, _states = model.predict(obs, deterministic=True)
|
||||
obs, reward, done, info = env.step(action)
|
||||
env.render()
|
||||
if done:
|
||||
|
|
@ -104,6 +104,9 @@ SAC Policies
|
|||
:members:
|
||||
:inherited-members:
|
||||
|
||||
.. .. autoclass:: CnnPolicy
|
||||
.. :members:
|
||||
.. :inherited-members:
|
||||
.. autoclass:: stable_baselines3.sac.policies.SACPolicy
|
||||
:members:
|
||||
:noindex:
|
||||
|
||||
.. autoclass:: CnnPolicy
|
||||
:members:
|
||||
|
|
|
|||
|
|
@ -101,7 +101,9 @@ TD3 Policies
|
|||
:members:
|
||||
:inherited-members:
|
||||
|
||||
.. autoclass:: stable_baselines3.td3.policies.TD3Policy
|
||||
:members:
|
||||
:noindex:
|
||||
|
||||
.. .. autoclass:: CnnPolicy
|
||||
.. :members:
|
||||
.. :inherited-members:
|
||||
.. autoclass:: CnnPolicy
|
||||
:members:
|
||||
|
|
|
|||
|
|
@ -29,6 +29,7 @@ per-file-ignores =
|
|||
./stable_baselines3/a2c/__init__.py:F401
|
||||
./stable_baselines3/ddpg/__init__.py:F401
|
||||
./stable_baselines3/dqn/__init__.py:F401
|
||||
./stable_baselines3/her/__init__.py:F401
|
||||
./stable_baselines3/ppo/__init__.py:F401
|
||||
./stable_baselines3/sac/__init__.py:F401
|
||||
./stable_baselines3/td3/__init__.py:F401
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@ import os
|
|||
from stable_baselines3.a2c import A2C
|
||||
from stable_baselines3.ddpg import DDPG
|
||||
from stable_baselines3.dqn import DQN
|
||||
from stable_baselines3.her import HER
|
||||
from stable_baselines3.ppo import PPO
|
||||
from stable_baselines3.sac import SAC
|
||||
from stable_baselines3.td3 import TD3
|
||||
|
|
|
|||
|
|
@ -34,7 +34,7 @@ class NoopResetEnv(gym.Wrapper):
|
|||
else:
|
||||
noops = self.unwrapped.np_random.randint(1, self.noop_max + 1)
|
||||
assert noops > 0
|
||||
obs = np.empty(0)
|
||||
obs = np.zeros(0)
|
||||
for _ in range(noops):
|
||||
obs, _, done, _ = self.env.step(self.noop_action)
|
||||
if done:
|
||||
|
|
|
|||
|
|
@ -26,7 +26,15 @@ from stable_baselines3.common.utils import (
|
|||
set_random_seed,
|
||||
update_learning_rate,
|
||||
)
|
||||
from stable_baselines3.common.vec_env import DummyVecEnv, VecEnv, VecNormalize, VecTransposeImage, unwrap_vec_normalize
|
||||
from stable_baselines3.common.vec_env import (
|
||||
DummyVecEnv,
|
||||
VecEnv,
|
||||
VecNormalize,
|
||||
VecTransposeImage,
|
||||
is_wrapped,
|
||||
unwrap_vec_normalize,
|
||||
)
|
||||
from stable_baselines3.common.vec_env.obs_dict_wrapper import ObsDictWrapper
|
||||
|
||||
|
||||
def maybe_make_env(env: Union[GymEnv, str, None], monitor_wrapper: bool, verbose: int) -> Optional[GymEnv]:
|
||||
|
|
@ -146,7 +154,7 @@ class BaseAlgorithm(ABC):
|
|||
self.eval_env = maybe_make_env(env, monitor_wrapper, self.verbose)
|
||||
|
||||
env = maybe_make_env(env, monitor_wrapper, self.verbose)
|
||||
env = self._wrap_env(env)
|
||||
env = self._wrap_env(env, self.verbose)
|
||||
|
||||
self.observation_space = env.observation_space
|
||||
self.action_space = env.action_space
|
||||
|
|
@ -158,19 +166,25 @@ class BaseAlgorithm(ABC):
|
|||
"Error: the model does not support multiple envs; it requires " "a single vectorized environment."
|
||||
)
|
||||
|
||||
if self.use_sde and not isinstance(self.observation_space, gym.spaces.Box):
|
||||
if self.use_sde and not isinstance(self.action_space, gym.spaces.Box):
|
||||
raise ValueError("generalized State-Dependent Exploration (gSDE) can only be used with continuous actions.")
|
||||
|
||||
def _wrap_env(self, env: GymEnv) -> VecEnv:
|
||||
@staticmethod
|
||||
def _wrap_env(env: GymEnv, verbose: int = 0) -> VecEnv:
|
||||
if not isinstance(env, VecEnv):
|
||||
if self.verbose >= 1:
|
||||
if verbose >= 1:
|
||||
print("Wrapping the env in a DummyVecEnv.")
|
||||
env = DummyVecEnv([lambda: env])
|
||||
|
||||
if is_image_space(env.observation_space) and not isinstance(env, VecTransposeImage):
|
||||
if self.verbose >= 1:
|
||||
if is_image_space(env.observation_space) and not is_wrapped(env, VecTransposeImage):
|
||||
if verbose >= 1:
|
||||
print("Wrapping the env in a VecTransposeImage.")
|
||||
env = VecTransposeImage(env)
|
||||
|
||||
# check if wrapper for dict support is needed when using HER
|
||||
if isinstance(env.observation_space, gym.spaces.dict.Dict):
|
||||
env = ObsDictWrapper(env)
|
||||
|
||||
return env
|
||||
|
||||
@abstractmethod
|
||||
|
|
@ -188,7 +202,7 @@ class BaseAlgorithm(ABC):
|
|||
eval_env = self.eval_env
|
||||
|
||||
if eval_env is not None:
|
||||
eval_env = self._wrap_env(eval_env)
|
||||
eval_env = self._wrap_env(eval_env, self.verbose)
|
||||
assert eval_env.num_envs == 1
|
||||
return eval_env
|
||||
|
||||
|
|
@ -402,10 +416,11 @@ class BaseAlgorithm(ABC):
|
|||
|
||||
:param env: The environment for learning a policy
|
||||
"""
|
||||
check_for_correct_spaces(env, self.observation_space, self.action_space)
|
||||
# it must be coherent now
|
||||
# if it is not a VecEnv, make it a VecEnv
|
||||
env = self._wrap_env(env)
|
||||
# and do other transformations (dict obs, image transpose) if needed
|
||||
env = self._wrap_env(env, self.verbose)
|
||||
# Check that the observation spaces match
|
||||
check_for_correct_spaces(env, self.observation_space, self.action_space)
|
||||
|
||||
self.n_envs = env.num_envs
|
||||
self.env = env
|
||||
|
|
@ -576,6 +591,8 @@ class BaseAlgorithm(ABC):
|
|||
raise KeyError("The observation_space and action_space were not given, can't verify new environments")
|
||||
|
||||
if env is not None:
|
||||
# Wrap first if needed
|
||||
cls._wrap_env(env, data["verbose"])
|
||||
# Check if given env is valid
|
||||
check_for_correct_spaces(env, data["observation_space"], data["action_space"])
|
||||
else:
|
||||
|
|
|
|||
|
|
@ -1,8 +1,9 @@
|
|||
from collections import OrderedDict
|
||||
from typing import Dict, Optional, Union
|
||||
from typing import Any, Dict, Optional, Union
|
||||
|
||||
import numpy as np
|
||||
from gym import GoalEnv, spaces
|
||||
from gym.envs.registration import EnvSpec
|
||||
|
||||
from stable_baselines3.common.type_aliases import GymStepReturn
|
||||
|
||||
|
|
@ -22,6 +23,8 @@ class BitFlippingEnv(GoalEnv):
|
|||
version or not, by default, it uses the MultiBinary one
|
||||
"""
|
||||
|
||||
spec = EnvSpec("BitFlippingEnv-v0")
|
||||
|
||||
def __init__(
|
||||
self, n_bits: int = 10, continuous: bool = False, max_steps: Optional[int] = None, discrete_obs_space: bool = False
|
||||
):
|
||||
|
|
@ -61,7 +64,9 @@ class BitFlippingEnv(GoalEnv):
|
|||
max_steps = n_bits
|
||||
self.max_steps = max_steps
|
||||
self.current_step = 0
|
||||
self.reset()
|
||||
|
||||
def seed(self, seed: int) -> None:
|
||||
self.obs_space.seed(seed)
|
||||
|
||||
def convert_if_needed(self, state: np.ndarray) -> Union[int, np.ndarray]:
|
||||
"""
|
||||
|
|
@ -101,7 +106,7 @@ class BitFlippingEnv(GoalEnv):
|
|||
else:
|
||||
self.state[action] = 1 - self.state[action]
|
||||
obs = self._get_obs()
|
||||
reward = self.compute_reward(obs["achieved_goal"], obs["desired_goal"], None)
|
||||
reward = float(self.compute_reward(obs["achieved_goal"], obs["desired_goal"], None))
|
||||
done = reward == 0
|
||||
self.current_step += 1
|
||||
# Episode terminate when we reached the goal or the max number of steps
|
||||
|
|
@ -109,11 +114,13 @@ class BitFlippingEnv(GoalEnv):
|
|||
done = done or self.current_step >= self.max_steps
|
||||
return obs, reward, done, info
|
||||
|
||||
def compute_reward(self, achieved_goal: np.ndarray, desired_goal: np.ndarray, _info) -> float:
|
||||
def compute_reward(
|
||||
self, achieved_goal: Union[int, np.ndarray], desired_goal: Union[int, np.ndarray], _info: Optional[Dict[str, Any]]
|
||||
) -> np.float32:
|
||||
# Deceptive reward: it is positive only when the goal is achieved
|
||||
if self.discrete_obs_space:
|
||||
return 0.0 if achieved_goal == desired_goal else -1.0
|
||||
return 0.0 if (achieved_goal == desired_goal).all() else -1.0
|
||||
# vectorized version
|
||||
distance = np.linalg.norm(achieved_goal - desired_goal, axis=-1)
|
||||
return -(distance > 0).astype(np.float32)
|
||||
|
||||
def render(self, mode: str = "human") -> Optional[np.ndarray]:
|
||||
if mode == "rgb_array":
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
import warnings
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Generator, Optional, Union
|
||||
from typing import Dict, Generator, Optional, Union
|
||||
|
||||
import numpy as np
|
||||
import torch as th
|
||||
|
|
@ -129,9 +129,11 @@ class BaseBuffer(ABC):
|
|||
return th.as_tensor(array).to(self.device)
|
||||
|
||||
@staticmethod
|
||||
def _normalize_obs(obs: np.ndarray, env: Optional[VecNormalize] = None) -> np.ndarray:
|
||||
def _normalize_obs(
|
||||
obs: Union[np.ndarray, Dict[str, np.ndarray]], env: Optional[VecNormalize] = None
|
||||
) -> Union[np.ndarray, Dict[str, np.ndarray]]:
|
||||
if env is not None:
|
||||
return env.normalize_obs(obs).astype(np.float32)
|
||||
return env.normalize_obs(obs)
|
||||
return obs
|
||||
|
||||
@staticmethod
|
||||
|
|
|
|||
|
|
@ -67,6 +67,8 @@ class OffPolicyAlgorithm(BaseAlgorithm):
|
|||
:param use_sde_at_warmup: Whether to use gSDE instead of uniform sampling
|
||||
during the warm up phase (before learning starts)
|
||||
:param sde_support: Whether the model support gSDE or not
|
||||
:param remove_time_limit_termination: Remove terminations (dones) that are due to time limit.
|
||||
See https://github.com/hill-a/stable-baselines/issues/863
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
|
|
@ -97,6 +99,7 @@ class OffPolicyAlgorithm(BaseAlgorithm):
|
|||
sde_sample_freq: int = -1,
|
||||
use_sde_at_warmup: bool = False,
|
||||
sde_support: bool = True,
|
||||
remove_time_limit_termination: bool = False,
|
||||
):
|
||||
|
||||
super(OffPolicyAlgorithm, self).__init__(
|
||||
|
|
@ -126,6 +129,10 @@ class OffPolicyAlgorithm(BaseAlgorithm):
|
|||
self.action_noise = action_noise
|
||||
self.optimize_memory_usage = optimize_memory_usage
|
||||
|
||||
# Remove terminations (dones) that are due to time limit
|
||||
# see https://github.com/hill-a/stable-baselines/issues/863
|
||||
self.remove_time_limit_termination = remove_time_limit_termination
|
||||
|
||||
if train_freq > 0 and n_episodes_rollout > 0:
|
||||
warnings.warn(
|
||||
"You passed a positive value for `train_freq` and `n_episodes_rollout`."
|
||||
|
|
|
|||
|
|
@ -23,6 +23,7 @@ from stable_baselines3.common.preprocessing import get_action_dim, is_image_spac
|
|||
from stable_baselines3.common.torch_layers import BaseFeaturesExtractor, FlattenExtractor, MlpExtractor, NatureCNN, create_mlp
|
||||
from stable_baselines3.common.utils import get_device, is_vectorized_observation
|
||||
from stable_baselines3.common.vec_env import VecTransposeImage
|
||||
from stable_baselines3.common.vec_env.obs_dict_wrapper import ObsDictWrapper
|
||||
|
||||
|
||||
class BaseModel(nn.Module, ABC):
|
||||
|
|
@ -234,7 +235,10 @@ class BasePolicy(BaseModel):
|
|||
# state = self.initial_state
|
||||
# if mask is None:
|
||||
# mask = [False for _ in range(self.n_envs)]
|
||||
observation = np.array(observation)
|
||||
if isinstance(observation, dict):
|
||||
observation = ObsDictWrapper.convert_dict(observation)
|
||||
else:
|
||||
observation = np.array(observation)
|
||||
|
||||
# Handle the different cases for images
|
||||
# as PyTorch use channel first format
|
||||
|
|
|
|||
|
|
@ -16,9 +16,7 @@ except ImportError:
|
|||
SummaryWriter = None
|
||||
|
||||
from stable_baselines3.common import logger
|
||||
from stable_baselines3.common.preprocessing import is_image_space
|
||||
from stable_baselines3.common.type_aliases import GymEnv
|
||||
from stable_baselines3.common.vec_env import VecTransposeImage
|
||||
|
||||
|
||||
def set_random_seed(seed: int, using_cuda: bool = False) -> None:
|
||||
|
|
@ -204,14 +202,7 @@ def check_for_correct_spaces(env: GymEnv, observation_space: gym.spaces.Space, a
|
|||
:param observation_space: Observation space to check against
|
||||
:param action_space: Action space to check against
|
||||
"""
|
||||
if (
|
||||
observation_space != env.observation_space
|
||||
# Special cases for images that need to be transposed
|
||||
and not (
|
||||
is_image_space(env.observation_space)
|
||||
and observation_space == VecTransposeImage.transpose_space(env.observation_space)
|
||||
)
|
||||
):
|
||||
if observation_space != env.observation_space:
|
||||
raise ValueError(f"Observation spaces do not match: {observation_space} != {env.observation_space}")
|
||||
if action_space != env.action_space:
|
||||
raise ValueError(f"Action spaces do not match: {action_space} != {env.action_space}")
|
||||
|
|
|
|||
|
|
@ -41,6 +41,17 @@ def unwrap_vec_normalize(env: Union["GymEnv", VecEnv]) -> Optional[VecNormalize]
|
|||
return unwrap_vec_wrapper(env, VecNormalize) # pytype:disable=bad-return-type
|
||||
|
||||
|
||||
def is_wrapped(env: Union["GymEnv", VecEnv], vec_wrapper_class: Type[VecEnvWrapper]) -> bool:
|
||||
"""
|
||||
Check if an environment is already wrapped by a given ``VecEnvWrapper``.
|
||||
|
||||
:param env:
|
||||
:param vec_wrapper_class:
|
||||
:return:
|
||||
"""
|
||||
return unwrap_vec_wrapper(env, vec_wrapper_class) is not None
|
||||
|
||||
|
||||
# Define here to avoid circular import
|
||||
def sync_envs_normalization(env: "GymEnv", eval_env: "GymEnv") -> None:
|
||||
"""
|
||||
|
|
|
|||
68
stable_baselines3/common/vec_env/obs_dict_wrapper.py
Normal file
68
stable_baselines3/common/vec_env/obs_dict_wrapper.py
Normal 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)
|
||||
|
|
@ -1,8 +1,11 @@
|
|||
import pickle
|
||||
from typing import Any, Dict
|
||||
from copy import deepcopy
|
||||
from typing import Any, Dict, Union
|
||||
|
||||
import gym
|
||||
import numpy as np
|
||||
|
||||
from stable_baselines3.common import utils
|
||||
from stable_baselines3.common.running_mean_std import RunningMeanStd
|
||||
from stable_baselines3.common.vec_env.base_vec_env import VecEnv, VecEnvStepReturn, VecEnvWrapper
|
||||
|
||||
|
|
@ -34,7 +37,19 @@ class VecNormalize(VecEnvWrapper):
|
|||
epsilon: float = 1e-8,
|
||||
):
|
||||
VecEnvWrapper.__init__(self, venv)
|
||||
self.obs_rms = RunningMeanStd(shape=self.observation_space.shape)
|
||||
|
||||
assert isinstance(
|
||||
self.observation_space, (gym.spaces.Box, gym.spaces.Dict)
|
||||
), "VecNormalize only support `gym.spaces.Box` and `gym.spaces.Dict` observation spaces"
|
||||
|
||||
if isinstance(self.observation_space, gym.spaces.Dict):
|
||||
self.obs_keys = set(self.observation_space.spaces.keys())
|
||||
self.obs_spaces = self.observation_space.spaces
|
||||
self.obs_rms = {key: RunningMeanStd(shape=space.shape) for key, space in self.obs_spaces.items()}
|
||||
else:
|
||||
self.obs_keys, self.obs_spaces = None, None
|
||||
self.obs_rms = RunningMeanStd(shape=self.observation_space.shape)
|
||||
|
||||
self.ret_rms = RunningMeanStd(shape=())
|
||||
self.clip_obs = clip_obs
|
||||
self.clip_reward = clip_reward
|
||||
|
|
@ -83,8 +98,9 @@ class VecNormalize(VecEnvWrapper):
|
|||
if self.venv is not None:
|
||||
raise ValueError("Trying to set venv of already initialized VecNormalize wrapper.")
|
||||
VecEnvWrapper.__init__(self, venv)
|
||||
if self.obs_rms.mean.shape != self.observation_space.shape:
|
||||
raise ValueError("venv is incompatible with current statistics.")
|
||||
|
||||
# Check only that the observation_space match
|
||||
utils.check_for_correct_spaces(venv, self.observation_space, venv.action_space)
|
||||
self.ret = np.zeros(self.num_envs)
|
||||
|
||||
def step_wait(self) -> VecEnvStepReturn:
|
||||
|
|
@ -99,7 +115,12 @@ class VecNormalize(VecEnvWrapper):
|
|||
self.old_reward = rews
|
||||
|
||||
if self.training:
|
||||
self.obs_rms.update(obs)
|
||||
if isinstance(obs, dict) and isinstance(self.obs_rms, dict):
|
||||
for key in self.obs_rms.keys():
|
||||
self.obs_rms[key].update(obs[key])
|
||||
else:
|
||||
self.obs_rms.update(obs)
|
||||
|
||||
obs = self.normalize_obs(obs)
|
||||
|
||||
if self.training:
|
||||
|
|
@ -114,14 +135,38 @@ class VecNormalize(VecEnvWrapper):
|
|||
self.ret = self.ret * self.gamma + reward
|
||||
self.ret_rms.update(self.ret)
|
||||
|
||||
def normalize_obs(self, obs: np.ndarray) -> np.ndarray:
|
||||
def _normalize_obs(self, obs: np.ndarray, obs_rms: RunningMeanStd) -> np.ndarray:
|
||||
"""
|
||||
Helper to normalize observation.
|
||||
:param obs:
|
||||
:param obs_rms: associated statistics
|
||||
:return: normalized observation
|
||||
"""
|
||||
return np.clip((obs - obs_rms.mean) / np.sqrt(obs_rms.var + self.epsilon), -self.clip_obs, self.clip_obs)
|
||||
|
||||
def _unnormalize_obs(self, obs: np.ndarray, obs_rms: RunningMeanStd) -> np.ndarray:
|
||||
"""
|
||||
Helper to unnormalize observation.
|
||||
:param obs:
|
||||
:param obs_rms: associated statistics
|
||||
:return: unnormalized observation
|
||||
"""
|
||||
return (obs * np.sqrt(obs_rms.var + self.epsilon)) + obs_rms.mean
|
||||
|
||||
def normalize_obs(self, obs: Union[np.ndarray, Dict[str, np.ndarray]]) -> Union[np.ndarray, Dict[str, np.ndarray]]:
|
||||
"""
|
||||
Normalize observations using this VecNormalize's observations statistics.
|
||||
Calling this method does not update statistics.
|
||||
"""
|
||||
# Avoid modifying by reference the original object
|
||||
obs_ = deepcopy(obs)
|
||||
if self.norm_obs:
|
||||
obs = np.clip((obs - self.obs_rms.mean) / np.sqrt(self.obs_rms.var + self.epsilon), -self.clip_obs, self.clip_obs)
|
||||
return obs
|
||||
if isinstance(obs, dict) and isinstance(self.obs_rms, dict):
|
||||
for key in self.obs_rms.keys():
|
||||
obs_[key] = self._normalize_obs(obs[key], self.obs_rms[key]).astype(np.float32)
|
||||
else:
|
||||
obs_ = self._normalize_obs(obs, self.obs_rms).astype(np.float32)
|
||||
return obs_
|
||||
|
||||
def normalize_reward(self, reward: np.ndarray) -> np.ndarray:
|
||||
"""
|
||||
|
|
@ -132,22 +177,28 @@ class VecNormalize(VecEnvWrapper):
|
|||
reward = np.clip(reward / np.sqrt(self.ret_rms.var + self.epsilon), -self.clip_reward, self.clip_reward)
|
||||
return reward
|
||||
|
||||
def unnormalize_obs(self, obs: np.ndarray) -> np.ndarray:
|
||||
def unnormalize_obs(self, obs: Union[np.ndarray, Dict[str, np.ndarray]]) -> Union[np.ndarray, Dict[str, np.ndarray]]:
|
||||
# Avoid modifying by reference the original object
|
||||
obs_ = deepcopy(obs)
|
||||
if self.norm_obs:
|
||||
return (obs * np.sqrt(self.obs_rms.var + self.epsilon)) + self.obs_rms.mean
|
||||
return obs
|
||||
if isinstance(obs, dict) and isinstance(self.obs_rms, dict):
|
||||
for key in self.obs_rms.keys():
|
||||
obs_[key] = self._unnormalize_obs(obs[key], self.obs_rms[key])
|
||||
else:
|
||||
obs_ = self._unnormalize_obs(obs, self.obs_rms)
|
||||
return obs_
|
||||
|
||||
def unnormalize_reward(self, reward: np.ndarray) -> np.ndarray:
|
||||
if self.norm_reward:
|
||||
return reward * np.sqrt(self.ret_rms.var + self.epsilon)
|
||||
return reward
|
||||
|
||||
def get_original_obs(self) -> np.ndarray:
|
||||
def get_original_obs(self) -> Union[np.ndarray, Dict[str, np.ndarray]]:
|
||||
"""
|
||||
Returns an unnormalized version of the observations from the most recent
|
||||
step or reset.
|
||||
"""
|
||||
return self.old_obs.copy()
|
||||
return deepcopy(self.old_obs)
|
||||
|
||||
def get_original_reward(self) -> np.ndarray:
|
||||
"""
|
||||
|
|
@ -155,9 +206,10 @@ class VecNormalize(VecEnvWrapper):
|
|||
"""
|
||||
return self.old_reward.copy()
|
||||
|
||||
def reset(self) -> np.ndarray:
|
||||
def reset(self) -> Union[np.ndarray, Dict[str, np.ndarray]]:
|
||||
"""
|
||||
Reset all environments
|
||||
:return: first observation of the episode
|
||||
"""
|
||||
obs = self.venv.reset()
|
||||
self.old_obs = obs
|
||||
|
|
|
|||
4
stable_baselines3/her/__init__.py
Normal file
4
stable_baselines3/her/__init__.py
Normal 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
26
stable_baselines3/her/goal_selection_strategy.py
Normal 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
567
stable_baselines3/her/her.py
Normal 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
370
stable_baselines3/her/her_replay_buffer.py
Normal 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)
|
||||
|
|
@ -1 +1 @@
|
|||
0.10.0a0
|
||||
0.10.0a1
|
||||
|
|
|
|||
318
tests/test_her.py
Normal file
318
tests/test_her.py
Normal 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
|
||||
|
|
@ -231,8 +231,10 @@ def test_exclude_include_saved_params(tmp_path, model_class):
|
|||
def test_save_load_replay_buffer(tmp_path, model_class):
|
||||
path = pathlib.Path(tmp_path / "logs/replay_buffer.pkl")
|
||||
path.parent.mkdir(exist_ok=True, parents=True) # to not raise a warning
|
||||
model = model_class("MlpPolicy", select_env(model_class), buffer_size=1000)
|
||||
model.learn(500)
|
||||
model = model_class(
|
||||
"MlpPolicy", select_env(model_class), buffer_size=1000, policy_kwargs=dict(net_arch=[64]), learning_starts=200
|
||||
)
|
||||
model.learn(300)
|
||||
old_replay_buffer = deepcopy(model.replay_buffer)
|
||||
model.save_replay_buffer(path)
|
||||
model.replay_buffer = None
|
||||
|
|
@ -305,7 +307,7 @@ def test_save_load_policy(tmp_path, model_class, policy_str):
|
|||
if policy_str == "MlpPolicy":
|
||||
env = select_env(model_class)
|
||||
else:
|
||||
if model_class in [SAC, TD3, DQN]:
|
||||
if model_class in [SAC, TD3, DQN, DDPG]:
|
||||
# Avoid memory error when using replay buffer
|
||||
# Reduce the size of the features
|
||||
kwargs = dict(buffer_size=250)
|
||||
|
|
|
|||
|
|
@ -54,6 +54,11 @@ def test_state_dependent_exploration_grad():
|
|||
assert sigma_hat.grad.allclose(grad)
|
||||
|
||||
|
||||
def test_sde_check():
|
||||
with pytest.raises(ValueError):
|
||||
PPO("MlpPolicy", "CartPole-v1", use_sde=True)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_class", [SAC, A2C, PPO])
|
||||
@pytest.mark.parametrize("sde_net_arch", [None, [32, 16], []])
|
||||
@pytest.mark.parametrize("use_expln", [False, True])
|
||||
|
|
@ -65,9 +70,9 @@ def test_state_dependent_offpolicy_noise(model_class, sde_net_arch, use_expln):
|
|||
seed=None,
|
||||
create_eval_env=True,
|
||||
verbose=1,
|
||||
policy_kwargs=dict(log_std_init=-2, sde_net_arch=sde_net_arch, use_expln=use_expln),
|
||||
policy_kwargs=dict(log_std_init=-2, sde_net_arch=sde_net_arch, use_expln=use_expln, net_arch=[64]),
|
||||
)
|
||||
model.learn(total_timesteps=int(500), eval_freq=250)
|
||||
model.learn(total_timesteps=int(300), eval_freq=250)
|
||||
model.policy.reset_noise()
|
||||
if model_class == SAC:
|
||||
model.policy.actor.get_std()
|
||||
|
|
|
|||
|
|
@ -1,8 +1,9 @@
|
|||
import gym
|
||||
import numpy as np
|
||||
import pytest
|
||||
from gym import spaces
|
||||
|
||||
from stable_baselines3 import SAC, TD3
|
||||
from stable_baselines3 import HER, SAC, TD3
|
||||
from stable_baselines3.common.running_mean_std import RunningMeanStd
|
||||
from stable_baselines3.common.vec_env import (
|
||||
DummyVecEnv,
|
||||
|
|
@ -15,14 +16,68 @@ from stable_baselines3.common.vec_env import (
|
|||
ENV_ID = "Pendulum-v0"
|
||||
|
||||
|
||||
class DummyDictEnv(gym.GoalEnv):
|
||||
"""
|
||||
Dummy gym goal env for testing purposes
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super(DummyDictEnv, self).__init__()
|
||||
self.observation_space = spaces.Dict(
|
||||
{
|
||||
"observation": spaces.Box(low=-20.0, high=20.0, shape=(4,), dtype=np.float32),
|
||||
"achieved_goal": spaces.Box(low=-20.0, high=20.0, shape=(4,), dtype=np.float32),
|
||||
"desired_goal": spaces.Box(low=-20.0, high=20.0, shape=(4,), dtype=np.float32),
|
||||
}
|
||||
)
|
||||
self.action_space = spaces.Box(low=-1, high=1, shape=(3,), dtype=np.float32)
|
||||
|
||||
def reset(self):
|
||||
return self.observation_space.sample()
|
||||
|
||||
def step(self, action):
|
||||
obs = self.observation_space.sample()
|
||||
reward = self.compute_reward(obs["achieved_goal"], obs["desired_goal"], {})
|
||||
done = np.random.rand() > 0.8
|
||||
return obs, reward, done, {}
|
||||
|
||||
def compute_reward(self, achieved_goal: np.ndarray, desired_goal: np.ndarray, _info) -> np.float32:
|
||||
distance = np.linalg.norm(achieved_goal - desired_goal, axis=-1)
|
||||
return -(distance > 0).astype(np.float32)
|
||||
|
||||
|
||||
def allclose(obs_1, obs_2):
|
||||
"""
|
||||
Generalized np.allclose() to work with dict spaces.
|
||||
"""
|
||||
if isinstance(obs_1, dict):
|
||||
all_close = True
|
||||
for key in obs_1.keys():
|
||||
if not np.allclose(obs_1[key], obs_2[key]):
|
||||
all_close = False
|
||||
break
|
||||
return all_close
|
||||
return np.allclose(obs_1, obs_2)
|
||||
|
||||
|
||||
def make_env():
|
||||
return gym.make(ENV_ID)
|
||||
|
||||
|
||||
def make_dict_env():
|
||||
return DummyDictEnv()
|
||||
|
||||
|
||||
def check_rms_equal(rmsa, rmsb):
|
||||
assert np.all(rmsa.mean == rmsb.mean)
|
||||
assert np.all(rmsa.var == rmsb.var)
|
||||
assert np.all(rmsa.count == rmsb.count)
|
||||
if isinstance(rmsa, dict):
|
||||
for key in rmsa.keys():
|
||||
assert np.all(rmsa[key].mean == rmsb[key].mean)
|
||||
assert np.all(rmsa[key].var == rmsb[key].var)
|
||||
assert np.all(rmsa[key].count == rmsb[key].count)
|
||||
else:
|
||||
assert np.all(rmsa.mean == rmsb.mean)
|
||||
assert np.all(rmsa.var == rmsb.var)
|
||||
assert np.all(rmsa.count == rmsb.count)
|
||||
|
||||
|
||||
def check_vec_norm_equal(norma, normb):
|
||||
|
|
@ -56,6 +111,19 @@ def _make_warmstart_cartpole():
|
|||
return venv
|
||||
|
||||
|
||||
def _make_warmstart_dict_env():
|
||||
"""Warm-start VecNormalize by stepping through BitFlippingEnv"""
|
||||
venv = DummyVecEnv([make_dict_env])
|
||||
venv = VecNormalize(venv)
|
||||
venv.reset()
|
||||
venv.get_original_obs()
|
||||
|
||||
for _ in range(100):
|
||||
actions = [venv.action_space.sample()]
|
||||
venv.step(actions)
|
||||
return venv
|
||||
|
||||
|
||||
def test_runningmeanstd():
|
||||
"""Test RunningMeanStd object"""
|
||||
for (x_1, x_2, x_3) in [
|
||||
|
|
@ -74,7 +142,8 @@ def test_runningmeanstd():
|
|||
assert np.allclose(moments_1, moments_2)
|
||||
|
||||
|
||||
def test_vec_env(tmp_path):
|
||||
@pytest.mark.parametrize("make_env", [make_env, make_dict_env])
|
||||
def test_vec_env(tmp_path, make_env):
|
||||
"""Test VecNormalize Object"""
|
||||
clip_obs = 0.5
|
||||
clip_reward = 5.0
|
||||
|
|
@ -85,7 +154,11 @@ def test_vec_env(tmp_path):
|
|||
while not done[0]:
|
||||
actions = [norm_venv.action_space.sample()]
|
||||
obs, rew, done, _ = norm_venv.step(actions)
|
||||
assert np.max(np.abs(obs)) <= clip_obs
|
||||
if isinstance(obs, dict):
|
||||
for key in obs.keys():
|
||||
assert np.max(np.abs(obs[key])) <= clip_obs
|
||||
else:
|
||||
assert np.max(np.abs(obs)) <= clip_obs
|
||||
assert np.max(np.abs(rew)) <= clip_reward
|
||||
|
||||
path = tmp_path / "vec_normalize"
|
||||
|
|
@ -113,6 +186,26 @@ def test_get_original():
|
|||
np.testing.assert_allclose(venv.normalize_reward(orig_rewards), rewards)
|
||||
|
||||
|
||||
def test_get_original_dict():
|
||||
venv = _make_warmstart_dict_env()
|
||||
for _ in range(3):
|
||||
actions = [venv.action_space.sample()]
|
||||
obs, rewards, _, _ = venv.step(actions)
|
||||
# obs = obs[0]
|
||||
orig_obs = venv.get_original_obs()
|
||||
rewards = rewards[0]
|
||||
orig_rewards = venv.get_original_reward()[0]
|
||||
|
||||
for key in orig_obs.keys():
|
||||
assert orig_obs[key].shape == obs[key].shape
|
||||
assert orig_rewards.dtype == rewards.dtype
|
||||
|
||||
assert not allclose(orig_obs, obs)
|
||||
assert not np.array_equal(orig_rewards, rewards)
|
||||
assert allclose(venv.normalize_obs(orig_obs), obs)
|
||||
np.testing.assert_allclose(venv.normalize_reward(orig_rewards), rewards)
|
||||
|
||||
|
||||
def test_normalize_external():
|
||||
venv = _make_warmstart_cartpole()
|
||||
|
||||
|
|
@ -123,21 +216,24 @@ def test_normalize_external():
|
|||
assert np.all(norm_rewards < 1)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_class", [SAC, TD3])
|
||||
@pytest.mark.parametrize("model_class", [SAC, TD3, HER])
|
||||
def test_offpolicy_normalization(model_class):
|
||||
env = DummyVecEnv([make_env])
|
||||
make_env_ = make_dict_env if model_class == HER else make_env
|
||||
env = DummyVecEnv([make_env_])
|
||||
env = VecNormalize(env, norm_obs=True, norm_reward=True, clip_obs=10.0, clip_reward=10.0)
|
||||
|
||||
eval_env = DummyVecEnv([make_env])
|
||||
eval_env = DummyVecEnv([make_env_])
|
||||
eval_env = VecNormalize(eval_env, training=False, norm_obs=True, norm_reward=False, clip_obs=10.0, clip_reward=10.0)
|
||||
|
||||
model = model_class("MlpPolicy", env, verbose=1, policy_kwargs=dict(net_arch=[64]))
|
||||
model.learn(total_timesteps=1000, eval_env=eval_env, eval_freq=500)
|
||||
kwargs = dict(model_class=SAC, max_episode_length=200, online_sampling=True) if model_class == HER else {}
|
||||
model = model_class("MlpPolicy", env, verbose=1, learning_starts=100, policy_kwargs=dict(net_arch=[64]), **kwargs)
|
||||
model.learn(total_timesteps=500, eval_env=eval_env, eval_freq=250)
|
||||
# Check getter
|
||||
assert isinstance(model.get_vec_normalize_env(), VecNormalize)
|
||||
|
||||
|
||||
def test_sync_vec_normalize():
|
||||
@pytest.mark.parametrize("make_env", [make_env, make_dict_env])
|
||||
def test_sync_vec_normalize(make_env):
|
||||
env = DummyVecEnv([make_env])
|
||||
|
||||
assert unwrap_vec_normalize(env) is None
|
||||
|
|
@ -146,13 +242,15 @@ def test_sync_vec_normalize():
|
|||
|
||||
assert isinstance(unwrap_vec_normalize(env), VecNormalize)
|
||||
|
||||
env = VecFrameStack(env, 1)
|
||||
|
||||
assert isinstance(unwrap_vec_normalize(env), VecNormalize)
|
||||
if not isinstance(env.observation_space, spaces.Dict):
|
||||
env = VecFrameStack(env, 1)
|
||||
assert isinstance(unwrap_vec_normalize(env), VecNormalize)
|
||||
|
||||
eval_env = DummyVecEnv([make_env])
|
||||
eval_env = VecNormalize(eval_env, training=False, norm_obs=True, norm_reward=True, clip_obs=100.0, clip_reward=100.0)
|
||||
eval_env = VecFrameStack(eval_env, 1)
|
||||
|
||||
if not isinstance(env.observation_space, spaces.Dict):
|
||||
eval_env = VecFrameStack(eval_env, 1)
|
||||
|
||||
env.seed(0)
|
||||
env.action_space.seed(0)
|
||||
|
|
@ -171,12 +269,12 @@ def test_sync_vec_normalize():
|
|||
dummy_rewards = np.random.rand(10)
|
||||
original_obs = env.get_original_obs()
|
||||
# Check that unnormalization works
|
||||
assert np.allclose(original_obs, env.unnormalize_obs(obs))
|
||||
assert allclose(original_obs, env.unnormalize_obs(obs))
|
||||
# Normalization must be different (between different environments)
|
||||
assert not np.allclose(obs, eval_env.normalize_obs(original_obs))
|
||||
assert not allclose(obs, eval_env.normalize_obs(original_obs))
|
||||
|
||||
# Test syncing of parameters
|
||||
sync_envs_normalization(env, eval_env)
|
||||
# Now they must be synced
|
||||
assert np.allclose(obs, eval_env.normalize_obs(original_obs))
|
||||
assert np.allclose(env.normalize_reward(dummy_rewards), eval_env.normalize_reward(dummy_rewards))
|
||||
assert allclose(obs, eval_env.normalize_obs(original_obs))
|
||||
assert allclose(env.normalize_reward(dummy_rewards), eval_env.normalize_reward(dummy_rewards))
|
||||
|
|
|
|||
Loading…
Reference in a new issue