Completed refactoring

pytorch · vmoens · Nov 21, 2022 · Nov 15, 2022 · Nov 16, 2022 · Nov 18, 2022
commit 5aed5411d522f953af8f08a7db95e289bba681a5
diff --git a/test/_utils_internal.py b/test/_utils_internal.py
@@ -12,13 +12,42 @@
 import pytest
 import torch.cuda
 from tensordict.tensordict import TensorDictBase
-from torchrl._utils import seed_generator
+from torchrl._utils import seed_generator, implement_for
 from torchrl.envs import EnvBase
 
 
 # Specified for test_utils.py
 __version__ = "0.3"
 
+# Default versions of the environments.
+CARTPOLE_VERSIONED = "CartPole-v1"
+PENDULUM_VERSIONED = "Pendulum-v1"
+PONG_VERSIONED = "ALE/Pong-v5"
+HALFCHEETAH_VERSIONED = "HalfCheetah-v4"
+
+
+@implement_for("gym", None, "0.20.0")
+def _set_gym_environments():  # noqa: F811
+    global CARTPOLE_VERSIONED, PENDULUM_VERSIONED, PONG_VERSIONED, HALFCHEETAH_VERSIONED
+
+    PENDULUM_VERSIONED = "Pendulum-v0"
+    CARTPOLE_VERSIONED = "CartPole-v0"
+    PONG_VERSIONED = "Pong-v4"
+    HALFCHEETAH_VERSIONED = "HalfCheetah-v2"
+
+
+@implement_for("gym", "0.20.0", None)
+def _set_gym_environments():  # noqa: F811
+    global CARTPOLE_VERSIONED, PENDULUM_VERSIONED, PONG_VERSIONED, HALFCHEETAH_VERSIONED
+
+    CARTPOLE_VERSIONED = "CartPole-v1"
+    PENDULUM_VERSIONED = "Pendulum-v1"
+    PONG_VERSIONED = "ALE/Pong-v5"
+    HALFCHEETAH_VERSIONED = "HalfCheetah-v4"
+
+
+_set_gym_environments()
+
 
 def get_relative_path(curr_file, *path_components):
     return os.path.join(os.path.dirname(curr_file), *path_components)

diff --git a/test/test_collector.py b/test/test_collector.py
@@ -8,7 +8,7 @@
 import numpy as np
 import pytest
 import torch
-from _utils_internal import generate_seeds
+from _utils_internal import generate_seeds, PENDULUM_VERSIONED, PONG_VERSIONED
 from mocking_classes import (
     ContinuousActionVecMockEnv,
     DiscreteActionConvMockEnv,
@@ -42,22 +42,6 @@
     TensorDictModule,
 )
 
-if _has_gym:
-    import gym
-    from packaging import version
-
-    gym_version = version.parse(gym.__version__)
-    PENDULUM_VERSIONED = (
-        "Pendulum-v1" if gym_version > version.parse("0.20.0") else "Pendulum-v0"
-    )
-    PONG_VERSIONED = (
-        "ALE/Pong-v5" if gym_version > version.parse("0.20.0") else "Pong-v4"
-    )
-else:
-    # placeholders
-    PENDULUM_VERSIONED = "Pendulum-v1"
-    PONG_VERSIONED = "ALE/Pong-v5"
-
 # torch.set_default_dtype(torch.double)
 
 

diff --git a/test/test_env.py b/test/test_env.py
@@ -11,7 +11,13 @@
 import pytest
 import torch
 import yaml
