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), )