Grid

Array helpers for grid geometry - the low-level toolkit the environments, navix.actions and navix.observations are built from.

Everything here is a pure JAX function over plain arrays; nothing knows about State or Entity. Groups:

• coordinate <-> flat-index conversion (coordinates, idx_from_coordinates, ...);
• movement and rotation of a (row, col) position or a direction (translate, rotate, translate_forward/left/right);
• rotating an image patch to align it with a direction (align, rotate_tile);
• random placement (random_positions, random_distinct_positions, random_position_far_from, random_directions, random_colour);
• building maps (room, two_rooms, vertical_wall, horizontal_wall, from_ascii_map) and the multi-room grid helpers (room_grid*, room_*, RoomsGrid);
• cropping and first-person rendering (crop, view_cone, draw_grid_lines, apply_minigrid_opacity).

Convention: positions are (row, col), directions are 0 east, 1 south, 2 west, 3 north, and a "grid" is i32[H, W] with 0 = floor and -1 = wall.

Coordinates = Tuple[Array, Array] module-attribute

A (rows, cols) pair of index arrays, as returned by coordinates - the shape jnp.where / advanced indexing expect.

RoomsGrid

Bases: PyTreeNode

A grid of rooms. Each room is represented as a 2D grid of shape (room_height, room_width), with walls set to -1 and empty tiles set to 0. The grid of rooms is represented as a 2D grid of shape (rows * (room_height + 1), cols * (room_width + 1)), with walls set to -1 and empty tiles set to 0. The grid of rooms is represented as a 2D grid of shape (rows * (room_height + 1), cols * (room_width + 1)), with walls set to -1 and empty tiles set to 0.

create(num_rows, num_cols, room_size) classmethod

Creates a grid of rooms with the given number of rows and columns, and the given room size.

Parameters:

Name	Type	Description	Default
num_rows	int	The number of rows.	required
num_cols	int	The number of columns.	required
room_size	Tuple[int, int]	The size of each room (height, width).	required

Returns:

Name	Type	Description
RoomsGrid	RoomsGrid	A grid of rooms.

get_grid(occupied_positions=None)

Computes the array representation of the grid of rooms, with walls set to -1 and empty tiles set to 0.

Parameters:

Name	Type	Description	Default
occupied_positions	Array	A batch of extra occupied positions of shape (n, 2). Defaults to None.	None

Returns:

Name	Type	Description
Array	Array	A 2D grid of shape (rows * (room_height + 1), cols * (room_width + 1)).

position_in_room(row, col, *, key)

Generates a random position in a given room.

Parameters:

Name	Type	Description	Default
row	Array	The row index of the room.	required
col	Array	The column index of the room.	required
key	Array	A random key.	required

Returns:

Name	Type	Description
Array	Array	A random position in the given room.

position_on_border(row, col, side, *, key)

Generates a random position on the border of a given room. Side is 0: west, 1: east, 2: north, 3: south (like padding)

Parameters:

Name	Type	Description	Default
row	Array	The row index of the room.	required
col	Array	The column index of the room.	required
side	int	The side of the room.	required
key	Array	A random key.	required

Returns:

Name	Type	Description
Array	Array	A random position on the border of the given room.

align(patch, current_direction, desired_direction)

Aligns a patch of the grid from the current direction to the desired direction.

Parameters:

Name	Type	Description	Default
patch	Array	A patch of the grid.	required
current_direction	Array	The current direction in the range [0, 1, 2, 3] representing the cardinal directions [east, south, west, north].	required
desired_direction	Array	The desired direction in the range [0, 1, 2, 3] representing the cardinal directions [east, south, west, north].	required

Returns:

Name	Type	Description
Array	Array	A patch of the grid aligned to the desired direction.

apply_minigrid_opacity(image, opacity=jnp.asarray(0.7))

Applies minigrid opacity to the given image, used in minigrid.wrappers.RGBImgPartialObsWrapper. The default MiniGrid opacity is 0.7.

Parameters:

Name	Type	Description	Default
image	Array	The input image to which opacity is applied.	required
opacity	Array	The opacity value to apply. Defaults to 0.7.	asarray(0.7)

Returns:

Name	Type	Description
Array	Array	The input image with applied opacity.

coordinates(grid)

Returns a tuple of 2D coordinates [(col, row), ...] for each cell in the grid. A grid array of shape i32[height, width] will return a tuple of length (height * width), containing two arrays, each of shape i32[2].

