multi-agent-rlenv 3.2.2__py3-none-any.whl → 3.3.1__py3-none-any.whl

marlenv/__init__.py

@@ -1,19 +1,68 @@
 """
 `marlenv` is a strongly typed library for multi-agent and multi-objective reinforcement learning.
 
-It aims to
-    - provide a simple and consistent interface for reinforcement learning environments
-    - provide fundamental models such as `Observation`s, `Episode`s, `Transition`s, ...
-    - work with gymnasium, pettingzoo and SMAC out of the box
-    - work with multi-objective environments
-    - provide helpful wrappers to add intrinsic rewards, agent ids, record videos, ...
+It aims to provide a simple and consistent interface for reinforcement learning environments by providing abstraction models such as `Observation`s or `Episode`s. `marlenv` provides adapters for popular libraries such as `gym` or `pettingzoo` and provides utility wrappers to add functionalities such as video recording or limiting the number of steps.
 
+Almost every class is a dataclassto enable seemless serialiation with the `orjson` library.
 
-A design choice is taht almost every class is a dataclass. This makes it easy to
-serialize and deserialize classes, for instance to json with the `orjson` library.
+# Existing environments
+The `MARLEnv` class represents a multi-agent RL environment and is at the center of this library, and `marlenv` provides an adapted implementation of multiple common MARL environments (gym, pettingzoo, smac and overcooked) in `marlenv.adapters`. Note that these adapters will only work if you have the corresponding library installed.
+
+```python
+from marlenv.adapters import Gym, PettingZoo, SMAC, Overcooked
+import marlenv
+
+env1 = Gym("CartPole-v1")
+env2 = marlenv.make("CartPole-v1")
+env3 = PettingZoo("prospector_v4")
+env4 = SMAC("3m")
+env5 = Overcooked.from_layout("cramped_room")
+```
+
+# Wrappers & Builder
+To facilitate the create of an environment with common wrappers, `marlenv` provides a `Builder` class that can be used to chain the creation of multiple wrappers.
+
+```python
+from marlenv import make, Builder
+
+env = <your env>
+env = Builder(env).agent_id().time_limit(50).record("videos").build()
+```
+
+# Using the library
+A typical environment loop would look like this:
+
+```python
+from marlenv import DiscreteMockEnv, Builder, Episode
+
+env = Builder(DicreteMockEnv()).agent_id().build()
+obs, state = env.reset()
+terminated = False
+episode = Episode.new(obs, state)
+while not episode.is_finished:
+    action = env.sample_action() # a valid random action
+    step = env.step(action) # Step data `step.obs`, `step.reward`, ...
+    episode.add(step, action) # Progressively build the episode
+```
+
+# Extras
+To cope with complex observation spaces, `marlenv` distinguishes the "main" observation data from the "extra" observation data. A typical example would be the observation of a gridworld environment with a time limit. In that case, the main observation has shape (height, width), i.e. the content of the grid, but the current time is an extra observation data of shape (1, ).
+
+```python
+env = GridWorldEnv()
+print(env.observation_shape) # (height, width)
+print(env.extras_shape) # (0, )
+
+env = Builder(env).time_limit(25).build()
+print(env.observation_shape) # (height, width)
+print(env.extras_shape) # (1, )
+```
+
+# Creating a new environment
+If you want to create a new environment, you can simply create a class that inherits from `MARLEnv`. If you want to create a wrapper around an existing `MARLEnv`, you probably want to subclass `RLEnvWrapper` which implements a default behaviour for every method.
 """
 
-__version__ = "3.2.2"
+__version__ = "3.3.1"
 
 from . import models
 from . import wrappers

marlenv/adapters/__init__.py

@@ -1,24 +1,42 @@
+from importlib.util import find_spec
 from .pymarl_adapter import PymarlAdapter
-from typing import Any
 
-__all__ = ["PymarlAdapter"]
-try:
+HAS_GYM = False
+if find_spec("gymnasium") is not None:
     from .gym_adapter import Gym
 
-    __all__.append("Gym")
-except ImportError:
-    Gym = Any
+    HAS_GYM = True
 
-try:
+HAS_PETTINGZOO = False
+if find_spec("pettingzoo") is not None:
     from .pettingzoo_adapter import PettingZoo
 
-    __all__.append("PettingZoo")
-except ImportError:
-    PettingZoo = Any
+    HAS_PETTINGZOO = True
 
-try:
+HAS_SMAC = False
+if find_spec("smac") is not None:
     from .smac_adapter import SMAC
 
-    __all__.append("SMAC")
-except ImportError:
-    SMAC = Any
+    HAS_SMAC = True
+
+HAS_OVERCOOKED = False
+if find_spec("overcooked_ai_py.mdp") is not None:
+    import numpy
+
+    # Overcooked assumes a version of numpy <2.0 where np.Inf is available.
+    setattr(numpy, "Inf", numpy.inf)
+    from .overcooked_adapter import Overcooked
+
+    HAS_OVERCOOKED = True
+
+__all__ = [
+    "PymarlAdapter",
+    "Gym",
+    "PettingZoo",
+    "SMAC",
+    "Overcooked",
+    "HAS_GYM",
+    "HAS_PETTINGZOO",
+    "HAS_SMAC",
+    "HAS_OVERCOOKED",
+]

