Environments API¶

Environment APIs cover the PettingZoo-compatible environments, labelled wrappers, and episode logs used by training and experiment notebooks.

Shared Wrappers and Logs¶

src.environments.episode_log.EpisodeLog `dataclass` ¶

EpisodeLog(
    episode: int,
    t_end: int,
    ep_len: int,
    cum_reward: float,
    cum_violations: int,
    start_step: int | None = None,
    unsafe_episode: bool = False,
    violation_rate: float = 0.0,
    first_violation_step: int | None = None,
    shield_level_raises: list[tuple[int, int]] = list(),
    shield_overrides: int = 0,
    shield_unsafe_proposals: int = 0,
    shield_missing_coverage: int = 0,
    shield_mean_proposed_risk: float | None = None,
    shield_max_proposed_risk: float | None = None,
    shield_mean_final_risk: float | None = None,
    shield_max_final_risk: float | None = None,
    shield_override_rate: float = 0.0,
    shield_unsafe_proposal_rate: float = 0.0,
    shield_missing_coverage_rate: float = 0.0,
    shield_infeasible_budget_steps: int = 0,
    shield_mean_current_budget: float | None = None,
    shield_min_current_budget: float | None = None,
    shield_mean_next_budget: float | None = None,
    shield_min_next_budget: float | None = None,
    shield_mean_final_safety_margin: float | None = None,
    shield_min_final_safety_margin: float | None = None,
    shield_mean_admissible_actions: float | None = None,
    shield_min_admissible_actions: int | None = None,
    shield_mean_action_availability: float | None = None,
    shield_mean_proposed_immediate_risk: float
    | None = None,
    shield_mean_final_immediate_risk: float | None = None,
    shield_immediate_risk_brier: float | None = None,
    shield_immediate_risk_ece: float | None = None,
    shield_immediate_risk_calibration_bins: list[
        dict[str, float | int | None]
    ] = list(),
    shield_initial_reachability_risk: float | None = None,
    shield_finite_episode_reachability_brier: float
    | None = None,
    shield_mean_proposed_reward: float | None = None,
    shield_mean_final_reward: float | None = None,
    shield_total_predicted_reward_delta: float = 0.0,
    shield_mean_level_floor: float | None = None,
    shield_final_level_floor: int | None = None,
    shield_credible_steps: int = 0,
    shield_mean_credible_delta_t: float | None = None,
    shield_mean_credible_excluded_mass: float | None = None,
    shield_max_credible_excluded_mass: float | None = None,
    shield_mean_credible_retained_mass: float | None = None,
    shield_posterior_tail_constraint_miss_steps: int = 0,
    shield_posterior_tail_constraint_miss_rate: float = 0.0,
    shield_decision_reason_counts: dict[str, int] = dict(),
    shield_proposed_action_counts: dict[int, int] = dict(),
    shield_final_action_counts: dict[int, int] = dict(),
    shield_replacement_counts: dict[str, int] = dict(),
    om_bayesian_level_changes: list[
        tuple[int, int, int]
    ] = list(),
    om_bayesian_start_level: int | None = None,
    om_bayesian_final_level: int | None = None,
    om_bayesian_mean_level: float | None = None,
    om_bayesian_level_change_count: int = 0,
    om_bayesian_level_decreases: int = 0,
    om_bayesian_max_decrease: int = 0,
    om_bayesian_level_occupancy: dict[int, int] = dict(),
    om_bayesian_mean_confidence: float | None = None,
    om_bayesian_min_confidence: float | None = None,
    om_bayesian_mean_entropy: float | None = None,
    om_bayesian_below_floor_steps: int = 0,
    om_bayesian_below_floor_rate: float = 0.0,
    om_mixture_mean_observed_action_probability: float
    | None = None,
    om_mixture_observed_action_nll: float | None = None,
    om_mixture_observed_action_brier: float | None = None,
    om_mixture_action_ece: float | None = None,
    om_mixture_action_calibration_bins: list[
        dict[str, float | int | None]
    ] = list(),
    om_per_level_observed_action_probability: list[
        float
    ] = list(),
    om_per_level_observed_action_nll: list[float] = list(),
    om_per_level_observed_action_brier: list[
        float
    ] = list(),
    om_per_level_action_ece: list[float] = list(),
)

src.environments.labelled_wrapper.LabelledEnv ¶

LabelledEnv(
    env,
    agent_id: str = "player_0",
    *,
    safety_model: SafetyAbstraction[Any],
)

Bases: BaseParallelWrapper

Generic episode logger and cost tracker backed by a safety abstraction.

