multi-agent-rlenv 3.5.5__py3-none-any.whl → 3.6.1__py3-none-any.whl

marlenv/__init__.py

@@ -62,7 +62,13 @@ 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.5.5"
+from importlib.metadata import version, PackageNotFoundError
+
+try:
+    __version__ = version("overcooked")
+except PackageNotFoundError:
+    __version__ = "0.0.0"  # fallback pratique en dev/CI
+
 
 from . import models
 from .models import (
@@ -82,16 +88,19 @@ from .models import (
 
 from . import wrappers
 from . import adapters
-from .env_builder import make, Builder
+from .env_builder import Builder
 from .wrappers import RLEnvWrapper
 from .mock_env import DiscreteMockEnv, DiscreteMOMockEnv
+from . import catalog
+from .adapters import make
 
 __all__ = [
     "models",
+    "make",
+    "catalog",
     "wrappers",
     "adapters",
     "spaces",
-    "make",
     "Builder",
     "MARLEnv",
     "Step",

marlenv/adapters/__init__.py

@@ -1,42 +1,33 @@
 from importlib.util import find_spec
 from .pymarl_adapter import PymarlAdapter
+from marlenv.utils import dummy_type, dummy_function
 
-HAS_GYM = False
-if find_spec("gymnasium") is not None:
-    from .gym_adapter import Gym
+HAS_GYM = find_spec("gymnasium") is not None
+if HAS_GYM:
+    from .gym_adapter import Gym, make
+else:
+    Gym = dummy_type("gymnasium")
+    make = dummy_function("gymnasium")
 
-    HAS_GYM = True
-
-HAS_PETTINGZOO = False
-if find_spec("pettingzoo") is not None:
+HAS_PETTINGZOO = find_spec("pettingzoo") is not None
+if HAS_PETTINGZOO:
     from .pettingzoo_adapter import PettingZoo
+else:
+    PettingZoo = dummy_type("pettingzoo")
 
-    HAS_PETTINGZOO = True
-
-HAS_SMAC = False
-if find_spec("smac") is not None:
+HAS_SMAC = find_spec("smac") is not None
+if HAS_SMAC:
     from .smac_adapter import SMAC
-
-    HAS_SMAC = True
-
-HAS_OVERCOOKED = False
-if find_spec("overcooked_ai_py") is not None and find_spec("overcooked_ai_py.mdp") is not None:
-    import numpy
-
-    # Overcooked assumes a version of numpy <2.0 where np.Inf is available.
-    setattr(numpy, "Inf", numpy.inf)
-    from .overcooked_adapter import Overcooked
-
-    HAS_OVERCOOKED = True
+else:
+    SMAC = dummy_type("smac", "https://github.com/oxwhirl/smac.git")
 
 __all__ = [
     "PymarlAdapter",
     "Gym",
+    "make",
     "PettingZoo",
     "SMAC",
-    "Overcooked",
     "HAS_GYM",
     "HAS_PETTINGZOO",
     "HAS_SMAC",
-    "HAS_OVERCOOKED",
 ]

marlenv/adapters/gym_adapter.py

@@ -78,3 +78,9 @@ class Gym(MARLEnv[Space]):
 
     def seed(self, seed_value: int):
         self._gym_env.reset(seed=seed_value)
+
+
+def make(env_id: str, **kwargs):
+    """Make an RLEnv from str (Gym) or PettingZoo"""
+    gym_env = gym.make(env_id, render_mode="rgb_array", **kwargs)
+    return Gym(gym_env)

marlenv/catalog/__init__.py

@@ -0,0 +1,26 @@
+from importlib.util import find_spec
+from ..utils.import_placeholders import dummy_type
+from marlenv.adapters import SMAC
+from .deepsea import DeepSea
+
+
+HAS_LLE = find_spec("lle") is not None
+if HAS_LLE:
+    from lle import LLE  # pyright: ignore[reportMissingImports]
+else:
+    LLE = dummy_type("lle", "laser-learning-environment")
+
+HAS_OVERCOOKED = find_spec("overcooked") is not None
+if HAS_OVERCOOKED:
+    from overcooked import Overcooked  # pyright: ignore[reportMissingImports]
+else:
+    Overcooked = dummy_type("overcooked", "overcooked")
+
+__all__ = [
+    "Overcooked",
+    "SMAC",
+    "LLE",
+    "DeepSea",
+    "HAS_LLE",
+    "HAS_OVERCOOKED",
+]

marlenv/catalog/deepsea.py

@@ -0,0 +1,73 @@
+from typing import Sequence
+import numpy as np
+from marlenv import MARLEnv, MultiDiscreteSpace, DiscreteSpace, Observation, State, Step
+from dataclasses import dataclass
+
+
+LEFT = 0
+RIGHT = 1
+
+
+@dataclass
+class DeepSea(MARLEnv[MultiDiscreteSpace]):
+    """
+    Deep Sea single-agent environment to test for deep exploration. The probability of reaching the goal state under random exploration is 2^(-max_depth).
+
+    The agent explores a 2D grid where the bottom-right corner (max_depth, max_depth) is the goal and is the only state to yield a reward.
+    The agent starts in the top-left corner (0, 0).
+    The agent has two actions: left or right, and taking an action makes the agent dive one row deeper. The agent can not go beyond the grid boundaries.
+    Going right gives a penalty of (0.01 / max_depth).
+    """
+
+    max_depth: int
+
+    def __init__(self, max_depth: int):
+        super().__init__(
+            n_agents=1,
+            action_space=DiscreteSpace(size=2, labels=["left", "right"]).repeat(1),
+            observation_shape=(2,),
+            state_shape=(2,),
+        )
+        self.max_depth = max_depth
+        self._row = 0
+        self._col = 0
+        self._step_right_penalty = -0.01 / self.max_depth
+
+    def get_observation(self) -> Observation:
+        return Observation(np.array([self._row, self._col], dtype=np.float32), self.available_actions())
+
+    def get_state(self) -> State:
+        return State(np.array([self._row, self._col], dtype=np.float32))
+
+    def reset(self):
+        self._row = 0
+        self._col = 0
+        return self.get_observation(), self.get_state()
+
+    def step(self, action: Sequence[int]):
+        self._row += 1
+        if action[0] == LEFT:
+            self._col -= 1
+        else:
+            self._col += 1
+        self._col = max(0, self._col)
+        if action[0] == RIGHT:
+            if self._row == self.max_depth:
+                reward = 1.0
+            else:
+                reward = self._step_right_penalty
+        else:
+            reward = 0.0
+        return Step(
+            self.get_observation(),
+            self.get_state(),
+            reward,
+            done=self._row == self.max_depth,
+        )
+
+    def set_state(self, state: State):
+        self._row, self._col = state.data
+
+    @property
+    def agent_state_size(self):
+        return 2

marlenv/env_builder.py

@@ -1,73 +1,13 @@
 from dataclasses import dataclass
-from typing import Generic, Literal, Optional, TypeVar, overload
+from typing import Generic, Literal, Optional, TypeVar
 import numpy as np
 import numpy.typing as npt
 
 from . import wrappers
-from marlenv import adapters
 from .models import Space, MARLEnv
 
 AS = TypeVar("AS", bound=Space)
 
-if adapters.HAS_PETTINGZOO:
-    from .adapters import PettingZoo
-    from pettingzoo import ParallelEnv
-
-    @overload
-    def make(env: ParallelEnv) -> PettingZoo: ...
-
-
-if adapters.HAS_GYM:
-    from .adapters import Gym
-    from gymnasium import Env
-    import gymnasium
-
-    @overload
-    def make(env: Env) -> Gym: ...
-
-    @overload
-    def make(env: str, **kwargs) -> Gym:
-        """
-        Make an RLEnv from the `gymnasium` registry (e.g: "CartPole-v1").
-        """
-
-
-if adapters.HAS_SMAC:
-    from .adapters import SMAC
-    from smac.env import StarCraft2Env
-
-    @overload
-    def make(env: StarCraft2Env) -> SMAC: ...
-
-
-if adapters.HAS_OVERCOOKED:
-    from .adapters import Overcooked
-    from overcooked_ai_py.mdp.overcooked_env import OvercookedEnv
-
-    @overload
-    def make(env: OvercookedEnv) -> Overcooked: ...
-
-
-def make(env, **kwargs):
-    """Make an RLEnv from str (Gym) or PettingZoo"""
-    match env:
-        case MARLEnv():
-            return env
-        case str(env_id):
-            if adapters.HAS_GYM:
-                gym_env = gymnasium.make(env_id, render_mode="rgb_array", **kwargs)
-                return Gym(gym_env)
-
-    if adapters.HAS_PETTINGZOO and isinstance(env, ParallelEnv):
-        return PettingZoo(env)  # type: ignore
-    if adapters.HAS_SMAC and isinstance(env, StarCraft2Env):
-        return SMAC(env)
-    if adapters.HAS_OVERCOOKED and isinstance(env, OvercookedEnv):
-        return Overcooked(env)  # type: ignore
-    if adapters.HAS_GYM and isinstance(env, Env):
-        return Gym(env)
-    raise ValueError(f"Unknown environment type: {type(env)}")
-
 
 @dataclass
 class Builder(Generic[AS]):

marlenv/models/spaces.py

@@ -105,7 +105,7 @@ class MultiDiscreteSpace(Space):
     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=mask) for mask, space in zip(mask, self.spaces)], dtype=np.int32)
