multi-agent-rlenv 3.5.1__tar.gz → 3.5.4__tar.gz

{multi_agent_rlenv-3.5.1 → multi_agent_rlenv-3.5.4}/.github/workflows/ci.yaml

@@ -46,7 +46,7 @@ jobs:
           uv-version: 0.6.4
       - name: Install dependencies and run pytest
         run: |
-          uv sync --extra overcooked --extra gym --extra pettingzoo
+          uv sync --extra overcooked --extra gym --extra pettingzoo --extra torch
           uv run pytest
 
   build:

{multi_agent_rlenv-3.5.1 → multi_agent_rlenv-3.5.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: multi-agent-rlenv
-Version: 3.5.1
+Version: 3.5.4
 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>
@@ -19,6 +19,7 @@ Requires-Dist: pymunk>=6.0; extra == 'all'
 Requires-Dist: pysc2; extra == 'all'
 Requires-Dist: scipy>=1.10; extra == 'all'
 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: overcooked
@@ -31,6 +32,8 @@ Requires-Dist: scipy>=1.10; extra == 'pettingzoo'
 Provides-Extra: smac
 Requires-Dist: pysc2; extra == 'smac'
 Requires-Dist: smac; extra == 'smac'
+Provides-Extra: torch
+Requires-Dist: torch>=2.0; extra == 'torch'
 Description-Content-Type: text/markdown
 
 # `marlenv` - A unified framework for muti-agent reinforcement learning

{multi_agent_rlenv-3.5.1 → multi_agent_rlenv-3.5.4}/pyproject.toml

@@ -20,6 +20,7 @@ gym = ["gymnasium>=0.29.1"]
 smac = ["smac", "pysc2"]
 pettingzoo = ["pettingzoo>=1.20", "pymunk>=6.0", "scipy>=1.10"]
 overcooked = ["overcooked-ai>=1.1.0", "scipy>=1.10"]
+torch = ["torch>=2.0"]
 all = [
     "gymnasium>0.29.1",
     "pettingzoo>=1.20",
@@ -28,6 +29,7 @@ all = [
     "pysc2",
     "pymunk>=6.0",
     "scipy>=1.10",
+    "torch>=2.0",
 ]
 
 [build-system]

{multi_agent_rlenv-3.5.1 → multi_agent_rlenv-3.5.4}/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.5.1"
+__version__ = "3.5.4"
 
 from . import models
 from .models import (

{multi_agent_rlenv-3.5.1 → multi_agent_rlenv-3.5.4}/src/marlenv/models/env.py

@@ -199,7 +199,7 @@ class MARLEnv(ABC, Generic[ActionSpaceType]):
             episode.add(step, action)
         return episode
 
-    def has_same_inouts(self, other) -> bool:
+    def has_same_inouts(self, other: "MARLEnv[ActionSpaceType]") -> bool:
         """Alias for `have_same_inouts(self, other)`."""
         if not isinstance(other, MARLEnv):
             return False

{multi_agent_rlenv-3.5.1 → multi_agent_rlenv-3.5.4}/src/marlenv/models/episode.py

@@ -22,10 +22,10 @@ class Episode:
     all_extras: list[npt.NDArray[np.float32]]
     actions: list[npt.NDArray]
     rewards: list[npt.NDArray[np.float32]]
-    all_available_actions: list[npt.NDArray[np.bool_]]
+    all_available_actions: list[npt.NDArray[np.bool]]
     all_states: list[npt.NDArray[np.float32]]
     all_states_extras: list[npt.NDArray[np.float32]]
-    metrics: dict[str, float]
+    metrics: dict[str, Any]
     episode_len: int
     other: dict[str, list[Any]]
     is_done: bool = False
@@ -33,7 +33,7 @@ class Episode:
     """Whether the episode did reach a terminal state (different from truncated)"""
 
     @staticmethod
-    def new(obs: Observation, state: State, metrics: Optional[dict[str, float]] = None) -> "Episode":
+    def new(obs: Observation, state: State, metrics: Optional[dict[str, Any]] = None) -> "Episode":
         if metrics is None:
             metrics = {}
         return Episode(
@@ -66,13 +66,13 @@ class Episode:
         if target_len < self.episode_len:
             raise ValueError(f"Cannot pad episode to a smaller size: {target_len} < {self.episode_len}")
         padding_size = target_len - self.episode_len
-        obs = self.all_observations + [self.all_observations[0]] * padding_size
-        extras = self.all_extras + [self.all_extras[0]] * padding_size
-        actions = self.actions + [self.actions[0]] * padding_size
-        rewards = self.rewards + [self.rewards[0]] * padding_size
+        obs = self.all_observations + [np.zeros_like(self.all_observations[0])] * padding_size
+        extras = self.all_extras + [np.zeros_like(self.all_extras[0])] * padding_size
+        actions = self.actions + [np.zeros_like(self.actions[0])] * padding_size
+        rewards = self.rewards + [np.zeros_like(self.rewards[0])] * padding_size
         availables = self.all_available_actions + [self.all_available_actions[0]] * padding_size
-        states = self.all_states + [self.all_states[0]] * padding_size
-        states_extras = self.all_states_extras + [self.all_states_extras[0]] * padding_size
+        states = self.all_states + [np.zeros_like(self.all_states[0])] * padding_size
+        states_extras = self.all_states_extras + [np.zeros_like(self.all_states_extras[0])] * padding_size
         other = {key: value + [value[0]] * padding_size for key, value in self.other.items()}
         return Episode(
             all_observations=obs,
@@ -363,51 +363,6 @@ class Episode:
             for i, s in enumerate(scores):
                 self.metrics[f"score-{i}"] = float(s)
 
-    # def add_data(
-    #     self,
-    #     new_obs: Observation,
-    #     new_state: State,
-    #     action: A,
-    #     reward: np.ndarray,
-    #     done: bool,
-    #     truncated: bool,
-    #     info: dict[str, Any],
-    #     **kwargs,
-    # ):
-    #     """Add a new transition to the episode"""
-    #     self.episode_len += 1
-    #     self.all_observations.append(new_obs.data)
-    #     self.all_extras.append(new_obs.extras)
-    #     self.all_available_actions.append(new_obs.available_actions)
-    #     self.all_states.append(new_state.data)
-    #     self.all_states_extras.append(new_state.extras)
-    #     match action:
-    #         case np.ndarray() as action:
-    #             self.actions.append(action)
-    #         case other:
-    #             self.actions.append(np.array(other))
-    #     self.rewards.append(reward)
-    #     for key, value in kwargs.items():
-    #         current = self.other.get(key, [])
-    #         current.append(value)
-    #         self.other[key] = current
-
-    #     if done:
-    #         # Only set the truncated flag if the episode is not done (both could happen with a time limit)
-    #         self.is_truncated = truncated
-    #         self.is_done = done
-    #         # Add metrics that can be plotted
-    #         for key, value in info.items():
-    #             if isinstance(value, bool):
-    #                 value = int(value)
-    #             self.metrics[key] = value
-    #         self.metrics["episode_len"] = self.episode_len
-
-    #         rewards = np.array(self.rewards)
-    #         scores = np.sum(rewards, axis=0)
-    #         for i, s in enumerate(scores):
-    #             self.metrics[f"score-{i}"] = float(s)
-
-    def add_metrics(self, metrics: dict[str, float]):
+    def add_metrics(self, metrics: dict[str, Any]):
         """Add metrics to the episode"""
         self.metrics.update(metrics)

{multi_agent_rlenv-3.5.1 → multi_agent_rlenv-3.5.4}/src/marlenv/models/observation.py

@@ -87,3 +87,13 @@ class Observation:
             if not np.array_equal(self.data, other.data):
                 return False
         return np.array_equal(self.extras, other.extras) and np.array_equal(self.available_actions, other.available_actions)
+
+    def as_tensors(self, device=None):
+        """
+        Convert the observation to a tuple of tensors of shape (1, n_agents, <dim>).
+        """
+        import torch
+
+        data = torch.from_numpy(self.data).unsqueeze(0).to(device, non_blocking=True)
+        extras = torch.from_numpy(self.extras).unsqueeze(0).to(device, non_blocking=True)
+        return data, extras

{multi_agent_rlenv-3.5.1 → multi_agent_rlenv-3.5.4}/src/marlenv/models/state.py

@@ -52,3 +52,11 @@ class State(Generic[StateType]):
         if not np.array_equal(self.extras, value.extras):
             return False
         return True
+
+    def as_tensors(self, device=None):
+        """Convert the state to a tuple of tensors of shape (1, <dim>)."""
+        import torch
+
+        data = torch.from_numpy(self.data).unsqueeze(0).to(device, non_blocking=True)
+        extras = torch.from_numpy(self.extras).unsqueeze(0).to(device, non_blocking=True)
+        return data, extras

{multi_agent_rlenv-3.5.1 → multi_agent_rlenv-3.5.4}/src/marlenv/wrappers/potential_shaping.py

@@ -2,10 +2,15 @@ from abc import abstractmethod, ABC
 from .rlenv_wrapper import RLEnvWrapper
 from marlenv import Space, MARLEnv, Observation
 from typing import TypeVar, Optional
+import numpy as np
+import numpy.typing as npt
+
+from dataclasses import dataclass
 
 A = TypeVar("A", bound=Space)
 
 
+@dataclass
 class PotentialShaping(RLEnvWrapper[A], ABC):
     """
     Potential shaping for the Laser Learning Environment (LLE).
@@ -23,7 +28,7 @@ class PotentialShaping(RLEnvWrapper[A], ABC):
     ):
         super().__init__(env, extra_shape=extra_shape)
         self.gamma = gamma
-        self.current_potential = self.compute_potential()
+        self._current_potential = self.compute_potential()
 
     def add_extras(self, obs: Observation) -> Observation:
         """Add the extras related to potential shaping. Does nothing by default."""
@@ -31,19 +36,19 @@ class PotentialShaping(RLEnvWrapper[A], ABC):
 
     def reset(self):
         obs, state = super().reset()
-        self.current_potential = self.compute_potential()
+        self._current_potential = self.compute_potential()
         return self.add_extras(obs), state
 
     def step(self, actions):
-        phi_t = self.current_potential
+        prev_potential = self._current_potential
         step = super().step(actions)
 
-        self.current_potential = self.compute_potential()
-        shaped_reward = self.gamma * self.current_potential - phi_t
+        self._current_potential = self.compute_potential()
+        shaped_reward = self.gamma * self._current_potential - prev_potential
         step.obs = self.add_extras(step.obs)
         step.reward += shaped_reward
         return step
 
     @abstractmethod
-    def compute_potential(self) -> float:
+    def compute_potential(self) -> float | npt.NDArray[np.float32]:
         """Compute the potential of the current state of the environment."""

{multi_agent_rlenv-3.5.1 → multi_agent_rlenv-3.5.4}/tests/test_models.py

@@ -1,8 +1,11 @@
 from marlenv import Observation, Transition, DiscreteMockEnv, DiscreteMOMockEnv, Builder, State, Episode, MARLEnv, DiscreteSpace
 import numpy as np
-
+import pytest
+from importlib.util import find_spec
 from .utils import generate_episode
 
+HAS_PYTORCH = find_spec("torch") is not None
+
 
 def test_obs_eq():
     obs1 = Observation(
@@ -422,3 +425,33 @@ def test_wrong_extras_meanings_length():
         assert False, "This should raise a ValueError because the length of extras_meanings is different from the actual number of extras"
     except ValueError:
         pass
+
+
+@pytest.mark.skipif(not HAS_PYTORCH, reason="torch is not installed")
+def test_observation_as_tensor():
+    import torch
+
+    env = DiscreteMockEnv(4)
+    obs = env.reset()[0]
+    data, extras = obs.as_tensors()
+    assert isinstance(data, torch.Tensor)
+    assert data.shape == (1, env.n_agents, *env.observation_shape)
+    assert data.dtype == torch.float32
+    assert isinstance(extras, torch.Tensor)
+    assert extras.shape == (1, env.n_agents, *env.extras_shape)
+    assert extras.dtype == torch.float32
+
+
+@pytest.mark.skipif(not HAS_PYTORCH, reason="torch is not installed")
+def test_state_as_tensor():
+    import torch
+
+    env = DiscreteMockEnv(4)
+    state = env.reset()[1]
+    data, extras = state.as_tensors()
+    assert isinstance(data, torch.Tensor)
+    assert data.shape == (1, *env.state_shape)
+    assert data.dtype == torch.float32
+    assert isinstance(extras, torch.Tensor)
+    assert extras.shape == (1, *env.state_extra_shape)
+    assert extras.dtype == torch.float32