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

marlenv/utils/schedule.py

@@ -0,0 +1,281 @@
+from abc import abstractmethod
+from dataclasses import dataclass
+from typing import Callable, Optional, TypeVar
+
+T = TypeVar("T")
+
+
+@dataclass
+class Schedule:
+    """
+    Schedules the value of a varaible over time.
+    """
+
+    name: str
+    start_value: float
+    end_value: float
+    _t: int
+    n_steps: int
+
+    def __init__(self, start_value: float, end_value: float, n_steps: int):
+        self.start_value = start_value
+        self.end_value = end_value
+        self.n_steps = n_steps
+        self.name = self.__class__.__name__
+        self._t = 0
+        self._current_value = self.start_value
+
+    def update(self, step: Optional[int] = None):
+        """Update the value of the schedule. Force a step if given."""
+        if step is not None:
+            self._t = step
+        else:
+            self._t += 1
+        if self._t >= self.n_steps:
+            self._current_value = self.end_value
+        else:
+            self._current_value = self._compute()
+
+    @abstractmethod
+    def _compute(self) -> float:
+        """Compute the value of the schedule"""
+
+    @property
+    def value(self) -> float:
+        """Returns the current value of the schedule"""
+        return self._current_value
+
+    @staticmethod
+    def constant(value: float):
+        return ConstantSchedule(value)
+
+    @staticmethod
+    def linear(start_value: float, end_value: float, n_steps: int):
+        return LinearSchedule(start_value, end_value, n_steps)
+
+    @staticmethod
+    def exp(start_value: float, end_value: float, n_steps: int):
+        return ExpSchedule(start_value, end_value, n_steps)
+
+    @staticmethod
+    def arbitrary(func: Callable[[int], float], n_steps: Optional[int] = None):
+        if n_steps is None:
+            n_steps = 0
+        return ArbitrarySchedule(func, n_steps)
+
+    def rounded(self, n_digits: int = 0) -> "RoundedSchedule":
+        return RoundedSchedule(self, n_digits)
+
+    # Operator overloading
+    def __mul__(self, other: T) -> T:
+        return self.value * other  # type: ignore
+
+    def __rmul__(self, other: T) -> T:
+        return self.value * other  # type: ignore
+
+    def __pow__(self, exp: float) -> float:
+        return self.value**exp
+
+    def __rpow__(self, other: T) -> T:
+        return other**self.value  # type: ignore
+
+    def __add__(self, other: T) -> T:
+        return self.value + other  # type: ignore
+
+    def __radd__(self, other: T) -> T:
+        return self.value + other  # type: ignore
+
+    def __neg__(self):
+        return -self.value
+
+    def __pos__(self):
+        return +self.value
+
+    def __sub__(self, other: T) -> T:
+        return self.value - other  # type: ignore
+
+    def __rsub__(self, other: T) -> T:
+        return other - self.value  # type: ignore
+
+    def __div__(self, other: T) -> T:
+        return self.value // other  # type: ignore
+
+    def __rdiv__(self, other: T) -> T:
+        return other // self.value  # type: ignore
+
+    def __truediv__(self, other: T) -> T:
+        return self.value / other  # type: ignore
+
+    def __rtruediv__(self, other: T) -> T:
+        return other / self.value  # type: ignore
+
+    def __lt__(self, other) -> bool:
+        return self.value < other
+
+    def __le__(self, other) -> bool:
+        return self.value <= other
+
+    def __gt__(self, other) -> bool:
+        return self.value > other
+
+    def __ge__(self, other) -> bool:
+        return self.value >= other
+
+    def __eq__(self, other) -> bool:
+        if isinstance(other, Schedule):
+            if self.start_value != other.start_value:
+                return False
+            if self.end_value != other.end_value:
+                return False
+            if self.n_steps != other.n_steps:
+                return False
+            if type(self) is not type(other):
+                return False
+        return self.value == other
+
+    def __ne__(self, other) -> bool:
+        return not (self.__eq__(other))
+
+    def __float__(self):
+        return self.value
+
+    def __int__(self) -> int:
+        return int(self.value)
+
+
+@dataclass(eq=False)
+class LinearSchedule(Schedule):
+    a: float
+    b: float
+
+    def __init__(self, start_value: float, end_value: float, n_steps: int):
+        super().__init__(start_value, end_value, n_steps)
+        self._current_value = self.start_value
+        # y = ax + b
+        self.a = (self.end_value - self.start_value) / self.n_steps
+        self.b = self.start_value
+
+    def _compute(self):
+        return self.a * (self._t) + self.b
+
+    @property
+    def value(self) -> float:
+        return self._current_value
+
+
+@dataclass(eq=False)
+class ExpSchedule(Schedule):
+    """Exponential schedule. After n_steps, the value will be min_value.
+
+    Update formula is next_value = start_value * (min_value / start_value) ** (step / (n - 1))
+    """
+
+    n_steps: int
+    base: float
+    last_update_step: int
+
+    def __init__(self, start_value: float, min_value: float, n_steps: int):
+        super().__init__(start_value, min_value, n_steps)
+        self.base = self.end_value / self.start_value
+        self.last_update_step = self.n_steps - 1
+
+    def _compute(self):
+        return self.start_value * (self.base) ** (self._t / (self.n_steps - 1))
+
+    @property
+    def value(self) -> float:
+        return self._current_value
+
+
+@dataclass(eq=False)
+class ConstantSchedule(Schedule):
+    def __init__(self, value: float):
+        super().__init__(value, value, 0)
+        self._value = value
+
+    def update(self, step=None):
+        return
+
+    @property
+    def value(self) -> float:
+        return self._value
+
+
+@dataclass(eq=False)
+class RoundedSchedule(Schedule):
+    def __init__(self, schedule: Schedule, n_digits: int):
+        super().__init__(schedule.start_value, schedule.end_value, schedule.n_steps)
+        self.schedule = schedule
+        self.n_digits = n_digits
+
+    def update(self, step: int | None = None):
+        return self.schedule.update(step)
+
+    def _compute(self) -> float:
+        return self.schedule._compute()
+
+    @property
+    def value(self) -> float:
+        return round(self.schedule.value, self.n_digits)
+
+
+@dataclass(eq=False)
+class MultiSchedule(Schedule):
+    def __init__(self, schedules: dict[int, Schedule]):
+        ordered_schedules, ordered_steps = MultiSchedule._verify(schedules)
+        n_steps = ordered_steps[-1] + ordered_schedules[-1].n_steps
+        super().__init__(ordered_schedules[0].start_value, ordered_schedules[-1].end_value, n_steps)
+        self.schedules = iter(ordered_schedules)
+        self.current_schedule = next(self.schedules)
+        self.offset = 0
+        self.current_end = ordered_steps[1]
+
+    @staticmethod
+    def _verify(schedules: dict[int, Schedule]):
+        sorted_steps = sorted(schedules.keys())
+        sorted_schedules = [schedules[t] for t in sorted_steps]
+        if sorted_steps[0] != 0:
+            raise ValueError("First schedule must start at t=0")
+        current_step = 0
+        for i in range(len(sorted_steps)):
+            # Artificially set the end step of ConstantSchedules to the next step
+            if isinstance(sorted_schedules[i], ConstantSchedule):
+                if i + 1 < len(sorted_steps):
+                    sorted_schedules[i].n_steps = sorted_steps[i + 1] - sorted_steps[i]
+            if sorted_steps[i] != current_step:
+                raise ValueError(f"Schedules are not contiguous at t={current_step}")
+            current_step += sorted_schedules[i].n_steps
+        return sorted_schedules, sorted_steps
+
+    def update(self, step: int | None = None):
+        if step is not None:
+            raise NotImplementedError("Cannot update MultiSchedule with a specific step")
+        super().update(step)
+        # If we reach the end of the current schedule, update to the next one
+        # except if we are at the end.
+        if self._t == self.current_end and self._t < self.n_steps:
+            self.current_schedule = next(self.schedules)
+            self.offset = self._t
+            self.current_end += self.current_schedule.n_steps
+        self.current_schedule.update(self.relative_step)
+
+    @property
+    def relative_step(self):
+        return self._t - self.offset
+
+    def _compute(self) -> float:
+        return self.current_schedule._compute()
+
+    @property
+    def value(self):
+        return self.current_schedule.value
+
+
+@dataclass(eq=False)
+class ArbitrarySchedule(Schedule):
+    def __init__(self, fn: Callable[[int], float], n_steps: int):
+        super().__init__(fn(0), fn(n_steps), n_steps)
+        self._func = fn
+
+    def _compute(self) -> float:
+        return self._func(self._t)

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