+        return np.array([space.sample(mask=mask) for mask, space in zip(mask, self.spaces)], dtype=np.int32)  # type: ignore
 
     def __eq__(self, value: object) -> bool:
         if not isinstance(value, MultiDiscreteSpace):

marlenv/utils/__init__.py

@@ -1,5 +1,6 @@
 from .cached_property_collector import CachedPropertyCollector, CachedPropertyInvalidator
 from .schedule import ExpSchedule, LinearSchedule, MultiSchedule, RoundedSchedule, Schedule
+from .import_placeholders import dummy_function, dummy_type
 
 __all__ = [
     "Schedule",
@@ -9,4 +10,6 @@ __all__ = [
     "RoundedSchedule",
     "CachedPropertyCollector",
     "CachedPropertyInvalidator",
+    "dummy_function",
+    "dummy_type",
 ]

marlenv/utils/import_placeholders.py

@@ -0,0 +1,30 @@
+from typing import Optional, Any
+from types import SimpleNamespace
+
+
+def _raise_error(module_name: str, package_name: Optional[str] = None):
+    raise ImportError(
+        f"The optional dependency `{module_name}` is not installed.\nInstall the `{package_name}` package (e.g. pip install {package_name})."
+    )
+
+
+def dummy_type(module_name: str, package_name: Optional[str] = None):
+    class DummyType(type):
+        def __getattr__(cls, _) -> Any:
+            _raise_error(module_name, package_name)
+
+        def __call__(self, *args: Any, **kwds: Any) -> Any:
+            _raise_error(module_name, package_name)
+
+    class DummyClass(SimpleNamespace, metaclass=DummyType):
+        def __getattr__(self, _) -> Any:
+            _raise_error(module_name, package_name)
+
+    return DummyClass
+
+
+def dummy_function(module_name: str, package_name: Optional[str] = None):
+    def fail(*args, **kwargs) -> Any:
+        _raise_error(module_name, package_name)
+
+    return fail

marlenv/utils/schedule.py

@@ -1,6 +1,6 @@
 from abc import abstractmethod
 from dataclasses import dataclass
-from typing import Callable, Optional, TypeVar
+from typing import Any, Callable, Optional, TypeVar
 
 T = TypeVar("T")
 
@@ -142,6 +142,21 @@ class Schedule:
     def __int__(self) -> int:
         return int(self.value)
 
+    @staticmethod
+    def from_json(data: dict[str, Any]):
+        """Create a Schedule from a JSON-like dictionary."""
+        classname = data.get("name")
+        if classname == "LinearSchedule":
+            return LinearSchedule(data["start_value"], data["end_value"], data["n_steps"])
+        elif classname == "ExpSchedule":
+            return ExpSchedule(data["start_value"], data["end_value"], data["n_steps"])
+        elif classname == "ConstantSchedule":
+            return ConstantSchedule(data["value"])
+        elif classname == "ArbitrarySchedule":
+            raise NotImplementedError("ArbitrarySchedule cannot be deserialized from JSON")
+        else:
+            raise ValueError(f"Unknown schedule type: {classname}")
+
 
 @dataclass(eq=False)
 class LinearSchedule(Schedule):

{multi_agent_rlenv-3.5.5.dist-info → multi_agent_rlenv-3.6.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: multi-agent-rlenv
-Version: 3.5.5
+Version: 3.6.1
 Summary: A strongly typed Multi-Agent Reinforcement Learning framework
 Project-URL: repository, https://github.com/yamoling/multi-agent-rlenv
 Author-email: Yannick Molinghen <yannick.molinghen@ulb.be>
@@ -13,7 +13,8 @@ Requires-Dist: opencv-python>=4.0
 Requires-Dist: typing-extensions>=4.0
 Provides-Extra: all
 Requires-Dist: gymnasium>0.29.1; extra == 'all'
-Requires-Dist: overcooked-ai; extra == 'all'
+Requires-Dist: laser-learning-environment>=2.6.1; extra == 'all'
+Requires-Dist: overcooked>=0.1.0; extra == 'all'
 Requires-Dist: pettingzoo>=1.20; extra == 'all'
 Requires-Dist: pymunk>=6.0; extra == 'all'
 Requires-Dist: pysc2; extra == 'all'
@@ -22,9 +23,10 @@ Requires-Dist: smac; extra == 'all'
 Requires-Dist: torch>=2.0; extra == 'all'
 Provides-Extra: gym
 Requires-Dist: gymnasium>=0.29.1; extra == 'gym'
+Provides-Extra: lle
+Requires-Dist: laser-learning-environment>=2.6.1; extra == 'lle'
 Provides-Extra: overcooked
-Requires-Dist: overcooked-ai>=1.1.0; extra == 'overcooked'
-Requires-Dist: scipy>=1.10; extra == 'overcooked'
+Requires-Dist: overcooked>=0.1.0; extra == 'overcooked'
 Provides-Extra: pettingzoo
 Requires-Dist: pettingzoo>=1.20; extra == 'pettingzoo'
 Requires-Dist: pymunk>=6.0; extra == 'pettingzoo'
@@ -62,9 +64,24 @@ $ pip install marlenv[smac] # Install SMAC
 $ pip install marlenv[gym,smac]  # Install Gym & smac support
 ```
 
+## Using the `marlenv` environment catalog
+Some environments are registered in the `marlenv` and can be easily instantiated via its catalog. 
+
+```python
+from marlenv import catalog
+
+env1 = catalog.Overcooked.from_layout("scenario4")
+env2 = catalog.LLE.level(6)
+env3 = catalog.DeepSea(mex_depth=5)
+```
+Note that using the catalog requires the corresponding environment package to be installed. For instance you need to install the `laser-learning-environment` package to use `catalog.LLE`, which can be done by using the corresponding feature when at installation as shown below.
+```bash
+pip install multi-agent-rlenv[lle]
+```
+
 
 ## Using `marlenv` with existing libraries
-`marlenv` unifies multiple popular libraries under a single interface. Namely, `marlenv` supports `smac`, `gymnasium` and `pettingzoo`.
+`marlenv` provides adapters from most popular libraries to unify them under a single interface. Namely, `marlenv` supports `smac`, `gymnasium` and `pettingzoo`.
 
 ```python
 import marlenv
@@ -74,7 +91,7 @@ gym_env = marlenv.make("CartPole-v1", seed=25)
 
 # You can seemlessly instanciate a SMAC environment and directly pass your required arguments
 from marlenv.adapters import SMAC
-smac_env = env2 = SMAC("3m", debug=True, difficulty="9")
+smac_env = SMAC("3m", debug=True, difficulty="9")
 
 # pettingzoo is also supported
 from pettingzoo.sisl import pursuit_v4

{multi_agent_rlenv-3.5.5.dist-info → multi_agent_rlenv-3.6.1.dist-info}/RECORD

@@ -1,26 +1,28 @@
-marlenv/__init__.py,sha256=bX76JknjwfVJ6IOKql_y4rIqvvx9raepD7u2lB9CgGo,3656
-marlenv/env_builder.py,sha256=RJoHJLYAUE1ausAoJiRC3fUxyxpH1WRJf7Sdm2ml-uk,5517
+marlenv/__init__.py,sha256=MJgaW73zWYJKTNMWE8V3hTvrcMk-WEX3RaG-K_oIDD8,3886
+marlenv/env_builder.py,sha256=RUMFvW7dAJtHMLm8-oPVpjBefDtNliZtjlHci97Xj-Q,3874
 marlenv/env_pool.py,sha256=nCEBkGQU62fcvCAANyAqY8gCFjYlVnSCg-V3Fhx00yc,933
 marlenv/exceptions.py,sha256=gJUC_2rVAvOfK_ypVFc7Myh-pIfSU3To38VBVS_0rZA,1179
 marlenv/mock_env.py,sha256=kKvTdZl4_xSTTI9V6otZ1P709sfPYrqZSbbZaTip9iI,4684
 marlenv/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-marlenv/adapters/__init__.py,sha256=rWiqQOqTx3kVL5ZkPo3rkczrlQBBhQbU55zGI26SEeY,929
-marlenv/adapters/gym_adapter.py,sha256=Rx8ZnHW0XEwZzRT53BdDP1f4CtNp2tATAYZ0BbtBbd0,2863
-marlenv/adapters/overcooked_adapter.py,sha256=0-5sfKHGjmM4eYtbydsENMXV2Qx4WVGlvJl7MFZSaj8,9123
+marlenv/adapters/__init__.py,sha256=JsKWaeXvUjWEg3JC9TOtyjtZlTI9AAkLyab-jDa5yzM,783
+marlenv/adapters/gym_adapter.py,sha256=5HZF3g0QD4n7K4GQoMis4q0zj97uFTLdzdxMYHzM_UE,3041
 marlenv/adapters/pettingzoo_adapter.py,sha256=w9Ta-X4L_6ZXdDGmREOdcU0vpLR8lGP__s49DyK3dk8,2852
 marlenv/adapters/pymarl_adapter.py,sha256=2s7EY31s1hrml3q-BBaXo_eDMXTjkebozZPvzsgrb9c,3353
 marlenv/adapters/smac_adapter.py,sha256=8uWC7YKsaSXeTS8AUhpGOKvrWMbVEQT2-pml5BaFUB0,8343
+marlenv/catalog/__init__.py,sha256=VMAlot7aHjCXJPH6iycDHuYIEUyK_giOIO8qWTN8H_Q,655
+marlenv/catalog/deepsea.py,sha256=VPRV6tjPIkj-TMqp2L8U9COUK3FBlodoRW6gpsnwZ9Y,2332
 marlenv/models/__init__.py,sha256=uihmRs71Gg5z7Bvau_xtaQVg7xEtX8sTzi74bIHL5P0,443
 marlenv/models/env.py,sha256=BG1iVHxGD_p827mF0ewyOBn6wU2gtFsHLW1b4UtW-V0,7841
 marlenv/models/episode.py,sha256=zsyxsW4LIioPKyY4DZKn64A31e5ZvlwOf3HIGuRUzhs,13531
 marlenv/models/observation.py,sha256=RhvKvmys4bu3UwwVsvu7fJ7TMKt2QkKnBD1e0hw2r7s,3528
-marlenv/models/spaces.py,sha256=v7jnhPfj7vq7DFFJpSbQEYe4NGLLlj_bj2pzvvSBX7Y,7777
+marlenv/models/spaces.py,sha256=1aPmTcoOTU9nlwlcN7aswNrORwghOYAGqCLAMpk39SA,7793
 marlenv/models/state.py,sha256=LbP--JxBzRwMFpEAaZyxCX13xKQ27xPE2fabohaq9YI,2058
 marlenv/models/step.py,sha256=00PhD_ccdCIYAY1SVJdJU91weU0Y_tNIJwK16TN_53I,3056
 marlenv/models/transition.py,sha256=UkJVRNxZoyRkjE7YmKtUf_4xA7cOEh20O60dTldbvys,5070
-marlenv/utils/__init__.py,sha256=36pNw0r4V3xsqPZ5ljM29o96dfPAFq8WvMwggyv41fI,362
+marlenv/utils/__init__.py,sha256=ky5mz_T7EF65YNaEN1UDCUYZVlz7hFyKResgIJlE_1Q,462
 marlenv/utils/cached_property_collector.py,sha256=IOjbr61f0DqLhcidXKrl7MhN1BOEGiTzCANIKQCxaF0,600
-marlenv/utils/schedule.py,sha256=slhtpQiBHSUNyPmSkKb2yBgiHJqPhoPxa33GxvyV8Jc,8565
+marlenv/utils/import_placeholders.py,sha256=QN7gsfbFgSP2Lh-7YBC1RH-SNjbFacvRFmBgNs4Eb90,972
+marlenv/utils/schedule.py,sha256=BdjefYgAtGlh1wWGHENid4WNnPOU78kkNiRvR5A9GEA,9308
 marlenv/wrappers/__init__.py,sha256=Z4_M-mxRNKQeu52tkmQ4B2m3-zrsmjfXXL5NsWQ4vu4,952
 marlenv/wrappers/action_randomizer.py,sha256=A1kejqGOTA0sc_RQL0EOd6sMSbcIdiV5zlscjKUlzdY,474
 marlenv/wrappers/agent_id_wrapper.py,sha256=9qHV3LMQ4AjcDCSuvQhz5h9hUf7Xtrdi2sIxmNZk5NA,1126
@@ -36,7 +38,7 @@ marlenv/wrappers/potential_shaping.py,sha256=T_QvnmWReCgpyoInxRw2UXbmdvcBD5U-vV1
 marlenv/wrappers/rlenv_wrapper.py,sha256=S6G1VjFklTEzU6bj0AXrTDXnsTQJARq8VB4uUH6AXe4,2993
 marlenv/wrappers/time_limit.py,sha256=GxbxcbfFyuVg14ylQU2C_cjmV9q4uDAt5wepfgX_PyM,3976
 marlenv/wrappers/video_recorder.py,sha256=ucBQSNRPqDr-2mYxrTCqlrWcxSWtSJ7XlRC9-LdukBM,2535
-multi_agent_rlenv-3.5.5.dist-info/METADATA,sha256=WKf56Bb7PqZFrw2B6Sx8zulM-h7aZRqXMTvYHrSxEtQ,5005
-multi_agent_rlenv-3.5.5.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-multi_agent_rlenv-3.5.5.dist-info/licenses/LICENSE,sha256=_eeiGVoIJ7kYt6l1zbIvSBQppTnw0mjnYk1lQ4FxEjE,1074
-multi_agent_rlenv-3.5.5.dist-info/RECORD,,
+multi_agent_rlenv-3.6.1.dist-info/METADATA,sha256=FS_Hq7KmVDmZTHze01jijOSlEUf_QTA0CsZMcO1pzkY,5751
+multi_agent_rlenv-3.6.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+multi_agent_rlenv-3.6.1.dist-info/licenses/LICENSE,sha256=_eeiGVoIJ7kYt6l1zbIvSBQppTnw0mjnYk1lQ4FxEjE,1074
+multi_agent_rlenv-3.6.1.dist-info/RECORD,,

marlenv/adapters/overcooked_adapter.py

@@ -1,241 +0,0 @@
-import sys
-from dataclasses import dataclass
-from typing import Literal, Sequence, Optional
-from copy import deepcopy
-
-import cv2
-import numpy as np
-import numpy.typing as npt
-import pygame
-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
-from overcooked_ai_py.visualization.state_visualizer import StateVisualizer
-
-
-@dataclass
-class Overcooked(MARLEnv[MultiDiscreteSpace]):
-    horizon: int
-    shaping_factor: Schedule
-
-    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
-        self._visualizer = StateVisualizer()
-        width, height, layers = tuple(self._mdp.lossless_state_encoding_shape)
-        # -1 because we extract the "urgent" layer to the extras
-        shape = (int(layers - 1), int(width), int(height))
-        super().__init__(
-            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,),
-            extras_meanings=["timestep", "urgent"],
-            state_shape=shape,
-            state_extra_shape=(2,),
-            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:
-        """Current state of the environment"""
-        return self._oenv.state
-
-    def set_state(self, state: State):
-        raise NotImplementedError("Not yet implemented")
-
-    @property
-    def time_step(self):
-        return self.state.timestep
-
-    def _state_data(self):
-        players_layers = self._mdp.lossless_state_encoding(self.state)
-        state = np.array(players_layers, dtype=np.float32)
-        # Use axes (agents, channels, height, width) instead of (agents, height, width, channels)
-        state = np.transpose(state, (0, 3, 1, 2))
-        # The last last layer is for "urgency", put it in the extras
-        urgency = float(np.all(state[:, -1]))
-        state = state[:, :-1]
-        return state, urgency
-
-    def get_state(self):
-        data, is_urgent = self._state_data()
-        return State(data[0], np.array([self.time_step / self.horizon, is_urgent], dtype=np.float32))
-
-    def get_observation(self) -> Observation:
-        data, is_urgent = self._state_data()
-        return Observation(
-            data=data,
-            available_actions=self.available_actions(),
-            extras=np.array([[self.time_step / self.horizon, is_urgent]] * self.n_agents, dtype=np.float32),
-        )
-
-    def available_actions(self):
-        available_actions = np.full((self.n_agents, self.n_actions), False)
-        actions = self._mdp.get_actions(self._oenv.state)
-        for agent_num, agent_actions in enumerate(actions):
-            for action in agent_actions:
-                available_actions[agent_num, Action.ACTION_TO_INDEX[action]] = True
-        return np.array(available_actions, dtype=np.bool)
-
-    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)
-
-        reward += sum(info["shaped_r_by_agent"]) * self.shaping_factor
-        return Step(
-            obs=self.get_observation(),
-            state=self.get_state(),
-            reward=np.array([reward], dtype=np.float32),
-            done=done,
-            truncated=False,
-            info=info,
-        )
-
-    def reset(self):
-        self._oenv.reset()
-        return self.get_observation(), self.get_state()
-
-    def __deepcopy__(self, _):
-        """
-        Note: a specific implementation is needed because `pygame.font.Font` objects are not deep-copiable by default.
-        """
-        mdp = deepcopy(self._mdp)
-        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, "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"]), shaping_factor=state["schedule"])
-        self.name = state["name"]
-
-    def get_image(self):
-        rewards_dict = {}  # dictionary of details you want rendered in the UI
-        for key, value in self._oenv.game_stats.items():
-            if key in [
-                "cumulative_shaped_rewards_by_agent",
-                "cumulative_sparse_rewards_by_agent",
-            ]:
-                rewards_dict[key] = value
-
-        image = self._visualizer.render_state(
-            state=self._oenv.state,
-            grid=self._mdp.terrain_mtx,
-            hud_data=StateVisualizer.default_hud_data(self._oenv.state, **rewards_dict),
-        )
-
-        image = pygame.surfarray.array3d(image)
-        image = np.flip(np.rot90(image, 3), 1)
-        # Depending on the platform, the image may need to be converted to RGB
-        if sys.platform in ("linux", "linux2"):
-            image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
-        return image
-
-    @staticmethod
-    def from_layout(
-        layout: Literal[
-            "asymmetric_advantages",
-            "asymmetric_advantages_tomato",
-            "bonus_order_test",
-            "bottleneck",
-            "centre_objects",
-            "centre_pots",
-            "coordination_ring",
-            "corridor",
-            "counter_circuit",
-            "counter_circuit_o_1order",
-            "cramped_corridor",
-            "cramped_room",
-            "cramped_room_o_3orders",
-            "cramped_room_single",
-            "cramped_room_tomato",
-            "five_by_five",
-            "forced_coordination",
-            "forced_coordination_tomato",
-            "inverse_marshmallow_experiment",
-            "large_room",
-            "long_cook_time",
-            "marshmallow_experiment_coordination",
-            "marshmallow_experiment",
-            "mdp_test",
-            "m_shaped_s",
-            "multiplayer_schelling",
-            "pipeline",
-            "scenario1_s",
-            "scenario2",
-            "scenario2_s",
-            "scenario3",
-            "scenario4",
-            "schelling",
-            "schelling_s",
-            "simple_o",
-            "simple_o_t",
-            "simple_tomato",
-            "small_corridor",
-            "soup_coordination",
-            "tutorial_0",
-            "tutorial_1",
-            "tutorial_2",
-            "tutorial_3",
-            "unident",
-            "you_shall_not_pass",
-        ],
-        horizon: int = 400,
-        reward_shaping_factor: float | Schedule = 1.0,
-    ):
-        mdp = OvercookedGridworld.from_layout_name(layout)
-        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)