S3 maps generation and loading

configs/game/map_builder/load.yaml

@@ -0,0 +1,7 @@
+_target_: mettagrid.map.load.Load
+
+uri: ???
+
+# Optional scene to render on top of the loaded map.
+extra_root:
+  _target_: mettagrid.map.scenes.make_connected.MakeConnected

configs/game/map_builder/load_random.yaml

@@ -0,0 +1,18 @@
+_target_: mettagrid.map.load_random.LoadRandom
+
+dir: ???
+
+# Optional scene to render on top of the loaded map.
+# The following example shows how to patch the map if its agent count doesn't match `game.num_agents`.
+extra_root:
+  _target_: mettagrid.map.scenes.nop.Nop
+
+  children:
+    - where: full
+      scene:
+        _target_: mettagrid.map.scenes.remove_agents.RemoveAgents
+
+    - where: full
+      scene:
+        _target_: mettagrid.map.scenes.random.Random
+        agents: 40

configs/game/map_builder/mapgen_auto.yaml

@@ -11,7 +11,7 @@ root:
 
   config:
     # How many agents to generate? These are placed randomly over the map.
-    num_agents: 10
+    num_agents: 0
 
     # These will be placed anywhere, randomly distributed over the entire map.
     # Values are absolute counts. (TODO: make them percentages?)
@@ -70,14 +70,14 @@ root:
           _target_: mettagrid.map.scenes.random_scene_from_dir.RandomSceneFromDir
           dir: /scenes/dcss/wfc
         weight: 20
-      - scene: /scenes/wfc/blob
-      - scene: /scenes/wfc/blob2
-      - scene: /scenes/wfc/blob3
-      - scene: /scenes/wfc/blocks
-      - scene: /scenes/wfc/dungeons
-      - scene: /scenes/wfc/mazelike1
-      - scene: /scenes/wfc/mazelike2
-      - scene: /scenes/wfc/simple
+      - scene: ./configs/scenes/wfc/blob.yaml
+      - scene: ./configs/scenes/wfc/blob2.yaml
+      - scene: ./configs/scenes/wfc/blob3.yaml
+      - scene: ./configs/scenes/wfc/blocks.yaml
+      - scene: ./configs/scenes/wfc/dungeons.yaml
+      - scene: ./configs/scenes/wfc/mazelike1.yaml
+      - scene: ./configs/scenes/wfc/mazelike2.yaml
+      - scene: ./configs/scenes/wfc/simple.yaml
       - scene:
           _target_: mettagrid.map.scenes.maze.MazeKruskal
           room_size: ["uniform", 1, 3]

configs/game/map_builder/mapgen_convchain.yaml

@@ -13,22 +13,22 @@ root:
     - limit: 1
       order_by: first
       lock: lock1
-      scene: /scenes/convchain/blob
+      scene: ./configs/scenes/convchain/blob.yaml
 
     - limit: 1
       order_by: first
       lock: lock1
-      scene: /scenes/convchain/c_shape
+      scene: ./configs/scenes/convchain/c_shape.yaml
 
     - limit: 1
       order_by: first
       lock: lock1
-      scene: /scenes/convchain/diagonal
+      scene: ./configs/scenes/convchain/diagonal.yaml
 
     - limit: 1
       order_by: last
       lock: lock1
-      scene: /scenes/convchain/dungeon
+      scene: ./configs/scenes/convchain/dungeon.yaml
 
     - where: full
       scene:

configs/game/map_builder/mapgen_maze.yaml

@@ -1,7 +1,7 @@
 _target_: mettagrid.map.mapgen.MapGen
 
-width: ${uniform:20,80,40}
-height: ${uniform:20,80,40}
+width: ${sampling:20,80,40}
+height: ${sampling:20,80,40}
 
 root:
   _target_: mettagrid.map.scenes.room_grid.RoomGrid

configs/game/map_builder/mapgen_simple.yaml

@@ -1,17 +1,16 @@
-# Reproduce old simple.yaml config with MapGen.
-
+# Similar to simple.yaml, but using MapGen.
 _target_: mettagrid.map.mapgen.MapGen
 
-width: ${uniform:20,200,50}
-height: ${uniform:20,200,50}
+width: ${sampling:20,200,50}
+height: ${sampling:20,200,50}
 
 border_width: 6
 
 root:
   _target_: mettagrid.map.scenes.room_grid.RoomGrid
 
   rows: 2
-  columns: ${div:${...num_agents},12}
+  columns: 3 # simple.yaml referred to num_agents here, but we're using OmegaConf now without Hydra
 
   border_width: 0
 
