Update WFCEnv documentation (#438)

Co-authored-by: Mark Towers <mark.m.towers@gmail.com>

docs/_scripts/gen_env_docs.py

@@ -19,6 +19,7 @@ all_envs = list(registry.values())
 filtered_envs_by_type = {}
 env_names = []
 babyai_envs = {}
+wfc_envs = {}
 
 # Obtain filtered list
 for env_spec in tqdm(all_envs):
@@ -32,6 +33,9 @@ for env_spec in tqdm(all_envs):
             curr_babyai_env = split[2]
             babyai_env_name = curr_babyai_env.split(":")[1]
             babyai_envs[babyai_env_name] = env_spec
+        elif len(split) > 2 and "wfc" in split[2]:
+            curr_wfc_env = env_spec.kwargs["wfc_config"]
+            wfc_envs[curr_wfc_env] = env_spec
         elif env_module == "minigrid":
             env_name = split[1]
             filtered_envs_by_type[env_name] = env_spec
@@ -56,7 +60,13 @@ filtered_babyai_envs = {
     )
 }
 
-for env_name, env_spec in chain(filtered_envs.items(), filtered_babyai_envs.items()):
+# Because they share a class, only the default (MazeSimple) environment should be kept
+canonical_wfc_env_name = "MazeSimple"
+filtered_wfc_envs = {canonical_wfc_env_name: wfc_envs[canonical_wfc_env_name]}
+
+for env_name, env_spec in chain(
+    filtered_envs.items(), filtered_babyai_envs.items(), filtered_wfc_envs.items()
+):
     env = env_spec.make()
 
     docstring = trim(env.unwrapped.__doc__)

docs/_scripts/gen_envs_display.py

@@ -4,16 +4,41 @@ import os
 
 import gymnasium
 
+# Display bonus WFC presets
+from minigrid.envs.wfc import WFCEnv
+from minigrid.envs.wfc.config import (
+    WFC_PRESETS_INCONSISTENT,
+    WFC_PRESETS_SLOW,
+    register_wfc_presets,
+)
+from utils import env_name_format
+
+register_wfc_presets(WFC_PRESETS_INCONSISTENT, gymnasium.register)
+register_wfc_presets(WFC_PRESETS_SLOW, gymnasium.register)
+
+# Read name from the actual class so it is updated if the class name changes
+WFCENV_NAME = WFCEnv.__name__
+
+
+def title_from_id(env_id):
+    words = []
+    for chunk in env_id.split("_"):
+        words.extend(env_name_format(chunk).split(" "))
+
+    return " ".join(w.title() for w in words)
+
 
 def create_grid_cell(type_id, env_id, base_path):
