Cleanup PR, bump minimal gymnasium version

Author: ffelten

mo_gymnasium/wrappers/vector/__init__.py

@@ -1,6 +1,7 @@
 """Vector wrappers."""
 
 from mo_gymnasium.wrappers.vector.wrappers import (
+    MOAsyncVectorEnv,
     MORecordEpisodeStatistics,
     MOSyncVectorEnv,
 )

mo_gymnasium/wrappers/vector/wrappers.py

@@ -1,27 +1,29 @@
 """Vector wrappers."""
 
-import os
+import multiprocessing
 import sys
 import time
 import traceback
 from copy import deepcopy
-from typing import Any, Callable, Dict, Iterator, Tuple, Sequence, Union
-
-import multiprocessing
 from multiprocessing import Array, Queue
 from multiprocessing.connection import Connection
-import numpy as np
+from typing import Any, Callable, Dict, Iterator, Sequence, Tuple, Union
 
 import gymnasium as gym
 import numpy as np
 from gymnasium.core import ActType, ObsType
-from gymnasium.vector import SyncVectorEnv, AsyncVectorEnv
-from gymnasium.vector.async_vector_env import AsyncState
-from gymnasium.vector.utils import concatenate, iterate, create_empty_array, write_to_shared_memory
-from gymnasium.vector.vector_env import ArrayType, VectorEnv, AutoresetMode
-from gymnasium.wrappers.vector import RecordEpisodeStatistics
 from gymnasium.error import NoAsyncCallError
 from gymnasium.spaces.utils import is_space_dtype_shape_equiv
+from gymnasium.vector import AsyncVectorEnv, SyncVectorEnv
+from gymnasium.vector.async_vector_env import AsyncState
+from gymnasium.vector.utils import (
+    concatenate,
+    create_empty_array,
+    iterate,
+    write_to_shared_memory,
+)
+from gymnasium.vector.vector_env import ArrayType, AutoresetMode, VectorEnv
+from gymnasium.wrappers.vector import RecordEpisodeStatistics
 
 
 class MOSyncVectorEnv(SyncVectorEnv):
@@ -112,7 +114,8 @@ def step(self, actions: ActType) -> Tuple[ObsType, ArrayType, ArrayType, ArrayTy
             np.copy(self._truncations),
             infos,
         )
