mettagrid 0.0.0__cp311-cp311-macosx_11_0_arm64.whl

mettagrid/__init__.py

@@ -0,0 +1,70 @@
+"""
+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
+
+# Import mettascope submodule (optional, for visualization)
+try:
+    from mettagrid import mettascope
+except ImportError:
+    mettascope = None
+
+__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",
+    # Optional visualization module
+    "mettascope",
+]

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,176 @@
+from typing import Literal
+
+from mettagrid.config.mettagrid_config import AssemblerConfig, ChestConfig, 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,
+            ),
+        )
+    ],
+)
+
+
+# Chest building definitions
+def make_chest(
+    resource_type: str,
+    type_id: int,
+    deposit_positions: list[Literal["NW", "N", "NE", "W", "E", "SW", "S", "SE"]] | None = None,
+    withdrawal_positions: list[Literal["NW", "N", "NE", "W", "E", "SW", "S", "SE"]] | None = None,
+) -> ChestConfig:
+    """Create a chest configuration for a specific resource type."""
+    if deposit_positions is None:
+        deposit_positions = ["N", "S", "E", "W"]  # Default to cardinal directions
+    if withdrawal_positions is None:
+        withdrawal_positions = ["N", "S", "E", "W"]  # Default to cardinal directions
+
+    return ChestConfig(
+        type_id=type_id,
+        resource_type=resource_type,
+        deposit_positions=deposit_positions,
+        withdrawal_positions=withdrawal_positions,
+    )
+
+
+# Example chest configurations
+chest_ore_red = make_chest("ore_red", 20)
+chest_ore_blue = make_chest("ore_blue", 21)
+chest_ore_green = make_chest("ore_green", 22)
+chest_battery_red = make_chest("battery_red", 23)
+chest_battery_blue = make_chest("battery_blue", 24)
+chest_battery_green = make_chest("battery_green", 25)
+
+# Special chest that only allows deposits from North and withdrawals from South
+chest_depot = ChestConfig(
+    type_id=26,
+    resource_type="ore_red",
+    deposit_positions=["N"],
+    withdrawal_positions=["S"],
+)

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,261 @@
+from typing import Optional
+
+import mettagrid.mapgen.scenes.random
+
+# Local import moved to factory usage to avoid forbidden cross-package dependency at import time
+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_in_context_chains(
+    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
+
+
+def make_icl_with_numpy(
+    num_agents: int,
+    num_instances: int,
+    max_steps,
+    game_objects: dict,
+    instance_map: MapBuilderConfig,
+) -> MettaGridConfig:
+    game_objects["wall"] = empty_converters.wall
+    cfg = MettaGridConfig(
+        game=GameConfig(
+            max_steps=max_steps,
+            num_agents=num_agents * num_instances,
+            objects=game_objects,
+            map_builder=MapGen.Config(
+                instances=num_instances,
+                instance_map=instance_map,
+            ),
+            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 @@
+from .config import Config
+from .mettagrid_c_config import from_mettagrid_config
+from .mettagrid_config import (
+    ActionConfig,
+    ActionsConfig,
+    AgentConfig,
+    AgentRewards,
+    AssemblerConfig,
+    AttackActionConfig,
+    ChangeGlyphActionConfig,
+    ConverterConfig,
+    GameConfig,
+    GlobalObsConfig,
+    MettaGridConfig,
+    RecipeConfig,
+    WallConfig,
+)
+
+__all__ = [
+    "ActionConfig",
+    "ActionsConfig",
+    "AgentConfig",
+    "AgentRewards",
+    "AssemblerConfig",
+    "AttackActionConfig",
+    "ChangeGlyphActionConfig",
+    "Config",
+    "ConverterConfig",
+    "from_mettagrid_config",
+    "GameConfig",
+    "GlobalObsConfig",
+    "MettaGridConfig",
+    "RecipeConfig",
+    "WallConfig",
+]