[Doc] Document multi-agent _step TensorDict layout

[YeonwooSung]

Description

Extends docs/source/reference/envs_multiagent.rst with a Multiple agent groups section so a native EnvBase._step() is documented beyond the single "agents" group used by VMAS.

The new section:

• Shows the key tree a custom _step() must write for several groups ("red" / "blue", or "agents" / "adversaries"): per-group nested tensordicts with a stacked agent dimension.
• States where action (input only), observation, and reward live, and where done / terminated / truncated live when they are shared (root reset signal), per-group / per-agent (stacked), or one-group-per-agent.
• Requires specs to mirror that nesting, with a two-group spec example and an optional per-group done spec.
• Includes a copy-paste native TwoTeamEnv(EnvBase) sketch (not a PettingZoo wrapper) that reads (group, "action") and returns observations, rewards, and shared done flags without a "next" wrapper.
• Cross-links the competitive MADDPG tutorial as the canonical multi-group training loop.

The existing VMAS single-group example is kept as the simple case. The intro now treats "agents" as one group name, not a required key.

Code example

Example return value from a custom _step: two groups with different agent counts and shared environment termination. EnvBase.step adds the outer "next" key.

import torch
from tensordict import TensorDict

batch_size = 4
out = TensorDict({
    group: TensorDict({
        "observation": torch.zeros(batch_size, n_agents, 8),
        "reward": torch.zeros(batch_size, n_agents, 1),
    }, batch_size=[batch_size, n_agents])
    for group, n_agents in {"red": 2, "blue": 3}.items()
}, batch_size=[batch_size])
out["done"] = torch.zeros(batch_size, 1, dtype=torch.bool)
out["terminated"] = torch.zeros_like(out["done"])
assert out["red", "reward"].shape == (4, 2, 1)
assert out["blue", "observation"].shape == (4, 3, 8)
# return out from _step; group actions are read from the input TensorDict.

Motivation and Context

close #2425

The current multi-agent env page only documented the single "agents" group. Matteo Bettini agreed it needed a more general multi-group description; a later comment pointed at the competitive DDPG tutorial rollout as the pattern.

• I have raised an issue to propose this change (required for new features and bug fixes)

Types of changes

• Documentation (update in the documentation)

Checklist

• I have read the CONTRIBUTION guide (required)
• My change requires a change to the documentation.
• I have updated the tests accordingly (required for a bug fix or a new feature).
• I have updated the documentation accordingly.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/rl/4232

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 New Failure, 15 Unclassified Failures

As of commit 7ac9875 with merge base 1d3de3d ():

NEW FAILURE - The following job has failed:

• Lint (gh)

UNCLASSIFIED FAILURES - DrCI could not classify the following jobs because the workflow did not run on the merge base. The failures may be pre-existing on trunk or introduced by this PR:

• Build Aarch64 Linux Wheels (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
• Build Linux Wheels (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
• Build M1 Wheels (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
• Build Windows Wheels (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
• Continuous Benchmark (PR) (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
• Examples Tests on Linux (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
• Generate documentation (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
• Habitat Tests on Linux (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
• Libs Tests on Linux (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
• LLM Tests on Linux (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
• Push Binary Nightly (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
• SOTA Tests on Linux (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
• Tutorials Tests on Linux (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
• Unit-tests on Linux (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
• Validate Test Partitioning (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[vmoens]

The TwoTeamEnv example runs and passes check_env_specs. Two statements in the new layout contract do not match the supported environment/spec behavior.

[vmoens]

[P2] Do not require a root done flag for every native env

EnvBase also supports environments whose done/terminated/reset keys are entirely nested; NestedCountingEnv(nest_done=True) and the existing nested StepCounter tests exercise that contract. Root shared flags are the convention for this example and some wrappers, not a universal native-env requirement. Qualify this paragraph so implementers are not told to add a second termination hierarchy unnecessarily.

[YeonwooSung]

Addressed in 7ac9875. The paragraph now states that root done/terminated/truncated is the convention used by the TwoTeamEnv sketch and wrappers such as VMAS, not a native-env requirement. EnvBase already resets from nested done/_reset keys (NestedCountingEnv(nest_done=True, has_root_done=False) and the nested StepCounter tests), so the text no longer tells implementers to add a second termination hierarchy.

[vmoens]

[P2] Keep the singleton agent dimension for one-agent groups

ONE_GROUP_PER_AGENT changes group membership, not the group TensorDict layout. PettingZoo constructs each group Composite with shape [n_agents] and each done leaf with [n_agents, 1], so a one-agent group still has trailing shape [1, 1] (or [*batch, 1, 1] with an env batch). Saying there is no agent dimension will make readers build incompatible specs/policies. Retain that singleton dimension in the documented shape.

[YeonwooSung]

Addressed in 7ac9875. ONE_GROUP_PER_AGENT is documented as a grouping change only. PettingZoo still builds each group Composite with shape [n_agents] and each done leaf with [n_agents, 1], so a one-agent group keeps trailing shape [*batch, 1, 1]. The previous "no extra agent dimension" wording is gone.

[YeonwooSung]

Updated the layout contract in 7ac9875 for the two CHANGES_REQUESTED items:

• Root done/terminated/truncated is documented as a convention (TwoTeamEnv, VMAS), not a native-env requirement. EnvBase already supports fully nested done/_reset keys.
• ONE_GROUP_PER_AGENT keeps the singleton agent dimension: Composite shape (*batch, 1), done leaf (*batch, 1, 1).

The TwoTeamEnv sketch is unchanged and still uses shared root done flags.

[vmoens]

Re-reviewed 7ac9875. The text now allows nested-only done specs instead of requiring a root done universally, and the one-group-per-agent example retains the required singleton done dimension.

I reran the TwoTeamEnv example and check_env_specs successfully. The revised examples and shape descriptions are consistent. No remaining actionables.