Add Diffusion Miner teacher checkpoint, Tile VAE, live infer, and gameplay GIF

b4b7ca9 verified about 1 month ago

24.5 kB

from __future__ import annotations

import numpy as np

from livediffusion.config import VideoConfig

AIR = 0
DIRT = 1
STONE = 2
SAND = 3
GRAVEL = 4
WATER = 5
ORE = 6
BEDROCK = 7

TILE_COUNT = 8

# Flat, high-contrast palette. No thin lines anywhere, so an 8x VAE keeps the edges.
TILE_COLORS = np.array(
    [
        (16, 18, 30),
        (126, 84, 48),
        (108, 112, 122),
        (214, 182, 96),
        (52, 54, 62),
        (44, 104, 216),
        (72, 240, 236),
        (74, 62, 102),
    ],
    dtype=np.int16,
)
# Every tile hue is separable from every other one and from the player, so a
# reconstruction error is never ambiguous about which material it was.
_DIG_FLASH_COLOR = (255, 84, 240)
_PLAYER_HEAD_COLOR = (248, 252, 255)
_PLAYER_BODY_COLOR = (232, 96, 72)

# Water is passable; everything else with a body blocks the player.
TILE_SOLID = np.array([0, 1, 1, 1, 1, 0, 1, 1], dtype=bool)
# Sand and gravel fall; bedrock never moves; ore and stone stay until mined.
TILE_FALLS = np.array([0, 0, 0, 1, 1, 0, 0, 0], dtype=bool)
TILE_DIGGABLE = np.array([0, 1, 1, 1, 1, 0, 1, 0], dtype=bool)
# Per-tile shade jitter is a pure function of position, so it is memorizable structure
# rather than noise the model cannot predict.
_SHADE_RANGE = 10
_PLAYER_EYE_COLOR = (24, 28, 40)


