Multi agent env

Jittable abstract base class for all JaxMARL Environments.

agent_classes property

Returns a dictionary with agent classes

Format

agent_names: [agent_base_name_1, agent_base_name_2, ...]

name property

Environment name.

__init__(num_agents)

Parameters:

Name	Type	Description	Default
num_agents	int	maximum number of agents within the environment, used to set array dimensions	required

action_space(agent)

Action space for a given agent.

Parameters:

Name	Type	Description	Default
agent	str	agent name	required

Returns:

Name	Type	Description
space	Space	action space

get_avail_actions(state)

Returns the available actions for each agent.

Parameters:

Name	Type	Description	Default
state	State	environment state	required

Returns:

Type	Description
Dict[str, Array]	available actions (Dict[str, chex.Array]): available actions keyed by agent name

get_obs(state)

Applies observation function to state.

Parameters:

Name	Type	Description	Default
State	state	Environment state	required

Returns:

Name	Type	Description
Observations	Dict[str, Array]	observations keyed by agent names

observation_space(agent)

Observation space for a given agent.

Parameters:

Name	Type	Description	Default
agent	str	agent name	required

Returns:

Name	Type	Description
space	Space	observation space

reset(key)

Performs resetting of the environment.

Parameters:

Name	Type	Description	Default
key	PRNGKey	random key	required

Returns:

Name	Type	Description
Observations	Dict[str, Array]	observations for each agent, keyed by agent name
State	State	environment state

step(key, state, actions, reset_state=None)

Performs step transitions in the environment. Resets the environment if done. To control the reset state, pass reset_state. Otherwise, the environment will reset using self.reset.

Parameters:

Name	Type	Description	Default
key	PRNGKey	random key	required
state	State	environment state	required
actions	Dict[str, Array]	agent actions, keyed by agent name	required
reset_state	Optional[State]	Optional environment state to reset to on episode completion. Defaults to None.	None

Returns:

Name	Type	Description
Observations	Dict[str, Array]	next observations
State	State	next environment state
Rewards	Dict[str, float]	rewards, keyed by agent name
Dones	Dict[str, bool]	dones, keyed by agent name:
Info	Dict	info dictionary

step_env(key, state, actions)

Environment-specific step transition.

Parameters:

Name	Type	Description	Default
key	PRNGKey	random key	required
state	State	environment state	required
actions	Dict[str, Array]	agent actions, keyed by agent name	required

Returns:

Name	Type	Description
Observations	Dict[str, Array]	next observations
State	State	next environment state
Rewards	Dict[str, float]	rewards, keyed by agent name
Dones	Dict[str, bool]	dones, keyed by agent name:
Info	Dict	info dictionary