diff --git a/docs/source/reference/envs_multiagent.rst b/docs/source/reference/envs_multiagent.rst index 13f0a7cb9ca..fc63420bfbd 100644 --- a/docs/source/reference/envs_multiagent.rst +++ b/docs/source/reference/envs_multiagent.rst @@ -20,11 +20,16 @@ Some of the main differences between these paradigms include: - **done** (and ``"truncated"`` or ``"terminated"``) can be per-agent or shared. TorchRL accommodates all these possible paradigms thanks to its :class:`tensordict.TensorDict` data carrier. -In particular, in multi-agent environments, per-agent keys will be carried in a nested "agents" TensorDict. -This TensorDict will have the additional agent dimension and thus group data that is different for each agent. -The shared keys, on the other hand, will be kept in the first level, as in single-agent cases. +Per-agent keys live in nested **group** tensordicts. Each group has an extra agent dimension +so that data that differs across agents can be stacked. Shared keys stay at the root, +as in single-agent cases. -Let's look at an example to understand this better. For this example we are going to use +The simplest layout uses a single group named ``"agents"`` (the VMAS default shown +below). When agents belong to several groups -- for example two competing teams -- +each group is its own nested tensordict. See :ref:`MARL-multiple-groups` for the +general contract a native :class:`~torchrl.envs.EnvBase` ``_step()`` must implement. + +Let's look at the single-group case first. For this example we are going to use `VMAS `_, a multi-robot task simulator also based on PyTorch, which runs parallel batched simulation on device. @@ -67,7 +72,8 @@ and **info** are present in the nested "agents" tensordict with batch size `(num which represents the additional agent dimension. Multi-agent tensor specs will follow the same style as in tensordicts. -Specs relating to values that vary between agents will need to be nested in the "agents" entry. +Specs relating to values that vary between agents will need to be nested in the +group entry (here, ``"agents"``). Here is an example of how specs can be created in a multi-agent environment where only the done flag is shared across agents (as in VMAS): @@ -114,11 +120,14 @@ As you can see, it is very simple! Per-agent keys will have the nested composite single agent standards. .. note:: - Since reward, done and action keys may have the additional "agent" prefix (e.g., `("agents","action")`), - the default keys used in the arguments of other TorchRL components (e.g. "action") will not match exactly. - Therefore, TorchRL provides the `env.action_key`, `env.reward_key`, and `env.done_key` attributes, + Since reward, done and action keys may have the additional group prefix (e.g., ``("agents", "action")``), + the default keys used in the arguments of other TorchRL components (e.g. ``"action"``) will not match exactly. + Therefore, TorchRL provides the ``env.action_key``, ``env.reward_key``, and ``env.done_key`` attributes, which will automatically point to the right key to use. Make sure you pass these attributes to the various - components in TorchRL to inform them of the right key (e.g., the `loss.set_keys()` function). + components in TorchRL to inform them of the right key (e.g., the ``loss.set_keys()`` function). + When there is more than one action, reward or done key (as with multiple groups), + use the plural ``env.action_keys``, ``env.reward_keys`` and ``env.done_keys`` instead -- + the singular attributes raise ``KeyError``. .. note:: TorchRL abstracts these nested specs away for ease of use. @@ -130,6 +139,367 @@ single agent standards. Note that `env.reward_spec == env.output_spec["full_reward_spec"][env.reward_key]`. +.. _MARL-multiple-groups: + +Multiple agent groups +--------------------- + +``"agents"`` is just a group name, not a required key. A native +:class:`~torchrl.envs.EnvBase` can expose any number of groups -- ``"red"`` / +``"blue"``, or the ``"agents"`` / ``"adversaries"`` teams of *simple_tag* -- +as long as each group is a nested tensordict whose last batch dimension +indexes the agents in that group. + +Agents that share a policy (and typically a spec) belong in the same group so +their tensors can be stacked. Heterogeneous or competing teams go in separate +groups and are processed by separate modules. The grouping is the +``group_map`` dict ``{group_name: [agent_name, ...]}``; see +:class:`~torchrl.envs.MarlGroupMapType` and +:func:`~torchrl.envs.check_marl_grouping`. + +The canonical trained example of this layout is the +:doc:`competitive MADDPG tutorial <../tutorials/multiagent_competitive_ddpg>` +(in particular the *Rollout* section), which consumes the ``"agents"`` / +``"adversaries"`` groups of *simple_tag*. + +What ``_step()`` must return +~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +:meth:`~torchrl.envs.EnvBase.step` is a thin wrapper: it calls the private +``_step()``, then stores that result under the ``"next"`` key of the input +tensordict. Consequently a custom ``_step()`` must: + +- **read** actions from the *input* tensordict at ``(group, "action")``; +- **write** next observations, rewards and done flags into a **new** + tensordict (out-of-place); +- **not** wrap that tensordict in ``"next"`` -- the public + :meth:`~torchrl.envs.EnvBase.step` does that; +- **not** write actions onto the output. + +``_reset()`` uses the same key tree minus rewards (there is no reward at +reset time). After a public ``step()`` the input (with actions) stays at the +root and the ``_step()`` output sits under ``"next"``. + +For two groups ``"red"`` (2 agents) and ``"blue"`` (3 agents) and an +environment batch ``B``, the tensordict ``_step()`` returns looks like this: + +.. code-block:: + :caption: Multi-group key tree returned by a native ``_step()`` + + TensorDict( + fields={ + red: TensorDict( + fields={ + observation: Tensor(shape=torch.Size([*B, 2, obs_red])), + reward: Tensor(shape=torch.Size([*B, 2, 1]))}, + batch_size=torch.Size([*B, 2])), + blue: TensorDict( + fields={ + observation: Tensor(shape=torch.Size([*B, 3, obs_blue])), + reward: Tensor(shape=torch.Size([*B, 3, 1]))}, + batch_size=torch.Size([*B, 3])), + done: Tensor(shape=torch.Size([*B, 1])), + terminated: Tensor(shape=torch.Size([*B, 1])), + truncated: Tensor(shape=torch.Size([*B, 1]))}, + batch_size=torch.Size([*B,])) + +The corresponding public ``step()`` / ``rand_step()`` output is the input +tensordict (root ``red`` / ``blue`` entries hold ``action`` only) plus a +``next`` tensordict that is exactly the tree above. + +Where each field lives +~~~~~~~~~~~~~~~~~~~~~~ + +- **Action.** Input of ``_step()`` only, at ``(group, "action")``, shape + ``(*batch, n_agents_in_group, *action_shape)``. Never written by + ``_step()``. +- **Observation.** Output of ``_step()`` and ``_reset()``, typically at + ``(group, "observation")`` with shape + ``(*batch, n_agents_in_group, *obs_shape)``. A shared / global observation + (for example a global ``"state"``) is a root-level key with no extra agent + dimension, exactly as in the single-agent case. +- **Reward.** Output of ``_step()`` only (not ``_reset()``). Per-agent or + per-group rewards live at ``(group, "reward")`` with shape + ``(*batch, n_agents_in_group, 1)``. A fully shared reward is a root + ``"reward"`` with shape ``(*batch, 1)``. +- **Done / terminated / truncated** can sit at three levels, which can be + combined: + + * **Shared (root).** ``"done"``, ``"terminated"`` and (if used) + ``"truncated"`` at the root, shape ``(*batch, 1)``. This is the + convention used by the example below and by wrappers such as VMAS + (which writes only this). It is not required: + :class:`~torchrl.envs.EnvBase` also accepts environments whose done / + terminated / truncated / ``_reset`` keys live only in nested + tensordicts. Write a root flag when the episode ends for the whole + environment at once; do not add one just to create a second + termination hierarchy. + * **Per-group / per-agent, stacked.** ``(group, "done")`` (and the + ``terminated`` / ``truncated`` siblings) with shape + ``(*batch, n_agents_in_group, 1)``. PettingZoo writes these *and* + aggregates them into the root flag (``any`` or ``all``, controlled by + ``done_on_any``). + * **One group per agent.** + :attr:`~torchrl.envs.MarlGroupMapType.ONE_GROUP_PER_AGENT` changes + group membership (each agent is its own group), not the group + TensorDict layout. Each group is still a nested tensordict with a + stacked agent dimension, so a one-agent group has Composite shape + ``(*batch, 1)`` and ``(agent_name, "done")`` has shape + ``(*batch, 1, 1)``. + +Specs must mirror the nesting +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Every leaf that ``_step()`` (or ``_reset()``) writes must have a spec at the +same nested key, with the same leading shape. The group-level +:class:`~torchrl.data.Composite` carries shape ``(*batch, n_agents_in_group)``; +root-level (shared) specs carry shape ``(*batch,)`` or ``(*batch, 1)``. + +.. code-block:: + :caption: Specs for a two-group env (shared done, per-group obs/reward/action) + + >>> from torchrl.data import Bounded, Categorical, Composite, Unbounded + >>> n_red, n_blue = 2, 3 + >>> obs_dim, act_dim = 8, 2 + >>> bs = env.batch_size + >>> env.action_spec = Composite( + ... { + ... "red": Composite( + ... {"action": Bounded(-1, 1, shape=(*bs, n_red, act_dim))}, + ... shape=(*bs, n_red), + ... ), + ... "blue": Composite( + ... {"action": Bounded(-1, 1, shape=(*bs, n_blue, act_dim))}, + ... shape=(*bs, n_blue), + ... ), + ... }, + ... shape=bs, + ... ) + >>> env.observation_spec = Composite( + ... { + ... "red": Composite( + ... {"observation": Unbounded(shape=(*bs, n_red, obs_dim))}, + ... shape=(*bs, n_red), + ... ), + ... "blue": Composite( + ... {"observation": Unbounded(shape=(*bs, n_blue, obs_dim))}, + ... shape=(*bs, n_blue), + ... ), + ... }, + ... shape=bs, + ... ) + >>> env.reward_spec = Composite( + ... { + ... "red": Composite( + ... {"reward": Unbounded(shape=(*bs, n_red, 1))}, + ... shape=(*bs, n_red), + ... ), + ... "blue": Composite( + ... {"reward": Unbounded(shape=(*bs, n_blue, 1))}, + ... shape=(*bs, n_blue), + ... ), + ... }, + ... shape=bs, + ... ) + >>> env.done_spec = Composite( + ... { + ... "done": Categorical(n=2, shape=(*bs, 1), dtype=torch.bool), + ... "terminated": Categorical(n=2, shape=(*bs, 1), dtype=torch.bool), + ... "truncated": Categorical(n=2, shape=(*bs, 1), dtype=torch.bool), + ... }, + ... shape=bs, + ... ) + +To also expose per-agent done flags, nest them in each group the same way +as ``reward``. If the environment also uses a shared root done (as in the +example below), keep those root flags as the environment-wide reset signal: + +.. code-block:: + + >>> env.done_spec["red"] = Composite( + ... { + ... "done": Categorical(n=2, shape=(*bs, n_red, 1), dtype=torch.bool), + ... "terminated": Categorical(n=2, shape=(*bs, n_red, 1), dtype=torch.bool), + ... "truncated": Categorical(n=2, shape=(*bs, n_red, 1), dtype=torch.bool), + ... }, + ... shape=(*bs, n_red), + ... ) + +Call :func:`~torchrl.envs.check_env_specs` after construction: it runs a +short rollout and checks that every key ``_step()`` / ``_reset()`` writes +matches the spec tree. + +Native ``_step()`` sketch +~~~~~~~~~~~~~~~~~~~~~~~~~ + +The following is a copy-paste starting point for a two-team +:class:`~torchrl.envs.EnvBase`. It is not a PettingZoo or VMAS wrapper: +actions are read from the group tensordicts and the returned tensordict +follows the tree above. ``_initial_obs`` / ``_apply_dynamics`` are the +only placeholders -- replace them with a real reset distribution and +transition. In ``_step``, take devices from the action tensors. + +.. code-block:: python + :caption: Native multi-group ``EnvBase._step()`` + + import torch + from tensordict import TensorDict, TensorDictBase + from torchrl.data import Bounded, Categorical, Composite, Unbounded + from torchrl.envs import EnvBase, check_env_specs + + + class TwoTeamEnv(EnvBase): + """Minimal two-group env. Replace the two helpers with real physics.""" + + def __init__(self, n_red=2, n_blue=3, obs_dim=8, act_dim=2, **kwargs): + super().__init__(**kwargs) + self.n_red = n_red + self.n_blue = n_blue + self.obs_dim = obs_dim + self.act_dim = act_dim + self.group_map = { + "red": [f"red_{i}" for i in range(n_red)], + "blue": [f"blue_{i}" for i in range(n_blue)], + } + bs = self.batch_size + self.action_spec = Composite( + { + "red": Composite( + {"action": Bounded(-1, 1, shape=(*bs, n_red, act_dim))}, + shape=(*bs, n_red), + ), + "blue": Composite( + {"action": Bounded(-1, 1, shape=(*bs, n_blue, act_dim))}, + shape=(*bs, n_blue), + ), + }, + shape=bs, + ) + self.observation_spec = Composite( + { + "red": Composite( + {"observation": Unbounded(shape=(*bs, n_red, obs_dim))}, + shape=(*bs, n_red), + ), + "blue": Composite( + {"observation": Unbounded(shape=(*bs, n_blue, obs_dim))}, + shape=(*bs, n_blue), + ), + }, + shape=bs, + ) + self.reward_spec = Composite( + { + "red": Composite( + {"reward": Unbounded(shape=(*bs, n_red, 1))}, + shape=(*bs, n_red), + ), + "blue": Composite( + {"reward": Unbounded(shape=(*bs, n_blue, 1))}, + shape=(*bs, n_blue), + ), + }, + shape=bs, + ) + self.done_spec = Composite( + { + "done": Categorical(n=2, shape=(*bs, 1), dtype=torch.bool), + "terminated": Categorical(n=2, shape=(*bs, 1), dtype=torch.bool), + "truncated": Categorical(n=2, shape=(*bs, 1), dtype=torch.bool), + }, + shape=bs, + ) + + def _set_seed(self, seed): + if seed is not None: + torch.manual_seed(seed) + + def _initial_obs(self): + # Replace with a real reset distribution. + red_obs = torch.zeros( + *self.batch_size, self.n_red, self.obs_dim, device=self.device + ) + blue_obs = torch.zeros( + *self.batch_size, self.n_blue, self.obs_dim, device=self.device + ) + return red_obs, blue_obs + + def _apply_dynamics(self, red_action, blue_action): + # Replace with a real transition. Devices come from the actions. + red_obs = red_action.new_zeros(*self.batch_size, self.n_red, self.obs_dim) + blue_obs = blue_action.new_zeros(*self.batch_size, self.n_blue, self.obs_dim) + red_rew = red_action.new_zeros(*self.batch_size, self.n_red, 1) + blue_rew = blue_action.new_zeros(*self.batch_size, self.n_blue, 1) + done = red_action.new_zeros(*self.batch_size, 1, dtype=torch.bool) + return red_obs, blue_obs, red_rew, blue_rew, done + + def _reset(self, tensordict): + red_obs, blue_obs = self._initial_obs() + done = red_obs.new_zeros(*self.batch_size, 1, dtype=torch.bool) + return TensorDict( + { + "red": TensorDict( + {"observation": red_obs}, + batch_size=(*self.batch_size, self.n_red), + ), + "blue": TensorDict( + {"observation": blue_obs}, + batch_size=(*self.batch_size, self.n_blue), + ), + "done": done, + "terminated": done.clone(), + "truncated": done.clone(), + }, + batch_size=self.batch_size, + ) + + def _step(self, tensordict: TensorDictBase) -> TensorDictBase: + # Actions live on the *input* tensordict, never on the output. + red_action = tensordict["red", "action"] # (*batch, n_red, act_dim) + blue_action = tensordict["blue", "action"] # (*batch, n_blue, act_dim) + + red_obs, blue_obs, red_rew, blue_rew, done = self._apply_dynamics( + red_action, blue_action + ) + # Optional per-agent flags go on this returned tensordict, + # e.g. out["red", "done"] with shape (*batch, n_red, 1). + + return TensorDict( + { + "red": TensorDict( + {"observation": red_obs, "reward": red_rew}, + batch_size=(*self.batch_size, self.n_red), + ), + "blue": TensorDict( + {"observation": blue_obs, "reward": blue_rew}, + batch_size=(*self.batch_size, self.n_blue), + ), + "done": done, + "terminated": done.clone(), + "truncated": done.new_zeros(done.shape, dtype=torch.bool), + }, + batch_size=self.batch_size, + ) + + + env = TwoTeamEnv() + check_env_specs(env) + # env.action_keys == [("blue", "action"), ("red", "action")] + # env.reward_keys == [("blue", "reward"), ("red", "reward")] + # env.done_keys == ["done", "terminated", "truncated"] + +Collectors and replay buffers transport these nested keys automatically; they +do not take ``env.action_keys``, ``env.reward_keys`` or ``env.done_keys`` as +configuration arguments. Compose the group policies so that each reads and +writes its group's nested keys, then use one loss per group. Configure each loss +with the individual :class:`~tensordict.NestedKey` values accepted by its +``set_keys()`` method. For example, a loss for the red group in the environment +above can use ``reward=("red", "reward")`` together with the root +``done="done"`` and ``terminated="terminated"`` keys. The +:doc:`competitive MADDPG tutorial <../tutorials/multiagent_competitive_ddpg>` +shows a full two-group training loop. + + .. autosummary:: :toctree: generated/ :template: rl_template_fun.rst