Section 10.7: PettingZoo for multi-agent

"The moment a second agent enters the loop, the environment is no longer a function of your actions alone."

A Multi-Agent Environment Designer

What This Section Builds

PettingZoo becomes operational when a multi-agent task names its agents, timing, and per-agent returns. Gymnasium assumes one learning agent acts at each step, following the same agent-environment interface formalized for single-agent settings. PettingZoo adds agent identities, per-agent observations, per-agent rewards, and per-agent episode endings.

The goal is to choose the right multi-agent API before writing the task. Use AEC when turn order matters. Use the Parallel API when agents act simultaneously and the environment resolves their joint action at the end of a cycle.

By the end of this section you will be able to do three concrete things: implement a minimal PettingZoo Parallel environment with per-agent dictionaries (Code Fragment 10.7.1), drive a turn-based environment through the official AEC loop with agent_iter() and last() (Code Fragment 10.7.2), and decide which API a new multi-agent task needs before writing any code, using the timing test in the Practical Recipe below.

Theory

PettingZoo's AEC API models an Agent Environment Cycle: the environment selects one active agent, the agent observes, acts, and the environment advances to the next agent. That is one Gymnasium contract per agent, sequenced by a scheduler (a piece of environment logic that decides which agent's turn is next), and it fits turn-based or order-sensitive settings, such as a robot handing off an object after another robot clears space.

The Parallel API steps all live agents together through dictionaries keyed by agent id. It fits simultaneous control, such as two mobile robots moving in the same timestep. Both APIs keep terminations, truncations, rewards, observations, and infos per agent, because different agents may leave the task for different reasons. Figure 10.7B compares the two timing models side by side.

Worked Example

With the scheduler distinction from the theory in hand, the fastest way to make the per-agent dictionary contract concrete is to build the smallest environment that exercises it.

Code Fragment 10.7.1 builds a tiny Parallel API environment with two agents on a one-dimensional line. It is intentionally small so the return dictionaries are easy to inspect.

# Implement the minimal shape of a PettingZoo ParallelEnv.
# Observations, rewards, terminations, truncations, and infos are keyed by agent.
from gymnasium import spaces
from pettingzoo.utils.env import ParallelEnv

class TwoRobotLine(ParallelEnv):
    metadata = {"name": "two_robot_line_v0"}

    def __init__(self):
        self.possible_agents = ["picker", "carrier"]
        self.observation_spaces = {agent: spaces.Box(0, 4, shape=(1,), dtype=int) for agent in self.possible_agents}
        self.action_spaces = {agent: spaces.Discrete(3) for agent in self.possible_agents}

    def reset(self, seed=None, options=None):
        self.agents = self.possible_agents[:]
        self.positions = {"picker": 0, "carrier": 4}
        observations = {agent: [pos] for agent, pos in self.positions.items()}
        infos = {agent: {} for agent in self.agents}
        return observations, infos

    def step(self, actions):
        moves = {0: -1, 1: 0, 2: 1}
        for agent, action in actions.items():
            self.positions[agent] = min(4, max(0, self.positions[agent] + moves[int(action)]))
        observations = {agent: [pos] for agent, pos in self.positions.items()}
        rewards = {agent: float(self.positions["picker"] == self.positions["carrier"]) for agent in self.agents}
        terminations = {agent: rewards[agent] == 1.0 for agent in self.agents}
        truncations = {agent: False for agent in self.agents}
        infos = {agent: {"position": self.positions[agent]} for agent in self.agents}
        if any(terminations.values()):
            self.agents = []
        return observations, rewards, terminations, truncations, infos

env = TwoRobotLine()
observations, infos = env.reset(seed=42)
print(observations)
print(env.step({"picker": 2, "carrier": 0})[1:4])

The expected output exposes two independent agent views at reset, then three dictionaries keyed by agent id for rewards, terminations, and truncations after one joint step. Readers should interpret the all-false ending dictionaries as evidence that neither robot has yet reached the meeting condition.

Practical Recipe

1. Choose AEC when turn order affects state, legality, or reward assignment.
2. Choose Parallel API when all live agents submit actions for the same environment tick.
3. Define possible_agents, per-agent spaces, and live agents explicitly.
4. Return observations, rewards, terminations, truncations, and infos keyed by agent id.
5. Log agent-specific failure labels, because one agent can truncate or terminate before another.