Parameters:

Name	Type	Description	Default
grid	Array	A 2D grid of shape (height, width).	required

Returns:

Type	Description
Coordinates	Tuple[Array, Array]: A tuple of two arrays containing the 2D coordinates of each cell in the grid.

coordinates_from_idx(grid, idx)

Converts a flat index of shape i32[] into a 2D coordinate i32[2] containing (col, row) data. The index is calculated as idx = row * width + col.

Parameters:

Name	Type	Description	Default
grid	Array	A 2D grid of shape (height, width).	required
idx	Array	A flat index of shape i32[].	required

Returns:

Name	Type	Description
Array	Array	A 2D coordinate of shape i32[2] containing the (col, row) data.

crop(grid, origin, direction, radius, padding_value=100)

Crops a grid around a given origin, facing a given direction, with a given radius.

Parameters:

Name	Type	Description	Default
grid	Array	A 2D grid of shape (height, width).	required
origin	Array	The origin of the crop.	required
direction	Array	The direction the crop is facing.	required
radius	int	The radius of the crop.	required
padding_value	int	The padding value. Defaults to 0.	100

Returns:

Name	Type	Description
Array	Array	A cropped grid.

draw_grid_lines(tile, luminosity=jnp.asarray(100), corner_luminosity=None)

Draws grid lines on the given tile.

luminosity defaults to MiniGrid's own raw grey constant (COLORS["grey"]), used as-is for e.g. walls. But MiniGrid's grid lines specifically are drawn as a sub-pixel-width, anti-aliased strip - at typical tile sizes only partially covered by that raw colour, blending it with the background - whereas this function does a hard, fully-opaque fill of line_thickness whole pixels with no antialiasing. Using the raw 100 here renders visibly bolder/brighter than MiniGrid's actual output; pass a lower value (e.g. ~32, empirically matched against a real MiniGrid render - see rendering/cache.py::render_background) to compensate.

corner_luminosity fills just the line_thickness x line_thickness corner block with a separate value, defaulting to luminosity (i.e. no distinct corner treatment) when omitted. In MiniGrid, the top and left line strips are drawn independently and both get anti-aliased, so the corner - covered by both - ends up brighter than either strip alone; a single flat luminosity for the whole line can't reproduce that without also being tuned to a higher value at just the corner.

Parameters:

Name	Type	Description	Default
tile	Array	The input tile to which grid lines are drawn.	required

Returns:

Name	Type	Description
Array	Array	The tile with drawn grid lines.

from_ascii_map(ascii_map, mapping={})

Converts an ASCII map into a 2D grid. The ASCII map is a string where each character represents a tile in the grid. The mapping dictionary can be used to map ASCII characters to integer values. By default, the mapping is as follows: - # is mapped to -1 - . is mapped to 0

Parameters:

Name	Type	Description	Default
ascii_map	str	The ASCII map.	required
mapping	Dict[str, int]	A dictionary mapping ASCII characters to integer values. Defaults to {}.	{}

Returns:

Name	Type	Description
Array	Array	A 2D grid representing the ASCII map.

horizontal_wall(grid, col_idx, opening_row_idx=None)

Creates a horizontal wall in the grid at the given column index, with an opening at the given row index.

Parameters:

Name	Type	Description	Default
grid	Array	A 2D grid of shape (height, width).	required
col_idx	int	The column index where the wall is placed.	required
opening_row_idx	Array	The row index where the opening is placed. Defaults to None.	None

Returns:

Name	Type	Description
Array	Array	A 2D grid of shape (height, width) with a horizontal wall.

idx_from_coordinates(grid, coordinates)

Converts a batch of 2D coordinates [(col, row), ...] into a flat index

Parameters:

Name	Type	Description	Default
grid	Array	A 2D grid of shape (height, width).	required
coordinates	Array	A batch of 2D coordinates of shape (batch_size, 2).	required

Returns:

Name	Type	Description
Array	Array	A flat index of shape i32[batch_size] for each coordinate in the batch.

mask_by_coordinates(grid, address, comparison_fn=jnp.greater_equal)

This is a workaround to compute dynamicly-sized masks in XLA, which would not be possible otherwise. Returns a mask of the same shape as grid where the value is 1 if the corresponding element in grid satisfies the comparison_fn with the corresponding element in address (col, row) and 0 otherwise.

Parameters:

Name	Type	Description	Default
grid	Array	A 2D grid of shape (height, width).	required
address	Coordinates	A tuple of 2D coordinates (col, row).	required
comparison_fn	Callable[[Array, Array], Array]	A comparison function. Defaults to jnp.greater_equal.	greater_equal

Returns:

Name	Type	Description
Array	Array	A boolean mask of the same shape as grid.

open_wall(grid, position)

Removes a single wall cell (e.g. to place a door through it) by setting it to floor (0) - the door's own entity (not the grid) then controls whether the player can actually pass through.

Parameters:

Name	Type	Description	Default
grid	Array	A 2D grid of shape (height, width).	required
position	Array	The (row, col) position to open.	required

Returns:

Name	Type	Description
Array	Array	The updated grid.

positions_equal(a, b)

Checks if two points are equal.

Parameters:

Name	Type	Description	Default
a	Array	A 2D coordinate of shape i32[2] containing the (col, row) data.	required
b	Array	A 2D coordinate of shape i32[2] containing the (col, row) data.	required

Returns:

random_colour(key, n=1)

Generates n random colours in the range [0, 1, 2, 3, 4, 5].

Parameters:

Name	Type	Description	Default
key	Array	A random key.	required
n	int	The number of random colours to generate. Defaults to 1.	1

Returns:

Name	Type	Description
Array	Array	A batch of random colours of shape u8[n].

random_directions(key, n=1)

Generates n random directions in the range [0, 1, 2, 3] representing the cardinal directions [east, south, west, north].

Parameters:

Name	Type	Description	Default
key	Array	A random key.	required
n	int	The number of random directions to generate. Defaults to 1.	1

Returns:

Name	Type	Description
Array	Array	A batch of random directions of shape i32[n].

random_distinct_positions(key, grid, n, exclude=jnp.asarray((-1, -1)))

Generates n mutually distinct random positions in the grid, each also excluding exclude - unlike random_positions(..., n=n), whose n draws are i.i.d. (jax.random.categorical samples with replacement) and so can collide with each other (see issue #172's PR: GoToObject/Fetch/PutNear all need genuinely distinct object positions to track a mission by position alone). Draws sequentially, each excluding every position drawn so far in addition to exclude - n is a small static Python int in every current caller, so the unrolled loop is fine under jit (same pattern as jax.random.split being called with a static count elsewhere in this codebase).

Parameters:

Name	Type	Description	Default
key	Array	A random key.	required
grid	Array	A 2D grid of shape (height, width).	required
n	int	The number of distinct positions to generate.	required
exclude	Array	Position(s) to also exclude, shape (2,) or (k, 2). Defaults to jnp.asarray((-1, -1)).	asarray((-1, -1))

Returns:

Name	Type	Description
Array	Array	n mutually distinct positions of shape i32[n, 2].

random_position_far_from(key, grid, reference, min_distance=2, exclude=jnp.asarray((-1, -1)))

Generates one random position at Chebyshev distance >= min_distance from reference, also excluding exclude.

For PutNear-style tasks: quantified directly (500 seeds each), random_distinct_positions alone let 36% of Navix-PutNear-6x6-N2-v0 episodes spawn with the "move" object already within Chebyshev distance 1 of the "drop near" target - trivially "solved" with no real navigation needed, unlike real MiniGrid's PutNearEnv, which explicitly rejects that via reject_fn=near_obj (see PR #191 review's "New risks" section).

Parameters:

Name	Type	Description	Default
key	Array	A random key.	required
grid	Array	A 2D grid of shape (height, width).	required
reference	Array	The (row, col) position to stay away from.	required
min_distance	int	Minimum Chebyshev distance from reference. Defaults to 2 (i.e. not orthogonally/ diagonally adjacent, and not the same cell).	2
exclude	Array	Position(s) to also exclude, shape (2,) or (k, 2). Defaults to jnp.asarray((-1, -1)).	asarray((-1, -1))

Returns:

Name	Type	Description
Array	Array	A position of shape i32[2].

random_positions(key, grid, n=1, exclude=jnp.asarray((-1, -1)))

Generates n random positions in the grid, excluding the exclude position.

Parameters:

Name	Type	Description	Default
key	Array	A random key.	required
grid	Array	A 2D grid of shape (height, width).	required
n	int	The number of random positions to generate. Defaults to 1.	1
exclude	Array	The position to exclude. Defaults to jnp.asarray((-1, -1)).	asarray((-1, -1))

Returns:

Name	Type	Description
Array	Array	A batch of random positions of shape i32[n, 2].

room(height, width)

Creates an array representing a room of size height x width, including a set of walls around the room. The room is represented as a 2D grid of shape (height, width), including walls, with walls set to -1 and empty tiles set to 0.

Parameters:

Name	Type	Description	Default
height	int	The height of the room.	required
width	int	The width of the room.	required

Returns:

Name	Type	Description
Array	Array	A 2D grid of shape (height, width) representing a room.

room_grid(room_size, num_rows, num_cols)

Creates the base occupancy grid for a num_rows x num_cols layout of room_size x room_size rooms: the outer border plus every internal room-dividing wall, with no doors punched through yet (every room is fully sealed off from its neighbours) - callers open specific cells with open_wall to place doors. room_size, num_rows, num_cols are always static (environment-registration- time) values, never per-episode traced ones, so this is plain Python control flow, not vectorised.

Parameters:

Name	Type	Description	Default
room_size	int	The size of one room, including its own walls.	required
num_rows	int	Number of rooms stacked vertically.	required
num_cols	int	Number of rooms stacked horizontally.	required

Returns:

Name	Type	Description
Array	Array	A 2D grid of shape (height, width).

room_grid_dims(room_size, num_rows, num_cols)

The (height, width) of a room_grid layout - adjacent rooms share a single-cell-thick dividing wall (verified against MiniGrid's actual RoomGrid.__init__), so the grid is smaller than num_rows * room_size naively.

Parameters:

Name	Type	Description	Default
room_size	int	The size (both height and width) of one room, including its own walls. Must be >= 3 (at least a 1-cell interior).	required
num_rows	int	Number of rooms stacked vertically.	required
num_cols	int	Number of rooms stacked horizontally.	required

Returns:

Type	Description
Tuple[int, int]	Tuple[int, int]: The (height, width) of the full grid.

room_grid_door_position(key, room_size, i, j, side)

A random position along room (i, j)'s shared wall on the given side - where add_door would place a door in MiniGrid's actual RoomGrid. Verified against MiniGrid's own door placement: a uniform random offset along the wall, excluding the corners.

Parameters:

Name	Type	Description	Default
key	Array	A random key.	required
room_size	int	The size of one room, including its own walls.	required
i	int	The room's row index.	required
j	int	The room's column index.	required
side	int	The wall side, using navix's own direction convention (entities.Directions): 0=east, 1=south, 2=west, 3=north. Always a static (registration-time) value, never a per-episode traced one, so this is a plain if, not jax.lax.switch.	required

Returns:

Name	Type	Description
Array	Array	An i32[2] (row, col) position on the wall.

room_interior_bounds(room_size, i, j)

The inclusive (row_min, col_min, row_max, col_max) interior bounds of room (i, j) in a room_grid layout - excludes the room's own walls, the range a position can be sampled from.

Parameters:

Name	Type	Description	Default
room_size	int	The size of one room, including its own walls.	required
i	int	The room's row index.	required
j	int	The room's column index.	required

Returns:

Type	Description
Tuple[int, int, int, int]	Tuple[int, int, int, int]: The inclusive interior bounds.

room_mask(grid, room_size, i, j)

A boolean mask of grid's shape, True only within room (i, j)'s interior (its own walls excluded) - unlike unlock.py's mask_by_coordinates (which only tests "before a single row/col", fine for a 1-row-of-rooms layout corner-anchored at the origin), rooms in a general room_grid need all four bounds, since they aren't corner-anchored. Useful for scoping random_positions' sampling to one room.

Parameters:

Name	Type	Description	Default
grid	Array	A 2D grid of shape (height, width).	required
room_size	int	The size of one room, including its own walls.	required
i	int	The room's row index.	required
j	int	The room's column index.	required

Returns:

Name	Type	Description
Array	Array	A boolean mask of shape (height, width).

room_top_left(room_size, i, j)

The (row, col) of room (i, j)'s top-left corner (its own wall, not its interior) in a room_grid layout.

Parameters:

Name	Type	Description	Default
room_size	int	The size of one room, including its own walls.	required
i	int	The room's row index.	required
j	int	The room's column index.	required

Returns:

Type	Description
Tuple[int, int]	Tuple[int, int]: The (row, col) of the room's top-left corner.

rotate(direction, spin)

Changes a direction vectory by a given number of spins.

Parameters:

Name	Type	Description	Default
direction	Array	A direction vector of shape i32[] in the range [0, 3] representing the cardinal directions [east, south, west, north].	required
spin	int	The number of spins to apply.	required

Returns:

Name	Type	Description
Array	Array	A direction vector of shape i32[] in the range [0, 3] representing the cardinal directions [east, south, west, north].

rotate_tile(patch, num_times_90)

Rotates a patch of the grid by a given number of 90-degree rotations.

Parameters:

Name	Type	Description	Default
patch	Array	A patch of the grid.	required
num_times_90	int	The number of 90-degree rotations to apply.	required

Returns:

Name	Type	Description
Array	Array	A patch of the grid rotated by the given number of 90-degree rotations.

translate(position, direction, modulus=jnp.asarray(1))

Translates a point in a grid by a given direction and modulus.

Parameters:

Name	Type	Description	Default
position	Array	A 2D coordinate of shape i32[2] containing the (col, row) data.	required
direction	Array	A direction in the range [0, 1, 2, 3] representing the cardinal directions [east, south, west, north].	required
modulus	Array	The modulus of the translation. Defaults to jnp.asarray(1).	asarray(1)

Returns:

Name	Type	Description
Array	Array	A 2D coordinate of shape i32[2] containing the (col, row) data.

translate_forward(position, forward_direction, modulus)

Translates a point in a grid by a given forward direction and modulus.

Parameters:

Name	Type	Description	Default
position	Array	A 2D coordinate of shape i32[2] containing the (col, row) data.	required
forward_direction	Array	A direction in the range [0, 1, 2, 3] representing the cardinal directions [east, south, west, north].	required
modulus	Array	The modulus of the translation.	required

Returns:

Name	Type	Description
Array		A 2D coordinate of shape i32[2] containing the (col, row) data.

translate_left(position, forward_direction, modulus)

Translates a point in a grid by a given left direction and modulus.

Parameters:

Name	Type	Description	Default
position	Array	A 2D coordinate of shape i32[2] containing the (col, row) data.	required
forward_direction	Array	A direction in the range [0, 1, 2, 3] representing the cardinal directions [east, south, west, north].	required
modulus	Array	The modulus of the translation.	required

Returns:

Name	Type	Description
Array		A 2D coordinate of shape i32[2] containing the (col, row) data.

translate_right(position, forward_direction, modulus)

Translates a point in a grid by a given right direction and modulus.

Parameters:

Name	Type	Description	Default
position	Array	A 2D coordinate of shape i32[2] containing the (col, row) data.	required
forward_direction	Array	A direction in the range [0, 1, 2, 3] representing the cardinal directions [east, south, west, north].	required
modulus	Array	The modulus of the translation.	required

Returns:

Name	Type	Description
Array		A 2D coordinate of shape i32[2] containing the (col, row) data.

two_rooms(height, width, key)

Creates a 2D grid representing two rooms of size height x width, separated by a wall. The rooms are represented as a 2D grid of shape (height, width), including walls, with walls set to -1 and empty tiles set to 0.

Parameters:

Name	Type	Description	Default
height	int	The height of the rooms.	required
width	int	The width of the rooms.	required
key	Array	A random key, determining the position of the wall separating the rooms.	required

Returns:

Type	Description
Tuple[Array, Array]	Tuple[Array, Array]: A tuple containing the 2D grid representing the rooms and the column index of the wall separating the rooms.

vertical_wall(grid, row_idx, opening_col_idx=None)

Creates a vertical wall in the grid at the given row index, with an opening at the given column index.

Parameters:

Name	Type	Description	Default
grid	Array	A 2D grid of shape (height, width).	required
row_idx	int	The row index where the wall is placed.	required
opening_col_idx	Array	The column index where the opening is placed. Defaults to None.	None

Returns:

Name	Type	Description
Array	Array	A 2D grid of shape (height, width) with a vertical wall.

view_cone(transparency_map, origin, radius)

Computes the view cone of a given origin in a grid with a given radius. The view cone is a boolean map of transparent (1) and opaque (0) tiles, indicating whether a tile is visible from the origin or not.

Parameters:

Name	Type	Description	Default
transparency_map	Array	A boolean map of transparent (1) and opaque (0) tiles.	required
origin	Array	The origin of the view cone.	required
radius	int	The radius of the view cone.	required

Returns:

Name	Type	Description
Array	Array	The view cone of the given origin in the grid with the given radius.