`ctde_v0.controller`¶

controller ¶

L1 controller + goal-candidate stencil + the mission-safety mechanism.

The multi-level (L3->L1) action stack lives here. L3 (the goal head in nets.Actor) picks one of K candidate relative waypoints; this module turns that choice into the actual env move, emitting ONLY valid 1-step moves (STAY fallback). The simulator therefore still sees movement-only actions and the 100-step budget is unchanged (agent_architecture.md L2/L1; EXPERIMENT_PLAN 1a').

Pieces¶

:func:goal_stencil — the fixed K relative offsets (center + 8 compass dirs at stride, in ABSOLUTE cells so the goal geometry is scale-invariant).
:func:goal_targets — absolute goal cell for every agent × candidate.
:func:greedy_move — the L1 controller: of the env-valid moves (dynamics.targets / action_mask), take the one that most reduces Chebyshev distance to the chosen goal; STAY if none helps.
:func:occupied_cell_mask — (N,A) bool flagging actions whose committed cell is the CURRENT cell of another agent (the hard collision-mask signal); both controllers can forbid_collision to remove those actions before argmin/max.
:func:team_lambda2_after_action — true λ₂ the team would have after a proposed joint move (used by the action-mask mechanism to score candidates).
:func:candidate_first_moves — for each agent × candidate, the env move the greedy controller would take FIRST (so the mechanism can mask a candidate by the connectivity it would cause).

Everything is pure JAX (vmap/scan-safe). The env (its dynamics / action_mask tables) is the single source of truth for wall-awareness.

goal_stencil ¶

goal_stencil(K, stride)

(K, 2) int32 relative offsets — the first K of the compass stencil, each scaled by stride (cell 0 = here, unscaled).

Source code in experiments/ctde_v0/controller.py

def goal_stencil(K: int, stride: int) -> jax.Array:
    """(K, 2) int32 relative offsets — the first ``K`` of the compass stencil,
    each scaled by ``stride`` (cell 0 = here, unscaled)."""
    if K < 1 or K > _COMPASS.shape[0]:
        raise ValueError(f"K must be in 1..{_COMPASS.shape[0]}, got {K}")
    base = _COMPASS[:K]
    scale = jnp.where((jnp.arange(K) == 0)[:, None], 1, stride)
    return (base * scale).astype(jnp.int32)

goal_targets ¶

goal_targets(pos, stencil, h, w)

(N, K, 2) int32 absolute goal cells (clipped in-bounds) for each agent × candidate. pos (N,2), stencil (K,2).

Source code in experiments/ctde_v0/controller.py

def goal_targets(pos: jax.Array, stencil: jax.Array, h: int, w: int) -> jax.Array:
    """(N, K, 2) int32 absolute goal cells (clipped in-bounds) for each agent ×
    candidate. ``pos`` (N,2), ``stencil`` (K,2)."""
    g = pos[:, None, :] + stencil[None, :, :]                  # (N,K,2)
    r = jnp.clip(g[..., 0], 0, h - 1)
    c = jnp.clip(g[..., 1], 0, w - 1)
    return jnp.stack([r, c], axis=-1).astype(jnp.int32)

occupied_cell_mask ¶

occupied_cell_mask(pos, valid_targets)

(N, A) bool — True where action a's committed target valid_targets[i,a] lands on the CURRENT cell of ANY OTHER agent j != i (the hard collision signal). Exact integer-cell match, broadcast over agents.

STAY (valid_targets[i, STAY] == pos[i], the agent's own cell) is NEVER flagged: an agent's own current cell does not count as "occupied by another", and we force STAY's column off regardless so it always remains selectable.

NOTE: this forbids moving onto a cell another agent occupies RIGHT NOW. The remaining case — two agents simultaneously claiming the same EMPTY cell under the NoCollision env semantics — is closed by :func:resolve_target_conflicts (deterministic index priority), applied to the chosen moves after selection. The two together GUARANTEE no two agents share a cell after the step (verified by tests/test_collision_free.py).

Source code in experiments/ctde_v0/controller.py

def occupied_cell_mask(pos: jax.Array, valid_targets: jax.Array) -> jax.Array:
    """(N, A) bool — True where action ``a``'s committed target ``valid_targets[i,a]``
    lands on the CURRENT cell of ANY OTHER agent ``j != i`` (the hard collision
    signal). Exact integer-cell match, broadcast over agents.

    STAY (``valid_targets[i, STAY] == pos[i]``, the agent's own cell) is NEVER
    flagged: an agent's own current cell does not count as "occupied by another",
    and we force STAY's column off regardless so it always remains selectable.

    NOTE: this forbids moving onto a cell another agent occupies RIGHT NOW. The remaining
    case — two agents simultaneously claiming the same EMPTY cell under the NoCollision env
    semantics — is closed by :func:`resolve_target_conflicts` (deterministic index priority),
    applied to the chosen moves after selection. The two together GUARANTEE no two agents share
    a cell after the step (verified by tests/test_collision_free.py).
    """
    n = pos.shape[0]
    # tgt[i,a] == pos[j] for some j != i  (compare every target to every cur cell).
    same = jnp.all(valid_targets[:, :, None, :] == pos[None, None, :, :], axis=-1)  # (N,A,N)
    other = ~jnp.eye(n, dtype=bool)                                  # (N,N) j != i
    occ = jnp.any(same & other[:, None, :], axis=-1)                 # (N,A) any other on it
    stay = int(ActionId.STAY)
    return occ.at[:, stay].set(False)                               # STAY never blocked

resolve_target_conflicts ¶

resolve_target_conflicts(move, valid_targets)

(N,) resolved move — the hard-collision COMPLETION of :func:occupied_cell_mask.

occupied_cell_mask forbids stepping onto a cell occupied RIGHT NOW, but two agents can still commit to the same EMPTY cell in a single tick (the documented NoCollision hole). This closes it by DETERMINISTIC INDEX PRIORITY: among all agents whose committed target cell valid_targets[i, move[i]] is identical, only the LOWEST-index agent proceeds; every higher-index agent reverts to STAY (its own current cell).

Together with occupied_cell_mask (no move targets a currently-occupied cell) and distinct start cells this GUARANTEES no two agents share a cell after the step: convergence is broken here; moving-onto-occupied is masked upstream; and a reverted agent lands on its own cell, which no other agent can target (that cell was occupied, hence already masked). Deterministic and jit-stable (no data-dependent shapes).

Source code in experiments/ctde_v0/controller.py

def resolve_target_conflicts(move: jax.Array, valid_targets: jax.Array) -> jax.Array:
    """(N,) resolved move — the hard-collision COMPLETION of :func:`occupied_cell_mask`.

    ``occupied_cell_mask`` forbids stepping onto a cell occupied RIGHT NOW, but two agents can
    still commit to the same EMPTY cell in a single tick (the documented NoCollision hole). This
    closes it by DETERMINISTIC INDEX PRIORITY: among all agents whose committed target cell
    ``valid_targets[i, move[i]]`` is identical, only the LOWEST-index agent proceeds; every
    higher-index agent reverts to STAY (its own current cell).

    Together with ``occupied_cell_mask`` (no move targets a currently-occupied cell) and distinct
    start cells this GUARANTEES no two agents share a cell after the step: convergence is broken
    here; moving-onto-occupied is masked upstream; and a reverted agent lands on its own cell,
    which no other agent can target (that cell was occupied, hence already masked). Deterministic
    and jit-stable (no data-dependent shapes)."""
    n = move.shape[0]
    tgt = valid_targets[jnp.arange(n), move]                          # (N,2) committed target cell
    same = jnp.all(tgt[:, None, :] == tgt[None, :, :], axis=-1)       # (N,N) i,j claim same cell
    lower = jnp.tril(jnp.ones((n, n), bool), -1)                      # (N,N) j < i (priority)
    yield_i = jnp.any(same & lower, axis=1)                           # (N,) a higher-priority agent claims it
    return jnp.where(yield_i, jnp.int32(int(ActionId.STAY)), move)

greedy_move ¶

greedy_move(pos, goal, valid_targets, action_valid, forbid_collision=False)

(N,) int32 — L1 greedy controller move toward goal.

For each agent: among env-VALID actions (action_valid (N,A) bool, from env.action_mask), pick the one whose committed cell (valid_targets (N,A,2), from dynamics.targets) minimizes Chebyshev distance to the goal. Ties and "no move helps" fall back to STAY (STAY is always valid). The result is always a valid move, so :class:SequentialClaim/the env never reverts it.

goal (N,2) int32 — the chosen absolute goal cell per agent. When forbid_collision is True the hard collision-mask (:func:occupied_cell_mask) removes actions whose target sits on another agent's CURRENT cell (their distance is set to +inf before the argmin); STAY always stays selectable.

Source code in experiments/ctde_v0/controller.py

def greedy_move(pos: jax.Array, goal: jax.Array, valid_targets: jax.Array,
                action_valid: jax.Array, forbid_collision: bool = False) -> jax.Array:
    """(N,) int32 — L1 greedy controller move toward ``goal``.

    For each agent: among env-VALID actions (``action_valid`` (N,A) bool, from
    ``env.action_mask``), pick the one whose committed cell (``valid_targets``
    (N,A,2), from ``dynamics.targets``) minimizes Chebyshev distance to the goal.
    Ties and "no move helps" fall back to STAY (STAY is always valid). The result
    is always a valid move, so :class:`SequentialClaim`/the env never reverts it.

    ``goal`` (N,2) int32 — the chosen absolute goal cell per agent. When
    ``forbid_collision`` is True the hard collision-mask (:func:`occupied_cell_mask`)
    removes actions whose target sits on another agent's CURRENT cell (their
    distance is set to +inf before the argmin); STAY always stays selectable.
    """
    d = _cheby(valid_targets, goal[:, None, :]).astype(jnp.float32)   # (N,A) dist if taken
    # forbid invalid actions by a large distance so argmin never selects them.
    d = jnp.where(action_valid, d, jnp.inf)
    if forbid_collision:
        # hard collision mask: same +inf-distance trick removes occupied-cell moves.
        d = jnp.where(occupied_cell_mask(pos, valid_targets), jnp.inf, d)
    # current distance (= STAY distance, STAY target == own cell).
    stay = int(ActionId.STAY)
    best = jnp.argmin(d, axis=-1).astype(jnp.int32)                   # (N,)
    # if the best valid move doesn't strictly improve on STAY, STAY.
    d_best = jnp.take_along_axis(d, best[:, None], axis=-1)[:, 0]
    d_stay = d[:, stay]
    move = jnp.where(d_best < d_stay, best, jnp.int32(stay))
    return move

nav_distance_field ¶

nav_distance_field(goal, blocked, planner='wavefront')

(H,W) float32 distance-to-goal field over the FREE cells of blocked (True = obstacle), the L2 nav-field. planner (Python-static) selects the solver:

"wavefront" / "bfs" — BFS distance transform (Jacobi min-plus relaxation).
"astar" — reuses the BFS field. Descending a shortest-path field reconstructs the FIRST move of an A shortest path; a real priority-queue A is not vmap-safe, and the emitted 1-step move is identical, so the field is shared.
"fmm" — fast-sweeping Eikonal approximation (:func:_fmm_...).

Iterations are fixed from the grid size (2*(H+W) BFS sweeps / 4 fast-sweep rounds) — enough to fully propagate an open grid and route moderate wall detours; the grid is small (agent_architecture.md L2). Pure JAX (vmap/scan/jit-safe).

Source code in experiments/ctde_v0/controller.py

def nav_distance_field(goal: jax.Array, blocked: jax.Array, planner: str = "wavefront") -> jax.Array:
    """(H,W) float32 distance-to-``goal`` field over the FREE cells of ``blocked`` (True =
    obstacle), the L2 nav-field. ``planner`` (Python-static) selects the solver:

    * ``"wavefront"`` / ``"bfs"`` — BFS distance transform (Jacobi min-plus relaxation).
    * ``"astar"``                 — reuses the BFS field. Descending a shortest-path field
      reconstructs the FIRST move of an A* shortest path; a real priority-queue A* is not
      vmap-safe, and the emitted 1-step move is identical, so the field is shared.
    * ``"fmm"``                   — fast-sweeping Eikonal approximation (:func:`_fmm_...`).

    Iterations are fixed from the grid size (``2*(H+W)`` BFS sweeps / 4 fast-sweep rounds) —
    enough to fully propagate an open grid and route moderate wall detours; the grid is
    small (agent_architecture.md L2). Pure JAX (vmap/scan/jit-safe)."""
    h, w = blocked.shape
    if planner == "fmm":
        return _fmm_distance_field(goal, blocked, n_rounds=4)
    # wavefront / bfs / astar -> BFS distance transform (down-gradient == shortest-path move)
    return _bfs_distance_field(goal, blocked, n_iters=2 * (h + w))

navfield_move ¶

navfield_move(pos, goal, blocked, valid_targets, action_valid, planner='wavefront', forbid_collision=True)

(N,) int32 — the L2 nav-field + L1 reactive controller move toward goal.

For each agent: plan a distance field to its goal over its OWN KB occupancy blocked[i] (:func:nav_distance_field), then among the env-VALID actions take the one whose committed cell (valid_targets) most reduces the field — descending the nav-field gradient one step. The reactive collision veto (the hard :func:occupied_cell_mask, on by default) removes moves onto a cell another agent occupies NOW; STAY is always selectable and is the fallback when no valid move strictly improves on staying (so the emitted move is always env-valid, exactly like :func:greedy_move). On open terrain the field is Chebyshev distance -> this reproduces the greedy move; around known walls it routes around them.

blocked (N,H,W) bool — per-agent obstacle map (known walls; unknown/known-free are traversable, optimistic). planner (Python-static) picks the field solver.

Source code in experiments/ctde_v0/controller.py

def navfield_move(pos: jax.Array, goal: jax.Array, blocked: jax.Array,
                  valid_targets: jax.Array, action_valid: jax.Array,
                  planner: str = "wavefront", forbid_collision: bool = True) -> jax.Array:
    """(N,) int32 — the L2 nav-field + L1 reactive controller move toward ``goal``.

    For each agent: plan a distance field to its goal over its OWN KB occupancy
    ``blocked[i]`` (:func:`nav_distance_field`), then among the env-VALID actions take the
    one whose committed cell (``valid_targets``) most reduces the field — descending the
    nav-field gradient one step. The reactive collision veto (the hard
    :func:`occupied_cell_mask`, on by default) removes moves onto a cell another agent
    occupies NOW; STAY is always selectable and is the fallback when no valid move strictly
    improves on staying (so the emitted move is always env-valid, exactly like
    :func:`greedy_move`). On open terrain the field is Chebyshev distance -> this reproduces
    the greedy move; around known walls it routes around them.

    ``blocked`` (N,H,W) bool — per-agent obstacle map (known walls; unknown/known-free are
    traversable, optimistic). ``planner`` (Python-static) picks the field solver."""
    n = pos.shape[0]
    fields = jax.vmap(lambda g, b: nav_distance_field(g, b, planner))(goal, blocked)  # (N,H,W)

    def gather_agent(i):
        D = fields[i]                                              # (H,W) agent i's field
        return D[valid_targets[i, :, 0], valid_targets[i, :, 1]]  # (A,) field if action taken

    d = jax.vmap(gather_agent)(jnp.arange(n))                     # (N,A) down-gradient score
    d = jnp.where(action_valid, d, _FAR)                         # forbid invalid actions
    if forbid_collision:
        d = jnp.where(occupied_cell_mask(pos, valid_targets), _FAR, d)  # reactive veto
    stay = int(ActionId.STAY)
    best = jnp.argmin(d, axis=-1).astype(jnp.int32)              # (N,) steepest descent
    d_best = jnp.take_along_axis(d, best[:, None], axis=-1)[:, 0]
    d_stay = d[:, stay]
    return jnp.where(d_best < d_stay, best, jnp.int32(stay))      # STAY unless a move helps

candidate_first_moves ¶

candidate_first_moves(pos, goal_cells, valid_targets, action_valid)

(N, K) int32 — the FIRST greedy move each agent would take for every candidate goal. Used by the action-mask mechanism to evaluate the connectivity each candidate would cause.

goal_cells (N,K,2). Vmaps :func:greedy_move over the K candidates.

Source code in experiments/ctde_v0/controller.py

def candidate_first_moves(pos: jax.Array, goal_cells: jax.Array,
                          valid_targets: jax.Array, action_valid: jax.Array) -> jax.Array:
    """(N, K) int32 — the FIRST greedy move each agent would take for every
    candidate goal. Used by the action-mask mechanism to evaluate the
    connectivity each candidate would cause.

    ``goal_cells`` (N,K,2). Vmaps :func:`greedy_move` over the K candidates.
    """
    def for_candidate(goal_k):                                       # goal_k (N,2)
        return greedy_move(pos, goal_k, valid_targets, action_valid)
    # vmap over K (axis 1 of goal_cells) -> (K, N) -> transpose
    moves = jax.vmap(for_candidate, in_axes=1, out_axes=0)(goal_cells)  # (K,N)
    return moves.T                                                    # (N,K)

relay_move ¶

relay_move(pos, valid_targets, action_valid, comm_r, sharp, forbid_collision=False)

(N,) int32 — relay L1 controller: each relay agent takes the env-VALID move that MAXIMIZES its own local connectivity proxy (soft incident-edge mass), others held at their current cell while it is scored. STAY is the safe default (it is always valid and is scored at the agent's current anchoring).

For agent i and candidate action a, we move ONLY i to valid_targets[i,a] (everyone else stays) and read i's soft degree at the resulting layout; the valid action with the largest value wins. This makes the relay hold/strengthen the bridge from purely local information (agent_architecture.md: relay = the λ̂₂-anchor tool). Result is always a valid move, so the env never reverts it.

When forbid_collision is True the hard collision-mask (:func:occupied_cell_mask) removes actions whose target sits on another agent's CURRENT cell (their score is set to -inf before the argmax); STAY always stays selectable.

Source code in experiments/ctde_v0/controller.py

def relay_move(pos: jax.Array, valid_targets: jax.Array, action_valid: jax.Array,
               comm_r: int, sharp: float, forbid_collision: bool = False) -> jax.Array:
    """(N,) int32 — relay L1 controller: each relay agent takes the env-VALID move
    that MAXIMIZES its own local connectivity proxy (soft incident-edge mass),
    others held at their current cell while it is scored. STAY is the safe default
    (it is always valid and is scored at the agent's current anchoring).

    For agent i and candidate action a, we move ONLY i to ``valid_targets[i,a]``
    (everyone else stays) and read i's soft degree at the resulting layout; the
    valid action with the largest value wins. This makes the relay hold/strengthen
    the bridge from purely local information (agent_architecture.md: relay = the
    λ̂₂-anchor tool). Result is always a valid move, so the env never reverts it.

    When ``forbid_collision`` is True the hard collision-mask
    (:func:`occupied_cell_mask`) removes actions whose target sits on another
    agent's CURRENT cell (their score is set to -inf before the argmax); STAY
    always stays selectable.
    """
    n, A = action_valid.shape
    blocked = occupied_cell_mask(pos, valid_targets) if forbid_collision else None

    def score_agent(i):
        # for each action a: i -> its committed cell, others stay at pos.
        def for_action(a):
            tgt = valid_targets[i, a]                                # (2,)
            pos_next = pos.at[i].set(tgt)                            # only i moves
            return _local_conn_score(pos_next, comm_r, sharp)[i]     # scalar
        s = jax.vmap(for_action)(jnp.arange(A))                      # (A,)
        s = jnp.where(action_valid[i], s, -jnp.inf)                  # forbid invalid
        if blocked is not None:
            s = jnp.where(blocked[i], -jnp.inf, s)                   # forbid collisions
        return jnp.argmax(s).astype(jnp.int32)                      # best valid action

    return jax.vmap(score_agent)(jnp.arange(n))                      # (N,)

relay_hold_move ¶

relay_hold_move(pos, valid_targets, action_valid, comm_r, sharp, hold_target=0.5, forbid_collision=False)

(N,) int32 — the "hold" relay L1 controller: a STATIC BEACON. Each relay agent simply STAYS where it is — a low-energy "don't wander, keep the bridge from where you stand" tool — UNLESS staying would leave it isolated, in which case it takes the SINGLE env-valid move that best restores a neighbour (minimal movement to re-anchor). STAY is the safe default and is always valid.

This is the complement of :func:relay_move (the lambda2_anchor tool, which ACTIVELY climbs local connectivity every step): relay_hold_move moves ONLY when its anchoring drops below the floor, so a well-connected relay never wanders off its post (agent_architecture.md: relay = "stop & hold the connection"; Relay-tool axis = hold-connection heuristic).

Decision (per relay i, from purely LOCAL information): * soft_deg_i(STAY) = i's soft incident-edge mass at its CURRENT cell (:func:_local_conn_score, the same proxy relay_move / the edge-margin signal use, so "anchored enough" agrees across the agent). * if soft_deg_i(STAY) >= hold_target -> STAY (it is comfortably anchored). * else (about to isolate) -> among the env-VALID moves, take the one that MAXIMIZES soft_deg_i at the resulting layout (others held), i.e. the minimal step that best re-establishes a neighbour. STAY is included in the argmax, so if no move improves anchoring the agent still STAYs (never an invalid move).

hold_target is the soft-degree floor that defines "isolated" (default 0.5 — a single in-range neighbour at the comm edge already carries ~0.5 soft mass at the default sharpness, so the relay holds as long as it has roughly one live link). When forbid_collision is True the hard collision-mask (:func:occupied_cell_mask) removes actions whose target sits on another agent's CURRENT cell (their score is set to -inf before the argmax); STAY always stays selectable. Result is always a valid move, so the env never reverts it. Pure JAX (vmap/scan/jit-safe) — same signature/contract as :func:relay_move.

Source code in experiments/ctde_v0/controller.py

def relay_hold_move(pos: jax.Array, valid_targets: jax.Array, action_valid: jax.Array,
                    comm_r: int, sharp: float, hold_target: float = 0.5,
                    forbid_collision: bool = False) -> jax.Array:
    """(N,) int32 — the "hold" relay L1 controller: a STATIC BEACON. Each relay
    agent simply STAYS where it is — a low-energy "don't wander, keep the bridge
    from where you stand" tool — UNLESS staying would leave it isolated, in which
    case it takes the SINGLE env-valid move that best restores a neighbour (minimal
    movement to re-anchor). STAY is the safe default and is always valid.

    This is the complement of :func:`relay_move` (the ``lambda2_anchor`` tool, which
    ACTIVELY climbs local connectivity every step): ``relay_hold_move`` moves ONLY
    when its anchoring drops below the floor, so a well-connected relay never wanders
    off its post (agent_architecture.md: relay = "stop & hold the connection";
    Relay-tool axis = hold-connection *heuristic*).

    Decision (per relay i, from purely LOCAL information):
      * soft_deg_i(STAY) = i's soft incident-edge mass at its CURRENT cell
        (:func:`_local_conn_score`, the same proxy ``relay_move`` / the edge-margin
        signal use, so "anchored enough" agrees across the agent).
      * if soft_deg_i(STAY) >= ``hold_target`` -> STAY (it is comfortably anchored).
      * else (about to isolate) -> among the env-VALID moves, take the one that
        MAXIMIZES soft_deg_i at the resulting layout (others held), i.e. the minimal
        step that best re-establishes a neighbour. STAY is included in the argmax, so
        if no move improves anchoring the agent still STAYs (never an invalid move).

    ``hold_target`` is the soft-degree floor that defines "isolated" (default 0.5 — a
    single in-range neighbour at the comm edge already carries ~0.5 soft mass at the
    default sharpness, so the relay holds as long as it has roughly one live link).
    When ``forbid_collision`` is True the hard collision-mask
    (:func:`occupied_cell_mask`) removes actions whose target sits on another agent's
    CURRENT cell (their score is set to -inf before the argmax); STAY always stays
    selectable. Result is always a valid move, so the env never reverts it. Pure JAX
    (vmap/scan/jit-safe) — same signature/contract as :func:`relay_move`.
    """
    n, A = action_valid.shape
    stay = int(ActionId.STAY)
    blocked = occupied_cell_mask(pos, valid_targets) if forbid_collision else None
    # current soft degree at the agent's present cell (the STAY anchoring).
    deg_now = _local_conn_score(pos, comm_r, sharp)                  # (N,)

    def score_agent(i):
        # for each action a: move ONLY i to its committed cell, read i's soft degree.
        def for_action(a):
            tgt = valid_targets[i, a]                                # (2,)
            pos_next = pos.at[i].set(tgt)                            # only i moves
            return _local_conn_score(pos_next, comm_r, sharp)[i]     # scalar
        s = jax.vmap(for_action)(jnp.arange(A))                      # (A,)
        s = jnp.where(action_valid[i], s, -jnp.inf)                  # forbid invalid
        if blocked is not None:
            s = jnp.where(blocked[i], -jnp.inf, s)                   # forbid collisions
        best = jnp.argmax(s).astype(jnp.int32)                      # best valid re-anchor
        # HOLD unless isolated: stay put while comfortably anchored, only move to
        # re-establish a neighbour when the current cell falls below the floor.
        isolated = deg_now[i] < hold_target
        return jnp.where(isolated, best, jnp.int32(stay))           # (,) int32

    return jax.vmap(score_agent)(jnp.arange(n))                      # (N,)

positions_after ¶

positions_after(pos, actions, valid_targets)

(N,2) int32 — committed positions if every agent took actions (N,).

Reads the env's valid_targets (N,A,2) table (already wall/boundary resolved). NoCollision semantics (agents may share a cell) — matches the comm-coverage recipe default; the mechanism only needs the comm graph, which is collision-agnostic.

Source code in experiments/ctde_v0/controller.py

def positions_after(pos: jax.Array, actions: jax.Array,
                    valid_targets: jax.Array) -> jax.Array:
    """(N,2) int32 — committed positions if every agent took ``actions`` (N,).

    Reads the env's ``valid_targets`` (N,A,2) table (already wall/boundary
    resolved). NoCollision semantics (agents may share a cell) — matches the
    comm-coverage recipe default; the mechanism only needs the comm graph, which
    is collision-agnostic.
    """
    n = actions.shape[0]
    return valid_targets[jnp.arange(n), actions]                     # (N,2)

team_lambda2_after ¶

team_lambda2_after(pos_next, comm_r, sharp)

Scalar true λ₂ of the soft comm-graph at pos_next (N,2).

Source code in experiments/ctde_v0/controller.py

def team_lambda2_after(pos_next: jax.Array, comm_r: int, sharp: float) -> jax.Array:
    """Scalar true λ₂ of the soft comm-graph at ``pos_next`` (N,2)."""
    return _lambda2(pos_next, comm_r, sharp)

safe_goal_mask ¶

safe_goal_mask(pos, goal_cells, valid_targets, action_valid, comm_r, sharp, min_lambda2)

(N, K) bool — action-mask mechanism: a candidate is SAFE iff, when its agent alone takes the greedy first move toward it (others STAY), the team's true λ₂ stays >= min_lambda2.

This is a LOCAL, per-agent guardrail (each agent screens its own candidates independently against the connectivity floor) — the "forbid goals that would disconnect" mechanism. If a row would mask ALL candidates, the "here" candidate (index 0, a STAY goal) is force-unmasked so a valid goal always exists.

Source code in experiments/ctde_v0/controller.py

def safe_goal_mask(pos: jax.Array, goal_cells: jax.Array, valid_targets: jax.Array,
                   action_valid: jax.Array, comm_r: int, sharp: float,
                   min_lambda2: float) -> jax.Array:
    """(N, K) bool — action-mask mechanism: a candidate is SAFE iff, when its
    agent alone takes the greedy first move toward it (others STAY), the team's
    true λ₂ stays >= ``min_lambda2``.

    This is a LOCAL, per-agent guardrail (each agent screens its own candidates
    independently against the connectivity floor) — the "forbid goals that would
    disconnect" mechanism. If a row would mask ALL candidates, the "here"
    candidate (index 0, a STAY goal) is force-unmasked so a valid goal always
    exists.
    """
    n = pos.shape[0]
    first_moves = candidate_first_moves(pos, goal_cells, valid_targets, action_valid)  # (N,K)
    K = goal_cells.shape[1]
    stay = jnp.full((n,), int(ActionId.STAY), dtype=jnp.int32)

    def lambda2_for_agent_candidate(i, kk):
        actions = stay.at[i].set(first_moves[i, kk])                 # only agent i moves
        pos_next = positions_after(pos, actions, valid_targets)
        return team_lambda2_after(pos_next, comm_r, sharp)

    # vectorize over (N agents) x (K candidates) -> (N, K) true-λ₂-if-taken.
    l2 = jax.vmap(jax.vmap(lambda2_for_agent_candidate, in_axes=(None, 0)),
                  in_axes=(0, None))(jnp.arange(n), jnp.arange(K))   # (N,K)
    safe = l2 >= min_lambda2
    # guarantee >=1 safe candidate per agent: force "here" (index 0) on.
    safe = safe.at[:, 0].set(True)
    return safe

ctde_v0.controller¶

controller ¶

Pieces¶

goal_stencil ¶

goal_targets ¶

occupied_cell_mask ¶

resolve_target_conflicts ¶

greedy_move ¶

nav_distance_field ¶

navfield_move ¶

candidate_first_moves ¶

relay_move ¶

relay_hold_move ¶

positions_after ¶

team_lambda2_after ¶

safe_goal_mask ¶

`ctde_v0.controller`¶