diff --git a/configs/game/map_builder/load.yaml b/configs/game/map_builder/load.yaml new file mode 100644 index 00000000..c5c24ee5 --- /dev/null +++ b/configs/game/map_builder/load.yaml @@ -0,0 +1,7 @@ +_target_: mettagrid.map.load.Load + +uri: ??? + +# Optional scene to render on top of the loaded map. +extra_root: + _target_: mettagrid.map.scenes.make_connected.MakeConnected diff --git a/configs/game/map_builder/load_random.yaml b/configs/game/map_builder/load_random.yaml new file mode 100644 index 00000000..9d78f5ec --- /dev/null +++ b/configs/game/map_builder/load_random.yaml @@ -0,0 +1,18 @@ +_target_: mettagrid.map.load_random.LoadRandom + +dir: ??? + +# Optional scene to render on top of the loaded map. +# The following example shows how to patch the map if its agent count doesn't match `game.num_agents`. +extra_root: + _target_: mettagrid.map.scenes.nop.Nop + + children: + - where: full + scene: + _target_: mettagrid.map.scenes.remove_agents.RemoveAgents + + - where: full + scene: + _target_: mettagrid.map.scenes.random.Random + agents: 40 diff --git a/configs/game/map_builder/mapgen_auto.yaml b/configs/game/map_builder/mapgen_auto.yaml index 657e562b..8b3e73e7 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?) @@ -70,14 +70,14 @@ root: _target_: mettagrid.map.scenes.random_scene_from_dir.RandomSceneFromDir dir: /scenes/dcss/wfc weight: 20 - - scene: /scenes/wfc/blob - - scene: /scenes/wfc/blob2 - - scene: /scenes/wfc/blob3 - - scene: /scenes/wfc/blocks - - scene: /scenes/wfc/dungeons - - scene: /scenes/wfc/mazelike1 - - scene: /scenes/wfc/mazelike2 - - scene: /scenes/wfc/simple + - scene: ./configs/scenes/wfc/blob.yaml + - scene: ./configs/scenes/wfc/blob2.yaml + - scene: ./configs/scenes/wfc/blob3.yaml + - scene: ./configs/scenes/wfc/blocks.yaml + - scene: ./configs/scenes/wfc/dungeons.yaml + - scene: ./configs/scenes/wfc/mazelike1.yaml + - scene: ./configs/scenes/wfc/mazelike2.yaml + - scene: ./configs/scenes/wfc/simple.yaml - scene: _target_: mettagrid.map.scenes.maze.MazeKruskal room_size: ["uniform", 1, 3] 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_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/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/configs/mapgen.yaml b/configs/mapgen.yaml deleted file mode 100644 index fc00f1b6..00000000 --- a/configs/mapgen.yaml +++ /dev/null @@ -1,17 +0,0 @@ -defaults: - - mettagrid - - 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/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/docs/mapgen.md b/docs/mapgen.md new file mode 100644 index 00000000..e521ae4c --- /dev/null +++ b/docs/mapgen.md @@ -0,0 +1,61 @@ +# Map generation + +## S3 maps + +To produce maps in bulk and store them in S3, use the following commands: + +### Creating maps + +```bash +python -m tools.map.gen --output-uri=s3://BUCKET/DIR ./configs/game/map_builder/mapgen_auto.yaml +``` + +`mapgen_auto` builder is an example. You can use any YAML config that can be parsed by OmegaConf. + +If `--output-uri` looks like a file (ends with `.yaml` or other extension), the map will be saved to that file. + +Otherwise, the map will be saved to a file with a random suffix in that directory. + +If `--output-uri` is not specified, the map won't be saved, only shown on screen. + +To create maps in bulk, use `--count=N` option. + +See `python -m tools.map.gen --help` for more options. + +### Viewing maps + +You can view a single map by running: + +```bash +python -m tools.map.view s3://BUCKET/PATH/TO/MAP.yaml +``` + +The following command will show a random map from an S3 directory: + +```bash +python -m tools.map.view s3://BUCKET/DIR +``` + +Same heuristics about detecting if the URI is a file apply here. + +### Loading maps in map_builder configs + +You can load a random map from an S3 directory in your YAML configs by using `mettagrid.map.load_random.LoadRandom` as a map builder. + +`LoadRandom` allows you to modify the map by applying additional scenes to it. Check out `configs/game/map_builder/load_random.yaml` for an example config that modifies the number of agents in the map. + +### Indexing maps + +Optionally, you can index your maps to make loading them faster. + +This is intended to speed up reading from S3. It shouldn't change any functionality, and you should skip playing with this unless you find map loading from S3 is slow. + +Index is a plain text file that lists URIs of all the maps. You can assemble it manually, or use the following script: + +```bash +python -m tools.index_s3_maps --dir=s3://BUCKET/DIR --target=s3://BUCKET/DIR/index.txt +``` + +`--target` is optional. If not provided, the index will be saved to `{--dir}/index.txt`. + +You can then use `mettagrid.map.load_random_from_index.LoadRandomFromIndex` to load a random map from the index. diff --git a/mettagrid/map/load.py b/mettagrid/map/load.py new file mode 100644 index 00000000..af18e3c6 --- /dev/null +++ b/mettagrid/map/load.py @@ -0,0 +1,32 @@ +from mettagrid.config.room.room import Room +from mettagrid.map.utils.storable_map import StorableMap + +from .scene import SceneCfg, make_scene + + +# Note that this class can't be a scene, because the width and height come from the stored data. +class Load(Room): + """ + Load a pregenerated map from a URI (file or S3 object). + + See also: `FromS3Dir` for picking a random map from a directory of pregenerated maps. + """ + + def __init__(self, uri: str, extra_root: SceneCfg | None = None): + super().__init__() + self._uri = uri + self._storable_map = StorableMap.from_uri(uri) + + if extra_root is not None: + self._root = make_scene(extra_root) + else: + self._root = None + + def build(self): + grid = self._storable_map.grid + + if self._root is not None: + root_node = self._root.make_node(grid) + root_node.render() + + return grid diff --git a/mettagrid/map/load_random.py b/mettagrid/map/load_random.py new file mode 100644 index 00000000..77257078 --- /dev/null +++ b/mettagrid/map/load_random.py @@ -0,0 +1,39 @@ +import os +import random +from pathlib import Path + +from mettagrid.map.load import Load +from mettagrid.map.utils import s3utils +from mettagrid.map.utils.storage import parse_file_uri + +from .scene import SceneCfg + + +def get_random_map_uri(dir_uri: str) -> str: + if dir_uri.startswith("s3://"): + filenames = s3utils.list_objects(dir_uri) + filenames = [uri for uri in filenames if uri.endswith(".yaml")] + return random.choice(filenames) + else: + dirname = parse_file_uri(dir_uri) + if not os.path.isdir(dirname): + raise ValueError(f"Directory {dirname} does not exist") + + filenames = os.listdir(dirname) + filenames = [Path(dirname) / Path(filename) for filename in filenames if filename.endswith(".yaml")] + return str(random.choice(filenames)) + + +class LoadRandom(Load): + """ + Load a random map from a directory, local or S3. + + See also: `LoadRandomFromIndex` for a version that loads a random map from a pre-generated index. + """ + + def __init__(self, dir: str, extra_root: SceneCfg | None = None): + self._dir_uri = dir + + random_map_uri = get_random_map_uri(self._dir_uri) + + super().__init__(random_map_uri, extra_root) diff --git a/mettagrid/map/load_random_from_index.py b/mettagrid/map/load_random_from_index.py new file mode 100644 index 00000000..3d800c91 --- /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 storage + +from .scene import SceneCfg + + +class LoadRandomFromIndex(Load): + """ + Load a random map from a list of pregenerated maps. + + The index file can be produced with the following command: + python -m tools.index_s3_maps --dir=s3://... + + See also: `LoadRandom` for a version that loads a random map from an S3 directory. + """ + + def __init__(self, index_uri: str, extra_root: SceneCfg | None = None): + self._index_uri = index_uri + + # For 10k maps in a directory we'd have to fetch 100Kb of index data. + # (Can we optimize this further by caching?) + index = storage.load_from_uri(self._index_uri) + index = index.split("\n") + random_map_uri = random.choice(index) + + super().__init__(random_map_uri, extra_root) diff --git a/mettagrid/map/mapgen.py b/mettagrid/map/mapgen.py index a4064f7b..ef5809d1 100644 --- a/mettagrid/map/mapgen.py +++ b/mettagrid/map/mapgen.py @@ -1,11 +1,19 @@ -# 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 -class MapGen: 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/node.py b/mettagrid/map/node.py index 81d932b6..b569e0d1 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, 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/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/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/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/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" diff --git a/mettagrid/map/utils/s3utils.py b/mettagrid/map/utils/s3utils.py new file mode 100644 index 00000000..a1c2376f --- /dev/null +++ b/mettagrid/map/utils/s3utils.py @@ -0,0 +1,44 @@ +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") + + (bucket, key) = uri[5:].split("/", 1) + return bucket, key + + +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 + # (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 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/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/mettagrid/map/utils/storable_map.py b/mettagrid/map/utils/storable_map.py new file mode 100644 index 00000000..50f0d74d --- /dev/null +++ b/mettagrid/map/utils/storable_map.py @@ -0,0 +1,128 @@ +import logging +from dataclasses import dataclass + +import numpy as np +from omegaconf import DictConfig, OmegaConf + +from mettagrid.map.mapgen import MapGrid + +from . import storage + +logger = logging.getLogger(__name__) + +ascii_symbols = { + "empty": " ", + "wall": "#", + "agent.agent": "A", + "mine": "g", + "generator": "c", + "altar": "a", + "armory": "r", + "lasery": "l", + "lab": "b", + "factory": "f", + "temple": "t", +} + +reverse_ascii_symbols = {v: k for k, v in ascii_symbols.items()} + + +def grid_object_to_ascii(name: str) -> 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 = storage.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): + 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/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.py b/tests/mapgen.py deleted file mode 100644 index 587b87b5..00000000 --- a/tests/mapgen.py +++ /dev/null @@ -1,92 +0,0 @@ -import os -import random -import signal # Aggressively exit on ctrl+c -import string -import time -from datetime import datetime - -import hydra -from omegaconf import OmegaConf - -from mettagrid.mettagrid_env import MettaGridEnv - -signal.signal(signal.SIGINT, lambda sig, frame: os._exit(0)) - - -def env_to_ascii(env): - grid = env._c_env.render_ascii() - # convert to strings - return ["".join(row) for row in grid] - - -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(), - } - - 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") - - print(f"Saved map to {target_file}") - - -@hydra.main(version_base=None, config_path="../configs", config_name="mapgen") -def main(cfg): - start = time.time() - env = MettaGridEnv(cfg, render_mode="human") - gen_time = time.time() - start - print(f"Time taken to create env: {gen_time} seconds") - - 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 - - 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}") - - -if __name__ == "__main__": - main() diff --git a/tests/mapgen/test_cli.py b/tests/mapgen/test_cli.py new file mode 100644 index 00000000..c1528810 --- /dev/null +++ b/tests/mapgen/test_cli.py @@ -0,0 +1,133 @@ +import os +import subprocess + + +def test_gen_basic(): + subprocess.check_call( + [ + "python", + "-m", + "tools.map.gen", + "--show-mode", + "ascii", + "./configs/game/map_builder/mapgen_simple.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 new file mode 100644 index 00000000..65bc0942 --- /dev/null +++ b/tools/map/gen.py @@ -0,0 +1,115 @@ +import argparse +import logging +import os +import random +import signal +import string +import time +from datetime import datetime +from typing import cast, get_args + +import hydra +from omegaconf import DictConfig, OmegaConf + +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)) + +logger = logging.getLogger(__name__) +logging.basicConfig(level=logging.INFO) + + +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() + map_builder = hydra.utils.instantiate(cfg, _recursive_=False) + grid = map_builder.build() + gen_time = time.time() - start + 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 + + +# 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 + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--output-uri", type=str, help="Output URI") + 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_mode + if not show_mode and not args.output_uri: + # if not asked to save, show the map + show_mode = "raylib" + + 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 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") + + # 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] + + 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") + + def make_output_uri() -> str | None: + if not output_uri: + return None # the map won't be saved + + if output_is_file: + return 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: + 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__": + main() diff --git a/tools/map/index_s3.py b/tools/map/index_s3.py new file mode 100644 index 00000000..3786c571 --- /dev/null +++ b/tools/map/index_s3.py @@ -0,0 +1,30 @@ +import argparse +import logging +import os +import signal # Aggressively exit on ctrl+c + +from mettagrid.map.utils import s3utils, storage + +signal.signal(signal.SIGINT, lambda sig, frame: os._exit(0)) + +logger = logging.getLogger(__name__) +logging.basicConfig(level=logging.INFO) + + +def main(): + parser = argparse.ArgumentParser() + 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() + + s3_dir = args.dir + target = args.target or f"{s3_dir}/index.txt" + + uri_list = s3utils.list_objects(s3_dir) + + storage.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__": + main() diff --git a/tools/map/view.py b/tools/map/view.py new file mode 100644 index 00000000..8ef00fc0 --- /dev/null +++ b/tools/map/view.py @@ -0,0 +1,38 @@ +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 + while uri.endswith("/"): + uri = uri[:-1] + 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()