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

marlenv/models/spaces.py

@@ -1,13 +1,11 @@
 import math
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
-from typing import Any, Generic, Optional, TypeVar
+from typing import Optional
 
 import numpy as np
 import numpy.typing as npt
 
-S = TypeVar("S", bound="Space")
-
 
 @dataclass
 class Space(ABC):
@@ -23,7 +21,7 @@ class Space(ABC):
         self.labels = labels
 
     @abstractmethod
-    def sample(self, mask: Optional[npt.NDArray[np.bool_]] = None) -> Any:
+    def sample(self, mask: Optional[npt.NDArray[np.bool_]] = None) -> npt.NDArray[np.float32]:
         """Sample a value from the space."""
 
     def __eq__(self, value: object) -> bool:
@@ -34,6 +32,16 @@ class Space(ABC):
     def __ne__(self, value: object) -> bool:
         return not self.__eq__(value)
 
+    @property
+    @abstractmethod
+    def is_discrete(self) -> bool:
+        """Whether the space is discrete."""
+
+    @property
+    def is_continuous(self) -> bool:
+        """Whether the space is continuous."""
+        return not self.is_discrete
+
 
 @dataclass
 class DiscreteSpace(Space):
@@ -45,8 +53,8 @@ class DiscreteSpace(Space):
         self.size = size
         self.space = np.arange(size)
 
-    def sample(self, mask: Optional[npt.NDArray[np.bool_]] = None) -> int:
-        space = self.space
+    def sample(self, mask: Optional[npt.NDArray[np.bool]] = None):
+        space = self.space.copy()
         if mask is not None:
             space = space[mask]
         return int(np.random.choice(space))
@@ -58,6 +66,25 @@ class DiscreteSpace(Space):
             return False
         return super().__eq__(value)
 
+    @property
+    def is_discrete(self) -> bool:
+        return True
+
+    @staticmethod
+    def action(size, labels: Optional[list[str]] = None):
+        """
+        Create a discrete action space where the default labels are set to "Action-n".
+        """
+        if labels is None:
+            labels = [f"Action {i}" for i in range(size)]
+        return DiscreteSpace(size, labels)
+
+    def repeat(self, n: int):
+        """
+        Repeat the discrete space n times.
+        """
+        return MultiDiscreteSpace(*([self] * n), labels=self.labels)
+
 
 @dataclass
 class MultiDiscreteSpace(Space):
@@ -75,10 +102,10 @@ class MultiDiscreteSpace(Space):
     def from_sizes(cls, *sizes: int):
         return cls(*(DiscreteSpace(size) for size in sizes))
 
-    def sample(self, masks: Optional[npt.NDArray[np.bool_] | list[npt.NDArray[np.bool_]]] = None):
-        if masks is None:
+    def sample(self, mask: Optional[npt.NDArray[np.bool] | list[npt.NDArray[np.bool]]] = None):
+        if mask is None:
             return np.array([space.sample() for space in self.spaces], dtype=np.int32)
-        return np.array([space.sample(mask) for mask, space in zip(masks, self.spaces)], dtype=np.int32)
+        return np.array([space.sample(mask=mask) for mask, space in zip(mask, self.spaces)], dtype=np.int32)
 
     def __eq__(self, value: object) -> bool:
         if not isinstance(value, MultiDiscreteSpace):
@@ -90,6 +117,10 @@ class MultiDiscreteSpace(Space):
                 return False
         return super().__eq__(value)
 
+    @property
+    def is_discrete(self) -> bool:
+        return True
+
 
 @dataclass
 class ContinuousSpace(Space):
@@ -100,23 +131,35 @@ 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],
+    def __init__(
+        self,
+        low: int | float | list | npt.NDArray[np.float32] | None,
+        high: int | float | list | npt.NDArray[np.float32] | None,
         labels: Optional[list[str]] = None,
     ):
         match low:
+            case None:
+                assert high is not None, "If low is None, high must be set to infer the shape."
+                shape = ContinuousSpace.get_shape(high)
+                low = np.full(shape, -np.inf, dtype=np.float32)
             case list():
                 low = np.array(low, dtype=np.float32)
             case float() | int():
                 low = np.array([low], dtype=np.float32)
         match high:
