from __future__ import annotations
import warnings
from collections.abc import Iterable, Iterator
from typing import Any, Generic, TypeVar
import gymnasium.spaces
import numpy as np
from typing_extensions import override
ObsType = TypeVar("ObsType")
ActionType = TypeVar("ActionType")
AgentID = TypeVar("AgentID")
# deprecated
ObsDict = dict[AgentID, ObsType]
# deprecated
ActionDict = dict[AgentID, ActionType]
"""
Base environment definitions
See docs/api.md for api documentation
See docs/dev_docs.md for additional documentation and an example environment.
"""
[docs]
class AECEnv(Generic[AgentID, ObsType, ActionType]):
"""The AECEnv steps agents one at a time.
If you are unsure if you have implemented a AECEnv correctly, try running
the `api_test` documented in the Developer documentation on the website.
"""
metadata: dict[str, Any] # Metadata for the environment
# All agents that may appear in the environment
possible_agents: list[AgentID]
# Agents active at any given time. The base class never modifies this list;
# keeping it up to date is the environment's job. Environments are expected
# to follow the convention that a terminated or truncated agent stays listed
# until it has taken its final vacuous step(None), which is what
# `_was_dead_step` performs. `api_test` checks both halves of that
# convention, and for an environment that follows it `not env.agents` means
# the episode is over.
agents: list[AgentID]
observation_spaces: dict[
AgentID, gymnasium.spaces.Space[ObsType]
] # Observation space for each agent
# Action space for each agent
action_spaces: dict[AgentID, gymnasium.spaces.Space[ActionType]]
# Whether each agent has just reached a terminal state
terminations: dict[AgentID, bool]
truncations: dict[AgentID, bool]
rewards: dict[AgentID, float] # Reward from the last step for each agent
# Cumulative rewards for each agent
_cumulative_rewards: dict[AgentID, float]
infos: dict[
AgentID, dict[str, Any]
] # Additional information from the last step for each agent
agent_selection: AgentID # The agent currently being stepped
def __init__(self):
pass
[docs]
def step(self, action: ActionType) -> None:
"""Accepts and executes the action of the current agent_selection in the environment.
Automatically switches control to the next agent.
"""
raise NotImplementedError
[docs]
def reset(
self,
seed: int | None = None,
options: dict[str, Any] | None = None,
) -> None:
"""Resets the environment to a starting state."""
raise NotImplementedError
# TODO: Remove `Optional` type below
[docs]
def observe(self, agent: AgentID) -> ObsType | None:
"""Returns the observation an agent currently can make.
`last()` calls this function.
"""
raise NotImplementedError
[docs]
def render(self) -> None | np.ndarray | str | list[Any]:
"""Renders the environment as specified by self.render_mode.
Render mode can be `human` to display a window.
Other render modes in the default environments are `'rgb_array'`
which returns a numpy array and is supported by all environments outside of classic,
and `'ansi'` which returns the strings printed (specific to classic environments).
"""
raise NotImplementedError
def state(self) -> np.ndarray:
"""State returns a global view of the environment.
It is appropriate for centralized training decentralized execution methods like QMIX
"""
raise NotImplementedError(
"state() method has not been implemented in the environment {}.".format(
self.metadata.get("name", self.__class__.__name__)
)
)
[docs]
def close(self):
"""Closes any resources that should be released.
Closes the rendering window, subprocesses, network connections,
or any other resources that should be released.
"""
def observation_space(self, agent: AgentID) -> gymnasium.spaces.Space[ObsType]:
"""Takes in agent and returns the observation space for that agent.
MUST return the same value for the same agent name
Default implementation is to return the observation_spaces dict
"""
warnings.warn(
"Your environment should override the observation_space function. Attempting to use the observation_spaces dict attribute."
)
return self.observation_spaces[agent]
def action_space(self, agent: AgentID) -> gymnasium.spaces.Space[ActionType]:
"""Takes in agent and returns the action space for that agent.
MUST return the same value for the same agent name
Default implementation is to return the action_spaces dict
"""
warnings.warn(
"Your environment should override the action_space function. Attempting to use the action_spaces dict attribute."
)
return self.action_spaces[agent]
@property
def num_agents(self) -> int:
return len(self.agents)
@property
def max_num_agents(self) -> int:
return len(self.possible_agents)
def _deads_step_first(self) -> AgentID:
"""Makes .agent_selection point to first terminated agent.
Stores old value of agent_selection so that _was_dead_step can restore the variable after the dead agent steps.
"""
_deads_order = [
agent
for agent in self.agents
if (self.terminations[agent] or self.truncations[agent])
]
if _deads_order:
self._skip_agent_selection = self.agent_selection
self.agent_selection = _deads_order[0]
return self.agent_selection
def _clear_rewards(self) -> None:
"""Clears all items in .rewards."""
for agent in self.rewards:
self.rewards[agent] = 0
def _accumulate_rewards(self) -> None:
"""Adds .rewards dictionary to ._cumulative_rewards dictionary.
Typically called near the end of a step() method
"""
for agent, reward in self.rewards.items():
self._cumulative_rewards[agent] += reward
def agent_iter(
self, max_iter: int = 2**63
) -> AECIterable[AgentID, ObsType, ActionType]:
"""Yields the current agent (self.agent_selection).
Needs to be used in a loop where you step() each iteration.
"""
return AECIterable(self, max_iter)
def last(
self, observe: bool = True
) -> tuple[ObsType | None, float, bool, bool, dict[str, Any]]:
"""Returns observation, cumulative reward, terminated, truncated, info for the current agent (specified by self.agent_selection)."""
agent = self.agent_selection
assert agent is not None
observation = self.observe(agent) if observe else None
return (
observation,
self._cumulative_rewards[agent],
self.terminations[agent],
self.truncations[agent],
self.infos[agent],
)
[docs]
def _was_dead_step(self, action: ActionType) -> None:
"""Helper function that performs step() for dead agents.
Does the following:
1. Removes dead agent from .agents, .terminations, .truncations, .rewards, ._cumulative_rewards, and .infos
2. Loads next agent into .agent_selection: if another agent is dead, loads that one, otherwise load next live agent
3. Clear the rewards dict
Why the extra step is needed: an agent that dies must still be given one
more turn so that the user can call `last()` and see its final
observation, accumulated reward and termination/truncation flag. If the
agent were dropped from `.agents` as soon as it died, `agent_iter` would
never select it again and that final transition would be lost. The dead
step is where the agent is retired instead, which is why a dead agent is
still listed in `.agents` when `last()` first reports it done, and is
gone right after its `step(None)`.
Environments are responsible for calling this; the base class cannot do
it for them. Environments that skip it break the usual `not env.agents`
done check and can leave `agent_iter` looping forever, so implementing
it is effectively required rather than optional.
Examples:
Highly recommended to use at the beginning of step as follows::
def step(self, action):
if (
self.terminations[self.agent_selection]
or self.truncations[self.agent_selection]
):
self._was_dead_step(action)
return
# main contents of step
"""
if action is not None:
raise ValueError("when an agent is dead, the only valid action is None")
# removes dead agent
agent = self.agent_selection
assert self.terminations[agent] or self.truncations[agent], (
"an agent that was not dead as attempted to be removed"
)
del self.terminations[agent]
del self.truncations[agent]
del self.rewards[agent]
del self._cumulative_rewards[agent]
del self.infos[agent]
self.agents.remove(agent)
# finds next dead agent or loads next live agent (Stored in _skip_agent_selection)
_deads_order = [
agent
for agent in self.agents
if (self.terminations[agent] or self.truncations[agent])
]
if _deads_order:
if getattr(self, "_skip_agent_selection", None) is None:
self._skip_agent_selection = self.agent_selection
self.agent_selection = _deads_order[0]
else:
if getattr(self, "_skip_agent_selection", None) is not None:
assert self._skip_agent_selection is not None
self.agent_selection = self._skip_agent_selection
self._skip_agent_selection = None
self._clear_rewards()
@override
def __str__(self) -> str:
"""Returns a name which looks like: `space_invaders_v1`."""
if hasattr(self, "metadata"):
return self.metadata.get("name", self.__class__.__name__)
return self.__class__.__name__
@property
def unwrapped(self) -> AECEnv[AgentID, ObsType, ActionType]:
return self
class AECIterable(Iterable[AgentID], Generic[AgentID, ObsType, ActionType]):
def __init__(self, env, max_iter):
self.env = env
self.max_iter = max_iter
@override
def __iter__(self) -> AECIterator[AgentID, ObsType, ActionType]:
return AECIterator(self.env, self.max_iter)
class AECIterator(Iterator[AgentID], Generic[AgentID, ObsType, ActionType]):
def __init__(self, env: AECEnv[AgentID, ObsType, ActionType], max_iter: int):
self.env = env
self.iters_til_term = max_iter
@override
def __next__(self) -> AgentID:
if not self.env.agents or self.iters_til_term <= 0:
raise StopIteration
self.iters_til_term -= 1
return self.env.agent_selection
@override
def __iter__(self) -> AECIterator[AgentID, ObsType, ActionType]:
return self
[docs]
class ParallelEnv(Generic[AgentID, ObsType, ActionType]):
"""Parallel environment class.
It steps every live agent at once. If you are unsure if you
have implemented a ParallelEnv correctly, try running the `parallel_api_test` in
the Developer documentation on the website.
"""
metadata: dict[str, Any]
agents: list[AgentID]
possible_agents: list[AgentID]
observation_spaces: dict[
AgentID, gymnasium.spaces.Space[ObsType]
] # Observation space for each agent
action_spaces: dict[
AgentID, gymnasium.spaces.Space[ActionType]
] # Action space for each agent
[docs]
def reset(
self,
seed: int | None = None,
options: dict[str, Any] | None = None,
) -> tuple[dict[AgentID, ObsType], dict[AgentID, dict[str, Any]]]:
"""Resets the environment.
And returns a dictionary of observations (keyed by the agent name)
"""
raise NotImplementedError
[docs]
def step(
self, actions: dict[AgentID, ActionType]
) -> tuple[
dict[AgentID, ObsType],
dict[AgentID, float],
dict[AgentID, bool],
dict[AgentID, bool],
dict[AgentID, dict[str, Any]],
]:
"""Receives a dictionary of actions keyed by the agent name.
Returns the observation dictionary, reward dictionary, terminated dictionary, truncated dictionary
and info dictionary, where each dictionary is keyed by the agent.
"""
raise NotImplementedError
[docs]
def render(self) -> None | np.ndarray | str | list[Any]:
"""Displays a rendered frame from the environment, if supported.
Alternate render modes in the default environments are `'rgb_array'`
which returns a numpy array and is supported by all environments outside
of classic, and `'ansi'` which returns the strings printed
(specific to classic environments).
"""
raise NotImplementedError
[docs]
def close(self):
"""Closes any resources that should be released.
Closes the rendering window, subprocesses, network connections,
or any other resources that should be released.
"""
[docs]
def state(self) -> np.ndarray:
"""Returns the state.
State returns a global view of the environment appropriate for
centralized training decentralized execution methods like QMIX
"""
raise NotImplementedError(
"state() method has not been implemented in the environment {}.".format(
self.metadata.get("name", self.__class__.__name__)
)
)
[docs]
def observation_space(self, agent: AgentID) -> gymnasium.spaces.Space[ObsType]:
"""Takes in agent and returns the observation space for that agent.
MUST return the same value for the same agent name
Default implementation is to return the observation_spaces dict
"""
warnings.warn(
"Your environment should override the observation_space function. Attempting to use the observation_spaces dict attribute."
)
return self.observation_spaces[agent]
[docs]
def action_space(self, agent: AgentID) -> gymnasium.spaces.Space[ActionType]:
"""Takes in agent and returns the action space for that agent.
MUST return the same value for the same agent name
Default implementation is to return the action_spaces dict
"""
warnings.warn(
"Your environment should override the action_space function. Attempting to use the action_spaces dict attribute."
)
return self.action_spaces[agent]
@property
def num_agents(self) -> int:
return len(self.agents)
@property
def max_num_agents(self) -> int:
return len(self.possible_agents)
@override
def __str__(self) -> str:
"""Returns the name.
Which looks like: "space_invaders_v1" by default
"""
if hasattr(self, "metadata"):
return self.metadata.get("name", self.__class__.__name__)
return self.__class__.__name__
@property
def unwrapped(self) -> ParallelEnv[AgentID, ObsType, ActionType]:
return self
