[Cleanup] Removing gym-retro interface (#444)

Author: vmoens

README.md

@@ -457,7 +457,7 @@ pip3 install moviepy
 pip3 install dm_control 
 
 # gym, atari games
-pip3 install gym[atari] "gym[accept-rom-license]" pygame gym_retro
+pip3 install gym[atari] "gym[accept-rom-license]" pygame
 
 # tests
 pip3 install pytest pyyaml pytest-instafail

torchrl/envs/libs/gym.py

@@ -4,7 +4,7 @@
 # LICENSE file in the root directory of this source tree.
 import warnings
 from types import ModuleType
-from typing import List, Optional, Sequence, Dict
+from typing import List, Dict
 
 import torch
 from packaging import version
@@ -50,14 +50,8 @@
     gym_version = version.parse(gym.__version__)
     if gym_version >= version.parse("0.26.0"):
         from gym.wrappers.compatibility import EnvCompatibility
-try:
-    import retro
-
-    _has_retro = True
-except ImportError:
-    _has_retro = False
 
-__all__ = ["GymWrapper", "GymEnv", "RetroEnv"]
+__all__ = ["GymWrapper", "GymEnv"]
 
 
 def _gym_to_torchrl_spec_transform(spec, dtype=None, device="cpu") -> TensorSpec:
@@ -79,12 +73,15 @@ def _gym_to_torchrl_spec_transform(spec, dtype=None, device="cpu") -> TensorSpec
             dtype=dtype,
             device=device,
         )
-    elif isinstance(spec, (dict, gym.spaces.dict.Dict)):
-        spec = {
-            "next_" + k: _gym_to_torchrl_spec_transform(spec[k], device=device)
-            for k in spec
-        }
-        return CompositeSpec(**spec)
+    elif isinstance(spec, (Dict,)):
+        spec_out = {}
+        for k in spec.keys():
+            spec_out["next_" + k] = _gym_to_torchrl_spec_transform(
+                spec[k], device=device
+            )
+        return CompositeSpec(**spec_out)
+    elif isinstance(spec, gym.spaces.dict.Dict):
+        return _gym_to_torchrl_spec_transform(spec.spaces, device=device)
     else:
         raise NotImplementedError(
             f"spec of type {type(spec).__name__} is currently unaccounted for"
@@ -111,9 +108,12 @@ def _get_gym():
 
 def _is_from_pixels(env):
     observation_spec = env.observation_space
-    if isinstance(observation_spec, (Dict, gym.spaces.dict.Dict)):
+    if isinstance(observation_spec, (Dict,)):
         if "pixels" in set(observation_spec.keys()):
             return True
+    if isinstance(observation_spec, (gym.spaces.dict.Dict,)):
+        if "pixels" in set(observation_spec.spaces.keys()):
+            return True
     elif (
         isinstance(observation_spec, gym.spaces.Box)
         and (observation_spec.low == 0).all()
@@ -337,23 +337,3 @@ def _check_kwargs(self, kwargs: Dict):
 
     def __repr__(self) -> str:
         return f"{self.__class__.__name__}(env={self.env_name}, batch_size={self.batch_size}, device={self.device})"
-
-
-def _get_retro_envs() -> Sequence:
-    if not _has_retro:
-        return tuple()
-    else:
-        return retro.data.list_games()
-
-
-def _get_retro() -> Optional[ModuleType]:
-    if _has_retro:
-        return retro
-    else:
-        return None
-
-
-class RetroEnv(GymEnv):
-    available_envs = _get_retro_envs()
-    lib = "retro"
-    lib = _get_retro()

torchrl/trainers/helpers/envs.py

@@ -13,7 +13,7 @@
 from torchrl.envs.common import EnvBase
 from torchrl.envs.env_creator import env_creator, EnvCreator
 from torchrl.envs.libs.dm_control import DMControlEnv
-from torchrl.envs.libs.gym import GymEnv, RetroEnv
+from torchrl.envs.libs.gym import GymEnv
 from torchrl.envs.transforms import (
     CatFrames,
     CatTensors,
@@ -42,7 +42,6 @@
 
 LIBS = {
     "gym": GymEnv,
-    "retro": RetroEnv,
     "dm_control": DMControlEnv,
 }