Prioritized environment replay

mettagrid/config/room/box_world.py

@@ -0,0 +1,147 @@
+"""
+BoxWorld
+========
+
+Densely fills the map with rectangular *boxes*:
+
+* Interior rectangle; one side is guaranteed ≥ 8 cells.
+* Walls 1‑cell thick.
+* A single *entrance* (one gap) sits at a randomly chosen corner.
+* A single heart **altar** sits in the *diagonally opposite* corner.
+* Boxes are stamped until 10 consecutive placement failures, then agents
+  are spawned in the remaining open corridors.
+
+Parameters (override in YAML `room:` block if desired):
+
+    short_range  – inclusive (min, max) for interior short side
+    elong_range  – inclusive (min_add, max_add) added to short → long side
+"""
+
+from typing import List, Optional, Tuple
+
+import numpy as np
+from mettagrid.config.room.room import Room
+
+
+class BoxWorld(Room):
+    STYLE_PARAMETERS = {"box_world": {"hearts_count": 0}}
+
+    # ------------------------------------------------------------------ #
+    # Init
+    # ------------------------------------------------------------------ #
+    def __init__(
+        self,
+        width: int = 120,
+        height: int = 120,
+        agents: int | dict = 0,
+        short_range: Tuple[int, int] = (4, 6),    # interior short side
+        elong_range: Tuple[int, int] = (4, 8),    # + add to short → long
+        seed: Optional[int] = 42,
+        border_width: int = 2,
+        border_object: str = "wall",
+    ):
+        super().__init__(border_width=border_width, border_object=border_object)
+        self._rng = np.random.default_rng(seed)
+        self._width, self._height = np.random.randint(40,100), np.random.randint(40,100)
+        self._agents = agents
+        self._occ = np.zeros((height, width), dtype=bool)
+        self.short_min, self.short_max = short_range
+        self.elong_min, self.elong_max = elong_range
+
+    # ------------------------------------------------------------------ #
+    # Public build
+    # ------------------------------------------------------------------ #
+    def _build(self):
+        grid = np.full((self._height, self._width), "empty", dtype=object)
+        fails = 0
+        while fails < 10:                               # pack until space gone
+            if self._place_box(grid, clearance=1):
+                fails = 0
+            else:
+                fails += 1
+        return self._place_agents(grid)
+
+    # ------------------------------------------------------------------ #
+    # Box helpers
+    # ------------------------------------------------------------------ #
+    def _place_box(self, grid, clearance: int) -> bool:
+        pattern = self._generate_box()
+        return self._place_region(grid, pattern, clearance)
+
+    def _generate_box(self) -> np.ndarray:
+        # --- choose interior dimensions -------------------------------- #
+        short = int(self._rng.integers(self.short_min, self.short_max + 1))
+        long_add = int(self._rng.integers(self.elong_min, self.elong_max + 1))
+        long_side = max(short + long_add, 8)            # ensure ≥ 8 somewhere
+
+        horiz = self._rng.random() < 0.5
+        h_int, w_int = (short, long_side) if horiz else (long_side, short)
+
+        h, w = h_int + 2, w_int + 2                     # +2 for walls
+        pat = np.full((h, w), "wall", dtype=object)
+        pat[1:-1, 1:-1] = "empty"
+
+        # pick entrance corner, altar opposite
+        corners = [(1, 1), (1, w - 2), (h - 2, 1), (h - 2, w - 2)]
+        entrance = corners[self._rng.integers(4)]
+        altar = corners[(corners.index(entrance) + 2) % 4]
+        er, ec = entrance
+
+        pat[altar] = "altar"
+
+        # carve a single gap adjacent to the entrance corner
+        gaps = []
+        if er == 1:          gaps.append((0, ec))          # top wall
+        if er == h - 2:      gaps.append((h - 1, ec))      # bottom wall
+        if ec == 1:          gaps.append((er, 0))          # left wall
+        if ec == w - 2:      gaps.append((er, w - 1))      # right wall
+        gap = gaps[self._rng.integers(len(gaps))]
+        pat[gap] = "empty"
+        return pat
+
+    # ------------------------------------------------------------------ #
+    # Agent placement
+    # ------------------------------------------------------------------ #
+    def _place_agents(self, grid):
+        tags = (
+            ["agent.agent"] * self._agents
+            if isinstance(self._agents, int)
+            else ["agent." + n for n, k in self._agents.items() for _ in range(k)]
+        )
+        for t in tags:
+            pos = self._rand_empty()
+            if pos:
+                grid[pos] = t
+                self._occ[pos] = True
+        return grid
+
+    # ------------------------------------------------------------------ #
+    # Generic helpers (same as LabyrinthWorld)
+    # ------------------------------------------------------------------ #
+    def _place_region(self, grid, pat, clearance: int) -> bool:
+        ph, pw = pat.shape
+        for r, c in self._free_windows((ph + 2 * clearance, pw + 2 * clearance)):
+            grid[r + clearance : r + clearance + ph,
+                 c + clearance : c + clearance + pw] = pat
+            self._occ[r : r + ph + 2 * clearance,
+                      c : c + pw + 2 * clearance] = True
+            return True
+        return False
+
+    def _free_windows(self, shape) -> List[Tuple[int, int]]:
+        h, w = shape
+        H, W = self._occ.shape
+        if h > H or w > W:
+            return []
+        view = np.lib.stride_tricks.as_strided(
+            self._occ, (H - h + 1, W - w + 1, h, w), self._occ.strides * 2
+        )
+        coords = np.argwhere(view.sum(axis=(2, 3)) == 0)
+        self._rng.shuffle(coords)
+        return [tuple(t) for t in coords]
+
+    def _rand_empty(self):
+        empties = np.flatnonzero(~self._occ)
+        return None if empties.size == 0 else tuple(
+            np.unravel_index(self._rng.integers(empties.size), self._occ.shape)
+        )

mettagrid/config/room/cubicle_world.py

@@ -0,0 +1,126 @@
+"""
+CubicleWorld
+=============
+
+Creates a cubicle of full‑length walls: vertical lines from top‑to‑bottom and
+horizontal lines from left‑to‑right.  Interior rectangles of empty space are the
+“cubicles.”  Door‑sized openings are carved in each interior wall.
+
+The outermost perimeter is always a solid wall, ensuring the arena is fully enclosed.
+
+Episode‑time random sampling
+----------------------------
+* env_width,  env_height  ← uniform int samples from width_range / height_range
+* gap_x, gap_y            ← uniform int samples from gap_range (4‑14)
+
+Wall coordinates:
+    vertical walls at 0, gap_x+1, 2*(gap_x+1), …, width‑1
+    horizontal walls at 0, gap_y+1, 2*(gap_y+1), …, height‑1
+"""
+
+from __future__ import annotations
+from typing import Sequence
+import numpy as np
+from mettagrid.config.room.room import Room
+
+
+
+class CubicleWorld(Room):
+    STYLE_PARAMETERS = {"cubicle": {"hearts_count": 0}}
+
+    def __init__(
+        self,
+        width_range: Sequence[int] = (120, 320),
+        height_range: Sequence[int] = (120, 320),
+        gap_range: Sequence[int] = (4, 14),
+        altars_count: int = 50,
+        agents: int | dict = 20,
+        seed: int | None = 42,
+        border_object: str = "wall",
+        border_width: int = 6
+    ):
+        """
+        Initialize a CubicleWorld environment.
+        """
+        super().__init__(border_width = border_width, border_object=border_object)
+        rng = np.random.default_rng(seed)
+
+        self.width = np.random.randint(width_range[0], width_range[1])
+        self.height = np.random.randint(height_range[0], height_range[1])
+        self.gap_x = np.random.randint(gap_range[0], gap_range[1])
+        self.gap_y =  np.random.randint(gap_range[0], gap_range[1])
+
+        self.altars_count = altars_count
+        self._agents = agents
+        self._rng = rng
+
+    # --------------------------------------------------------------- #
+    def _build(self):
+        grid = self._generate_grid()
+        self._place_altars(grid)
+        self._place_agents(grid)
+        return grid
+
+    # --------------------------------------------------------------- #
+    def _generate_grid(self) -> np.ndarray:
+        H, W = self.height, self.width
+        grid = np.full((H, W), "empty", dtype=object)
+
+        # Draw vertical full‑height walls
+        col = 0
+        while col < W:
+            grid[:, col] = "wall"
+            col += self.gap_x + 1
+
+        # Draw horizontal full‑width walls
+        row = 0
+        while row < H:
+            grid[row, :] = "wall"
+            row += self.gap_y + 1
+
+        # Carve door‑sized breaks in every interior wall
+        # Vertical walls: skip borders (col == 0 or col == W‑1)
+        for col in range(self.gap_x + 1, W - 1, self.gap_x + 1):
+            for row_start in range(0, H, self.gap_y + 1):
+                if row_start + self.gap_y < H:
+                    door_row = row_start + 1 + self._rng.integers(self.gap_y)
+                    grid[door_row, col] = "empty"
+
+        # Horizontal walls: skip borders (row == 0 or row == H‑1)
+        for row in range(self.gap_y + 1, H - 1, self.gap_y + 1):
+            for col_start in range(0, W, self.gap_x + 1):
+                if col_start + self.gap_x < W:
+                    door_col = col_start + 1 + self._rng.integers(self.gap_x)
+                    grid[row, door_col] = "empty"
+
+        # Ensure the outer perimeter is a solid wall
+        grid[0, :] = "wall"
+        grid[-1, :] = "wall"
+        grid[:, 0] = "wall"
+        grid[:, -1] = "wall"
+
+        # Occupancy mask for placement
+        self._occ = np.zeros((H, W), dtype=bool)
+        return grid
+
+    # --------------------------------------------------------------- #
+    def _place_altars(self, grid: np.ndarray):
+        empties = list(zip(*np.where(grid == "empty")))
+        self._rng.shuffle(empties)
+        for pos in empties[: self.altars_count]:
+            grid[pos] = "altar"
+            self._occ[pos] = True
+
+    def _place_agents(self, grid: np.ndarray):
+        # build tag list
+        if isinstance(self._agents, int):
+            tags = ["agent.agent"] * self._agents
+        else:
+            tags = ["agent." + t for t, n in self._agents.items() for _ in range(n)]
+
+        empties = np.flatnonzero((grid == "empty") & (~self._occ))
+        self._rng.shuffle(empties)
+        for tag, idx in zip(tags, empties):
+            pos = tuple(np.unravel_index(idx, grid.shape))
+            grid[pos] = tag
+            self._occ[pos] = True

mettagrid/config/room/cubicle_world_uniform.py

@@ -0,0 +1,111 @@
+"""
+CubicleUniformWorld
+===================
+
+Variant of *CubicleWorld* in which every interior wall has a single door
+**exactly at the midpoint** of the wall segment, producing perfectly aligned
+corridors between adjacent cubicles.
+"""
+
+from __future__ import annotations
+from typing import Sequence
+import numpy as np
+
+from mettagrid.config.room.room import Room
+
+
+class CubicleUniformWorld(Room):
+    STYLE_PARAMETERS = {"cubicle_uniform": {"hearts_count": 0}}
+
+    # ------------------------------------------------------------------ #
+    def __init__(
+        self,
+        width_range: Sequence[int] = (120, 320),
+        height_range: Sequence[int] = (120, 320),
+        gap_range: Sequence[int] = (4, 14),
+        altars_count: int = 50,
+        agents: int | dict = 20,
+        seed: int | None = 42,
+        border_object: str = "wall",
+    ):
+        super().__init__(border_object=border_object)
+        rng = np.random.default_rng(seed)
+
+        self.width = rng.integers(width_range[0], width_range[1])
+        self.height = rng.integers(height_range[0], height_range[1])
+        self.gap_x = rng.integers(gap_range[0], gap_range[1])
+        self.gap_y = rng.integers(gap_range[0], gap_range[1])
+
+        self.altars_count = altars_count
+        self._agents = agents
+        self._rng = rng
+
+    # ------------------------------------------------------------------ #
+    def _build(self):
+        grid = self._generate_grid()
+        self._place_altars(grid)
+        self._place_agents(grid)
+        return grid
+
+    # ------------------------------------------------------------------ #
+    def _generate_grid(self) -> np.ndarray:
+        """Create the cubicle lattice with centred doors."""
+        H, W = self.height, self.width
+        grid = np.full((H, W), "empty", dtype=object)
+
+        # Draw vertical full‑height walls
+        col = 0
+        while col < W:
+            grid[:, col] = "wall"
+            col += self.gap_x + 1
+
+        # Draw horizontal full‑width walls
+        row = 0
+        while row < H:
+            grid[row, :] = "wall"
+            row += self.gap_y + 1
+
+        # Carve centred doors in vertical walls (skip outer perimeter)
+        for col in range(self.gap_x + 1, W - 1, self.gap_x + 1):
+            row_start = 0
+            while row_start + self.gap_y < H:
+                door_row = row_start + 1 + self.gap_y // 2
+                grid[door_row, col] = "empty"
+                row_start += self.gap_y + 1
+
+        # Carve centred doors in horizontal walls
+        for row in range(self.gap_y + 1, H - 1, self.gap_y + 1):
+            col_start = 0
+            while col_start + self.gap_x < W:
+                door_col = col_start + 1 + self.gap_x // 2
+                grid[row, door_col] = "empty"
+                col_start += self.gap_x + 1
+
+        # Ensure solid outer perimeter
+        grid[0, :] = grid[-1, :] = "wall"
+        grid[:, 0] = grid[:, -1] = "wall"
+
+        # Occupancy mask
+        self._occ = np.zeros((H, W), dtype=bool)
+        return grid
+
+    # ------------------------------------------------------------------ #
+    def _place_altars(self, grid: np.ndarray):
+        empties = list(zip(*np.where(grid == "empty")))
+        self._rng.shuffle(empties)
+        for pos in empties[: self.altars_count]:
+            grid[pos] = "altar"
+            self._occ[pos] = True
+
+    def _place_agents(self, grid: np.ndarray):
+        tags = (
+            ["agent.agent"] * self._agents
+            if isinstance(self._agents, int)
+            else ["agent." + t for t, n in self._agents.items() for _ in range(n)]
+        )
+        empties = np.flatnonzero((grid == "empty") & (~self._occ))
+        self._rng.shuffle(empties)
+        for tag, idx in zip(tags, empties):
+            pos = tuple(np.unravel_index(idx, grid.shape))
+            grid[pos] = tag
+            self._occ[pos] = True