[Test] Fix MuJoCo macro shape and gym Atari setup (#3811)

Author: vmoens

.github/unittest/linux_libs/scripts_gym/run_all.sh

@@ -231,6 +231,7 @@ printf "* Testing gym 0.20\n"
 pip install "pip<24.1" setuptools==65.3.0 wheel==0.38.4
 pip install 'gym[atari]==0.20'
 pip install 'ale-py==0.7.4'
+ale-import-roms Roms
 run_tests "gym==0.20" || true
 pip uninstall -y gym ale-py wheel || true
 pip install --upgrade pip setuptools wheel  # restore latest versions

test/libs/test_mujoco.py

@@ -117,7 +117,9 @@ def test_satellite_generic_macro_sequence(self):
         sequence = transform.action_sequence(
             td, MacroPrimitive.MOVE, target_qpos=target
         )
-        assert sequence.shape == torch.Size([2, 4, env.action_spec.shape[-1]])
+        assert sequence.shape == td.batch_size + torch.Size(
+            [4, env.action_spec.shape[-1]]
+        )
         env.close()
 
     def test_humanoid_generic_macro_sequence(self):

torchrl/envs/libs/gym.py

@@ -75,6 +75,26 @@
 """
 
 
+def _patch_legacy_ale_py_gym_env(env_name: str) -> None:
+    return
+
+
+@implement_for("gym", "0.20.0", "0.21.0")
+def _patch_legacy_ale_py_gym_env(env_name: str) -> None:  # noqa: F811
+    """Expose the legacy ALE entry point used by gym 0.20 Atari specs."""
+    del env_name
+    try:
+        ale_gym = importlib.import_module("ale_py.gym")
+    except ImportError:
+        return
+    if hasattr(ale_gym, "ALGymEnv"):
+        return
+    try:
+        ale_gym.ALGymEnv = gym_backend("envs.atari").AtariEnv
+    except (AttributeError, ImportError):
+        return
+
+
 def _gymnasium_reward_space(env):
     reward_space = getattr(env, "__dict__", {}).get("reward_space", None)
     if reward_space is not None:
@@ -1986,6 +2006,7 @@ def _build_env(
                             torchrl_logger.warning(
                                 f"ale_py not found, this may cause issues with ALE environments: {err}"
                             )
+                    _patch_legacy_ale_py_gym_env(env_name)
                     # we catch warnings as they may cause silent bugs
                     env = self.lib.make(env_name, **kwargs)
                     if len(w) and "frameskip" in str(w[-1].message):