pettingzoo.py

# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
from __future__ import annotations

import copy
import importlib
import warnings
from typing import Dict, List, Tuple, Union

import packaging
import torch
from tensordict import TensorDictBase

from torchrl.data.tensor_specs import (
    CompositeSpec,
    DiscreteTensorSpec,
    OneHotDiscreteTensorSpec,
    UnboundedContinuousTensorSpec,
)
from torchrl.envs.common import _EnvWrapper
from torchrl.envs.libs.gym import _gym_to_torchrl_spec_transform, set_gym_backend
from torchrl.envs.utils import _classproperty, check_marl_grouping, MarlGroupMapType

_has_pettingzoo = importlib.util.find_spec("pettingzoo") is not None


def _get_envs():
    if not _has_pettingzoo:
        raise ImportError("PettingZoo is not installed in your virtual environment.")
    try:
        from pettingzoo.utils.all_modules import all_environments
    except ModuleNotFoundError as err:
        warnings.warn(
            f"PettingZoo failed to load all modules with error message {err}, trying to load individual modules."
        )
        all_environments = _load_available_envs()

    return list(all_environments.keys())


def _load_available_envs() -> Dict:
    all_environments = {}
    try:
        from pettingzoo.mpe.all_modules import mpe_environments

        all_environments.update(mpe_environments)
    except ModuleNotFoundError as err:
        warnings.warn(f"MPE environments failed to load with error message {err}.")
    try:
        from pettingzoo.sisl.all_modules import sisl_environments

        all_environments.update(sisl_environments)
    except ModuleNotFoundError as err:
        warnings.warn(f"SISL environments failed to load with error message {err}.")
    try:
        from pettingzoo.classic.all_modules import classic_environments

        all_environments.update(classic_environments)
    except ModuleNotFoundError as err:
        warnings.warn(f"Classic environments failed to load with error message {err}.")
    try:
        from pettingzoo.atari.all_modules import atari_environments

        all_environments.update(atari_environments)
    except ModuleNotFoundError as err:
        warnings.warn(f"Atari environments failed to load with error message {err}.")
    try:
        from pettingzoo.butterfly.all_modules import butterfly_environments

        all_environments.update(butterfly_environments)
    except ModuleNotFoundError as err:
        warnings.warn(
            f"Butterfly environments failed to load with error message {err}."
        )
    return all_environments


