fixed tests

Author: AdrianOrenstein

src/ale/python/vector_env.py

@@ -126,6 +126,12 @@ def __init__(
         self.full_action_space = full_action_space
         self.continuous = continuous
         self.continuous_action_threshold = continuous_action_threshold
+        self.map_action_idx = np.zeros((3, 3, 2), dtype=np.int32)
+        for h in (-1, 0, 1):
+            for v in (-1, 0, 1):
+                for f in (0, 1):
+                    action = AtariEnv.map_action_idx(h, v, bool(f)).value
+                    self.map_action_idx[h + 1, v + 1, f] = action
         self.single_action_space, self.action_space = (
             self._setup_continuous_action()
             if self.continuous
@@ -141,12 +147,6 @@ def _setup_obs(self, stack_num: int, img_height: int, img_width: int) -> tuple:
         return single, gymnasium.vector.utils.batch_space(single, self.batch_size)
 
     def _setup_continuous_action(self) -> tuple:
-        self.map_action_idx = np.zeros((3, 3, 2), dtype=np.int32)
-        for h in (-1, 0, 1):
-            for v in (-1, 0, 1):
-                for f in (0, 1):
-                    action = AtariEnv.map_action_idx(h, v, bool(f)).value
-                    self.map_action_idx[h + 1, v + 1, f] = action
         # Actions are radius, theta, and fire, where first two are the parameters of polar coordinates.
         single = Box(
             low=np.array([0.0, -np.pi, 0.0]).astype(np.float32),

tests/python/test_atari_vector_env.py

@@ -666,8 +666,12 @@ def test_same_step_autoreset_mode(
         gym_envs.close()
         ale_envs.close()
 
+
+class TestMultiRomVectorEnv:
+    """Tests for multi-ROM support in AtariVectorEnv."""
+
     @pytest.mark.parametrize("game,n_minimal", [("breakout", 4), ("pong", 6)])
-    def test_single_action_space(game, n_minimal):
+    def test_single_action_space(self, game, n_minimal):
         """single_action_space reflects the real action-set size, not a constant."""
         env = AtariVectorEnv(game=game, num_envs=3)
         assert env.single_action_space == gym.spaces.Discrete(n_minimal)
@@ -684,10 +688,6 @@ def test_single_action_space(game, n_minimal):
         assert np.all(env.action_space.nvec == 18)
         env.close()
 
-
-class TestMultiRomVectorEnv:
-    """Tests for multi-ROM support in AtariVectorEnv."""
-
     def test_single_action_space_none_for_different_roms(self):
         """Different ROMs (different action-set sizes) have no single space."""
         # num_envs repeats each ROM, so this is [pong, pong, breakout, breakout].