mettagrid 0.2.0.3__cp311-cp311-macosx_11_0_arm64.whl

mettagrid/__init__.py

@@ -0,0 +1,62 @@
+"""
+MettaGrid - Multi-agent reinforcement learning grid environments.
+
+This module provides various environment adapters for different RL frameworks:
+- MettaGridCore: Core C++ wrapper (no training features)
+- MettaGridEnv: Training environment (PufferLib-based with stats/replay)
+- MettaGridGymEnv: Gymnasium adapter
+- MettaGridPettingZooEnv: PettingZoo adapter
+
+All adapters inherit from MettaGridCore and provide framework-specific interfaces.
+For PufferLib integration, use PufferLib's MettaPuff wrapper directly.
+"""
+
+from __future__ import annotations
+
+from mettagrid.config.mettagrid_config import MettaGridConfig
+
+# Import environment classes
+from mettagrid.core import MettaGridCore
+
+# Import other commonly used classes
+from mettagrid.envs.gym_env import MettaGridGymEnv
+from mettagrid.envs.mettagrid_env import MettaGridEnv
+from mettagrid.envs.pettingzoo_env import MettaGridPettingZooEnv
+from mettagrid.map_builder.map_builder import GameMap
+
+# Import data types from C++ module (source of truth)
+from mettagrid.mettagrid_c import (
+    dtype_actions,
+    dtype_masks,
+    dtype_observations,
+    dtype_rewards,
+    dtype_success,
+    dtype_terminals,
+    dtype_truncations,
+)
+from mettagrid.util.replay_writer import ReplayWriter
+from mettagrid.util.stats_writer import StatsWriter
+
+__all__ = [
+    # Config
+    "MettaGridConfig",
+    # Core classes
+    "MettaGridCore",
+    # Main environment (backward compatible)
+    "MettaGridEnv",
+    # Environment adapters
+    "MettaGridGymEnv",
+    "MettaGridPettingZooEnv",
+    # Data types
+    "dtype_actions",
+    "dtype_observations",
+    "dtype_rewards",
+    "dtype_terminals",
+    "dtype_truncations",
+    "dtype_masks",
+    "dtype_success",
+    # Supporting classes
+    "GameMap",
+    "ReplayWriter",
+    "StatsWriter",
+]

mettagrid/builder/__init__.py

@@ -0,0 +1,7 @@
+"""
+Configuration builders for Metta environments.
+"""
+
+from . import building, empty_converters, envs
+
+__all__ = ["building", "envs", "empty_converters"]

mettagrid/builder/building.py

