From 550b3ad88b9930347f79e869f27a3995a8f44390 Mon Sep 17 00:00:00 2001
From: Matteo Bettini <mb2389@cl.cam.ac.uk>
Date: Fri, 29 Mar 2024 19:03:11 +0100
Subject: [PATCH 01/16] amend

---
 torchrl/envs/libs/pettingzoo.py | 49 +++++++++++++++++++++++++--------
 1 file changed, 38 insertions(+), 11 deletions(-)

diff --git a/torchrl/envs/libs/pettingzoo.py b/torchrl/envs/libs/pettingzoo.py
index a1470776f10..4a059c7621d 100644
--- a/torchrl/envs/libs/pettingzoo.py
+++ b/torchrl/envs/libs/pettingzoo.py
@@ -90,12 +90,12 @@ class PettingZooWrapper(_EnvWrapper):
     If the number of agents during the task varies, please set ``use_mask=True``.
     ``"mask"`` will be provided
     as an output in each group and should be used to mask out dead agents.
-    The environment will be reset as soon as one agent is done.
+    The environment will be reset as soon as one agent is done (unless ``done_on_any`` is ``False``).
 
     In wrapped ``pettingzoo.AECEnv``, at each step only one agent will act.
     For this reason, it is compulsory to set ``use_mask=True`` for this type of environment.
     ``"mask"`` will be provided as an output for each group and can be used to mask out non-acting agents.
-    The environment will be reset only when all agents are done.
+    The environment will be reset only when all agents are done (unless ``done_on_any`` is ``True``).
 
     If there are any unavailable actions for an agent,
     the environment will also automatically update the mask of its ``action_spec`` and output an ``"action_mask"``
@@ -156,6 +156,9 @@ class PettingZooWrapper(_EnvWrapper):
         categorical_actions (bool, optional): if the enviornments actions are discrete, whether to transform
             them to categorical or one-hot.
         seed (int, optional): the seed. Defaults to ``None``.
+        done_on_any (bool, optional): whether the environment's done keys are set by aggregating the agent keys
+            using ``any()`` (when True) or ``all()`` (when False). Default (``None``) is to use ``any()`` for
+            parallel environments and ``all()`` for AEC ones.
 
     Examples:
         >>> # Parallel env
@@ -204,6 +207,7 @@ def __init__(
         use_mask: bool = False,
         categorical_actions: bool = True,
         seed: int | None = None,
+        done_on_any: bool = None,
         **kwargs,
     ):
         if env is not None:
@@ -214,6 +218,7 @@ def __init__(
         self.seed = seed
         self.use_mask = use_mask
         self.categorical_actions = categorical_actions
+        self.done_on_any = done_on_any
 
         super().__init__(**kwargs, allow_done_after_reset=True)
 
@@ -283,6 +288,9 @@ def _make_specs(
             "pettingzoo.utils.env.AECEnv",  # noqa: F821
         ],
     ) -> None:
+        # Set default for done on any or all
+        if self.done_on_any is None:
+            self.done_on_any = self.parallel
 
         # Create and check group map
         if self.group_map is None:
@@ -582,7 +590,6 @@ def _step(
         self,
         tensordict: TensorDictBase,
     ) -> TensorDictBase:
