Code cleanup. Add tests.

gym/vector/async_vector_env.py

@@ -12,7 +12,6 @@
     AlreadyPendingCallError,
     ClosedEnvironmentError,
     CustomSpaceError,
-    InvalidInfoFormat,
     NoAsyncCallError,
 )
 from gym.vector.utils import (

gym/wrappers/__init__.py

@@ -1,4 +1,4 @@
-"""Module of wrapper classes."""
+from gym import error
 from gym.wrappers.atari_preprocessing import AtariPreprocessing
 from gym.wrappers.autoreset import AutoResetWrapper
 from gym.wrappers.clip_action import ClipAction
@@ -8,8 +8,12 @@
 from gym.wrappers.gray_scale_observation import GrayScaleObservation
 from gym.wrappers.normalize import NormalizeObservation, NormalizeReward
 from gym.wrappers.order_enforcing import OrderEnforcing
-from gym.wrappers.pixel_observation import PixelObservationWrapper
-from gym.wrappers.record_episode_statistics import RecordEpisodeStatistics
+from gym.wrappers.record_episode_statistics import (
+    BraxVecEnvStatsInfoStrategy,
+    ClassicVecEnvStatsInfoStrategy,
+    RecordEpisodeStatistics,
+    get_statistic_info_strategy,
+)
 from gym.wrappers.record_video import RecordVideo, capped_cubic_video_schedule
 from gym.wrappers.rescale_action import RescaleAction
 from gym.wrappers.resize_observation import ResizeObservation

gym/wrappers/record_episode_statistics.py

@@ -66,7 +66,7 @@ def get_statistic_info_strategy(wrapped_env_strategy: str):
     }
     if wrapped_env_strategy not in strategies:
         raise NoMatchingInfoStrategy(
-            "Wrapped environment has an info format of type %s which is not a processable format by this wrappers. Please use one in %s"
+            "Wrapped environment has an info format of type %s which is not a processable format by this wrapper. Please use one in %s"
             % (wrapped_env_strategy, list(strategies.keys()))
         )
     return strategies[wrapped_env_strategy]

tests/vector/test_info_format_strategies.py

@@ -1,40 +1,57 @@
 import pytest
 
-from gym.vector.utils import BraxVecEnvInfoStrategy, ClassicVecEnvInfoStrategy
-
-
-def test_classic_vec_env_info_strategy():
-    NUM_ENVS = 3
-
-    infos = ClassicVecEnvInfoStrategy(NUM_ENVS)
-    for i in range(NUM_ENVS):
+from gym.error import InvalidInfoFormat
+from gym.vector.utils import (
+    BraxVecEnvInfoStrategy,
+    ClassicVecEnvInfoStrategy,
+    get_info_strategy,
+)
+
+
+@pytest.mark.parametrize(("num_envs"), [3])
+def test_classic_vec_env_info_strategy(num_envs):
+    infos = ClassicVecEnvInfoStrategy(num_envs)
+    for i in range(num_envs):
         info = {"example_info": i}
         infos.add_info(info, i)
 
     expected_info = [{"example_info": 0}, {"example_info": 1}, {"example_info": 2}]
     assert expected_info == infos.get_info()
 
 
-def test_brax_vec_env_info_strategy():
+@pytest.mark.parametrize(("num_envs"), [3])
+def test_brax_vec_env_info_strategy(num_envs):
     NUM_ENVS = 3
 
-    infos = BraxVecEnvInfoStrategy(NUM_ENVS)
-    for i in range(NUM_ENVS):
+    infos = BraxVecEnvInfoStrategy(num_envs)
+    for i in range(num_envs):
         info = {"example_info": i}
         infos.add_info(info, i)
 
     expected_info = {"example_info": [0, 1, 2]}
     assert expected_info == infos.get_info()
 
 
-def test_brax_vec_env_info_strategy_with_nones():
-    NUM_ENVS = 5
-
-    infos = BraxVecEnvInfoStrategy(NUM_ENVS)
-    for i in range(NUM_ENVS):
+@pytest.mark.parametrize(("num_envs"), [5])
+def test_brax_vec_env_info_strategy_with_nones(num_envs):
+    infos = BraxVecEnvInfoStrategy(num_envs)
+    for i in range(num_envs):
         if i % 2 == 0:
             info = {"example_info": i}
             infos.add_info(info, i)
 
     expected_info = {"example_info": [0, None, 2, None, 4]}
     assert expected_info == infos.get_info()
+
+
+@pytest.mark.parametrize(("info_format"), [("classic"), ("brax"), ("non_existent")])
+def test_get_info_strategy(info_format):
+    if info_format == "classic":
+        InfoStrategy = get_info_strategy(info_format)
+        assert InfoStrategy == ClassicVecEnvInfoStrategy
+    elif info_format == "brax":
+        InfoStrategy = get_info_strategy(info_format)
+        assert InfoStrategy == BraxVecEnvInfoStrategy
+    else:
+        with pytest.raises(InvalidInfoFormat):
+            get_info_strategy(info_format)

tests/wrappers/test_record_episode_statistics.py

@@ -1,7 +1,13 @@
 import pytest
 
 import gym
-from gym.wrappers import RecordEpisodeStatistics
+from gym.error import NoMatchingInfoStrategy
+from gym.wrappers import (
+    BraxVecEnvStatsInfoStrategy,
+    ClassicVecEnvStatsInfoStrategy,
+    RecordEpisodeStatistics,
+    get_statistic_info_strategy,
+)
 
 
 @pytest.mark.parametrize("env_id", ["CartPole-v1", "Pendulum-v1"])
@@ -55,3 +61,38 @@ def test_record_episode_statistics_with_vectorenv(num_envs, asynchronous):
                 assert "episode" in info
                 assert all([item in info["episode"] for item in ["r", "l", "t"]])
                 break
+
+
+@pytest.mark.parametrize(("num_envs", "asynchronous"), [(3, False), (3, True)])
+def test_episode_statistics_brax_info(num_envs, asynchronous):
+    envs = gym.vector.make(
+        "CartPole-v1", asynchronous=asynchronous, num_envs=num_envs, info_format="brax"
+    )
+    envs = RecordEpisodeStatistics(envs)
+    envs.reset()
+    dones = [False for _ in range(num_envs)]
+    actions = np.array([1, 0, 1])
+    while not any(dones):
+        _, _, dones, infos = envs.step(actions)
+
+    assert "episode" in infos
+    assert len(infos["episode"]) == num_envs
+    assert "terminal_observation" in infos
+    for i in range(num_envs):
+        if dones[i]:
+            assert infos["terminal_observation"][i] is not None
+        else:
+            assert infos["terminal_observation"][i] is None
+
+
+@pytest.mark.parametrize(("info_format"), [("classic"), ("brax"), ("non_existent")])
+def test_get_statistic_info_strategy(info_format):
+    if info_format == "classic":
+        InfoStrategy = get_statistic_info_strategy(info_format)
+        assert InfoStrategy == ClassicVecEnvStatsInfoStrategy
+    elif info_format == "brax":
+        InfoStrategy = get_statistic_info_strategy(info_format)
+        assert InfoStrategy == BraxVecEnvStatsInfoStrategy
+    else:
+        with pytest.raises(NoMatchingInfoStrategy):
+            get_statistic_info_strategy(info_format)