+            case None:
+                assert low is not None, "If high is None, low must be set to infer the shape."
+                shape = ContinuousSpace.get_shape(low)
+                high = np.full(shape, np.inf, dtype=np.float32)
             case list():
                 high = np.array(high, dtype=np.float32)
             case float() | int():
                 high = np.array([high], dtype=np.float32)
-        return ContinuousSpace(low, high, labels)
+        assert low.shape == high.shape, f"Low and high must have the same shape. Low shape: {low.shape}, high shape: {high.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)
+        self.low = low
+        self.high = high
 
     @staticmethod
     def from_shape(
@@ -143,20 +186,24 @@ class ContinuousSpace(Space):
                 high = np.array(high, dtype=np.float32)
         return ContinuousSpace(low, high, labels)
 
-    def __init__(
-        self,
-        low: npt.NDArray[np.float32],
-        high: npt.NDArray[np.float32],
-        labels: Optional[list[str]] = None,
-    ):
-        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)
-        self.low = low
-        self.high = high
+    def clamp(self, action: np.ndarray | list):
+        """Clamp the action to the bounds of the space."""
+        if isinstance(action, list):
+            action = np.array(action)
+        return np.clip(action, self.low, self.high)
+
+    def sample(self) -> npt.NDArray[np.float32]:
+        r = np.random.random(self.shape) * (self.high - self.low) + self.low
+        return r.astype(np.float32)
 
-    def sample(self, *_):
-        return np.random.random(self.shape) * (self.high - self.low) + self.low
+    @staticmethod
+    def get_shape(item: float | int | list | npt.NDArray[np.float32]) -> tuple[int, ...]:
+        """Get the shape of the item."""
+        if isinstance(item, list):
+            item = np.array(item)
+        if isinstance(item, np.ndarray):
+            return item.shape
+        return (1,)
 
     def __eq__(self, value: object) -> bool:
         if not isinstance(value, ContinuousSpace):
@@ -167,59 +214,19 @@ class ContinuousSpace(Space):
             return False
         return super().__eq__(value)
 
-
-@dataclass
-class ActionSpace(Space, Generic[S]):
-    n_agents: int
-    """Number of agents."""
-    action_names: list[str]
-    """The meaning of each action."""
-    n_actions: int
-    individual_action_space: S
-
-    def __init__(self, n_agents: int, individual_action_space: S, action_names: Optional[list] = None):
-        Space.__init__(self, (n_agents, *individual_action_space.shape), action_names)
-        self.n_agents = n_agents
-        self.individual_action_space = individual_action_space
-        self.n_actions = math.prod(individual_action_space.shape)
-        self.action_names = action_names or [f"Action {i}" for i in range(self.n_actions)]
-
-    def sample(self, mask: np.ndarray | None = None):
-        res = []
-        for i in range(self.n_agents):
-            if mask is not None:
-                m = mask[i]
-            else:
-                m = None
-            res.append(self.individual_action_space.sample(m))
-        return np.array(res)
-
-    def __eq__(self, value: object) -> bool:
-        if not isinstance(value, ActionSpace):
-            return False
-        if self.n_agents != value.n_agents:
-            return False
-        if self.n_actions != value.n_actions:
-            return False
-        if self.individual_action_space != value.individual_action_space:
-            return False
-        return super().__eq__(value)
-
-
-@dataclass
-class DiscreteActionSpace(ActionSpace[DiscreteSpace]):
-    def __init__(self, n_agents: int, n_actions: int, action_names: Optional[list[str]] = None):
-        individual_action_space = DiscreteSpace(n_actions, action_names)
-        super().__init__(n_agents, individual_action_space, action_names)
-
-
-@dataclass
-class MultiDiscreteActionSpace(ActionSpace[MultiDiscreteSpace]):
-    pass
-
-
-@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.from_bounds(low, high, action_names)
-        super().__init__(n_agents, space, action_names)
+    def repeat(self, n: int):
+        """
+        Repeat the continuous space n times to become of shape (n, *shape).
+        """
+        low = np.tile(self.low, (n, 1))
+        high = np.tile(self.high, (n, 1))
+        return ContinuousSpace.from_shape(
+            (n, *self.shape),
+            low=low,
+            high=high,
+            labels=self.labels,
+        )
+
+    @property
+    def is_discrete(self) -> bool:
+        return False

marlenv/models/transition.py

@@ -1,6 +1,5 @@
 from dataclasses import dataclass
