Wave Function Collapse Environments (#371)

Co-authored-by: Isaac Karth <[email protected]> Co-authored-by: Valentin Valls <[email protected]> Co-authored-by: Isaac Karth <[email protected]> Co-authored-by: Kyle Benesch <[email protected]> Co-authored-by: Samuel Garcin <[email protected]> Co-authored-by: Mark Towers <[email protected]>
Farama-Foundation · Oct 11, 2023 · e726259 · e726259
1 parent f84fc69
commit e726259
Show file tree

Hide file tree

Showing 49 changed files with 2,421 additions and 3 deletions.
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -16,7 +16,7 @@ repos:
       - id: flake8
         args:
           - '--per-file-ignores=*/__init__.py:F401'
-#          - --ignore=
+          - --ignore=E203, W503
           - --max-complexity=30
           - --max-line-length=456
           - --show-source

diff --git a/minigrid/__init__.py b/minigrid/__init__.py
@@ -5,6 +5,7 @@
 from minigrid import minigrid_env, wrappers
 from minigrid.core import roomgrid
 from minigrid.core.world_object import Wall
+from minigrid.envs.wfc.config import WFC_PRESETS
 
 __version__ = "2.3.1"
 
@@ -565,6 +566,15 @@ def register_minigrid_envs():
         entry_point="minigrid.envs:UnlockPickupEnv",
     )
 
+    # WaveFunctionCollapse
+    # ----------------------------------------
+    for name in WFC_PRESETS.keys():
+        register(
+            id=f"MiniGrid-WFC-{name}-v0",
+            entry_point="minigrid.envs.wfc:WFCEnv",
+            kwargs={"wfc_config": name},
+        )
+
     # BabyAI - Language based levels - GoTo
     # ----------------------------------------
 

diff --git a/minigrid/core/world_object.py b/minigrid/core/world_object.py
@@ -74,7 +74,7 @@ def decode(type_idx: int, color_idx: int, state: int) -> WorldObj | None:
         obj_type = IDX_TO_OBJECT[type_idx]
         color = IDX_TO_COLOR[color_idx]
 
-        if obj_type == "empty" or obj_type == "unseen":
+        if obj_type == "empty" or obj_type == "unseen" or obj_type == "agent":
             return None
 
         # State, 0: open, 1: closed, 2: locked

diff --git a/minigrid/envs/wfc/__init__.py b/minigrid/envs/wfc/__init__.py
@@ -0,0 +1,24 @@
+from __future__ import annotations
+
+from minigrid.envs.wfc.config import (
+    WFC_PRESETS,
+    WFC_PRESETS_INCONSISTENT,
+    WFC_PRESETS_SLOW,
+    WFCConfig,
+)
+
+# This is wrapped in a try-except block so the presets can be accessed for registration
+# Otherwise, importing here will fail when networkx is not installed
+try:
+    from minigrid.envs.wfc.wfcenv import WFCEnv
+except ImportError:
+
+    class WFCEnv:
+        """Dummy class to give a helpful error message when dependencies are missing"""
+
+        def __init__(self, *args, **kwargs):
+            from gymnasium.error import DependencyNotInstalled
+
+            raise DependencyNotInstalled(
+                'WFC dependencies are missing, please run `pip install "minigrid[wfc]"`'
+            )
diff --git a/minigrid/envs/wfc/config.py b/minigrid/envs/wfc/config.py
@@ -0,0 +1,220 @@
+from __future__ import annotations
+
+from dataclasses import asdict, dataclass
+from pathlib import Path
+
+from typing_extensions import Literal
+
+PATTERN_PATH = Path(__file__).parent / "patterns"
+
+
+@dataclass
+class WFCConfig:
+    """Dataclass for holding WFC configuration parameters.
+
+    This controls the behavior of the WFC algorithm. The parameters are passed directly to the WFC solver.
+
+    Attributes:
+        pattern_path: Path to the pattern image that will be automatically loaded.
+        tile_size: Size of the tiles in pixels to create from the pattern image.
+        pattern_width: Size of the patterns in tiles to take from the pattern image. (greater than 3 is quite slow)
+        rotations: Number of rotations for each tile.
+        output_periodic: Whether the output should be periodic (wraps over edges).
+        input_periodic: Whether the input should be periodic (wraps over edges).
+        loc_heuristic: Heuristic for choosing the next tile location to collapse.
+        choice_heuristic: Heuristic for choosing the next tile to use between possible tiles.
+        backtracking: Whether to backtrack when contradictions are discovered.
+    """
+
+    pattern_path: Path
+    tile_size: int = 1
+    pattern_width: int = 2
+    rotations: int = 8
+    output_periodic: bool = False
+    input_periodic: bool = False
+    loc_heuristic: Literal[
+        "lexical", "spiral", "entropy", "anti-entropy", "simple", "random"
+    ] = "entropy"
+    choice_heuristic: Literal["lexical", "rarest", "weighted", "random"] = "weighted"
+    backtracking: bool = False
+
+    @property
+    def wfc_kwargs(self):
+        try:
+            from imageio.v2 import imread
+        except ImportError as e:
+            from gymnasium.error import DependencyNotInstalled
+
+            raise DependencyNotInstalled(
+                'imageio is missing, please run `pip install "minigrid[wfc]"`'
+            ) from e
+        kwargs = asdict(self)
+        kwargs["image"] = imread(kwargs.pop("pattern_path"))[:, :, :3]
+        return kwargs
+
+
+# Basic presets for WFC configurations (that should generate in <1 min)
+WFC_PRESETS = {
+    "MazeSimple": WFCConfig(
+        pattern_path=PATTERN_PATH / "SimpleMaze.png",
+        tile_size=1,
+        pattern_width=2,
+        output_periodic=False,
+        input_periodic=False,
+    ),
+    "DungeonMazeScaled": WFCConfig(
+        pattern_path=PATTERN_PATH / "ScaledMaze.png",
+        tile_size=1,
+        pattern_width=2,
+        output_periodic=True,
+        input_periodic=True,
+    ),
+    "RoomsFabric": WFCConfig(
+        pattern_path=PATTERN_PATH / "Fabric.png",
+        tile_size=1,
+        pattern_width=3,
+        output_periodic=False,
+        input_periodic=False,
+    ),
+    "ObstaclesBlackdots": WFCConfig(
+        pattern_path=PATTERN_PATH / "Blackdots.png",
+        tile_size=1,
+        pattern_width=2,
+        output_periodic=False,
+        input_periodic=False,
+    ),
+    "ObstaclesAngular": WFCConfig(
+        pattern_path=PATTERN_PATH / "Angular.png",
+        tile_size=1,
+        pattern_width=3,
+        output_periodic=True,
+        input_periodic=True,
+    ),
+    "ObstaclesHogs3": WFCConfig(
+        pattern_path=PATTERN_PATH / "Hogs.png",
+        tile_size=1,
+        pattern_width=3,
+        output_periodic=True,
+        input_periodic=True,
+    ),
+}
+
+# Presets that take a large number of attempts to generate a consistent environment
+WFC_PRESETS_INCONSISTENT = {
+    "MazeKnot": WFCConfig(
+        pattern_path=PATTERN_PATH / "Knot.png",
+        tile_size=1,
+        pattern_width=3,
+        output_periodic=True,
+        input_periodic=True,
+    ),  # This is not too inconsistent (often 10 attempts is enough)
+    "MazeWall": WFCConfig(
+        pattern_path=PATTERN_PATH / "SimpleWall.png",
+        tile_size=1,
+        pattern_width=2,
+        output_periodic=True,
+        input_periodic=True,
+    ),
+    "RoomsOffice": WFCConfig(
+        pattern_path=PATTERN_PATH / "Office.png",
+        tile_size=1,
+        pattern_width=3,
+        output_periodic=True,
+        input_periodic=True,
+    ),
+    "ObstaclesHogs2": WFCConfig(
+        pattern_path=PATTERN_PATH / "Hogs.png",
+        tile_size=1,
+        pattern_width=2,
+        output_periodic=True,
+        input_periodic=True,
+    ),
+    "Skew2": WFCConfig(
+        pattern_path=PATTERN_PATH / "Skew2.png",
+        tile_size=1,
+        pattern_width=3,
+        output_periodic=True,
+        input_periodic=True,
+    ),
+}
+
+# Slow presets for WFC configurations (Most take about 2-4 min but some take 10+ min)
+WFC_PRESETS_SLOW = {
+    "Maze": WFCConfig(
+        pattern_path=PATTERN_PATH / "Maze.png",
+        tile_size=1,
+        pattern_width=3,
+        output_periodic=True,
+        input_periodic=True,
+    ),  # This is unusually slow: ~20min per 25x25 room
+    "MazeSpirals": WFCConfig(
+        pattern_path=PATTERN_PATH / "Spirals.png",
+        tile_size=1,
+        pattern_width=3,
+        output_periodic=True,
+        input_periodic=True,
+    ),
+    "MazePaths": WFCConfig(
+        pattern_path=PATTERN_PATH / "Paths.png",
+        tile_size=1,
+        pattern_width=3,
+        output_periodic=True,
+        input_periodic=True,
+    ),
+    "Mazelike": WFCConfig(
+        pattern_path=PATTERN_PATH / "Mazelike.png",
+        tile_size=1,
+        pattern_width=3,
+        output_periodic=True,
+        input_periodic=True,
+    ),
+    "Dungeon": WFCConfig(
+        pattern_path=PATTERN_PATH / "DungeonExtr.png",
+        tile_size=1,
+        pattern_width=3,
+        output_periodic=True,
+        input_periodic=True,
+    ),  # ~10 mins
+    "DungeonRooms": WFCConfig(
+        pattern_path=PATTERN_PATH / "Rooms.png",
+        tile_size=1,
+        pattern_width=3,
+        output_periodic=True,
+        input_periodic=True,
+    ),
+    "DungeonLessRooms": WFCConfig(
+        pattern_path=PATTERN_PATH / "LessRooms.png",
+        tile_size=1,
+        pattern_width=3,
+        output_periodic=True,
+        input_periodic=True,
+    ),
+    "DungeonSpirals": WFCConfig(
+        pattern_path=PATTERN_PATH / "SpiralsNeg.png",
+        tile_size=1,
+        pattern_width=3,
+        output_periodic=True,
+        input_periodic=True,
+    ),
+    "RoomsMagicOffice": WFCConfig(
+        pattern_path=PATTERN_PATH / "MagicOffice.png",
+        tile_size=1,
+        pattern_width=3,
+        output_periodic=True,
+        input_periodic=True,
+    ),
+    "SkewCave": WFCConfig(
+        pattern_path=PATTERN_PATH / "Cave.png",
+        tile_size=1,
+        pattern_width=3,
+        output_periodic=False,
+        input_periodic=False,
+    ),
+    "SkewLake": WFCConfig(
+        pattern_path=PATTERN_PATH / "Lake.png",
+        tile_size=1,
+        pattern_width=3,
+        output_periodic=True,
+        input_periodic=True,
+    ),  # ~10 mins
+}