multi-agent-rlenv 3.3.7__py3-none-any.whl → 3.5.0__py3-none-any.whl

marlenv/__init__.py

@@ -62,16 +62,11 @@ print(env.extras_shape) # (1, )
 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.3.7"
+__version__ = "3.5.0"
 
 from . import models
-from . import wrappers
-from . import adapters
-from .models import spaces
-
-
-from .env_builder import make, Builder
 from .models import (
+    spaces,
     MARLEnv,
     State,
     Step,
@@ -80,10 +75,14 @@ from .models import (
     Transition,
     DiscreteSpace,
     ContinuousSpace,
-    ActionSpace,
-    DiscreteActionSpace,
-    ContinuousActionSpace,
+    Space,
+    MultiDiscreteSpace,
 )
+
+
+from . import wrappers
+from . import adapters
+from .env_builder import make, Builder
 from .wrappers import RLEnvWrapper
 from .mock_env import DiscreteMockEnv, DiscreteMOMockEnv
 
@@ -100,12 +99,11 @@ __all__ = [
     "Observation",
     "Episode",
     "Transition",
-    "ActionSpace",
     "DiscreteSpace",
     "ContinuousSpace",
-    "DiscreteActionSpace",
-    "ContinuousActionSpace",
     "DiscreteMockEnv",
     "DiscreteMOMockEnv",
     "RLEnvWrapper",
+    "Space",
+    "MultiDiscreteSpace",
 ]

marlenv/adapters/gym_adapter.py

@@ -1,26 +1,16 @@
 import sys
-import cv2
 from dataclasses import dataclass
-from typing import Sequence
 
+import cv2
 import gymnasium as gym
 import numpy as np
-import numpy.typing as npt
 from gymnasium import Env, spaces
 
-from marlenv.models import (
-    ActionSpace,
-    ContinuousActionSpace,
-    DiscreteActionSpace,
-    MARLEnv,
-    Observation,
-    State,
-    Step,
-)
+from marlenv import ContinuousSpace, DiscreteSpace, MARLEnv, Observation, Space, State, Step
 
 
 @dataclass
-class Gym(MARLEnv[Sequence | npt.NDArray, ActionSpace]):
+class Gym(MARLEnv[Space]):
     """Wraps a gym envronment in an RLEnv"""
 
     def __init__(self, env: Env | str, **kwargs):
@@ -30,7 +20,7 @@ class Gym(MARLEnv[Sequence | npt.NDArray, ActionSpace]):
             raise NotImplementedError("Observation space must have a shape")
         match env.action_space:
             case spaces.Discrete() as s:
-                space = DiscreteActionSpace(1, int(s.n))
+                space = DiscreteSpace(int(s.n), labels=[f"Action {i}" for i in range(s.n)]).repeat(1)
             case spaces.Box() as s:
                 low = s.low.astype(np.float32)
                 high = s.high.astype(np.float32)
@@ -38,10 +28,10 @@ class Gym(MARLEnv[Sequence | npt.NDArray, ActionSpace]):
                     low = np.full(s.shape, s.low, dtype=np.float32)
                 if not isinstance(high, np.ndarray):
                     high = np.full(s.shape, s.high, dtype=np.float32)
-                space = ContinuousActionSpace(1, low, high)
+                space = ContinuousSpace(low, high, labels=[f"Action {i}" for i in range(s.shape[0])]).repeat(1)
             case other:
                 raise NotImplementedError(f"Action space {other} not supported")
-        super().__init__(space, env.observation_space.shape, (1,))
+        super().__init__(1, space, env.observation_space.shape, (1,))
         self._gym_env = env
         if self._gym_env.unwrapped.spec is not None:
             self.name = self._gym_env.unwrapped.spec.id

marlenv/adapters/overcooked_adapter.py

@@ -1,14 +1,14 @@
 import sys
 from dataclasses import dataclass
-from typing import Literal, Sequence
+from typing import Literal, Sequence, Optional
 from copy import deepcopy
-from time import time
 
 import cv2
 import numpy as np
 import numpy.typing as npt
 import pygame
-from marlenv.models import ContinuousSpace, DiscreteActionSpace, MARLEnv, Observation, State, Step
+from marlenv.models import ContinuousSpace, DiscreteSpace, MARLEnv, Observation, State, Step, MultiDiscreteSpace
+from marlenv.utils import Schedule
 
 from overcooked_ai_py.mdp.overcooked_env import OvercookedEnv
 from overcooked_ai_py.mdp.overcooked_mdp import Action, OvercookedGridworld, OvercookedState
@@ -16,12 +16,19 @@ from overcooked_ai_py.visualization.state_visualizer import StateVisualizer
 
 
 @dataclass
-class Overcooked(MARLEnv[Sequence[int] | npt.NDArray, DiscreteActionSpace]):
+class Overcooked(MARLEnv[MultiDiscreteSpace]):
     horizon: int
-    reward_shaping: bool
+    shaping_factor: Schedule
 
-    def __init__(self, oenv: OvercookedEnv, reward_shaping: bool = True):
-        self.reward_shaping = reward_shaping
+    def __init__(
+        self,
+        oenv: OvercookedEnv,
+        shaping_factor: float | Schedule = 1.0,
+        name_suffix: Optional[str] = None,
+    ):
+        if isinstance(shaping_factor, (int, float)):
+            shaping_factor = Schedule.constant(shaping_factor)
+        self.shaping_factor = shaping_factor
         self._oenv = oenv
         assert isinstance(oenv.mdp, OvercookedGridworld)
         self._mdp = oenv.mdp
@@ -30,10 +37,9 @@ class Overcooked(MARLEnv[Sequence[int] | npt.NDArray, DiscreteActionSpace]):
         # -1 because we extract the "urgent" layer to the extras
         shape = (int(layers - 1), int(width), int(height))
         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],
