t5lab.distill

distill

Distil the MVProp planner toward the classical wavefront (controller.nav_distance_field).

WHY (probe_mvprop.py): RL-only MVProp never sharpens — at init the flood dies within ~5 cells, the field is flat where agents are, and the gradient through K max-prop sweeps vanishes. The fix is a dense supervised target: for random (walls, goal) samples, the classical BFS wavefront gives the exact distance-to-goal field; MVProp learns to reproduce gamma^distance everywhere, which makes its argmax-value move equal the classical shortest-path move. The teacher is used ONLY at training time — at inference the planner is pure MVProp (no classical code runs), so the system stays fully learned.

Output: a frozen MVPropPlanner checkpoint plugged into ctde_v0 as action_head.controller='mvprop'.

Run (GPU): PY=$HOME/ZymeraLab/.venv/bin/python PYTHONPATH=. GRID=32 STEPS=4000 OUT=mvprop_distilled.eqx $PY -m t5lab.distill

gen_wall

gen_wall(key, grid, max_rects=5, p_open=0.3)

(H,W) bool — jittable random terrain: with prob p_open fully open, else up to max_rects random filled axis-aligned rectangles (varied clutter for general routing).

Source code in experiments/t5lab/distill.py

def gen_wall(key, grid, max_rects=5, p_open=0.30):
    """(H,W) bool — jittable random terrain: with prob ``p_open`` fully open, else up to
    ``max_rects`` random filled axis-aligned rectangles (varied clutter for general routing)."""
    ko, kr = jax.random.split(key)
    rows = jnp.arange(grid)[:, None]
    cols = jnp.arange(grid)[None, :]

    def add_rect(carry, rk):
        wall, = carry
        k1, k2, k3, k4, k5 = jax.random.split(rk, 5)
        r0 = jax.random.randint(k1, (), 0, grid)
        c0 = jax.random.randint(k2, (), 0, grid)
        rh = jax.random.randint(k3, (), 1, grid // 3 + 2)
        rw = jax.random.randint(k4, (), 1, grid // 3 + 2)
        place = jax.random.uniform(k5) < 0.8                        # some rect slots empty
        rect = (rows >= r0) & (rows < r0 + rh) & (cols >= c0) & (cols < c0 + rw)
        return (wall | (rect & place),), None

    wall0 = jnp.zeros((grid, grid), bool)
    (wall,), _ = jax.lax.scan(add_rect, (wall0,), jax.random.split(kr, max_rects))
    return jnp.where(jax.random.uniform(ko) < p_open, jnp.zeros_like(wall), wall)

sample_free_cell

sample_free_cell(key, wall)

(2,) int32 — a uniformly random FREE cell (jittable categorical over free cells).

Source code in experiments/t5lab/distill.py

def sample_free_cell(key, wall):
    """(2,) int32 — a uniformly random FREE cell (jittable categorical over free cells)."""
    free = (~wall).astype(jnp.float32).ravel()
    logits = jnp.where(free > 0, 0.0, -1e9)
    idx = jax.random.categorical(key, logits)
    return jnp.stack([idx // wall.shape[1], idx % wall.shape[1]]).astype(jnp.int32)

oracle_field

oracle_field(wall, goal, grid)

(H,W) target value field = the 4-connected geodesic gamma^dist to goal, computed by propagate with ORACLE passability w = 1 - wall. This is the EXACT field a 4-move agent should descend (von-Neumann geometry, same as the env action space and as MVProp's own flood) — so the student can reproduce it exactly by learning w -> 1-wall. Walls and unreached cells sit at ~0.

Source code in experiments/t5lab/distill.py

def oracle_field(wall, goal, grid):
    """(H,W) target value field = the 4-connected geodesic gamma^dist to ``goal``, computed by
    ``propagate`` with ORACLE passability ``w = 1 - wall``. This is the EXACT field a 4-move
    agent should descend (von-Neumann geometry, same as the env action space and as MVProp's
    own flood) — so the student can reproduce it exactly by learning ``w -> 1-wall``. Walls and
    unreached cells sit at ~0."""
    goal_oh = jnp.zeros((grid, grid)).at[goal[0], goal[1]].set(1.0)
    w = (~wall).astype(jnp.float32)
    return propagate(goal_oh, w, KPROP, GAMMA)                       # (H,W) gamma^(4-conn dist)