From 86473e0172944b849eb899d82675609f47b60238 Mon Sep 17 00:00:00 2001 From: Vyacheslav Matyukhin Date: Sat, 19 Apr 2025 20:09:23 -0300 Subject: [PATCH 01/19] tests.mapgen can store AsciiMap files to s3 --- configs/mapgen.yaml | 10 +-- configs/mapgen/local_viewer.yaml | 7 ++ configs/mapgen/s3.yaml | 7 ++ configs/mapgen/show.yaml | 2 + mettagrid/map/utils/serialization.py | 123 +++++++++++++++++++++++++++ tests/mapgen.py | 87 +++++++------------ 6 files changed, 172 insertions(+), 64 deletions(-) create mode 100644 configs/mapgen/local_viewer.yaml create mode 100644 configs/mapgen/s3.yaml create mode 100644 configs/mapgen/show.yaml create mode 100644 mettagrid/map/utils/serialization.py diff --git a/configs/mapgen.yaml b/configs/mapgen.yaml index fc00f1b6..578e6948 100644 --- a/configs/mapgen.yaml +++ b/configs/mapgen.yaml @@ -1,17 +1,9 @@ defaults: - mettagrid + - mapgen: show - game/map_builder: ??? - _self_ -mapgen: - show: raylib - metadata: {} - - save: false - target: - dir: ./maps - # name: ??? - will be generated if not set - game: num_agents: 0 recursive_map_builder: false diff --git a/configs/mapgen/local_viewer.yaml b/configs/mapgen/local_viewer.yaml new file mode 100644 index 00000000..aeed45fb --- /dev/null +++ b/configs/mapgen/local_viewer.yaml @@ -0,0 +1,7 @@ +save: true +show: none + +target: + # this is specific to my dir layout -- Slava + # (this config will go away when we fully move to s3) + dir: ../../../mettamap/viewer/public/maps diff --git a/configs/mapgen/s3.yaml b/configs/mapgen/s3.yaml new file mode 100644 index 00000000..bac2197e --- /dev/null +++ b/configs/mapgen/s3.yaml @@ -0,0 +1,7 @@ +save: true +show: none + +s3_dir: ??? + +target: + dir: s3://softmax-public/maps/${..s3_dir} diff --git a/configs/mapgen/show.yaml b/configs/mapgen/show.yaml new file mode 100644 index 00000000..f2866e72 --- /dev/null +++ b/configs/mapgen/show.yaml @@ -0,0 +1,2 @@ +show: raylib +save: false diff --git a/mettagrid/map/utils/serialization.py b/mettagrid/map/utils/serialization.py new file mode 100644 index 00000000..e1c31e0b --- /dev/null +++ b/mettagrid/map/utils/serialization.py @@ -0,0 +1,123 @@ +import logging +import os +from dataclasses import dataclass +from datetime import datetime + +from omegaconf import OmegaConf + +from mettagrid.mettagrid_env import MettaGridEnv + +logger = logging.getLogger(__name__) + + +def env_to_ascii(env: MettaGridEnv, border: bool = False) -> list[str]: + grid = env._c_env.render_ascii() + + # convert to strings + lines = ["".join(row) for row in grid] + + if border: + width = len(lines[0]) + lines = ["┌" + "─" * width + "┐"] + for row in lines: + lines.append("│" + row + "│") + lines.append("└" + "─" * width + "┘") + + return lines + + +@dataclass +class AsciiMap: + metadata: dict + lines: list[str] + config: dict # config that was used to generate the map; can contain unresolved OmegaConf resolvers + resolved_config: dict # resolved config + + def __str__(self) -> str: + frontmatter = OmegaConf.to_yaml( + { + "metadata": self.metadata, + "config": self.config, + "resolved_config": self.resolved_config, + } + ) + content = frontmatter + "\n---\n" + "\n".join(self.lines) + "\n" + return content + + @staticmethod + def from_env(env: MettaGridEnv, gen_time: float) -> "AsciiMap": + ascii_lines = env_to_ascii(env) + + resolved_config = env._env_cfg.game.map_builder + config = env._cfg_template.game.map_builder + metadata = { + **env._env_cfg.mapgen.get("metadata", {}), + "gen_time": gen_time, + "timestamp": datetime.now().isoformat(), + } + return AsciiMap(metadata, ascii_lines, config, resolved_config) + + @staticmethod + def from_uri(uri: str) -> "AsciiMap": + content = load_from_uri(uri) + + # TODO - validate content in a more principled way + (frontmatter, content) = content.split("---\n", 1) + + frontmatter = OmegaConf.load(frontmatter) + metadata = frontmatter.metadata + config = frontmatter.config + resolved_config = frontmatter.resolved_config + lines = content.split("\n") + + return AsciiMap(metadata, lines, config, resolved_config) + + def save(self, uri: str): + save_to_uri(str(self), uri) + + +# The following functions are pretty generic, they can save or load any text to local filesystem or S3. + + +def is_s3_uri(uri: str) -> bool: + return uri.startswith("s3://") + + +def parse_s3_uri(uri: str) -> tuple[str, str]: + if not uri.startswith("s3://"): + raise ValueError(f"URI {uri} is not an S3 URI") + + s3_parts = uri[5:].split("/", 1) + bucket = s3_parts[0] + key = s3_parts[1] + return bucket, key + + +def get_s3_client(): + import boto3 + + # AWS_PROFILE won't be neceesary for most people, but some envirnoments can have multiple profiles + # (Boto3 doesn't pick up the env variable automatically) + session = boto3.Session(profile_name=os.environ.get("AWS_PROFILE", None)) + return session.client("s3") + + +def save_to_uri(text: str, uri: str): + if is_s3_uri(uri): + bucket, key = parse_s3_uri(uri) + s3 = get_s3_client() + s3.put_object(Bucket=bucket, Key=key, Body=text) + else: + with open(uri, "w") as f: + f.write(text) + + +def load_from_uri(uri: str) -> str: + if is_s3_uri(uri): + bucket, key = parse_s3_uri(uri) + s3 = get_s3_client() + response = s3.get_object(Bucket=bucket, Key=key) + return response["Body"].read().decode("utf-8") + else: + with open(uri, "r") as f: + return f.read() diff --git a/tests/mapgen.py b/tests/mapgen.py index 587b87b5..d6bdcd3a 100644 --- a/tests/mapgen.py +++ b/tests/mapgen.py @@ -1,91 +1,68 @@ +import logging import os import random import signal # Aggressively exit on ctrl+c import string import time -from datetime import datetime +from typing import Literal import hydra -from omegaconf import OmegaConf +from mettagrid.map.utils.serialization import AsciiMap, env_to_ascii from mettagrid.mettagrid_env import MettaGridEnv signal.signal(signal.SIGINT, lambda sig, frame: os._exit(0)) +logger = logging.getLogger(__name__) -def env_to_ascii(env): - grid = env._c_env.render_ascii() - # convert to strings - return ["".join(row) for row in grid] +ShowMode = Literal["raylib", "ascii", "ascii_border", "none"] -def save_env_map(env, target_file, gen_time): - ascii_grid = env_to_ascii(env) - resolved_config = env._env_cfg.game.map_builder - config = env._cfg_template.game.map_builder - metadata = { - **env._env_cfg.mapgen.metadata, - "gen_time": gen_time, - "timestamp": datetime.now().isoformat(), - } +def show_env(env: MettaGridEnv, mode: ShowMode): + if mode == "raylib": + from mettagrid.renderer.raylib.raylib_renderer import MettaGridRaylibRenderer + + renderer = MettaGridRaylibRenderer(env._c_env, env._env_cfg.game) + while True: + renderer.render_and_wait() + + elif mode == "ascii": + ascii_lines = env_to_ascii(env) + print("\n".join(ascii_lines)) - with open(target_file, "w") as f: - # Note: OmegaConf messes up multiline strings (adds extra newlines). - # But we take care of it in the mettamap viewer. - frontmatter = OmegaConf.to_yaml( - { - "metadata": metadata, - "config": config, - "resolved_config": resolved_config, - } - ) - f.write(frontmatter) - f.write("\n---\n") - f.write("\n".join(ascii_grid) + "\n") + elif mode == "ascii_border": + ascii_lines = env_to_ascii(env, border=True) + print("\n".join(ascii_lines)) + + elif mode == "none": + pass - print(f"Saved map to {target_file}") + else: + raise ValueError(f"Invalid show mode: {mode}") @hydra.main(version_base=None, config_path="../configs", config_name="mapgen") def main(cfg): + # Generate and measure time taken start = time.time() env = MettaGridEnv(cfg, render_mode="human") gen_time = time.time() - start - print(f"Time taken to create env: {gen_time} seconds") + logger.info(f"Time taken to create env: {gen_time} seconds") + # Save the map if requested if cfg.mapgen.save: target_name = cfg.mapgen.target.get("name", None) if target_name is None: random_suffix = "".join(random.choices(string.ascii_lowercase + string.digits, k=8)) target_name = f"map_{random_suffix}.yaml" - target_file = os.path.join(cfg.mapgen.target.dir, target_name) - save_env_map(env, target_file, gen_time=gen_time) - - show = cfg.mapgen.show - if show == "raylib": - from mettagrid.renderer.raylib.raylib_renderer import MettaGridRaylibRenderer + target_uri = os.path.join(cfg.mapgen.target.dir, target_name) + ascii_map = AsciiMap.from_env(env, gen_time=gen_time) + ascii_map.save(target_uri) - renderer = MettaGridRaylibRenderer(env._c_env, env._env_cfg.game) - while True: - renderer.render_and_wait() - elif show == "ascii": - ascii_grid = env_to_ascii(env) - print("\n".join(ascii_grid)) - elif show == "ascii_border": - ascii_grid = env_to_ascii(env) - # Useful for generating examples for docstrings in code. - width = len(ascii_grid[0]) - lines = ["┌" + "─" * width + "┐"] - for row in ascii_grid: - lines.append("│" + row + "│") - lines.append("└" + "─" * width + "┘") - print("\n".join(lines)) - elif show == "none": - pass - else: - raise ValueError(f"Invalid show mode: {show}") + # Show the map if requested + show_env(env, cfg.mapgen.show) if __name__ == "__main__": From b30704a520962b8592c494c0029273f232517aea Mon Sep 17 00:00:00 2001 From: Vyacheslav Matyukhin Date: Sat, 19 Apr 2025 22:04:47 -0300 Subject: [PATCH 02/19] mettagrid.map.from_uri, mettagrid.map.from_s3_dir --- configs/game/map_builder/from_s3_dir.yaml | 7 ++++ configs/game/map_builder/load.yaml | 7 ++++ configs/index_s3_maps.yaml | 2 ++ mettagrid/map/from_s3_dir.py | 29 +++++++++++++++ mettagrid/map/from_uri.py | 34 ++++++++++++++++++ mettagrid/map/mapgen.py | 10 +++--- mettagrid/map/scenes/ascii.py | 1 + mettagrid/map/utils/serialization.py | 43 +++++++++++++++++++---- mettagrid/objects/constants.pyx | 7 ++++ tools/index_s3_maps.py | 28 +++++++++++++++ 10 files changed, 158 insertions(+), 10 deletions(-) create mode 100644 configs/game/map_builder/from_s3_dir.yaml create mode 100644 configs/game/map_builder/load.yaml create mode 100644 configs/index_s3_maps.yaml create mode 100644 mettagrid/map/from_s3_dir.py create mode 100644 mettagrid/map/from_uri.py create mode 100644 tools/index_s3_maps.py diff --git a/configs/game/map_builder/from_s3_dir.yaml b/configs/game/map_builder/from_s3_dir.yaml new file mode 100644 index 00000000..d0856c70 --- /dev/null +++ b/configs/game/map_builder/from_s3_dir.yaml @@ -0,0 +1,7 @@ +_target_: mettagrid.map.from_s3_dir.FromS3Dir + +s3_dir: ??? + +# Optional scene to render on top of the loaded map. +extra_root: + _target_: mettagrid.map.scenes.make_connected.MakeConnected diff --git a/configs/game/map_builder/load.yaml b/configs/game/map_builder/load.yaml new file mode 100644 index 00000000..b2a6de8a --- /dev/null +++ b/configs/game/map_builder/load.yaml @@ -0,0 +1,7 @@ +_target_: mettagrid.map.from_uri.FromUri + +uri: ??? + +# Optional scene to render on top of the loaded map. +extra_root: + _target_: mettagrid.map.scenes.make_connected.MakeConnected diff --git a/configs/index_s3_maps.yaml b/configs/index_s3_maps.yaml new file mode 100644 index 00000000..a57da86e --- /dev/null +++ b/configs/index_s3_maps.yaml @@ -0,0 +1,2 @@ +index_s3_maps: + dir: ??? diff --git a/mettagrid/map/from_s3_dir.py b/mettagrid/map/from_s3_dir.py new file mode 100644 index 00000000..6b202b8b --- /dev/null +++ b/mettagrid/map/from_s3_dir.py @@ -0,0 +1,29 @@ +# Root map generator, based on nodes. + +import random + +from mettagrid.map.from_uri import FromUri +from mettagrid.map.utils.serialization import load_from_uri, parse_s3_uri + +from .scene import SceneCfg + + +class FromS3Dir(FromUri): + """ + Load a random map from a directory of pregenerated maps. + + The directory must contain a file `index.txt` that lists all the maps in the + directory. (Listing S3 objects would be too slow because of pagination.) + """ + + def __init__(self, s3_dir: str, extra_root: SceneCfg | None = None): + self._s3_dir = s3_dir + + # For 10k maps in a directory we'd have to fetch 100Kb of index data. + # (Can we optimize this further by caching?) + index_uri = load_from_uri(self._s3_dir + "/index.txt") + index = index_uri.split("\n") + random_map = random.choice(index) + + bucket, _ = parse_s3_uri(s3_dir) + super().__init__(f"s3://{bucket}/{random_map}", extra_root) diff --git a/mettagrid/map/from_uri.py b/mettagrid/map/from_uri.py new file mode 100644 index 00000000..6d0a7a21 --- /dev/null +++ b/mettagrid/map/from_uri.py @@ -0,0 +1,34 @@ +# Root map generator, based on nodes. + +from mettagrid.config.room.room import Room +from mettagrid.map.utils.serialization 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 FromUri(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._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._map.to_grid() + + if self._root is not None: + root_node = self._root.make_node(grid) + root_node.render() + + return grid diff --git a/mettagrid/map/mapgen.py b/mettagrid/map/mapgen.py index 00cb326c..5a87d361 100644 --- a/mettagrid/map/mapgen.py +++ b/mettagrid/map/mapgen.py @@ -1,12 +1,14 @@ # Root map generator, based on nodes. import numpy as np -from omegaconf import DictConfig -from .scene import make_scene +from mettagrid.config.room.room import Room +from .scene import SceneCfg, make_scene -class MapGen: - def __init__(self, width: int, height: int, root: DictConfig, border_width: int = 1): + +class MapGen(Room): + def __init__(self, width: int, height: int, root: SceneCfg, border_width: int = 1): + super().__init__() self._width = width self._height = height self._border_width = border_width diff --git a/mettagrid/map/scenes/ascii.py b/mettagrid/map/scenes/ascii.py index cbb5b919..72912a49 100644 --- a/mettagrid/map/scenes/ascii.py +++ b/mettagrid/map/scenes/ascii.py @@ -1,5 +1,6 @@ from mettagrid.map.scenes.inline_ascii import InlineAscii +# Note that these symbols don't match the encoding used by C env (see constants.hpp for those) SYMBOLS = { "A": "agent.agent", "Ap": "agent.prey", diff --git a/mettagrid/map/utils/serialization.py b/mettagrid/map/utils/serialization.py index e1c31e0b..cf30f658 100644 --- a/mettagrid/map/utils/serialization.py +++ b/mettagrid/map/utils/serialization.py @@ -3,9 +3,11 @@ from dataclasses import dataclass from datetime import datetime +import numpy as np from omegaconf import OmegaConf from mettagrid.mettagrid_env import MettaGridEnv +from mettagrid.objects.constants import get_object_type_ascii, get_object_type_names logger = logging.getLogger(__name__) @@ -27,7 +29,11 @@ def env_to_ascii(env: MettaGridEnv, border: bool = False) -> list[str]: @dataclass -class AsciiMap: +class StorableMap: + """ + Serialized map that can be saved to a file or S3. + """ + metadata: dict lines: list[str] config: dict # config that was used to generate the map; can contain unresolved OmegaConf resolvers @@ -44,8 +50,14 @@ def __str__(self) -> str: content = frontmatter + "\n---\n" + "\n".join(self.lines) + "\n" return content + def width(self) -> int: + return len(self.lines[0]) + + def height(self) -> int: + return len(self.lines) + @staticmethod - def from_env(env: MettaGridEnv, gen_time: float) -> "AsciiMap": + def from_env(env: MettaGridEnv, gen_time: float) -> "StorableMap": ascii_lines = env_to_ascii(env) resolved_config = env._env_cfg.game.map_builder @@ -55,25 +67,44 @@ def from_env(env: MettaGridEnv, gen_time: float) -> "AsciiMap": "gen_time": gen_time, "timestamp": datetime.now().isoformat(), } - return AsciiMap(metadata, ascii_lines, config, resolved_config) + return StorableMap(metadata, ascii_lines, config, resolved_config) @staticmethod - def from_uri(uri: str) -> "AsciiMap": + def from_uri(uri: str) -> "StorableMap": + logger.info(f"Loading map from {uri}") content = load_from_uri(uri) # TODO - validate content in a more principled way (frontmatter, content) = content.split("---\n", 1) - frontmatter = OmegaConf.load(frontmatter) + frontmatter = OmegaConf.create(frontmatter) metadata = frontmatter.metadata config = frontmatter.config resolved_config = frontmatter.resolved_config lines = content.split("\n") - return AsciiMap(metadata, lines, config, resolved_config) + return StorableMap(metadata, lines, config, resolved_config) def save(self, uri: str): save_to_uri(str(self), uri) + logger.info(f"Saved map to {uri}") + + def to_grid(self) -> np.ndarray: + grid = np.full((self.height(), self.width()), "empty", dtype=" Date: Sat, 19 Apr 2025 22:04:58 -0300 Subject: [PATCH 03/19] move mapgen script to tools/ --- {tests => tools}/mapgen.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) rename {tests => tools}/mapgen.py (92%) diff --git a/tests/mapgen.py b/tools/mapgen.py similarity index 92% rename from tests/mapgen.py rename to tools/mapgen.py index d6bdcd3a..52e44f41 100644 --- a/tests/mapgen.py +++ b/tools/mapgen.py @@ -8,7 +8,7 @@ import hydra -from mettagrid.map.utils.serialization import AsciiMap, env_to_ascii +from mettagrid.map.utils.serialization import StorableMap, env_to_ascii from mettagrid.mettagrid_env import MettaGridEnv signal.signal(signal.SIGINT, lambda sig, frame: os._exit(0)) @@ -58,7 +58,7 @@ def main(cfg): target_name = f"map_{random_suffix}.yaml" target_uri = os.path.join(cfg.mapgen.target.dir, target_name) - ascii_map = AsciiMap.from_env(env, gen_time=gen_time) + ascii_map = StorableMap.from_env(env, gen_time=gen_time) ascii_map.save(target_uri) # Show the map if requested From 5f1ac6d098b15a72452451e119220e1bc523defc Mon Sep 17 00:00:00 2001 From: Vyacheslav Matyukhin Date: Sat, 19 Apr 2025 22:16:21 -0300 Subject: [PATCH 04/19] nop and remove_agents scenes --- configs/game/map_builder/from_s3_dir.yaml | 13 +++++++++++- mettagrid/map/from_s3_dir.py | 3 +++ mettagrid/map/scenes/nop.py | 14 +++++++++++++ mettagrid/map/scenes/remove_agents.py | 25 +++++++++++++++++++++++ 4 files changed, 54 insertions(+), 1 deletion(-) create mode 100644 mettagrid/map/scenes/nop.py create mode 100644 mettagrid/map/scenes/remove_agents.py diff --git a/configs/game/map_builder/from_s3_dir.yaml b/configs/game/map_builder/from_s3_dir.yaml index d0856c70..3e471e71 100644 --- a/configs/game/map_builder/from_s3_dir.yaml +++ b/configs/game/map_builder/from_s3_dir.yaml @@ -3,5 +3,16 @@ _target_: mettagrid.map.from_s3_dir.FromS3Dir s3_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.make_connected.MakeConnected + _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 diff --git a/mettagrid/map/from_s3_dir.py b/mettagrid/map/from_s3_dir.py index 6b202b8b..4e7b3284 100644 --- a/mettagrid/map/from_s3_dir.py +++ b/mettagrid/map/from_s3_dir.py @@ -14,6 +14,9 @@ class FromS3Dir(FromUri): The directory must contain a file `index.txt` that lists all the maps in the directory. (Listing S3 objects would be too slow because of pagination.) + + The index file can be produced with the following command: + python -m tools.index_s3_maps index_s3_maps.dir=s3://... """ def __init__(self, s3_dir: str, extra_root: SceneCfg | None = None): diff --git a/mettagrid/map/scenes/nop.py b/mettagrid/map/scenes/nop.py new file mode 100644 index 00000000..5f833d1b --- /dev/null +++ b/mettagrid/map/scenes/nop.py @@ -0,0 +1,14 @@ +from mettagrid.map.node import Node +from mettagrid.map.scene import Scene, TypedChild + + +class Nop(Scene): + """ + This scene doesn't do anything. + """ + + def __init__(self, children: list[TypedChild] | None = None): + super().__init__(children=children) + + def _render(self, node: Node): + pass diff --git a/mettagrid/map/scenes/remove_agents.py b/mettagrid/map/scenes/remove_agents.py new file mode 100644 index 00000000..1c5b79d9 --- /dev/null +++ b/mettagrid/map/scenes/remove_agents.py @@ -0,0 +1,25 @@ +from mettagrid.map.node import Node +from mettagrid.map.scene import Scene, TypedChild + + +class RemoveAgents(Scene): + """ + This class solves a frequent problem: `game.num_agents` must match the + number of agents in the map. + + You can use this scene to remove agents from the map. Then apply `Random` + scene to place as many agents as you want. + + (TODO - it might be better to remove `game.num_agents` from the config + entirely, and just use the number of agents in the map.) + """ + + def __init__(self, children: list[TypedChild] | None = None): + super().__init__(children=children) + + def _render(self, node: Node): + for i in range(node.height): + for j in range(node.width): + value = node.grid[i, j] + if value.startswith("agent.") or value == "agent": + node.grid[i, j] = "empty" From ce7a53015f0d2ae2c73850c8ab041b5074578ebe Mon Sep 17 00:00:00 2001 From: Vyacheslav Matyukhin Date: Sat, 19 Apr 2025 22:55:26 -0300 Subject: [PATCH 05/19] rename builders, add documentation --- configs/game/map_builder/load.yaml | 2 +- .../{from_s3_dir.yaml => load_random.yaml} | 4 +- configs/game/map_builder/mapgen_auto.yaml | 2 +- configs/index_s3_maps.yaml | 9 ++++ configs/mapgen.yaml | 1 + configs/mapgen/s3.yaml | 2 +- configs/s3/softmax.yaml | 1 + docs/mapgen.md | 45 +++++++++++++++++++ mettagrid/map/from_s3_dir.py | 32 ------------- mettagrid/map/{from_uri.py => load.py} | 4 +- mettagrid/map/load_random.py | 31 +++++++++++++ tools/index_s3_maps.py | 32 ++++++++++--- 12 files changed, 119 insertions(+), 46 deletions(-) rename configs/game/map_builder/{from_s3_dir.yaml => load_random.yaml} (87%) create mode 100644 configs/s3/softmax.yaml create mode 100644 docs/mapgen.md delete mode 100644 mettagrid/map/from_s3_dir.py rename mettagrid/map/{from_uri.py => load.py} (93%) create mode 100644 mettagrid/map/load_random.py diff --git a/configs/game/map_builder/load.yaml b/configs/game/map_builder/load.yaml index b2a6de8a..c5c24ee5 100644 --- a/configs/game/map_builder/load.yaml +++ b/configs/game/map_builder/load.yaml @@ -1,4 +1,4 @@ -_target_: mettagrid.map.from_uri.FromUri +_target_: mettagrid.map.load.Load uri: ??? diff --git a/configs/game/map_builder/from_s3_dir.yaml b/configs/game/map_builder/load_random.yaml similarity index 87% rename from configs/game/map_builder/from_s3_dir.yaml rename to configs/game/map_builder/load_random.yaml index 3e471e71..5e64918b 100644 --- a/configs/game/map_builder/from_s3_dir.yaml +++ b/configs/game/map_builder/load_random.yaml @@ -1,6 +1,6 @@ -_target_: mettagrid.map.from_s3_dir.FromS3Dir +_target_: mettagrid.map.load_random.LoadRandom -s3_dir: ??? +index_uri: ??? # 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`. diff --git a/configs/game/map_builder/mapgen_auto.yaml b/configs/game/map_builder/mapgen_auto.yaml index 59c2b980..c9079de0 100644 --- a/configs/game/map_builder/mapgen_auto.yaml +++ b/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?) diff --git a/configs/index_s3_maps.yaml b/configs/index_s3_maps.yaml index a57da86e..c0f013b6 100644 --- a/configs/index_s3_maps.yaml +++ b/configs/index_s3_maps.yaml @@ -1,2 +1,11 @@ +defaults: + - s3: softmax + index_s3_maps: + # S3 directory. + # Can be a full URI, `s3://.../dir/`, or a relative path, `dir/`. Relative paths are relative to the `s3.maps_root` config value. dir: ??? + # By default the index will be saved on S3 to `dir/index.txt`. + # But you can provide a local filename instead. + # If you want to save the index to a custom S3 destination, you'll need to provide the full URI, `s3://.../index_name.txt`. + target: null diff --git a/configs/mapgen.yaml b/configs/mapgen.yaml index 578e6948..52acd607 100644 --- a/configs/mapgen.yaml +++ b/configs/mapgen.yaml @@ -1,6 +1,7 @@ defaults: - mettagrid - mapgen: show + - s3: softmax - game/map_builder: ??? - _self_ diff --git a/configs/mapgen/s3.yaml b/configs/mapgen/s3.yaml index bac2197e..b395f493 100644 --- a/configs/mapgen/s3.yaml +++ b/configs/mapgen/s3.yaml @@ -4,4 +4,4 @@ show: none s3_dir: ??? target: - dir: s3://softmax-public/maps/${..s3_dir} + dir: ${...s3.maps_root}/${..s3_dir} diff --git a/configs/s3/softmax.yaml b/configs/s3/softmax.yaml new file mode 100644 index 00000000..a026ecb2 --- /dev/null +++ b/configs/s3/softmax.yaml @@ -0,0 +1 @@ +maps_root: s3://softmax-public/maps diff --git a/docs/mapgen.md b/docs/mapgen.md new file mode 100644 index 00000000..4ce2e5de --- /dev/null +++ b/docs/mapgen.md @@ -0,0 +1,45 @@ +# Map generation + +## S3 maps + +To produce maps in bulk and store them in S3, use the following commands: + +### Creating maps + +```bash +python -m tools.mapgen game/map_builder=mapgen_auto mapgen=s3 mapgen.s3_dir=MY_DIR +``` + +`mapgen_auto` builder is an example. You can use any config from `configs/game/map_builder/`, or write your own. + +**Replace `MY_DIR` with the S3 directory to store the maps.** + +The directory will be created under `s3://softmax-public/maps/` by default; see `configs/mapgen/s3.yaml` for details. + +To create maps in bulk, you can run this in a loop, or use the Hydra's `-m` flag with a range parameter (`x='range(1,100)'`) to [run multiple instances](https://hydra.cc/docs/tutorials/basic/running_your_app/multi-run/). + +### Indexing maps + +The common scenario is to load a random map from the set of pre-generated maps. + +To do this, you need to create an index file. + +Index is a plain text file that lists all the maps. You can assemble it manually, or use the following script: + +```bash +python -m tools.index_s3_maps index_s3_maps.dir=... index_s3_maps.target=... +``` + +`index_s3_maps.target` is optional. If not provided, the index will be saved to `{index_s3_maps.dir}/index.txt`. + +### Loading maps + +By now, all your maps are stored in S3. You can load a random map by using `mettagrid.map.load_random.LoadRandom` as a map builder. + +Check out `configs/game/map_builder/load_random.yaml` for an example config. + +Preview a random map: + +```bash +python -m tests.mapgen game/map_builder=load_random game.map_builder.index_uri=... +``` diff --git a/mettagrid/map/from_s3_dir.py b/mettagrid/map/from_s3_dir.py deleted file mode 100644 index 4e7b3284..00000000 --- a/mettagrid/map/from_s3_dir.py +++ /dev/null @@ -1,32 +0,0 @@ -# Root map generator, based on nodes. - -import random - -from mettagrid.map.from_uri import FromUri -from mettagrid.map.utils.serialization import load_from_uri, parse_s3_uri - -from .scene import SceneCfg - - -class FromS3Dir(FromUri): - """ - Load a random map from a directory of pregenerated maps. - - The directory must contain a file `index.txt` that lists all the maps in the - directory. (Listing S3 objects would be too slow because of pagination.) - - The index file can be produced with the following command: - python -m tools.index_s3_maps index_s3_maps.dir=s3://... - """ - - def __init__(self, s3_dir: str, extra_root: SceneCfg | None = None): - self._s3_dir = s3_dir - - # For 10k maps in a directory we'd have to fetch 100Kb of index data. - # (Can we optimize this further by caching?) - index_uri = load_from_uri(self._s3_dir + "/index.txt") - index = index_uri.split("\n") - random_map = random.choice(index) - - bucket, _ = parse_s3_uri(s3_dir) - super().__init__(f"s3://{bucket}/{random_map}", extra_root) diff --git a/mettagrid/map/from_uri.py b/mettagrid/map/load.py similarity index 93% rename from mettagrid/map/from_uri.py rename to mettagrid/map/load.py index 6d0a7a21..74b57554 100644 --- a/mettagrid/map/from_uri.py +++ b/mettagrid/map/load.py @@ -1,5 +1,3 @@ -# Root map generator, based on nodes. - from mettagrid.config.room.room import Room from mettagrid.map.utils.serialization import StorableMap @@ -7,7 +5,7 @@ # Note that this class can't be a scene, because the width and height come from the stored data. -class FromUri(Room): +class Load(Room): """ Load a pregenerated map from a URI (file or S3 object). diff --git a/mettagrid/map/load_random.py b/mettagrid/map/load_random.py new file mode 100644 index 00000000..e7ca8e2c --- /dev/null +++ b/mettagrid/map/load_random.py @@ -0,0 +1,31 @@ +# Root map generator, based on nodes. + +import random + +from mettagrid.map.load import Load +from mettagrid.map.utils.serialization import load_from_uri + +from .scene import SceneCfg + + +class LoadRandom(Load): + """ + Load a random map from a list of pregenerated maps. + + The directory must contain a file `index.txt` that lists all the maps in the + directory. (Listing S3 objects would be too slow because of pagination.) + + The index file can be produced with the following command: + python -m tools.index_s3_maps index_s3_maps.dir=s3://... + """ + + 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 = load_from_uri(self._index_uri) + index = index.split("\n") + random_map_uri = random.choice(index) + + super().__init__(random_map_uri, extra_root) diff --git a/tools/index_s3_maps.py b/tools/index_s3_maps.py index 9725aca3..7d2b2ef2 100644 --- a/tools/index_s3_maps.py +++ b/tools/index_s3_maps.py @@ -4,7 +4,7 @@ import hydra -from mettagrid.map.utils.serialization import get_s3_client, parse_s3_uri, save_to_uri +from mettagrid.map.utils.serialization import get_s3_client, is_s3_uri, parse_s3_uri, save_to_uri signal.signal(signal.SIGINT, lambda sig, frame: os._exit(0)) logger = logging.getLogger(__name__) @@ -13,15 +13,35 @@ @hydra.main(version_base=None, config_path="../configs", config_name="index_s3_maps") def main(cfg): s3_dir = cfg.index_s3_maps.dir + if not is_s3_uri(s3_dir): + s3_dir = f"{cfg.s3.maps_root}/{s3_dir}" + s3 = get_s3_client() bucket, key = parse_s3_uri(s3_dir) - response = s3.list_objects_v2(Bucket=bucket, Prefix=key) - keys: list[str] = [] - for obj in response["Contents"]: - keys.append(obj["Key"]) + paginator = s3.get_paginator("list_objects_v2") + pages = paginator.paginate(Bucket=bucket, Prefix=key) + + uri_list: list[str] = [] + logger.info(f"Listing objects in s3://{bucket}/{key}...") + for page in pages: + if "Contents" not in page: + continue + + for obj in page["Contents"]: + obj_key = obj["Key"] + if obj_key == key or not obj_key.endswith(".yaml"): + continue + uri_list.append(f"s3://{bucket}/{obj_key}") + + logger.info("Finished listing objects.") + + target = cfg.index_s3_maps.target + if target is None: + target = f"{s3_dir}/index.txt" - save_to_uri(text="\n".join(keys), uri=f"{s3_dir}/index.txt") + save_to_uri(text="\n".join(uri_list), uri=target) + logger.info(f"Index with {len(uri_list)} maps saved to {target}") if __name__ == "__main__": From 446eb2d4276b8c01b8b9170a76050491a8a3a613 Mon Sep 17 00:00:00 2001 From: Vyacheslav Matyukhin Date: Sat, 19 Apr 2025 22:59:25 -0300 Subject: [PATCH 06/19] doc update --- docs/mapgen.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/mapgen.md b/docs/mapgen.md index 4ce2e5de..f89b2139 100644 --- a/docs/mapgen.md +++ b/docs/mapgen.md @@ -36,7 +36,7 @@ python -m tools.index_s3_maps index_s3_maps.dir=... index_s3_maps.target=... By now, all your maps are stored in S3. You can load a random map by using `mettagrid.map.load_random.LoadRandom` as a map builder. -Check out `configs/game/map_builder/load_random.yaml` for an example config. +Check out `configs/game/map_builder/load_random.yaml` for an example config and for how to tune the number of agents in the map. Preview a random map: From 86a1fc47ddc7d6d9a41a56ffb5076fd560a73b80 Mon Sep 17 00:00:00 2001 From: Vyacheslav Matyukhin Date: Mon, 21 Apr 2025 13:18:14 -0300 Subject: [PATCH 07/19] simplify hydra; load_random doesn't need index; refactorings --- configs/game/map_builder/load_random.yaml | 2 +- configs/game/map_builder/mapgen_maze.yaml | 4 +- configs/index_s3_maps.yaml | 6 +-- configs/mapgen.yaml | 1 - configs/mapgen/s3.yaml | 4 +- configs/s3/softmax.yaml | 1 - docs/mapgen.md | 32 ++++++----- mettagrid/map/load_random.py | 24 +++------ mettagrid/map/load_random_from_index.py | 28 ++++++++++ mettagrid/map/node.py | 5 +- mettagrid/map/utils/s3utils.py | 66 +++++++++++++++++++++++ mettagrid/map/utils/serialization.py | 54 ++----------------- tools/index_s3_maps.py | 26 ++------- 13 files changed, 133 insertions(+), 120 deletions(-) delete mode 100644 configs/s3/softmax.yaml create mode 100644 mettagrid/map/load_random_from_index.py create mode 100644 mettagrid/map/utils/s3utils.py diff --git a/configs/game/map_builder/load_random.yaml b/configs/game/map_builder/load_random.yaml index 5e64918b..9d78f5ec 100644 --- a/configs/game/map_builder/load_random.yaml +++ b/configs/game/map_builder/load_random.yaml @@ -1,6 +1,6 @@ _target_: mettagrid.map.load_random.LoadRandom -index_uri: ??? +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`. diff --git a/configs/game/map_builder/mapgen_maze.yaml b/configs/game/map_builder/mapgen_maze.yaml index 4bf865a5..b9d8cbcc 100644 --- a/configs/game/map_builder/mapgen_maze.yaml +++ b/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 diff --git a/configs/index_s3_maps.yaml b/configs/index_s3_maps.yaml index c0f013b6..60669613 100644 --- a/configs/index_s3_maps.yaml +++ b/configs/index_s3_maps.yaml @@ -1,9 +1,5 @@ -defaults: - - s3: softmax - index_s3_maps: - # S3 directory. - # Can be a full URI, `s3://.../dir/`, or a relative path, `dir/`. Relative paths are relative to the `s3.maps_root` config value. + # S3 directory, `s3://.../dir/` dir: ??? # By default the index will be saved on S3 to `dir/index.txt`. # But you can provide a local filename instead. diff --git a/configs/mapgen.yaml b/configs/mapgen.yaml index 52acd607..578e6948 100644 --- a/configs/mapgen.yaml +++ b/configs/mapgen.yaml @@ -1,7 +1,6 @@ defaults: - mettagrid - mapgen: show - - s3: softmax - game/map_builder: ??? - _self_ diff --git a/configs/mapgen/s3.yaml b/configs/mapgen/s3.yaml index b395f493..595a82f4 100644 --- a/configs/mapgen/s3.yaml +++ b/configs/mapgen/s3.yaml @@ -1,7 +1,5 @@ save: true show: none -s3_dir: ??? - target: - dir: ${...s3.maps_root}/${..s3_dir} + dir: ??? diff --git a/configs/s3/softmax.yaml b/configs/s3/softmax.yaml deleted file mode 100644 index a026ecb2..00000000 --- a/configs/s3/softmax.yaml +++ /dev/null @@ -1 +0,0 @@ -maps_root: s3://softmax-public/maps diff --git a/docs/mapgen.md b/docs/mapgen.md index f89b2139..edebe3f8 100644 --- a/docs/mapgen.md +++ b/docs/mapgen.md @@ -7,39 +7,37 @@ To produce maps in bulk and store them in S3, use the following commands: ### Creating maps ```bash -python -m tools.mapgen game/map_builder=mapgen_auto mapgen=s3 mapgen.s3_dir=MY_DIR +python -m tools.mapgen game/map_builder=mapgen_auto mapgen=s3 mapgen.s3_dir=s3://BUCKET/DIR ``` `mapgen_auto` builder is an example. You can use any config from `configs/game/map_builder/`, or write your own. -**Replace `MY_DIR` with the S3 directory to store the maps.** - -The directory will be created under `s3://softmax-public/maps/` by default; see `configs/mapgen/s3.yaml` for details. +**Replace `s3://BUCKET/DIR` with the S3 directory to store the maps.** To create maps in bulk, you can run this in a loop, or use the Hydra's `-m` flag with a range parameter (`x='range(1,100)'`) to [run multiple instances](https://hydra.cc/docs/tutorials/basic/running_your_app/multi-run/). -### Indexing maps +### Loading maps -The common scenario is to load a random map from the set of pre-generated maps. +You can load a random map from an S3 directory by using `mettagrid.map.load_random.LoadRandom` as a map builder. -To do this, you need to create an index file. +Check out `configs/game/map_builder/load_random.yaml` for an example config and for how to tune the number of agents in the map. -Index is a plain text file that lists all the maps. You can assemble it manually, or use the following script: +Preview a random map: ```bash -python -m tools.index_s3_maps index_s3_maps.dir=... index_s3_maps.target=... +python -m tools.mapgen game/map_builder=load_random game.map_builder.dir=... ``` -`index_s3_maps.target` is optional. If not provided, the index will be saved to `{index_s3_maps.dir}/index.txt`. - -### Loading maps - -By now, all your maps are stored in S3. You can load a random map by using `mettagrid.map.load_random.LoadRandom` as a map builder. +### Indexing maps -Check out `configs/game/map_builder/load_random.yaml` for an example config and for how to tune the number of agents in the map. +Optionally, you can index your maps to make loading them faster. -Preview a random map: +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 tests.mapgen game/map_builder=load_random game.map_builder.index_uri=... +python -m tools.index_s3_maps index_s3_maps.dir=s3://BUCKET/DIR index_s3_maps.target=s3://BUCKET/DIR/index.txt ``` + +`index_s3_maps.target` is optional. If not provided, the index will be saved to `{index_s3_maps.dir}/index.txt`. + +You can then use `mettagrid.map.load_random_from_index.LoadRandomFromIndex` to load a random map from the index. diff --git a/mettagrid/map/load_random.py b/mettagrid/map/load_random.py index e7ca8e2c..c56c7035 100644 --- a/mettagrid/map/load_random.py +++ b/mettagrid/map/load_random.py @@ -1,31 +1,23 @@ -# Root map generator, based on nodes. - import random from mettagrid.map.load import Load -from mettagrid.map.utils.serialization import load_from_uri +from mettagrid.map.utils import s3utils from .scene import SceneCfg class LoadRandom(Load): """ - Load a random map from a list of pregenerated maps. - - The directory must contain a file `index.txt` that lists all the maps in the - directory. (Listing S3 objects would be too slow because of pagination.) + Load a random map from S3 directory. - The index file can be produced with the following command: - python -m tools.index_s3_maps index_s3_maps.dir=s3://... + See also: `LoadRandomFromIndex` for a version that loads a random map from a pre-generated index. """ - def __init__(self, index_uri: str, extra_root: SceneCfg | None = None): - self._index_uri = index_uri + def __init__(self, dir: str, extra_root: SceneCfg | None = None): + self._dir = dir - # For 10k maps in a directory we'd have to fetch 100Kb of index data. - # (Can we optimize this further by caching?) - index = load_from_uri(self._index_uri) - index = index.split("\n") - random_map_uri = random.choice(index) + uris = s3utils.list_objects(self._dir) + uris = [uri for uri in uris if uri.endswith(".yaml")] + random_map_uri = random.choice(uris) super().__init__(random_map_uri, extra_root) diff --git a/mettagrid/map/load_random_from_index.py b/mettagrid/map/load_random_from_index.py new file mode 100644 index 00000000..37801029 --- /dev/null +++ b/mettagrid/map/load_random_from_index.py @@ -0,0 +1,28 @@ +import random + +from mettagrid.map.load import Load +from mettagrid.map.utils import s3utils + +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 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 = s3utils.load_from_uri(self._index_uri) + index = index.split("\n") + random_map_uri = random.choice(index) + + super().__init__(random_map_uri, extra_root) diff --git a/mettagrid/map/node.py b/mettagrid/map/node.py index 81d932b6..e73f8528 100644 --- a/mettagrid/map/node.py +++ b/mettagrid/map/node.py @@ -4,6 +4,7 @@ import numpy as np import numpy.typing as npt +from omegaconf import DictConfig, ListConfig @dataclass @@ -58,9 +59,9 @@ def select_areas(self, query) -> list[Area]: selected_areas = [self._full_area] else: # Type check and handling - if isinstance(where, dict) and "tags" in where: + if (isinstance(where, dict) or isinstance(where, DictConfig)) and "tags" in where: tags = where.get("tags", []) - if isinstance(tags, list): + if isinstance(tags, list) or isinstance(tags, ListConfig): for area in areas: match = True for tag in tags: diff --git a/mettagrid/map/utils/s3utils.py b/mettagrid/map/utils/s3utils.py new file mode 100644 index 00000000..88b7e183 --- /dev/null +++ b/mettagrid/map/utils/s3utils.py @@ -0,0 +1,66 @@ +import os + + +def is_s3_uri(uri: str) -> bool: + return uri.startswith("s3://") + + +def parse_s3_uri(uri: str) -> tuple[str, str]: + if not uri.startswith("s3://"): + raise ValueError(f"URI {uri} is not an S3 URI") + + s3_parts = uri[5:].split("/", 1) + bucket = s3_parts[0] + key = s3_parts[1] + return bucket, key + + +def get_s3_client(): + import boto3 + + # AWS_PROFILE won't be neceesary for most people, but some envirnoments can have multiple profiles + # (Boto3 doesn't pick up the env variable automatically) + session = boto3.Session(profile_name=os.environ.get("AWS_PROFILE", None)) + return session.client("s3") + + +def save_to_uri(text: str, uri: str): + if is_s3_uri(uri): + bucket, key = parse_s3_uri(uri) + s3 = get_s3_client() + s3.put_object(Bucket=bucket, Key=key, Body=text) + else: + with open(uri, "w") as f: + f.write(text) + + +def load_from_uri(uri: str) -> str: + if is_s3_uri(uri): + bucket, key = parse_s3_uri(uri) + s3 = get_s3_client() + response = s3.get_object(Bucket=bucket, Key=key) + return response["Body"].read().decode("utf-8") + else: + with open(uri, "r") as f: + return f.read() + + +def list_objects(dir: str) -> list[str]: + s3 = get_s3_client() + bucket, key = parse_s3_uri(dir) + + paginator = s3.get_paginator("list_objects_v2") + pages = paginator.paginate(Bucket=bucket, Prefix=key) + + uri_list: list[str] = [] + for page in pages: + if "Contents" not in page: + continue + + for obj in page["Contents"]: + obj_key = obj["Key"] + if obj_key == key or not obj_key.endswith(".yaml"): + continue + uri_list.append(f"s3://{bucket}/{obj_key}") + + return uri_list diff --git a/mettagrid/map/utils/serialization.py b/mettagrid/map/utils/serialization.py index cf30f658..aa9a6c09 100644 --- a/mettagrid/map/utils/serialization.py +++ b/mettagrid/map/utils/serialization.py @@ -1,5 +1,4 @@ import logging -import os from dataclasses import dataclass from datetime import datetime @@ -9,6 +8,8 @@ from mettagrid.mettagrid_env import MettaGridEnv from mettagrid.objects.constants import get_object_type_ascii, get_object_type_names +from . import s3utils + logger = logging.getLogger(__name__) @@ -72,7 +73,7 @@ def from_env(env: MettaGridEnv, gen_time: float) -> "StorableMap": @staticmethod def from_uri(uri: str) -> "StorableMap": logger.info(f"Loading map from {uri}") - content = load_from_uri(uri) + content = s3utils.load_from_uri(uri) # TODO - validate content in a more principled way (frontmatter, content) = content.split("---\n", 1) @@ -86,7 +87,7 @@ def from_uri(uri: str) -> "StorableMap": return StorableMap(metadata, lines, config, resolved_config) def save(self, uri: str): - save_to_uri(str(self), uri) + s3utils.save_to_uri(str(self), uri) logger.info(f"Saved map to {uri}") def to_grid(self) -> np.ndarray: @@ -105,50 +106,3 @@ def to_grid(self) -> np.ndarray: raise ValueError(f"Unknown character: {char}") return grid - - -# The following functions are pretty generic, they can save or load any text to local filesystem or S3. - - -def is_s3_uri(uri: str) -> bool: - return uri.startswith("s3://") - - -def parse_s3_uri(uri: str) -> tuple[str, str]: - if not uri.startswith("s3://"): - raise ValueError(f"URI {uri} is not an S3 URI") - - s3_parts = uri[5:].split("/", 1) - bucket = s3_parts[0] - key = s3_parts[1] - return bucket, key - - -def get_s3_client(): - import boto3 - - # AWS_PROFILE won't be neceesary for most people, but some envirnoments can have multiple profiles - # (Boto3 doesn't pick up the env variable automatically) - session = boto3.Session(profile_name=os.environ.get("AWS_PROFILE", None)) - return session.client("s3") - - -def save_to_uri(text: str, uri: str): - if is_s3_uri(uri): - bucket, key = parse_s3_uri(uri) - s3 = get_s3_client() - s3.put_object(Bucket=bucket, Key=key, Body=text) - else: - with open(uri, "w") as f: - f.write(text) - - -def load_from_uri(uri: str) -> str: - if is_s3_uri(uri): - bucket, key = parse_s3_uri(uri) - s3 = get_s3_client() - response = s3.get_object(Bucket=bucket, Key=key) - return response["Body"].read().decode("utf-8") - else: - with open(uri, "r") as f: - return f.read() diff --git a/tools/index_s3_maps.py b/tools/index_s3_maps.py index 7d2b2ef2..be0d6a2b 100644 --- a/tools/index_s3_maps.py +++ b/tools/index_s3_maps.py @@ -4,7 +4,7 @@ import hydra -from mettagrid.map.utils.serialization import get_s3_client, is_s3_uri, parse_s3_uri, save_to_uri +from mettagrid.map.utils import s3utils signal.signal(signal.SIGINT, lambda sig, frame: os._exit(0)) logger = logging.getLogger(__name__) @@ -13,34 +13,16 @@ @hydra.main(version_base=None, config_path="../configs", config_name="index_s3_maps") def main(cfg): s3_dir = cfg.index_s3_maps.dir - if not is_s3_uri(s3_dir): + if not s3utils.is_s3_uri(s3_dir): s3_dir = f"{cfg.s3.maps_root}/{s3_dir}" - s3 = get_s3_client() - bucket, key = parse_s3_uri(s3_dir) - - paginator = s3.get_paginator("list_objects_v2") - pages = paginator.paginate(Bucket=bucket, Prefix=key) - - uri_list: list[str] = [] - logger.info(f"Listing objects in s3://{bucket}/{key}...") - for page in pages: - if "Contents" not in page: - continue - - for obj in page["Contents"]: - obj_key = obj["Key"] - if obj_key == key or not obj_key.endswith(".yaml"): - continue - uri_list.append(f"s3://{bucket}/{obj_key}") - - logger.info("Finished listing objects.") + uri_list = s3utils.list_objects(s3_dir) target = cfg.index_s3_maps.target if target is None: target = f"{s3_dir}/index.txt" - save_to_uri(text="\n".join(uri_list), uri=target) + s3utils.save_to_uri(text="\n".join(uri_list), uri=target) logger.info(f"Index with {len(uri_list)} maps saved to {target}") From 70307e998f58db70b267e4722ad5eb1fbc8a6ff8 Mon Sep 17 00:00:00 2001 From: Vyacheslav Matyukhin Date: Mon, 21 Apr 2025 13:24:30 -0300 Subject: [PATCH 08/19] index_s3_maps argparse --- tools/index_s3_maps.py | 22 +++++++++++----------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/tools/index_s3_maps.py b/tools/index_s3_maps.py index be0d6a2b..eaa3312d 100644 --- a/tools/index_s3_maps.py +++ b/tools/index_s3_maps.py @@ -1,26 +1,26 @@ +import argparse import logging import os import signal # Aggressively exit on ctrl+c -import hydra - from mettagrid.map.utils import s3utils signal.signal(signal.SIGINT, lambda sig, frame: os._exit(0)) + logger = logging.getLogger(__name__) +logging.basicConfig(level=logging.INFO) -@hydra.main(version_base=None, config_path="../configs", config_name="index_s3_maps") -def main(cfg): - s3_dir = cfg.index_s3_maps.dir - if not s3utils.is_s3_uri(s3_dir): - s3_dir = f"{cfg.s3.maps_root}/{s3_dir}" +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--dir", type=str) + parser.add_argument("--target", type=str, default=None) + args = parser.parse_args() - uri_list = s3utils.list_objects(s3_dir) + s3_dir = args.dir + target = args.target or f"{s3_dir}/index.txt" - target = cfg.index_s3_maps.target - if target is None: - target = f"{s3_dir}/index.txt" + uri_list = s3utils.list_objects(s3_dir) s3utils.save_to_uri(text="\n".join(uri_list), uri=target) logger.info(f"Index with {len(uri_list)} maps saved to {target}") From c793fd309ea0e9673745fc6561351d745bcca4b6 Mon Sep 17 00:00:00 2001 From: Vyacheslav Matyukhin Date: Mon, 21 Apr 2025 19:13:07 -0300 Subject: [PATCH 09/19] rewrite mapgen without hydra; generate maps without env --- configs/game/map_builder/mapgen_auto.yaml | 16 +-- .../game/map_builder/mapgen_convchain.yaml | 8 +- configs/game/map_builder/mapgen_simple.yaml | 31 ++--- configs/game/map_builder/mapgen_wfc_demo.yaml | 16 +-- .../game/map_builder/mapgen_wfc_simple.yaml | 2 +- configs/game/map_builder/random_scene.yaml | 4 +- mettagrid/map/load.py | 6 +- mettagrid/map/mapgen.py | 7 +- mettagrid/map/scene.py | 15 +- mettagrid/map/scenes/auto.py | 3 +- mettagrid/map/scenes/bsp.py | 3 +- mettagrid/map/utils/serialization.py | 108 --------------- mettagrid/map/utils/storable_map.py | 128 ++++++++++++++++++ mettagrid/mettagrid_env.py | 24 +++- tests/mapgen/test_cli.py | 9 ++ tools/mapgen.py | 126 +++++++++++++---- 16 files changed, 305 insertions(+), 201 deletions(-) delete mode 100644 mettagrid/map/utils/serialization.py create mode 100644 mettagrid/map/utils/storable_map.py create mode 100644 tests/mapgen/test_cli.py diff --git a/configs/game/map_builder/mapgen_auto.yaml b/configs/game/map_builder/mapgen_auto.yaml index c9079de0..edbc3697 100644 --- a/configs/game/map_builder/mapgen_auto.yaml +++ b/configs/game/map_builder/mapgen_auto.yaml @@ -66,14 +66,14 @@ root: # Which content to use? You can list any scenes here. The list will be passed to a RandomScene. # Weights affect the probability of each scene being chosen. Default weight is 1. content: - - 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] diff --git a/configs/game/map_builder/mapgen_convchain.yaml b/configs/game/map_builder/mapgen_convchain.yaml index eebfd897..e4ae6a42 100644 --- a/configs/game/map_builder/mapgen_convchain.yaml +++ b/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: diff --git a/configs/game/map_builder/mapgen_simple.yaml b/configs/game/map_builder/mapgen_simple.yaml index d28bbc3b..29fe57d2 100644 --- a/configs/game/map_builder/mapgen_simple.yaml +++ b/configs/game/map_builder/mapgen_simple.yaml @@ -1,9 +1,8 @@ -# 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 @@ -11,7 +10,7 @@ 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 diff --git a/configs/game/map_builder/mapgen_wfc_demo.yaml b/configs/game/map_builder/mapgen_wfc_demo.yaml index 68d15646..2f9896b5 100644 --- a/configs/game/map_builder/mapgen_wfc_demo.yaml +++ b/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 diff --git a/configs/game/map_builder/mapgen_wfc_simple.yaml b/configs/game/map_builder/mapgen_wfc_simple.yaml index 07b7284f..ba0aac30 100644 --- a/configs/game/map_builder/mapgen_wfc_simple.yaml +++ b/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 diff --git a/configs/game/map_builder/random_scene.yaml b/configs/game/map_builder/random_scene.yaml index 2eea792a..a71af9bd 100644 --- a/configs/game/map_builder/random_scene.yaml +++ b/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 diff --git a/mettagrid/map/load.py b/mettagrid/map/load.py index 74b57554..af18e3c6 100644 --- a/mettagrid/map/load.py +++ b/mettagrid/map/load.py @@ -1,5 +1,5 @@ from mettagrid.config.room.room import Room -from mettagrid.map.utils.serialization import StorableMap +from mettagrid.map.utils.storable_map import StorableMap from .scene import SceneCfg, make_scene @@ -15,7 +15,7 @@ class Load(Room): def __init__(self, uri: str, extra_root: SceneCfg | None = None): super().__init__() self._uri = uri - self._map = StorableMap.from_uri(uri) + self._storable_map = StorableMap.from_uri(uri) if extra_root is not None: self._root = make_scene(extra_root) @@ -23,7 +23,7 @@ def __init__(self, uri: str, extra_root: SceneCfg | None = None): self._root = None def build(self): - grid = self._map.to_grid() + grid = self._storable_map.grid if self._root is not None: root_node = self._root.make_node(grid) diff --git a/mettagrid/map/mapgen.py b/mettagrid/map/mapgen.py index 5a87d361..ef5809d1 100644 --- a/mettagrid/map/mapgen.py +++ b/mettagrid/map/mapgen.py @@ -1,12 +1,17 @@ -# Root map generator, based on nodes. import numpy as np +import numpy.typing as npt from mettagrid.config.room.room import Room from .scene import SceneCfg, make_scene +MapGrid = npt.NDArray[np.str_] + +# Root map generator, based on nodes. class MapGen(Room): + _grid: MapGrid + def __init__(self, width: int, height: int, root: SceneCfg, border_width: int = 1): super().__init__() self._width = width diff --git a/mettagrid/map/scene.py b/mettagrid/map/scene.py index 7b5f5512..086be267 100644 --- a/mettagrid/map/scene.py +++ b/mettagrid/map/scene.py @@ -1,4 +1,4 @@ -from typing import Any, List, Optional, TypedDict, Union +from typing import Any, List, Optional, TypedDict, Union, cast import hydra import numpy as np @@ -16,20 +16,9 @@ class TypedChild(TypedDict): # TODO - more props; use dataclasses instead, or structured configs? -def config_from_path(config_path: str) -> DictConfig: - """Copy-pasted from metta.util.config.config_from_path""" - env_cfg = hydra.compose(config_name=config_path) - if config_path.startswith("/"): - config_path = config_path[1:] - path = config_path.split("/") - for p in path[:-1]: - env_cfg = env_cfg[p] - return env_cfg - - def make_scene(cfg: SceneCfg) -> "Scene": if isinstance(cfg, str): - cfg = config_from_path(cfg) + cfg = cast(SceneCfg, OmegaConf.load(cfg)) if isinstance(cfg, Scene): return cfg diff --git a/mettagrid/map/scenes/auto.py b/mettagrid/map/scenes/auto.py index 2a90ae88..6b827a68 100644 --- a/mettagrid/map/scenes/auto.py +++ b/mettagrid/map/scenes/auto.py @@ -16,8 +16,7 @@ # Global config for convenience. -# Never instantiated, used as a duck type for hydra configs. -# (Not registered as a structured config with Hydra yet.) +# Never instantiated, used as a duck type for OmegaConf configs. @dataclass class AutoConfig: num_agents: int = 0 diff --git a/mettagrid/map/scenes/bsp.py b/mettagrid/map/scenes/bsp.py index 08ac20e5..53cb36f0 100644 --- a/mettagrid/map/scenes/bsp.py +++ b/mettagrid/map/scenes/bsp.py @@ -3,6 +3,7 @@ import numpy as np +from mettagrid.map.mapgen import MapGrid from mettagrid.map.node import Node from mettagrid.map.scene import Scene from mettagrid.map.utils.random import MaybeSeed @@ -263,7 +264,7 @@ def behind(y1, y2) -> bool: return (x + self.min_x, self.ys[x]) @staticmethod - def from_zone(grid: np.ndarray, zone: Zone, side: Literal["up", "down"]) -> "Surface": + def from_zone(grid: MapGrid, zone: Zone, side: Literal["up", "down"]) -> "Surface": # Scan the entire zone, starting from the top or bottom, and collect all the y values that are part of # the surface. min_x = None diff --git a/mettagrid/map/utils/serialization.py b/mettagrid/map/utils/serialization.py deleted file mode 100644 index aa9a6c09..00000000 --- a/mettagrid/map/utils/serialization.py +++ /dev/null @@ -1,108 +0,0 @@ -import logging -from dataclasses import dataclass -from datetime import datetime - -import numpy as np -from omegaconf import OmegaConf - -from mettagrid.mettagrid_env import MettaGridEnv -from mettagrid.objects.constants import get_object_type_ascii, get_object_type_names - -from . import s3utils - -logger = logging.getLogger(__name__) - - -def env_to_ascii(env: MettaGridEnv, border: bool = False) -> list[str]: - grid = env._c_env.render_ascii() - - # convert to strings - lines = ["".join(row) for row in grid] - - if border: - width = len(lines[0]) - lines = ["┌" + "─" * width + "┐"] - for row in lines: - lines.append("│" + row + "│") - lines.append("└" + "─" * width + "┘") - - return lines - - -@dataclass -class StorableMap: - """ - Serialized map that can be saved to a file or S3. - """ - - metadata: dict - lines: list[str] - config: dict # config that was used to generate the map; can contain unresolved OmegaConf resolvers - resolved_config: dict # resolved config - - def __str__(self) -> str: - frontmatter = OmegaConf.to_yaml( - { - "metadata": self.metadata, - "config": self.config, - "resolved_config": self.resolved_config, - } - ) - content = frontmatter + "\n---\n" + "\n".join(self.lines) + "\n" - return content - - def width(self) -> int: - return len(self.lines[0]) - - def height(self) -> int: - return len(self.lines) - - @staticmethod - def from_env(env: MettaGridEnv, gen_time: float) -> "StorableMap": - ascii_lines = env_to_ascii(env) - - resolved_config = env._env_cfg.game.map_builder - config = env._cfg_template.game.map_builder - metadata = { - **env._env_cfg.mapgen.get("metadata", {}), - "gen_time": gen_time, - "timestamp": datetime.now().isoformat(), - } - return StorableMap(metadata, ascii_lines, config, resolved_config) - - @staticmethod - def from_uri(uri: str) -> "StorableMap": - logger.info(f"Loading map from {uri}") - content = s3utils.load_from_uri(uri) - - # TODO - validate content in a more principled way - (frontmatter, content) = content.split("---\n", 1) - - frontmatter = OmegaConf.create(frontmatter) - metadata = frontmatter.metadata - config = frontmatter.config - resolved_config = frontmatter.resolved_config - lines = content.split("\n") - - return StorableMap(metadata, lines, config, resolved_config) - - def save(self, uri: str): - s3utils.save_to_uri(str(self), uri) - logger.info(f"Saved map to {uri}") - - def to_grid(self) -> np.ndarray: - grid = np.full((self.height(), self.width()), "empty", dtype=" str: + if name in ascii_symbols: + return ascii_symbols[name] + + if name == "block": + # FIXME - store maps in a different format, or pick a different character + raise ValueError("Block is not supported in ASCII mode") + + if name.startswith("mine."): + raise ValueError("Colored mines are not supported in ASCII mode") + + if name.startswith("agent."): + raise ValueError("Agent groups are not supported in ASCII mode") + + raise ValueError(f"Unknown object type: {name}") + + +def ascii_to_grid_object(ascii: str) -> str: + if ascii in reverse_ascii_symbols: + return reverse_ascii_symbols[ascii] + + raise ValueError(f"Unknown character: {ascii}") + + +def grid_to_ascii(grid: MapGrid, border: bool = False) -> list[str]: + lines: list[str] = [] + for r in range(grid.shape[0]): + row = [] + for c in range(grid.shape[1]): + row.append(grid_object_to_ascii(grid[r, c])) + lines.append("".join(row)) + + if border: + width = len(lines[0]) + border_lines = ["┌" + "─" * width + "┐"] + for row in lines: + border_lines.append("│" + row + "│") + border_lines.append("└" + "─" * width + "┘") + lines = border_lines + + return lines + + +def ascii_to_grid(lines: list[str]) -> MapGrid: + grid = np.full((len(lines), len(lines[0])), "empty", dtype=" str: + frontmatter = OmegaConf.to_yaml( + { + "metadata": self.metadata, + "config": self.config, + } + ) + content = frontmatter + "\n---\n" + "\n".join(grid_to_ascii(self.grid)) + "\n" + return content + + def width(self) -> int: + return self.grid.shape[1] + + def height(self) -> int: + return self.grid.shape[0] + + @staticmethod + def from_uri(uri: str) -> "StorableMap": + logger.info(f"Loading map from {uri}") + content = s3utils.load_from_uri(uri) + + # TODO - validate content in a more principled way + (frontmatter, content) = content.split("---\n", 1) + + frontmatter = OmegaConf.create(frontmatter) + metadata = frontmatter.metadata + config = frontmatter.config + lines = content.split("\n") + + # make sure we didn't add extra lines because of newlines in the content + lines = [line for line in lines if line] + + return StorableMap(ascii_to_grid(lines), metadata=metadata, config=config) + + def save(self, uri: str): + s3utils.save_to_uri(str(self), uri) + logger.info(f"Saved map to {uri}") diff --git a/mettagrid/mettagrid_env.py b/mettagrid/mettagrid_env.py index 56f03293..c198b858 100644 --- a/mettagrid/mettagrid_env.py +++ b/mettagrid/mettagrid_env.py @@ -12,10 +12,18 @@ class MettaGridEnv(pufferlib.PufferEnv, gym.Env): - def __init__(self, env_cfg: DictConfig, render_mode: Optional[str], buf=None, **kwargs): + def __init__( + self, + env_cfg: DictConfig, + render_mode: Optional[str], + env_map: Optional[np.ndarray] = None, + buf=None, + **kwargs, + ): self._render_mode = render_mode self._cfg_template = env_cfg self._env_cfg = self._get_new_env_cfg() + self._env_map = env_map self._reset_env() self.should_reset = False self._renderer = None @@ -28,11 +36,15 @@ def _get_new_env_cfg(self): return env_cfg def _reset_env(self): - self._map_builder = hydra.utils.instantiate( - self._env_cfg.game.map_builder, - _recursive_=self._env_cfg.game.recursive_map_builder, - ) - env_map = self._map_builder.build() + if self._env_map is None: + self._map_builder = hydra.utils.instantiate( + self._env_cfg.game.map_builder, + _recursive_=self._env_cfg.game.recursive_map_builder, + ) + env_map = self._map_builder.build() + else: + env_map = self._env_map + map_agents = np.count_nonzero(np.char.startswith(env_map, "agent")) assert self._env_cfg.game.num_agents == map_agents, ( f"Number of agents {self._env_cfg.game.num_agents} does not match number of agents in map {map_agents}" diff --git a/tests/mapgen/test_cli.py b/tests/mapgen/test_cli.py new file mode 100644 index 00000000..48ed77aa --- /dev/null +++ b/tests/mapgen/test_cli.py @@ -0,0 +1,9 @@ +import os + + +def test_cli_smoke(): + exit_status = os.system("python -m tools.mapgen --show ascii ./configs/game/map_builder/mapgen_simple.yaml") + assert exit_status == 0 + + exit_status = os.system("python -m tools.mapgen --show ascii ./configs/game/map_builder/NOT_A_CONFIG.yaml") + assert exit_status != 0 diff --git a/tools/mapgen.py b/tools/mapgen.py index 52e44f41..4ab71823 100644 --- a/tools/mapgen.py +++ b/tools/mapgen.py @@ -1,26 +1,41 @@ +import argparse import logging import os import random -import signal # Aggressively exit on ctrl+c +import signal import string import time -from typing import Literal +from datetime import datetime +from typing import Any, Literal, cast, get_args import hydra +import numpy as np +from omegaconf import DictConfig, OmegaConf -from mettagrid.map.utils.serialization import StorableMap, env_to_ascii +from mettagrid.map.utils.storable_map import StorableMap, grid_to_ascii from mettagrid.mettagrid_env import MettaGridEnv +# Aggressively exit on ctrl+c signal.signal(signal.SIGINT, lambda sig, frame: os._exit(0)) logger = logging.getLogger(__name__) +logging.basicConfig(level=logging.INFO) +ShowMode = Literal["raylib", "ascii", "ascii_border"] -ShowMode = Literal["raylib", "ascii", "ascii_border", "none"] +def show_map(storable_map: StorableMap, mode: ShowMode | None): + if not mode: + return -def show_env(env: MettaGridEnv, mode: ShowMode): if mode == "raylib": + num_agents = np.count_nonzero(np.char.startswith(storable_map.grid, "agent")) + + env_cfg = OmegaConf.load("configs/mettagrid.yaml") + env_cfg.game.num_agents = num_agents + + env = MettaGridEnv(cast(Any, env_cfg), env_map=storable_map.grid, render_mode="none") + from mettagrid.renderer.raylib.raylib_renderer import MettaGridRaylibRenderer renderer = MettaGridRaylibRenderer(env._c_env, env._env_cfg.game) @@ -28,41 +43,96 @@ def show_env(env: MettaGridEnv, mode: ShowMode): renderer.render_and_wait() elif mode == "ascii": - ascii_lines = env_to_ascii(env) + ascii_lines = grid_to_ascii(storable_map.grid) print("\n".join(ascii_lines)) elif mode == "ascii_border": - ascii_lines = env_to_ascii(env, border=True) + ascii_lines = grid_to_ascii(storable_map.grid, border=True) print("\n".join(ascii_lines)) - elif mode == "none": - pass - else: raise ValueError(f"Invalid show mode: {mode}") -@hydra.main(version_base=None, config_path="../configs", config_name="mapgen") -def main(cfg): +def make_map(cfg_path: str, overrides: DictConfig | None = None): + cfg: DictConfig = cast(DictConfig, OmegaConf.merge(OmegaConf.load(cfg_path), overrides)) + if not OmegaConf.is_dict(cfg): + raise ValueError(f"Invalid config type: {type(cfg)}") + # Generate and measure time taken start = time.time() - env = MettaGridEnv(cfg, render_mode="human") + map_builder = hydra.utils.instantiate(cfg, _recursive_=False) + grid = map_builder.build() gen_time = time.time() - start - logger.info(f"Time taken to create env: {gen_time} seconds") - - # Save the map if requested - if cfg.mapgen.save: - target_name = cfg.mapgen.target.get("name", None) - if target_name is None: - random_suffix = "".join(random.choices(string.ascii_lowercase + string.digits, k=8)) - target_name = f"map_{random_suffix}.yaml" - - target_uri = os.path.join(cfg.mapgen.target.dir, target_name) - ascii_map = StorableMap.from_env(env, gen_time=gen_time) - ascii_map.save(target_uri) - - # Show the map if requested - show_env(env, cfg.mapgen.show) + logger.info(f"Time taken to build map: {gen_time}s") + + storable_map = StorableMap( + grid=grid, + metadata={ + "gen_time": gen_time, + "timestamp": datetime.now().isoformat(), + }, + config=cfg, + ) + return storable_map + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--output-dir", type=str) + parser.add_argument("--output-uri", type=str) + parser.add_argument("--show", choices=get_args(ShowMode)) + parser.add_argument("--count", type=int, default=1) + parser.add_argument("--overrides", type=str, default="") + parser.add_argument("cfg_path", type=str) + args = parser.parse_args() + + show_mode = args.show + if not show_mode and not args.output_dir and not args.output_uri: + # if not asked to save, show the map + show_mode = "raylib" + + output_dir = args.output_dir + output_uri = args.output_uri + count = args.count + cfg_path = args.cfg_path + overrides = args.overrides + + overrides_cfg = OmegaConf.from_cli([override for override in overrides.split(" ") if override]) + + if output_uri and count > 1: + raise ValueError("Cannot provide both output_uri and count > 1") + + if output_dir and output_uri: + raise ValueError("Cannot provide both output_dir and output_uri") + + def make_output_uri() -> str | None: + if output_uri: + return output_uri + + if not output_dir: + return None # the map won't be saved + + random_suffix = "".join(random.choices(string.ascii_lowercase + string.digits, k=8)) + return f"{output_dir}/map_{random_suffix}.yaml" + + if output_uri: + output_dir = os.path.dirname(output_uri) + + for i in range(count): + if count > 1: + logger.info(f"Generating map {i + 1} of {count}") + + # Generate and measure time taken + storable_map = make_map(cfg_path, overrides_cfg) + + # Save the map if requested + target_uri = make_output_uri() + if target_uri: + storable_map.save(target_uri) + + # Show the map if requested + show_map(storable_map, show_mode) if __name__ == "__main__": From 758e427ae3dfbaa7f837b809516c8c5bcf9e618d Mon Sep 17 00:00:00 2001 From: Vyacheslav Matyukhin Date: Mon, 21 Apr 2025 19:14:22 -0300 Subject: [PATCH 10/19] revert constants.pyx --- mettagrid/objects/constants.pyx | 7 ------- 1 file changed, 7 deletions(-) diff --git a/mettagrid/objects/constants.pyx b/mettagrid/objects/constants.pyx index 253dae8a..e69de29b 100644 --- a/mettagrid/objects/constants.pyx +++ b/mettagrid/objects/constants.pyx @@ -1,7 +0,0 @@ -from mettagrid.objects.constants cimport ObjectTypeNames, ObjectTypeAscii - -def get_object_type_names(): - return ObjectTypeNames - -def get_object_type_ascii(): - return ObjectTypeAscii From 95aec6212e8c7dadea35675de19bb3f3fd738f2e Mon Sep 17 00:00:00 2001 From: Vyacheslav Matyukhin Date: Tue, 22 Apr 2025 10:04:58 -0300 Subject: [PATCH 11/19] simplify type checking for "where" --- mettagrid/map/node.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mettagrid/map/node.py b/mettagrid/map/node.py index e73f8528..b569e0d1 100644 --- a/mettagrid/map/node.py +++ b/mettagrid/map/node.py @@ -59,7 +59,7 @@ def select_areas(self, query) -> list[Area]: selected_areas = [self._full_area] else: # Type check and handling - if (isinstance(where, dict) or isinstance(where, DictConfig)) and "tags" in where: + if isinstance(where, DictConfig) and "tags" in where: tags = where.get("tags", []) if isinstance(tags, list) or isinstance(tags, ListConfig): for area in areas: From ceb5a01c49d8ec4a38708d015c9ba77e8aa7ffd5 Mon Sep 17 00:00:00 2001 From: Vyacheslav Matyukhin Date: Tue, 22 Apr 2025 10:08:45 -0300 Subject: [PATCH 12/19] comment about boto3 --- mettagrid/map/utils/s3utils.py | 1 + 1 file changed, 1 insertion(+) diff --git a/mettagrid/map/utils/s3utils.py b/mettagrid/map/utils/s3utils.py index 88b7e183..147d72c7 100644 --- a/mettagrid/map/utils/s3utils.py +++ b/mettagrid/map/utils/s3utils.py @@ -16,6 +16,7 @@ def parse_s3_uri(uri: str) -> tuple[str, str]: def get_s3_client(): + # no strict dependency on boto3 in mettagrid, install if you need it import boto3 # AWS_PROFILE won't be neceesary for most people, but some envirnoments can have multiple profiles From 448a0b3f4c31692353053e6983bf86e9165b72cf Mon Sep 17 00:00:00 2001 From: Vyacheslav Matyukhin Date: Tue, 22 Apr 2025 10:34:35 -0300 Subject: [PATCH 13/19] extract mettagrid.map.utils.storage code; handle file:// uris --- mettagrid/map/load_random_from_index.py | 4 +-- mettagrid/map/utils/s3utils.py | 25 +--------------- mettagrid/map/utils/storable_map.py | 6 ++-- mettagrid/map/utils/storage.py | 38 +++++++++++++++++++++++++ tools/index_s3_maps.py | 4 +-- 5 files changed, 46 insertions(+), 31 deletions(-) create mode 100644 mettagrid/map/utils/storage.py diff --git a/mettagrid/map/load_random_from_index.py b/mettagrid/map/load_random_from_index.py index 37801029..29131dc0 100644 --- a/mettagrid/map/load_random_from_index.py +++ b/mettagrid/map/load_random_from_index.py @@ -1,7 +1,7 @@ import random from mettagrid.map.load import Load -from mettagrid.map.utils import s3utils +from mettagrid.map.utils import storage from .scene import SceneCfg @@ -21,7 +21,7 @@ def __init__(self, index_uri: str, extra_root: SceneCfg | None = None): # For 10k maps in a directory we'd have to fetch 100Kb of index data. # (Can we optimize this further by caching?) - index = s3utils.load_from_uri(self._index_uri) + index = storage.load_from_uri(self._index_uri) index = index.split("\n") random_map_uri = random.choice(index) diff --git a/mettagrid/map/utils/s3utils.py b/mettagrid/map/utils/s3utils.py index 147d72c7..a1c2376f 100644 --- a/mettagrid/map/utils/s3utils.py +++ b/mettagrid/map/utils/s3utils.py @@ -9,9 +9,7 @@ def parse_s3_uri(uri: str) -> tuple[str, str]: if not uri.startswith("s3://"): raise ValueError(f"URI {uri} is not an S3 URI") - s3_parts = uri[5:].split("/", 1) - bucket = s3_parts[0] - key = s3_parts[1] + (bucket, key) = uri[5:].split("/", 1) return bucket, key @@ -25,27 +23,6 @@ def get_s3_client(): return session.client("s3") -def save_to_uri(text: str, uri: str): - if is_s3_uri(uri): - bucket, key = parse_s3_uri(uri) - s3 = get_s3_client() - s3.put_object(Bucket=bucket, Key=key, Body=text) - else: - with open(uri, "w") as f: - f.write(text) - - -def load_from_uri(uri: str) -> str: - if is_s3_uri(uri): - bucket, key = parse_s3_uri(uri) - s3 = get_s3_client() - response = s3.get_object(Bucket=bucket, Key=key) - return response["Body"].read().decode("utf-8") - else: - with open(uri, "r") as f: - return f.read() - - def list_objects(dir: str) -> list[str]: s3 = get_s3_client() bucket, key = parse_s3_uri(dir) diff --git a/mettagrid/map/utils/storable_map.py b/mettagrid/map/utils/storable_map.py index 1e9a149f..50f0d74d 100644 --- a/mettagrid/map/utils/storable_map.py +++ b/mettagrid/map/utils/storable_map.py @@ -6,7 +6,7 @@ from mettagrid.map.mapgen import MapGrid -from . import s3utils +from . import storage logger = logging.getLogger(__name__) @@ -108,7 +108,7 @@ def height(self) -> int: @staticmethod def from_uri(uri: str) -> "StorableMap": logger.info(f"Loading map from {uri}") - content = s3utils.load_from_uri(uri) + content = storage.load_from_uri(uri) # TODO - validate content in a more principled way (frontmatter, content) = content.split("---\n", 1) @@ -124,5 +124,5 @@ def from_uri(uri: str) -> "StorableMap": return StorableMap(ascii_to_grid(lines), metadata=metadata, config=config) def save(self, uri: str): - s3utils.save_to_uri(str(self), uri) + storage.save_to_uri(str(self), uri) logger.info(f"Saved map to {uri}") diff --git a/mettagrid/map/utils/storage.py b/mettagrid/map/utils/storage.py new file mode 100644 index 00000000..59037e0f --- /dev/null +++ b/mettagrid/map/utils/storage.py @@ -0,0 +1,38 @@ +from mettagrid.map.utils.s3utils import get_s3_client, is_s3_uri, parse_s3_uri + + +def parse_file_uri(uri: str) -> str: + if uri.startswith("file://"): + return uri.split("file://")[1] + + # we don't support any other schemes + if "://" in uri: + raise ValueError(f"Invalid URI: {uri}") + + # probably a local file name + return uri + + +def save_to_uri(text: str, uri: str): + if is_s3_uri(uri): + bucket, key = parse_s3_uri(uri) + s3 = get_s3_client() + s3.put_object(Bucket=bucket, Key=key, Body=text) + + filename = parse_file_uri(uri) + + with open(filename, "w") as f: + f.write(text) + + +def load_from_uri(uri: str) -> str: + if is_s3_uri(uri): + bucket, key = parse_s3_uri(uri) + s3 = get_s3_client() + response = s3.get_object(Bucket=bucket, Key=key) + return response["Body"].read().decode("utf-8") + + filename = parse_file_uri(uri) + + with open(filename, "r") as f: + return f.read() diff --git a/tools/index_s3_maps.py b/tools/index_s3_maps.py index eaa3312d..bc3e3cf0 100644 --- a/tools/index_s3_maps.py +++ b/tools/index_s3_maps.py @@ -3,7 +3,7 @@ import os import signal # Aggressively exit on ctrl+c -from mettagrid.map.utils import s3utils +from mettagrid.map.utils import s3utils, storage signal.signal(signal.SIGINT, lambda sig, frame: os._exit(0)) @@ -22,7 +22,7 @@ def main(): uri_list = s3utils.list_objects(s3_dir) - s3utils.save_to_uri(text="\n".join(uri_list), uri=target) + storage.save_to_uri(text="\n".join(uri_list), uri=target) logger.info(f"Index with {len(uri_list)} maps saved to {target}") From 98440fef5f52853c6d36d0358eb2dd5f3c7cbfd1 Mon Sep 17 00:00:00 2001 From: Vyacheslav Matyukhin Date: Tue, 22 Apr 2025 10:40:04 -0300 Subject: [PATCH 14/19] remove old script configs --- configs/index_s3_maps.yaml | 7 ------- configs/mapgen.yaml | 9 --------- configs/mapgen/for_viewer.yaml | 5 ----- configs/mapgen/local_viewer.yaml | 7 ------- configs/mapgen/s3.yaml | 5 ----- configs/mapgen/show.yaml | 2 -- 6 files changed, 35 deletions(-) delete mode 100644 configs/index_s3_maps.yaml delete mode 100644 configs/mapgen.yaml delete mode 100644 configs/mapgen/for_viewer.yaml delete mode 100644 configs/mapgen/local_viewer.yaml delete mode 100644 configs/mapgen/s3.yaml delete mode 100644 configs/mapgen/show.yaml diff --git a/configs/index_s3_maps.yaml b/configs/index_s3_maps.yaml deleted file mode 100644 index 60669613..00000000 --- a/configs/index_s3_maps.yaml +++ /dev/null @@ -1,7 +0,0 @@ -index_s3_maps: - # S3 directory, `s3://.../dir/` - dir: ??? - # By default the index will be saved on S3 to `dir/index.txt`. - # But you can provide a local filename instead. - # If you want to save the index to a custom S3 destination, you'll need to provide the full URI, `s3://.../index_name.txt`. - target: null diff --git a/configs/mapgen.yaml b/configs/mapgen.yaml deleted file mode 100644 index 578e6948..00000000 --- a/configs/mapgen.yaml +++ /dev/null @@ -1,9 +0,0 @@ -defaults: - - mettagrid - - mapgen: show - - game/map_builder: ??? - - _self_ - -game: - num_agents: 0 - recursive_map_builder: false diff --git a/configs/mapgen/for_viewer.yaml b/configs/mapgen/for_viewer.yaml deleted file mode 100644 index bb24539f..00000000 --- a/configs/mapgen/for_viewer.yaml +++ /dev/null @@ -1,5 +0,0 @@ -save: true -show: none - -target: - dir: ../../../mettamap/viewer/public/maps diff --git a/configs/mapgen/local_viewer.yaml b/configs/mapgen/local_viewer.yaml deleted file mode 100644 index aeed45fb..00000000 --- a/configs/mapgen/local_viewer.yaml +++ /dev/null @@ -1,7 +0,0 @@ -save: true -show: none - -target: - # this is specific to my dir layout -- Slava - # (this config will go away when we fully move to s3) - dir: ../../../mettamap/viewer/public/maps diff --git a/configs/mapgen/s3.yaml b/configs/mapgen/s3.yaml deleted file mode 100644 index 595a82f4..00000000 --- a/configs/mapgen/s3.yaml +++ /dev/null @@ -1,5 +0,0 @@ -save: true -show: none - -target: - dir: ??? diff --git a/configs/mapgen/show.yaml b/configs/mapgen/show.yaml deleted file mode 100644 index f2866e72..00000000 --- a/configs/mapgen/show.yaml +++ /dev/null @@ -1,2 +0,0 @@ -show: raylib -save: false From 3b3c6d0e72a1326dbecd40d91f7b6b9ea2a4b271 Mon Sep 17 00:00:00 2001 From: Vyacheslav Matyukhin Date: Tue, 22 Apr 2025 10:40:21 -0300 Subject: [PATCH 15/19] better --help --- tools/index_s3_maps.py | 2 +- tools/mapgen.py | 12 ++++++------ 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/tools/index_s3_maps.py b/tools/index_s3_maps.py index bc3e3cf0..3786c571 100644 --- a/tools/index_s3_maps.py +++ b/tools/index_s3_maps.py @@ -13,7 +13,7 @@ def main(): parser = argparse.ArgumentParser() - parser.add_argument("--dir", type=str) + parser.add_argument("--dir", type=str, required=True, help="S3 directory, e.g. s3://.../dir") parser.add_argument("--target", type=str, default=None) args = parser.parse_args() diff --git a/tools/mapgen.py b/tools/mapgen.py index 4ab71823..955e788a 100644 --- a/tools/mapgen.py +++ b/tools/mapgen.py @@ -79,12 +79,12 @@ def make_map(cfg_path: str, overrides: DictConfig | None = None): def main(): parser = argparse.ArgumentParser() - parser.add_argument("--output-dir", type=str) - parser.add_argument("--output-uri", type=str) - parser.add_argument("--show", choices=get_args(ShowMode)) - parser.add_argument("--count", type=int, default=1) - parser.add_argument("--overrides", type=str, default="") - parser.add_argument("cfg_path", type=str) + parser.add_argument("--output-dir", type=str, help="Output directory, e.g. ./maps or s3://.../dir") + parser.add_argument("--output-uri", type=str, help="Output URI if you're generating a single map") + parser.add_argument("--show", choices=get_args(ShowMode), help="Show the map in the specified mode") + parser.add_argument("--count", type=int, default=1, help="Number of maps to generate") + parser.add_argument("--overrides", type=str, default="", help="OmniConf overrides for the map config") + parser.add_argument("cfg_path", type=str, help="Path to the map config file") args = parser.parse_args() show_mode = args.show From 484765f926c371161d97240da1d31c41c4f85018 Mon Sep 17 00:00:00 2001 From: Vyacheslav Matyukhin Date: Tue, 22 Apr 2025 10:43:52 -0300 Subject: [PATCH 16/19] update docs --- docs/mapgen.md | 12 ++++++------ mettagrid/map/load_random_from_index.py | 2 +- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/docs/mapgen.md b/docs/mapgen.md index edebe3f8..fa1ea544 100644 --- a/docs/mapgen.md +++ b/docs/mapgen.md @@ -7,14 +7,14 @@ To produce maps in bulk and store them in S3, use the following commands: ### Creating maps ```bash -python -m tools.mapgen game/map_builder=mapgen_auto mapgen=s3 mapgen.s3_dir=s3://BUCKET/DIR +python -m tools.mapgen --output-dir=s3://BUCKET/DIR ./configs/game/map_builder/mapgen_auto.yaml ``` -`mapgen_auto` builder is an example. You can use any config from `configs/game/map_builder/`, or write your own. +`mapgen_auto` builder is an example. You can use any YAML config that can be parsed by OmegaConf. **Replace `s3://BUCKET/DIR` with the S3 directory to store the maps.** -To create maps in bulk, you can run this in a loop, or use the Hydra's `-m` flag with a range parameter (`x='range(1,100)'`) to [run multiple instances](https://hydra.cc/docs/tutorials/basic/running_your_app/multi-run/). +To create maps in bulk, you can run this in a loop, or use `--count` option to generate multiple maps at once. ### Loading maps @@ -25,7 +25,7 @@ Check out `configs/game/map_builder/load_random.yaml` for an example config and Preview a random map: ```bash -python -m tools.mapgen game/map_builder=load_random game.map_builder.dir=... +python -m tools.mapgen ./configs/game/map_builder/load_random.yaml --overrides='dir=s3://BUCKET/DIR' ``` ### Indexing maps @@ -35,9 +35,9 @@ Optionally, you can index your maps to make loading them faster. 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 index_s3_maps.dir=s3://BUCKET/DIR index_s3_maps.target=s3://BUCKET/DIR/index.txt +python -m tools.index_s3_maps --dir=s3://BUCKET/DIR --target=s3://BUCKET/DIR/index.txt ``` -`index_s3_maps.target` is optional. If not provided, the index will be saved to `{index_s3_maps.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. diff --git a/mettagrid/map/load_random_from_index.py b/mettagrid/map/load_random_from_index.py index 29131dc0..3d800c91 100644 --- a/mettagrid/map/load_random_from_index.py +++ b/mettagrid/map/load_random_from_index.py @@ -11,7 +11,7 @@ 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 index_s3_maps.dir=s3://... + python -m tools.index_s3_maps --dir=s3://... See also: `LoadRandom` for a version that loads a random map from an S3 directory. """ From b48645433e8b35342257ad03f0b5cb7d397d0419 Mon Sep 17 00:00:00 2001 From: Vyacheslav Matyukhin Date: Tue, 22 Apr 2025 18:08:55 -0300 Subject: [PATCH 17/19] change mapgen options according to https://github.com/Metta-AI/mettagrid/pull/108#discussion_r2054699842 --- tools/{mapgen.py => map/gen.py} | 40 ++++++++++++--------- tools/{index_s3_maps.py => map/index_s3.py} | 0 2 files changed, 24 insertions(+), 16 deletions(-) rename tools/{mapgen.py => map/gen.py} (80%) rename tools/{index_s3_maps.py => map/index_s3.py} (100%) diff --git a/tools/mapgen.py b/tools/map/gen.py similarity index 80% rename from tools/mapgen.py rename to tools/map/gen.py index 955e788a..9908e6de 100644 --- a/tools/mapgen.py +++ b/tools/map/gen.py @@ -77,10 +77,14 @@ def make_map(cfg_path: str, overrides: DictConfig | None = None): return storable_map +def uri_is_file(uri: str) -> bool: + last_part = uri.split("/")[-1] + return "." in last_part and len(last_part.split(".")[-1]) <= 4 + + def main(): parser = argparse.ArgumentParser() - parser.add_argument("--output-dir", type=str, help="Output directory, e.g. ./maps or s3://.../dir") - parser.add_argument("--output-uri", type=str, help="Output URI if you're generating a single map") + parser.add_argument("--output-uri", type=str, help="Output URI") parser.add_argument("--show", choices=get_args(ShowMode), help="Show the map in the specified mode") parser.add_argument("--count", type=int, default=1, help="Number of maps to generate") parser.add_argument("--overrides", type=str, default="", help="OmniConf overrides for the map config") @@ -88,11 +92,10 @@ def main(): args = parser.parse_args() show_mode = args.show - if not show_mode and not args.output_dir and not args.output_uri: + if not show_mode and not args.output_uri: # if not asked to save, show the map show_mode = "raylib" - output_dir = args.output_dir output_uri = args.output_uri count = args.count cfg_path = args.cfg_path @@ -100,24 +103,29 @@ def main(): overrides_cfg = OmegaConf.from_cli([override for override in overrides.split(" ") if override]) - if output_uri and count > 1: - raise ValueError("Cannot provide both output_uri and count > 1") + if count > 1 and not output_uri: + # requested multiple maps, let's check that output_uri is a directory + if not output_uri: + raise ValueError("Cannot generate more than one map without providing output_uri") - if output_dir and output_uri: - raise ValueError("Cannot provide both output_dir and output_uri") + # s3 can store things at s3://.../foo////file, so we need to remove trailing slashes + while output_uri and output_uri.endswith("/"): + output_uri = output_uri[:-1] - def make_output_uri() -> str | None: - if output_uri: - return output_uri + output_is_file = output_uri and uri_is_file(output_uri) + + if count > 1 and output_is_file: + raise ValueError(f"{output_uri} looks like a file, cannot generate multiple maps in a single file") - if not output_dir: + def make_output_uri() -> str | None: + if not output_uri: return None # the map won't be saved - random_suffix = "".join(random.choices(string.ascii_lowercase + string.digits, k=8)) - return f"{output_dir}/map_{random_suffix}.yaml" + if output_is_file: + return output_uri - if output_uri: - output_dir = os.path.dirname(output_uri) + random_suffix = "".join(random.choices(string.ascii_lowercase + string.digits, k=8)) + return f"{output_uri}/map_{random_suffix}.yaml" for i in range(count): if count > 1: diff --git a/tools/index_s3_maps.py b/tools/map/index_s3.py similarity index 100% rename from tools/index_s3_maps.py rename to tools/map/index_s3.py From 3cb907d4e7679617488d7e57f96df7c65d52a79f Mon Sep 17 00:00:00 2001 From: Vyacheslav Matyukhin Date: Tue, 22 Apr 2025 19:16:40 -0300 Subject: [PATCH 18/19] tools.map.view script; tests; update docs --- docs/mapgen.md | 34 ++++++--- mettagrid/map/load_random.py | 26 +++++-- mettagrid/map/utils/show.py | 39 +++++++++++ tests/mapgen/test_cli.py | 132 +++++++++++++++++++++++++++++++++-- tools/map/gen.py | 44 ++---------- tools/map/view.py | 36 ++++++++++ 6 files changed, 256 insertions(+), 55 deletions(-) create mode 100644 mettagrid/map/utils/show.py create mode 100644 tools/map/view.py diff --git a/docs/mapgen.md b/docs/mapgen.md index fa1ea544..e521ae4c 100644 --- a/docs/mapgen.md +++ b/docs/mapgen.md @@ -7,31 +7,49 @@ To produce maps in bulk and store them in S3, use the following commands: ### Creating maps ```bash -python -m tools.mapgen --output-dir=s3://BUCKET/DIR ./configs/game/map_builder/mapgen_auto.yaml +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. -**Replace `s3://BUCKET/DIR` with the S3 directory to store the maps.** +If `--output-uri` looks like a file (ends with `.yaml` or other extension), the map will be saved to that file. -To create maps in bulk, you can run this in a loop, or use `--count` option to generate multiple maps at once. +Otherwise, the map will be saved to a file with a random suffix in that directory. -### Loading maps +If `--output-uri` is not specified, the map won't be saved, only shown on screen. -You can load a random map from an S3 directory by using `mettagrid.map.load_random.LoadRandom` as a map builder. +To create maps in bulk, use `--count=N` option. -Check out `configs/game/map_builder/load_random.yaml` for an example config and for how to tune the number of agents in the map. +See `python -m tools.map.gen --help` for more options. -Preview a random map: +### 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.mapgen ./configs/game/map_builder/load_random.yaml --overrides='dir=s3://BUCKET/DIR' +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 diff --git a/mettagrid/map/load_random.py b/mettagrid/map/load_random.py index c56c7035..77257078 100644 --- a/mettagrid/map/load_random.py +++ b/mettagrid/map/load_random.py @@ -1,23 +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 S3 directory. + 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 = dir + self._dir_uri = dir - uris = s3utils.list_objects(self._dir) - uris = [uri for uri in uris if uri.endswith(".yaml")] - random_map_uri = random.choice(uris) + random_map_uri = get_random_map_uri(self._dir_uri) super().__init__(random_map_uri, extra_root) diff --git a/mettagrid/map/utils/show.py b/mettagrid/map/utils/show.py new file mode 100644 index 00000000..167b90f5 --- /dev/null +++ b/mettagrid/map/utils/show.py @@ -0,0 +1,39 @@ +from typing import Any, Literal, cast + +import numpy as np +from omegaconf import OmegaConf + +from mettagrid.map.utils.storable_map import StorableMap, grid_to_ascii +from mettagrid.mettagrid_env import MettaGridEnv + +ShowMode = Literal["raylib", "ascii", "ascii_border", "none"] + + +def show_map(storable_map: StorableMap, mode: ShowMode | None): + if not mode or mode == "none": + return + + if mode == "raylib": + num_agents = np.count_nonzero(np.char.startswith(storable_map.grid, "agent")) + + env_cfg = OmegaConf.load("configs/mettagrid.yaml") + env_cfg.game.num_agents = num_agents + + env = MettaGridEnv(cast(Any, env_cfg), env_map=storable_map.grid, render_mode="none") + + from mettagrid.renderer.raylib.raylib_renderer import MettaGridRaylibRenderer + + renderer = MettaGridRaylibRenderer(env._c_env, env._env_cfg.game) + while True: + renderer.render_and_wait() + + elif mode == "ascii": + ascii_lines = grid_to_ascii(storable_map.grid) + print("\n".join(ascii_lines)) + + elif mode == "ascii_border": + ascii_lines = grid_to_ascii(storable_map.grid, border=True) + print("\n".join(ascii_lines)) + + else: + raise ValueError(f"Invalid show mode: {mode}") diff --git a/tests/mapgen/test_cli.py b/tests/mapgen/test_cli.py index 48ed77aa..c1528810 100644 --- a/tests/mapgen/test_cli.py +++ b/tests/mapgen/test_cli.py @@ -1,9 +1,133 @@ import os +import subprocess -def test_cli_smoke(): - exit_status = os.system("python -m tools.mapgen --show ascii ./configs/game/map_builder/mapgen_simple.yaml") - assert exit_status == 0 +def test_gen_basic(): + subprocess.check_call( + [ + "python", + "-m", + "tools.map.gen", + "--show-mode", + "ascii", + "./configs/game/map_builder/mapgen_simple.yaml", + ] + ) - exit_status = os.system("python -m tools.mapgen --show ascii ./configs/game/map_builder/NOT_A_CONFIG.yaml") + +def test_gen_missing_config(): + exit_status = subprocess.call( + [ + "python", + "-m", + "tools.map.gen", + "--show-mode", + "ascii", + "./NON_EXISTENT_CONFIG.yaml", + ] + ) assert exit_status != 0 + + +def test_save(tmpdir): + subprocess.check_call( + [ + "python", + "-m", + "tools.map.gen", + "--output-uri", + tmpdir, + "--show-mode", + "ascii", + "./configs/game/map_builder/mapgen_simple.yaml", + ] + ) + files = os.listdir(tmpdir) + assert len(files) == 1 + assert files[0].endswith(".yaml") + + +def test_save_exact_file(tmpdir): + subprocess.check_call( + [ + "python", + "-m", + "tools.map.gen", + "--output-uri", + tmpdir + "/map.yaml", + "--show-mode", + "ascii", + "./configs/game/map_builder/mapgen_simple.yaml", + ] + ) + files = os.listdir(tmpdir) + assert len(files) == 1 + assert files[0] == "map.yaml" + + +def test_save_multiple(tmpdir): + count = 3 + subprocess.check_call( + [ + "python", + "-m", + "tools.map.gen", + "--output-uri", + tmpdir, + "--show-mode", + "none", + "./configs/game/map_builder/mapgen_maze.yaml", + "--count", + str(count), + ] + ) + files = os.listdir(tmpdir) + assert len(files) == count + for file in files: + assert file.endswith(".yaml") + + +def test_view(tmpdir): + subprocess.check_call( + [ + "python", + "-m", + "tools.map.gen", + "--output-uri", + tmpdir + "/map.yaml", + "--show-mode", + "none", + "./configs/game/map_builder/mapgen_simple.yaml", + ] + ) + subprocess.check_call( + [ + "python", + "-m", + "tools.map.view", + "--show-mode", + "ascii", + tmpdir + "/map.yaml", + ] + ) + + +def test_view_random(tmpdir): + subprocess.check_call( + [ + "python", + "-m", + "tools.map.gen", + "--output-uri", + tmpdir, + "--show-mode", + "none", + "--count", + "3", + "./configs/game/map_builder/mapgen_simple.yaml", + ], + ) + view_output = subprocess.check_output( + ["python", "-m", "tools.map.view", "--show-mode", "ascii", tmpdir], stderr=subprocess.STDOUT, text=True + ) + assert "Loading random map" in view_output diff --git a/tools/map/gen.py b/tools/map/gen.py index 9908e6de..65bc0942 100644 --- a/tools/map/gen.py +++ b/tools/map/gen.py @@ -6,14 +6,13 @@ import string import time from datetime import datetime -from typing import Any, Literal, cast, get_args +from typing import cast, get_args import hydra -import numpy as np from omegaconf import DictConfig, OmegaConf -from mettagrid.map.utils.storable_map import StorableMap, grid_to_ascii -from mettagrid.mettagrid_env import MettaGridEnv +from mettagrid.map.utils.show import ShowMode, show_map +from mettagrid.map.utils.storable_map import StorableMap # Aggressively exit on ctrl+c signal.signal(signal.SIGINT, lambda sig, frame: os._exit(0)) @@ -21,38 +20,6 @@ logger = logging.getLogger(__name__) logging.basicConfig(level=logging.INFO) -ShowMode = Literal["raylib", "ascii", "ascii_border"] - - -def show_map(storable_map: StorableMap, mode: ShowMode | None): - if not mode: - return - - if mode == "raylib": - num_agents = np.count_nonzero(np.char.startswith(storable_map.grid, "agent")) - - env_cfg = OmegaConf.load("configs/mettagrid.yaml") - env_cfg.game.num_agents = num_agents - - env = MettaGridEnv(cast(Any, env_cfg), env_map=storable_map.grid, render_mode="none") - - from mettagrid.renderer.raylib.raylib_renderer import MettaGridRaylibRenderer - - renderer = MettaGridRaylibRenderer(env._c_env, env._env_cfg.game) - while True: - renderer.render_and_wait() - - elif mode == "ascii": - ascii_lines = grid_to_ascii(storable_map.grid) - print("\n".join(ascii_lines)) - - elif mode == "ascii_border": - ascii_lines = grid_to_ascii(storable_map.grid, border=True) - print("\n".join(ascii_lines)) - - else: - raise ValueError(f"Invalid show mode: {mode}") - def make_map(cfg_path: str, overrides: DictConfig | None = None): cfg: DictConfig = cast(DictConfig, OmegaConf.merge(OmegaConf.load(cfg_path), overrides)) @@ -77,6 +44,7 @@ def make_map(cfg_path: str, overrides: DictConfig | None = None): return storable_map +# Based on heuristics, see https://github.com/Metta-AI/mettagrid/pull/108#discussion_r2054699842 def uri_is_file(uri: str) -> bool: last_part = uri.split("/")[-1] return "." in last_part and len(last_part.split(".")[-1]) <= 4 @@ -85,13 +53,13 @@ def uri_is_file(uri: str) -> bool: def main(): parser = argparse.ArgumentParser() parser.add_argument("--output-uri", type=str, help="Output URI") - parser.add_argument("--show", choices=get_args(ShowMode), help="Show the map in the specified mode") + parser.add_argument("--show-mode", choices=get_args(ShowMode), help="Show the map in the specified mode") parser.add_argument("--count", type=int, default=1, help="Number of maps to generate") parser.add_argument("--overrides", type=str, default="", help="OmniConf overrides for the map config") parser.add_argument("cfg_path", type=str, help="Path to the map config file") args = parser.parse_args() - show_mode = args.show + show_mode = args.show_mode if not show_mode and not args.output_uri: # if not asked to save, show the map show_mode = "raylib" diff --git a/tools/map/view.py b/tools/map/view.py new file mode 100644 index 00000000..baf08cb6 --- /dev/null +++ b/tools/map/view.py @@ -0,0 +1,36 @@ +import argparse +import logging +from typing import get_args + +from mettagrid.map.load_random import get_random_map_uri +from mettagrid.map.utils.show import ShowMode, show_map +from mettagrid.map.utils.storable_map import StorableMap +from tools.map.gen import uri_is_file + +logger = logging.getLogger(__name__) +logging.basicConfig(level=logging.INFO) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument( + "--show-mode", choices=get_args(ShowMode), help="Show the map in the specified mode", default="raylib" + ) + parser.add_argument("uri", type=str, help="URI of the map to view") + args = parser.parse_args() + + uri = args.uri + + if not uri_is_file(uri): + # probably a directory + logger.info(f"Loading random map from directory {uri}") + uri = get_random_map_uri(uri) + + logger.info(f"Loading map from {uri}") + storable_map = StorableMap.from_uri(uri) + + show_map(storable_map, args.show_mode) + + +if __name__ == "__main__": + main() From 8b8baa435b022aaeed588cd0b5914381b4c22e0d Mon Sep 17 00:00:00 2001 From: Vyacheslav Matyukhin Date: Tue, 22 Apr 2025 19:21:51 -0300 Subject: [PATCH 19/19] normalize uri --- tools/map/view.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tools/map/view.py b/tools/map/view.py index baf08cb6..8ef00fc0 100644 --- a/tools/map/view.py +++ b/tools/map/view.py @@ -23,6 +23,8 @@ def main(): if not uri_is_file(uri): # probably a directory + while uri.endswith("/"): + uri = uri[:-1] logger.info(f"Loading random map from directory {uri}") uri = get_random_map_uri(uri)