marlenv/adapters/gym_adapter.py

@@ -1,3 +1,5 @@
+import sys
+import cv2
 from dataclasses import dataclass
 from typing import Sequence
 
@@ -79,7 +81,10 @@ class Gym(MARLEnv[Sequence | npt.NDArray, ActionSpace]):
         return self.last_obs, self.get_state()
 
     def get_image(self):
-        return self.env.render()
+        image = np.array(self.env.render())
+        if sys.platform in ("linux", "linux2"):
+            image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
+        return image
 
     def seed(self, seed_value: int):
         self.env.reset(seed=seed_value)

marlenv/adapters/overcooked_adapter.py

@@ -0,0 +1,164 @@
+import sys
+from dataclasses import dataclass
+from typing import Literal, Sequence
+
+import cv2
+import numpy as np
+import numpy.typing as npt
+import pygame
+from marlenv.models import ContinuousSpace, DiscreteActionSpace, MARLEnv, Observation, State, Step
+
+from overcooked_ai_py.mdp.overcooked_env import OvercookedEnv
+from overcooked_ai_py.mdp.overcooked_mdp import Action, OvercookedGridworld, OvercookedState
+from overcooked_ai_py.visualization.state_visualizer import StateVisualizer
+
+
+@dataclass
+class Overcooked(MARLEnv[Sequence[int] | npt.NDArray, DiscreteActionSpace]):
+    horizon: int
+
+    def __init__(self, oenv: OvercookedEnv):
+        self._oenv = oenv
+        assert isinstance(oenv.mdp, OvercookedGridworld)
+        self._mdp = oenv.mdp
+        self.visualizer = StateVisualizer()
+        shape = tuple(int(s) for s in self._mdp.get_lossless_state_encoding_shape())
+        shape = (shape[2], shape[0], shape[1])
+        super().__init__(
+            action_space=DiscreteActionSpace(
+                n_agents=self._mdp.num_players,
+                n_actions=Action.NUM_ACTIONS,
+                action_names=[Action.ACTION_TO_CHAR[a] for a in Action.ALL_ACTIONS],
+            ),
+            observation_shape=shape,
+            extras_shape=(1,),
+            extras_meanings=["timestep"],
+            state_shape=shape,
+            state_extra_shape=(1,),
+            reward_space=ContinuousSpace.from_shape(1),
+        )
+        self.horizon = int(self._oenv.horizon)
+
+    @property
+    def state(self) -> OvercookedState:
+        """Current state of the environment"""
+        return self._oenv.state
+
+    def set_state(self, state: State):
+        raise NotImplementedError("Not yet implemented")
+
+    @property
+    def time_step(self):
+        return self.state.timestep
+
+    def _state_data(self):
+        state = np.array(self._mdp.lossless_state_encoding(self.state))
+        # Use axes (agents, channels, height, width) instead of (agents, height, width, channels)
+        state = np.transpose(state, (0, 3, 1, 2))
+        return state
+
+    def get_state(self):
+        return State(self._state_data()[0], np.array([self.time_step / self.horizon]))
+
+    def get_observation(self) -> Observation:
+        return Observation(
+            data=self._state_data(),
+            available_actions=self.available_actions(),
+            extras=np.array([[self.time_step / self.horizon]] * self.n_agents),
+        )
+
+    def available_actions(self):
+        available_actions = np.full((self.n_agents, self.n_actions), False)
+        actions = self._mdp.get_actions(self._oenv.state)
+        for agent_num, agent_actions in enumerate(actions):
+            for action in agent_actions:
+                available_actions[agent_num, Action.ACTION_TO_INDEX[action]] = True
+        return np.array(available_actions)
+
+    def step(self, actions: Sequence[int] | npt.NDArray[np.int32 | np.int64]) -> Step:
+        actions = [Action.ALL_ACTIONS[a] for a in actions]
+        _, reward, done, info = self._oenv.step(actions, display_phi=True)
+        return Step(
+            obs=self.get_observation(),
+            state=self.get_state(),
+            reward=np.array([reward]),
+            done=done,
+            truncated=False,
+            info=info,
+        )
+
+    def get_image(self):
+        rewards_dict = {}  # dictionary of details you want rendered in the UI
+        for key, value in self._oenv.game_stats.items():
+            if key in [
+                "cumulative_shaped_rewards_by_agent",
+                "cumulative_sparse_rewards_by_agent",
+            ]:
+                rewards_dict[key] = value
+
+        image = self.visualizer.render_state(
+            state=self._oenv.state,
+            grid=self._mdp.terrain_mtx,
+            hud_data=StateVisualizer.default_hud_data(self._oenv.state, **rewards_dict),
+        )
+
+        image = pygame.surfarray.array3d(image)
+        image = np.flip(np.rot90(image, 3), 1)
+        # Depending on the platform, the image may need to be converted to RGB
+        if sys.platform in ("linux", "linux2"):
+            image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
+        return image
+
+    @staticmethod
+    def from_layout(
+        layout: Literal[
+            "asymmetric_advantages",
+            "asymmetric_advantages_tomato",
+            "bonus_order_test",
+            "bottleneck",
+            "centre_objects",
+            "centre_pots",
+            "coordination_ring",
+            "corridor",
+            "counter_circuit",
+            "counter_circuit_o_1order",
+            "cramped_corridor",
+            "cramped_room",
+            "cramped_room_o_3orders",
+            "cramped_room_single",
+            "cramped_room_tomato",
+            "five_by_five",
+            "forced_coordination",
+            "forced_coordination_tomato",
+            "inverse_marshmallow_experiment",
+            "large_room",
+            "long_cook_time",
+            "marshmallow_experiment_coordination",
+            "marshmallow_experiment",
+            "mdp_test",
+            "m_shaped_s",
+            "multiplayer_schelling",
+            "pipeline",
+            "scenario1_s",
+            "scenario2",
+            "scenario2_s",
+            "scenario3",
+            "scenario4",
+            "schelling",
+            "schelling_s",
+            "simple_o",
+            "simple_o_t",
+            "simple_tomato",
+            "small_corridor",
+            "soup_coordination",
+            "tutorial_0",
+            "tutorial_1",
+            "tutorial_2",
+            "tutorial_3",
+            "unident",
+            "you_shall_not_pass",
+        ],
+        horizon: int = 400,
+    ):
+        mdp = OvercookedGridworld.from_layout_name(layout)
+        return Overcooked(OvercookedEnv.from_mdp(mdp, horizon=horizon))