Every wrapper supplies an explicit safety abstraction.

abstract_state ¶

abstract_state(observation: object) -> Any

Project a public observation onto its safety-relevant quotient.

record_shield_raise ¶

record_shield_raise(
    prev_floor: int, new_floor: int
) -> None

Called by an outer TransitionShield when its level rises.

record_om_bayesian_step ¶

record_om_bayesian_step(
    *,
    timestep: int,
    posterior_mode: int,
    posterior_alpha: list[float] | tuple[float, ...],
    level_floor: int,
    observed_action_probability: float,
    observed_action_nll: float,
    observed_action_brier: float,
    prediction_confidence: float,
    prediction_accuracy: float,
    per_level_observed_action_probability: list[float]
    | tuple[float, ...],
    per_level_observed_action_nll: list[float]
    | tuple[float, ...],
    per_level_observed_action_brier: list[float]
    | tuple[float, ...],
    per_level_prediction_confidence: list[float]
    | tuple[float, ...],
    per_level_prediction_accuracy: list[float]
    | tuple[float, ...],
) -> None

Record the unconstrained Bayesian OM mode and predictive scores.

Safety Contract¶

src.environments.safety.SafetyAbstraction ¶

Bases: Protocol[SafetyState]

Finite safety-relevant projection of a public observation.

Environment models own safety semantics. Wrappers may use a model for runtime cost accounting, while graph-based shields can use the same model to label exported states.

src.environments.safety.AgentSafetyState `dataclass` ¶

AgentSafetyState(
    unsafe_agents: frozenset[str] = frozenset(),
)

The quotient needed by OMSH's current per-agent safety objectives.

src.environments.safety.AgentSafetyModel ¶

AgentSafetyModel(agent_ids: tuple[str, ...] | list[str])

Shared agent validation and binary safety labelling.

src.environments.safety.UnsafeFlagSafetyModel ¶

UnsafeFlagSafetyModel(
    agent_ids: tuple[str, ...] | list[str],
    *,
    unsafe_flag_index: int,
    observation_name: str,
)

Bases: AgentSafetyModel

Safety model for observations containing one global failure flag.

Gridworld Environments¶

src.environments.gridworlds.congestion.env.CongestionNetwork ¶

CongestionNetwork(
    *,
    n_agents: int = 2,
    max_steps: int = 64,
    direct_reward: float = 1.0,
    bypass_reward: float = 0.2,
    collision_reward: float = -4.0,
    signalling_reward: float = 0.0,
    flatten_observations: bool = True,
    render_mode: str | None = None,
    seed: int | None = None,
)

Bases: ParallelEnv

A continuing ring of vehicles making one public merge commitment.

Every episode begins with a cost-free public signalling round. On the second round each vehicle commits to the early slot, late slot, or an always-safe slow bypass. Adjacent direct-route vehicles collide exactly when they select the same slot. The resulting safe operating state or collision state is absorbing, so a compatible commitment gives genuinely high-reward, zero-risk infinite-horizon operation rather than a fresh mismatch lottery at every timestep.

The ring has bounded physical degree two for every even fleet size. The public observation remains labelled by vehicle identity, and the default OmSh implementation still treats all non-focal vehicles as one joint opponent.

src.environments.gridworlds.markov_stag_hunt.env.MarkovStagHunt ¶

MarkovStagHunt(
    grid_size: tuple[int, int] = (4, 4),
    n_agents: int = 2,
    n_plants: int = 2,
    n_stags: int = 1,
    max_steps: int = 500,
    render_mode: str | None = None,
    stag_move_prob: float = 0.2,
    flatten_observations: bool = True,
    seed: int | None = None,
    reward_fn_by_agent: RewardFn
    | None = stag_hunt_base_rewards,
    min_hunters_for_stag: int = 2,
    damage_steps_on_maul: int = 1,
    respawn_stag_on_maul: bool = False,
    bounce_stag_on_maul: bool = True,
    plant_respawn_points: list[tuple[int, int]]
    | None = None,
    stag_respawn_points: list[tuple[int, int]]
    | None = None,
)

Bases: ParallelEnv

state ¶

state()

Global, agent-agnostic state (same content as obs), dtype=uint8.

num_cells ¶

num_cells() -> int

Number of spatial cells in the layered observation grid.

channel_names ¶

channel_names() -> list[str]

Human-friendly names for each observation channel index.

Returns a list of length n_obs_types such that index i in the list corresponds to channel i in the (H, W, C) observation tensor.

action_names ¶

action_names(action: int) -> str

Human-friendly name for an action integer ID.

get_state ¶

get_state()

Tuple layout