-
         if self.parallel:
             (
                 observation_dict,
@@ -651,16 +658,33 @@ def _step(
                                 value, device=self.device
                             )
 
-                elif not self.use_action_mask:
+                elif self.use_mask:
+                    if agent in self.agents:
+                        raise ValueError(
+                            f"Dead agent {agent} not found in step observation but still available in {self.agents}"
+                        )
+                    # Dead agent
+                    terminated = (
+                        terminations_dict[agent] if agent in terminations_dict else True
+                    )
+                    truncated = (
+                        truncations_dict[agent] if agent in truncations_dict else True
+                    )
+                    done = terminated or truncated
+                    group_done[index] = done
+                    group_terminated[index] = terminated
+                    group_truncated[index] = truncated
+
+                else:
                     # Dead agent, if we are not masking it out, this is not allowed
                     raise ValueError(
                         "Dead agents found in the environment,"
-                        " you need to set use_action_mask=True to allow this."
+                        " you need to set use_mask=True to allow this."
                     )
 
         # set done values
         done, terminated, truncated = self._aggregate_done(
-            tensordict_out, use_any=self.parallel
+            tensordict_out, use_any=self.done_on_any
         )
 
         tensordict_out.set("done", done)
@@ -673,7 +697,7 @@ def _aggregate_done(self, tensordict_out, use_any):
         truncated = False if use_any else True
         terminated = False if use_any else True
         for key in self.done_keys:
-            if isinstance(key, tuple):
+            if isinstance(key, tuple):  # Only look at group keys
                 if use_any:
                     if key[-1] == "done":
                         done = done | tensordict_out.get(key).any()
@@ -719,7 +743,6 @@ def _step_aec(
         self,
         tensordict: TensorDictBase,
     ) -> Tuple[Dict, Dict, Dict, Dict, Dict]:
-
         for group, agents in self.group_map.items():
             if self.agent_selection in agents:
                 agent_index = agents.index(self._env.agent_selection)
@@ -747,7 +770,6 @@ def _step_aec(
         )
 
     def _update_action_mask(self, td, observation_dict, info_dict):
-
         # Since we remove the action_mask keys we need to copy the data
         observation_dict = copy.deepcopy(observation_dict)
         info_dict = copy.deepcopy(info_dict)
@@ -821,7 +843,7 @@ class PettingZooEnv(PettingZooWrapper):
     If the number of agents during the task varies, please set ``use_mask=True``.
     ``"mask"`` will be provided
     as an output in each group and should be used to mask out dead agents.
-    The environment will be reset as soon as one agent is done.
+    The environment will be reset as soon as one agent is done (unless ``done_on_any`` is ``False``).
 
     For wrapping ``pettingzoo.AECEnv`` provide the name of your petting zoo task (in the ``task`` argument)
     and specify ``parallel=False``. This will construct the ``pettingzoo.AECEnv`` version of that task
@@ -829,7 +851,7 @@ class PettingZooEnv(PettingZooWrapper):
     In wrapped ``pettingzoo.AECEnv``, at each step only one agent will act.
     For this reason, it is compulsory to set ``use_mask=True`` for this type of environment.
     ``"mask"`` will be provided as an output for each group and can be used to mask out non-acting agents.
-    The environment will be reset only when all agents are done.
+    The environment will be reset only when all agents are done (unless ``done_on_any`` is ``True``).
 
     If there are any unavailable actions for an agent,
     the environment will also automatically update the mask of its ``action_spec`` and output an ``"action_mask"``
@@ -892,6 +914,9 @@ class PettingZooEnv(PettingZooWrapper):
         categorical_actions (bool, optional): if the enviornments actions are discrete, whether to transform
             them to categorical or one-hot.
         seed (int, optional): the seed.  Defaults to ``None``.
+        done_on_any (bool, optional): whether the environment's done keys are set by aggregating the agent keys
+            using ``any()`` (when True) or ``all()`` (when False). Default (``None``) is to use ``any()`` for
+            parallel environments and ``all()`` for AEC ones.
 
     Examples:
         >>> # Parallel env
@@ -930,6 +955,7 @@ def __init__(
         use_mask: bool = False,
         categorical_actions: bool = True,
         seed: int | None = None,
+        done_on_any: bool = None,
         **kwargs,
     ):
         if not _has_pettingzoo:
@@ -944,6 +970,7 @@ def __init__(
         kwargs["use_mask"] = use_mask
         kwargs["categorical_actions"] = categorical_actions
         kwargs["seed"] = seed
+        kwargs["done_on_any"] = done_on_any
 
         super().__init__(**kwargs)
 

From 7f7324f9bed7812e0399f3a9c9cebd0d94ef859d Mon Sep 17 00:00:00 2001
From: Matteo Bettini <mb2389@cl.cam.ac.uk>
Date: Tue, 2 Apr 2024 10:25:24 +0200
Subject: [PATCH 02/16] test

---
 test/test_libs.py | 52 +++++++++++++++++++++++++++++++++++++++++++----
 1 file changed, 48 insertions(+), 4 deletions(-)

diff --git a/test/test_libs.py b/test/test_libs.py
index c72fcc562e6..f9151bf535b 100644
--- a/test/test_libs.py
+++ b/test/test_libs.py
@@ -278,7 +278,6 @@ def _make_spec(  # noqa: F811
 
     @pytest.mark.parametrize("categorical", [True, False])
     def test_gym_spec_cast(self, categorical):
-
         batch_size = [3, 4]
         cat = DiscreteTensorSpec if categorical else OneHotDiscreteTensorSpec
         cat_shape = batch_size if categorical else (*batch_size, 5)
@@ -543,7 +542,6 @@ def test_torchrl_to_gym(self, backend, numpy):
         ],
     )
     def test_gym(self, env_name, frame_skip, from_pixels, pixels_only):
-
         if env_name == PONG_VERSIONED() and not from_pixels:
             # raise pytest.skip("already pixel")
             # we don't skip because that would raise an exception
@@ -3126,7 +3124,6 @@ class TestPettingZoo:
     def test_pistonball(
         self, parallel, continuous_actions, use_mask, return_state, group_map
     ):
-
         kwargs = {"n_pistons": 21, "continuous": continuous_actions}
 
         env = PettingZooEnv(
@@ -3141,6 +3138,54 @@ def test_pistonball(
 
         check_env_specs(env)
 
+    def test_dead_agents_done(self, seed=0):
+        scenario_args = {"n_walkers": 3, "terminate_on_fall": False}
+
+        env = PettingZooEnv(
+            task="multiwalker_v9",
+            parallel=True,
+            seed=0,
+            use_mask=False,
+            done_on_any=False,
+            **scenario_args,
+        )
+        with pytest.raises(
+            ValueError,
+            match="Dead agents found in the environment, "
+            "you need to set use_mask=True to allow this.",
+        ):
+            env.rollout(
+                max_steps=100,
+                break_when_any_done=True,  # This looks at root done set with done_on_any
+            )
+
+        for done_on_any in [True, False]:
+            env = PettingZooEnv(
+                task="multiwalker_v9",
+                parallel=True,
+                seed=0,
+                use_mask=True,
+                done_on_any=done_on_any,
+                **scenario_args,
+            )
+            td = env.rollout(
+                max_steps=100,
+                break_when_any_done=True,  # This looks at root done set with done_on_any
+            )
+            done = td.get(("next", "walker", "done"))
+            mask = td.get(("next", "walker", "mask"))
+
+            if done_on_any:
+                assert not done[-1].all()  # Done triggered on any
+            else:
+                assert done[-1].all()  # Done triggered on all
+            assert not done[
+                mask
+            ].any()  # When mask is true (alive agent), all agents are not done
+            assert done[
+                ~mask
+            ].all()  # When mask is false (dead agent), all agents are done
+
     @pytest.mark.parametrize(
         "wins_player_0",
         [True, False],
@@ -3156,7 +3201,6 @@ def test_tic_tac_toe(self, wins_player_0):
         )
 
         class Policy:
-
             action = 0
             t = 0
 

From 4b6f348bd8db5264b4edab13315a583521fc30b2 Mon Sep 17 00:00:00 2001
From: Matteo Bettini <mb2389@cl.cam.ac.uk>
Date: Tue, 2 Apr 2024 10:26:25 +0200
Subject: [PATCH 03/16] test

---
 test/test_libs.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/test/test_libs.py b/test/test_libs.py
index f9151bf535b..34f22ae9eac 100644
--- a/test/test_libs.py
+++ b/test/test_libs.py
@@ -3144,7 +3144,7 @@ def test_dead_agents_done(self, seed=0):
         env = PettingZooEnv(
             task="multiwalker_v9",
             parallel=True,
-            seed=0,
+            seed=seed,
             use_mask=False,
             done_on_any=False,
             **scenario_args,
@@ -3163,7 +3163,7 @@ def test_dead_agents_done(self, seed=0):
             env = PettingZooEnv(
                 task="multiwalker_v9",
                 parallel=True,
-                seed=0,
+                seed=seed,
                 use_mask=True,
                 done_on_any=done_on_any,
                 **scenario_args,

From 12f28273fd3700cfc4a88f8daebc2d117dc13265 Mon Sep 17 00:00:00 2001
From: Matteo Bettini <mb2389@cl.cam.ac.uk>
Date: Tue, 2 Apr 2024 10:27:16 +0200
Subject: [PATCH 04/16] deps

---
 .github/unittest/linux_libs/scripts_pettingzoo/environment.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/unittest/linux_libs/scripts_pettingzoo/environment.yml b/.github/unittest/linux_libs/scripts_pettingzoo/environment.yml
index 76f97355f7a..66f79eb2c24 100644
--- a/.github/unittest/linux_libs/scripts_pettingzoo/environment.yml
+++ b/.github/unittest/linux_libs/scripts_pettingzoo/environment.yml
@@ -20,4 +20,4 @@ dependencies:
     - expecttest
     - pyyaml
     - autorom[accept-rom-license]
-    - pettingzoo[all]==1.24.1
+    - pettingzoo[all]==1.24.3

From 120b25f73113343dd653b5d93f4a2ccc741d1941 Mon Sep 17 00:00:00 2001
From: Matteo Bettini <mb2389@cl.cam.ac.uk>
Date: Tue, 2 Apr 2024 15:06:53 +0200
Subject: [PATCH 05/16] amend

---
 test/test_libs.py | 12 +++++++++---
 1 file changed, 9 insertions(+), 3 deletions(-)

diff --git a/test/test_libs.py b/test/test_libs.py
index 34f22ae9eac..2d1d416e176 100644
--- a/test/test_libs.py
+++ b/test/test_libs.py
@@ -3149,14 +3149,17 @@ def test_dead_agents_done(self, seed=0):
             done_on_any=False,
             **scenario_args,
         )
+        td_reset = env.reset(seed=seed)
         with pytest.raises(
             ValueError,
             match="Dead agents found in the environment, "
             "you need to set use_mask=True to allow this.",
         ):
             env.rollout(
-                max_steps=100,
+                max_steps=500,
                 break_when_any_done=True,  # This looks at root done set with done_on_any
+                auto_reset=False,
+                tensordict=td_reset,
             )
 
         for done_on_any in [True, False]:
@@ -3168,11 +3171,14 @@ def test_dead_agents_done(self, seed=0):
                 done_on_any=done_on_any,
                 **scenario_args,
             )
+            td_reset = env.reset(seed=seed)
             td = env.rollout(
-                max_steps=100,
+                max_steps=500,
                 break_when_any_done=True,  # This looks at root done set with done_on_any
+                auto_reset=False,
+                tensordict=td_reset,
             )
-            done = td.get(("next", "walker", "done"))
+            done = td.get(("next", "walker", "done")).clone()
             mask = td.get(("next", "walker", "mask"))
 
             if done_on_any:

From 9fd9725fa44457fae76574008056406165fb51c4 Mon Sep 17 00:00:00 2001
From: Matteo Bettini <mb2389@cl.cam.ac.uk>
Date: Tue, 2 Apr 2024 15:07:53 +0200
Subject: [PATCH 06/16] amend

---
 test/test_libs.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test/test_libs.py b/test/test_libs.py
index 2d1d416e176..a19f23d720f 100644
--- a/test/test_libs.py
+++ b/test/test_libs.py
@@ -3178,7 +3178,7 @@ def test_dead_agents_done(self, seed=0):
                 auto_reset=False,
                 tensordict=td_reset,
             )
-            done = td.get(("next", "walker", "done")).clone()
+            done = td.get(("next", "walker", "done"))
             mask = td.get(("next", "walker", "mask"))
 
             if done_on_any:

From 96d7c5ae545815b86fff45d72a01cb402a180bd8 Mon Sep 17 00:00:00 2001
From: Matteo Bettini <55539777+matteobettini@users.noreply.github.com>
Date: Wed, 3 Apr 2024 07:39:19 +0100
Subject: [PATCH 07/16] Update torchrl/envs/libs/pettingzoo.py

Co-authored-by: Vincent Moens <vincentmoens@gmail.com>
---
 torchrl/envs/libs/pettingzoo.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/torchrl/envs/libs/pettingzoo.py b/torchrl/envs/libs/pettingzoo.py
index 4a059c7621d..0d3f8736ecc 100644
--- a/torchrl/envs/libs/pettingzoo.py
+++ b/torchrl/envs/libs/pettingzoo.py
@@ -955,7 +955,7 @@ def __init__(
         use_mask: bool = False,
         categorical_actions: bool = True,
         seed: int | None = None,
-        done_on_any: bool = None,
+        done_on_any: bool | None = None,
         **kwargs,
     ):
         if not _has_pettingzoo:

From de11c7ab3254f3e29a12e7031bd8d47c2b6b81cc Mon Sep 17 00:00:00 2001
From: Matteo Bettini <55539777+matteobettini@users.noreply.github.com>
Date: Wed, 3 Apr 2024 07:39:26 +0100
Subject: [PATCH 08/16] Update torchrl/envs/libs/pettingzoo.py

Co-authored-by: Vincent Moens <vincentmoens@gmail.com>
---
 torchrl/envs/libs/pettingzoo.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/torchrl/envs/libs/pettingzoo.py b/torchrl/envs/libs/pettingzoo.py
index 0d3f8736ecc..8cd41251d29 100644
--- a/torchrl/envs/libs/pettingzoo.py
+++ b/torchrl/envs/libs/pettingzoo.py
@@ -915,7 +915,7 @@ class PettingZooEnv(PettingZooWrapper):
             them to categorical or one-hot.
         seed (int, optional): the seed.  Defaults to ``None``.
         done_on_any (bool, optional): whether the environment's done keys are set by aggregating the agent keys
-            using ``any()`` (when True) or ``all()`` (when False). Default (``None``) is to use ``any()`` for
+            using ``any()`` (when ``True``) or ``all()`` (when ``False``). Default (``None``) is to use ``any()`` for
             parallel environments and ``all()`` for AEC ones.
 
     Examples:

From 16685ce03794f196af46c3e181d04be95ffbdc19 Mon Sep 17 00:00:00 2001
From: Matteo Bettini <55539777+matteobettini@users.noreply.github.com>
Date: Wed, 3 Apr 2024 07:39:31 +0100
Subject: [PATCH 09/16] Update torchrl/envs/libs/pettingzoo.py

Co-authored-by: Vincent Moens <vincentmoens@gmail.com>
---
 torchrl/envs/libs/pettingzoo.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/torchrl/envs/libs/pettingzoo.py b/torchrl/envs/libs/pettingzoo.py
index 8cd41251d29..0d8d2ffe687 100644
--- a/torchrl/envs/libs/pettingzoo.py
+++ b/torchrl/envs/libs/pettingzoo.py
@@ -207,7 +207,7 @@ def __init__(
         use_mask: bool = False,
         categorical_actions: bool = True,
         seed: int | None = None,
-        done_on_any: bool = None,
+        done_on_any: bool | None = None,
         **kwargs,
     ):
         if env is not None:

From 301790961f1c4f6b2faeda167ba925daedb4b940 Mon Sep 17 00:00:00 2001
From: Matteo Bettini <55539777+matteobettini@users.noreply.github.com>
Date: Wed, 3 Apr 2024 07:39:38 +0100
Subject: [PATCH 10/16] Update torchrl/envs/libs/pettingzoo.py

Co-authored-by: Vincent Moens <vincentmoens@gmail.com>
---
 torchrl/envs/libs/pettingzoo.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/torchrl/envs/libs/pettingzoo.py b/torchrl/envs/libs/pettingzoo.py
index 0d8d2ffe687..812ccdb07fa 100644
--- a/torchrl/envs/libs/pettingzoo.py
+++ b/torchrl/envs/libs/pettingzoo.py
@@ -157,7 +157,7 @@ class PettingZooWrapper(_EnvWrapper):
             them to categorical or one-hot.
         seed (int, optional): the seed. Defaults to ``None``.
         done_on_any (bool, optional): whether the environment's done keys are set by aggregating the agent keys
-            using ``any()`` (when True) or ``all()`` (when False). Default (``None``) is to use ``any()`` for
+            using ``any()`` (when ``True``) or ``all()`` (when ``False``). Default (``None``) is to use ``any()`` for
             parallel environments and ``all()`` for AEC ones.
 
     Examples:

From 4b00e30ec53bcbb3135c441ae5a2f5da97afabcb Mon Sep 17 00:00:00 2001
From: Matteo Bettini <mb2389@cl.cam.ac.uk>
Date: Mon, 8 Apr 2024 11:24:20 +0200
Subject: [PATCH 11/16] warn on version

---
 torchrl/envs/libs/pettingzoo.py | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/torchrl/envs/libs/pettingzoo.py b/torchrl/envs/libs/pettingzoo.py
index 812ccdb07fa..ea9d5b12f3b 100644
--- a/torchrl/envs/libs/pettingzoo.py
+++ b/torchrl/envs/libs/pettingzoo.py
@@ -270,6 +270,13 @@ def _build_env(
     ):
         import pettingzoo
 
+        if importlib.metadata.version("pettingzoo") > "1.24.3":
+            warnings.warn(
+                "PettingZoo in TorchRL is tested using version == 1.24.3."
+                "If you are using a later version and are experiencing backwards compatibility issues,"
+                "please raise an issue in the TorchRL github"
+            )
+
         self.parallel = isinstance(env, pettingzoo.utils.env.ParallelEnv)
         if not self.parallel and not self.use_mask:
             raise ValueError("For AEC environments you need to set use_mask=True")

From 5c44a39c5198b60344c663f61cc0d5743c555844 Mon Sep 17 00:00:00 2001
From: Matteo Bettini <mb2389@cl.cam.ac.uk>
Date: Mon, 8 Apr 2024 11:29:19 +0200
Subject: [PATCH 12/16] typos

---
 torchrl/envs/libs/pettingzoo.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/torchrl/envs/libs/pettingzoo.py b/torchrl/envs/libs/pettingzoo.py
index ea9d5b12f3b..48bf813dc46 100644
--- a/torchrl/envs/libs/pettingzoo.py
+++ b/torchrl/envs/libs/pettingzoo.py
@@ -272,9 +272,9 @@ def _build_env(
 
         if importlib.metadata.version("pettingzoo") > "1.24.3":
             warnings.warn(
-                "PettingZoo in TorchRL is tested using version == 1.24.3."
+                "PettingZoo in TorchRL is tested using version == 1.24.3 , "
                 "If you are using a later version and are experiencing backwards compatibility issues,"
-                "please raise an issue in the TorchRL github"
+                "please raise an issue in the TorchRL github."
             )
 
         self.parallel = isinstance(env, pettingzoo.utils.env.ParallelEnv)

From ebd612d9978d2036b68f91feb2c6788c31987972 Mon Sep 17 00:00:00 2001
From: Matteo Bettini <mb2389@cl.cam.ac.uk>
Date: Mon, 8 Apr 2024 11:44:03 +0200
Subject: [PATCH 13/16] amend

---
 torchrl/envs/libs/pettingzoo.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/torchrl/envs/libs/pettingzoo.py b/torchrl/envs/libs/pettingzoo.py
index 48bf813dc46..a3a0343ddbf 100644
--- a/torchrl/envs/libs/pettingzoo.py
+++ b/torchrl/envs/libs/pettingzoo.py
@@ -270,10 +270,10 @@ def _build_env(
     ):
         import pettingzoo
 
-        if importlib.metadata.version("pettingzoo") > "1.24.3":
+        if importlib.metadata.version("pettingzoo") != "1.24.3":
             warnings.warn(
                 "PettingZoo in TorchRL is tested using version == 1.24.3 , "
-                "If you are using a later version and are experiencing backwards compatibility issues,"
+                "If you are using a different version and are experiencing compatibility issues,"
                 "please raise an issue in the TorchRL github."
             )
 

From 57dedcf4f6011dc0e4b58b63b9b34927b196a00a Mon Sep 17 00:00:00 2001
From: Matteo Bettini <mb2389@cl.cam.ac.uk>
Date: Mon, 8 Apr 2024 12:36:27 +0200
Subject: [PATCH 14/16] typo importlib metadata

---
 .github/unittest/linux_libs/scripts_pettingzoo/environment.yml | 2 +-
 .github/unittest/linux_libs/scripts_smacv2/environment.yml     | 2 +-
 .github/unittest/linux_libs/scripts_vmas/environment.yml       | 2 +-
 3 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/.github/unittest/linux_libs/scripts_pettingzoo/environment.yml b/.github/unittest/linux_libs/scripts_pettingzoo/environment.yml
index 66f79eb2c24..90bd75469eb 100644
--- a/.github/unittest/linux_libs/scripts_pettingzoo/environment.yml
+++ b/.github/unittest/linux_libs/scripts_pettingzoo/environment.yml
@@ -8,7 +8,7 @@ dependencies:
     - cloudpickle
     - gym
     - gym-notices
-    - importlib-metadata
+    - importlib_metadata
     - six
     - zipp
     - pytest
diff --git a/.github/unittest/linux_libs/scripts_smacv2/environment.yml b/.github/unittest/linux_libs/scripts_smacv2/environment.yml
index d1e1e1f5edc..7ad0696f8d5 100644
--- a/.github/unittest/linux_libs/scripts_smacv2/environment.yml
+++ b/.github/unittest/linux_libs/scripts_smacv2/environment.yml
@@ -7,7 +7,7 @@ dependencies:
     - cloudpickle
     - gym
     - gym-notices
-    - importlib-metadata
+    - importlib_metadata
     - zipp
     - pytest
     - pytest-cov
diff --git a/.github/unittest/linux_libs/scripts_vmas/environment.yml b/.github/unittest/linux_libs/scripts_vmas/environment.yml
index c77fd3a0bdc..6c2e8cf52fa 100644
--- a/.github/unittest/linux_libs/scripts_vmas/environment.yml
+++ b/.github/unittest/linux_libs/scripts_vmas/environment.yml
@@ -7,7 +7,7 @@ dependencies:
     - cloudpickle
     - gym
     - gym-notices
-    - importlib-metadata
+    - importlib_metadata
     - numpy
     - pyglet==1.5.27
     - six

From 0ed8f170c02381bb6b4236e8ba0fd870fae36408 Mon Sep 17 00:00:00 2001
From: Matteo Bettini <mb2389@cl.cam.ac.uk>
Date: Mon, 8 Apr 2024 15:02:43 +0200
Subject: [PATCH 15/16] amend

---
 torchrl/envs/libs/pettingzoo.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/torchrl/envs/libs/pettingzoo.py b/torchrl/envs/libs/pettingzoo.py
index a3a0343ddbf..eb94a27cbba 100644
--- a/torchrl/envs/libs/pettingzoo.py
+++ b/torchrl/envs/libs/pettingzoo.py
@@ -9,6 +9,7 @@
 import warnings
 from typing import Dict, List, Tuple, Union
 
+import packaging
 import torch
 from tensordict import TensorDictBase
 
@@ -270,7 +271,7 @@ def _build_env(
     ):
         import pettingzoo
 
-        if importlib.metadata.version("pettingzoo") != "1.24.3":
+        if packaging.version.parse(pettingzoo.__version__).base_version != "1.24.3":
             warnings.warn(
                 "PettingZoo in TorchRL is tested using version == 1.24.3 , "
                 "If you are using a different version and are experiencing compatibility issues,"

From 9d7d5154b8a56f65ecd20a3100864d0007caa444 Mon Sep 17 00:00:00 2001
From: Matteo Bettini <mb2389@cl.cam.ac.uk>
Date: Mon, 8 Apr 2024 15:35:50 +0200
Subject: [PATCH 16/16] remove importlib_metadata

---
 .github/unittest/linux_libs/scripts_pettingzoo/environment.yml | 1 -
 .github/unittest/linux_libs/scripts_smacv2/environment.yml     | 1 -
 .github/unittest/linux_libs/scripts_vmas/environment.yml       | 1 -
 3 files changed, 3 deletions(-)

diff --git a/.github/unittest/linux_libs/scripts_pettingzoo/environment.yml b/.github/unittest/linux_libs/scripts_pettingzoo/environment.yml
index 90bd75469eb..f6c79784a0b 100644
--- a/.github/unittest/linux_libs/scripts_pettingzoo/environment.yml
+++ b/.github/unittest/linux_libs/scripts_pettingzoo/environment.yml
@@ -8,7 +8,6 @@ dependencies:
     - cloudpickle
     - gym
     - gym-notices
-    - importlib_metadata
     - six
     - zipp
     - pytest
diff --git a/.github/unittest/linux_libs/scripts_smacv2/environment.yml b/.github/unittest/linux_libs/scripts_smacv2/environment.yml
index 7ad0696f8d5..f6dae836c91 100644
--- a/.github/unittest/linux_libs/scripts_smacv2/environment.yml
+++ b/.github/unittest/linux_libs/scripts_smacv2/environment.yml
@@ -7,7 +7,6 @@ dependencies:
     - cloudpickle
     - gym
     - gym-notices
-    - importlib_metadata
     - zipp
     - pytest
     - pytest-cov
diff --git a/.github/unittest/linux_libs/scripts_vmas/environment.yml b/.github/unittest/linux_libs/scripts_vmas/environment.yml
index 6c2e8cf52fa..58fa50e828d 100644
--- a/.github/unittest/linux_libs/scripts_vmas/environment.yml
+++ b/.github/unittest/linux_libs/scripts_vmas/environment.yml
@@ -7,7 +7,6 @@ dependencies:
     - cloudpickle
     - gym
     - gym-notices
-    - importlib_metadata
     - numpy
     - pyglet==1.5.27
     - six