-from typing import Any, Generic, Sequence
-from typing_extensions import TypeVar
+from typing import Any, Sequence
 
 import numpy as np
 import numpy.typing as npt
@@ -10,16 +9,13 @@ from .state import State
 from .step import Step
 
 
-A = TypeVar("A", default=np.ndarray)
-
-
 @dataclass
-class Transition(Generic[A]):
+class Transition:
     """Transition model"""
 
     obs: Observation
     state: State
-    action: A
+    action: np.ndarray
     reward: npt.NDArray[np.float32]
     done: bool
     info: dict[str, Any]
@@ -32,7 +28,7 @@ class Transition(Generic[A]):
         self,
         obs: Observation,
         state: State,
-        action: A,
+        action: np.ndarray | Sequence[float],
         reward: npt.NDArray[np.float32] | float | Sequence[float],
         done: bool,
         info: dict[str, Any],
@@ -65,14 +61,14 @@ class Transition(Generic[A]):
     def from_step(
         prev_obs: Observation,
         prev_state: State,
-        actions: A,
+        action: np.ndarray | Sequence[float],
         step: Step,
         **kwargs,
     ):
         return Transition(
             obs=prev_obs,
             state=prev_state,
-            action=actions,
+            action=action,
             reward=step.reward,
             done=step.done,
             info=step.info,

marlenv/wrappers/agent_id_wrapper.py

@@ -1,19 +1,18 @@
 import numpy as np
-from marlenv.models import MARLEnv, ActionSpace
+from marlenv.models import MARLEnv, Space
 from dataclasses import dataclass
 from .rlenv_wrapper import RLEnvWrapper
 
 from typing_extensions import TypeVar
 
-A = TypeVar("A", default=np.ndarray)
-AS = TypeVar("AS", bound=ActionSpace, default=ActionSpace)
+AS = TypeVar("AS", bound=Space, default=Space)
 
 
 @dataclass
-class AgentId(RLEnvWrapper[A, AS]):
+class AgentId(RLEnvWrapper[AS]):
     """RLEnv wrapper that adds a one-hot encoding of the agent id."""
 
-    def __init__(self, env: MARLEnv[A, AS]):
+    def __init__(self, env: MARLEnv[AS]):
         assert len(env.extras_shape) == 1, "AgentIdWrapper only works with single dimension extras"
         meanings = env.extras_meanings + [f"Agent ID-{i}" for i in range(env.n_agents)]
         super().__init__(env, extra_shape=(env.n_agents + env.extras_shape[0],), extra_meanings=meanings)

marlenv/wrappers/available_actions_mask.py

@@ -2,20 +2,19 @@ import numpy as np
 import numpy.typing as npt
 from typing_extensions import TypeVar
 from .rlenv_wrapper import MARLEnv, RLEnvWrapper
-from marlenv.models import ActionSpace
+from marlenv.models import Space
 from dataclasses import dataclass
 
-A = TypeVar("A", default=npt.NDArray)
-AS = TypeVar("AS", bound=ActionSpace, default=ActionSpace)
+AS = TypeVar("AS", bound=Space, default=Space)
 
 
 @dataclass
-class AvailableActionsMask(RLEnvWrapper[A, AS]):
+class AvailableActionsMask(RLEnvWrapper[AS]):
     """Permanently masks a subset of the available actions."""
 
     action_mask: npt.NDArray[np.bool_]
 
-    def __init__(self, env: MARLEnv[A, AS], action_mask: npt.NDArray[np.bool_]):
+    def __init__(self, env: MARLEnv[AS], action_mask: npt.NDArray[np.bool_]):
         super().__init__(env)
         assert action_mask.shape == (env.n_agents, env.n_actions), "Action mask must have shape (n_agents, n_actions)."
         n_available_action_per_agent = action_mask.sum(axis=-1)
@@ -27,8 +26,8 @@ class AvailableActionsMask(RLEnvWrapper[A, AS]):
         obs.available_actions = self.available_actions()
         return obs, state
 
-    def step(self, actions):
-        step = self.wrapped.step(actions)
+    def step(self, action):
+        step = self.wrapped.step(action)
         step.obs.available_actions = self.available_actions()
         return step
 

marlenv/wrappers/available_actions_wrapper.py

@@ -1,21 +1,19 @@
 import numpy as np
-import numpy.typing as npt
 from typing_extensions import TypeVar
-from marlenv.models import ActionSpace, MARLEnv
+from marlenv.models import Space, MARLEnv
 from .rlenv_wrapper import RLEnvWrapper
 from dataclasses import dataclass
 
 
-A = TypeVar("A", default=npt.NDArray)
-AS = TypeVar("AS", bound=ActionSpace, default=ActionSpace)
+AS = TypeVar("AS", bound=Space, default=Space)
 
 
 @dataclass
-class AvailableActions(RLEnvWrapper[A, AS]):
+class AvailableActions(RLEnvWrapper[AS]):
     """Adds the available actions (one-hot) as an extra feature to the observation."""
 
-    def __init__(self, env: MARLEnv[A, AS]):
-        meanings = env.extras_meanings + [f"{a} available" for a in env.action_space.action_names]
+    def __init__(self, env: MARLEnv[AS]):
+        meanings = env.extras_meanings + [f"{a} available" for a in env.action_space.labels]
         super().__init__(env, extra_shape=(env.extras_shape[0] + env.n_actions,), extra_meanings=meanings)
 
     def reset(self):
@@ -23,7 +21,7 @@ class AvailableActions(RLEnvWrapper[A, AS]):
         obs.add_extra(self.available_actions().astype(np.float32))
         return obs, state
 
-    def step(self, actions: A):
-        step = self.wrapped.step(actions)
+    def step(self, action):
+        step = self.wrapped.step(action)
         step.obs.add_extra(self.available_actions().astype(np.float32))
         return step

marlenv/wrappers/blind_wrapper.py

@@ -1,26 +1,24 @@
 import random
 from typing_extensions import TypeVar
 import numpy as np
-import numpy.typing as npt
 from dataclasses import dataclass
 
-from marlenv.models import MARLEnv, ActionSpace
+from marlenv.models import MARLEnv, Space
 from .rlenv_wrapper import RLEnvWrapper
 
 
-A = TypeVar("A", default=npt.NDArray)
-AS = TypeVar("AS", bound=ActionSpace, default=ActionSpace)
+AS = TypeVar("AS", bound=Space, default=Space)
 
 
 @dataclass
-class Blind(RLEnvWrapper[A, AS]):
+class Blind(RLEnvWrapper[AS]):
     p: float
 
-    def __init__(self, env: MARLEnv[A, AS], p: float | int):
+    def __init__(self, env: MARLEnv[AS], p: float | int):
         super().__init__(env)
         self.p = float(p)
 
-    def step(self, actions: A):
+    def step(self, actions):
         step = super().step(actions)
         if random.random() < self.p:
             step.obs.data = np.zeros_like(step.obs.data)

marlenv/wrappers/centralised.py

@@ -4,28 +4,26 @@ from typing import Sequence
 
 import numpy as np
 import numpy.typing as npt
-from typing_extensions import TypeVar
 
-from marlenv.models import ActionSpace, DiscreteActionSpace, DiscreteSpace, MARLEnv, Observation
+from marlenv.models import DiscreteSpace, MARLEnv, MultiDiscreteSpace, Observation
 
 from .rlenv_wrapper import RLEnvWrapper
 
-A = TypeVar("A", bound=npt.NDArray | Sequence[int] | Sequence[Sequence[float]])
-
 
 @dataclass
-class Centralized(RLEnvWrapper[A, DiscreteActionSpace]):
-    joint_action_space: ActionSpace
+class Centralized(RLEnvWrapper[MultiDiscreteSpace]):
+    joint_action_space: DiscreteSpace
 
-    def __init__(self, env: MARLEnv[A, DiscreteActionSpace]):
-        if not isinstance(env.action_space.individual_action_space, DiscreteSpace):
+    def __init__(self, env: MARLEnv[MultiDiscreteSpace]):
+        if not isinstance(env.action_space, MultiDiscreteSpace):
             raise NotImplementedError(f"Action space {env.action_space} not supported")
         joint_observation_shape = (env.observation_shape[0] * env.n_agents, *env.observation_shape[1:])
         super().__init__(
             env,
-            joint_observation_shape,
-            env.state_shape,
-            env.extras_shape,
+            n_agents=1,
+            observation_shape=joint_observation_shape,
+            state_shape=env.state_shape,
+            state_extra_shape=env.extras_shape,
             action_space=self._make_joint_action_space(env),
         )
 
@@ -37,12 +35,12 @@ class Centralized(RLEnvWrapper[A, DiscreteActionSpace]):
         obs = super().get_observation()
         return self._joint_observation(obs)
 
-    def _make_joint_action_space(self, env: MARLEnv[A, DiscreteActionSpace]):
+    def _make_joint_action_space(self, env: MARLEnv[MultiDiscreteSpace]):
         agent_actions = list[list[str]]()
         for agent in range(env.n_agents):
-            agent_actions.append([f"{agent}-{action}" for action in env.action_space.action_names])
+            agent_actions.append([f"{agent}-{action}" for action in env.action_space.labels])
         action_names = [str(a) for a in product(*agent_actions)]
-        return DiscreteActionSpace(1, env.n_actions**env.n_agents, action_names)
+        return DiscreteSpace(env.n_actions**env.n_agents, action_names).repeat(1)
 
     def step(self, actions: npt.NDArray | Sequence):
         action = actions[0]

marlenv/wrappers/delayed_rewards.py

@@ -1,20 +1,22 @@
-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
+from dataclasses import dataclass
+
+import numpy as np
+import numpy.typing as npt
+from typing_extensions import TypeVar
+
+from marlenv.models import Space
+
+from .rlenv_wrapper import MARLEnv, RLEnvWrapper
 
-A = TypeVar("A", default=npt.NDArray)
-AS = TypeVar("AS", bound=ActionSpace, default=ActionSpace)
+AS = TypeVar("AS", bound=Space, default=Space)
 
 
 @dataclass
-class DelayedReward(RLEnvWrapper[A, AS]):
+class DelayedReward(RLEnvWrapper[AS]):
     delay: int
 
-    def __init__(self, env: MARLEnv[A, AS], delay: int):
+    def __init__(self, env: MARLEnv[AS], delay: int):
         super().__init__(env)
         self.delay = delay
         self.reward_queue = deque[npt.NDArray[np.float32]](maxlen=delay + 1)
@@ -25,7 +27,7 @@ class DelayedReward(RLEnvWrapper[A, AS]):
             self.reward_queue.append(np.zeros(self.reward_space.shape, dtype=np.float32))
         return super().reset()
 
-    def step(self, actions: A):
+    def step(self, actions):
         step = super().step(actions)
         self.reward_queue.append(step.reward)
         # If the step is terminal, we sum all the remaining rewards

marlenv/wrappers/last_action_wrapper.py

@@ -1,25 +1,21 @@
 from dataclasses import dataclass
-from typing_extensions import TypeVar
-from typing import Sequence
 
 import numpy as np
 import numpy.typing as npt
+from typing_extensions import TypeVar
 
-from marlenv.models import State, ActionSpace, ContinuousActionSpace, DiscreteActionSpace
+from marlenv.models import ContinuousSpace, DiscreteSpace, MultiDiscreteSpace, Space, State
 
 from .rlenv_wrapper import MARLEnv, RLEnvWrapper
 
-AS = TypeVar("AS", bound=ActionSpace, default=ActionSpace)
-DiscreteActionType = npt.NDArray[np.int64 | np.int32] | Sequence[int]
-ContinuousActionType = npt.NDArray[np.float32] | Sequence[Sequence[float]]
-A = TypeVar("A", bound=DiscreteActionType | ContinuousActionType)
+AS = TypeVar("AS", bound=Space, default=Space)
 
 
 @dataclass
-class LastAction(RLEnvWrapper[A, AS]):
+class LastAction(RLEnvWrapper[AS]):
     """Env wrapper that adds the last action taken by the agents to the extra features."""
 
-    def __init__(self, env: MARLEnv[A, AS]):
+    def __init__(self, env: MARLEnv[AS]):
         assert len(env.extras_shape) == 1, "Adding last action is only possible with 1D extras"
         super().__init__(
             env,
@@ -37,13 +33,13 @@ class LastAction(RLEnvWrapper[A, AS]):
         state.add_extra(self.last_one_hot_actions.flatten())
         return obs, state
 
-    def step(self, actions: A):
+    def step(self, actions):
         step = super().step(actions)
         match self.wrapped.action_space:
-            case ContinuousActionSpace():
+            case ContinuousSpace():
                 self.last_actions = actions
-            case DiscreteActionSpace():
-                self.last_one_hot_actions = self.compute_one_hot_actions(actions)  # type: ignore
+            case DiscreteSpace() | MultiDiscreteSpace():
+                self.last_one_hot_actions = self.compute_one_hot_actions(actions)
             case other:
                 raise NotImplementedError(f"Action space {other} not supported")
         step.obs.add_extra(self.last_one_hot_actions)
@@ -60,7 +56,7 @@ class LastAction(RLEnvWrapper[A, AS]):
         self.last_one_hot_actions = flattened_one_hots.reshape(self.n_agents, self.n_actions)
         return super().set_state(state)
 
-    def compute_one_hot_actions(self, actions: DiscreteActionType) -> npt.NDArray:
+    def compute_one_hot_actions(self, actions) -> npt.NDArray:
         one_hot_actions = np.zeros((self.n_agents, self.n_actions), dtype=np.float32)
         index = np.arange(self.n_agents)
         one_hot_actions[index, actions] = 1.0

marlenv/wrappers/paddings.py

@@ -1,22 +1,20 @@
 import numpy as np
-import numpy.typing as npt
 from dataclasses import dataclass
-from marlenv.models import Observation, ActionSpace
+from marlenv.models import Observation, Space
 from .rlenv_wrapper import RLEnvWrapper, MARLEnv
 from typing_extensions import TypeVar
 
 
-A = TypeVar("A", default=npt.NDArray)
-AS = TypeVar("AS", bound=ActionSpace, default=ActionSpace)
+AS = TypeVar("AS", bound=Space, default=Space)
 
 
 @dataclass
-class PadExtras(RLEnvWrapper[A, AS]):
+class PadExtras(RLEnvWrapper[AS]):
     """RLEnv wrapper that adds extra zeros at the end of the observation extras."""
 
     n: int
 
-    def __init__(self, env: MARLEnv[A, AS], n_added: int):
+    def __init__(self, env: MARLEnv[AS], n_added: int):
         assert len(env.extras_shape) == 1, "PadExtras only accepts 1D extras"
         meanings = env.extras_meanings + [f"Padding-{i}" for i in range(n_added)]
         super().__init__(
@@ -42,10 +40,10 @@ class PadExtras(RLEnvWrapper[A, AS]):
 
 
 @dataclass
-class PadObservations(RLEnvWrapper[A, AS]):
+class PadObservations(RLEnvWrapper[AS]):
     """RLEnv wrapper that adds extra zeros at the end of the observation data."""
 
-    def __init__(self, env: MARLEnv[A, AS], n_added: int) -> None:
+    def __init__(self, env: MARLEnv[AS], n_added: int) -> None:
         assert len(env.observation_shape) == 1, "PadObservations only accepts 1D observations"
         super().__init__(env, observation_shape=(env.observation_shape[0] + n_added,))
         self.n = n_added

marlenv/wrappers/penalty_wrapper.py

@@ -1,21 +1,18 @@
 from dataclasses import dataclass
 import numpy as np
-import numpy.typing as npt
-from marlenv.models import ActionSpace
+from marlenv.models import Space
 from .rlenv_wrapper import RLEnvWrapper, MARLEnv
-# from ..models.rl_env import MOMARLEnv
 
 from typing_extensions import TypeVar
 
-A = TypeVar("A", default=npt.NDArray)
-AS = TypeVar("AS", bound=ActionSpace, default=ActionSpace)
+AS = TypeVar("AS", bound=Space, default=Space)
 
 
 @dataclass
-class TimePenalty(RLEnvWrapper[A, AS]):
+class TimePenalty(RLEnvWrapper[AS]):
     penalty: float | np.ndarray
 
-    def __init__(self, env: MARLEnv[A, AS], penalty: float | list[float]):
+    def __init__(self, env: MARLEnv[AS], penalty: float | list[float]):
         super().__init__(env)
 
         if env.is_multi_objective:
@@ -26,7 +23,7 @@ class TimePenalty(RLEnvWrapper[A, AS]):
             assert isinstance(penalty, (float, int))
             self.penalty = penalty
 
-    def step(self, action: A):
+    def step(self, action):
         step = self.wrapped.step(action)
         step.reward = step.reward - self.penalty
         return step