Allow AtariPreprocessing non-square observations (#1312)

Author: li-plus

gymnasium/wrappers/atari_preprocessing.py

@@ -142,7 +142,7 @@ def __init__(
         self.game_over = False
 
         _low, _high, _dtype = (0, 1, np.float32) if scale_obs else (0, 255, np.uint8)
-        _shape = self.screen_size + (1 if grayscale_obs else 3,)
+        _shape = (self.screen_size[1], self.screen_size[0], 1 if grayscale_obs else 3)
         if grayscale_obs and not grayscale_newaxis:
             _shape = _shape[:-1]  # Remove channel axis
         self.observation_space = Box(low=_low, high=_high, shape=_shape, dtype=_dtype)

tests/wrappers/test_atari_preprocessing.py

@@ -47,6 +47,17 @@
             ),
             (84, 84, 1),
         ),
+        (
+            AtariPreprocessing(
+                gym.make("ALE/Pong-v5"),
+                screen_size=(160, 210),
+                grayscale_obs=False,
+                frame_skip=1,
+                noop_max=0,
+                grayscale_newaxis=True,
+            ),
+            (210, 160, 3),
+        ),
     ],
 )
 def test_atari_preprocessing_grayscale(env, expected_obs_shape):