Multi-Agent Constraints¶

def __getattr__(self, name: str):
    if name == "env":
        raise AttributeError(name)
    return getattr(self.env, name)

def reset(self, seed: int | None = None, options: Dict[str, Any] | None = None):
    obs, info = self.env.reset(seed=seed, options=options)
    self.agents = list(getattr(self.env, "agents", self.possible_agents))
    for a in obs:
        lf = self.label_fn[a] if isinstance(self.label_fn, dict) else self.label_fn
        agent_info = info.setdefault(a, {})
        if agent_info is None:
            agent_info = {}
            info[a] = agent_info
        agent_info["labels"] = set(lf(obs[a]))
    return obs, info

def step(self, actions):
    obs, rewards, term, trunc, infos = self.env.step(actions)
    self.agents = list(getattr(self.env, "agents", self.possible_agents))
    for a in obs:
        lf = self.label_fn[a] if isinstance(self.label_fn, dict) else self.label_fn
        labels = set(lf(obs[a]))
        agent_info = infos.setdefault(a, {})
        if agent_info is None:
            agent_info = {}
            infos[a] = agent_info
        agent_info["labels"] = labels

    return obs, rewards, term, trunc, infos

def observation_space(self, agent):
    return self.env.observation_space(agent)

def action_space(self, agent):
    return self.env.action_space(agent)

def state(self):
    return self.env.state()

def render(self):
    return self.env.render()

def close(self):
    return self.env.close()

def __post_init__(self) -> None:
    if isinstance(self.agents, (str, bytes)):
        raise TypeError("Coalition agents must be a sequence of agent names.")
    agents = tuple(self.agents)
    if not agents:
        raise ValueError("A coalition must contain at least one agent.")
    if any(not isinstance(agent, str) or not agent for agent in agents):
        raise TypeError("Coalition agents must be non-empty strings.")
    if len(set(agents)) != len(agents):
        raise ValueError("Coalition agents must be unique.")
    if self.name is not None and (
        not isinstance(self.name, str) or not self.name
    ):
        raise TypeError("Coalition name must be a non-empty string or None.")
    # Store a canonical order so equality and hashing follow set membership;
    # the optional display name is not part of coalition identity.
    object.__setattr__(self, "agents", tuple(sorted(agents)))

def ordered(self, possible_agents: Sequence[str]) -> tuple[str, ...]:
    """Return members in the environment's canonical agent order."""
    possible = tuple(possible_agents)
    if len(set(possible)) != len(possible):
        raise ValueError("possible_agents must be unique.")
    members = set(self.agents)
    unknown = members - set(possible)
    if unknown:
        raise ValueError(
            f"Coalition contains unknown agents: {sorted(unknown)}."
        )
    return tuple(agent for agent in possible if agent in members)

def __contains__(self, agent: str) -> bool:
    return agent in self.agents

def get_transition_matrix(self) -> np.ndarray | None:
    """Return ``P[next_state, state, joint_action_index]``, when available."""
    return self._transition_matrix

def get_successor_states_dict(
    self,
) -> tuple[
    Mapping[int, Sequence[int]],
    Mapping[tuple[int, int], Sequence[float]],
] | None:
    """Return the sparse state-successor/probability representation.

    Probability vectors are aligned with ``successor_states[state]`` and
    keyed by ``(state, joint_action_index)``.
    """
    if not self.has_successor_states_dict:
        return None
    assert self._successor_states is not None
    assert self._transition_probs is not None
    return self._successor_states, self._transition_probs

def _check_state(self, state: int) -> int:
    if not isinstance(state, Integral) or isinstance(state, bool):
        raise TypeError(f"State ID must be an integer, got {state!r}.")
    state = int(state)
    if not 0 <= state < self.n_states:
        raise ValueError(
            f"State ID {state} is outside [0, {self.n_states})."
        )
    return state

def _joint_action_tuple(
    self, actions: Mapping[str, int] | Sequence[int]
) -> JointAction:
    agents = tuple(self.possible_agents)
    if isinstance(actions, Mapping):
        supplied = set(actions)
        expected = set(agents)
        if supplied != expected:
            raise ValueError(
                "Joint action must contain exactly possible_agents; "
                f"missing={sorted(expected - supplied)}, "
                f"extra={sorted(supplied - expected)}."
            )
        raw = tuple(actions[agent] for agent in agents)
    else:
        if isinstance(actions, (str, bytes)):
            raise TypeError("Joint action must be a mapping or action sequence.")
        raw = tuple(actions)
        if len(raw) != len(agents):
            raise ValueError(
                f"Expected {len(agents)} primitive actions, got {len(raw)}."
            )

    sizes = self.action_sizes
    result: list[int] = []
    for agent, primitive in zip(agents, raw):
        if not isinstance(primitive, Integral) or isinstance(primitive, bool):
            raise TypeError(
                f"Action for {agent!r} must be an integer, got {primitive!r}."
            )
        primitive = int(primitive)
        if not 0 <= primitive < sizes[agent]:
            raise ValueError(
                f"Action {primitive} is outside the action space of {agent!r}."
            )
        result.append(primitive)
    return tuple(result)

def encode_joint_action(
    self, actions: Mapping[str, int] | Sequence[int]
) -> int:
    """Encode a full joint action in canonical Cartesian-product order."""
    action = self._joint_action_tuple(actions)
    index = 0
    for primitive, size in zip(action, self.action_sizes.values()):
        index = index * size + primitive
    return int(index)

