zymera.missions

missions

Missions — reward as data, plus group routing for the adversarial roadmap.

A :class:Mission is a frozen composition of :class:RewardTerm\ s: the env asks it for reward (Σ weighted terms + the UNWEIGHTED per-term dict that feeds info["reward_terms"]), done, mission-owned metrics, and drawable annotations. :class:GroupedMission routes per-group missions over a shared :class:~zymera.metrics.StepCtx so k-of-N red agents can score a different objective than the blue team (design spec §3.1).

JAX contract (the static-object rule):

• Missions and terms are frozen, hashable, closure-captured trace-time constants. Term SET is static; weights may become traced later — never gate on a weight's value (doctrine #3).
• init_state / update must keep an IDENTICAL pytree structure between reset and every step (doctrine #6). The default mission state is ().
• done gating on max_steps is Python-time (static config), never on traced data.

Sign convention: term functions return UNSIGNED magnitudes; penalties get their sign from the :class:RewardTerm weight (v0 subtracted w_coll * collisions — here the default collision term carries -4.0).

RewardTerm dataclass

RewardTerm(name, weight, fn, requires=frozenset())

One named reward component: weight · fn(prev, world, action, ctx).

requires names the :class:~zymera.metrics.StepCtx fields the term reads — the env unions these across terms/obs at __init__ so only the requested machinery compiles. fn must return the UNWEIGHTED (N,) value; the weighted sum happens in :meth:Mission.reward.

Point dataclass

Point(pos, tag='')

A marked cell — VIP, intruder, rally point. pos is (2,) (row, col).

Path dataclass

Path(cells, tag='')

An ordered cell sequence — patrol route, planned path. cells is (K, 2).

Region dataclass

Region(mask, tag='')

A cell mask — jammed zone, goal area. mask is (H, W) bool.

Mission dataclass

Mission(terms=(), max_steps=None)

A reward-term bundle with the full mission protocol surface.

Defaults: () mission state, identity update, timeout-only done (all-False when max_steps is None), sum-of-weighted-terms reward returning the unweighted per-term dict.

requires property

requires

Union of the terms' StepCtx requirements.

init_state

init_state(key, world)

Mission-owned pytree at reset. Default: () (uses no key).

Source code in zymera/missions.py

def init_state(self, key: jax.Array, world) -> Any:
    """Mission-owned pytree at reset. Default: ``()`` (uses no key)."""
    del key, world
    return ()

update

update(prev, world, ctx, mstate, key)

Advance mission-owned state (scripted NPCs, waypoints). Default: identity — MUST preserve pytree structure (doctrine #6).

Source code in zymera/missions.py

def update(self, prev, world, ctx, mstate, key: jax.Array) -> Any:
    """Advance mission-owned state (scripted NPCs, waypoints). Default:
    identity — MUST preserve pytree structure (doctrine #6)."""
    del prev, world, ctx, key
    return mstate

done

done(world, ctx, mstate)

(N,) bool. All-False, or step_count >= max_steps broadcast when max_steps is set (Python-time gate on static config).

Source code in zymera/missions.py

def done(self, world, ctx, mstate) -> chex.Array:
    """(N,) bool. All-False, or ``step_count >= max_steps`` broadcast when
    ``max_steps`` is set (Python-time gate on static config)."""
    del ctx, mstate
    n = world.n_agents
    if self.max_steps is None:
        return jnp.zeros((n,), dtype=jnp.bool_)
    return jnp.broadcast_to(world.step_count >= self.max_steps, (n,))

reward

reward(prev, world, action, ctx, mstate)

(Σ wᵢ·termᵢ (N,) f32, {name: UNWEIGHTED (N,) f32}).

The unweighted dict feeds info["reward_terms"] so analysis can re-weight post-hoc without re-running.

Source code in zymera/missions.py

def reward(
    self, prev, world, action, ctx, mstate,
) -> Tuple[chex.Array, Dict[str, chex.Array]]:
    """``(Σ wᵢ·termᵢ  (N,) f32,  {name: UNWEIGHTED (N,) f32})``.

    The unweighted dict feeds ``info["reward_terms"]`` so analysis can
    re-weight post-hoc without re-running.
    """
    del mstate
    n = world.n_agents
    unweighted = {
        t.name: t.fn(prev, world, action, ctx).astype(jnp.float32)
        for t in self.terms
    }
    total = jnp.zeros((n,), jnp.float32)
    for t in self.terms:
        total = total + t.weight * unweighted[t.name]
    return total.astype(jnp.float32), unweighted

metrics

metrics(world, ctx, mstate)

Mission-owned success metrics: team coverage fraction, plus the giant-component fraction when the ctx carries reach (Python-time gate — ctx structure is a function of env config only).

Source code in zymera/missions.py

def metrics(self, world, ctx, mstate) -> Dict[str, chex.Array]:
    """Mission-owned success metrics: team coverage fraction, plus the
    giant-component fraction when the ctx carries ``reach`` (Python-time
    gate — ctx structure is a function of env config only)."""
    del mstate
    out = {"coverage": coverage_fraction(world.covered)}
    if ctx is not None and ctx.reach is not None:
        out["giant_fraction"] = (
            ctx.reach.sum(-1).max() / world.n_agents
        ).astype(jnp.float32)
    return out

annotations

annotations(world, mstate)

Drawable primitives for zymera.viz. Default: none.

Source code in zymera/missions.py

def annotations(self, world, mstate) -> Tuple[Annotation, ...]:
    """Drawable primitives for zymera.viz. Default: none."""
    del world, mstate
    return ()

Assignment

Bases: Protocol

Reset-time group-id assignment: assign(key, n_agents) -> (N,) int32.

FixedAssignment dataclass

FixedAssignment(groups=None)

Deterministic group ids. groups=None → everyone in group 0. Ignores its key.

RandomKofN dataclass

RandomKofN(k, group=1)

A random k-subset of agents gets group (default 1); the rest stay 0.

Membership re-randomizes per reset WITHOUT retracing — the draw is pure JAX on the reset key (the red-within-blue graft, spec §3.1).

GroupedMission dataclass

GroupedMission(assignment, missions)

Per-group objectives over a shared world: missions[g] scores the agents with world.group == g.

Fixed-shape routing: every group-mission's reward/done is computed over ALL N agents, then where-selected by group id — no dynamic shapes, no retrace when :class:RandomKofN re-rolls membership. Per-term/metric names are namespaced g{i}/<name>. Agents whose group id has no mission (out of range) score 0 and never finish — keep ids in range.

NOTE (judge-panel mandate): the StepCtx is SHARED across groups — coverage counts every agent's footprint and connectivity reads the union graph. Per-group ctx semantics (e.g. blue-only coverage) must be decided explicitly per term when red training lands.

requires property

requires

Union of all group-missions' StepCtx requirements.

terms property

terms

All groups' terms, names namespaced g{i}/<name> (protocol compatibility; uniqueness holds when each sub-mission's does).

done

done(world, ctx, mstate)

(N,) bool — each agent reports its OWN group's mission done.

Source code in zymera/missions.py

def done(self, world, ctx, mstate) -> chex.Array:
    """(N,) bool — each agent reports its OWN group's mission done."""
    out = jnp.zeros((world.n_agents,), dtype=jnp.bool_)
    for g, (m, ms) in enumerate(zip(self.missions, mstate)):
        out = jnp.where(world.group == g, m.done(world, ctx, ms), out)
    return out

reward

reward(prev, world, action, ctx, mstate)

Each group-mission's total is computed over ALL N (fixed shape), then routed by where(group == g, ...) and summed. The per-term dict keeps the UNMASKED unweighted values under g{i}/<name> — mask post-hoc with world.group in analysis.

Source code in zymera/missions.py

def reward(
    self, prev, world, action, ctx, mstate,
) -> Tuple[chex.Array, Dict[str, chex.Array]]:
    """Each group-mission's total is computed over ALL N (fixed shape),
    then routed by ``where(group == g, ...)`` and summed. The per-term
    dict keeps the UNMASKED unweighted values under ``g{i}/<name>`` —
    mask post-hoc with ``world.group`` in analysis."""
    total = jnp.zeros((world.n_agents,), jnp.float32)
    unweighted: Dict[str, chex.Array] = {}
    for g, (m, ms) in enumerate(zip(self.missions, mstate)):
        r_g, terms_g = m.reward(prev, world, action, ctx, ms)
        total = total + jnp.where(world.group == g, r_g, 0.0)
        for name, val in terms_g.items():
            unweighted[f"g{g}/{name}"] = val
    return total.astype(jnp.float32), unweighted