Add testing

gymnasium/wrappers/atari_preprocessing.py

@@ -105,7 +105,7 @@ def __init__(
             isinstance(screen_size, tuple)
             and len(screen_size) == 2
             and all(isinstance(size, int) and size > 0 for size in screen_size)
-        )
+        ), f"Expect the `screen_size` to be positive, actually: {screen_size}"
         assert noop_max >= 0
         if frame_skip > 1 and getattr(env.unwrapped, "_frameskip", None) != 1:
             raise ValueError(
@@ -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 = screen_size + (1 if grayscale_obs else 3,)
+        _shape = self.screen_size + (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

@@ -1,5 +1,7 @@
 """Test suite for AtariProcessing wrapper."""
 
+import re
+
 import numpy as np
 import pytest
 
@@ -84,3 +86,24 @@ def test_atari_preprocessing_scale(grayscale, scaled, max_test_steps=10):
 
         step_i += 1
     env.close()
+
+
+def test_screen_size():
+    env = gym.make("ALE/Pong-v5", frameskip=1)
+
+    assert AtariPreprocessing(env).screen_size == (84, 84)
+    assert AtariPreprocessing(env, screen_size=50).screen_size == (50, 50)
+    assert AtariPreprocessing(env, screen_size=(100, 120)).screen_size == (100, 120)
+
+    with pytest.raises(
+        AssertionError, match="Expect the `screen_size` to be positive, actually: -1"
+    ):
+        AtariPreprocessing(env, screen_size=-1)
+
+    with pytest.raises(
+        AssertionError,
+        match=re.escape("Expect the `screen_size` to be positive, actually: (-1, 10)"),
+    ):
+        AtariPreprocessing(env, screen_size=(-1, 10))
+
+    env.close()