zymera.env

env

Env contract, state schema, and registry — the spine of zymera.

This module owns the FROZEN contracts everything else conforms to (see docs/specs/2026-06-11-zymera-design.md):

• :class:ActionId / ACTION_DELTAS — the movement vocabulary.
• :class:Body / :class:World — the state pytree every component reads.
• :class:Env — gym-style base: reset(key) / step(state, action, key).
• :class:GridEnv — the orchestrator over the five components (worldgen / dynamics / comms / obs / missions).
• make / make_from / register_env / list_envs — registry.

Key protocol (FROZEN — parity and reproducibility ride on it):

• reset(key): wkey, skey = split(key) → terrain(wkey), spawn(skey); group assignment gets fold_in(key, 1); mission init gets fold_in(key, 2).
• step(state, action, key): k_chan, k_mis = split(key).

ActionId

Bases: IntEnum

Movement vocabulary on the square grid.

Integer values are stable — extend by appending so existing checkpoints keep their meaning.

Body

Per-agent physical state. SoA — every field is shape (N, ...).

World

The simulated world state. Immutable JAX pytree.

Field semantics (the committed contract components and user code read):

• explored — (H, W) int32 per-cell visit counts (heatmaps, redundancy).
• seen_by — (N, H, W) bool, each agent's OWN covered/sensed cells. The team-coverage metric reads covered = seen_by.any(0).
• comm_graph — (N, N) bool, edges that DELIVERED this step (realized, post-dropout). Potential topology lives in :class:zymera.metrics.StepCtx.
• channel — channel-owned pytree (ring buffers, beliefs); () when the env has no channel.
• mission — mission-owned pytree (waypoints, NPC positions); () by default. Structure must be identical between reset and every step.
• group — (N,) int32 group ids, assigned at reset (red-within-blue).

visited property

visited

(H, W) bool — any agent has stepped here.

covered property

covered

(H, W) bool — covered by any agent's footprint. THE coverage source.

Env

Gym-style base.

Contract::

obs, state = env.reset(key)
obs, state, reward, done, info = env.step(state, action, key)

state is a :class:World (or compatible pytree); action is (N,) int32; reward/done are (N,). info has a fixed keyset per env configuration (scan-stackable).

GridEnv

GridEnv(*, grid_h=8, grid_w=8, n_agents=1, cover_r=0, wall_sense_r=0, sense_free=False, terrain=None, spawn=None, dynamics=None, channel=None, obs=None, mission=None)

Bases: Env

Square-grid env composed from the five swappable components.

::

env = GridEnv(grid_h=16, grid_w=16, n_agents=4,
              spawn=ClusterSpawn(2),
              dynamics=GridDynamics(collision=SequentialClaim()),
              channel=GossipChannel(DiskTopology(5)),
              obs=GridObs(("known", "own_pos", "known_walls",
                           "neighbors", "local_frontier")),
              mission=Mission(terms=DEFAULT_TERMS))

One step runs: dynamics -> visit counts -> cover footprint -> channel delivery (comm_graph := delivered edges) -> metrics.derive (once) -> mission update/reward/done -> obs. Per-term UNWEIGHTED rewards land in info["reward_terms"]; mission success metrics in info["metrics"].

cover_r is the coverage footprint radius (core grid physics): the cells within Chebyshev cover_r of an agent count as covered by it (World.seen_by). cover_r=0 covers only the agent's own cell.

Source code in zymera/env.py

def __init__(self, *, grid_h: int = 8, grid_w: int = 8, n_agents: int = 1,
             cover_r: int = 0, wall_sense_r: int = 0, sense_free: bool = False,
             terrain=None, spawn=None,
             dynamics=None, channel=None, obs=None, mission=None):
    self.grid_h, self.grid_w = int(grid_h), int(grid_w)
    self.n_agents = int(n_agents)
    self.cover_r = int(cover_r)
    # Perception radius for walls (SLAM-style occupancy). 0 = wall-blind
    # (v0 parity): walls never enter the belief. >0 folds sensed walls into
    # the gossip outbox so the shared belief becomes a true occupancy map.
    self.wall_sense_r = int(wall_sense_r)
    # Occupancy belief: when True (with wall_sense_r>0), the gossip belief folds the
    # FULL sensed region (free + walls) within wall_sense_r, not just walls — so
    # `known` becomes a true free/occupied/unknown occupancy map and `local_frontier`
    # a real band-edge frontier (not the degenerate adjacent-cell signal at cover_r=0).
    self.sense_free = bool(sense_free)
    self.terrain = terrain if terrain is not None else OpenTerrain()
    self.spawn = spawn if spawn is not None else ScatterSpawn()
    self.dynamics = dynamics if dynamics is not None else GridDynamics()
    self.channel = channel if channel is not None else NullChannel()
    self.obs = obs if obs is not None else VectorObs()
    self.mission = mission if mission is not None else Mission(terms=())
    # Topology comes from the channel; group assignment from the mission.
    self._topology = getattr(self.channel, "topology", None)
    self._assignment = getattr(self.mission, "assignment", None)
    self.requires = frozenset(self.obs.requires) | frozenset(self.mission.requires)
    self._validate()

action_mask

action_mask(state)

(N, A) bool physical validity — delegates to dynamics.

Source code in zymera/env.py

def action_mask(self, state: World) -> jax.Array:
    """(N, A) bool physical validity — delegates to dynamics."""
    return self.dynamics.action_mask(state)

central_obs

central_obs(state)

Centralized critic view — delegates to the obs builder.

Source code in zymera/env.py

def central_obs(self, state: World) -> jax.Array:
    """Centralized critic view — delegates to the obs builder."""
    return self.obs.central_obs(state, None)

annotations

annotations(state)

Mission overlay primitives for viz.

Source code in zymera/env.py

def annotations(self, state: World):
    """Mission overlay primitives for viz."""
    return self.mission.annotations(state, state.mission)

register_env

register_env(name, factory)

Register factory(**kwargs) -> Env under name.

Source code in zymera/env.py

def register_env(name: str, factory: Callable[..., Env]) -> None:
    """Register ``factory(**kwargs) -> Env`` under ``name``."""
    if name in _REGISTRY:
        raise ValueError(f"env name already registered: {name!r}")
    _REGISTRY[name] = factory

make

make(name, **kwargs)

Construct a registered env (recipe) by name.

The returned env remembers (name, kwargs) so env.spec() / env.replace(...) / make_from round-trip.

Source code in zymera/env.py

def make(name: str, **kwargs) -> Env:
    """Construct a registered env (recipe) by name.

    The returned env remembers ``(name, kwargs)`` so ``env.spec()`` /
    ``env.replace(...)`` / ``make_from`` round-trip.
    """
    if name not in _REGISTRY:
        raise ValueError(f"unknown env: {name!r}; available: {list_envs()}")
    env = _REGISTRY[name](**kwargs)
    env._recipe = (name, dict(kwargs))
    return env

make_from

make_from(spec)

Rebuild an env from env.spec() output: {"recipe": name, **kwargs}.

Source code in zymera/env.py

def make_from(spec: Dict[str, Any]) -> Env:
    """Rebuild an env from ``env.spec()`` output: ``{"recipe": name, **kwargs}``."""
    spec = dict(spec)
    name = spec.pop("recipe")
    return make(name, **spec)