@@ -0,0 +1,136 @@
+from mettagrid.config.mettagrid_config import AssemblerConfig, ConverterConfig, RecipeConfig, WallConfig
+
+wall = WallConfig(type_id=1)
+block = WallConfig(type_id=14, swappable=True)
+
+altar = ConverterConfig(
+    type_id=8,
+    input_resources={"battery_red": 3},
+    output_resources={"heart": 1},
+    cooldown=10,
+)
+
+
+def make_mine(color: str, type_id: int) -> ConverterConfig:
+    return ConverterConfig(
+        type_id=type_id,
+        output_resources={f"ore_{color}": 1},
+        cooldown=50,
+    )
+
+
+mine_red = make_mine("red", 2)
+mine_blue = make_mine("blue", 3)
+mine_green = make_mine("green", 4)
+
+
+def make_generator(color: str, type_id: int) -> ConverterConfig:
+    return ConverterConfig(
+        type_id=type_id,
+        input_resources={f"ore_{color}": 1},
+        output_resources={f"battery_{color}": 1},
+        cooldown=25,
+    )
+
+
+generator_red = make_generator("red", 5)
+generator_blue = make_generator("blue", 6)
+generator_green = make_generator("green", 7)
+
+lasery = ConverterConfig(
+    type_id=15,
+    input_resources={"battery_red": 1, "ore_red": 2},
+    output_resources={"laser": 1},
+    cooldown=10,
+)
+
+armory = ConverterConfig(
+    type_id=16,
+    input_resources={"ore_red": 3},
+    output_resources={"armor": 1},
+    cooldown=10,
+)
+
+# Assembler building definitions
+assembler_altar = AssemblerConfig(
+    type_id=8,
+    recipes=[
+        (
+            ["Any"],
+            RecipeConfig(
+                input_resources={"battery_red": 3},
+                output_resources={"heart": 1},
+                cooldown=10,
+            ),
+        )
+    ],
+)
+
+
+def make_assembler_mine(color: str, type_id: int) -> AssemblerConfig:
+    return AssemblerConfig(
+        type_id=type_id,
+        recipes=[
+            (
+                ["Any"],
+                RecipeConfig(
+                    output_resources={f"ore_{color}": 1},
+                    cooldown=50,
+                ),
+            )
+        ],
+    )
+
+
+assembler_mine_red = make_assembler_mine("red", 2)
+assembler_mine_blue = make_assembler_mine("blue", 3)
+assembler_mine_green = make_assembler_mine("green", 4)
+
+
+def make_assembler_generator(color: str, type_id: int) -> AssemblerConfig:
+    return AssemblerConfig(
+        type_id=type_id,
+        recipes=[
+            (
+                ["Any"],
+                RecipeConfig(
+                    input_resources={f"ore_{color}": 1},
+                    output_resources={f"battery_{color}": 1},
+                    cooldown=25,
+                ),
+            )
+        ],
+    )
+
+
+assembler_generator_red = make_assembler_generator("red", 5)
+assembler_generator_blue = make_assembler_generator("blue", 6)
+assembler_generator_green = make_assembler_generator("green", 7)
+
+assembler_lasery = AssemblerConfig(
+    type_id=15,
+    recipes=[
+        (
+            ["Any"],
+            RecipeConfig(
+                input_resources={"battery_red": 1, "ore_red": 2},
+                output_resources={"laser": 1},
+                cooldown=10,
+            ),
+        )
+    ],
+)
+
+assembler_armory = AssemblerConfig(
+    type_id=16,
+    recipes=[
+        (
+            ["Any"],
+            RecipeConfig(
+                input_resources={"ore_red": 3},
+                output_resources={"armor": 1},
+                cooldown=10,
+            ),
+        )
+    ],
+)

mettagrid/builder/empty_converters.py

@@ -0,0 +1,85 @@
+from mettagrid.config.mettagrid_config import ConverterConfig, WallConfig
+
+wall = WallConfig(type_id=1)
+block = WallConfig(type_id=14, swappable=True)
+
+mine_red = ConverterConfig(
+    type_id=2,
+    input_resources={},
+    output_resources={},
+    cooldown=5,
+)
+
+mine_blue = ConverterConfig(
+    type_id=3,
+    input_resources={},
+    output_resources={},
+    cooldown=5,
+)
+
+mine_green = ConverterConfig(
+    type_id=4,
+    input_resources={},
+    output_resources={},
+    cooldown=5,
+)
+
+generator_red = ConverterConfig(
+    type_id=5,
+    input_resources={},
+    output_resources={},
+    cooldown=5,
+)
+
+generator_blue = ConverterConfig(
+    type_id=6,
+    input_resources={},
+    output_resources={},
+    cooldown=5,
+)
+
+generator_green = ConverterConfig(
+    type_id=7,
+    input_resources={},
+    output_resources={},
+    cooldown=5,
+)
+
+altar = ConverterConfig(
+    type_id=8,
+    input_resources={},
+    output_resources={},
+    cooldown=5,
+)
+
+
+lasery = ConverterConfig(
+    type_id=15,
+    input_resources={},
+    output_resources={},
+    cooldown=5,
+)
+armory = ConverterConfig(
+    type_id=16,
+    input_resources={},
+    output_resources={},
+    cooldown=5,
+)
+lab = ConverterConfig(
+    type_id=17,
+    input_resources={},
+    output_resources={},
+    cooldown=5,
+)
+factory = ConverterConfig(
+    type_id=18,
+    input_resources={},
+    output_resources={},
+    cooldown=5,
+)
+temple = ConverterConfig(
+    type_id=19,
+    input_resources={},
+    output_resources={},
+    cooldown=5,
+)

