multi-agent-rlenv 3.3.6__tar.gz → 3.4.0__tar.gz

{multi_agent_rlenv-3.3.6 → multi_agent_rlenv-3.4.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: multi-agent-rlenv
-Version: 3.3.6
+Version: 3.4.0
 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>

{multi_agent_rlenv-3.3.6 → multi_agent_rlenv-3.4.0}/src/marlenv/__init__.py

@@ -62,7 +62,7 @@ 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.6"
+__version__ = "3.4.0"
 
 from . import models
 from . import wrappers

{multi_agent_rlenv-3.3.6 → multi_agent_rlenv-3.4.0}/src/marlenv/adapters/overcooked_adapter.py

@@ -1,6 +1,6 @@
 import sys
 from dataclasses import dataclass
-from typing import Literal, Sequence
+from typing import Literal, Sequence, Optional
 from copy import deepcopy
 
 import cv2
@@ -8,6 +8,7 @@ import numpy as np
 import numpy.typing as npt
 import pygame
 from marlenv.models import ContinuousSpace, DiscreteActionSpace, MARLEnv, Observation, State, Step
+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
@@ -17,8 +18,17 @@ from overcooked_ai_py.visualization.state_visualizer import StateVisualizer
 @dataclass
 class Overcooked(MARLEnv[Sequence[int] | npt.NDArray, DiscreteActionSpace]):
     horizon: int
+    shaping_factor: Schedule
 
-    def __init__(self, oenv: OvercookedEnv):
+    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
@@ -40,6 +50,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:
@@ -84,12 +96,15 @@ class Overcooked(MARLEnv[Sequence[int] | npt.NDArray, DiscreteActionSpace]):
         return np.array(available_actions, dtype=np.bool)
 
     def step(self, actions: Sequence[int] | npt.NDArray[np.int32 | np.int64]) -> Step:
+        self.shaping_factor.update()
         actions = [Action.ALL_ACTIONS[a] for a in actions]
         _, reward, done, info = self._oenv.step(actions, display_phi=True)
+
+        reward += sum(info["shaped_r_by_agent"]) * self.shaping_factor
         return Step(
             obs=self.get_observation(),
             state=self.get_state(),
-            reward=np.array([reward]),
+            reward=np.array([reward], dtype=np.float32),
             done=done,
             truncated=False,
             info=info,
@@ -99,19 +114,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
@@ -185,6 +206,37 @@ class Overcooked(MARLEnv[Sequence[int] | npt.NDArray, DiscreteActionSpace]):
             "you_shall_not_pass",
         ],
         horizon: int = 400,
+        reward_shaping_factor: float | Schedule = 1.0,
     ):
         mdp = OvercookedGridworld.from_layout_name(layout)