-from _utils_internal import get_available_devices
+from _utils_internal import (
+    get_available_devices,
+    CARTPOLE_VERSIONED,
+    PENDULUM_VERSIONED,
+    PONG_VERSIONED,
+    HALFCHEETAH_VERSIONED,
+)
 from mocking_classes import (
     ActionObsMergeLinear,
     DiscreteActionConvMockEnv,
@@ -53,26 +59,6 @@
     import gym
 
     gym_version = version.parse(gym.__version__)
-    PENDULUM_VERSIONED = (
-        "Pendulum-v1" if gym_version > version.parse("0.20.0") else "Pendulum-v0"
-    )
-    CARTPOLE_VERSIONED = (
-        "CartPole-v1" if gym_version > version.parse("0.20.0") else "CartPole-v0"
-    )
-    PONG_VERSIONED = (
-        "ALE/Pong-v5" if gym_version > version.parse("0.20.0") else "Pong-v4"
-    )
-    HALFCHEETAH_VERSIONED = (
-        "HalfCheetah-v4" if gym_version > version.parse("0.20.0") else "HalfCheetah-v2"
-    )
-else:
-    # placeholder
-    gym_version = version.parse("0.0.1")
-
-    # placeholders
-    PENDULUM_VERSIONED = "Pendulum-v1"
-    CARTPOLE_VERSIONED = "CartPole-v1"
-    PONG_VERSIONED = "ALE/Pong-v5"
 
 try:
     this_dir = os.path.dirname(os.path.realpath(__file__))

diff --git a/test/test_libs.py b/test/test_libs.py
@@ -8,8 +8,13 @@
 import numpy as np
 import pytest
 import torch
-from _utils_internal import _test_fake_tensordict
-from _utils_internal import get_available_devices
+from _utils_internal import (
+    _test_fake_tensordict,
+    get_available_devices,
+    HALFCHEETAH_VERSIONED,
+    PONG_VERSIONED,
+    PENDULUM_VERSIONED,
+)
 from packaging import version
 from tensordict.tensordict import assert_allclose_td
 from torchrl._utils import implement_for
@@ -22,28 +27,16 @@
 from torchrl.envs.libs.gym import _has_gym, _is_from_pixels
 from torchrl.envs.libs.habitat import HabitatEnv, _has_habitat
 
-
-@implement_for("gym", None, "0.20")
-def _import_pixel_observation_wrapper():  # noqa: F811
-    from torchrl.envs.libs.utils import (
-        GymPixelObservationWrapper as PixelObservationWrapper,
-    )
-
-    return PixelObservationWrapper
-
-
-@implement_for("gym", "0.20", None)
-def _import_pixel_observation_wrapper():  # noqa: F811
-    from gym.wrappers.pixel_observation import PixelObservationWrapper
-
-    return PixelObservationWrapper
-
-
 if _has_gym:
     import gym
 
     gym_version = version.parse(gym.__version__)
-    PixelObservationWrapper = _import_pixel_observation_wrapper()
+    if gym_version > version.parse("0.19"):
+        from gym.wrappers.pixel_observation import PixelObservationWrapper
+    else:
+        from torchrl.envs.libs.utils import (
+            GymPixelObservationWrapper as PixelObservationWrapper,
+        )
 
 if _has_dmc:
     from dm_control import suite
@@ -52,42 +45,12 @@ def _import_pixel_observation_wrapper():  # noqa: F811
 IS_OSX = platform == "darwin"
 
 
-@implement_for("gym", None, "0.21")
-def _pendulum_versioned():  # noqa: F811
-    return "Pendulum-v0"
-
-
-@implement_for("gym", "0.21", None)
-def _pendulum_versioned():  # noqa: F811
-    return "Pendulum-v1"
-
-
-@implement_for("gym", None, "0.21")
-def _hc_versioned():  # noqa: F811
-    return "HalfCheetah-v2"
-
-
-@implement_for("gym", "0.21", None)
-def _hc_versioned():  # noqa: F811
-    return "HalfCheetah-v4"
-
-
-@implement_for("gym", None, "0.21")
-def _pong_versioned():  # noqa: F811
-    return "Pong-v4"
-
-
-@implement_for("gym", "0.21", None)
-def _pong_versioned():  # noqa: F811
-    return "ALE/Pong-v5"
-
-
 @pytest.mark.skipif(not _has_gym, reason="no gym library found")
 @pytest.mark.parametrize(
     "env_name",
     [
-        _pong_versioned(),
-        _pendulum_versioned(),
+        PONG_VERSIONED,
+        PENDULUM_VERSIONED,
     ],
 )
 @pytest.mark.parametrize("frame_skip", [1, 3])
@@ -101,10 +64,10 @@ def _pong_versioned():  # noqa: F811
 )
 class TestGym:
     def test_gym(self, env_name, frame_skip, from_pixels, pixels_only):
-        if env_name == _pong_versioned() and not from_pixels:
+        if env_name == PONG_VERSIONED and not from_pixels:
             raise pytest.skip("already pixel")
         elif (
-            env_name != _pong_versioned()
+            env_name != PONG_VERSIONED
             and from_pixels
             and (not torch.has_cuda or not torch.cuda.device_count())
         ):
@@ -135,7 +98,7 @@ def test_gym(self, env_name, frame_skip, from_pixels, pixels_only):
         final_seed0, final_seed1 = final_seed
         assert final_seed0 == final_seed1
 
-        if env_name == _pong_versioned():
+        if env_name == PONG_VERSIONED:
             base_env = gym.make(env_name, frameskip=frame_skip)
             frame_skip = 1
         else:
@@ -159,10 +122,10 @@ def test_gym(self, env_name, frame_skip, from_pixels, pixels_only):
         assert_allclose_td(tdrollout[0], rollout2, rtol=1e-4, atol=1e-4)
 
     def test_gym_fake_td(self, env_name, frame_skip, from_pixels, pixels_only):
-        if env_name == _pong_versioned() and not from_pixels:
+        if env_name == PONG_VERSIONED and not from_pixels:
             raise pytest.skip("already pixel")
         elif (
-            env_name != _pong_versioned()
+            env_name != PONG_VERSIONED
             and from_pixels
             and (not torch.has_cuda or not torch.cuda.device_count())
         ):
@@ -283,10 +246,10 @@ def test_faketd(self, env_name, task, frame_skip, from_pixels, pixels_only):
     "env_lib,env_args,env_kwargs",
     [
         [DMControlEnv, ("cheetah", "run"), {"from_pixels": True}],
-        [GymEnv, (_hc_versioned(),), {"from_pixels": True}],
+        [GymEnv, (HALFCHEETAH_VERSIONED,), {"from_pixels": True}],
         [DMControlEnv, ("cheetah", "run"), {"from_pixels": False}],
-        [GymEnv, (_hc_versioned(),), {"from_pixels": False}],
-        [GymEnv, (_pong_versioned(),), {}],
+        [GymEnv, (HALFCHEETAH_VERSIONED,), {"from_pixels": False}],
+        [GymEnv, (PONG_VERSIONED,), {}],
     ],
 )
 def test_td_creation_from_spec(env_lib, env_args, env_kwargs):
@@ -320,10 +283,10 @@ def test_td_creation_from_spec(env_lib, env_args, env_kwargs):
     "env_lib,env_args,env_kwargs",
     [
         [DMControlEnv, ("cheetah", "run"), {"from_pixels": True}],
-        [GymEnv, (_hc_versioned(),), {"from_pixels": True}],
+        [GymEnv, (HALFCHEETAH_VERSIONED,), {"from_pixels": True}],
         [DMControlEnv, ("cheetah", "run"), {"from_pixels": False}],
-        [GymEnv, (_hc_versioned(),), {"from_pixels": False}],
-        [GymEnv, (_pong_versioned(),), {}],
+        [GymEnv, (HALFCHEETAH_VERSIONED,), {"from_pixels": False}],
+        [GymEnv, (PONG_VERSIONED,), {}],
     ],
 )
 @pytest.mark.parametrize("device", get_available_devices())

diff --git a/test/test_transforms.py b/test/test_transforms.py
@@ -8,7 +8,12 @@
 import numpy as np
 import pytest
 import torch
-from _utils_internal import get_available_devices, retry, dtype_fixture  # noqa
+from _utils_internal import (  # noqa
+    get_available_devices,
+    retry,
+    dtype_fixture,
+    PENDULUM_VERSIONED,
+)
 from mocking_classes import (
     ContinuousActionVecMockEnv,
     DiscreteActionConvMockEnvNumpy,
@@ -59,18 +64,6 @@
 )
 from torchrl.envs.transforms.vip import _VIPNet, VIPRewardTransform
 
-if _has_gym:
-    import gym
-    from packaging import version
-
-    gym_version = version.parse(gym.__version__)
-    PENDULUM_VERSIONED = (
-        "Pendulum-v1" if gym_version > version.parse("0.20.0") else "Pendulum-v0"
-    )
-else:
-    # placeholders
-    PENDULUM_VERSIONED = "Pendulum-v1"
-
 TIMEOUT = 10.0
 
 

diff --git a/torchrl/envs/libs/gym.py b/torchrl/envs/libs/gym.py
@@ -26,7 +26,6 @@
 
 try:
     import gym
-    from packaging import version
 
     _has_gym = True
 except ImportError:
@@ -49,9 +48,6 @@
         from torchrl.envs.libs.utils import (
             GymPixelObservationWrapper as PixelObservationWrapper,
         )
-    gym_version = version.parse(gym.__version__)
-    if gym_version >= version.parse("0.26.0"):
-        from gym.wrappers.compatibility import EnvCompatibility
 
 __all__ = ["GymWrapper", "GymEnv"]
 
@@ -202,6 +198,8 @@ def _build_gym_env(self, env, pixels_only):  # noqa: F811
 
     @implement_for("gym", "0.26.0", None)
     def _build_gym_env(self, env, pixels_only):  # noqa: F811
+        from gym.wrappers.compatibility import EnvCompatibility
+
         if env.render_mode:
             return PixelObservationWrapper(env, pixels_only=pixels_only)