+            n_agents=self._mdp.num_players,
+            action_space=DiscreteSpace(Action.NUM_ACTIONS, labels=[Action.ACTION_TO_CHAR[a] for a in Action.ALL_ACTIONS]).repeat(
+                self._mdp.num_players
             ),
             observation_shape=shape,
             extras_shape=(2,),
@@ -43,6 +49,8 @@ class Overcooked(MARLEnv[Sequence[int] | npt.NDArray, DiscreteActionSpace]):
             reward_space=ContinuousSpace.from_shape(1),
         )
         self.horizon = int(self._oenv.horizon)
+        if name_suffix is not None:
+            self.name = f"{self.name}-{name_suffix}"
 
     @property
     def state(self) -> OvercookedState:
@@ -86,11 +94,12 @@ class Overcooked(MARLEnv[Sequence[int] | npt.NDArray, DiscreteActionSpace]):
                 available_actions[agent_num, Action.ACTION_TO_INDEX[action]] = True
         return np.array(available_actions, dtype=np.bool)
 
-    def step(self, actions: Sequence[int] | npt.NDArray[np.int32 | np.int64]) -> Step:
+    def step(self, actions: Sequence[int] | np.ndarray) -> Step:
+        self.shaping_factor.update()
         actions = [Action.ALL_ACTIONS[a] for a in actions]
         _, reward, done, info = self._oenv.step(actions, display_phi=True)
-        if self.reward_shaping:
-            reward += sum(info["shaped_r_by_agent"])
+
+        reward += sum(info["shaped_r_by_agent"]) * self.shaping_factor
         return Step(
             obs=self.get_observation(),
             state=self.get_state(),
@@ -104,19 +113,25 @@ class Overcooked(MARLEnv[Sequence[int] | npt.NDArray, DiscreteActionSpace]):
         self._oenv.reset()
         return self.get_observation(), self.get_state()
 
-    def __deepcopy__(self, memo: dict):
+    def __deepcopy__(self, _):
+        """
+        Note: a specific implementation is needed because `pygame.font.Font` objects are not deep-copiable by default.
+        """
         mdp = deepcopy(self._mdp)
-        return Overcooked(OvercookedEnv.from_mdp(mdp, horizon=self.horizon))
+        copy = Overcooked(OvercookedEnv.from_mdp(mdp, horizon=self.horizon), deepcopy(self.shaping_factor))
+        copy.name = self.name
+        return copy
 
     def __getstate__(self):
-        return {"horizon": self.horizon, "mdp": self._mdp}
+        return {"horizon": self.horizon, "mdp": self._mdp, "name": self.name, "schedule": self.shaping_factor}
 
     def __setstate__(self, state: dict):
         from overcooked_ai_py.mdp.overcooked_mdp import Recipe
 
         mdp = state["mdp"]
         Recipe.configure(mdp.recipe_config)
-        self.__init__(OvercookedEnv.from_mdp(state["mdp"], horizon=state["horizon"]))
+        self.__init__(OvercookedEnv.from_mdp(state["mdp"], horizon=state["horizon"]), shaping_factor=state["schedule"])
+        self.name = state["name"]
 
     def get_image(self):
         rewards_dict = {}  # dictionary of details you want rendered in the UI
@@ -190,16 +205,17 @@ class Overcooked(MARLEnv[Sequence[int] | npt.NDArray, DiscreteActionSpace]):
             "you_shall_not_pass",
         ],
         horizon: int = 400,