-    
+
+
 def _mo_async_worker(
     index: int,
     env_fn: callable,
@@ -138,9 +141,7 @@ def _mo_async_worker(
             if command == "reset":
                 observation, info = env.reset(**data)
                 if shared_memory:
-                    write_to_shared_memory(
-                        observation_space, index, observation, shared_memory
-                    )
+                    write_to_shared_memory(observation_space, index, observation, shared_memory)
                     observation = None
                     autoreset = False
                 pipe.send(((observation, info), True))
@@ -150,7 +151,11 @@ def _mo_async_worker(
                 if autoreset_mode == AutoresetMode.NEXT_STEP:
                     if autoreset:
                         observation, info = env.reset()
-                        reward, terminated, truncated = np.zeros(reward_space.shape[0], dtype=np.float32), False, False
+                        reward, terminated, truncated = (
+                            np.zeros(reward_space.shape[0], dtype=np.float32),
+                            False,
+                            False,
+                        )
                     else:
                         (
                             observation,
@@ -191,9 +196,7 @@ def _mo_async_worker(
                     raise ValueError(f"Unexpected autoreset_mode: {autoreset_mode}")
 
                 if shared_memory:
-                    write_to_shared_memory(
-                        observation_space, index, observation, shared_memory
-                    )
+                    write_to_shared_memory(observation_space, index, observation, shared_memory)
                     observation = None
 
                 pipe.send(((observation, reward, terminated, truncated, info), True))
@@ -203,9 +206,7 @@ def _mo_async_worker(
             elif command == "_call":
                 name, args, kwargs = data
                 if name in ["reset", "step", "close", "_setattr", "_check_spaces"]:
-                    raise ValueError(
-                        f"Trying to call function `{name}` with `call`, use `{name}` directly instead."
-                    )
+                    raise ValueError(f"Trying to call function `{name}` with `call`, use `{name}` directly instead.")
 
                 attr = env.get_wrapper_attr(name)
                 if callable(attr):
@@ -225,9 +226,7 @@ def _mo_async_worker(
                             (
                                 single_obs_space == observation_space
                                 if obs_mode == "same"
-                                else is_space_dtype_shape_equiv(
-                                    single_obs_space, observation_space
-                                )
+                                else is_space_dtype_shape_equiv(single_obs_space, observation_space)
                             ),
                             single_action_space == action_space,
                         ),
@@ -246,14 +245,15 @@ def _mo_async_worker(
         pipe.send((None, False))
     finally:
         env.close()
-    
+
+
 class MOAsyncVectorEnv(AsyncVectorEnv):
     """Vectorized environment that runs multiple environments in parallel.
 
     It uses ``multiprocessing`` processes, and pipes for communication.
 
-    Mofified from gymnasium.vector.async_vector_env.AsyncVectorEnv to allow for multi-objective rewards.
-    
+    Modified from gymnasium.vector.async_vector_env.AsyncVectorEnv to allow for multi-objective rewards.
+
     Example:
         >>> import mo_gymnasium as mo_gym
         >>> envs = mo_gym.wrappers.vector.MOAsyncVectorEnv([
@@ -274,11 +274,13 @@ class MOAsyncVectorEnv(AsyncVectorEnv):
         >>> terminateds
         array([False,  True, False, False])
     """
-    def __init__(
-        self,
-        env_fns: Sequence[Callable[[], gym.Env]],
-        **kwargs
-    ):
+
+    def __init__(self, env_fns: Sequence[Callable[[], gym.Env]], **kwargs):
+        """Vectorized environment that runs multiple environments in parallel.
+
+        Args:
+            env_fns: env constructors
+        """
         super().__init__(env_fns=env_fns, worker=_mo_async_worker, **kwargs)
 
         # extract reward space from first vector env and create 2d array to store vector rewards
@@ -288,10 +290,7 @@ def __init__(
         del dummy_env
         self.rewards = create_empty_array(self.reward_space, n=self.num_envs, fn=np.zeros)
 
-    
-    def step_wait(
-        self, timeout: int | float | None = None
-    ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, dict]:
+    def step_wait(self, timeout: int | float | None = None) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, dict]:
         """Wait for the calls to :obj:`step` in each sub-environment to finish.
 
         Args:
@@ -308,15 +307,13 @@ def step_wait(
         self._assert_is_running()
         if self._state != AsyncState.WAITING_STEP:
             raise NoAsyncCallError(
-                "Calling `step_wait` without any prior call " "to `step_async`.",
+                "Calling `step_wait` without any prior call to `step_async`.",
                 AsyncState.WAITING_STEP.value,
             )
 
         if not self._poll_pipe_envs(timeout):
             self._state = AsyncState.DEFAULT
-            raise multiprocessing.TimeoutError(
-                f"The call to `step_wait` has timed out after {timeout} second(s)."
-            )
+            raise multiprocessing.TimeoutError(f"The call to `step_wait` has timed out after {timeout} second(s).")
 
         observations, rewards, terminations, truncations, infos = [], [], [], [], {}
         successes = []
@@ -339,7 +336,7 @@ def step_wait(
                 observations,
                 self.observations,
             )
-        
+
         # modify to allow return of vector rewards
         self.rewards = concatenate(
             self.reward_space,

pyproject.toml

@@ -22,7 +22,7 @@ classifiers = [
     'Topic :: Scientific/Engineering :: Artificial Intelligence',
 ]
 dependencies = [
-    "gymnasium >=1.0.0",
+    "gymnasium >=1.1.0",
     "numpy >=1.21.0,<2.0",
     "pygame >=2.1.3",
     "scipy >=1.7.3",

tests/test_vector_wrappers.py

@@ -2,26 +2,39 @@
 import numpy as np
 
 import mo_gymnasium as mo_gym
-from mo_gymnasium.wrappers.vector import MORecordEpisodeStatistics, MOSyncVectorEnv
+from mo_gymnasium.wrappers.vector import (
+    MOAsyncVectorEnv,
+    MORecordEpisodeStatistics,
+    MOSyncVectorEnv,
+)
 
 
-def test_mo_sync_wrapper():
-    num_envs = 3
-    envs = MOSyncVectorEnv([lambda: mo_gym.make("deep-sea-treasure-v0") for _ in range(num_envs)])
-
-    envs.reset()
+def _test_logic(envs, num_envs: int):
     obs, rewards, terminateds, truncateds, infos = envs.step(envs.action_space.sample())
     assert len(obs) == num_envs, "Number of observations do not match the number of envs"
     assert len(rewards) == num_envs, "Number of rewards do not match the number of envs"
     assert len(terminateds) == num_envs, "Number of terminateds do not match the number of envs"
     assert len(truncateds) == num_envs, "Number of truncateds do not match the number of envs"
-    envs.close()
 
 
-def test_mo_sync_autoreset():
-    num_envs = 2
+def test_mo_sync_wrapper():
+    num_envs = 3
     envs = MOSyncVectorEnv([lambda: mo_gym.make("deep-sea-treasure-v0") for _ in range(num_envs)])
 
+    envs.reset()
+    _test_logic(envs, num_envs)
+    envs.close()
+
+
+def test_mo_async_wrapper():
+    num_envs = 3
+    envs = MOAsyncVectorEnv([lambda: mo_gym.make("deep-sea-treasure-v0") for _ in range(num_envs)])
+    envs.reset()
+    _test_logic(envs, num_envs)
+    envs.close()
+
+
+def _test_autoreset_logic(envs):
     obs, infos = envs.reset()
     assert (obs[0] == [0, 0]).all()
     assert (obs[1] == [0, 0]).all()
@@ -42,14 +55,25 @@ def test_mo_sync_autoreset():
     assert (rewards[1] == [0.0, 0.0]).all()  # Reset step
     assert not terminateds[0]
     assert not terminateds[1]  # Not done anymore
-    envs.close()
 
 
-def test_mo_record_ep_statistic_vector_env():
+def test_mo_sync_autoreset():
     num_envs = 2
     envs = MOSyncVectorEnv([lambda: mo_gym.make("deep-sea-treasure-v0") for _ in range(num_envs)])
-    envs = MORecordEpisodeStatistics(envs, gamma=0.97)
 
+    _test_autoreset_logic(envs)
+    envs.close()
+
+
+def test_mo_async_autoreset():
+    num_envs = 2
+    envs = MOAsyncVectorEnv([lambda: mo_gym.make("deep-sea-treasure-v0") for _ in range(num_envs)])
+
+    _test_autoreset_logic(envs)
+    envs.close()
+
+
+def _test_record_ep_statistic_logic(envs, num_envs: int):
     envs.reset()
     terminateds = np.array([False] * num_envs)
     info = {}
@@ -69,21 +93,48 @@ def test_mo_record_ep_statistic_vector_env():
     assert isinstance(info["episode"]["l"], np.ndarray)
     np.testing.assert_almost_equal(info["episode"]["l"], np.array([0, 3], dtype=np.float32), decimal=2)
     assert isinstance(info["episode"]["t"], np.ndarray)
+
+
+def test_mo_record_ep_statistic_vector_env():
+    num_envs = 2
+    envs = MOSyncVectorEnv([lambda: mo_gym.make("deep-sea-treasure-v0") for _ in range(num_envs)])
+    envs = MORecordEpisodeStatistics(envs, gamma=0.97)
+    _test_record_ep_statistic_logic(envs, num_envs)
     envs.close()
 
 
-def test_gym_wrapper_and_vector():
-    # This tests the integration of gym-wrapped envs with MO-Gymnasium vectorized envs
+def test_mo_record_ep_statistic_vector_env_async():
     num_envs = 2
-    envs = MOSyncVectorEnv(
-        [lambda: gym.wrappers.NormalizeObservation(mo_gym.make("deep-sea-treasure-v0")) for _ in range(num_envs)]
-    )
+    envs = MOAsyncVectorEnv([lambda: mo_gym.make("deep-sea-treasure-v0") for _ in range(num_envs)])
+    envs = MORecordEpisodeStatistics(envs, gamma=0.97)
+    _test_record_ep_statistic_logic(envs, num_envs)
+    envs.close()
+
 
+def _test_gym_wrapper_and_vector_logic(envs, num_envs: int):
     envs.reset()
     for i in range(30):
         obs, rewards, terminateds, truncateds, infos = envs.step(envs.action_space.sample())
     assert len(obs) == num_envs, "Number of observations do not match the number of envs"
     assert len(rewards) == num_envs, "Number of rewards do not match the number of envs"
     assert len(terminateds) == num_envs, "Number of terminateds do not match the number of envs"
     assert len(truncateds) == num_envs, "Number of truncateds do not match the number of envs"
+
+
+def test_gym_wrapper_and_vector():
+    # This tests the integration of gym-wrapped envs with MO-Gymnasium vectorized envs
+    num_envs = 2
+    envs = MOSyncVectorEnv(
+        [lambda: gym.wrappers.NormalizeObservation(mo_gym.make("deep-sea-treasure-v0")) for _ in range(num_envs)]
+    )
+    _test_gym_wrapper_and_vector_logic(envs, num_envs)
+    envs.close()
+
+
+def test_gym_wrapper_and_vector_async():
+    num_envs = 2
+    envs = MOAsyncVectorEnv(
+        [lambda: gym.wrappers.NormalizeObservation(mo_gym.make("deep-sea-treasure-v0")) for _ in range(num_envs)]
+    )
+    _test_gym_wrapper_and_vector_logic(envs, num_envs)
     envs.close()