Skip to content

Multi room

MultiRoom (issue #182): a chain of num_rooms randomly-sized, randomly-positioned rooms, each connected to the next by an unlocked Door, winding through the grid in a genuinely random 2D path - a Goal sits in the last room.

MiniGrid's own generation algorithm is retry/backtrack-based: it repeatedly tries a random room placement and rejects it (retrying with a new random size/position, up to 8 attempts) if it overlaps an already-placed room or falls outside the grid. That's inherently data-dependent, unbounded control flow - doesn't trace under JAX the way every other navix environment's generation does. This file keeps the same algorithm (bounded to a fixed number of retries per room, jax.lax.while_loop instead of Python recursion) rather than approximating it with an easier-but-different generation scheme - verified end to end against MiniGrid's actual _placeRoom/_gen_grid (the four wall-relative positioning formulas, the entry/exit wall bookkeeping, the retry count).

Faithful to MiniGrid in one more respect worth flagging plainly: Navix-MultiRoom-N6-v0 uses a 25x25 grid (matching MiniGrid's own fixed default, used for every registration regardless of room count/ size - its retry logic is what keeps a random layout compact enough to fit that grid, not a smaller canvas). 625 cells is the biggest grid in navix by a wide margin (ObstructedMaze's Full is 16x16 = 256), and every per-step operation in this codebase scans the whole grid/entity set - so N6 is genuinely, permanently more expensive to run than any other navix environment, not just slower to compile. Kept anyway, deliberately, for fidelity - see this session's own discussion on issue #182 for the (rejected) cheaper alternative (a straight room chain) and the actual cost numbers.

One consequence worth calling out explicitly: Environment.step embeds a full Environment.reset call as its jax.lax.cond autoreset branch, so calling env.step(...) un-jitted in a Python loop retraces and recompiles that whole branch - reset included - on every single call, since eager lax.cond doesn't cache across separate top-level calls the way jax.jit does. For most navix environments that's an unnoticeable cost, because their reset is cheap. For MultiRoom, and especially N6, reset is exactly the heavy nested- retry search described above, so this pattern is a real footgun: it was measured to balloon to minutes of compile time and multiple GB of memory before crashing outright. Always jax.jit (or jax.vmap over a jitted function) env.step/env.reset before calling either in a loop against any MultiRoom variant.

Bases: Environment

See this module's own docstring for the generation algorithm and the N6 grid-size/cost note. num_rooms/max_room_size are static (pytree_node=False) - each registration gets its own traced generation graph, same convention as every other navix environment's structural parameters.

One attempt at placing all num_rooms rooms - a pure function of key, returning fixed-shape stacked arrays (not the growing Python lists place_room's own docstring describes - those are fine within one room's placement, where the call site's list length is static per room index, but this function's own output must be a single fixed-shape value so _reset can retry it whole via jax.lax.while_loop). The last return value is all_valid: whether every room actually found a legal placement (see place_room's docstring on why a single room's own retry can't always salvage a bad earlier choice) - _reset restarts this whole function with a fresh key when it's False, rather than ever using a layout that silently contains an invalid room.

A random position on room top/size's wall side, excluding the two corners (matching grid.room_grid_door_position's own corner-excluding convention - the same reason: a door needs a real interior cell on both sides, not the room's own corner).

A deterministic straight chain of n minimum-size rooms, centred as a whole and extending east - always valid by construction (GRID_SIZE=25 comfortably fits the chain's total span, MIN_ROOM_SIZE + (n-1) * (MIN_ROOM_SIZE-1) = 19 for the largest real n=6, for every real registration), unlike the random search _reset's own jax.lax.while_loop runs, which only usually finds a valid layout (see MAX_LAYOUT_RESTARTS). _reset uses this as the guaranteed-safe fallback for the rare episode where that random search doesn't resolve within its budget, rather than ever using a broken layout. Room positions are fixed (not random) here, but door colours still are, so it isn't the exact same episode every time this path is taken - assumes n >= 2 (true of every real registration; a hypothetical n=1 custom instantiation would need its own, simpler no-doors handling, not provided here).

Centring the chain's total span, not just room 0's own position, matters: consecutive rooms share one wall cell each, so the chain only grows by MIN_ROOM_SIZE - 1 per additional room, not a full MIN_ROOM_SIZE - anchoring room 0 alone at the grid centre and extending purely eastward from there runs the last room off the grid's edge for n=6 (confirmed directly: room 0 at column 10, chain reaching column 25 - already outside the valid 0-24 range, before that room's own width is even added).

Whether two rooms' interiors (each room's own 1-cell wall border excluded) overlap. Interiors, not full bounds: two rooms connected by a door are deliberately placed edge-to-edge, sharing exactly the 1-cell wall between them - a full-bounds overlap test would flag that legitimate adjacency as a collision. Interiors never legitimately touch, so this only rejects real collisions.

Bounded-retry room placement - jax.lax.while_loop's JAX- traceable equivalent of MiniGrid's own recursive retry (up to MAX_PLACEMENT_TRIES random (size, position) draws, accepting the first that's in-bounds with a margin (see below) and doesn't overlap any already-placed room). existing_tops/existing_sizes are plain Python lists, not padded arrays - safe here because the outer room-by-room loop in _reset/generate_layout is itself a static Python for (room count is a registration-time constant), so each call site's list length is static, same as ObstructedMazeFull's own door_positions accumulation.

BOUNDARY_MARGIN keeps every room a few cells clear of the grid edges (not just the two relevant to its own entry wall), lowering the odds a room lands flush against a boundary - which would make the next room's placement out-of-bounds by construction, unfixable by any amount of retrying here (confirmed directly: Navix-MultiRoom-N6-v0 PRNGKey(1), before this margin existed - room 3 landed at column 0, and room 4's own entry-wall formula, extend further left from a door already at the boundary, can only ever produce a negative column). The margin lowers the odds, but doesn't eliminate them - a genuinely tight, unluckily-drawn layout can still exhaust every retry (confirmed too: PRNGKey(2) still failed with the margin in place, room 3's door only 4 cells from an edge, too tight for any legal room 4 size to fit). That's what found (the third return value) is for: MiniGrid's true retry backtracks across multiple rooms when a placement can't be salvaged locally, which a single room's own retry loop structurally can't replicate - so on total failure here, this function reports it honestly instead of silently returning an invalid placement, and generate_layout's own outer retry restarts the entire chain with a fresh key instead - the closest JAX-traceable equivalent of "back up and reconsider an earlier choice", since which earlier room to blame isn't something this function can know from inside a single room's own placement attempt.

The new room's top-left corner, positioned so its wall side touches (entry_row, entry_col) and it extends away from there - verified against MiniGrid's actual _placeRoom's 4 wall-relative formulas. wall is a per-episode traced value (unlike every other navix room-placement helper's static side), hence jax.lax. switch rather than a plain Python if.