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.

With an even fleet, alternating early and late slots lets every vehicle take the direct route safely. An odd ring cannot alternate, so every safe commitment routes at least one vehicle through the bypass; with three vehicles all pairs are adjacent and the best safe outcome has exactly one bypass vehicle.

The ring has bounded physical degree two for every 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: RepeatedStageGame)

Bases: StageGameSafetyModel

Safety quotient for repeated Bertrand (see StageGameSafetyModel).

src.environments.ChickenSafetyModel ¶

ChickenSafetyModel(env: RepeatedStageGame)

Bases: StageGameSafetyModel

Safety quotient for repeated Chicken (see StageGameSafetyModel).

Matrix Environments¶

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

BertrandMatrix(
    *,
    num_agents: int = 2,
    max_steps: int = 200,
    T: float | None = None,
    R: float | None = None,
    S: float | None = None,
    P: float | None = None,
    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: RepeatedStageGame

Repeated n-player Bertrand price competition.

High is passive and Low aggressive. A firm that undercuts alone takes the market and earns T, two or more firms undercutting are in a price war and earn P, and a firm that prices High earns R when nobody undercut and S otherwise. Defaults: T=8, R=5, S=0, P=0. See RepeatedStageGame for the observation layout.

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

ChickenMatrix(
    *,
    num_agents: int = 2,
    max_steps: int = 200,
    T: float | None = None,
    R: float | None = None,
    S: float | None = None,
    P: float | None = None,
    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: RepeatedStageGame

Repeated n-player Chicken: a single-lane crossing.

Swerve is passive and Straight aggressive. A driver who goes Straight alone earns T, two or more drivers going Straight crash and earn P, and a driver who swerves earns R when nobody went Straight and S otherwise. Defaults: T=3, R=2, S=1, P=0. See RepeatedStageGame for the observation layout.

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 ¶

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

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` ¶