mettagrid/builder/envs.py

@@ -0,0 +1,221 @@
+from typing import Optional
+
+import mettagrid.mapgen.scenes.random
+from mettagrid.config.mettagrid_config import (
+    ActionConfig,
+    ActionsConfig,
+    AgentConfig,
+    AgentRewards,
+    AttackActionConfig,
+    GameConfig,
+    MettaGridConfig,
+)
+from mettagrid.map_builder.map_builder import MapBuilderConfig
+from mettagrid.map_builder.perimeter_incontext import PerimeterInContextMapBuilder
+from mettagrid.map_builder.random import RandomMapBuilder
+from mettagrid.mapgen.mapgen import MapGen
+
+from . import building, empty_converters
+
+
+def make_arena(
+    num_agents: int,
+    combat: bool = True,
+    map_builder: MapBuilderConfig | None = None,  # custom map builder; must match num_agents
+) -> MettaGridConfig:
+    objects = {
+        "wall": building.wall,
+        "altar": building.altar,
+        "mine_red": building.mine_red,
+        "generator_red": building.generator_red,
+        "lasery": building.lasery,
+        "armory": building.armory,
+    }
+
+    actions = ActionsConfig(
+        noop=ActionConfig(),
+        move=ActionConfig(),
+        rotate=ActionConfig(enabled=False),  # Disabled for unified movement system
+        put_items=ActionConfig(),
+        get_items=ActionConfig(),
+        attack=AttackActionConfig(
+            consumed_resources={
+                "laser": 1,
+            },
+            defense_resources={
+                "armor": 1,
+            },
+        ),
+        swap=ActionConfig(enabled=False),
+        change_color=ActionConfig(enabled=False),
+    )
+
+    if not combat:
+        actions.attack.consumed_resources = {"laser": 100}
+
+    if map_builder is None:
+        map_builder = MapGen.Config(
+            num_agents=num_agents,
+            width=25,
+            height=25,
+            border_width=6,
+            instance_border_width=0,
+            root=mettagrid.mapgen.scenes.random.Random.factory(
+                params=mettagrid.mapgen.scenes.random.Random.Params(
+                    agents=6,
+                    objects={
+                        "wall": 10,
+                        "altar": 5,
+                        "mine_red": 10,
+                        "generator_red": 5,
+                        "lasery": 1,
+                        "armory": 1,
+                    },
+                ),
+            ),
+        )
+
+    return MettaGridConfig(
+        label="arena" + (".combat" if combat else ""),
+        game=GameConfig(
+            num_agents=num_agents,
+            actions=actions,
+            objects=objects,
+            agent=AgentConfig(
+                default_resource_limit=50,
+                resource_limits={
+                    "heart": 255,
+                },
+                rewards=AgentRewards(
+                    inventory={
+                        "heart": 1,
+                    },
+                ),
+            ),
+            map_builder=map_builder,
+        ),
+    )
+
+
+def make_navigation(num_agents: int) -> MettaGridConfig:
+    altar = empty_converters.altar.model_copy()
+    altar.cooldown = 255  # Maximum cooldown
+    altar.initial_resource_count = 1
+    altar.max_conversions = 0
+    altar.input_resources = {}
+    altar.output_resources = {"heart": 1}
+    cfg = MettaGridConfig(
+        game=GameConfig(
+            num_agents=num_agents,
+            objects={
+                "altar": altar,
+                "wall": building.wall,
+            },
+            resource_names=["heart"],
+            actions=ActionsConfig(
+                move=ActionConfig(),
+                rotate=ActionConfig(enabled=False),
+                get_items=ActionConfig(),
+            ),
+            agent=AgentConfig(
+                rewards=AgentRewards(
+                    inventory={
+                        "heart": 1,
+                    },
+                ),
+            ),
+            # Always provide a concrete map builder config so tests can set width/height
+            map_builder=RandomMapBuilder.Config(agents=num_agents),
+        )
+    )
+    return cfg
+
+
+def make_navigation_sequence(num_agents: int) -> MettaGridConfig:
+    altar = building.altar.model_copy()
+    altar.input_resources = {"battery_red": 1}
+    altar.cooldown = 15
+    mine = building.mine_red.model_copy()
+    mine.cooldown = 15
+    generator = building.generator_red.model_copy()
+    generator.cooldown = 15
+    cfg = MettaGridConfig(
+        game=GameConfig(
+            num_agents=num_agents,
+            objects={
+                "altar": altar,
+                "wall": building.wall,
+                "mine_red": mine,
+                "generator_red": generator,
+            },
+            resource_names=["heart", "ore_red", "battery_red"],
+            actions=ActionsConfig(
+                move=ActionConfig(),
+                rotate=ActionConfig(enabled=False),
+                get_items=ActionConfig(),
+            ),
+            agent=AgentConfig(
+                rewards=AgentRewards(
+                    inventory={
+                        "heart": 1,
+                        "ore_red": 0.001,
+                        "battery_red": 0.01,
+                    },
+                ),
+                default_resource_limit=1,
+                resource_limits={
+                    "heart": 100,
+                },
+            ),
+            # Always provide a concrete map builder config so tests can set width/height
+            map_builder=RandomMapBuilder.Config(agents=num_agents),
+        )
+    )
+    return cfg
+
+
+def make_icl_resource_chain(
+    num_agents: int,
+    max_steps,
+    game_objects: dict,
+    map_builder_objects: dict,
+    width: int = 6,
+    height: int = 6,
+    obstacle_type: Optional[str] = None,
+    density: Optional[str] = None,
+) -> MettaGridConfig:
+    game_objects["wall"] = empty_converters.wall
+    cfg = MettaGridConfig(
+        game=GameConfig(
+            max_steps=max_steps,
+            num_agents=num_agents,
+            objects=game_objects,
+            map_builder=MapGen.Config(
+                instances=num_agents,
+                instance_map=PerimeterInContextMapBuilder.Config(
+                    agents=1,
+                    width=width,
+                    height=height,
+                    objects=map_builder_objects,
+                    obstacle_type=obstacle_type,
+                    density=density,
+                ),
+            ),
+            actions=ActionsConfig(
+                move=ActionConfig(),
+                rotate=ActionConfig(enabled=False),  # Disabled for unified movement system
+                get_items=ActionConfig(),
+                put_items=ActionConfig(),
+            ),
+            agent=AgentConfig(
+                rewards=AgentRewards(
+                    inventory={
+                        "heart": 1,
+                    },
+                ),
+                default_resource_limit=1,
+                resource_limits={"heart": 15},
+            ),
+        )
+    )
+    return cfg