def decode_joint_action(self, index: int) -> dict[str, int]:
    """Decode a canonical joint-action index to an agent-action mapping."""
    if not isinstance(index, Integral) or isinstance(index, bool):
        raise TypeError("Joint-action index must be an integer.")
    index = int(index)
    if not 0 <= index < self.n_joint_actions:
        raise ValueError(
            f"Joint-action index {index} is outside [0, {self.n_joint_actions})."
        )

    sizes = tuple(self.action_sizes.values())
    primitives = [0] * len(sizes)
    remainder = index
    for position in range(len(sizes) - 1, -1, -1):
        remainder, primitives[position] = divmod(remainder, sizes[position])
    return dict(zip(self.possible_agents, primitives))

def legal_actions(self, state: int, agent: str) -> tuple[int, ...]:
    """Legal primitive actions for an agent in a model state.

    Override for state-dependent action availability. The default permits the
    entire action space.
    """
    self._check_state(state)
    try:
        size = self.action_sizes[agent]
    except KeyError as exc:
        raise ValueError(f"Unknown agent: {agent!r}.") from exc
    return tuple(range(size))

def get_legal_actions(self, state: int, agent: str) -> tuple[int, ...]:
    """Return a validated, sorted copy of :meth:`legal_actions`."""
    state = self._check_state(state)
    try:
        size = self.action_sizes[agent]
    except KeyError as exc:
        raise ValueError(f"Unknown agent: {agent!r}.") from exc
    raw = tuple(self.legal_actions(state, agent))
    if not raw:
        raise ValueError(
            f"Agent {agent!r} has no legal action in model state {state}."
        )
    if any(
        not isinstance(action, Integral)
        or isinstance(action, bool)
        or not 0 <= int(action) < size
        for action in raw
    ):
        raise ValueError(
            f"Invalid legal actions for {agent!r} in state {state}: {raw!r}."
        )
    return tuple(sorted(set(map(int, raw))))

def successors(
    self,
    state: int,
    actions: Mapping[str, int] | Sequence[int],
) -> tuple[int, ...]:
    """Return every non-zero-probability successor of a legal joint action."""
    state = self._check_state(state)
    action = self._joint_action_tuple(actions)
    for agent, primitive in zip(self.possible_agents, action):
        if primitive not in self.get_legal_actions(state, agent):
            raise ValueError(
                f"Action {primitive} is not legal for {agent!r} in state {state}."
            )
    action_index = self.encode_joint_action(action)

    if self.has_successor_states_dict:
        sparse = self.get_successor_states_dict()
        assert sparse is not None
        successor_states, transition_probs = sparse
        ids = np.asarray(successor_states.get(state, ()))
        probs = np.asarray(
            transition_probs.get((state, action_index), ()),
            dtype=np.float64,
        )
    elif self.has_transition_matrix:
        matrix = np.asarray(self.get_transition_matrix())
        expected = (self.n_states, self.n_states, self.n_joint_actions)
        if matrix.shape != expected:
            raise ValueError(
                "Expected transition shape "
                f"(next_state, state, joint_action)={expected}, got {matrix.shape}."
            )
        ids = np.arange(self.n_states, dtype=np.intp)
        probs = np.asarray(matrix[:, state, action_index], dtype=np.float64)
    else:
        raise ValueError(
            "TabularParallelEnv must expose a transition matrix or sparse "
            "successor-state dictionaries."
        )

    if ids.ndim != 1 or (ids.size and ids.dtype.kind not in "iu"):
        raise ValueError(f"Successors of state {state} must be integer IDs.")
    ids = ids.astype(np.intp, copy=False)
    if np.any((ids < 0) | (ids >= self.n_states)):
        raise ValueError(f"Successor of state {state} is out of range.")
    if probs.shape != (ids.size,):
        raise ValueError(
            f"Probability vector for state {state}, joint action {action} "
            "has the wrong shape."
        )
    if not np.all(np.isfinite(probs)) or np.any(probs < 0):
        raise ValueError(
            f"Invalid probabilities at state {state}, joint action {action}."
        )
    if not np.isclose(probs.sum(), 1.0, rtol=1e-6, atol=1e-8):
        raise ValueError(
            f"Transition probabilities for state {state}, joint action "
            f"{action} must sum to 1."
        )

    # No probability threshold: every positive-probability outcome matters.
    result = tuple(sorted(set(map(int, ids[probs > 0]))))
    if not result:
        raise ValueError(
            f"Missing transition support for state {state}, joint action {action}."
        )
    return result

def get_state_id(self) -> int:
    """Return the current finite model-state ID.

    Subclasses with a non-integer runtime representation may override this,
    but should still return the exact state indexing the transition model.
    """
    if self._state is None:
        raise RuntimeError("The environment has no active tabular state.")
    return self._check_state(self._state)

def observations_from_state(self, state: int) -> dict[str, Any]:
    """Return each possible agent's observation for a finite model state.

    The default supports fully observed environments where every agent's
    observation is the same zero-based ``Discrete(n_states)`` state ID.
    Environments with structured or local observations should override it.
    """
    state = self._check_state(state)
    observations: dict[str, Any] = {}
    for agent in self.possible_agents:
        space = self.observation_space(agent)
        if (
            not isinstance(space, spaces.Discrete)
            or int(space.start) != 0
            or int(space.n) != self.n_states
        ):
            raise NotImplementedError(
                f"{type(self).__name__}.observations_from_state() must be "
                "implemented for structured or non-state observations."
            )
        observations[agent] = state
    return observations