Fix discrete obs support (#296)

* Fixed discrete obs support

* Suggest new edit, fix failed test

* Revert "Suggest new edit, fix failed test"

This reverts commit 6892bf05506bb5ad0e87016d8d382705ab72e6a4.

* Fix test

* Special case for discrete obs

Co-authored-by: Anssi "Miffyli" Kanervisto <kaneran21@hotmail.com>

docs/misc/changelog.rst

@@ -3,7 +3,7 @@
 Changelog
 ==========
 
-Pre-Release 0.11.0a5 (WIP)
+Pre-Release 0.11.0a6 (WIP)
 -------------------------------
 
 Breaking Changes:
@@ -37,7 +37,7 @@ Bug Fixes:
 - Fixed bug where replay buffer could not be saved if it was too big (> 4 Gb) for python<3.8 (thanks @hn2)
 - Added informative ``PPO`` construction error in edge-case scenario where ``n_steps * n_envs = 1`` (size of rollout buffer),
   which otherwise causes downstream breaking errors in training (@decodyng)
-
+- Fixed discrete observation space support when using multiple envs with A2C/PPO (thanks @ardabbour)
 
 Deprecations:
 ^^^^^^^^^^^^^
@@ -52,6 +52,7 @@ Others:
 - Updated docker base image to Ubuntu 18.04
 - Set tensorboard min version to 2.2.0 (earlier version are apparently not working with PyTorch)
 - Added warning for ``PPO`` when ``n_steps * n_envs`` is not a multiple of ``batch_size`` (last mini-batch truncated) (@decodyng)
+- Removed some warnings in the tests
 
 Documentation:
 ^^^^^^^^^^^^^^
@@ -539,4 +540,4 @@ 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
-@tfederico @hn2 @LucasAlegre @AptX395 @zampanteymedio @decodyng
+@tfederico @hn2 @LucasAlegre @AptX395 @zampanteymedio @decodyng @ardabbour

stable_baselines3/common/buffers.py

@@ -358,6 +358,11 @@ class RolloutBuffer(BaseBuffer):
             # Reshape 0-d tensor to avoid error
             log_prob = log_prob.reshape(-1, 1)
 
+        # Reshape needed when using multiple envs with discrete observations
+        # as numpy cannot broadcast (n_discrete,) to (n_discrete, 1)
+        if isinstance(self.observation_space, spaces.Discrete):
+            obs = obs.reshape((self.n_envs,) + self.obs_shape)
+
         self.observations[self.pos] = np.array(obs).copy()
         self.actions[self.pos] = np.array(action).copy()
         self.rewards[self.pos] = np.array(reward).copy()

stable_baselines3/version.txt

@@ -1 +1 @@
-0.11.0a5
+0.11.0a6

tests/test_cnn.py

@@ -21,7 +21,7 @@ def test_cnn(tmp_path, model_class):
     # to check that the network handle it automatically
     env = FakeImageEnv(screen_height=40, screen_width=40, n_channels=1, discrete=model_class not in {SAC, TD3})
     if model_class in {A2C, PPO}:
-        kwargs = dict(n_steps=100)
+        kwargs = dict(n_steps=64)
     else:
         # Avoid memory error when using replay buffer
         # Reduce the size of the features

tests/test_custom_policy.py

@@ -19,7 +19,7 @@ from stable_baselines3.common.sb2_compat.rmsprop_tf_like import RMSpropTFLike
 )
 @pytest.mark.parametrize("model_class", [A2C, PPO])
 def test_flexible_mlp(model_class, net_arch):
-    _ = model_class("MlpPolicy", "CartPole-v1", policy_kwargs=dict(net_arch=net_arch), n_steps=100).learn(300)
+    _ = model_class("MlpPolicy", "CartPole-v1", policy_kwargs=dict(net_arch=net_arch), n_steps=64).learn(300)
 
 
 @pytest.mark.parametrize("net_arch", [[], [4], [4, 4], dict(qf=[8], pi=[8, 4])])
@@ -35,7 +35,7 @@ def test_custom_optimizer(model_class, optimizer_kwargs):
     if model_class in {DQN, SAC, TD3}:
         kwargs = dict(learning_starts=100)
     elif model_class in {A2C, PPO}:
-        kwargs = dict(n_steps=100)
+        kwargs = dict(n_steps=64)
 
     policy_kwargs = dict(optimizer_class=th.optim.AdamW, optimizer_kwargs=optimizer_kwargs, net_arch=[32])
     _ = model_class("MlpPolicy", "Pendulum-v0", policy_kwargs=policy_kwargs, **kwargs).learn(300)

tests/test_distributions.py

@@ -37,7 +37,7 @@ def test_squashed_gaussian(model_class):
     """
     Test run with squashed Gaussian (notably entropy computation)
     """
-    model = model_class("MlpPolicy", "Pendulum-v0", use_sde=True, n_steps=100, policy_kwargs=dict(squash_output=True))
+    model = model_class("MlpPolicy", "Pendulum-v0", use_sde=True, n_steps=64, policy_kwargs=dict(squash_output=True))
     model.learn(500)
 
     gaussian_mean = th.rand(N_SAMPLES, N_ACTIONS)

tests/test_save_load.py

@@ -178,7 +178,7 @@ def test_set_env(model_class):
     if model_class in {DQN, DDPG, SAC, TD3}:
         kwargs = dict(learning_starts=100)
     elif model_class in {A2C, PPO}:
-        kwargs = dict(n_steps=100)
+        kwargs = dict(n_steps=64)
 
     # create model
     model = model_class("MlpPolicy", env, policy_kwargs=dict(net_arch=[16]), **kwargs)

tests/test_spaces.py

@@ -3,6 +3,7 @@ import numpy as np
 import pytest
 
 from stable_baselines3 import A2C, DDPG, DQN, PPO, SAC, TD3
+from stable_baselines3.common.env_util import make_vec_env
 from stable_baselines3.common.evaluation import evaluate_policy
 
 
@@ -66,3 +67,10 @@ def test_action_spaces(model_class, env):
     else:
         with pytest.raises(AssertionError):
             model_class("MlpPolicy", env)
+
+
+@pytest.mark.parametrize("model_class", [A2C, PPO])
+@pytest.mark.parametrize("env", ["Taxi-v3"])
+def test_discrete_obs_space(model_class, env):
+    env = make_vec_env(env, n_envs=2, seed=0)
+    model_class("MlpPolicy", env, n_steps=256).learn(500)