( agents: tuple[str, ...], results: tuple[int,int,int], # (stags_hunted, plants_harvested, maulings) damaged: tuple[int, ...], # ordered by possible_agents agent_positions: tuple[tuple[int,int]|None, ...] # ordered by possible_agents agent_counts: np.ndarray, # (H,W) int16 plants: tuple[tuple[int,int], ...], # sorted coords stags: tuple[tuple[int,int], ...], # sorted coords rng_state: object )

src.environments.gridworlds.pursuit.env.Pursuit ¶

Pursuit(
    *,
    num_agents: int = 3,
    layout: Layout | None = default_layout(),
    terminates_after_goal: bool = False,
    terminates_after_collision: bool = True,
    max_steps: int = 10000,
    reward_on_goal: float = 10,
    reward_guard_failure: float = -10,
    reward_normal_agent_collision: float = -10,
    reward_guard_agent_collision: float = 10,
    slip_intruder_prob: float = 0.1,
    slip_guard_prob: float = 0.1,
    flatten_observations: bool = True,
    render_mode: str | None = None,
    seed: int | None = None,
)

Bases: ParallelEnv

Pursuit Evasion (Parallel PettingZoo)

Agents: 2 or 3 (0 is the Guard; 1 and 2 are normal).
Actions: {Left, Right, Up, Down, Stay}. Moves into walls/bounds are ignored.
Collision: any two agents share a cell after the simultaneous move.
Goal: true iff agent 0 OR 1 is in the goal region.

The guard receives +10 for intercepting an intruder and -10 when an intruder reaches the goal. Intruders receive the opposite event rewards. Observations (globalised): binary channels shared by all agents: [Wall, Goal, A0, A1, (A2)] with dtype=np.float32. If flatten_observations is True, this is flattened to (HWC,).

state_space `property` ¶

state_space

Describe the shape/dtype of state().

state ¶

state()

Global, agent-agnostic state (same content as obs), matches state_space().

num_cells ¶

num_cells() -> int

Number of spatial cells in the layered observation grid.

channel_names ¶

channel_names() -> list[str]

Human-friendly names for each observation channel index.

Channels are

0: Wall 1: Goal 2+: Agents (0 = guard, others = normal agents)

Hidden episode time is not part of the public state contract.

action_names ¶

action_names(action: int) -> str

Human-friendly name for an action integer ID.

get_state ¶

get_state()

Tuple layout

( agents: tuple[str, ...], collision_ever: bool, agent_pos: tuple[Coord, ...] # ordered by possible_agents rng_state: object )

src.environments.gridworlds.gathering.env.Gathering ¶

Gathering(
    grid_size: tuple[int, int] = (6, 6),
    n_agents: int = 2,
    n_apples: int | None = None,
    *,
    apple_respawn_prob: float = 0.3,
    tag_hits_to_remove: int = 1,
    tag_out_steps: int = 1,
    laser_range: int | None = None,
    flatten_observations: bool = True,
    max_steps: int = 500,
    render_mode: str | None = None,
    seed: int | None = None,
)

Bases: ParallelEnv

state_space `property` ¶

state_space

Global, agent-agnostic state spec (same as obs), dtype=uint8.

state ¶

state()

Global, agent-agnostic state (identical to obs), dtype=uint8.

num_cells ¶

num_cells() -> int

Number of spatial cells in the layered observation grid.

channel_names ¶

channel_names() -> list[str]

Human-friendly names for each observation channel index.

Channels

0 .. n_agents-1 : player_i_active n_agents .. 2n_agents-1 : player_i_frozen 2n_agents : apple

action_names ¶

action_names(action: int) -> str

Human-friendly name for an action integer ID.

get_state ¶

get_state()

Tuple layout

( agents: tuple[str, ...], agent_pos: tuple[Coord|None, ...] # ordered by possible_agents tag_hits: tuple[int, ...] # ordered by possible_agents frozen_timer: tuple[int, ...] # ordered by possible_agents apples: tuple[Coord, ...] # sorted coords rng_state: object )

Environment Safety Models¶

src.environments.CongestionSafetyModel ¶

CongestionSafetyModel(env)

Bases: AgentSafetyModel

Per-vehicle safety projection for persistent merge collisions.

src.environments.GatheringSafetyModel ¶

GatheringSafetyModel(env)

Bases: AgentSafetyModel

Safety quotient for Gathering's active/frozen agent channels.

src.environments.MarkovStagHuntSafetyModel ¶

MarkovStagHuntSafetyModel(env)

Bases: AgentSafetyModel

Safety quotient for the per-agent damage channels.

src.environments.PursuitSafetyModel ¶

PursuitSafetyModel(env)

Bases: AgentSafetyModel

Safety quotient for the intruder-goal objective.

The global state is unsafe once every intruder occupies a goal. The same label applies to all agents because it is the environment-level failure event used by OMSH's current Pursuit experiments.

src.environments.BertrandSafetyModel ¶

BertrandSafetyModel(env)

Bases: UnsafeFlagSafetyModel

Safety quotient for the repeated Bertrand price-war flag.

src.environments.ChickenSafetyModel ¶

ChickenSafetyModel(env)

Bases: UnsafeFlagSafetyModel

Safety quotient for the repeated Chicken crash flag.

Matrix Environments¶

src.environments.matrix.bertrand.env.BertrandMatrix ¶

BertrandMatrix(
    *,
    num_agents: int = 2,
    max_steps: int = 200,
    T: float = 8.0,
    R: float = 5.0,
    S: float = 0.0,
    P: float = 0.0,
    player_1_T: float | None = None,
    player_1_R: float | None = None,
    player_1_S: float | None = None,
    player_1_P: float | None = None,
    flatten_observations: bool = True,
    render_mode=None,
    seed: int | None = None,
)

Bases: ParallelEnv

Repeated 2-player Bertrand (Parallel PettingZoo)

Stage game (Row vs Column) with payoffs: - High vs Low : (S, T) - Low vs High: (T, S) - High vs High: (R, R) - Low vs Low : (P, P) # price war

Defaults: T=8, R=5, S=0, P=0 (all floats).

Observations (global, binary channels), shape is either (C,) if flattened or (1,1,C): For each agent i in {0,1}: - ch 2i + 0 = 1 if agent_i's last action was High, else 0 - ch 2i + 1 = 1 if agent_i's last action was Low, else 0 Plus: - ch 2*num_agents = 1 if last outcome was a price war (both Low), else 0

On the first round (after reset), all channels are 0.

Shared environment contract

n_agents
n_obs_types
flatten_observations
state() / state_space

num_cells ¶

num_cells() -> int

Number of spatial cells in the layered observation grid.

BertrandMatrix is non-spatial: the entire game state is a single global cell. This returns 1 regardless of whether observations are flattened or shaped (1, 1, C).

channel_names ¶

channel_names() -> list[str]

Human-friendly names for each observation channel index.

Index i of the returned list corresponds to channel i in the observation vector (or the last axis of the (1, 1, C) tensor).

action_names ¶

action_names(action: int) -> str

Human-friendly name for an action integer ID.

get_state ¶

get_state()

Fully restorable snapshot (no deep clone).

Tuple layout

( agents: tuple[str, ...], last_actions: tuple[int|None, int|None], # (player_0, player_1) last_war: bool, cum_rewards: tuple[float, float], # (player_0, player_1) rng_state: object # np.random.RandomState.get_state() )

set_state ¶

set_state(state)

Restore a snapshot produced by get_state().

src.environments.matrix.chicken.env.ChickenMatrix ¶

ChickenMatrix(
    *,
    num_agents: int = 2,
    max_steps: int = 200,
    T: float = 3.0,
    R: float = 2.0,
    S: float = 1.0,
    P: float = 0.0,
    player_1_T: float | None = None,
    player_1_R: float | None = None,
    player_1_S: float | None = None,
    player_1_P: float | None = None,
    flatten_observations: bool = True,
    render_mode=None,
    seed: int | None = None,
)

Bases: ParallelEnv

Repeated 2-player Chicken (Parallel PettingZoo)

Stage game (Row vs Column) with payoffs: - Straight vs Swerve: (T, S) - Swerve vs Straight: (S, T) - Swerve vs Swerve: (R, R) - Straight vs Straight:(P, P) # crash

Defaults: T=3, R=2, S=1, P=0 (all floats).

Observations (global, binary channels), shape is either (C,) if flattened or (1,1,C): For each agent i in {0,1}: - ch 2i + 0 = 1 if agent_i's last action was Swerve, else 0 - ch 2i + 1 = 1 if agent_i's last action was Straight, else 0 Plus: - ch 2*num_agents = 1 if last outcome was a crash (both Straight), else 0

On the first round (after reset), all channels are 0.

Shared environment contract

n_agents
n_obs_types
flatten_observations
state() / state_space

num_cells ¶

num_cells() -> int

Number of spatial cells in the layered observation grid.

ChickenMatrix is non-spatial: the entire game state is a single global cell. This returns 1 regardless of whether observations are flattened or shaped (1, 1, C).

channel_names ¶

channel_names() -> list[str]

Human-friendly names for each observation channel index.

Channels (by construction in _obs): ch0 = player_0 Swerve ch1 = player_0 Straight ch2 = player_1 Swerve ch3 = player_1 Straight ch4 = crash flag

action_names ¶

action_names(action: int) -> str

Human-friendly name for an action integer ID.

get_state ¶

get_state()

Fully restorable snapshot (no deep clone).

Tuple layout

( agents: tuple[str, ...], last_actions: tuple[int|None, int|None], # (player_0, player_1) last_crash: bool, cum_rewards: tuple[float, float], # (player_0, player_1) rng_state: object # np.random.RandomState.get_state() )

set_state ¶

set_state(state)

Restore a snapshot produced by get_state().

Notebook Environment Discovery¶

notebooks.env_discovery.NotebookEnvSpec `dataclass` ¶

NotebookEnvSpec(
    display_name: str,
    family: str,
    slug: str,
    env_cls: EnvClass,
    class_name: str,
    module_path: str,
    agent_param: str | None,
    default_agent_count: int | None,
    min_agent_count: int | None,
    max_agent_count: int | None,
    agent_count_options: tuple[int, ...] | None,
    export_includes_agent_count: bool,
    aliases: tuple[str, ...] = (),
)

notebooks.env_discovery.discover_env_specs `cached` ¶

discover_env_specs() -> tuple[NotebookEnvSpec, ...]

notebooks.env_discovery.get_env_spec ¶

get_env_spec(label: str) -> NotebookEnvSpec

notebooks.env_discovery.resolve_agent_count ¶

resolve_agent_count(
    spec: NotebookEnvSpec | str,
    num_agents: int | None = None,
) -> int | None

notebooks.env_discovery.instantiate_env ¶

instantiate_env(
    spec: NotebookEnvSpec | str,
    *,
    num_agents: int | None = None,
    render_mode=None,
    **overrides,
)

notebooks.env_discovery.export_root ¶

export_root(
    spec: NotebookEnvSpec | str,
    *,
    num_agents: int | None = None,
) -> Path

Environments API¶

Shared Wrappers and Logs¶

src.environments.episode_log.EpisodeLog dataclass ¶

src.environments.labelled_wrapper.LabelledEnv ¶

abstract_state ¶

record_shield_raise ¶

record_om_bayesian_step ¶

Safety Contract¶

src.environments.safety.SafetyAbstraction ¶

src.environments.safety.AgentSafetyState dataclass ¶

src.environments.safety.AgentSafetyModel ¶

src.environments.safety.UnsafeFlagSafetyModel ¶

Gridworld Environments¶

src.environments.gridworlds.congestion.env.CongestionNetwork ¶

src.environments.gridworlds.markov_stag_hunt.env.MarkovStagHunt ¶

state ¶

num_cells ¶

channel_names ¶

action_names ¶

get_state ¶

src.environments.gridworlds.pursuit.env.Pursuit ¶

state_space property ¶

state ¶

num_cells ¶

channel_names ¶

action_names ¶

get_state ¶

src.environments.gridworlds.gathering.env.Gathering ¶

state_space property ¶

state ¶

num_cells ¶

channel_names ¶

action_names ¶

get_state ¶

Environment Safety Models¶

src.environments.CongestionSafetyModel ¶

src.environments.GatheringSafetyModel ¶

src.environments.MarkovStagHuntSafetyModel ¶

src.environments.PursuitSafetyModel ¶

src.environments.BertrandSafetyModel ¶

src.environments.ChickenSafetyModel ¶

Matrix Environments¶

src.environments.matrix.bertrand.env.BertrandMatrix ¶

num_cells ¶

channel_names ¶

action_names ¶

get_state ¶

set_state ¶

src.environments.matrix.chicken.env.ChickenMatrix ¶

num_cells ¶

channel_names ¶

action_names ¶

get_state ¶

set_state ¶

Notebook Environment Discovery¶

notebooks.env_discovery.NotebookEnvSpec dataclass ¶

notebooks.env_discovery.discover_env_specs cached ¶

notebooks.env_discovery.get_env_spec ¶

notebooks.env_discovery.resolve_agent_count ¶

notebooks.env_discovery.instantiate_env ¶

notebooks.env_discovery.export_root ¶

src.environments.episode_log.EpisodeLog `dataclass` ¶

src.environments.safety.AgentSafetyState `dataclass` ¶

state_space `property` ¶

state_space `property` ¶

notebooks.env_discovery.NotebookEnvSpec `dataclass` ¶

notebooks.env_discovery.discover_env_specs `cached` ¶