+    # All WFCEnv environments should link to WFCEnv page
+    href = f"{base_path}{env_id if type_id != 'wfc' else WFCENV_NAME}"
     return f"""
-            <a href="{base_path}{env_id}">
+            <a href="{href}">
                 <div class="env-grid__cell">
                     <div class="cell__image-container">
                         <img src="/_static/videos/{type_id}/{env_id}.gif">
                     </div>
                     <div class="cell__title">
-                        <span>{' '.join(env_id.split('_')).title()}</span>
+                        <span>{title_from_id(env_id)}</span>
                     </div>
                 </div>
             </a>
@@ -64,6 +89,10 @@ if __name__ == "__main__":
             env_name = split[-1].split(":")[-1]
             env_type = env_module if len(split) == 2 else split[-1].split(":")[0]
 
+            if env_name == WFCENV_NAME:
+                env_name = env_spec.kwargs["wfc_config"]
+                assert isinstance(env_name, str)
+
             if env_module == "minigrid":
                 if env_type not in type_dict.keys():
                     type_dict[env_type] = []

docs/_scripts/gen_gifs.py

@@ -24,12 +24,18 @@ envs_completed = []
 # iterate through all envspecs
 for env_spec in tqdm(gymnasium.envs.registry.values()):
     # minigrid.envs:Env or minigrid.envs.babyai:Env
+    if not isinstance(env_spec.entry_point, str):
+        continue
     split = env_spec.entry_point.split(".")
     # ignore minigrid.envs.env_type:Env
     env_module = split[0]
     env_name = split[-1].split(":")[-1]
     env_type = env_module if len(split) == 2 else split[-1].split(":")[0]
 
+    # Override env_name for WFC to include the preset name
+    if env_name == "WFCEnv":
+        env_name = env_spec.kwargs["wfc_config"]
+
     if env_module == "minigrid" and env_name not in envs_completed:
         os.makedirs(os.path.join(output_dir, env_type), exist_ok=True)
         path = os.path.join(output_dir, env_type, env_name + ".gif")

docs/_scripts/utils.py

@@ -32,7 +32,9 @@ def trim(docstring):
 
 def env_name_format(str):
     # KeyCorridorEnv
-    split = re.findall(r"[A-Z](?:[a-z]+|[A-Z]*(?=[A-Z]|$))", str)
+    split = re.findall(r"[A-Z](?:[a-z]+[0-9]*|[A-Z]*[0-9]*(?=[A-Z]|$))", str)
+    if not split:
+        split.append(str)
     # ['Key', 'Corridor', 'Env']
     split = filter(lambda x: x.upper() != "ENV", split)
     return " ".join(split)

docs/content/publications.md

@@ -7,6 +7,7 @@ firstpage:
 
 
 List of publications & submissions using Minigrid or BabyAI (please open a pull request to add missing entries):
+- [DRED: Zero-Shot Transfer in Reinforcement Learning via Data-Regularised Environment Design](https://arxiv.org/abs/2402.03479) (University of Edinburgh, ICML 2024)
 - [Conviction-Based Planning for Sparse Reward Reinforcement Learning Problems](https://prl-theworkshop.github.io/prl2024-icaps/papers/8.pdf)(UQÀM, PRL @ ICAPS 2024)
 - [Learning from Active Human Involvement through Proxy Value Propagation](https://metadriverse.github.io/pvp/)(UCLA, NeurIPS Spotlight 2023)
 - [Go Beyond Imagination: Maximizing Episodic Reachability with World Models](https://arxiv.org/pdf/2308.13661.pdf) (UMich, ICML 2023)

docs/environments/wfc/index.md

@@ -0,0 +1,18 @@
+---
+firstpage:
+lastpage:
+---
+
+## WFC Environments
+
+```{raw} html
+   :file: list.html
+```
+
+```{toctree}
+:hidden:
+:caption: Wave Function Collapse Environments
+
+WFCEnv
+
+```

docs/index.md

@@ -74,6 +74,7 @@ api/wrapper
 
 environments/minigrid/index
 environments/babyai/index
+environments/wfc/index
 ```
 
 ```{toctree}

minigrid/__init__.py

@@ -5,7 +5,7 @@ from gymnasium.envs.registration import register
 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
+from minigrid.envs.wfc.config import WFC_PRESETS, register_wfc_presets
 
 __version__ = "2.3.1"
 
@@ -568,12 +568,7 @@ def register_minigrid_envs():
 
     # WaveFunctionCollapse
     # ----------------------------------------
-    for name in WFC_PRESETS.keys():
-        register(
-            id=f"MiniGrid-WFC-{name}-v0",
-            entry_point="minigrid.envs.wfc:WFCEnv",
-            kwargs={"wfc_config": name},
-        )
+    register_wfc_presets(WFC_PRESETS, register)
 
     # BabyAI - Language based levels - GoTo
     # ----------------------------------------

minigrid/envs/wfc/config.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from collections import ChainMap
 from dataclasses import asdict, dataclass
 from pathlib import Path
 
@@ -218,3 +219,15 @@ WFC_PRESETS_SLOW = {
         input_periodic=True,
     ),  # ~10 mins
 }
+
+WFC_PRESETS_ALL = ChainMap(WFC_PRESETS, WFC_PRESETS_INCONSISTENT, WFC_PRESETS_SLOW)
+
+
+def register_wfc_presets(wfc_presets: dict[str, WFCConfig], register_fn):
+    # Register fn needs to be provided to avoid a circular import
+    for name in wfc_presets.keys():
+        register_fn(
+            id=f"MiniGrid-WFC-{name}-v0",
+            entry_point="minigrid.envs.wfc:WFCEnv",
+            kwargs={"wfc_config": name},
+        )

minigrid/envs/wfc/wfcenv.py

@@ -8,7 +8,7 @@ import numpy as np
 from minigrid.core.constants import OBJECT_TO_IDX
 from minigrid.core.grid import Grid
 from minigrid.core.mission import MissionSpace
-from minigrid.envs.wfc.config import WFC_PRESETS, WFCConfig
+from minigrid.envs.wfc.config import WFC_PRESETS_ALL, WFCConfig
 from minigrid.envs.wfc.graphtransforms import EdgeDescriptor, GraphTransforms
 from minigrid.envs.wfc.wfclogic.control import execute_wfc
 from minigrid.minigrid_env import MiniGridEnv
@@ -33,6 +33,8 @@ class WFCEnv(MiniGridEnv):
 
     This environment procedurally generates a level using the Wave Function Collapse algorithm.
     The environment supports a variety of different level structures but the default is a simple maze.
+    See [WFC module page](index) for sample images of the available presets.
+
     Requires the optional dependencies `imageio` and `networkx` to be installed with `pip install minigrid[wfc]`.
 
     ## Mission Space
@@ -57,7 +59,7 @@ class WFCEnv(MiniGridEnv):
         `(OBJECT_IDX, COLOR_IDX, STATE)`
     - `OBJECT_TO_IDX` and `COLOR_TO_IDX` mapping can be found in
         [minigrid/core/constants.py](minigrid/core/constants.py)
-    - `STATE` refers to the door state with 0=open, 1=closed and 2=locked
+    - `STATE` refers to the door state with 0=open, 1=closed and 2=locked (unused)
 
     ## Rewards
 
@@ -72,8 +74,36 @@ class WFCEnv(MiniGridEnv):
 
     ## Registered Configurations
 
-    S: size of map SxS.
+    - `MiniGrid-WFC-MazeSimple-v0`
+    - `MiniGrid-WFC-DungeonMazeScaled-v0`
+    - `MiniGrid-WFC-RoomsFabric-v0`
+    - `MiniGrid-WFC-ObstaclesBlackdots-v0`
+    - `MiniGrid-WFC-ObstaclesAngular-v0`
+    - `MiniGrid-WFC-ObstaclesHogs3-v0`
+
+    Note: There are many more unregistered configuration presets but they may take a long time to generate a consistent environment.
+
+    They can be registered with the following snippet:
+    ```python
+    import gymnasium
+    from minigrid.envs.wfc.config import register_wfc_presets, WFC_PRESETS_INCONSISTENT, WFC_PRESETS_SLOW
+
+    register_wfc_presets(WFC_PRESETS_INCONSISTENT, gymnasium.register)
+    register_wfc_presets(WFC_PRESETS_SLOW, gymnasium.register)
+    ```
 
+    ## Research
+
+    Adapted for `Minigrid` by the following work.
+
+    ```bibtex
+    @inproceedings{garcin2024dred,
+      title = {DRED: Zero-Shot Transfer in Reinforcement Learning via Data-Regularised Environment Design},
+      author = {Garcin, Samuel and Doran, James and Guo, Shangmin and Lucas, Christopher G and Albrecht, Stefano V},
+      booktitle = {Forty-first International Conference on Machine Learning},
+      year = {2024},
+    }
+    ```
     """
 
     PATTERN_COLOR_CONFIG = {
@@ -90,7 +120,9 @@ class WFCEnv(MiniGridEnv):
         **kwargs,
     ):
         self.config = (
-            wfc_config if isinstance(wfc_config, WFCConfig) else WFC_PRESETS[wfc_config]
+            wfc_config
+            if isinstance(wfc_config, WFCConfig)
+            else WFC_PRESETS_ALL[wfc_config]
         )
         self.padding = 1
 

minigrid/envs/wfc/wfclogic/control.py

@@ -78,7 +78,6 @@ def execute_wfc(
     backtracking: bool = False,
     log_filename: str = "log",
     logging: bool = False,
-    global_constraints: None = None,
     log_stats_to_output: Callable[[dict[str, Any], str], None] | None = None,
     np_random: np.random.Generator | None = None,
 ) -> NDArray[np.integer]: