[Refactoring] Replace gym version checking with decorated functions (#)

Initial commit. Only tests.
pytorch · vmoens · Nov 21, 2022 · Nov 15, 2022 · Nov 16, 2022 · Nov 18, 2022
commit feea25d2b234dccce2db5527f724db0a002f61bb
diff --git a/test/smoke_test_deps.py b/test/smoke_test_deps.py
@@ -2,21 +2,10 @@
 import tempfile
 
 import pytest
+from torchrl._utils import implement_for
 from torchrl.envs.libs.dm_control import _has_dmc, DMControlEnv
 from torchrl.envs.libs.gym import _has_gym, GymEnv
 
-if _has_gym:
-    import gym
-    from packaging import version
-
-    gym_version = version.parse(gym.__version__)
-    PONG_VERSIONED = (
-        "ALE/Pong-v5" if gym_version > version.parse("0.20.0") else "Pong-v4"
-    )
-else:
-    # placeholders
-    PONG_VERSIONED = "ALE/Pong-v5"
-
 try:
     from torch.utils.tensorboard import SummaryWriter
 
@@ -46,7 +35,7 @@ def test_gym():
     import gym  # noqa: F401
 
     assert _has_gym
-    env = GymEnv(PONG_VERSIONED)
+    env = GymEnv(_pong_versioned())
     env.reset()
 
 
@@ -65,6 +54,16 @@ def test_tb():
             test_rounds -= 1
 
 
+@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"
+
+
 if __name__ == "__main__":
     args, unknown = argparse.ArgumentParser().parse_known_args()
     pytest.main([__file__, "--capture", "no", "--exitfirst"] + unknown)
diff --git a/test/test_libs.py b/test/test_libs.py
@@ -3,74 +3,91 @@
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 import argparse
+from sys import platform
 
 import numpy as np
 import pytest
 import torch
 from _utils_internal import _test_fake_tensordict
 from _utils_internal import get_available_devices
 from packaging import version
+from tensordict.tensordict import assert_allclose_td
+from torchrl._utils import implement_for
 from torchrl.collectors import MultiaSyncDataCollector
 from torchrl.collectors.collectors import RandomPolicy
+from torchrl.envs import EnvCreator, ParallelEnv
+from torchrl.envs.libs.dm_control import DMControlEnv, DMControlWrapper
 from torchrl.envs.libs.dm_control import _has_dmc
+from torchrl.envs.libs.gym import GymEnv, GymWrapper
 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__)
-    if gym_version > version.parse("0.19"):
-        from gym.wrappers.pixel_observation import PixelObservationWrapper
-    else:
-        from torchrl.envs.libs.utils import (
-            GymPixelObservationWrapper as PixelObservationWrapper,
-        )
+    PixelObservationWrapper = _import_pixel_observation_wrapper()
 
 if _has_dmc:
     from dm_control import suite
     from dm_control.suite.wrappers import pixels
 
-from sys import platform
+IS_OSX = platform == "darwin"
 
-from tensordict.tensordict import assert_allclose_td
-from torchrl.envs import EnvCreator, ParallelEnv
-from torchrl.envs.libs.dm_control import DMControlEnv, DMControlWrapper
-from torchrl.envs.libs.gym import GymEnv, GymWrapper
 
-IS_OSX = platform == "darwin"
+@implement_for("gym", None, "0.21")
+def _pendulum_versioned():  # noqa: F811
+    return "Pendulum-v0"
 
-if _has_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"
-    )
-    HC_VERSIONED = (
-        "HalfCheetah-v4" if gym_version > version.parse("0.20.0") else "HalfCheetah-v2"
-    )
-    PONG_VERSIONED = (
-        "ALE/Pong-v5" if gym_version > version.parse("0.20.0") else "Pong-v4"
-    )
+@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"
 
-    # if gym_version < version.parse("0.24.0") and torch.cuda.device_count() > 0:
-    #     from opengl_rendering import create_opengl_context
-    #
-    #     create_opengl_context()
-else:
-    # placeholders
-    PENDULUM_VERSIONED = "Pendulum-v1"
-    HC_VERSIONED = "HalfCheetah-v4"
-    PONG_VERSIONED = "ALE/Pong-v5"
+
+@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])
@@ -84,10 +101,10 @@
 )
 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())
         ):
@@ -118,14 +135,11 @@ 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:
-            if gym_version < version.parse("0.26.0"):
-                base_env = gym.make(env_name)
-            else:
-                base_env = gym.make(env_name, render_mode="rgb_array")
+            base_env = _make_gym_environment(env_name)
 
         if from_pixels and not _is_from_pixels(base_env):
             base_env = PixelObservationWrapper(base_env, pixels_only=pixels_only)
@@ -145,10 +159,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())
         ):
@@ -269,10 +283,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, (_hc_versioned(),), {"from_pixels": True}],
         [DMControlEnv, ("cheetah", "run"), {"from_pixels": False}],
-        [GymEnv, (HC_VERSIONED,), {"from_pixels": False}],
-        [GymEnv, (PONG_VERSIONED,), {}],
+        [GymEnv, (_hc_versioned(),), {"from_pixels": False}],
+        [GymEnv, (_pong_versioned(),), {}],
     ],
 )
 def test_td_creation_from_spec(env_lib, env_args, env_kwargs):
@@ -303,10 +317,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, (_hc_versioned(),), {"from_pixels": True}],
         [DMControlEnv, ("cheetah", "run"), {"from_pixels": False}],
-        [GymEnv, (HC_VERSIONED,), {"from_pixels": False}],
-        [GymEnv, (PONG_VERSIONED,), {}],
+        [GymEnv, (_hc_versioned(),), {"from_pixels": False}],
+        [GymEnv, (_pong_versioned(),), {}],
     ],
 )
 @pytest.mark.parametrize("device", get_available_devices())
@@ -351,6 +365,16 @@ def test_habitat(self, envname):
         _test_fake_tensordict(env)
 
 
+@implement_for("gym", None, "0.26")
+def _make_gym_environment(env_name):  # noqa: F811
+    return gym.make(env_name)
+
+
+@implement_for("gym", "0.26", None)
+def _make_gym_environment(env_name):  # noqa: F811
+    return gym.make(env_name, render_mode="rgb_array")
+
+
 if __name__ == "__main__":
     args, unknown = argparse.ArgumentParser().parse_known_args()
     pytest.main([__file__, "--capture", "no", "--exitfirst"] + unknown)