mettagrid/config/__init__.py

@@ -0,0 +1,35 @@
+"""Configuration module for mettagrid."""
+
+from .config import Config
+from .mettagrid_c_config import from_mettagrid_config
+from .mettagrid_config import (
+    ActionConfig,
+    ActionsConfig,
+    AgentConfig,
+    AgentRewards,
+    AttackActionConfig,
+    ChangeGlyphActionConfig,
+    ConverterConfig,
+    GameConfig,
+    GlobalObsConfig,
+    MettaGridConfig,
+    StatsRewards,
+    WallConfig,
+)
+
+__all__ = [
+    "Config",
+    "from_mettagrid_config",
+    "MettaGridConfig",
+    "ActionConfig",
+    "ActionsConfig",
+    "AgentConfig",
+    "AgentRewards",
+    "AttackActionConfig",
+    "ChangeGlyphActionConfig",
+    "ConverterConfig",
+    "GameConfig",
+    "GlobalObsConfig",
+    "StatsRewards",
+    "WallConfig",
+]

mettagrid/config/config.py

@@ -0,0 +1,106 @@
+from __future__ import annotations
+
+from typing import Any, NoReturn, Self, Union, get_args, get_origin
+
+from pydantic import BaseModel, ConfigDict, TypeAdapter
+
+
+class Config(BaseModel):
+    """
+    Common extension of Pydantic's BaseModel that:
+    - sets `extra="forbid"` by default
+    - adds `override` and `update` methods for overriding values based on `path.to.value` keys
+    """
+
+    model_config = ConfigDict(extra="forbid")
+
+    def _auto_initialize_field(self, parent_obj: "Config", field_name: str) -> "Config | None":
+        """Auto-initialize a None Config field if possible."""
+        field = type(parent_obj).model_fields.get(field_name)
+        if not field:
+            return None
+
+        field_type = self._unwrap_optional(field.annotation)
+        if not (isinstance(field_type, type) and issubclass(field_type, Config)):
+            return None
+
+        try:
+            new_instance = field_type()
+            setattr(parent_obj, field_name, new_instance)
+            return new_instance
+        except (TypeError, ValueError):
+            return None
+
+    def _unwrap_optional(self, field_type):
+        """Unwrap Optional[T] → T if applicable, else return original type."""
+        if get_origin(field_type) is Union:
+            non_none_types = [arg for arg in get_args(field_type) if arg is not type(None)]
+            return non_none_types[0] if len(non_none_types) == 1 else field_type
+        return field_type
+
+    def override(self, key: str, value: Any) -> Self:
+        """Override a value in the config."""
+        key_path = key.split(".")
+
+        def fail(error: str) -> NoReturn:
+            raise ValueError(
+                f"Override failed. Full config:\n {self.model_dump_json(indent=2)}\nOverride {key} failed: {error}"
+            )
+
+        inner_cfg: Config | dict[str, Any] = self
+        traversed_path: list[str] = []
+        for key_part in key_path[:-1]:
+            if isinstance(inner_cfg, dict):
+                if key_part not in inner_cfg:
+                    fail(f"key {key} not found")
+                inner_cfg = inner_cfg[key_part]
+                traversed_path.append(key_part)
+                continue
+
+            if not hasattr(inner_cfg, key_part):
+                failed_path = ".".join(traversed_path + [key_part])
+                fail(f"key {failed_path} not found")
+
+            next_inner_cfg = getattr(inner_cfg, key_part)
+            if next_inner_cfg is None:
+                # Auto-initialize None Config fields
+                next_inner_cfg = self._auto_initialize_field(inner_cfg, key_part)
+                if next_inner_cfg is None:
+                    failed_path = ".".join(traversed_path + [key_part])
+                    fail(f"Cannot auto-initialize None field {failed_path}")
+
+            if not isinstance(next_inner_cfg, (Config, dict)):
+                failed_path = ".".join(traversed_path + [key_part])
+                fail(f"key {failed_path} is not a Config object")
+
+            inner_cfg = next_inner_cfg
+            traversed_path.append(key_part)
+
+        # We allow dicts to get new keys, but not Configs. This is because we want to allow overrides like
+        # env_cfg.game.agent.rewards.inventory.ore_red = 0.1
+        # without requiring that "ore_red" was already in the inventory dict. Note that allowing overrides / updates
+        # to dicts like this leads to an obnoxious inconsistency in the way dicts are updated via overrides
+        # (foo.bar.baz = 1) vs how they're set in Python (foo.bar["baz"] = 1).
+        if isinstance(inner_cfg, Config):
+            if not hasattr(inner_cfg, key_path[-1]):
+                fail(f"key {key} not found")
+
+        if isinstance(inner_cfg, dict):
+            inner_cfg[key_path[-1]] = value
+            return self
+
+        cls = type(inner_cfg)
+        field = cls.model_fields.get(key_path[-1])
+        if field is None:
+            fail(f"key {key} is not a valid field")
+
+        value = TypeAdapter(field.annotation).validate_python(value)
+        setattr(inner_cfg, key_path[-1], value)
+
+        return self
+
+    def update(self, updates: dict[str, Any]) -> Self:
+        """Applies multiple overrides to the config."""
+        for key, value in updates.items():
+            self.override(key, value)
+        return self