Skip to content

Notebook Environment Discovery

Notebook environment selection is now centralized in notebooks/env_discovery.py.

Why

The repo had two drifting sources of notebook environment metadata:

  • an earlier selector helper discovered real environment classes dynamically from src.environments
  • notebooks/utils/env_registry.py kept a static label-to-export-path map

That split had already gone stale:

  • Pursuit Evasion histories were previously mapped through stale environment metadata
  • gathering, Pursuit Evasion, and congestion histories ignored their per-agent-count export directories
  • new notebook UIs had no shared way to instantiate environments safely

Current model

NotebookEnvSpec combines:

  • discovered class information from src.environments
  • stable notebook display metadata
  • export-path rules used by notebook tooling
  • agent-count validation used for notebook instantiation

Discovery still comes from introspection, but a small override table keeps notebook-facing names and export layouts stable where pure reflection is not enough.

Important fields

  • family: top-level notebook/export group, currently gridworlds or matrix
  • slug: notebook/export folder slug, e.g. pursuit, markov_stag_hunt, dpgg
  • agent_param: constructor parameter used for agent count, either num_agents or n_agents
  • agent_count_options: fixed allowed counts when the environment only supports a small discrete set
  • export_includes_agent_count: whether notebook artifacts or experiment histories live under .../<n_agents>/...

Shared helpers

  • list_env_specs(...): filtered notebook environment list
  • get_env_spec(label): label/alias lookup
  • resolve_agent_count(spec, num_agents=...): validates notebook agent-count input
  • instantiate_env(spec, ...): constructor wrapper that applies validated agent count plus render_mode
  • export_root(...) / experiments_dir(...): notebook-friendly export path builders

Current notebook consumers

  • notebooks/diagrams/common/notebook_support.py
  • notebooks/diagrams/wm/graphviz_mo.py
  • notebooks/diagrams/iop/graphviz_level_mo.py
  • notebooks/diagrams/shielding/graphviz_mo.py
  • notebooks/__init__.py

Tradeoffs

  • The helper is not fully “automatic”: export layout and display names still need explicit overrides for some envs
  • That explicitness is intentional because notebook labels and export folder structure are part of the repo’s working conventions, not just class structure

Follow-up opportunities

  • move other notebook env pickers onto NotebookEnvSpec
  • add a small helper for notebook artifact directories if more notebooks start saving ad-hoc analysis outputs under exports/tmp
  • if notebook UIs diversify further, split overrides into display metadata vs artifact metadata