zymera.worldgen

worldgen

WorldGen — Terrain and Spawn components (the un-baked halves of reset).

Two small protocols, each a family of frozen, hashable, trace-time-static dataclasses (the static-object rule, spec §3.0):

• :class:Terrain — walls(key, h, w) -> (H, W) bool obstacle mask. Implementations: :class:OpenTerrain, :class:RandomWalls, :class:MapFile, :class:Rooms.
• :class:Spawn — positions(key, wall, n_agents) -> (N, 2) int32 distinct free cells. Implementations: :class:ScatterSpawn, :class:ClusterSpawn, :class:FixedSpawn.

Parity lineage (bit-for-bit, gated by tests/test_worldgen.py): RandomWalls ports v0 zymera.env._random_wall; ScatterSpawn ports the v0 World.initial weighted-choice spawn; ClusterSpawn ports examples/comm_coverage.py::CommCoverageEnv._cluster_spawn verbatim — including its internal akey, ckey = split(key), the uniform tie-break noise, and the lax.top_k overflow behaviour.

All methods are pure JAX (jit/vmap/scan-safe); h/w/n_agents are static Python ints, only key/wall are traced.

Terrain

Bases: Protocol

Obstacle-mask generator. h/w are static; key may be ignored.

walls

walls(key, h, w)

Return an (H, W) bool mask — True where a wall blocks the cell.

Source code in zymera/worldgen.py

def walls(self, key: jax.Array, h: int, w: int) -> chex.Array:
    """Return an ``(H, W)`` bool mask — True where a wall blocks the cell."""
    ...

Spawn

Bases: Protocol

Initial-position generator over a terrain's free cells.

positions

positions(key, wall, n_agents)

Return (N, 2) int32 (row, col) — distinct, free cells.

Source code in zymera/worldgen.py

def positions(self, key: jax.Array, wall: chex.Array, n_agents: int) -> chex.Array:
    """Return ``(N, 2)`` int32 ``(row, col)`` — distinct, free cells."""
    ...

OpenTerrain dataclass

OpenTerrain()

No obstacles — the all-False mask. Ignores the key.

RandomWalls dataclass

RandomWalls(n_obstacles)

n_obstacles uniformly random wall cells, fresh each reset key.

v0 _random_wall verbatim: flat indices drawn by jax.random.choice without replacement, then scattered into the mask. n_obstacles is clamped to the cell count at trace time.

MapFile dataclass

MapFile(cells)

A fixed map, stored as nested tuples so the component stays hashable (spec §3.0 — components that conceptually hold arrays store tuples and materialize inside methods). Ignores the key.

Build via :meth:from_string ('#' = wall, '.' = free, whitespace-only rows skipped) or :meth:load (same format, from a file).

from_string classmethod

from_string(s)

Parse a map drawing: '#' = wall, '.' = free; blank rows skipped.

Source code in zymera/worldgen.py

@classmethod
def from_string(cls, s: str) -> "MapFile":
    """Parse a map drawing: ``'#'`` = wall, ``'.'`` = free; blank rows skipped."""
    rows = []
    for line in s.splitlines():
        line = line.strip()
        if not line:
            continue
        bad = set(line) - {"#", "."}
        if bad:
            raise ValueError(f"MapFile: unknown chars {sorted(bad)} (use '#' and '.')")
        rows.append(tuple(c == "#" for c in line))
    return cls(cells=tuple(rows))

load classmethod

load(path)

Read :meth:from_string format from path.

Source code in zymera/worldgen.py

@classmethod
def load(cls, path) -> "MapFile":
    """Read :meth:`from_string` format from ``path``."""
    with open(path, "r") as f:
        return cls.from_string(f.read())

Rooms dataclass

Rooms(rooms=2, door_w=1)

rooms equal-width rooms separated by full-height vertical walls, one door_w-tall door per wall at a key-driven row.

Deliberately simple (spec §3.1): wall i (of rooms - 1) sits at column (i + 1) * w // rooms; the door's top row is drawn uniformly from [0, h - door_w] per wall from the reset key. rooms=1 is the open grid. Doors keep the free cells connected, so any spawn works.

ScatterSpawn dataclass

ScatterSpawn()

Distinct free cells anywhere on the grid.

v0 World.initial verbatim: flat indices sampled without replacement, weighted to free cells (wall cells get probability zero).

ClusterSpawn dataclass

ClusterSpawn(radius)

N distinct free cells clustered within radius of a random interior anchor.

v0 _cluster_spawn verbatim: akey, ckey = split(key); anchor drawn from free ∩ interior cells; every cell scored block-and-free (noise + 1) ≫ free (noise) ≫ wall (−1) with uniform tie-break noise; lax.top_k takes the N best. Guarantees N distinct cells and degrades gracefully when obstacles crowd the patch (overflow spills to the highest-noise free cells anywhere on the grid).

FixedSpawn dataclass

FixedSpawn(cells)

Spawn at explicit (row, col) cells, in order. Ignores the key.

Cells must be distinct; the first n_agents of them are used (so one component can serve a ladder of team sizes). The caller is responsible for the cells being free on the paired terrain — checked by tests, not at trace time (the wall array is traced).