-        reward_shaping: bool = True,
+        reward_shaping_factor: float | Schedule = 1.0,
     ):
         mdp = OvercookedGridworld.from_layout_name(layout)
-        return Overcooked(OvercookedEnv.from_mdp(mdp, horizon=horizon), reward_shaping=reward_shaping)
+        return Overcooked(OvercookedEnv.from_mdp(mdp, horizon=horizon, info_level=0), reward_shaping_factor, layout)
 
     @staticmethod
     def from_grid(
         grid: Sequence[Sequence[Literal["S", "P", "X", "O", "D", "T", "1", "2", " "] | str]],
         horizon: int = 400,
-        reward_shaping: bool = True,
+        shaping_factor: float | Schedule = 1.0,
+        layout_name: Optional[str] = None,
     ):
         """
         Create an Overcooked environment from a grid layout where
@@ -212,10 +228,14 @@ class Overcooked(MARLEnv[Sequence[int] | npt.NDArray, DiscreteActionSpace]):
         - 1 is a player 1 starting location
         - 2 is a player 2 starting location
         - ' ' is a walkable space
+
+        If provided, `custom_name` is added to the environment name.
         """
         # It is necessary to add an explicit layout name because Overcooked saves some files under this
         # name. By default the name is a concatenation of the grid elements, which may include characters
         # such as white spaces, pipes ('|') and square brackets ('[' and ']') that are invalid Windows file paths.
-        layout_name = str(time())
+        if layout_name is None:
+            layout_name = "custom-layout"
         mdp = OvercookedGridworld.from_grid(grid, base_layout_params={"layout_name": layout_name})
-        return Overcooked(OvercookedEnv.from_mdp(mdp, horizon=horizon), reward_shaping=reward_shaping)
+
+        return Overcooked(OvercookedEnv.from_mdp(mdp, horizon=horizon, info_level=0), shaping_factor, layout_name)

marlenv/adapters/pettingzoo_adapter.py

@@ -6,17 +6,17 @@ import numpy.typing as npt
 from gymnasium import spaces  # pettingzoo uses gymnasium spaces
 from pettingzoo import ParallelEnv
 
-from marlenv.models import ActionSpace, ContinuousActionSpace, DiscreteActionSpace, MARLEnv, Observation, State, Step
+from marlenv.models import MARLEnv, Observation, State, Step, DiscreteSpace, ContinuousSpace, Space
 
 
 @dataclass
-class PettingZoo(MARLEnv[npt.NDArray, ActionSpace]):
+class PettingZoo(MARLEnv[Space]):
     def __init__(self, env: ParallelEnv):
         aspace = env.action_space(env.possible_agents[0])
         n_agents = len(env.possible_agents)
         match aspace:
             case spaces.Discrete() as s:
-                space = DiscreteActionSpace(n_agents, int(s.n))
+                space = DiscreteSpace.action(int(s.n)).repeat(n_agents)
 
             case spaces.Box() as s:
                 low = s.low.astype(np.float32)
@@ -25,7 +25,7 @@ class PettingZoo(MARLEnv[npt.NDArray, ActionSpace]):
                     low = np.full(s.shape, s.low, dtype=np.float32)
                 if not isinstance(high, np.ndarray):
                     high = np.full(s.shape, s.high, dtype=np.float32)
-                space = ContinuousActionSpace(n_agents, low, high=high)
+                space = ContinuousSpace(low, high=high).repeat(n_agents)
             case other:
                 raise NotImplementedError(f"Action space {other} not supported")
 
