Multi-Agent Base

Base abstractions for JAX-compatible multi-agent environments.

Abstract base class for multi-agent JAX environments.

Derived from JaxMARL: https://github.com/FLAIROx/JaxMARL/blob/main/jaxmarl/environments/multi_agent_env.py

class f1tenth_gym_jax.envs.multi_agent_env.State(done: jax.Array | numpy.ndarray | numpy.bool | numpy.number, step: int)

Bases: object

Parameters:

  • done (Array | ndarray | bool | number)

  • step (int)

done: Array | ndarray | bool | number
step: int
replace(**updates)

Returns a new object replacing the specified fields with new values.

class f1tenth_gym_jax.envs.multi_agent_env.MultiAgentEnv(num_agents)

Bases: object

Jittable abstract base class for multi-agent environments.

Parameters:

num_agents (int)

reset(key)

Performs resetting of the environment.

Parameters:

key (Array)

Return type:

Tuple[Dict[str, Array | ndarray | bool | number], 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 randomly.

Parameters:

  • key (Array)

  • state (State)

  • actions (Dict[str, Array | ndarray | bool | number])

  • reset_state (State | None)

Return type:

Tuple[Dict[str, Array | ndarray | bool | number], State, Dict[str, float], Dict[str, bool], Dict]

step_env(key, state, actions)

Environment-specific step transition.

Parameters:

  • key (Array)

  • state (State)

  • actions (Dict[str, Array | ndarray | bool | number])

Return type:

Tuple[Dict[str, Array | ndarray | bool | number], State, Dict[str, float], Dict[str, bool], Dict]

get_obs(state)

Applies observation function to state.

Parameters:

state (State)

Return type:

Dict[str, Array | ndarray | bool | number]

observation_space(agent)

Observation space for a given agent.

Parameters:

agent (str)

action_space(agent)

Action space for a given agent.

Parameters:

agent (str)

get_avail_actions(state)

Returns the available actions for each agent.

Parameters:

state (State)

Return type:

Dict[str, Array | ndarray | bool | number]

property name: str

Environment name.

property agent_classes: dict

Returns a dictionary with agent classes, used in environments with heterogeneous agents.

Format:

agent_base_name: [agent_base_name_1, agent_base_name_2, …]