marlenv/env_builder.py

@@ -1,32 +1,27 @@
 from dataclasses import dataclass
 from typing import Generic, Literal, Optional, TypeVar, overload
-
 import numpy as np
 import numpy.typing as npt
 
 from . import wrappers
+from marlenv import adapters
 from .models import ActionSpace, MARLEnv
-from .adapters import PettingZoo
 
 A = TypeVar("A")
 AS = TypeVar("AS", bound=ActionSpace)
 
-try:
+if adapters.HAS_PETTINGZOO:
+    from .adapters import PettingZoo
     from pettingzoo import ParallelEnv
 
     @overload
-    def make(
-        env: ParallelEnv,
-    ) -> PettingZoo: ...
-
-    HAS_PETTINGZOO = True
-except ImportError:
-    HAS_PETTINGZOO = False
+    def make(env: ParallelEnv) -> PettingZoo: ...
 
 
-try:
-    from gymnasium import Env
+if adapters.HAS_GYM:
     from .adapters import Gym
+    from gymnasium import Env
+    import gymnasium
 
     @overload
     def make(env: Env) -> Gym: ...
@@ -37,25 +32,21 @@ try:
         Make an RLEnv from the `gymnasium` registry (e.g: "CartPole-v1").
         """
 
-    HAS_GYM = True
-except ImportError:
-    HAS_GYM = False
 
-try:
-    from smac.env import StarCraft2Env
+if adapters.HAS_SMAC:
     from .adapters import SMAC
+    from smac.env import StarCraft2Env
 
     @overload
     def make(env: StarCraft2Env) -> SMAC: ...
 
-    HAS_SMAC = True
-except ImportError:
-    HAS_SMAC = False
 
+if adapters.HAS_OVERCOOKED:
+    from .adapters import Overcooked
+    from overcooked_ai_py.mdp.overcooked_env import OvercookedEnv
 
-@overload
-def make(env: MARLEnv[A, AS]) -> MARLEnv[A, AS]:
-    """Why would you do this ?"""
+    @overload
+    def make(env: OvercookedEnv) -> Overcooked: ...
 
 
 def make(env, **kwargs):
@@ -64,32 +55,18 @@ def make(env, **kwargs):
         case MARLEnv():
             return env
         case str(env_id):
-            try:
-                import gymnasium
-            except ImportError:
-                raise ImportError("Gymnasium is not installed !")
-            from marlenv.adapters import Gym
-
-            gym_env = gymnasium.make(env_id, render_mode="rgb_array", **kwargs)
-            return Gym(gym_env)
-
-    try:
-        from marlenv.adapters import PettingZoo
-
-        if isinstance(env, ParallelEnv):
-            return PettingZoo(env)
-    except ImportError:
-        pass
-    try:
-        from smac.env import StarCraft2Env
-
-        from marlenv.adapters import SMAC
-
-        if isinstance(env, StarCraft2Env):
-            return SMAC(env)
-    except ImportError:
-        pass
-
+            if adapters.HAS_GYM:
+                gym_env = gymnasium.make(env_id, render_mode="rgb_array", **kwargs)
+                return Gym(gym_env)
+
+    if adapters.HAS_PETTINGZOO and isinstance(env, ParallelEnv):
+        return PettingZoo(env)  # type: ignore
+    if adapters.HAS_SMAC and isinstance(env, StarCraft2Env):
+        return SMAC(env)
+    if adapters.HAS_OVERCOOKED and isinstance(env, OvercookedEnv):
+        return Overcooked(env)  # type: ignore
+    if adapters.HAS_GYM and isinstance(env, Env):
+        return Gym(env)
     raise ValueError(f"Unknown environment type: {type(env)}")
 
 
@@ -116,6 +93,11 @@ class Builder(Generic[A, AS]):
         self._env = wrappers.TimeLimit(self._env, n_steps, add_extra, truncation_penalty)
         return self
 
+    def delay_rewards(self, delay: int):
+        """Delays the rewards by `delay` steps"""
+        self._env = wrappers.DelayedReward(self._env, delay)
+        return self
+
     def pad(self, to_pad: Literal["obs", "extra"], n: int):
         match to_pad:
             case "obs":

marlenv/env_pool.py

@@ -1,6 +1,5 @@
 from typing import Sequence
 from dataclasses import dataclass
-import numpy as np
 import numpy.typing as npt
 from typing_extensions import TypeVar
 import random

marlenv/mock_env.py

@@ -2,7 +2,7 @@ from typing import Sequence
 import numpy as np
 import numpy.typing as npt
 from dataclasses import dataclass
-from marlenv import MARLEnv, Observation, DiscreteActionSpace, DiscreteSpace, Step, State
+from marlenv import MARLEnv, Observation, DiscreteActionSpace, ContinuousSpace, Step, State
 
 
 @dataclass
@@ -13,21 +13,31 @@ class DiscreteMockEnv(MARLEnv[Sequence[int] | npt.NDArray, DiscreteActionSpace])
         obs_size: int = 42,
         n_actions: int = 5,
         end_game: int = 30,
-        reward_step: int = 1,
+        reward_step: int | float | np.ndarray | list = 1,
         agent_state_size: int = 1,
         extras_size: int = 0,
     ) -> None:
+        match reward_step:
+            case int() | float():
+                reward_step = np.array([reward_step])
+            case list():
+                reward_step = np.array(reward_step)
+            case np.ndarray():
+                reward_step = reward_step
+            case _:
+                raise ValueError("reward_step must be an int, float or np.ndarray")
         super().__init__(
             DiscreteActionSpace(n_agents, n_actions),
             (obs_size,),
             (n_agents * agent_state_size,),
             extras_shape=(extras_size,),
+            reward_space=ContinuousSpace.from_shape(reward_step.shape),
         )
+        self.reward_step = reward_step
         self.obs_size = obs_size
         self.extra_size = extras_size
         self._agent_state_size = agent_state_size
         self.end_game = end_game
-        self.reward_step = reward_step
         self.t = 0
         self.actions_history = []
         self._seed = -1
@@ -70,7 +80,7 @@ class DiscreteMockEnv(MARLEnv[Sequence[int] | npt.NDArray, DiscreteActionSpace])
         return Step(
             self.get_observation(),
             self.get_state(),
-            np.array([self.reward_step]),
+            self.reward_step,
             self.t >= self.end_game,
         )
 
@@ -94,7 +104,7 @@ class DiscreteMOMockEnv(MARLEnv[Sequence[int] | npt.NDArray, DiscreteActionSpace
             (obs_size,),
             (n_agents * agent_state_size,),
             extras_shape=(extras_size,),
-            reward_space=DiscreteSpace(n_objectives),
+            reward_space=ContinuousSpace.from_shape(n_objectives),
         )
         self.obs_size = obs_size
         self.extra_size = extras_size

marlenv/models/env.py

@@ -1,17 +1,17 @@
 from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from itertools import product
 from typing import Generic, Optional, Sequence
-from typing_extensions import TypeVar
+
 import cv2
 import numpy as np
 import numpy.typing as npt
-from dataclasses import dataclass
-from itertools import product
-
+from typing_extensions import TypeVar
 
-from .step import Step
-from .state import State
-from .spaces import ActionSpace, DiscreteSpace
 from .observation import Observation
+from .spaces import ActionSpace, ContinuousSpace, Space
+from .state import State
+from .step import Step
 
 ActionType = TypeVar("ActionType", default=npt.NDArray)
 ActionSpaceType = TypeVar("ActionSpaceType", bound=ActionSpace, default=ActionSpace)
@@ -25,6 +25,34 @@ class MARLEnv(ABC, Generic[ActionType, ActionSpaceType]):
     This type is generic on
         - the action type
         - the action space
+
+    You can inherit from this class to create your own environemnt:
+    ```
+    import numpy as np
+    from marlenv import MARLEnv, DiscreteActionSpace, Observation
+
+    N_AGENTS = 3
+    N_ACTIONS = 5
+
+    class CustomEnv(MARLEnv[DiscreteActionSpace]):
+        def __init__(self, width: int, height: int):
+            super().__init__(
+                action_space=DiscreteActionSpace(N_AGENTS, N_ACTIONS),
+                observation_shape=(height, width),
+                state_shape=(1,),
+            )
+            self.time = 0
+
+        def reset(self) -> Observation:
+            self.time = 0
+            ...
+            return obs
+
+        def get_state(self):
+            return np.array([self.time])
+
+        ...
+    ```
     """
 
     action_space: ActionSpaceType