-        return Overcooked(OvercookedEnv.from_mdp(mdp, horizon=horizon))
+        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,
+        shaping_factor: float | Schedule = 1.0,
+        layout_name: Optional[str] = None,
+    ):
+        """
+        Create an Overcooked environment from a grid layout where
+        - S is a serving location
+        - P is a cooking pot
+        - X is a counter
+        - O is an onion dispenser
+        - D is a dish dispenser
+        - T is a tomato dispenser
+        - 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.
+        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, info_level=0), shaping_factor, layout_name)

{multi_agent_rlenv-3.3.6 → multi_agent_rlenv-3.4.0}/src/marlenv/models/env.py

@@ -1,7 +1,7 @@
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from itertools import product
-from typing import Generic, Optional, Sequence
+from typing import Any, Generic, Optional, Sequence
 
 import cv2
 import numpy as np
@@ -13,8 +13,8 @@ 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)
+ActionType = TypeVar("ActionType", default=Any)
+ActionSpaceType = TypeVar("ActionSpaceType", bound=ActionSpace, default=Any)
 
 
 @dataclass
@@ -108,6 +108,11 @@ class MARLEnv(ABC, Generic[ActionType, ActionSpaceType]):
         """Whether the environment is multi-objective."""
         return self.reward_space.size > 1
 
+    @property
+    def n_objectives(self) -> int:
+        """The number of objectives in the environment."""
+        return self.reward_space.size
+
     def sample_action(self) -> ActionType:
         """Sample an available action from the action space."""
         return self.action_space.sample(self.available_actions())  # type: ignore

{multi_agent_rlenv-3.3.6 → multi_agent_rlenv-3.4.0}/src/marlenv/models/episode.py

@@ -179,9 +179,9 @@ class Episode(Generic[A]):
     @cached_property
     def dones(self):
         """The done flags for each transition"""
-        dones = np.zeros_like(self.rewards, dtype=np.float32)
+        dones = np.zeros_like(self.rewards, dtype=np.bool)
         if self.is_done:
-            dones[self.episode_len - 1 :] = 1.0
+            dones[self.episode_len - 1 :] = True
         return dones
 
     @property

multi_agent_rlenv-3.4.0/src/marlenv/utils/__init__.py

@@ -0,0 +1,10 @@
+from .schedule import Schedule, MultiSchedule, RoundedSchedule, LinearSchedule, ExpSchedule
+
+
+__all__ = [
+    "Schedule",
+    "LinearSchedule",
+    "ExpSchedule",
+    "MultiSchedule",
+    "RoundedSchedule",
+]

multi_agent_rlenv-3.4.0/src/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)

{multi_agent_rlenv-3.3.6 → multi_agent_rlenv-3.4.0}/tests/test_adapters.py

@@ -1,5 +1,3 @@
-from importlib.util import find_spec
-
 import numpy as np
 import pytest
 
@@ -7,12 +5,12 @@ import marlenv
 from marlenv import ContinuousActionSpace, DiscreteActionSpace, DiscreteMockEnv, MARLEnv, Observation, State
 from marlenv.adapters import PymarlAdapter
 
-skip_gym = find_spec("gymnasium") is None
-skip_pettingzoo = find_spec("pettingzoo") is None
-skip_smac = find_spec("smac") is None
+skip_gym = not marlenv.adapters.HAS_GYM
+skip_pettingzoo = not marlenv.adapters.HAS_PETTINGZOO
+skip_smac = not marlenv.adapters.HAS_SMAC
 # Check for "overcooked_ai_py.mdp" specifically because after uninstalling, the package
 # can still be found because of some remaining .pkl file.
-skip_overcooked = find_spec("overcooked_ai_py.mdp") is None
+skip_overcooked = not marlenv.adapters.HAS_OVERCOOKED
 
 
 @pytest.mark.skipif(skip_gym, reason="Gymnasium is not installed")
@@ -183,6 +181,59 @@ def test_overcooked_obs_state():
             assert step.done
 
 
+@pytest.mark.skipif(skip_overcooked, reason="Overcooked is not installed")
+def test_overcooked_shaping():
+    from marlenv.adapters import Overcooked
+
+    UP = 0
+    RIGHT = 2
+    STAY = 4
+    INTERACT = 5
+    grid = [
+        ["X", "X", "X", "D", "X"],
+        ["X", "O", "S", "2", "X"],
+        ["X", "1", "P", " ", "X"],
+        ["X", "T", "S", " ", "X"],
+        ["X", "X", "X", "X", "X"],
+    ]
+
+    env = Overcooked.from_grid(grid)
+    env.reset()
+    actions_rewards = [
+        ([UP, STAY], False),
+        ([INTERACT, STAY], False),
+        ([RIGHT, STAY], False),
+        ([INTERACT, STAY], True),
+    ]
+
+    for action, expected_reward in actions_rewards:
+        step = env.step(action)
+        if expected_reward:
+            assert step.reward.item() > 0
+
+
+@pytest.mark.skipif(skip_overcooked, reason="Overcooked is not installed")
+def test_overcooked_name():
+    from marlenv.adapters import Overcooked
+
+    grid = [
+        ["X", "X", "X", "D", "X"],
+        ["X", "O", "S", "2", "X"],
+        ["X", "1", "P", " ", "X"],
+        ["X", "T", "S", " ", "X"],
+        ["X", "X", "X", "X", "X"],
+    ]
+
+    env = Overcooked.from_grid(grid)
+    assert env.name == "Overcooked-custom-layout"
+
+    env = Overcooked.from_grid(grid, layout_name="my incredible grid")
+    assert env.name == "Overcooked-my incredible grid"
+
+    env = Overcooked.from_layout("asymmetric_advantages")
+    assert env.name == "Overcooked-asymmetric_advantages"
+
+
 def test_pymarl():
     LIMIT = 20
     N_AGENTS = 2

multi_agent_rlenv-3.4.0/tests/test_schedules.py

@@ -0,0 +1,155 @@
+from marlenv.utils import Schedule, MultiSchedule
+
+
+def is_close(a: float, b: float, tol: float = 1e-6) -> bool:
+    return abs(a - b) < tol
+
+
+def test_linear_schedule_increasing():
+    s = Schedule.linear(0, 1, 10)
+    for i in range(10):
+        assert is_close(s.value, i / 10)
+        s.update()
+    for i in range(10):
+        assert s.value == 1.0
+        s.update()
+
+
+def test_linear_schedule_decreasing():
+    s = Schedule.linear(0, -1, 10)
+    for i in range(10):
+        assert is_close(s.value, -i / 10)
+        s.update()
+    for i in range(10):
+        assert s.value == -1.0
+        s.update()
+
+
+def test_linear_schedule_set_timestep():
+    s = Schedule.linear(0, 1, 10)
+    s.update(50)
+    assert is_close(s.value, 1)
+
+    s.update(0)
+    assert is_close(s.value, 0)
+
+    s.update(5)
+    assert is_close(s.value, 0.5)
+
+
+def test_exp_schedule_increasing():
+    s = Schedule.exp(1, 16, 5)
+    assert is_close(s.value, 1)
+    s.update()
+    assert is_close(s.value, 2)
+    s.update()
+    assert is_close(s.value, 4)
+    s.update()
+    assert is_close(s.value, 8)
+    s.update()
+    assert is_close(s.value, 16)
+    for _ in range(10):
+        s.update()
+        assert is_close(s.value, 16)
+
+
+def test_exp_schedule_set_timestep():
+    s = Schedule.exp(1, 16, 5)
+    s.update(50)
+    assert is_close(s.value, 16)
+
+    s.update(0)
+    assert is_close(s.value, 1)
+
+    s.update(5)
+    assert is_close(s.value, 16)
+
+    s.update(1)
+    assert is_close(s.value, 2)
+
+
+def test_exp_schedule_decreasing():
+    s = Schedule.exp(16, 1, 5)
+    assert is_close(s.value, 16)
+    s.update()
+    assert is_close(s.value, 8)
+    s.update()
+    assert is_close(s.value, 4)
+    s.update()
+    assert is_close(s.value, 2)
+    s.update()
+    assert is_close(s.value, 1)
+    for _ in range(10):
+        s.update()
+        assert is_close(s.value, 1)
+
+
+def test_const_schedule():
+    s = Schedule.constant(50)
+    for _ in range(10):
+        assert s.value == 50
+        s.update()
+
+
+def test_equality_linear():
+    s1 = Schedule.linear(0, 1, 10)
+    s2 = Schedule.linear(0, 1, 10)
+    s3 = Schedule.linear(0, 1, 5)
+    assert s1 == s2
+    assert s1 != s3
+
+    s1.update()
+    assert s1 != s2
+    s2.update()
+    assert s1 == s2
+
+
+def test_equality_exp():
+    s1 = Schedule.exp(1, 16, 5)
+    s2 = Schedule.exp(1, 16, 5)
+    s3 = Schedule.exp(1, 16, 10)
+    assert s1 != 5
+    assert s1 == 1
+    assert s1 == s2
+    assert s1 != s3
+
+    s1.update()
+    assert s1 != s2
+    assert s1 == 2
+    assert s2 == 1
+    s2.update()
+    assert s1 == s2
+    assert s2 == 2
+
+
+def test_multi_schedule():
+    s = MultiSchedule(
+        {
+            0: Schedule.constant(0),
+            10: Schedule.linear(0, 1, 10),
+            20: Schedule.exp(1, 16, 5),
+        }
+    )
+    expected_values = [0.0] * 10 + [i / 10 for i in range(10)] + [2**i for i in range(5)]
+    for i in range(25):
+        assert is_close(s.value, expected_values[i])
+        s.update()
+
+
+def test_equality_const():
+    s1 = Schedule.constant(50)
+    s2 = Schedule.constant(50)
+    assert s1 == s2
+
+    s1.update()
+    assert s1 == s2
+    s2.update()
+    assert s1 == s2
+
+
+def test_inequality_different_schedules():
+    s1 = Schedule.linear(1, 2, 10)
+    s2 = Schedule.exp(1, 2, 10)
+    s3 = Schedule.linear(1, 2, 10)
+    assert s1 != s2
+    assert not s1 != s3

{multi_agent_rlenv-3.3.6 → multi_agent_rlenv-3.4.0}/tests/test_serialization.py

@@ -7,6 +7,7 @@ from copy import deepcopy
 
 import marlenv
 from marlenv import DiscreteMockEnv, wrappers
+from marlenv.utils import Schedule
 
 
 def test_registry():
@@ -199,6 +200,16 @@ def test_deepcopy_overcooked():
     assert env == env2
 
 
+@pytest.mark.skipif(not marlenv.adapters.HAS_OVERCOOKED, reason="Overcooked is not installed")
+def test_deepcopy_overcooked_schedule():
+    env = marlenv.adapters.Overcooked.from_layout("scenario4", reward_shaping_factor=Schedule.linear(1, 0, 10))
+    env2 = deepcopy(env)
+    assert env == env2
+
+    env.random_step()
+    assert not env == env2, "The reward shaping factor should be different"
+
+
 @pytest.mark.skipif(not marlenv.adapters.HAS_OVERCOOKED, reason="Overcooked is not installed")
 def test_pickle_overcooked():
     env = marlenv.adapters.Overcooked.from_layout("scenario1_s", horizon=60)
@@ -281,3 +292,42 @@ def test_json_serialize_pettingzoo():
 def test_json_serialize_smac():
     env = marlenv.adapters.SMAC("3m")
     serde_and_check_key_values(env)
+
+
+class C:
+    def __call__(self, t):
+        return t + 1
+
+
+def test_serialize_schedule():
+    s = Schedule.linear(0, 1, 10)
+    orjson.dumps(s)
+    b = pickle.dumps(s)
+    s2 = pickle.loads(b)
+    assert s == s2
+
+    s = Schedule.exp(1, 16, 5)
+    orjson.dumps(s)
+    b = pickle.dumps(s)
+    s2 = pickle.loads(b)
+    assert s == s2
+
+    s = Schedule.constant(50)
+    orjson.dumps(s)
+    b = pickle.dumps(s)
+    s2 = pickle.loads(b)
+    assert s == s2
+
+    s = Schedule.arbitrary(lambda t: t + 1)
+    b = orjson.dumps(s)
+    try:
+        pickle.dumps(s)
+        assert False, "Should not be able to pickle arbitrary schedules because of the callable lambda"
+    except AttributeError:
+        pass
+
+    s = Schedule.arbitrary(C())
+    orjson.dumps(s)
+    b = pickle.dumps(s)
+    s2 = pickle.loads(b)
+    assert s == s2