@@ -20,16 +19,16 @@ root:
         _target_: mettagrid.map.scenes.random.Random
 
         objects:
-          mine: ${uniform:1,20,10}
-          generator: ${uniform:1,10,2}
-          altar: ${uniform:1,5,1}
-          armory: ${uniform:1,5,1}
-          lasery: ${uniform:1,5,1}
-          lab: ${uniform:1,5,1}
-          factory: ${uniform:1,5,1}
-          temple: ${uniform:1,5,1}
-
-          block: ${uniform:5,50,20}
-          wall: ${uniform:5,50,20}
+          mine: ${sampling:1,20,10}
+          generator: ${sampling:1,10,2}
+          altar: ${sampling:1,5,1}
+          armory: ${sampling:1,5,1}
+          lasery: ${sampling:1,5,1}
+          lab: ${sampling:1,5,1}
+          factory: ${sampling:1,5,1}
+          temple: ${sampling:1,5,1}
+
+          # block: ${sampling:5,50,20}
+          wall: ${sampling:5,50,20}
 
         agents: 6

configs/game/map_builder/mapgen_wfc_demo.yaml

@@ -12,39 +12,39 @@ root:
     - limit: 1
       lock: lock1
       order_by: first
-      scene: /scenes/wfc/blocks
+      scene: ./configs/scenes/wfc/blocks.yaml
 
     - limit: 1
       lock: lock1
       order_by: first
-      scene: /scenes/wfc/dungeons
+      scene: ./configs/scenes/wfc/dungeons.yaml
 
     - limit: 1
       lock: lock1
       order_by: first
-      scene: /scenes/wfc/simple
+      scene: ./configs/scenes/wfc/simple.yaml
 
     - limit: 1
       lock: lock1
       order_by: first
-      scene: /scenes/wfc/blob
+      scene: ./configs/scenes/wfc/blob.yaml
 
     - limit: 1
       lock: lock1
       order_by: first
-      scene: /scenes/wfc/blob2
+      scene: ./configs/scenes/wfc/blob2.yaml
 
     - limit: 1
       lock: lock1
       order_by: first
-      scene: /scenes/wfc/blob3
+      scene: ./configs/scenes/wfc/blob3.yaml
 
     - limit: 1
       lock: lock1
       order_by: first
-      scene: /scenes/wfc/mazelike1
+      scene: ./configs/scenes/wfc/mazelike1.yaml
 
     - limit: 1
       lock: lock1
       order_by: first
-      scene: /scenes/wfc/mazelike2
+      scene: ./configs/scenes/wfc/mazelike2.yaml

configs/game/map_builder/mapgen_wfc_simple.yaml

@@ -3,4 +3,4 @@ _target_: mettagrid.map.mapgen.MapGen
 width: 50
 height: 50
 
-root: /scenes/wfc/blocks
+root: ./configs/scenes/wfc/blocks.yaml

configs/game/map_builder/random_scene.yaml

@@ -8,7 +8,7 @@ root:
   _target_: mettagrid.map.scenes.random_scene.RandomScene
 
   candidates:
-    - scene: /scenes/convchain/dungeon
+    - scene: ./configs/scenes/convchain/dungeon.yaml
       weight: 2
-    - scene: /scenes/convchain/blob
+    - scene: ./configs/scenes/convchain/blob.yaml
       weight: 1

docs/mapgen.md

@@ -0,0 +1,61 @@
+# Map generation
+
+## S3 maps
+
+To produce maps in bulk and store them in S3, use the following commands:
+
+### Creating maps
+
+```bash
+python -m tools.map.gen --output-uri=s3://BUCKET/DIR ./configs/game/map_builder/mapgen_auto.yaml
+```
+
+`mapgen_auto` builder is an example. You can use any YAML config that can be parsed by OmegaConf.
+
+If `--output-uri` looks like a file (ends with `.yaml` or other extension), the map will be saved to that file.
+
+Otherwise, the map will be saved to a file with a random suffix in that directory.
+
+If `--output-uri` is not specified, the map won't be saved, only shown on screen.
+
+To create maps in bulk, use `--count=N` option.
+
+See `python -m tools.map.gen --help` for more options.
+
+### Viewing maps
+
+You can view a single map by running:
+
+```bash
+python -m tools.map.view s3://BUCKET/PATH/TO/MAP.yaml
+```
+
+The following command will show a random map from an S3 directory:
+
+```bash
+python -m tools.map.view s3://BUCKET/DIR
+```
+
+Same heuristics about detecting if the URI is a file apply here.
+
+### Loading maps in map_builder configs
+
+You can load a random map from an S3 directory in your YAML configs by using `mettagrid.map.load_random.LoadRandom` as a map builder.
+
+`LoadRandom` allows you to modify the map by applying additional scenes to it. Check out `configs/game/map_builder/load_random.yaml` for an example config that modifies the number of agents in the map.
+
+### Indexing maps
+
+Optionally, you can index your maps to make loading them faster.
+
+This is intended to speed up reading from S3. It shouldn't change any functionality, and you should skip playing with this unless you find map loading from S3 is slow.
+
+Index is a plain text file that lists URIs of all the maps. You can assemble it manually, or use the following script:
+
+```bash
+python -m tools.index_s3_maps --dir=s3://BUCKET/DIR --target=s3://BUCKET/DIR/index.txt
+```
+
+`--target` is optional. If not provided, the index will be saved to `{--dir}/index.txt`.
+
+You can then use `mettagrid.map.load_random_from_index.LoadRandomFromIndex` to load a random map from the index.

mettagrid/map/load.py

@@ -0,0 +1,32 @@
+from mettagrid.config.room.room import Room
+from mettagrid.map.utils.storable_map import StorableMap
+
+from .scene import SceneCfg, make_scene
+
+
+# Note that this class can't be a scene, because the width and height come from the stored data.
+class Load(Room):
+    """
+    Load a pregenerated map from a URI (file or S3 object).
+
+    See also: `FromS3Dir` for picking a random map from a directory of pregenerated maps.
+    """
+
+    def __init__(self, uri: str, extra_root: SceneCfg | None = None):
+        super().__init__()
+        self._uri = uri
+        self._storable_map = StorableMap.from_uri(uri)
+
+        if extra_root is not None:
+            self._root = make_scene(extra_root)
+        else:
+            self._root = None
+
+    def build(self):
+        grid = self._storable_map.grid
+
+        if self._root is not None:
+            root_node = self._root.make_node(grid)
+            root_node.render()
+
+        return grid

mettagrid/map/load_random.py

@@ -0,0 +1,39 @@
+import os
+import random
+from pathlib import Path
+
+from mettagrid.map.load import Load
+from mettagrid.map.utils import s3utils
+from mettagrid.map.utils.storage import parse_file_uri
+
+from .scene import SceneCfg
+
+
+def get_random_map_uri(dir_uri: str) -> str:
+    if dir_uri.startswith("s3://"):
+        filenames = s3utils.list_objects(dir_uri)
+        filenames = [uri for uri in filenames if uri.endswith(".yaml")]
+        return random.choice(filenames)
+    else:
+        dirname = parse_file_uri(dir_uri)
+        if not os.path.isdir(dirname):
+            raise ValueError(f"Directory {dirname} does not exist")
+
+        filenames = os.listdir(dirname)
+        filenames = [Path(dirname) / Path(filename) for filename in filenames if filename.endswith(".yaml")]
+        return str(random.choice(filenames))
+
+
+class LoadRandom(Load):
+    """
+    Load a random map from a directory, local or S3.
+
+    See also: `LoadRandomFromIndex` for a version that loads a random map from a pre-generated index.
+    """
+
+    def __init__(self, dir: str, extra_root: SceneCfg | None = None):
+        self._dir_uri = dir
+
+        random_map_uri = get_random_map_uri(self._dir_uri)
+
+        super().__init__(random_map_uri, extra_root)

mettagrid/map/load_random_from_index.py

@@ -0,0 +1,28 @@
+import random
+
+from mettagrid.map.load import Load
+from mettagrid.map.utils import storage
+
+from .scene import SceneCfg
+
+
+class LoadRandomFromIndex(Load):
+    """
+    Load a random map from a list of pregenerated maps.
+
+    The index file can be produced with the following command:
+        python -m tools.index_s3_maps --dir=s3://...
+
+    See also: `LoadRandom` for a version that loads a random map from an S3 directory.
+    """
+
+    def __init__(self, index_uri: str, extra_root: SceneCfg | None = None):
+        self._index_uri = index_uri
+
+        # For 10k maps in a directory we'd have to fetch 100Kb of index data.
+        # (Can we optimize this further by caching?)
+        index = storage.load_from_uri(self._index_uri)
+        index = index.split("\n")
+        random_map_uri = random.choice(index)
+
+        super().__init__(random_map_uri, extra_root)