@@ -38,6 +66,7 @@ class MARLEnv(ABC, Generic[ActionType, ActionSpaceType]):
     n_agents: int
     n_actions: int
     name: str
+    reward_space: Space
 
     def __init__(
         self,
@@ -46,7 +75,7 @@ class MARLEnv(ABC, Generic[ActionType, ActionSpaceType]):
         state_shape: tuple[int, ...],
         extras_shape: tuple[int, ...] = (0,),
         state_extra_shape: tuple[int, ...] = (0,),
-        reward_space: Optional[DiscreteSpace] = None,
+        reward_space: Optional[Space] = None,
         extras_meanings: Optional[list[str]] = None,
     ):
         super().__init__()
@@ -58,7 +87,9 @@ class MARLEnv(ABC, Generic[ActionType, ActionSpaceType]):
         self.state_shape = state_shape
         self.extras_shape = extras_shape
         self.state_extra_shape = state_extra_shape
-        self.reward_space = reward_space or DiscreteSpace(1, labels=["Reward"])
+        if reward_space is None:
+            reward_space = ContinuousSpace.from_shape(1, labels=["Reward"])
+        self.reward_space = reward_space
         if extras_meanings is None:
             extras_meanings = [f"{self.name}-extra-{i}" for i in range(extras_shape[0])]
         elif len(extras_meanings) != extras_shape[0]:
@@ -77,9 +108,9 @@ class MARLEnv(ABC, Generic[ActionType, ActionSpaceType]):
         """Whether the environment is multi-objective."""
         return self.reward_space.size > 1
 
-    def sample_action(self):
+    def sample_action(self) -> ActionType:
         """Sample an available action from the action space."""
-        return self.action_space.sample(self.available_actions())
+        return self.action_space.sample(self.available_actions())  # type: ignore
 
     def available_actions(self) -> npt.NDArray[np.bool]:
         """
@@ -123,6 +154,10 @@ class MARLEnv(ABC, Generic[ActionType, ActionSpaceType]):
         - info: Extra information
         """
 
+    def random_step(self) -> Step:
+        """Perform a random step in the environment."""
+        return self.step(self.sample_action())
+
     def reset(self) -> tuple[Observation, State]:
         """Reset the environment and return the initial observation and state."""
         return self.get_observation(), self.get_state()

marlenv/models/observation.py

@@ -58,9 +58,14 @@ class Observation:
             available_actions=self.available_actions[agent_id],
         )
 
+    @property
+    def shape(self) -> tuple[int, ...]:
+        """The individual shape of the observation data"""
+        return self.data[0].shape
+
     @property
     def extras_shape(self) -> tuple[int, ...]:
-        """The shape of the observation extras"""
+        """The individual shape of the observation extras"""
         return self.extras[0].shape
 
     def __hash__(self):

marlenv/models/spaces.py

@@ -12,14 +12,14 @@ S = TypeVar("S", bound="Space")
 @dataclass
 class Space(ABC):
     shape: tuple[int, ...]
-    n_dims: int
+    size: int
     labels: list[str]
 
     def __init__(self, shape: tuple[int, ...], labels: Optional[list[str]] = None):
         self.shape = shape
-        self.n_dims = len(shape)
+        self.size = math.prod(shape)
         if labels is None:
-            labels = [f"Dim {i}" for i in range(self.n_dims)]
+            labels = [f"Dim {i}" for i in range(self.size)]
         self.labels = labels
 
     @abstractmethod
@@ -100,16 +100,55 @@ class ContinuousSpace(Space):
     high: npt.NDArray[np.float32]
     """Upper bound of the space for each dimension."""
 
+    @staticmethod
+    def from_bounds(
+        low: int | float | list | npt.NDArray[np.float32],
+        high: int | float | list | npt.NDArray[np.float32],
+        labels: Optional[list[str]] = None,
+    ):
+        match low:
+            case list():
+                low = np.array(low, dtype=np.float32)
+            case float() | int():
+                low = np.array([low], dtype=np.float32)
+        match high:
+            case list():
+                high = np.array(high, dtype=np.float32)
+            case float() | int():
+                high = np.array([high], dtype=np.float32)
+        return ContinuousSpace(low, high, labels)
+
+    @staticmethod
+    def from_shape(
+        shape: int | tuple[int, ...],
+        low: Optional[int | float | list | npt.NDArray[np.float32]] = None,
+        high: Optional[int | float | list | npt.NDArray[np.float32]] = None,
+        labels: Optional[list[str]] = None,
+    ):
+        if isinstance(shape, int):
+            shape = (shape,)
+        match low:
+            case None:
+                low = np.full(shape, -np.inf, dtype=np.float32)
+            case float() | int():
+                low = np.full(shape, low, dtype=np.float32)
+            case list():
+                low = np.array(low, dtype=np.float32)
+        match high:
+            case None:
+                high = np.full(shape, np.inf, dtype=np.float32)
+            case float() | int():
+                high = np.full(shape, high, dtype=np.float32)
+            case list():
+                high = np.array(high, dtype=np.float32)
+        return ContinuousSpace(low, high, labels)
+
     def __init__(
         self,
-        low: list | npt.NDArray[np.float32],
-        high: list | npt.NDArray[np.float32],
+        low: npt.NDArray[np.float32],
+        high: npt.NDArray[np.float32],
         labels: Optional[list[str]] = None,
     ):
-        if isinstance(low, list):
-            low = np.array(low, dtype=np.float32)
-        if isinstance(high, list):
-            high = np.array(high, dtype=np.float32)
         assert low.shape == high.shape, "Low and high must have the same shape."
         assert np.all(low <= high), "All elements in low must be less than the corresponding elements in high."
         Space.__init__(self, low.shape, labels)
@@ -182,5 +221,5 @@ class MultiDiscreteActionSpace(ActionSpace[MultiDiscreteSpace]):
 @dataclass
 class ContinuousActionSpace(ActionSpace[ContinuousSpace]):
     def __init__(self, n_agents: int, low: np.ndarray | list, high: np.ndarray | list, action_names: list | None = None):
-        space = ContinuousSpace(low, high, action_names)
+        space = ContinuousSpace.from_bounds(low, high, action_names)
         super().__init__(n_agents, space, action_names)

marlenv/wrappers/__init__.py

@@ -9,6 +9,7 @@ from .available_actions_wrapper import AvailableActions
 from .blind_wrapper import Blind
 from .centralised import Centralised
 from .available_actions_mask import AvailableActionsMask
+from .delayed_rewards import DelayedReward
 
 __all__ = [
     "RLEnvWrapper",
@@ -24,4 +25,5 @@ __all__ = [
     "AvailableActions",
     "Blind",
     "Centralised",
+    "DelayedReward",
 ]

marlenv/wrappers/delayed_rewards.py

@@ -0,0 +1,36 @@
+from .rlenv_wrapper import RLEnvWrapper, MARLEnv
+from marlenv.models import ActionSpace
+from typing_extensions import TypeVar
+import numpy.typing as npt
+import numpy as np
+from dataclasses import dataclass
+from collections import deque
+
+A = TypeVar("A", default=npt.NDArray)
+AS = TypeVar("AS", bound=ActionSpace, default=ActionSpace)
+
+
+@dataclass
+class DelayedReward(RLEnvWrapper[A, AS]):
+    delay: int
+
+    def __init__(self, env: MARLEnv[A, AS], delay: int):
+        super().__init__(env)
+        self.delay = delay
+        self.reward_queue = deque[npt.NDArray[np.float32]](maxlen=delay + 1)
+
+    def reset(self):
+        self.reward_queue.clear()
+        for _ in range(self.delay):
+            self.reward_queue.append(np.zeros(self.reward_space.shape, dtype=np.float32))
+        return super().reset()
+
+    def step(self, actions: A):
+        step = super().step(actions)
+        self.reward_queue.append(step.reward)
+        # If the step is terminal, we sum all the remaining rewards
+        if step.is_terminal:
+            step.reward = np.sum(self.reward_queue, axis=0)
+        else:
+            step.reward = self.reward_queue.popleft()
+        return step

{multi_agent_rlenv-3.2.2.dist-info → multi_agent_rlenv-3.3.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: multi-agent-rlenv
-Version: 3.2.2
+Version: 3.3.1
 Summary: A strongly typed Multi-Agent Reinforcement Learning framework
 Project-URL: repository, https://github.com/yamoling/multi-agent-rlenv
 Author-email: Yannick Molinghen <yannick.molinghen@ulb.be>
@@ -8,16 +8,58 @@ License-File: LICENSE
 Classifier: Operating System :: OS Independent
 Classifier: Programming Language :: Python :: 3
 Requires-Python: <4,>=3.10
-Requires-Dist: gymnasium>=0.29.1
 Requires-Dist: numpy>=2.0.0
-Requires-Dist: opencv-python>=4.10.0.84
+Requires-Dist: opencv-python>=4.0
+Requires-Dist: typing-extensions>=4.0
+Provides-Extra: all
+Requires-Dist: gymnasium>0.29.1; extra == 'all'
+Requires-Dist: overcooked-ai; extra == 'all'
+Requires-Dist: pettingzoo>=1.20; extra == 'all'
+Requires-Dist: pymunk>=6.0; extra == 'all'
+Requires-Dist: pysc2; extra == 'all'
+Requires-Dist: scipy>=1.10; extra == 'all'
+Requires-Dist: smac; extra == 'all'
+Provides-Extra: gym
+Requires-Dist: gymnasium>=0.29.1; extra == 'gym'
+Provides-Extra: overcooked
+Requires-Dist: overcooked-ai>=1.1.0; extra == 'overcooked'
+Requires-Dist: scipy>=1.10; extra == 'overcooked'
+Provides-Extra: pettingzoo
+Requires-Dist: pettingzoo>=1.20; extra == 'pettingzoo'
+Requires-Dist: pymunk>=6.0; extra == 'pettingzoo'
+Requires-Dist: scipy>=1.10; extra == 'pettingzoo'
+Provides-Extra: smac
+Requires-Dist: pysc2; extra == 'smac'
+Requires-Dist: smac; extra == 'smac'
 Description-Content-Type: text/markdown
 
-# `marlenv` - A unified interface for muti-agent reinforcement learning
+# `marlenv` - A unified framework for muti-agent reinforcement learning
+**Documentation: [https://yamoling.github.io/multi-agent-rlenv](https://yamoling.github.io/multi-agent-rlenv)**
+
 The objective of `marlenv` is to provide a common (typed) interface for many different reinforcement learning environments.
 
 As such, `marlenv` provides high level abstractions of RL concepts such as `Observation`s or `Transition`s that are commonly represented as mere (confusing) lists or tuples.
 
+## Installation
+Install with you preferred package manager (`uv`, `pip`, `poetry`, ...):
+```bash
+$ pip install marlenv[all] # Enable all features
+$ pip install marlenv      # Basic installation
+```
+
+There are multiple optional dependencies if you want to support specific libraries and environments. Available options are:
+- `smac` for StarCraft II environments
+- `gym` for OpenAI Gym environments
+- `pettingzoo` for PettingZoo environments
+- `overcooked` for Overcooked environments
+
+Install them with:
+```bash
+$ pip install marlenv[smac] # Install SMAC
+$ pip install marlenv[gym,smac]  # Install Gym & smac support
+```
+
+
 ## Using `marlenv` with existing libraries
 `marlenv` unifies multiple popular libraries under a single interface. Namely, `marlenv` supports `smac`, `gymnasium` and `pettingzoo`.
 
@@ -47,7 +89,7 @@ from marlenv import RLEnv, DiscreteActionSpace, Observation
 N_AGENTS = 3
 N_ACTIONS = 5
 
-class CustomEnv(RLEnv[DiscreteActionSpace]):
+class CustomEnv(MARLEnv[DiscreteActionSpace]):
     def __init__(self, width: int, height: int):
         super().__init__(
             action_space=DiscreteActionSpace(N_AGENTS, N_ACTIONS),

{multi_agent_rlenv-3.2.2.dist-info → multi_agent_rlenv-3.3.1.dist-info}/RECORD

@@ -1,35 +1,37 @@
-marlenv/__init__.py,sha256=mA5XCI_yboGYESxCtbpp5oKHx1tplA-TwnyIjazx3nE,1499
-marlenv/env_builder.py,sha256=G2ZkTradBmbfN8U84ux_Bflv8S6pr2zJiQVrWaP6ygc,5527
-marlenv/env_pool.py,sha256=TSSYwD5-g4G473Ea097wFVbp3tyQrawywLIAFFEJCJY,1089
+marlenv/__init__.py,sha256=7HYiLVpZ4PQQWYjvc_78eXrnidv-FfUiY1pUkbxpm5U,3741
+marlenv/env_builder.py,sha256=_rdwcWRqnHP7i4M4Oje1Y2nrEBKH9EzTpqOuw_PNUyw,5560
+marlenv/env_pool.py,sha256=R3WIrnQ5Zvff4HR1ecfkDmuO2zl7v1ywQ0K2_nvWFzs,1070
 marlenv/exceptions.py,sha256=gJUC_2rVAvOfK_ypVFc7Myh-pIfSU3To38VBVS_0rZA,1179
-marlenv/mock_env.py,sha256=fGPTJ0GEYvHNraI9OJh3PN1My27equFZjf3o8iZnHJs,4289
+marlenv/mock_env.py,sha256=qB0fYFIfbopJf7Va8kCeVI5vsOy1-2JdEYe9gdV1Ruw,4761
 marlenv/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-marlenv/adapters/__init__.py,sha256=eJLY15xLDik9nO51Qa4Hp4Ak1_tLxm5C7kzFySNmUjs,425
-marlenv/adapters/gym_adapter.py,sha256=Nwp3gtuZQM1j9ify3h5QaA4enAJ1e9cTT_nPi7qNubE,2703
+marlenv/adapters/__init__.py,sha256=NEmuHPWz4SGQcgF7QuIeA0QaXK141JoYco-7mqj9Ghk,883
+marlenv/adapters/gym_adapter.py,sha256=Vx6ZrYI7kiNlJODmqyjXu9WCdbCr6trcMNot0pvYD74,2864
+marlenv/adapters/overcooked_adapter.py,sha256=exNYFQhtnrvXQtG7M7FuQUkcRJK60a8vu2uwhRLc0K4,5768
 marlenv/adapters/pettingzoo_adapter.py,sha256=9rwSc_b7qV3ChtEIevOkJvtIp7WoY3CVnu6L9DxlMB4,2852
 marlenv/adapters/pymarl_adapter.py,sha256=x__E90XpFbfSWhnBHtkcD6WYkmKki1LByNbUFoDBUcg,3416
 marlenv/adapters/smac_adapter.py,sha256=fOfKo1hL4ioKtM5qQGcwtfdkdwUEACjAZqaGmkoQUcU,8373
 marlenv/models/__init__.py,sha256=9M-rnj94nsdyO4zm_VEtyYBmde3iD2_eIY4bMB-IqCo,555
-marlenv/models/env.py,sha256=JOT00dFXRuRMQYrv6ZYb3siwNh9eMXKnGRw_dM_mgVU,6519
+marlenv/models/env.py,sha256=faezAKOIccBauOFeo9wu5sX32pFmP3AMmGyJzaTRJcM,7514
 marlenv/models/episode.py,sha256=ZGBx6lb2snrUhDgFEwHPV1dp-XvMA7k4quQVUNQxsP0,15140
-marlenv/models/observation.py,sha256=rTAesS_jaIyRlH4wjo2izEpWS0Hn5_UKjhbvdp0H4tA,2994
-marlenv/models/spaces.py,sha256=12PIoSWqgXVSfB5poqRTV4CdS-WNCrKiPmfALIGj7Mk,6226
+marlenv/models/observation.py,sha256=kAmh1hIoC2TGrZlGVzV0y4TXXCSrI7gcmG0raeoncYk,3153
+marlenv/models/spaces.py,sha256=pw8Sum_fHBkR-lyfTqUij4azMCNm8oBZrYZe4WVR7rA,7652
 marlenv/models/state.py,sha256=958PXTHadi3gtRnhGgcGtqBnF44R11kdcx62NN2gwxA,1717
 marlenv/models/step.py,sha256=LKGAV2Cu-k9Gz1hwrfvGx51l8axtQRqDE9WVL5r2A1Q,3037
 marlenv/models/transition.py,sha256=2vvuhSSq911weCXio9nuyfsLVh_7ORSU_znOqpLLdLg,5107
-marlenv/wrappers/__init__.py,sha256=E1IwrJjXGB6ZgQPiv-Huw7TGmnCSutFSkApHEi8dSQA,744
+marlenv/wrappers/__init__.py,sha256=P7YCK1KYJvE6BAlH--nOW9PSlrohhuw-1wlfgCTOl9U,808
 marlenv/wrappers/agent_id_wrapper.py,sha256=oTIAYxKD1JtHfrZN43mf-3e8pxjd0nxm07vxs3BfrGY,1187
 marlenv/wrappers/available_actions_mask.py,sha256=JoCJ9eqHlkY8wfY-oaceEi8yp1Efs1iK6IO2Ibf9oZA,1468
 marlenv/wrappers/available_actions_wrapper.py,sha256=9UTwP3LXvncBITJeQnEqwiP_lj-ipULACkGs-2QbMrI,1026
 marlenv/wrappers/blind_wrapper.py,sha256=YEayRf_dclhzx6LXsasZ-IM7C71kyPb1gV0pHYYfjig,857
 marlenv/wrappers/centralised.py,sha256=J4hOMRT2fit936LifANNJtP7UbBEb_xIyF4VL9-fZGw,3226
+marlenv/wrappers/delayed_rewards.py,sha256=6oGJe-L_gGI-pQMResbkjsMDvXpni2SQvnTQ6wsZqGo,1170
 marlenv/wrappers/last_action_wrapper.py,sha256=u7a3Da5sg_gMrwZ3SE7PAwt2m9xSYYDKjngQyOmcJ74,2886
 marlenv/wrappers/paddings.py,sha256=VQOF4zaP61R74tQ4XTTT-FkK2QSy31AukICnqCy6zB0,2119
 marlenv/wrappers/penalty_wrapper.py,sha256=v4_H8OEN2-yujLzRb6P7W7KwmXHtjAFsxcdp3SbnKpo,996
 marlenv/wrappers/rlenv_wrapper.py,sha256=C2XekgBIM4x3Wa2Mtsn7rihRD4ymC2hORI473Af0sfw,2962
 marlenv/wrappers/time_limit.py,sha256=CDIMMJPMyIDHSFxUJaC7nb7Kd86-07NgZeFhrpZm82o,3985
 marlenv/wrappers/video_recorder.py,sha256=d5AFu6qHqby9mOcBsYWYPxAPiK1vtnfMYdZ81AnCekI,2624
-multi_agent_rlenv-3.2.2.dist-info/METADATA,sha256=EX_3681ZauwzYUllRCn7QaX9ejiXihzO7ib7TZJBEiY,3357
-multi_agent_rlenv-3.2.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-multi_agent_rlenv-3.2.2.dist-info/licenses/LICENSE,sha256=_eeiGVoIJ7kYt6l1zbIvSBQppTnw0mjnYk1lQ4FxEjE,1074
-multi_agent_rlenv-3.2.2.dist-info/RECORD,,
+multi_agent_rlenv-3.3.1.dist-info/METADATA,sha256=lrE00NibOxZAoMD_DacNbfIbpEHqX7zw5wXc0xj8iiY,4897
+multi_agent_rlenv-3.3.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+multi_agent_rlenv-3.3.1.dist-info/licenses/LICENSE,sha256=_eeiGVoIJ7kYt6l1zbIvSBQppTnw0mjnYk1lQ4FxEjE,1074
+multi_agent_rlenv-3.3.1.dist-info/RECORD,,