class PettingZooWrapper(_EnvWrapper):
    """PettingZoo environment wrapper.

    To install petting zoo follow the guide `here <https://github.com/Farama-Foundation/PettingZoo#installation>__`.

    This class is a general torchrl wrapper for all PettingZoo environments.
    It is able to wrap both ``pettingzoo.AECEnv`` and ``pettingzoo.ParallelEnv``.

    Let's see how more in details:

    In wrapped ``pettingzoo.ParallelEnv`` all agents will step at each environment step.
    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 (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 (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"``
    for each group to reflect the latest available actions. This should be passed to a masked distribution during
    training.

    As a feature of torchrl multiagent, you are able to control the grouping of agents in your environment.
    You can group agents together (stacking their tensors) to leverage vectorization when passing them through the same
    neural network. You can split agents in different groups where they are heterogenous or should be processed by
    different neural networks. To group, you just need to pass a ``group_map`` at env constructiuon time.

    By default, agents in pettingzoo will be grouped by name.
    For example, with agents ``["agent_0","agent_1","agent_2","adversary_0"]``, the tensordicts will look like:

        >>> print(env.rand_action(env.reset()))
        TensorDict(
            fields={
                agent: TensorDict(
                    fields={
                        action: Tensor(shape=torch.Size([3, 9]), device=cpu, dtype=torch.int64, is_shared=False),
                        action_mask: Tensor(shape=torch.Size([3, 9]), device=cpu, dtype=torch.bool, is_shared=False),
                        done: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.bool, is_shared=False),
                        observation: Tensor(shape=torch.Size([3, 3, 3, 2]), device=cpu, dtype=torch.int8, is_shared=False),
                        terminated: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.bool, is_shared=False),
                        truncated: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
                    batch_size=torch.Size([3]))},
                adversary: TensorDict(
                    fields={
                        action: Tensor(shape=torch.Size([1, 9]), device=cpu, dtype=torch.int64, is_shared=False),
                        action_mask: Tensor(shape=torch.Size([1, 9]), device=cpu, dtype=torch.bool, is_shared=False),
                        done: Tensor(shape=torch.Size([1, 1]), device=cpu, dtype=torch.bool, is_shared=False),
                        observation: Tensor(shape=torch.Size([1, 3, 3, 2]), device=cpu, dtype=torch.int8, is_shared=False),
                        terminated: Tensor(shape=torch.Size([1, 1]), device=cpu, dtype=torch.bool, is_shared=False),
                        truncated: Tensor(shape=torch.Size([1, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
                    batch_size=torch.Size([1]))},
            batch_size=torch.Size([]))
        >>> print(env.group_map)
        {"agent": ["agent_0", "agent_1", "agent_2"], "adversary": ["adversary_0"]}

    Otherwise, a group map can be specified or selected from some premade options.
    See :class:`torchrl.envs.utils.MarlGroupMapType` for more info.
    For example, you can provide ``MarlGroupMapType.ONE_GROUP_PER_AGENT``, telling that each agent should
    have its own tensordict (similar to the pettingzoo parallel API).

    Grouping is useful for leveraging vectorisation among agents whose data goes through the same
    neural network.

    Args:
        env (``pettingzoo.utils.env.ParallelEnv`` or ``pettingzoo.utils.env.AECEnv``): the pettingzoo environment to wrap.
        return_state (bool, optional): whether to return the global state from pettingzoo
            (not available in all environments). Defaults to ``False``.
        group_map (MarlGroupMapType or Dict[str, List[str]]], optional): how to group agents in tensordicts for
            input/output. By default, agents will be grouped by their name. Otherwise, a group map can be specified
            or selected from some premade options. See :class:`torchrl.envs.utils.MarlGroupMapType` for more info.
        use_mask (bool, optional): whether the environment should output a ``"mask"``. This is compulsory in
            wrapped ``pettingzoo.AECEnv`` to mask out non-acting agents and should be also used
            for ``pettingzoo.ParallelEnv`` when the number of agents can vary. Defaults to ``False``.
        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
        >>> from torchrl.envs.libs.pettingzoo import PettingZooWrapper
        >>> from pettingzoo.butterfly import pistonball_v6
        >>> kwargs = {"n_pistons": 21, "continuous": True}
        >>> env = PettingZooWrapper(
        ...     env=pistonball_v6.parallel_env(**kwargs),
        ...     return_state=True,
        ...     group_map=None, # Use default for parallel (all pistons grouped together)
        ... )
        >>> print(env.group_map)
        ... {'piston': ['piston_0', 'piston_1', ..., 'piston_20']}
        >>> env.rollout(10)
        >>> # AEC env
        >>> from pettingzoo.classic import tictactoe_v3
        >>> from torchrl.envs.libs.pettingzoo import PettingZooWrapper
        >>> from torchrl.envs.utils import MarlGroupMapType
        >>> env = PettingZooWrapper(
        ...     env=tictactoe_v3.env(),
        ...     use_mask=True, # Must use it since one player plays at a time
        ...     group_map=None # # Use default for AEC (one group per player)
        ... )
        >>> print(env.group_map)
        ... {'player_1': ['player_1'], 'player_2': ['player_2']}
        >>> env.rollout(10)
    """

    git_url = "https://github.com/Farama-Foundation/PettingZoo"
    libname = "pettingzoo"

    @_classproperty
    def available_envs(cls):
        if not _has_pettingzoo:
            return []
        return list(_get_envs())

    def __init__(
        self,
        env: Union[
            "pettingzoo.utils.env.ParallelEnv",  # noqa: F821
            "pettingzoo.utils.env.AECEnv",  # noqa: F821
        ] = None,
        return_state: bool = False,
        group_map: MarlGroupMapType | Dict[str, List[str]] | None = None,
        use_mask: bool = False,
        categorical_actions: bool = True,
        seed: int | None = None,
        done_on_any: bool | None = None,
        **kwargs,
    ):
        if env is not None:
            kwargs["env"] = env

        self.group_map = group_map
        self.return_state = return_state
        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)

    def _get_default_group_map(self, agent_names: List[str]):
        # This function performs the default grouping in pettingzoo
        if not self.parallel:
            # In AEC envs we will have one group per agent by default
            group_map = MarlGroupMapType.ONE_GROUP_PER_AGENT.get_group_map(agent_names)
        else:
            # In parallel envs, by default
            # Agents with names "str_int" will be grouped in group name "str"
            group_map = {}
            for agent_name in agent_names:
                # See if the agent follows the convention "name_int"
                follows_convention = True
                agent_name_split = agent_name.split("_")
                if len(agent_name_split) == 1:
                    follows_convention = False
                try:
                    int(agent_name_split[-1])
                except ValueError:
                    follows_convention = False

                # If not, just put it in a single group
                if not follows_convention:
                    group_map[agent_name] = [agent_name]
                # Otherwise, group it with other agents that follow the same convention
                else:
                    group_name = "_".join(agent_name_split[:-1])
                    if group_name in group_map:
                        group_map[group_name].append(agent_name)
                    else:
                        group_map[group_name] = [agent_name]

        return group_map

    @property
    def lib(self):
        import pettingzoo

        return pettingzoo

    def _build_env(
        self,
        env: Union[
            "pettingzoo.utils.env.ParallelEnv",  # noqa: F821
            "pettingzoo.utils.env.AECEnv",  # noqa: F821
        ],
    ):
        import pettingzoo

        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,"
                "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")
        if len(self.batch_size):
            raise RuntimeError(
                f"PettingZoo does not support custom batch_size {self.batch_size}."
            )

        return env

    @set_gym_backend("gymnasium")
    def _make_specs(
        self,
        env: Union[
            "pettingzoo.utils.env.ParallelEnv",  # noqa: F821
            "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:
            self.group_map = self._get_default_group_map(self.possible_agents)
        elif isinstance(self.group_map, MarlGroupMapType):
            self.group_map = self.group_map.get_group_map(self.possible_agents)
        check_marl_grouping(self.group_map, self.possible_agents)
        self.has_action_mask = {group: False for group in self.group_map.keys()}

        action_spec = CompositeSpec()
        observation_spec = CompositeSpec()
        reward_spec = CompositeSpec()
        done_spec = CompositeSpec(
            {
                "done": DiscreteTensorSpec(
                    n=2,
                    shape=torch.Size((1,)),
                    dtype=torch.bool,
                    device=self.device,
                ),
                "terminated": DiscreteTensorSpec(
                    n=2,
                    shape=torch.Size((1,)),
                    dtype=torch.bool,
                    device=self.device,
                ),
                "truncated": DiscreteTensorSpec(
                    n=2,
                    shape=torch.Size((1,)),
                    dtype=torch.bool,
                    device=self.device,
                ),
            },
        )
        for group, agents in self.group_map.items():
            (
                group_observation_spec,
                group_action_spec,
                group_reward_spec,
                group_done_spec,
            ) = self._make_group_specs(group_name=group, agent_names=agents)
            action_spec[group] = group_action_spec
            observation_spec[group] = group_observation_spec
            reward_spec[group] = group_reward_spec
            done_spec[group] = group_done_spec

        self.action_spec = action_spec
        self.observation_spec = observation_spec
        self.reward_spec = reward_spec
        self.done_spec = done_spec

    def _make_group_specs(self, group_name: str, agent_names: List[str]):
        n_agents = len(agent_names)
        action_specs = []
        observation_specs = []
        for agent in agent_names:
            action_specs.append(
                CompositeSpec(
                    {
                        "action": _gym_to_torchrl_spec_transform(
                            self.action_space(agent),
                            remap_state_to_observation=False,
                            categorical_action_encoding=self.categorical_actions,
                            device=self.device,
                        )
                    },
                )
            )
            observation_specs.append(
                CompositeSpec(
                    {
                        "observation": _gym_to_torchrl_spec_transform(
                            self.observation_space(agent),
                            remap_state_to_observation=False,
                            device=self.device,
                        )
                    }
                )
            )
        group_action_spec = torch.stack(action_specs, dim=0)
        group_observation_spec = torch.stack(observation_specs, dim=0)

        # Sometimes the observation spec contains an action mask.
        # Or sometimes the info spec contains an action mask.
        # We uniform this by removing it from both places and optionally set it in a standard location.
        group_observation_inner_spec = group_observation_spec["observation"]
        if (
            isinstance(group_observation_inner_spec, CompositeSpec)
            and "action_mask" in group_observation_inner_spec.keys()
        ):
            self.has_action_mask[group_name] = True
            del group_observation_inner_spec["action_mask"]
            group_observation_spec["action_mask"] = DiscreteTensorSpec(
                n=2,
                shape=group_action_spec["action"].shape
                if not self.categorical_actions
                else (
                    *group_action_spec["action"].shape,
                    group_action_spec["action"].space.n,
                ),
                dtype=torch.bool,
                device=self.device,
            )

        if self.use_mask:
            group_observation_spec["mask"] = DiscreteTensorSpec(
                n=2,
                shape=torch.Size((n_agents,)),
                dtype=torch.bool,
                device=self.device,
            )

        group_reward_spec = CompositeSpec(
            {
                "reward": UnboundedContinuousTensorSpec(
                    shape=torch.Size((n_agents, 1)),
                    device=self.device,
                    dtype=torch.float32,
                )
            },
            shape=torch.Size((n_agents,)),
        )
        group_done_spec = CompositeSpec(
            {
                "done": DiscreteTensorSpec(
                    n=2,
                    shape=torch.Size((n_agents, 1)),
                    dtype=torch.bool,
                    device=self.device,
                ),
                "terminated": DiscreteTensorSpec(
                    n=2,
                    shape=torch.Size((n_agents, 1)),
                    dtype=torch.bool,
                    device=self.device,
                ),
                "truncated": DiscreteTensorSpec(
                    n=2,
                    shape=torch.Size((n_agents, 1)),
                    dtype=torch.bool,
                    device=self.device,
                ),
            },
            shape=torch.Size((n_agents,)),
        )
        return (
            group_observation_spec,
            group_action_spec,
            group_reward_spec,
            group_done_spec,
        )

    def _check_kwargs(self, kwargs: Dict):
        import pettingzoo

        if "env" not in kwargs:
            raise TypeError("Could not find environment key 'env' in kwargs.")
        env = kwargs["env"]
        if not isinstance(
            env, (pettingzoo.utils.env.ParallelEnv, pettingzoo.utils.env.AECEnv)
        ):
            raise TypeError("env is not of type expected.")

    def _init_env(self):
        # Add info
        if self.parallel:
            _, info_dict = self._reset_parallel(seed=self.seed)
        else:
            _, info_dict = self._reset_aec(seed=self.seed)

        for group, agents in self.group_map.items():
            info_specs = []
            for agent in agents:
                info_specs.append(
                    CompositeSpec(
                        {
                            "info": CompositeSpec(
                                {
                                    key: UnboundedContinuousTensorSpec(
                                        shape=torch.as_tensor(value).shape,
                                        device=self.device,
                                    )
                                    for key, value in info_dict[agent].items()
                                }
                            )
                        },
                        device=self.device,
                    )
                )
            info_specs = torch.stack(info_specs, dim=0)
            if ("info", "action_mask") in info_specs.keys(True, True):
                if not self.has_action_mask[group]:
                    self.has_action_mask[group] = True
                    group_action_spec = self.input_spec[
                        "full_action_spec", group, "action"
                    ]
                    self.observation_spec[group]["action_mask"] = DiscreteTensorSpec(
                        n=2,
                        shape=group_action_spec.shape
                        if not self.categorical_actions
                        else (*group_action_spec.shape, group_action_spec.space.n),
                        dtype=torch.bool,
                        device=self.device,
                    )
                group_inner_info_spec = info_specs["info"]
                del group_inner_info_spec["action_mask"]

            if len(info_specs["info"].keys()):
                self.observation_spec[group].update(info_specs)

        if self.return_state:
            try:
                state_spec = _gym_to_torchrl_spec_transform(
                    self.state_space,
                    remap_state_to_observation=False,
                    device=self.device,
                )
            except AttributeError:
                state_example = torch.as_tensor(self.state(), device=self.device)
                state_spec = UnboundedContinuousTensorSpec(
                    shape=state_example.shape,
                    dtype=state_example.dtype,
                    device=self.device,
                )
            self.observation_spec["state"] = state_spec

        # Caching
        self.cached_reset_output_zero = self.observation_spec.zero()
        self.cached_reset_output_zero.update(self.output_spec["full_done_spec"].zero())

        self.cached_step_output_zero = self.observation_spec.zero()
        self.cached_step_output_zero.update(self.output_spec["full_reward_spec"].zero())
        self.cached_step_output_zero.update(self.output_spec["full_done_spec"].zero())

    def _set_seed(self, seed: int):
        self.seed = seed
        self.reset(seed=self.seed)

    def _reset(
        self, tensordict: TensorDictBase | None = None, **kwargs
    ) -> TensorDictBase:
        if tensordict is not None:
            _reset = tensordict.get("_reset", None)
            if _reset is not None and not _reset.all():
                raise RuntimeError(
                    f"An attempt to call {type(self)}._reset was made when no "
                    f"reset signal could be found. Expected '_reset' entry to "
                    f"be `tensor(True)` or `None` but got `{_reset}`."
                )
        if self.parallel:
            # This resets when any is done
            observation_dict, info_dict = self._reset_parallel(**kwargs)
        else:
            # This resets when all are done
            observation_dict, info_dict = self._reset_aec(**kwargs)

        # We start with zeroed data and fill in the data for alive agents
        tensordict_out = self.cached_reset_output_zero.clone()
        # Update the "mask" for non-acting agents
        self._update_agent_mask(tensordict_out)
        # Update the "action_mask" for non-available actions
        observation_dict, info_dict = self._update_action_mask(
            tensordict_out, observation_dict, info_dict
        )

        # Now we get the data (obs and info)
        for group, agent_names in self.group_map.items():
            group_observation = tensordict_out.get((group, "observation"))
            group_info = tensordict_out.get((group, "info"), None)

            for index, agent in enumerate(agent_names):
                group_observation[index] = self.observation_spec[group, "observation"][
                    index
                ].encode(observation_dict[agent])
                if group_info is not None:
                    agent_info_dict = info_dict[agent]
                    for agent_info, value in agent_info_dict.items():
                        group_info.get(agent_info)[index] = torch.as_tensor(
                            value, device=self.device
                        )

        return tensordict_out

    def _reset_aec(self, **kwargs) -> Tuple[Dict, Dict]:
        self._env.reset(**kwargs)

        observation_dict = {
            agent: self._env.observe(agent) for agent in self.possible_agents
        }
        info_dict = self._env.infos
        return observation_dict, info_dict

    def _reset_parallel(self, **kwargs) -> Tuple[Dict, Dict]:
        return self._env.reset(**kwargs)

    def _step(
        self,
        tensordict: TensorDictBase,
    ) -> TensorDictBase:
        if self.parallel:
            (
                observation_dict,
                rewards_dict,
                terminations_dict,
                truncations_dict,
                info_dict,
            ) = self._step_parallel(tensordict)
        else:
            (
                observation_dict,
                rewards_dict,
                terminations_dict,
                truncations_dict,
                info_dict,
            ) = self._step_aec(tensordict)

        # We start with zeroed data and fill in the data for alive agents
        tensordict_out = self.cached_step_output_zero.clone()
        # Update the "mask" for non-acting agents
        self._update_agent_mask(tensordict_out)
        # Update the "action_mask" for non-available actions
        observation_dict, info_dict = self._update_action_mask(
            tensordict_out, observation_dict, info_dict
        )

        # Now we get the data
        for group, agent_names in self.group_map.items():
            group_observation = tensordict_out.get((group, "observation"))
            group_reward = tensordict_out.get((group, "reward"))
            group_done = tensordict_out.get((group, "done"))
            group_terminated = tensordict_out.get((group, "terminated"))
            group_truncated = tensordict_out.get((group, "truncated"))
            group_info = tensordict_out.get((group, "info"), None)

            for index, agent in enumerate(agent_names):
                if agent in observation_dict:  # Live agents
                    group_observation[index] = self.observation_spec[
                        group, "observation"
                    ][index].encode(observation_dict[agent])
                    group_reward[index] = torch.tensor(
                        rewards_dict[agent],
                        device=self.device,
                        dtype=torch.float32,
                    )
                    group_done[index] = torch.tensor(
                        terminations_dict[agent] or truncations_dict[agent],
                        device=self.device,
                        dtype=torch.bool,
                    )
                    group_truncated[index] = torch.tensor(
                        truncations_dict[agent],
                        device=self.device,
                        dtype=torch.bool,
                    )
                    group_terminated[index] = torch.tensor(
                        terminations_dict[agent],
                        device=self.device,
                        dtype=torch.bool,
                    )

                    if group_info is not None:
                        agent_info_dict = info_dict[agent]
                        for agent_info, value in agent_info_dict.items():
                            group_info.get(agent_info)[index] = torch.tensor(
                                value, device=self.device
                            )

                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_mask=True to allow this."
                    )

        # set done values
        done, terminated, truncated = self._aggregate_done(
            tensordict_out, use_any=self.done_on_any
        )

        tensordict_out.set("done", done)
        tensordict_out.set("terminated", terminated)
        tensordict_out.set("truncated", truncated)
        return tensordict_out

    def _aggregate_done(self, tensordict_out, use_any):
        done = False if use_any else True
        truncated = False if use_any else True
        terminated = False if use_any else True
        for key in self.done_keys:
            if isinstance(key, tuple):  # Only look at group keys
                if use_any:
                    if key[-1] == "done":
                        done = done | tensordict_out.get(key).any()
                    if key[-1] == "terminated":
                        terminated = terminated | tensordict_out.get(key).any()
                    if key[-1] == "truncated":
                        truncated = truncated | tensordict_out.get(key).any()
                    if done and terminated and truncated:
                        # no need to proceed further, all values are flipped
                        break
                else:
                    if key[-1] == "done":
                        done = done & tensordict_out.get(key).all()
                    if key[-1] == "terminated":
                        terminated = terminated & tensordict_out.get(key).all()
                    if key[-1] == "truncated":
                        truncated = truncated & tensordict_out.get(key).all()
                    if not done and not terminated and not truncated:
                        # no need to proceed further, all values are flipped
                        break
        return (
            torch.tensor([done], device=self.device),
            torch.tensor([terminated], device=self.device),
            torch.tensor([truncated], device=self.device),
        )

    def _step_parallel(
        self,
        tensordict: TensorDictBase,
    ) -> Tuple[Dict, Dict, Dict, Dict, Dict]:
        action_dict = {}
        for group, agents in self.group_map.items():
            group_action = tensordict.get((group, "action"))
            group_action_np = self.input_spec[
                "full_action_spec", group, "action"
            ].to_numpy(group_action)
            for index, agent in enumerate(agents):
                action_dict[agent] = group_action_np[index]

        return self._env.step(action_dict)

    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)
                group_action = tensordict.get((group, "action"))
                group_action_np = self.input_spec[
                    "full_action_spec", group, "action"
                ].to_numpy(group_action)
                action = group_action_np[agent_index]
                break

        self._env.step(action)
        terminations_dict = self._env.terminations
        truncations_dict = self._env.truncations
        info_dict = self._env.infos
        rewards_dict = self._env.rewards
        observation_dict = {
            agent: self._env.observe(agent) for agent in self.possible_agents
        }
        return (
            observation_dict,
            rewards_dict,
            terminations_dict,
            truncations_dict,
            info_dict,
        )

    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)
        # In AEC only one agent acts, in parallel env self.agents contains the agents alive
        agents_acting = self.agents if self.parallel else [self.agent_selection]

        for group, agents in self.group_map.items():
            if self.has_action_mask[group]:
                group_mask = td.get((group, "action_mask"))
                group_mask += True
                for index, agent in enumerate(agents):
                    agent_obs = observation_dict[agent]
                    agent_info = info_dict[agent]
                    if isinstance(agent_obs, Dict) and "action_mask" in agent_obs:
                        if agent in agents_acting:
                            group_mask[index] = torch.tensor(
                                agent_obs["action_mask"],
                                device=self.device,
                                dtype=torch.bool,
                            )
                        del agent_obs["action_mask"]
                    elif isinstance(agent_info, Dict) and "action_mask" in agent_info:
                        if agent in agents_acting:
                            group_mask[index] = torch.tensor(
                                agent_info["action_mask"],
                                device=self.device,
                                dtype=torch.bool,
                            )
                        del agent_info["action_mask"]

                group_action_spec = self.input_spec["full_action_spec", group, "action"]
                if isinstance(
                    group_action_spec, (DiscreteTensorSpec, OneHotDiscreteTensorSpec)
                ):
                    # We update the mask for available actions
                    group_action_spec.update_mask(group_mask.clone())

        return observation_dict, info_dict

    def _update_agent_mask(self, td):
        if self.use_mask:
            # In AEC only one agent acts, in parallel env self.agents contains the agents alive
            agents_acting = self.agents if self.parallel else [self.agent_selection]
            for group, agents in self.group_map.items():
                group_mask = td.get((group, "mask"))
                group_mask += True

                # We now add dead agents to the mask
                for index, agent in enumerate(agents):
                    if agent not in agents_acting:
                        group_mask[index] = False

    def close(self) -> None:
        self._env.close()


class PettingZooEnv(PettingZooWrapper):
    """PettingZoo Environment.

    To install petting zoo follow the guide `here <https://github.com/Farama-Foundation/PettingZoo#installation>__`.

    This class is a general torchrl wrapper for all PettingZoo environments.
    It is able to wrap both ``pettingzoo.AECEnv`` and ``pettingzoo.ParallelEnv``.

    Let's see how more in details:

    For wrapping ``pettingzoo.ParallelEnv`` provide the name of your petting zoo task (in the ``task`` argument)
    and specify ``parallel=True``. This will construct the ``pettingzoo.ParallelEnv`` version of that task
    (if it is supported in pettingzoo) and wrap it for torchrl.
    In wrapped ``pettingzoo.ParallelEnv`` all agents will step at each environment step.
    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 (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
    and wrap it for torchrl.
    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 (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"``
    for each group to reflect the latest available actions. This should be passed to a masked distribution during
    training.

    As a feature of torchrl multiagent, you are able to control the grouping of agents in your environment.
    You can group agents together (stacking their tensors) to leverage vectorization when passing them through the same
    neural network. You can split agents in different groups where they are heterogenous or should be processed by
    different neural networks. To group, you just need to pass a ``group_map`` at env constructiuon time.

    By default, agents in pettingzoo will be grouped by name.
    For example, with agents ``["agent_0","agent_1","agent_2","adversary_0"]``, the tensordicts will look like:

        >>> print(env.rand_action(env.reset()))
        TensorDict(
            fields={
                agent: TensorDict(
                    fields={
                        action: Tensor(shape=torch.Size([3, 9]), device=cpu, dtype=torch.int64, is_shared=False),
                        action_mask: Tensor(shape=torch.Size([3, 9]), device=cpu, dtype=torch.bool, is_shared=False),
                        done: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.bool, is_shared=False),
                        observation: Tensor(shape=torch.Size([3, 3, 3, 2]), device=cpu, dtype=torch.int8, is_shared=False),
                        terminated: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.bool, is_shared=False),
                        truncated: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
                    batch_size=torch.Size([3]))},
                adversary: TensorDict(
                    fields={
                        action: Tensor(shape=torch.Size([1, 9]), device=cpu, dtype=torch.int64, is_shared=False),
                        action_mask: Tensor(shape=torch.Size([1, 9]), device=cpu, dtype=torch.bool, is_shared=False),
                        done: Tensor(shape=torch.Size([1, 1]), device=cpu, dtype=torch.bool, is_shared=False),
                        observation: Tensor(shape=torch.Size([1, 3, 3, 2]), device=cpu, dtype=torch.int8, is_shared=False),
                        terminated: Tensor(shape=torch.Size([1, 1]), device=cpu, dtype=torch.bool, is_shared=False),
                        truncated: Tensor(shape=torch.Size([1, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
                    batch_size=torch.Size([1]))},
            batch_size=torch.Size([]))
        >>> print(env.group_map)
        {"agent": ["agent_0", "agent_1", "agent_2"], "adversary": ["adversary_0"]}

    Otherwise, a group map can be specified or selected from some premade options.
    See :class:`torchrl.envs.utils.MarlGroupMapType` for more info.
    For example, you can provide ``MarlGroupMapType.ONE_GROUP_PER_AGENT``, telling that each agent should
    have its own tensordict (similar to the pettingzoo parallel API).

    Grouping is useful for leveraging vectorisation among agents whose data goes through the same
    neural network.

    Args:
        task (str): the name of the pettingzoo task to create in the "<env>/<task>" format (for example, "sisl/multiwalker_v9")
            or "<task>" format (for example, "multiwalker_v9").
        parallel (bool): if to construct the ``pettingzoo.ParallelEnv`` version of the task or the ``pettingzoo.AECEnv``.
        return_state (bool, optional): whether to return the global state from pettingzoo
            (not available in all environments).  Defaults to ``False``.
        group_map (MarlGroupMapType or Dict[str, List[str]]], optional): how to group agents in tensordicts for
            input/output. By default, agents will be grouped by their name. Otherwise, a group map can be specified
            or selected from some premade options. See :class:`torchrl.envs.utils.MarlGroupMapType` for more info.
        use_mask (bool, optional): whether the environment should output an ``"mask"``. This is compulsory in
            wrapped ``pettingzoo.AECEnv`` to mask out non-acting agents and should be also used
            for ``pettingzoo.ParallelEnv`` when the number of agents can vary. Defaults to ``False``.
        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
        >>> from torchrl.envs.libs.pettingzoo import PettingZooEnv
        >>> kwargs = {"n_pistons": 21, "continuous": True}
        >>> env = PettingZooEnv(
        ...     task="pistonball_v6",
        ...     parallel=True,
        ...     return_state=True,
        ...     group_map=None, # Use default (all pistons grouped together)
        ...     **kwargs,
        ... )
        >>> print(env.group_map)
        ... {'piston': ['piston_0', 'piston_1', ..., 'piston_20']}
        >>> env.rollout(10)
        >>> # AEC env
        >>> from torchrl.envs.libs.pettingzoo import PettingZooEnv
        >>> from torchrl.envs.utils import MarlGroupMapType
        >>> env = PettingZooEnv(
        ...     task="tictactoe_v3",
        ...     parallel=False,
        ...     use_mask=True, # Must use it since one player plays at a time
        ...     group_map=None # # Use default for AEC (one group per player)
        ... )
        >>> print(env.group_map)
        ... {'player_1': ['player_1'], 'player_2': ['player_2']}
        >>> env.rollout(10)
    """

    def __init__(
        self,
        task: str,
        parallel: bool,
        return_state: bool = False,
        group_map: MarlGroupMapType | Dict[str, List[str]] | None = None,
        use_mask: bool = False,
        categorical_actions: bool = True,
        seed: int | None = None,
        done_on_any: bool | None = None,
        **kwargs,
    ):
        if not _has_pettingzoo:
            raise ImportError(
                f"pettingzoo python package was not found. Please install this dependency. "
                f"More info: {self.git_url}."
            )
        kwargs["task"] = task
        kwargs["parallel"] = parallel
        kwargs["return_state"] = return_state
        kwargs["group_map"] = group_map
        kwargs["use_mask"] = use_mask
        kwargs["categorical_actions"] = categorical_actions
        kwargs["seed"] = seed
        kwargs["done_on_any"] = done_on_any

        super().__init__(**kwargs)

    def _check_kwargs(self, kwargs: Dict):
        if "task" not in kwargs:
            raise TypeError("Could not find environment key 'task' in kwargs.")
        if "parallel" not in kwargs:
            raise TypeError("Could not find environment key 'parallel' in kwargs.")

    def _build_env(
        self,
        task: str,
        parallel: bool,
        **kwargs,
    ) -> Union[
        "pettingzoo.utils.env.ParallelEnv",  # noqa: F821
        "pettingzoo.utils.env.AECEnv",  # noqa: F821
    ]:
        self.task_name = task

        try:
            from pettingzoo.utils.all_modules import all_environments
        except ModuleNotFoundError as err:
            warnings.warn(
                f"PettingZoo failed to load all modules with error message {err}, trying to load individual modules."
            )
            all_environments = _load_available_envs()

        if task not in all_environments:
            # Try looking at the literal translation of values
            task_module = None
            for value in all_environments.values():
                if value.__name__.split(".")[-1] == task:
                    task_module = value
                    break
            if task_module is None:
                raise RuntimeError(
                    f"Specified task not in available environments {all_environments}"
                )
        else:
            task_module = all_environments[task]

        if parallel:
            petting_zoo_env = task_module.parallel_env(**kwargs)
        else:
            petting_zoo_env = task_module.env(**kwargs)

        return super()._build_env(env=petting_zoo_env)