Auto Environment Discovery

Reusable projects often need to answer a simple question: "what environments does this repo provide?" That answer is useful for experiment notebooks, batch jobs, smoke tests, docs tables, dashboards, and visualisation notebooks.

The reusable layer should not import project environments directly. Environment discovery is a project-level convention that can be written in a reusable style: small metadata records, lazy constructors, predictable filters, and no expensive side effects during discovery.

What It Is Useful For

Auto-discovery becomes valuable once a project has more than a couple of environments or environment variants.

Common uses:

• build a CLI command such as --env all, --env matrix/*, or --tags safety,gridworld;
• generate one experiment job per discovered environment;
• drive a marimo selector in a generic experiment, timing, or visualisation notebook;
• smoke-test that every registered environment can reset and step;
• render docs tables showing environment names, groups, tags, and capabilities;
• validate export directories by comparing discovered environment names against exports/<env_or_group>/...;
• keep paper scripts from drifting when a new environment is added.

The discovery code should return metadata first. Constructing the actual environment should be a separate, lazy step.

Recommended Shape

Use a small spec object in the project package:

from __future__ import annotations

from collections.abc import Callable, Mapping, Sequence
from dataclasses import dataclass, field
from typing import Any


@dataclass(frozen=True)
class EnvSpec:
    name: str
    group: str
    entry_point: str
    tags: tuple[str, ...] = ()
    default_config: Mapping[str, Any] = field(default_factory=dict)
    description: str = ""
    make: Callable[..., Any] | None = None

name should be stable and filesystem-safe enough to appear in export paths. group is the broader family, such as matrix, gridworld, omsh, csh, or pltlf. tags are for filtering. entry_point is a string reference to the constructor, such as:

my_project.envs.matrix.chicken:make_env

Keep make optional. Most discovery code should not need to import the real environment class.

Explicit Registry

For research repos, an explicit registry is usually the cleanest starting point:

ENV_SPECS = (
    EnvSpec(
        name="chicken",
        group="matrix",
        entry_point="my_project.envs.matrix.chicken:make_env",
        tags=("matrix", "safety", "two-agent"),
        default_config={"horizon": 100},
    ),
    EnvSpec(
        name="pursuit_evasion",
        group="gridworld",
        entry_point="my_project.envs.gridworlds.pursuit_evasion:make_env",
        tags=("gridworld", "safety", "multi-agent"),
        default_config={"map": "default"},
    ),
)


def discover_env_specs() -> tuple[EnvSpec, ...]:
    return ENV_SPECS

This looks boring, but it has excellent failure behavior. Adding an environment means adding one metadata entry. Discovery is deterministic, quick, and safe in docs builds, CI, and notebooks.

The helper that constructs the environment can stay separate:

from importlib import import_module


def load_entry_point(entry_point: str):
    module_name, function_name = entry_point.split(":", maxsplit=1)
    return getattr(import_module(module_name), function_name)


def make_env(spec: EnvSpec, **overrides):
    config = {**dict(spec.default_config), **overrides}
    return load_entry_point(spec.entry_point)(**config)

This separation is the important part. discover_env_specs() should be cheap; make_env(...) is where imports, wrappers, seeds, and heavy setup may happen.

Package Scanning

If the project has many environments and a consistent package layout, discovery can scan modules. For example, each environment package can expose ENV_SPEC:

src/my_project/envs/matrix/chicken.py
src/my_project/envs/matrix/inspection.py
src/my_project/envs/gridworlds/pursuit_evasion.py

Each module contains:

ENV_SPEC = EnvSpec(
    name="chicken",
    group="matrix",
    entry_point="my_project.envs.matrix.chicken:make_env",
    tags=("matrix", "two-agent"),
)

The scanner can collect those records:

from importlib import import_module
from pkgutil import walk_packages


def discover_env_specs(package_name: str = "my_project.envs") -> tuple[EnvSpec, ...]:
    package = import_module(package_name)
    specs: list[EnvSpec] = []
    for module_info in walk_packages(package.__path__, prefix=f"{package_name}."):
        if module_info.ispkg:
            continue
        module = import_module(module_info.name)
        spec = getattr(module, "ENV_SPEC", None)
        if spec is not None:
            specs.append(spec)
    return tuple(sorted(specs, key=lambda spec: (spec.group, spec.name)))

Only use this when environment modules are safe to import. If importing a module opens assets, initializes pygame, starts a simulator, downloads data, or checks for unavailable optional dependencies, prefer an explicit registry or a metadata file instead.

Metadata Files

For projects with heavy environment imports, keep metadata in JSON or YAML:

[
  {
    "name": "chicken",
    "group": "matrix",
    "entry_point": "my_project.envs.matrix.chicken:make_env",
    "tags": ["matrix", "two-agent"],
    "default_config": {"horizon": 100}
  }
]

Then discovery reads metadata without importing the environment package:

import json
from pathlib import Path


def discover_env_specs(path: str | Path = "envs.json") -> tuple[EnvSpec, ...]:
    payload = json.loads(Path(path).read_text(encoding="utf-8"))
    return tuple(EnvSpec(**item) for item in payload)

This is especially useful when docs or CI should list environments on machines that do not have every simulator dependency installed.

Filtering

Discovery is most useful when it accepts predictable filters:

def filter_env_specs(
    specs: Sequence[EnvSpec],
    *,
    groups: Sequence[str] | str = "all",
    envs: Sequence[str] | str = "all",
    tags: Sequence[str] | str = (),
) -> tuple[EnvSpec, ...]:
    group_set = _csvish(groups)
    env_set = _csvish(envs)
    tag_set = _csvish(tags)

    return tuple(
        spec
        for spec in specs
        if _matches(group_set, spec.group)
        and _matches(env_set, spec.name)
        and (not tag_set or tag_set.issubset(set(spec.tags)))
    )

Use the same CSV-ish conventions as the reusable notebooks:

--groups all
--groups matrix,gridworld
--envs chicken,pursuit_evasion
--tags safety

For UI notebooks, populate mo.ui.multiselect(...) from discover_env_specs(), then pass the selected names into the experiment suite configuration.

Pairing With Experiments

The reusable experiment runner assumes the notebook knows which environment it is running. Auto-discovery can sit one layer above that:

from rle.reusable import run_experiment_suite
from rle.reusable.experiments.models import ExperimentSuiteConfig

for env_spec in filter_env_specs(discover_env_specs(), tags="safety"):
    suite = ExperimentSuiteConfig(
        env_name=env_spec.name,
        export_dir=f"exports/{env_spec.group}/{env_spec.name}",
        env_config=env_spec.default_config,
        common_config=common_config,
        algorithm_configs=algorithm_configs,
    )
    run_experiment_suite(suite, algorithms)

This pattern is useful for batch jobs and smoke tests. For interactive research work, one notebook per environment is still often clearer. The notebook can use discovery to populate options, but the selected environment's config should be visible and editable in the notebook.

Validation

Add a small validation pass so discovery failures are caught early:

def validate_env_specs(specs: Sequence[EnvSpec]) -> None:
    seen: set[str] = set()
    for spec in specs:
        if spec.name in seen:
            raise ValueError(f"duplicate environment name: {spec.name}")
        seen.add(spec.name)
        if "/" in spec.name or "\\" in spec.name:
            raise ValueError(f"environment name is not path-safe: {spec.name}")
        if ":" not in spec.entry_point:
            raise ValueError(f"entry point must be module:function: {spec.entry_point}")

Good validation rules:

• names are unique;
• names and groups are path-safe;
• entry points have module:function form;
• tags use a consistent vocabulary;
• default configs are JSON-like if they will be stored in run records;
• optional dependency failures happen at make_env(...) time, not discovery time, unless the project explicitly wants stricter CI.

Adapting To Your Project

Add auto-discovery in the project package, not in rle.reusable, when the project knows what an environment is. A good location is something like:

src/<project>/envs/discovery.py
src/<project>/envs/registry.py

The code blocks most likely to change are:

• EnvSpec: add fields that matter for the project, such as num_agents, observation_kind, action_kind, supports_rgb_array, requires_assets, or paper_section;
• discover_env_specs(...): choose explicit registry, module scanning, metadata file loading, or a combination;
• filter_env_specs(...): define the filter vocabulary that CLI tools and notebooks should share;
• make_env(...): implement lazy construction and optional dependency handling;
• validation tests: assert every spec is unique, path-safe, and constructible in the project's intended CI environment.

Keep exports aligned with the discovered metadata. A simple convention is:

exports/<env_group>/<env_name>/<algorithm_slug>/timings.json
exports/<env_group>/<env_name>/runs/<algorithm_slug>_<run_idx>.json

Then reusable timings, legend export, and metric visualisation notebooks can treat <env_group>/<env_name> as the export group. That gives the project a consistent route from "discover all environments" to "run all selected algorithms" to "summarize timings and figures" without teaching reusable code about environment internals.