@@ -34,7 +34,7 @@ class PettingZoo(MARLEnv[npt.NDArray, ActionSpace]):
             raise NotImplementedError("Only discrete observation spaces are supported")
         self._pz_env = env
         env.reset()
-        super().__init__(space, obs_space.shape, self.get_state().shape)
+        super().__init__(n_agents, space, obs_space.shape, self.get_state().shape)
         self.agents = env.possible_agents
         self.last_observation = None
 

marlenv/adapters/pymarl_adapter.py

@@ -1,10 +1,9 @@
 from dataclasses import dataclass
-from typing import Any, Sequence
+from typing import Any
 
 import numpy as np
-import numpy.typing as npt
 
-from marlenv.models import DiscreteActionSpace, MARLEnv
+from marlenv.models import MARLEnv, MultiDiscreteSpace
 from marlenv.wrappers import TimeLimit
 
 
@@ -15,7 +14,7 @@ class PymarlAdapter:
     with the pymarl-qplex code base.
     """
 
-    def __init__(self, env: MARLEnv[Sequence | npt.NDArray, DiscreteActionSpace], episode_limit: int):
+    def __init__(self, env: MARLEnv[MultiDiscreteSpace], episode_limit: int):
         assert env.reward_space.size == 1, "Only single objective environments are supported."
         self.env = TimeLimit(env, episode_limit, add_extra=False)
         # Required by PyMarl

marlenv/adapters/smac_adapter.py

@@ -1,15 +1,15 @@
 from dataclasses import dataclass
-from typing import Sequence, overload
+from typing import overload
 
 import numpy as np
 import numpy.typing as npt
 from smac.env import StarCraft2Env
 
-from marlenv.models import DiscreteActionSpace, MARLEnv, Observation, State, Step
+from marlenv.models import MARLEnv, Observation, State, Step, MultiDiscreteSpace, DiscreteSpace
 
 
 @dataclass
-class SMAC(MARLEnv[Sequence[int] | npt.NDArray, DiscreteActionSpace]):
+class SMAC(MARLEnv[MultiDiscreteSpace]):
     """Wrapper for the SMAC environment to work with this framework"""
 
     @overload
@@ -157,10 +157,10 @@ class SMAC(MARLEnv[Sequence[int] | npt.NDArray, DiscreteActionSpace]):
             case other:
                 raise ValueError(f"Invalid argument type: {type(other)}")
         self._env = StarCraft2Env(map_name=map_name)
-        action_space = DiscreteActionSpace(self._env.n_agents, self._env.n_actions)
         self._env_info = self._env.get_env_info()
         super().__init__(
-            action_space=action_space,
+            self._env.n_agents,
+            action_space=DiscreteSpace(self._env.n_actions).repeat(self._env.n_agents),
             observation_shape=(self._env_info["obs_shape"],),
             state_shape=(self._env_info["state_shape"],),
         )
@@ -195,7 +195,7 @@ class SMAC(MARLEnv[Sequence[int] | npt.NDArray, DiscreteActionSpace]):
         )
         return step
 
-    def available_actions(self) -> npt.NDArray[np.bool_]:
+    def available_actions(self) -> npt.NDArray[np.bool]:
         return np.array(self._env.get_avail_actions()) == 1
 
     def get_image(self):

marlenv/env_builder.py

@@ -5,10 +5,9 @@ import numpy.typing as npt
 
 from . import wrappers
 from marlenv import adapters
-from .models import ActionSpace, MARLEnv
+from .models import Space, MARLEnv
 
-A = TypeVar("A")
-AS = TypeVar("AS", bound=ActionSpace)
+AS = TypeVar("AS", bound=Space)
 
 if adapters.HAS_PETTINGZOO:
     from .adapters import PettingZoo
@@ -71,12 +70,12 @@ def make(env, **kwargs):
 
 
 @dataclass
-class Builder(Generic[A, AS]):
+class Builder(Generic[AS]):
     """Builder for environments"""
 
-    _env: MARLEnv[A, AS]
+    _env: MARLEnv[AS]
 
-    def __init__(self, env: MARLEnv[A, AS]):
+    def __init__(self, env: MARLEnv[AS]):
         self._env = env
 
     def time_limit(self, n_steps: int, add_extra: bool = True, truncation_penalty: Optional[float] = None):
@@ -124,9 +123,9 @@ class Builder(Generic[A, AS]):
 
     def centralised(self):
         """Centralises the observations and actions"""
-        from marlenv.models import DiscreteActionSpace
+        from marlenv.models import MultiDiscreteSpace
 
-        assert isinstance(self._env.action_space, DiscreteActionSpace)
+        assert isinstance(self._env.action_space, MultiDiscreteSpace)
         self._env = wrappers.Centralized(self._env)  # type: ignore
         return self
 
@@ -159,6 +158,6 @@ class Builder(Generic[A, AS]):
         self._env = wrappers.TimePenalty(self._env, penalty)
         return self
 
-    def build(self) -> MARLEnv[A, AS]:
+    def build(self) -> MARLEnv[AS]:
         """Build and return the environment"""
         return self._env

marlenv/env_pool.py

@@ -1,21 +1,19 @@
 from typing import Sequence
 from dataclasses import dataclass
-import numpy.typing as npt
 from typing_extensions import TypeVar
 import random
 
 from marlenv import RLEnvWrapper, MARLEnv
-from marlenv.models import ActionSpace
+from marlenv.models import Space
 
-ActionType = TypeVar("ActionType", default=npt.NDArray)
-ActionSpaceType = TypeVar("ActionSpaceType", bound=ActionSpace, default=ActionSpace)
+ActionSpaceType = TypeVar("ActionSpaceType", bound=Space, default=Space)
 
 
 @dataclass
-class EnvPool(RLEnvWrapper[ActionType, ActionSpaceType]):
-    envs: Sequence[MARLEnv[ActionType, ActionSpaceType]]
+class EnvPool(RLEnvWrapper[ActionSpaceType]):
+    envs: Sequence[MARLEnv[ActionSpaceType]]
 
-    def __init__(self, envs: Sequence[MARLEnv[ActionType, ActionSpaceType]]):
+    def __init__(self, envs: Sequence[MARLEnv[ActionSpaceType]]):
         assert len(envs) > 0, "EnvPool must contain at least one environment"
         self.envs = envs
         for env in envs[1:]:

marlenv/mock_env.py

@@ -1,12 +1,10 @@
-from typing import Sequence
 import numpy as np
-import numpy.typing as npt
 from dataclasses import dataclass
-from marlenv import MARLEnv, Observation, DiscreteActionSpace, ContinuousSpace, Step, State
+from marlenv import MARLEnv, Observation, ContinuousSpace, Step, State, DiscreteSpace, MultiDiscreteSpace
 
 
 @dataclass
-class DiscreteMockEnv(MARLEnv[Sequence[int] | npt.NDArray, DiscreteActionSpace]):
+class DiscreteMockEnv(MARLEnv[MultiDiscreteSpace]):
     def __init__(
         self,
         n_agents: int = 4,
@@ -27,7 +25,8 @@ class DiscreteMockEnv(MARLEnv[Sequence[int] | npt.NDArray, DiscreteActionSpace])
             case _:
                 raise ValueError("reward_step must be an int, float or np.ndarray")
         super().__init__(
-            DiscreteActionSpace(n_agents, n_actions),
+            n_agents,
+            DiscreteSpace(n_actions).repeat(n_agents),
             (obs_size,),
             (n_agents * agent_state_size,),
             extras_shape=(extras_size,),
@@ -85,7 +84,7 @@ class DiscreteMockEnv(MARLEnv[Sequence[int] | npt.NDArray, DiscreteActionSpace])
         )
 
 
-class DiscreteMOMockEnv(MARLEnv[Sequence[int] | npt.NDArray, DiscreteActionSpace]):
+class DiscreteMOMockEnv(MARLEnv[DiscreteSpace]):
     """Multi-Objective Mock Environment"""
 
     def __init__(
@@ -100,7 +99,8 @@ class DiscreteMOMockEnv(MARLEnv[Sequence[int] | npt.NDArray, DiscreteActionSpace
         extras_size: int = 0,
     ) -> None:
         super().__init__(
-            DiscreteActionSpace(n_agents, n_actions),
+            n_agents,
+            DiscreteSpace(n_actions),
             (obs_size,),
             (n_agents * agent_state_size,),
             extras_shape=(extras_size,),

marlenv/models/__init__.py

@@ -1,4 +1,4 @@
-from .spaces import ActionSpace, DiscreteSpace, ContinuousSpace, MultiDiscreteSpace, DiscreteActionSpace, ContinuousActionSpace
+from .spaces import DiscreteSpace, ContinuousSpace, MultiDiscreteSpace, Space
 from .observation import Observation
 from .step import Step
 from .state import State
@@ -8,7 +8,6 @@ from .episode import Episode
 
 
 __all__ = [
-    "ActionSpace",
     "Step",
     "State",
     "DiscreteSpace",
@@ -18,6 +17,5 @@ __all__ = [
     "Transition",
     "Episode",
     "MultiDiscreteSpace",
-    "DiscreteActionSpace",
-    "ContinuousActionSpace",
+    "Space",
 ]

marlenv/models/env.py

@@ -1,24 +1,22 @@
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from itertools import product
-from typing import Generic, Optional, Sequence
+from typing import Generic, Optional, Sequence, TypeVar
 
 import cv2
 import numpy as np
 import numpy.typing as npt
-from typing_extensions import TypeVar
 
 from .observation import Observation
-from .spaces import ActionSpace, ContinuousSpace, Space
+from .spaces import ContinuousSpace, Space, DiscreteSpace, MultiDiscreteSpace
 from .state import State
 from .step import Step
 
-ActionType = TypeVar("ActionType", default=npt.NDArray)
-ActionSpaceType = TypeVar("ActionSpaceType", bound=ActionSpace, default=ActionSpace)
+ActionSpaceType = TypeVar("ActionSpaceType", bound=Space)
 
 
 @dataclass
-class MARLEnv(ABC, Generic[ActionType, ActionSpaceType]):
+class MARLEnv(ABC, Generic[ActionSpaceType]):
     """
     Multi-Agent Reinforcement Learning environment.
 
@@ -70,6 +68,7 @@ class MARLEnv(ABC, Generic[ActionType, ActionSpaceType]):
 
     def __init__(
         self,
+        n_agents: int,
         action_space: ActionSpaceType,
         observation_shape: tuple[int, ...],
         state_shape: tuple[int, ...],
@@ -81,8 +80,8 @@ class MARLEnv(ABC, Generic[ActionType, ActionSpaceType]):
         super().__init__()
         self.name = self.__class__.__name__
         self.action_space = action_space
-        self.n_actions = action_space.n_actions
-        self.n_agents = action_space.n_agents
+        self.n_actions = action_space.shape[-1]
+        self.n_agents = n_agents
         self.observation_shape = observation_shape
         self.state_shape = state_shape
         self.extras_shape = extras_shape
@@ -108,9 +107,21 @@ class MARLEnv(ABC, Generic[ActionType, ActionSpaceType]):
         """Whether the environment is multi-objective."""
         return self.reward_space.size > 1
 
-    def sample_action(self) -> ActionType:
+    @property
+    def n_objectives(self) -> int:
+        """The number of objectives in the environment."""
+        return self.reward_space.size
+
+    def sample_action(self):
         """Sample an available action from the action space."""
-        return self.action_space.sample(self.available_actions())  # type: ignore
+        match self.action_space:
+            case MultiDiscreteSpace() as aspace:
+                return aspace.sample(mask=self.available_actions())
+            case ContinuousSpace() as aspace:
+                return aspace.sample()
+            case DiscreteSpace() as aspace:
+                return np.array([aspace.sample(mask=self.available_actions())])
+        raise NotImplementedError("Action space not supported")
 
     def available_actions(self) -> npt.NDArray[np.bool]:
         """
@@ -142,7 +153,7 @@ class MARLEnv(ABC, Generic[ActionType, ActionSpaceType]):
         raise NotImplementedError("Method not implemented")
 
     @abstractmethod
-    def step(self, actions: ActionType) -> Step:
+    def step(self, action: Sequence | np.ndarray) -> Step:
         """Perform a step in the environment.
 
         Returns a Step object that can be unpacked as a 6-tuple containing:
@@ -175,7 +186,7 @@ class MARLEnv(ABC, Generic[ActionType, ActionSpaceType]):
         """Retrieve an image of the environment"""
         raise NotImplementedError("No image available for this environment")
 
-    def replay(self, actions: Sequence[ActionType], seed: Optional[int] = None):
+    def replay(self, actions: Sequence, seed: Optional[int] = None):
         """Replay a sequence of actions."""
         from .episode import Episode  # Avoid circular import