Cache

The rendering cache and the patch <-> image plumbing behind rgb observations.

An rgb observation is the grid drawn as TILE_SIZE x TILE_SIZE sprites. The static parts (walls, floor, grid lines) never change during an episode, so they are rendered once into a flat list of tile patches, stored on State.cache, and only the cells holding moving entities are re-drawn each step.

RenderingCache

Bases: PyTreeNode

The pre-rendered background, carried on State.cache so rgb observations are cheap. Build one with RenderingCache.init; it is valid for any state on the same-shaped grid.

patches instance-attribute

u8[H * W + 1, TILE_SIZE, TILE_SIZE, 3] - one RGB tile per grid cell (row-major), plus a trailing all-zero "discard pile" tile that off-grid entities render into and that rgb slices off.

init(grid) classmethod

Renders grid's static background (walls / floor / grid lines) once into the flat patch list.

Parameters:

Name	Type	Description	Default
grid	Array	i32[H, W] base map (0 floor, -1 wall).	required

Returns:

Name	Type	Description
RenderingCache	RenderingCache	the cache for that grid shape.

flatten_patches(image, patch_size=(TILE_SIZE, TILE_SIZE))

Splits an image into a row-major list of fixed-size tiles - the inverse of unflatten_patches.

Parameters:

Name	Type	Description	Default
image	Array	(H, W, C), with H/W multiples of patch_size.	required
patch_size	tuple[int, int]	tile (height, width); defaults to (TILE_SIZE, TILE_SIZE).	(TILE_SIZE, TILE_SIZE)

Returns:

Name	Type	Description
Array	Array	(H//ph * W//pw, ph, pw, C).

render_background(grid, sprites_registry=SPRITES_REGISTRY)

Draws the static layer of an rgb frame: a wall sprite on every -1 cell, a floor sprite (with MiniGrid-matched grid lines) on every 0 cell. No entities.

Parameters:

Name	Type	Description	Default
grid	Array	i32[H, W] base map.	required
sprites_registry	dict	sprite lookup; defaults to the global SPRITES_REGISTRY.	SPRITES_REGISTRY

Returns:

Name	Type	Description
Array	Array	u8[H * TILE_SIZE, W * TILE_SIZE, 3].

tile_grid(grid, tile)

Tiles a grid (H, W) with equal tiles tiles (w, h, 3) to get a final array of shape (H * h, W * w, 3) and dtype jnp.uint8

unflatten_patches(patches, image_size)

Reassembles a row-major list of tiles into an image - the inverse of flatten_patches.

Parameters:

Name	Type	Description	Default
patches	Array	(n_tiles, ph, pw, C).	required
image_size	tuple[int, int]	the target (H, W); H * W == n_tiles * ph * pw.	required

Returns:

Name	Type	Description
Array	Array	(H, W, C).