Remove _has_gym checks

lerobot/common/envs/aloha.py

@@ -17,7 +17,6 @@ from torchrl.data.tensor_specs import (
 from lerobot.common.envs.abstract import AbstractEnv
 from lerobot.common.utils import set_global_seed
 
-_has_gym = importlib.util.find_spec("gymnasium") is not None
 _has_aloha = importlib.util.find_spec("aloha") is not None
 
 
@@ -51,9 +50,6 @@ class AlohaEnv(AbstractEnv):
         )
 
     def _make_env(self):
-        if not _has_gym:
-            raise ImportError("Cannot import gymnasium.")
-
         if not self.from_pixels:
             raise NotImplementedError()
 

lerobot/common/envs/pusht.py

@@ -18,7 +18,6 @@ from torchrl.envs.libs.gym import _gym_to_torchrl_spec_transform
 from lerobot.common.envs.abstract import AbstractEnv
 from lerobot.common.utils import set_global_seed
 
-_has_gym = importlib.util.find_spec("gymnasium") is not None
 _has_pusht = importlib.util.find_spec("pusht") is not None
 
 
@@ -52,9 +51,6 @@ class PushtEnv(AbstractEnv):
         )
 
     def _make_env(self):
-        if not _has_gym:
-            raise ImportError("Cannot import gymnasium.")
-
         # TODO(rcadene) (PushTEnv is similar to PushTImageEnv, but without the image rendering, it's faster to iterate on)
         # from lerobot.common.envs.pusht.pusht_env import PushTEnv
         if not self.from_pixels:

lerobot/common/envs/xarm.py

@@ -53,15 +53,14 @@ class SimxarmEnv(AbstractEnv):
         )
 
     def _make_env(self):
-        if not _has_gym:
-            raise ImportError("Cannot import gymnasium.")
-
-        import gymnasium
-
         if not _has_xarm:
             raise ImportError(
                 "Cannot import xarm env. Please install it with `python -m pip install 'lerobot[xarm]'`"
             )
+        if not _has_gym:
+            raise ImportError("Cannot import gymnasium.")
+
+        import gymnasium
 
         from xarm import TASKS