Add Gym 0.26 support

setup.py

@@ -115,6 +115,7 @@
         "extra": [
             # For render
             "opencv-python",
+            "pygame",
             # For atari games,
             "ale-py~=0.7.4",
             "autorom[accept-rom-license]~=0.4.2",

stable_baselines3/common/env_checker.py

@@ -176,8 +176,8 @@ def _check_returned_values(env: gym.Env, observation_space: spaces.Space, action
     assert isinstance(done, bool), "The `done` signal must be a boolean"
     assert isinstance(info, dict), "The `info` returned by `step()` must be a python dictionary"
 
-    if isinstance(env, gym.GoalEnv):
-        # For a GoalEnv, the keys are checked at reset
+    # Goal conditioned env
+    if hasattr(env, "compute_reward"):
         assert reward == env.compute_reward(obs["achieved_goal"], obs["desired_goal"], info)
 
 

stable_baselines3/common/envs/bit_flipping_env.py

@@ -2,13 +2,13 @@
 from typing import Any, Dict, Optional, Union
 
 import numpy as np
-from gym import GoalEnv, spaces
+from gym import Env, spaces
 from gym.envs.registration import EnvSpec
 
 from stable_baselines3.common.type_aliases import GymStepReturn
 
 
-class BitFlippingEnv(GoalEnv):
+class BitFlippingEnv(Env):
     """
     Simple bit flipping env, useful to test HER.
     The goal is to flip all the bits to get a vector of ones.

tests/test_envs.py

@@ -27,6 +27,8 @@
     SimpleMultiObsEnv,
 ]
 
+GYM_MESSAGE = "Function `rng.randint"
+
 
 @pytest.mark.parametrize("env_id", ["CartPole-v0", "Pendulum-v1"])
 def test_env(env_id):
@@ -44,17 +46,17 @@ def test_env(env_id):
     if env_id == "Pendulum-v1":
         assert len(record) == 1
     else:
-        # The other environments must pass without warning
-        assert len(record) == 0
+        # The other environments are expected to raise warning introduced by Gym.
+        check_rand_warning(record)
 
 
 @pytest.mark.parametrize("env_class", ENV_CLASSES)
 def test_custom_envs(env_class):
     env = env_class()
     with warnings.catch_warnings(record=True) as record:
         check_env(env)
-    # No warnings for custom envs
-    assert len(record) == 0
+    # Only randint warning coming from gym
+    check_rand_warning(record)
 
 
 @pytest.mark.parametrize(
@@ -72,8 +74,12 @@ def test_bit_flipping(kwargs):
     with warnings.catch_warnings(record=True) as record:
         check_env(env)
 
-    # No warnings for custom envs
-    assert len(record) == 0
+    # Only randint warning coming from gym
+    check_rand_warning(record)
+
+
+def check_rand_warning(record):
+    assert all(GYM_MESSAGE in warning.message.args[0] for warning in record)
 
 
 def test_high_dimension_action_space():
@@ -151,13 +157,19 @@ def test_non_default_action_spaces(new_action_space):
     with warnings.catch_warnings(record=True) as record:
         check_env(env)
 
-    # No warnings for custom envs
-    assert len(record) == 0
+    # Only randint warning coming from gym
+    check_rand_warning(record)
+
     # Change the action space
     env.action_space = new_action_space
 
-    with pytest.warns(UserWarning):
-        check_env(env)
+    # Gym raises error for Boxed spaces if low > high
+    if env.action_space.low[0] > env.action_space.high[0]:
+        with pytest.raises(ValueError):
+            check_env(env)
+    else:
+        with pytest.warns(UserWarning):
+            check_env(env)
 
 
 def check_reset_assert_error(env, new_reset_return):

tests/test_save_load.py

@@ -27,7 +27,7 @@ def select_env(model_class: BaseAlgorithm) -> gym.Env:
     if model_class == DQN:
         return IdentityEnv(10)
     else:
-        return IdentityEnvBox(10)
+        return IdentityEnvBox(10, 10)
 
 
 @pytest.mark.parametrize("model_class", MODEL_LIST)

tests/test_vec_normalize.py

@@ -41,7 +41,7 @@ def reset(self):
         return np.array([self.returned_rewards[self.return_reward_idx]])
 
 
-class DummyDictEnv(gym.GoalEnv):
+class DummyDictEnv(gym.Env):
     """
     Dummy gym goal env for testing purposes
     """