Before reading this algorithm, ask yourself: if agent B's reward depends on what agent A did in the same tick, which API do you reach for, and what breaks silently if you pick the wrong one?

That under-specification check points at the deeper reason the per-agent contract exists in the first place.

Why One Action Vector Is Not Enough

PettingZoo matters because multi-agent environments are not only bigger Gymnasium environments. A single-agent Gymnasium wrapper applied to a two-robot task collapses both agents into one action vector. Credit for a successful handoff cannot be attributed to either robot individually, and a policy that works for robot A silently breaks robot B. This attribution gap has a concrete cost: without per-agent signal, isolating which robot is at fault reduces to a variance-based search over the single combined reward, which in practice can take orders of magnitude more rollouts than reading the fault directly off a per-agent reward trace, though the exact ratio depends on the task and noise level. With PettingZoo's per-agent dictionaries, the same rollout yields two reward traces, two termination signals, and two independent diagnostics. The unit of action can be one active agent or a dictionary of simultaneous agent actions. The unit of reward can be individual, shared, or both. The unit of termination can differ by agent.

The graduate-level habit is to state the game form (the full specification of who acts when, who observes what, and who is rewarded for what, before any code is written) before coding: agents, observation timing, action timing, reward ownership, termination ownership, and whether roles are symmetric. Once those are explicit, the PettingZoo API choice becomes a modeling decision rather than a software afterthought. A multi-agent policy trained without per-agent credit assignment is not a coordination strategy; it is a hope dressed up as a dictionary.

The possible_agents versus agents split matters in physical deployments because a real robot can fail mid-episode: a battery dies, a gripper jams, an e-stop (an emergency-stop switch that immediately halts a robot) fires. The environment must keep producing valid returns for the survivors without crashing the training loop, and a single fixed roster would corrupt every downstream dictionary lookup on the first early exit. So possible_agents is set once at construction and fixes the action and observation space definitions, while agents is a mutable copy that shrinks whenever a termination or truncation flag is set. Trainers iterate over this live list, skip a dead agent cleanly, and never query a key that no longer exists.

A robust implementation starts with the interaction schedule. After that, write the per-agent spaces and the step return structure. Only then should you connect policies or trainers.

1. List agents and roles before writing reward code.
2. Choose AEC or Parallel API from the action timing, not from trainer convenience.
3. Declare per-agent observation and action spaces.
4. Return per-agent dictionaries for observations, rewards, terminations, truncations, and infos.
5. Run PettingZoo API tests before trusting a custom environment.

# Sketch the official AEC loop shape for turn-based environments.
# The active agent receives last(), then submits one action with step().
def run_aec_policy(env, policy):
    env.reset(seed=42)
    for agent in env.agent_iter():
        observation, reward, termination, truncation, info = env.last()
        if termination or truncation:
            action = None
        else:
            action = policy(agent, observation, info)
        env.step(action)
    env.close()

print("AEC loop handles one active agent at a time.")

The expected output is intentionally verbal rather than numeric. It marks the core AEC interpretation: each iteration belongs to exactly one currently active agent, so per-agent observation, reward, and ending logic must be read turn by turn rather than as one simultaneous batch.

Once both loop shapes are in hand, the payoff of the per-agent contract shows up most clearly during debugging, when a rollout goes wrong.

When a two-robot rollout fails, clear the environment contract before you blame the policy. Assign the failure to one of the multi-agent-specific layers first. Check for an agent-id mismatch between the action dictionary and the live agents roster. Check for an AEC-versus-Parallel timing error where the carrier conditioned on a picker action that had not yet resolved. Check for a per-agent reward leak where a collision penalty meant for the picker arm landed on the carrier. Check for a stale state estimate shared across agents, or a roster-pruning bug that crashed the loop when one robot hit an e-stop. Then rerun one controlled perturbation to isolate the suspected layer, for example forcing a single agent to no-op for an entire episode. PettingZoo's per-agent dictionaries make this triage possible because each agent's reward and termination trace is recoverable separately rather than averaged into a team scalar.