DLR-RM · araffin · May 11, 2021 · Nov 24, 2020 · Nov 25, 2020 · Nov 26, 2020
diff --git a/README.md b/README.md
@@ -28,6 +28,7 @@ you can take a look at the issues [#48](https://github.com/DLR-RM/stable-baselin
 | Custom environments         | :heavy_check_mark: |
 | Custom policies             | :heavy_check_mark: |
 | Common interface            | :heavy_check_mark: |
+| `Dict` observation space support  | :heavy_check_mark: |
 | Ipython / Notebook friendly | :heavy_check_mark: |
 | Tensorboard support         | :heavy_check_mark: |
 | PEP8 code style             | :heavy_check_mark: |

diff --git a/docs/Makefile b/docs/Makefile
@@ -2,6 +2,7 @@
 #
 
 # You can set these variables from the command line.
+# For debug: SPHINXOPTS = -nWT --keep-going -vvv
 SPHINXOPTS    = -W  # make warnings fatal
 SPHINXBUILD   = sphinx-build
 SPHINXPROJ    = StableBaselines
@@ -17,4 +18,4 @@ help:
 # Catch-all target: route all unknown targets to Sphinx using the new
 # "make mode" option.  $(O) is meant as a shortcut for $(SPHINXOPTS).
 %: Makefile
-	@$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
+	@$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
diff --git a/docs/common/envs.rst b/docs/common/envs.rst
@@ -0,0 +1,24 @@
+.. _envs:
+
+.. automodule:: stable_baselines3.common.envs
+
+
+
+Custom Environments
+===================
+
+Those environments were created for testing purposes.
+
+
+BitFlippingEnv
+--------------
+
+.. autoclass:: BitFlippingEnv
+  :members:
+
+
+SimpleMultiObsEnv
+-----------------
+
+.. autoclass:: SimpleMultiObsEnv
+  :members:
diff --git a/docs/guide/algos.rst b/docs/guide/algos.rst
@@ -19,7 +19,9 @@ TD3          ✔️          ❌            ❌                ❌
 
 
 .. note::
-    Non-array spaces such as ``Dict`` or ``Tuple`` are not currently supported by any algorithm.
+  ``Tuple`` observation spaces are not supported by any environment
+  however single-level ``Dict`` spaces are (cf. :ref:`Examples <examples>`).
+
 
 Actions ``gym.spaces``:
 
@@ -41,6 +43,15 @@ Actions ``gym.spaces``:
   See `Issue #339 <https://github.com/hill-a/stable-baselines/issues/339>`_ for more info.
 
 
+.. note::
+
+  When using off-policy algorithms, `Time Limits <https://arxiv.org/abs/1712.00378>`_ (aka timeouts) are handled
+  properly (cf. `issue #284 <https://github.com/DLR-RM/stable-baselines3/issues/284>`_).
+  You can revert to SB3 < 2.1.0 behavior by passing ``handle_timeout_termination=False``
+  via the ``replay_buffer_kwargs`` argument.
+
+
+
 Reproducibility
 ---------------
 

diff --git a/docs/guide/custom_env.rst b/docs/guide/custom_env.rst
@@ -85,5 +85,5 @@ that will allow you to create the RL agent in one line (and use ``gym.make()`` t
 
 
 In the project, for testing purposes, we use a custom environment named ``IdentityEnv``
-defined `in this file <https://github.com/hill-a/stable-baselines/blob/master/stable_baselines/common/identity_env.py>`_.
-An example of how to use it can be found `here <https://github.com/hill-a/stable-baselines/blob/master/tests/test_identity.py>`_.
+defined `in this file <https://github.com/DLR-RM/stable-baselines3/blob/master/stable_baselines3/common/envs/identity_env.py>`_.
+An example of how to use it can be found `here <https://github.com/DLR-RM/stable-baselines3/blob/master/tests/test_identity.py>`_.
diff --git a/docs/guide/custom_policy.rst b/docs/guide/custom_policy.rst
@@ -3,8 +3,8 @@
 Custom Policy Network
 =====================
 
-Stable Baselines3 provides policy networks for images (CnnPolicies)
-and other type of input features (MlpPolicies).
+Stable Baselines3 provides policy networks for images (CnnPolicies),
+other type of input features (MlpPolicies) and multiple different inputs (MultiInputPolicies).
 
 
 .. warning::
@@ -149,6 +149,70 @@ that derives from ``BaseFeaturesExtractor`` and then pass it to the model when t
   model.learn(1000)
 
 
+Multiple Inputs and Dictionary Observations
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Stable Baselines3 supports handling of multiple inputs by using ``Dict`` Gym space. This can be done using
+``MultiInputPolicy``, which by default uses the ``CombinedExtractor`` feature extractor to turn multiple
+inputs into a single vector, handled by the ``net_arch`` network.
+
+By default, ``CombinedExtractor`` processes multiple inputs as follows:
+
+1. If input is an image (automatically detected, see ``common.preprocessing.is_image_space``), process image with Nature Atari CNN network and
+   output a latent vector of size ``256``.
+2. If input is not an image, flatten it (no layers).
+3. Concatenate all previous vectors into one long vector and pass it to policy.
+
+Much like above, you can define custom feature extractors. The following example assumes the environment has two keys in the
+observation space dictionary: "image" is a (1,H,W) image (channel first), and "vector" is a (D,) dimensional vector. We process "image" with a simple
+downsampling and "vector" with a single linear layer.
+
+.. code-block:: python
+
+  import gym
+  import torch as th
+  from torch import nn
+
+  from stable_baselines3.common.torch_layers import BaseFeaturesExtractor
+
+  class CustomCombinedExtractor(BaseFeaturesExtractor):
+      def __init__(self, observation_space: gym.spaces.Dict):
+          # We do not know features-dim here before going over all the items,
+          # so put something dummy for now. PyTorch requires calling
+          # nn.Module.__init__ before adding modules
+          super(CustomCombinedExtractor, self).__init__(observation_space, features_dim=1)
+
+          extractors = {}
+
+          total_concat_size = 0
+          # We need to know size of the output of this extractor,
+          # so go over all the spaces and compute output feature sizes
+          for key, subspace in observation_space.spaces.items():
+              if key == "image":
+                  # We will just downsample one channel of the image by 4x4 and flatten.
+                  # Assume the image is single-channel (subspace.shape[0] == 0)
+                  extractors[key] = nn.Sequential(nn.MaxPool2d(4), nn.Flatten())
+                  total_concat_size += subspace.shape[1] // 4 * subspace.shape[2] // 4
+              elif key == "vector":
+                  # Run through a simple MLP
+                  extractors[key] = nn.Linear(subspace.shape[0], 16)
+                  total_concat_size += 16
+
+          self.extractors = nn.ModuleDict(extractors)
+
+          # Update the features dim manually
+          self._features_dim = total_concat_size
+
+      def forward(self, observations) -> th.Tensor:
+          encoded_tensor_list = []
+
+          # self.extractors contain nn.Modules that do all the processing.
+          for key, extractor in self.extractors.items():
+              encoded_tensor_list.append(extractor(observations[key]))
+          # Return a (B, self._features_dim) PyTorch tensor, where B is batch dimension.
+          return th.cat(encoded_tensor_list, dim=1)
+
+
 
 On-Policy Algorithms
 ^^^^^^^^^^^^^^^^^^^^

diff --git a/docs/guide/examples.rst b/docs/guide/examples.rst
@@ -149,6 +149,27 @@ Multiprocessing: Unleashing the Power of Vectorized Environments
           env.render()
 
 
+Dict Observations
+-----------------
+
+You can use environments with dictionary observation spaces. This is useful in the case where one can't directly
+concatenate observations such as an image from a camera combined with a vector of servo sensor data (e.g., rotation angles).
+Stable Baselines3 provides ``SimpleMultiObsEnv`` as an example of this kind of of setting.
+The environment is a simple grid world but the observations for each cell come in the form of dictionaries.
+These dictionaries are randomly initilaized on the creation of the environment and contain a vector observation and an image observation.
+
+.. code-block:: python
+
+  from stable_baselines3 import PPO
+  from stable_baselines3.common.envs import SimpleMultiObsEnv
+
+
+  # Stable Baselines provides SimpleMultiObsEnv as an example environment with Dict observations
+  env = SimpleMultiObsEnv(random_start=False)
+
+  model = PPO("MultiInputPolicy", env, verbose=1)
+  model.learn(total_timesteps=1e5)
+
 
 Using Callback: Monitoring Training
 -----------------------------------
@@ -375,7 +396,7 @@ The parking env is a goal-conditioned continuous control task, in which the vehi
   import highway_env
   import numpy as np
 
-  from stable_baselines3 import HER, SAC, DDPG, TD3
+  from stable_baselines3 import HerReplayBuffer, SAC, DDPG, TD3
   from stable_baselines3.common.noise import NormalActionNoise
 
   env = gym.make("parking-v0")
@@ -384,21 +405,23 @@ The parking env is a goal-conditioned continuous control task, in which the vehi
   n_sampled_goal = 4
 
   # SAC hyperparams:
-  model = HER(
-      "MlpPolicy",
+  model = SAC(
+      "MultiInputPolicy",
       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,
+      replay_buffer_class=HerReplayBuffer,
+      replay_buffer_kwargs=dict(
+        n_sampled_goal=n_sampled_goal,
+        goal_selection_strategy="future",
+        # IMPORTANT: because the env is not wrapped with a TimeLimit wrapper
+        # we have to manually specify the max number of steps per episode
+        max_episode_length=100,
+        online_sampling=True,
+      )
       verbose=1,
       buffer_size=int(1e6),
       learning_rate=1e-3,
       gamma=0.95,
       batch_size=256,
-      online_sampling=True,
       policy_kwargs=dict(net_arch=[256, 256, 256]),
   )
 
@@ -408,7 +431,7 @@ The parking env is a goal-conditioned continuous control task, in which the vehi
   # Load saved model
   # Because it needs access to `env.compute_reward()`
   # HER must be loaded with the env
-  model = HER.load("her_sac_highway", env=env)
+  model = SAC.load("her_sac_highway", env=env)
 
   obs = env.reset()
 
@@ -658,12 +681,14 @@ to keep track of the agent progress.
 
   # ProcgenEnv is already vectorized
   venv = ProcgenEnv(num_envs=2, env_name='starpilot')
-  # PPO does not currently support Dict observations
-  # this will be solved in https://github.com/DLR-RM/stable-baselines3/pull/243
-  venv = VecExtractDictObs(venv, "rgb")
+
+  # To use only part of the observation:
+  # venv = VecExtractDictObs(venv, "rgb")
+
+  # Wrap with a VecMonitor to collect stats and avoid errors
   venv = VecMonitor(venv=venv)
 
-  model = PPO("MlpPolicy", venv, verbose=1)
+  model = PPO("MultiInputPolicy", venv, verbose=1)
   model.learn(10000)
 
 

diff --git a/docs/guide/vec_envs.rst b/docs/guide/vec_envs.rst
@@ -71,6 +71,17 @@ VecFrameStack
 .. autoclass:: VecFrameStack
   :members:
 
+StackedObservations
+~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: stable_baselines3.common.vec_env.stacked_observations.StackedObservations
+  :members:
+
+StackedDictObservations
+~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: stable_baselines3.common.vec_env.stacked_observations.StackedDictObservations
+  :members:
 
 VecNormalize
 ~~~~~~~~~~~~

diff --git a/docs/index.rst b/docs/index.rst
@@ -75,6 +75,7 @@ Main Features
 
   common/atari_wrappers
   common/env_util
+  common/envs
   common/distributions
   common/evaluation
   common/env_checker

diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst
@@ -4,23 +4,45 @@ Changelog
 ==========
 
 
-Release 1.1.0a5 (WIP)
+Release 1.1.0a6 (WIP)
 ---------------------------
 
+**Dict observation support, timeout handling and refactored HER**
+
 Breaking Changes:
 ^^^^^^^^^^^^^^^^^
+- All customs environments (e.g. the ``BitFlippingEnv`` or ``IdentityEnv``) were moved to ``stable_baselines3.common.envs`` folder
+- Refactored ``HER`` which is now the ``HerReplayBuffer`` class that can be passed to any off-policy algorithm
+- Handle timeout termination properly for off-policy algorithms (when using ``TimeLimit``)
 - Renamed ``_last_dones`` and ``dones`` to ``_last_episode_starts`` and ``episode_starts`` in ``RolloutBuffer``.
+- Removed ``ObsDictWrapper`` as ``Dict`` observation spaces are now supported
+
+.. code-block:: python
+
+  her_kwargs = dict(n_sampled_goal=2, goal_selection_strategy="future", online_sampling=True)
+  # SB3 < 1.1.0
+  # model = HER("MlpPolicy", env, model_class=SAC, **her_kwargs)
+  # SB3 >= 1.1.0:
+  model = SAC("MultiInputPolicy", env, replay_buffer_class=HerReplayBuffer, replay_buffer_kwargs=her_kwargs)
+
 - Updated the KL Divergence estimator in the PPO algorithm to be positive definite and have lower variance (@09tangriro)
 - Updated the KL Divergence check in the PPO algorithm to be before the gradient update step rather than after end of epoch (@09tangriro)
 
 New Features:
 ^^^^^^^^^^^^^
+- Added support for single-level ``Dict`` observation space (@JadenTravnik)
+- Added ``DictRolloutBuffer`` ``DictReplayBuffer`` to support dictionary observations (@JadenTravnik)
+- Added ``StackedObservations`` and ``StackedDictObservations`` that are used within ``VecFrameStack``
+- Added simple 4x4 room Dict test environments
+- ``HerReplayBuffer`` now supports ``VecNormalize`` when ``online_sampling=False``
 - Added `VecMonitor <https://github.com/DLR-RM/stable-baselines3/blob/master/stable_baselines3/common/vec_env/vec_monitor.py>`_ and
   `VecExtractDictObs <https://github.com/DLR-RM/stable-baselines3/blob/master/stable_baselines3/common/vec_env/vec_extract_dict_obs.py>`_ wrappers
   to handle gym3-style vectorized environments (@vwxyzjn)
 - Ignored the terminal observation if the it is not provided by the environment
   such as the gym3-style vectorized environments. (@vwxyzjn)
 - Add policy_base as input to the OnPolicyAlgorithm for more flexibility (@09tangriro)
+- Added support for image observation when using ``HER``
+- Added ``replay_buffer_class`` and ``replay_buffer_kwargs`` arguments to off-policy algorithms
 
 Bug Fixes:
 ^^^^^^^^^^
@@ -34,9 +56,9 @@ Deprecations:
 Others:
 ^^^^^^^
 - Added ``flake8-bugbear`` to tests dependencies to find likely bugs
+- Updated ``env_checker`` to reflect support of dict observation spaces
 - Added Code of Conduct
 - Added tests for GAE and lambda return computation
-- Updated docker image with newest black version
 
 Documentation:
 ^^^^^^^^^^^^^^
@@ -71,6 +93,7 @@ New Features:
 - Added support for ``custom_objects`` when loading models
 
 
+
 Bug Fixes:
 ^^^^^^^^^^
 - Fixed a bug with ``DQN`` predict method when using ``deterministic=False`` with image space
@@ -81,10 +104,14 @@ Documentation:
 - Added new project using SB3: rl_reach (@PierreExeter)
 - Added note about slow-down when switching to PyTorch
 - Add a note on continual learning and resetting environment
+
+Others:
+^^^^^^^
 - Updated RL-Zoo to reflect the fact that is it more than a collection of trained agents
 - Added images to illustrate the training loop and custom policies (created with https://excalidraw.com/)
 - Updated the custom policy section
 
+
 Pre-Release 0.11.1 (2021-02-27)
 -------------------------------
 
@@ -132,6 +159,7 @@ New Features:
 - Added new wrappers to log images and matplotlib figures to tensorboard. (@zampanteymedio)
 - Add support for text records to ``Logger``. (@lorenz-h)
 
+
 Bug Fixes:
 ^^^^^^^^^^
 - Fixed bug where code added VecTranspose on channel-first image environments (thanks @qxcv)
@@ -657,5 +685,9 @@ And all the contributors:
 @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 @megan-klaiber @thisray
+<<<<<<< HEAD
+@tfederico @hn2 @LucasAlegre @AptX395 @zampanteymedio @JadenTravnik @decodyng @ardabbour @lorenz-h @mschweizer @lorepieri8 @vwxyzjn
+=======
 @tfederico @hn2 @LucasAlegre @AptX395 @zampanteymedio @decodyng @ardabbour @lorenz-h @mschweizer @lorepieri8 @vwxyzjn
+>>>>>>> master
 @ShangqunYu @PierreExeter @JacopoPan @ltbd78 @tom-doerr @Atlis @liusida @09tangriro