def paint_player(
    frame: np.ndarray,
    left: int,
    top: int,
    width: int,
    height: int,
    facing: int = 1,
) -> None:
    """Draw the miner onto an RGB frame. Used by the sim and by infer click-spawn."""
    right = min(frame.shape[1], left + width)
    bottom = min(frame.shape[0], top + height)
    if right <= left or bottom <= top:
        return
    body_top = min(bottom, top + height // 2)
    frame[max(0, top) : body_top, max(0, left) : right] = _PLAYER_HEAD_COLOR
    frame[body_top:bottom, max(0, left) : right] = _PLAYER_BODY_COLOR
    eye = left + (width - width // 3 if facing > 0 else 0)
    eye_right = min(frame.shape[1], eye + max(1, width // 3))
    eye_bottom = min(bottom, max(0, top) + max(1, height // 6))
    frame[max(0, top) : eye_bottom, max(0, eye) : eye_right] = _PLAYER_EYE_COLOR


def _shift_down(mask: np.ndarray) -> np.ndarray:
    shifted = np.zeros_like(mask)
    shifted[1:] = mask[:-1]
    return shifted


def _shift_horizontal(mask: np.ndarray, step: int) -> np.ndarray:
    shifted = np.zeros_like(mask)
    if step > 0:
        shifted[:, step:] = mask[:, :-step]
    else:
        shifted[:, :step] = mask[:, -step:]
    return shifted


def _smooth(values: np.ndarray, passes: int) -> np.ndarray:
    for _ in range(passes):
        values = (values + np.roll(values, 1) + np.roll(values, -1)) / 3.0
    return values


class MiningSimulation:
    """Single-screen tile sandbox with irreversible digging, falling sand, and water."""

    # [left, right, jump, dig_left, dig_right, dig_up, dig_down]
    ACTION_DIM = 7

    def __init__(self, video_config: VideoConfig, seed: int) -> None:
        self.video_config = video_config
        self.rng = np.random.default_rng(seed)
        self.dt = 1.0 / video_config.fps
        self.tile = max(4, min(video_config.width, video_config.height) // 32)
        self.grid_width = video_config.width // self.tile
        self.grid_height = video_config.height // self.tile
        self.gravity = self.tile * 42.0
        self.move_speed = self.tile * 7.5
        self.jump_speed = self.tile * 17.0
        self.terminal_speed = self.tile * 30.0
        self.swim_drag = 0.45
        self.dig_cooldown = 0.16
        self.player_width = float(self.tile)
        self.player_height = float(self.tile * 2)
        self.frame_index = 0
        self.ore_collected = 0
        self.deaths = 0
        self.grid = np.zeros((self.grid_height, self.grid_width), dtype=np.uint8)
        self.shade = self._make_shade()
        self.position = np.zeros(2, dtype=np.float64)
        self.velocity = np.zeros(2, dtype=np.float64)
        self.facing = 1
        self.dig_timer = 0.0
        self.dig_target: tuple[tuple[int, int], ...] = ()
        self.dig_flash = 0.0
        self.drown_timer = 0.0
        self.pending_reward = 0.0
        self.pending_terminal = 0.0
        self._action = (0,) * self.ACTION_DIM
        self._hold_frames = 0
        self._stuck_frames = 0
        self._last_cell = (-1, -1)
        self._reset_round()
        for _ in range(6):
            self.step()

    @classmethod
    def from_seed(
        cls,
        video_config: VideoConfig,
        seed: int,
        *,
        at_start: bool = False,
    ) -> MiningSimulation:
        simulation = cls(video_config, seed)
        if at_start:
            simulation.reset_to_start()
        return simulation

    def reset_to_start(self) -> None:
        self.ore_collected = 0
        self.deaths = 0
        self.frame_index = 0
        self._reset_round()

    def _make_shade(self) -> np.ndarray:
        rows = np.arange(self.grid_height)[:, None]
        columns = np.arange(self.grid_width)[None, :]
        mixed = (rows * 73_856_093) ^ (columns * 19_349_663)
        return ((mixed % (2 * _SHADE_RANGE + 1)) - _SHADE_RANGE).astype(np.int16)

    def _generate_world(self) -> None:
        height = self.grid_height
        width = self.grid_width
        grid = np.full((height, width), STONE, dtype=np.uint8)

        walk = self.rng.normal(0.0, 1.0, size=width)
        surface = _smooth(walk, 6)
        surface = surface / max(float(np.abs(surface).max()), 1e-6)
        base = height * 0.28
        amplitude = height * 0.08
        surface_rows = np.clip(
            np.round(base + surface * amplitude).astype(int),
            2,
            height - 8,
        )

        rows = np.arange(height)[:, None]
        grid[rows < surface_rows[None, :]] = AIR
        dirt_depth = self.rng.integers(3, 6, size=width)
        dirt_limit = surface_rows + dirt_depth
        grid[(rows >= surface_rows[None, :]) & (rows < dirt_limit[None, :])] = DIRT

        self._carve_caves(grid, surface_rows)
        self._scatter(grid, ORE, count=self.rng.integers(9, 15), radius=1, in_stone=True)
        self._scatter(grid, SAND, count=self.rng.integers(3, 7), radius=2, in_stone=False)
        self._scatter(grid, GRAVEL, count=self.rng.integers(2, 5), radius=1, in_stone=True)
        self._pour_water(grid, surface_rows)

        grid[0, :] = BEDROCK
        grid[-1, :] = BEDROCK
        grid[:, 0] = BEDROCK
        grid[:, -1] = BEDROCK
        # Keep the spawn column open so the miner never starts entombed.
        spawn_column = int(self.rng.integers(3, width - 3))
        spawn_row = int(surface_rows[spawn_column])
        grid[max(1, spawn_row - 3) : spawn_row, spawn_column] = AIR
        self.grid = grid
        self._spawn_column = spawn_column
        self._spawn_row = spawn_row

    def _carve_caves(self, grid: np.ndarray, surface_rows: np.ndarray) -> None:
        cave_count = int(self.rng.integers(3, 7))
        for _ in range(cave_count):
            column = float(self.rng.integers(2, self.grid_width - 2))
            row = float(self.rng.integers(int(surface_rows.max()) + 3, self.grid_height - 3))
            length = int(self.rng.integers(6, 18))
            floor = int(surface_rows.max()) + 2
            for _ in range(length):
                radius = float(self.rng.uniform(1.0, 2.4))
                self._stamp(grid, row, column, radius, AIR, in_stone=False)
                angle = float(self.rng.uniform(0.0, np.pi * 2.0))
                column = float(np.clip(column + np.cos(angle) * 1.6, 2, self.grid_width - 3))
                row = float(np.clip(row + np.sin(angle) * 1.2, floor, self.grid_height - 3))

    def _stamp(
        self,
        grid: np.ndarray,
        row: float,
        column: float,
        radius: float,
        tile: int,
        in_stone: bool,
    ) -> None:
        top = max(1, int(row - radius))
        bottom = min(self.grid_height - 1, int(row + radius) + 1)
        left = max(1, int(column - radius))
        right = min(self.grid_width - 1, int(column + radius) + 1)
        if top >= bottom or left >= right:
            return
        rows = np.arange(top, bottom)[:, None]
        columns = np.arange(left, right)[None, :]
        inside = (rows - row) ** 2 + (columns - column) ** 2 <= radius**2
        region = grid[top:bottom, left:right]
        if in_stone:
            inside &= region == STONE
        region[inside] = tile

    def _scatter(
        self,
        grid: np.ndarray,
        tile: int,
        count: int,
        radius: int,
        in_stone: bool,
    ) -> None:
        for _ in range(int(count)):
            row = float(self.rng.integers(2, self.grid_height - 2))
            column = float(self.rng.integers(2, self.grid_width - 2))
            self._stamp(grid, row, column, float(radius), tile, in_stone)

    def _pour_water(self, grid: np.ndarray, surface_rows: np.ndarray) -> None:
        pool_count = int(self.rng.integers(1, 4))
        for _ in range(pool_count):
            column = int(self.rng.integers(2, self.grid_width - 2))
            row = int(self.rng.integers(int(surface_rows.max()) + 2, self.grid_height - 3))
            span = int(self.rng.integers(2, 5))
            left = max(1, column - span)
            right = min(self.grid_width - 1, column + span)
            region = grid[row : row + 2, left:right]
            region[region == AIR] = WATER

    def _reset_round(self) -> None:
        self._generate_world()
        self._respawn_player()

    def _respawn_player(self) -> None:
        # The world is never regenerated on death: tunnels the miner dug stay dug,
        # which is the persistence signal the model is meant to learn.
        column = self._spawn_column
        row = self._spawn_row - 2
        for candidate in range(1, self.grid_height - 2):
            if not TILE_SOLID[self.grid[candidate, column]] and not TILE_SOLID[
                self.grid[candidate + 1, column]
            ]:
                row = candidate
                break
        else:
            self.grid[row : row + 2, column] = AIR
        self.position[0] = float(column * self.tile)
        self.position[1] = float(row * self.tile)
        self.velocity[:] = 0.0
        self.facing = 1
        self.dig_timer = 0.0
        self.dig_target = ()
        self.dig_flash = 0.0
        self.drown_timer = 0.0
        self._action = (0,) * self.ACTION_DIM
        self._hold_frames = 0

    def _tile_at(self, row: int, column: int) -> int:
        if row < 0 or row >= self.grid_height or column < 0 or column >= self.grid_width:
            return BEDROCK
        return int(self.grid[row, column])

    def _solid_at(self, row: int, column: int) -> bool:
        return bool(TILE_SOLID[self._tile_at(row, column)])

    def _box_overlaps_solid(self, x: float, y: float) -> bool:
        left = int(np.floor(x / self.tile))
        right = int(np.floor((x + self.player_width - 1e-6) / self.tile))
        top = int(np.floor(y / self.tile))
        bottom = int(np.floor((y + self.player_height - 1e-6) / self.tile))
        for row in range(top, bottom + 1):
            for column in range(left, right + 1):
                if self._solid_at(row, column):
                    return True
        return False

    def _player_tiles(self) -> tuple[int, int, int]:
        column = int((self.position[0] + self.player_width * 0.5) // self.tile)
        head = int(self.position[1] // self.tile)
        feet = int((self.position[1] + self.player_height - 1e-6) // self.tile)
        return column, head, feet

    def _footprint(self) -> tuple[int, int, int, int]:
        left = int(self.position[0] // self.tile)
        right = int((self.position[0] + self.player_width - 1e-6) // self.tile)
        head = int(self.position[1] // self.tile)
        feet = int((self.position[1] + self.player_height - 1e-6) // self.tile)
        return left, right, head, feet

    def _on_ground(self) -> bool:
        return self._box_overlaps_solid(self.position[0], self.position[1] + 1.0)

    def _in_water(self) -> bool:
        column, head, feet = self._player_tiles()
        return self._tile_at(head, column) == WATER or self._tile_at(feet, column) == WATER

    def _step_physics(self, action: tuple[int, ...]) -> None:
        left, right, jump = action[0], action[1], action[2]
        submerged = self._in_water()
        direction = float(right) - float(left)
        if direction != 0.0:
            self.facing = 1 if direction > 0 else -1
        speed = self.move_speed * (0.6 if submerged else 1.0)
        self.velocity[0] = direction * speed
        grounded = self._on_ground()
        if jump and (grounded or submerged):
            self.velocity[1] = -self.jump_speed * (0.62 if submerged else 1.0)
        gravity = self.gravity * (0.35 if submerged else 1.0)
        self.velocity[1] = min(self.velocity[1] + gravity * self.dt, self.terminal_speed)
        if submerged:
            self.velocity[1] *= 1.0 - self.swim_drag * self.dt * self.video_config.fps * 0.1

        step_x = float(self.velocity[0] * self.dt)
        candidate = self.position[0] + step_x
        if self._box_overlaps_solid(candidate, self.position[1]):
            # Step up one tile so single-block ledges do not need a jump.
            lifted = self.position[1] - self.tile
            climbable = step_x != 0.0 and self._on_ground()
            if climbable and not self._box_overlaps_solid(candidate, lifted):
                self.position[0] = candidate
                self.position[1] = lifted
            else:
                self.velocity[0] = 0.0
        else:
            self.position[0] = candidate

        step_y = float(self.velocity[1] * self.dt)
        candidate = self.position[1] + step_y
        if self._box_overlaps_solid(self.position[0], candidate):
            while abs(step_y) > 0.5:
                step_y *= 0.5
                probe = self.position[1] + step_y
                if not self._box_overlaps_solid(self.position[0], probe):
                    self.position[1] = probe
            self.velocity[1] = 0.0
        else:
            self.position[1] = candidate

        self.position[0] = float(
            np.clip(self.position[0], self.tile, self.video_config.width - self.tile * 2)
        )
        self.position[1] = float(
            np.clip(self.position[1], self.tile, self.video_config.height - self.tile * 3)
        )

    def _dig(self, action: tuple[int, ...]) -> None:
        self.dig_flash = max(0.0, self.dig_flash - self.dt)
        self.dig_timer = max(0.0, self.dig_timer - self.dt)
        dig_left, dig_right, dig_up, dig_down = action[3], action[4], action[5], action[6]
        if not (dig_left or dig_right or dig_up or dig_down):
            self.dig_target = ()
            return
        left, right, head, feet = self._footprint()
        # A dig clears the player's whole cross-section, so the shaft it opens is
        # always wide enough to move through.
        if dig_left:
            self.facing = -1
            target = tuple((row, left - 1) for row in range(head, feet + 1))
        elif dig_right:
            self.facing = 1
            target = tuple((row, right + 1) for row in range(head, feet + 1))
        elif dig_up:
            target = tuple((head - 1, column) for column in range(left, right + 1))
        else:
            target = tuple((feet + 1, column) for column in range(left, right + 1))
        self.dig_target = target
        if self.dig_timer > 0.0:
            return
        removed = False
        for row, column in target:
            tile = self._tile_at(row, column)
            if not TILE_DIGGABLE[tile]:
                continue
            self.grid[row, column] = AIR
            removed = True
            if tile == ORE:
                self.ore_collected += 1
                self.pending_reward += 1.0
        if removed:
            self.dig_timer = self.dig_cooldown
            self.dig_flash = self.dig_cooldown

    def _step_tiles(self) -> None:
        grid = self.grid
        falls = TILE_FALLS[grid]
        below = np.zeros_like(grid)
        below[:-1] = grid[1:]
        below[-1] = BEDROCK
        open_below = (below == AIR) | (below == WATER)
        moving = falls & open_below
        moving[-1] = False
        if moving.any():
            destination = _shift_down(moving)
            carried = grid[moving]
            displaced = grid[destination]
            grid[destination] = carried
            grid[moving] = displaced

        # Sand slides off peaks so it forms cones instead of towers.
        side = 1 if self.frame_index % 2 == 0 else -1
        sand = grid == SAND
        below[:-1] = grid[1:]
        below[-1] = BEDROCK
        blocked = sand & ~((below == AIR) | (below == WATER))
        diagonal = np.zeros_like(grid)
        if side > 0:
            diagonal[:-1, :-1] = grid[1:, 1:]
        else:
            diagonal[:-1, 1:] = grid[1:, :-1]
        slide = blocked & (diagonal == AIR)
        slide[-1] = False
        if slide.any():
            destination = _shift_down(_shift_horizontal(slide, side))
            grid[destination] = grid[slide]
            grid[slide] = AIR

        water = grid == WATER
        below[:-1] = grid[1:]
        below[-1] = BEDROCK
        settled = water & (below != AIR)
        neighbour = np.zeros_like(grid)
        if side > 0:
            neighbour[:, :-1] = grid[:, 1:]
        else:
            neighbour[:, 1:] = grid[:, :-1]
        spread = settled & (neighbour == AIR)
        if spread.any():
            destination = _shift_horizontal(spread, side)
            grid[destination] = WATER
            grid[spread] = AIR

    def _check_fatal(self) -> None:
        buried = False
        if self._box_overlaps_solid(self.position[0], self.position[1]):
            # Material settling onto the miner shoves it upward; only a fully
            # sealed pocket actually kills.
            buried = True
            for offset in range(1, self.tile * 2 + 1):
                lifted = self.position[1] - offset
                if lifted < self.tile:
                    break
                if not self._box_overlaps_solid(self.position[0], lifted):
                    self.position[1] = float(lifted)
                    self.velocity[1] = 0.0
                    buried = False
                    break
        column, head, _ = self._player_tiles()
        if self._tile_at(head, column) == WATER:
            self.drown_timer += self.dt
        else:
            self.drown_timer = max(0.0, self.drown_timer - self.dt * 0.5)
        if buried or self.drown_timer > 2.5:
            self.deaths += 1
            self.pending_terminal = 1.0
            self.drown_timer = 0.0
            self._respawn_player()

    def scripted_action(self) -> tuple[int, ...]:
        """Next training-policy action. Does not step the world."""
        return self._scripted_action()

    def _scripted_action(self) -> tuple[int, ...]:
        column, _, feet = self._player_tiles()
        cell = (feet, column)
        if cell == self._last_cell:
            self._stuck_frames += 1
        else:
            self._stuck_frames = 0
            self._last_cell = cell
        if self._hold_frames > 0:
            self._hold_frames -= 1
            return self._action
        if self._stuck_frames > 18:
            # Break out of dead ends instead of grinding at an unreachable target.
            action = [0] * self.ACTION_DIM
            action[int(self.rng.integers(0, self.ACTION_DIM))] = 1
            action[2] = 1
            self._action = tuple(action)
            self._hold_frames = 4
            self._stuck_frames = 0
            return self._action
        action = self._plan_action()
        self._action = action
        self._hold_frames = int(self.rng.integers(0, 3))
        return action

    def _plan_action(self) -> tuple[int, ...]:
        if self.rng.random() < 0.06:
            choice = int(self.rng.integers(0, self.ACTION_DIM))
            action = [0] * self.ACTION_DIM
            action[choice] = 1
            return tuple(action)

        column, head, feet = self._player_tiles()
        action = [0] * self.ACTION_DIM

        # Escape falling material before anything else.
        if self._tile_at(head - 1, column) in {SAND, GRAVEL}:
            action[1 if self.facing > 0 else 0] = 1
            action[2] = 1
            return tuple(action)
        if self.drown_timer > 0.6:
            action[2] = 1
            action[5] = 1
            return tuple(action)

        ore = np.argwhere(self.grid == ORE)
        if ore.size == 0:
            action[6] = 1
            return tuple(action)
        distances = np.abs(ore[:, 0] - feet) + np.abs(ore[:, 1] - column)
        target_row, target_column = ore[int(np.argmin(distances))]
        delta_column = int(target_column) - column
        delta_row = int(target_row) - feet

        if abs(delta_column) <= 1 and abs(delta_row) <= 1:
            if delta_column < 0:
                action[3] = 1
            elif delta_column > 0:
                action[4] = 1
            elif delta_row < 0:
                action[5] = 1
            else:
                action[6] = 1
            return tuple(action)

        if delta_column != 0:
            forward = 1 if delta_column > 0 else 0
            action[forward] = 1
            ahead = column + (1 if delta_column > 0 else -1)
            if self._solid_at(feet, ahead):
                # Bedrock never yields, so climb over it instead of digging forever.
                if TILE_DIGGABLE[self._tile_at(feet, ahead)]:
                    action[4 if delta_column > 0 else 3] = 1
                else:
                    action[2] = 1
            elif self._solid_at(head, ahead):
                action[2] = 1
            return tuple(action)

        if delta_row > 0:
            action[6] = 1
        else:
            action[2] = 1
            action[5] = 1
        return tuple(action)

    def step(self, action: tuple[int, ...] | None = None) -> None:
        resolved = self._scripted_action() if action is None else tuple(action)
        if len(resolved) != self.ACTION_DIM:
            raise ValueError(f"Mining actions must have {self.ACTION_DIM} entries.")
        self.pending_reward = 0.0
        self.pending_terminal = 0.0
        self._dig(resolved)
        self._step_physics(resolved)
        self._step_tiles()
        self._check_fatal()
        self.frame_index += 1

    def render(self) -> np.ndarray:
        colors = TILE_COLORS[self.grid] + self.shade[:, :, None]
        if self.dig_flash > 0.0:
            for row, column in self.dig_target:
                if 0 <= row < self.grid_height and 0 <= column < self.grid_width:
                    colors[row, column] = _DIG_FLASH_COLOR
        frame = np.clip(colors, 0, 255).astype(np.uint8)
        frame = frame.repeat(self.tile, axis=0).repeat(self.tile, axis=1)

        paint_player(
            frame,
            round(float(self.position[0])),
            round(float(self.position[1])),
            int(self.player_width),
            int(self.player_height),
            self.facing,
        )

        # Ore pips give a fontless score readout that stays stable at training size.
        for index in range(min(self.ore_collected, 8)):
            x = self.tile + index * (self.tile + 2)
            pip = max(2, self.tile // 2)
            frame[2 : 2 + pip, x : x + pip] = TILE_COLORS[ORE]
        return frame

    def next_frame(self) -> np.ndarray:
        frame = self.render()
        self.step()
        return frame

    def sample_frames(self, count: int) -> np.ndarray:
        return np.stack([self.next_frame() for _ in range(count)])

    def sample_gameplay(
        self,
        count: int,
    ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
        frames: list[np.ndarray] = []
        actions: list[tuple[int, ...]] = []
        rewards: list[float] = []
        terminals: list[float] = []
        for _ in range(count):
            frames.append(self.render())
            action = self.scripted_action()
            actions.append(action)
            self.step(action)
            rewards.append(self.pending_reward)
            terminals.append(self.pending_terminal)
        return (
            np.stack(frames),
            np.asarray(actions, dtype=np.float32),
            np.asarray(rewards, dtype=np.float32),
            np.asarray(terminals, dtype=np.float32),
        )

	from __future__ import annotations

	import numpy as np

	from livediffusion.config import VideoConfig

	AIR = 0
	DIRT = 1
	STONE = 2
	SAND = 3
	GRAVEL = 4
	WATER = 5
	ORE = 6
	BEDROCK = 7

	TILE_COUNT = 8

	# Flat, high-contrast palette. No thin lines anywhere, so an 8x VAE keeps the edges.
	TILE_COLORS = np.array(
	[
	(16, 18, 30),
	(126, 84, 48),
	(108, 112, 122),
	(214, 182, 96),
	(52, 54, 62),
	(44, 104, 216),
	(72, 240, 236),
	(74, 62, 102),
	],
	dtype=np.int16,
	)
	# Every tile hue is separable from every other one and from the player, so a
	# reconstruction error is never ambiguous about which material it was.
	_DIG_FLASH_COLOR = (255, 84, 240)
	_PLAYER_HEAD_COLOR = (248, 252, 255)
	_PLAYER_BODY_COLOR = (232, 96, 72)

	# Water is passable; everything else with a body blocks the player.
	TILE_SOLID = np.array([0, 1, 1, 1, 1, 0, 1, 1], dtype=bool)
	# Sand and gravel fall; bedrock never moves; ore and stone stay until mined.
	TILE_FALLS = np.array([0, 0, 0, 1, 1, 0, 0, 0], dtype=bool)
	TILE_DIGGABLE = np.array([0, 1, 1, 1, 1, 0, 1, 0], dtype=bool)
	# Per-tile shade jitter is a pure function of position, so it is memorizable structure
	# rather than noise the model cannot predict.
	_SHADE_RANGE = 10
	_PLAYER_EYE_COLOR = (24, 28, 40)


	def paint_player(
	frame: np.ndarray,
	left: int,
	top: int,
	width: int,
	height: int,
	facing: int = 1,
	) -> None:
	"""Draw the miner onto an RGB frame. Used by the sim and by infer click-spawn."""
	right = min(frame.shape[1], left + width)
	bottom = min(frame.shape[0], top + height)
	if right <= left or bottom <= top:
	return
	body_top = min(bottom, top + height // 2)
	frame[max(0, top) : body_top, max(0, left) : right] = _PLAYER_HEAD_COLOR
	frame[body_top:bottom, max(0, left) : right] = _PLAYER_BODY_COLOR
	eye = left + (width - width // 3 if facing > 0 else 0)
	eye_right = min(frame.shape[1], eye + max(1, width // 3))
	eye_bottom = min(bottom, max(0, top) + max(1, height // 6))
	frame[max(0, top) : eye_bottom, max(0, eye) : eye_right] = _PLAYER_EYE_COLOR


	def _shift_down(mask: np.ndarray) -> np.ndarray:
	shifted = np.zeros_like(mask)
	shifted[1:] = mask[:-1]
	return shifted


	def _shift_horizontal(mask: np.ndarray, step: int) -> np.ndarray:
	shifted = np.zeros_like(mask)
	if step > 0:
	shifted[:, step:] = mask[:, :-step]
	else:
	shifted[:, :step] = mask[:, -step:]
	return shifted


	def _smooth(values: np.ndarray, passes: int) -> np.ndarray:
	for _ in range(passes):
	values = (values + np.roll(values, 1) + np.roll(values, -1)) / 3.0
	return values


	class MiningSimulation:
	"""Single-screen tile sandbox with irreversible digging, falling sand, and water."""

	# [left, right, jump, dig_left, dig_right, dig_up, dig_down]
	ACTION_DIM = 7

	def __init__(self, video_config: VideoConfig, seed: int) -> None:
	self.video_config = video_config
	self.rng = np.random.default_rng(seed)
	self.dt = 1.0 / video_config.fps
	self.tile = max(4, min(video_config.width, video_config.height) // 32)
	self.grid_width = video_config.width // self.tile
	self.grid_height = video_config.height // self.tile
	self.gravity = self.tile * 42.0
	self.move_speed = self.tile * 7.5
	self.jump_speed = self.tile * 17.0
	self.terminal_speed = self.tile * 30.0
	self.swim_drag = 0.45
	self.dig_cooldown = 0.16
	self.player_width = float(self.tile)
	self.player_height = float(self.tile * 2)
	self.frame_index = 0
	self.ore_collected = 0
	self.deaths = 0
	self.grid = np.zeros((self.grid_height, self.grid_width), dtype=np.uint8)
	self.shade = self._make_shade()
	self.position = np.zeros(2, dtype=np.float64)
	self.velocity = np.zeros(2, dtype=np.float64)
	self.facing = 1
	self.dig_timer = 0.0
	self.dig_target: tuple[tuple[int, int], ...] = ()
	self.dig_flash = 0.0
	self.drown_timer = 0.0
	self.pending_reward = 0.0
	self.pending_terminal = 0.0
	self._action = (0,) * self.ACTION_DIM
	self._hold_frames = 0
	self._stuck_frames = 0
	self._last_cell = (-1, -1)
	self._reset_round()
	for _ in range(6):
	self.step()

	@classmethod
	def from_seed(
	cls,
	video_config: VideoConfig,
	seed: int,
	*,
	at_start: bool = False,
	) -> MiningSimulation:
	simulation = cls(video_config, seed)
	if at_start:
	simulation.reset_to_start()
	return simulation

	def reset_to_start(self) -> None:
	self.ore_collected = 0
	self.deaths = 0
	self.frame_index = 0
	self._reset_round()

	def _make_shade(self) -> np.ndarray:
	rows = np.arange(self.grid_height)[:, None]
	columns = np.arange(self.grid_width)[None, :]
	mixed = (rows * 73_856_093) ^ (columns * 19_349_663)
	return ((mixed % (2 * _SHADE_RANGE + 1)) - _SHADE_RANGE).astype(np.int16)

	def _generate_world(self) -> None:
	height = self.grid_height
	width = self.grid_width
	grid = np.full((height, width), STONE, dtype=np.uint8)

	walk = self.rng.normal(0.0, 1.0, size=width)
	surface = _smooth(walk, 6)
	surface = surface / max(float(np.abs(surface).max()), 1e-6)
	base = height * 0.28
	amplitude = height * 0.08
	surface_rows = np.clip(
	np.round(base + surface * amplitude).astype(int),
	2,
	height - 8,
	)

	rows = np.arange(height)[:, None]
	grid[rows < surface_rows[None, :]] = AIR
	dirt_depth = self.rng.integers(3, 6, size=width)
	dirt_limit = surface_rows + dirt_depth
	grid[(rows >= surface_rows[None, :]) & (rows < dirt_limit[None, :])] = DIRT

	self._carve_caves(grid, surface_rows)
	self._scatter(grid, ORE, count=self.rng.integers(9, 15), radius=1, in_stone=True)
	self._scatter(grid, SAND, count=self.rng.integers(3, 7), radius=2, in_stone=False)
	self._scatter(grid, GRAVEL, count=self.rng.integers(2, 5), radius=1, in_stone=True)
	self._pour_water(grid, surface_rows)

	grid[0, :] = BEDROCK
	grid[-1, :] = BEDROCK
	grid[:, 0] = BEDROCK
	grid[:, -1] = BEDROCK
	# Keep the spawn column open so the miner never starts entombed.
	spawn_column = int(self.rng.integers(3, width - 3))
	spawn_row = int(surface_rows[spawn_column])
	grid[max(1, spawn_row - 3) : spawn_row, spawn_column] = AIR
	self.grid = grid
	self._spawn_column = spawn_column
	self._spawn_row = spawn_row

	def _carve_caves(self, grid: np.ndarray, surface_rows: np.ndarray) -> None:
	cave_count = int(self.rng.integers(3, 7))
	for _ in range(cave_count):
	column = float(self.rng.integers(2, self.grid_width - 2))
	row = float(self.rng.integers(int(surface_rows.max()) + 3, self.grid_height - 3))
	length = int(self.rng.integers(6, 18))
	floor = int(surface_rows.max()) + 2
	for _ in range(length):
	radius = float(self.rng.uniform(1.0, 2.4))
	self._stamp(grid, row, column, radius, AIR, in_stone=False)
	angle = float(self.rng.uniform(0.0, np.pi * 2.0))
	column = float(np.clip(column + np.cos(angle) * 1.6, 2, self.grid_width - 3))
	row = float(np.clip(row + np.sin(angle) * 1.2, floor, self.grid_height - 3))

	def _stamp(
	self,
	grid: np.ndarray,
	row: float,
	column: float,
	radius: float,
	tile: int,
	in_stone: bool,
	) -> None:
	top = max(1, int(row - radius))
	bottom = min(self.grid_height - 1, int(row + radius) + 1)
	left = max(1, int(column - radius))
	right = min(self.grid_width - 1, int(column + radius) + 1)
	if top >= bottom or left >= right:
	return
	rows = np.arange(top, bottom)[:, None]
	columns = np.arange(left, right)[None, :]
	inside = (rows - row) 2 + (columns - column) 2 <= radius**2
	region = grid[top:bottom, left:right]
	if in_stone:
	inside &= region == STONE
	region[inside] = tile

	def _scatter(
	self,
	grid: np.ndarray,
	tile: int,
	count: int,
	radius: int,
	in_stone: bool,
	) -> None:
	for _ in range(int(count)):
	row = float(self.rng.integers(2, self.grid_height - 2))
	column = float(self.rng.integers(2, self.grid_width - 2))
	self._stamp(grid, row, column, float(radius), tile, in_stone)

	def _pour_water(self, grid: np.ndarray, surface_rows: np.ndarray) -> None:
	pool_count = int(self.rng.integers(1, 4))
	for _ in range(pool_count):
	column = int(self.rng.integers(2, self.grid_width - 2))
	row = int(self.rng.integers(int(surface_rows.max()) + 2, self.grid_height - 3))
	span = int(self.rng.integers(2, 5))
	left = max(1, column - span)
	right = min(self.grid_width - 1, column + span)
	region = grid[row : row + 2, left:right]
	region[region == AIR] = WATER

	def _reset_round(self) -> None:
	self._generate_world()
	self._respawn_player()

	def _respawn_player(self) -> None:
	# The world is never regenerated on death: tunnels the miner dug stay dug,
	# which is the persistence signal the model is meant to learn.
	column = self._spawn_column
	row = self._spawn_row - 2
	for candidate in range(1, self.grid_height - 2):
	if not TILE_SOLID[self.grid[candidate, column]] and not TILE_SOLID[
	self.grid[candidate + 1, column]
	]:
	row = candidate
	break
	else:
	self.grid[row : row + 2, column] = AIR
	self.position[0] = float(column * self.tile)
	self.position[1] = float(row * self.tile)
	self.velocity[:] = 0.0
	self.facing = 1
	self.dig_timer = 0.0
	self.dig_target = ()
	self.dig_flash = 0.0
	self.drown_timer = 0.0
	self._action = (0,) * self.ACTION_DIM
	self._hold_frames = 0

	def _tile_at(self, row: int, column: int) -> int:
	if row < 0 or row >= self.grid_height or column < 0 or column >= self.grid_width:
	return BEDROCK
	return int(self.grid[row, column])

	def _solid_at(self, row: int, column: int) -> bool:
	return bool(TILE_SOLID[self._tile_at(row, column)])

	def _box_overlaps_solid(self, x: float, y: float) -> bool:
	left = int(np.floor(x / self.tile))
	right = int(np.floor((x + self.player_width - 1e-6) / self.tile))
	top = int(np.floor(y / self.tile))
	bottom = int(np.floor((y + self.player_height - 1e-6) / self.tile))
	for row in range(top, bottom + 1):
	for column in range(left, right + 1):
	if self._solid_at(row, column):
	return True
	return False

	def _player_tiles(self) -> tuple[int, int, int]:
	column = int((self.position[0] + self.player_width * 0.5) // self.tile)
	head = int(self.position[1] // self.tile)
	feet = int((self.position[1] + self.player_height - 1e-6) // self.tile)
	return column, head, feet

	def _footprint(self) -> tuple[int, int, int, int]:
	left = int(self.position[0] // self.tile)
	right = int((self.position[0] + self.player_width - 1e-6) // self.tile)
	head = int(self.position[1] // self.tile)
	feet = int((self.position[1] + self.player_height - 1e-6) // self.tile)
	return left, right, head, feet

	def _on_ground(self) -> bool:
	return self._box_overlaps_solid(self.position[0], self.position[1] + 1.0)

	def _in_water(self) -> bool:
	column, head, feet = self._player_tiles()
	return self._tile_at(head, column) == WATER or self._tile_at(feet, column) == WATER

	def _step_physics(self, action: tuple[int, ...]) -> None:
	left, right, jump = action[0], action[1], action[2]
	submerged = self._in_water()
	direction = float(right) - float(left)
	if direction != 0.0:
	self.facing = 1 if direction > 0 else -1
	speed = self.move_speed * (0.6 if submerged else 1.0)
	self.velocity[0] = direction * speed
	grounded = self._on_ground()
	if jump and (grounded or submerged):
	self.velocity[1] = -self.jump_speed * (0.62 if submerged else 1.0)
	gravity = self.gravity * (0.35 if submerged else 1.0)
	self.velocity[1] = min(self.velocity[1] + gravity * self.dt, self.terminal_speed)
	if submerged:
	self.velocity[1] = 1.0 - self.swim_drag self.dt * self.video_config.fps * 0.1

	step_x = float(self.velocity[0] * self.dt)
	candidate = self.position[0] + step_x
	if self._box_overlaps_solid(candidate, self.position[1]):
	# Step up one tile so single-block ledges do not need a jump.
	lifted = self.position[1] - self.tile
	climbable = step_x != 0.0 and self._on_ground()
	if climbable and not self._box_overlaps_solid(candidate, lifted):
	self.position[0] = candidate
	self.position[1] = lifted
	else:
	self.velocity[0] = 0.0
	else:
	self.position[0] = candidate

	step_y = float(self.velocity[1] * self.dt)
	candidate = self.position[1] + step_y
	if self._box_overlaps_solid(self.position[0], candidate):
	while abs(step_y) > 0.5:
	step_y *= 0.5
	probe = self.position[1] + step_y
	if not self._box_overlaps_solid(self.position[0], probe):
	self.position[1] = probe
	self.velocity[1] = 0.0
	else:
	self.position[1] = candidate

	self.position[0] = float(
	np.clip(self.position[0], self.tile, self.video_config.width - self.tile * 2)
	)
	self.position[1] = float(
	np.clip(self.position[1], self.tile, self.video_config.height - self.tile * 3)
	)

	def _dig(self, action: tuple[int, ...]) -> None:
	self.dig_flash = max(0.0, self.dig_flash - self.dt)
	self.dig_timer = max(0.0, self.dig_timer - self.dt)
	dig_left, dig_right, dig_up, dig_down = action[3], action[4], action[5], action[6]
	if not (dig_left or dig_right or dig_up or dig_down):
	self.dig_target = ()
	return
	left, right, head, feet = self._footprint()
	# A dig clears the player's whole cross-section, so the shaft it opens is
	# always wide enough to move through.
	if dig_left:
	self.facing = -1
	target = tuple((row, left - 1) for row in range(head, feet + 1))
	elif dig_right:
	self.facing = 1
	target = tuple((row, right + 1) for row in range(head, feet + 1))
	elif dig_up:
	target = tuple((head - 1, column) for column in range(left, right + 1))
	else:
	target = tuple((feet + 1, column) for column in range(left, right + 1))
	self.dig_target = target
	if self.dig_timer > 0.0:
	return
	removed = False
	for row, column in target:
	tile = self._tile_at(row, column)
	if not TILE_DIGGABLE[tile]:
	continue
	self.grid[row, column] = AIR
	removed = True
	if tile == ORE:
	self.ore_collected += 1
	self.pending_reward += 1.0
	if removed:
	self.dig_timer = self.dig_cooldown
	self.dig_flash = self.dig_cooldown

	def _step_tiles(self) -> None:
	grid = self.grid
	falls = TILE_FALLS[grid]
	below = np.zeros_like(grid)
	below[:-1] = grid[1:]
	below[-1] = BEDROCK
	open_below = (below == AIR) \| (below == WATER)
	moving = falls & open_below
	moving[-1] = False
	if moving.any():
	destination = _shift_down(moving)
	carried = grid[moving]
	displaced = grid[destination]
	grid[destination] = carried
	grid[moving] = displaced

	# Sand slides off peaks so it forms cones instead of towers.
	side = 1 if self.frame_index % 2 == 0 else -1
	sand = grid == SAND
	below[:-1] = grid[1:]
	below[-1] = BEDROCK
	blocked = sand & ~((below == AIR) \| (below == WATER))
	diagonal = np.zeros_like(grid)
	if side > 0:
	diagonal[:-1, :-1] = grid[1:, 1:]
	else:
	diagonal[:-1, 1:] = grid[1:, :-1]
	slide = blocked & (diagonal == AIR)
	slide[-1] = False
	if slide.any():
	destination = _shift_down(_shift_horizontal(slide, side))
	grid[destination] = grid[slide]
	grid[slide] = AIR

	water = grid == WATER
	below[:-1] = grid[1:]
	below[-1] = BEDROCK
	settled = water & (below != AIR)
	neighbour = np.zeros_like(grid)
	if side > 0:
	neighbour[:, :-1] = grid[:, 1:]
	else:
	neighbour[:, 1:] = grid[:, :-1]
	spread = settled & (neighbour == AIR)
	if spread.any():
	destination = _shift_horizontal(spread, side)
	grid[destination] = WATER
	grid[spread] = AIR

	def _check_fatal(self) -> None:
	buried = False
	if self._box_overlaps_solid(self.position[0], self.position[1]):
	# Material settling onto the miner shoves it upward; only a fully
	# sealed pocket actually kills.
	buried = True
	for offset in range(1, self.tile * 2 + 1):
	lifted = self.position[1] - offset
	if lifted < self.tile:
	break
	if not self._box_overlaps_solid(self.position[0], lifted):
	self.position[1] = float(lifted)
	self.velocity[1] = 0.0
	buried = False
	break
	column, head, _ = self._player_tiles()
	if self._tile_at(head, column) == WATER:
	self.drown_timer += self.dt
	else:
	self.drown_timer = max(0.0, self.drown_timer - self.dt * 0.5)
	if buried or self.drown_timer > 2.5:
	self.deaths += 1
	self.pending_terminal = 1.0
	self.drown_timer = 0.0
	self._respawn_player()

	def scripted_action(self) -> tuple[int, ...]:
	"""Next training-policy action. Does not step the world."""
	return self._scripted_action()

	def _scripted_action(self) -> tuple[int, ...]:
	column, _, feet = self._player_tiles()
	cell = (feet, column)
	if cell == self._last_cell:
	self._stuck_frames += 1
	else:
	self._stuck_frames = 0
	self._last_cell = cell
	if self._hold_frames > 0:
	self._hold_frames -= 1
	return self._action
	if self._stuck_frames > 18:
	# Break out of dead ends instead of grinding at an unreachable target.
	action = [0] * self.ACTION_DIM
	action[int(self.rng.integers(0, self.ACTION_DIM))] = 1
	action[2] = 1
	self._action = tuple(action)
	self._hold_frames = 4
	self._stuck_frames = 0
	return self._action
	action = self._plan_action()
	self._action = action
	self._hold_frames = int(self.rng.integers(0, 3))
	return action

	def _plan_action(self) -> tuple[int, ...]:
	if self.rng.random() < 0.06:
	choice = int(self.rng.integers(0, self.ACTION_DIM))
	action = [0] * self.ACTION_DIM
	action[choice] = 1
	return tuple(action)

	column, head, feet = self._player_tiles()
	action = [0] * self.ACTION_DIM

	# Escape falling material before anything else.
	if self._tile_at(head - 1, column) in {SAND, GRAVEL}:
	action[1 if self.facing > 0 else 0] = 1
	action[2] = 1
	return tuple(action)
	if self.drown_timer > 0.6:
	action[2] = 1
	action[5] = 1
	return tuple(action)

	ore = np.argwhere(self.grid == ORE)
	if ore.size == 0:
	action[6] = 1
	return tuple(action)
	distances = np.abs(ore[:, 0] - feet) + np.abs(ore[:, 1] - column)
	target_row, target_column = ore[int(np.argmin(distances))]
	delta_column = int(target_column) - column
	delta_row = int(target_row) - feet

	if abs(delta_column) <= 1 and abs(delta_row) <= 1:
	if delta_column < 0:
	action[3] = 1
	elif delta_column > 0:
	action[4] = 1
	elif delta_row < 0:
	action[5] = 1
	else:
	action[6] = 1
	return tuple(action)

	if delta_column != 0:
	forward = 1 if delta_column > 0 else 0
	action[forward] = 1
	ahead = column + (1 if delta_column > 0 else -1)
	if self._solid_at(feet, ahead):
	# Bedrock never yields, so climb over it instead of digging forever.
	if TILE_DIGGABLE[self._tile_at(feet, ahead)]:
	action[4 if delta_column > 0 else 3] = 1
	else:
	action[2] = 1
	elif self._solid_at(head, ahead):
	action[2] = 1
	return tuple(action)

	if delta_row > 0:
	action[6] = 1
	else:
	action[2] = 1
	action[5] = 1
	return tuple(action)

	def step(self, action: tuple[int, ...] \| None = None) -> None:
	resolved = self._scripted_action() if action is None else tuple(action)
	if len(resolved) != self.ACTION_DIM:
	raise ValueError(f"Mining actions must have {self.ACTION_DIM} entries.")
	self.pending_reward = 0.0
	self.pending_terminal = 0.0
	self._dig(resolved)
	self._step_physics(resolved)
	self._step_tiles()
	self._check_fatal()
	self.frame_index += 1

	def render(self) -> np.ndarray:
	colors = TILE_COLORS[self.grid] + self.shade[:, :, None]
	if self.dig_flash > 0.0:
	for row, column in self.dig_target:
	if 0 <= row < self.grid_height and 0 <= column < self.grid_width:
	colors[row, column] = _DIG_FLASH_COLOR
	frame = np.clip(colors, 0, 255).astype(np.uint8)
	frame = frame.repeat(self.tile, axis=0).repeat(self.tile, axis=1)

	paint_player(
	frame,
	round(float(self.position[0])),
	round(float(self.position[1])),
	int(self.player_width),
	int(self.player_height),
	self.facing,
	)

	# Ore pips give a fontless score readout that stays stable at training size.
	for index in range(min(self.ore_collected, 8)):
	x = self.tile + index * (self.tile + 2)
	pip = max(2, self.tile // 2)
	frame[2 : 2 + pip, x : x + pip] = TILE_COLORS[ORE]
	return frame

	def next_frame(self) -> np.ndarray:
	frame = self.render()
	self.step()
	return frame

	def sample_frames(self, count: int) -> np.ndarray:
	return np.stack([self.next_frame() for _ in range(count)])

	def sample_gameplay(
	self,
	count: int,
	) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
	frames: list[np.ndarray] = []
	actions: list[tuple[int, ...]] = []
	rewards: list[float] = []
	terminals: list[float] = []
	for _ in range(count):
	frames.append(self.render())
	action = self.scripted_action()
	actions.append(action)
	self.step(action)
	rewards.append(self.pending_reward)
	terminals.append(self.pending_terminal)
	return (
	np.stack(frames),
	np.asarray(actions, dtype=np.float32),
	np.asarray(rewards, dtype=np.float32),
	np.asarray(terminals, dtype=np.float32),
	)