mettagrid 0.2.0.15__cp311-cp311-manylinux_2_34_x86_64.whl

mettagrid/__init__.py

@@ -0,0 +1,146 @@
+"""
+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
+
+import importlib
+from typing import TYPE_CHECKING, Any, Dict, Tuple
+
+# Map attribute names to (module, attribute) for lazy loading.
+_LAZY_ATTRS: Dict[str, Tuple[str, str]] = {
+    # Config
+    "MettaGridConfig": ("mettagrid.config.mettagrid_config", "MettaGridConfig"),
+    # Core classes
+    "MettaGridCore": ("mettagrid.core", "MettaGridCore"),
+    "MettaGridAction": ("mettagrid.core", "MettaGridAction"),
+    "MettaGridObservation": ("mettagrid.core", "MettaGridObservation"),
+    # Main environment (backward compatible)
+    "MettaGridEnv": ("mettagrid.envs.mettagrid_env", "MettaGridEnv"),
+    # Environment adapters
+    "MettaGridGymEnv": ("mettagrid.envs.gym_env", "MettaGridGymEnv"),
+    "MettaGridPettingZooEnv": (
+        "mettagrid.envs.pettingzoo_env",
+        "MettaGridPettingZooEnv",
+    ),
+    # Data types (from C++)
+    "dtype_actions": ("mettagrid.mettagrid_c", "dtype_actions"),
+    "dtype_observations": ("mettagrid.mettagrid_c", "dtype_observations"),
+    "dtype_rewards": ("mettagrid.mettagrid_c", "dtype_rewards"),
+    "dtype_terminals": ("mettagrid.mettagrid_c", "dtype_terminals"),
+    "dtype_truncations": ("mettagrid.mettagrid_c", "dtype_truncations"),
+    "dtype_masks": ("mettagrid.mettagrid_c", "dtype_masks"),
+    "dtype_success": ("mettagrid.mettagrid_c", "dtype_success"),
+    # Type definitions
+    "validate_observation_space": (
+        "mettagrid.types",
+        "validate_observation_space",
+    ),
+    "validate_action_space": ("mettagrid.types", "validate_action_space"),
+    "get_observation_shape": ("mettagrid.types", "get_observation_shape"),
+    "get_action_nvec": ("mettagrid.types", "get_action_nvec"),
+    # Supporting classes
+    "GameMap": ("mettagrid.map_builder.map_builder", "GameMap"),
+    "ReplayWriter": ("mettagrid.util.replay_writer", "ReplayWriter"),
+    "StatsWriter": ("mettagrid.util.stats_writer", "StatsWriter"),
+}
+
+if TYPE_CHECKING:
+    from mettagrid.config.mettagrid_config import MettaGridConfig
+    from mettagrid.core import MettaGridAction, MettaGridCore, MettaGridObservation
+    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
+    from mettagrid.mettagrid_c import (
+        dtype_actions,
+        dtype_masks,
+        dtype_observations,
+        dtype_rewards,
+        dtype_success,
+        dtype_terminals,
+        dtype_truncations,
+    )
+    from mettagrid.types import (
+        get_action_nvec,
+        get_observation_shape,
+        validate_action_space,
+        validate_observation_space,
+    )
+    from mettagrid.util.replay_writer import ReplayWriter
+    from mettagrid.util.stats_writer import StatsWriter
+
+    try:
+        from mettagrid import mettascope
+    except (ImportError, OSError):
+        mettascope = None
+
+__all__ = [
+    # Config
+    "MettaGridConfig",
+    # Core classes
+    "MettaGridCore",
+    # Main environment (backward compatible)
+    "MettaGridEnv",
+    # Environment adapters
+    "MettaGridGymEnv",
+    "MettaGridPettingZooEnv",
+    # Data types (from C++)
+    "dtype_actions",
+    "dtype_observations",
+    "dtype_rewards",
+    "dtype_terminals",
+    "dtype_truncations",
+    "dtype_masks",
+    "dtype_success",
+    # Type definitions
+    "MettaGridObservation",
+    "MettaGridAction",
+    "validate_observation_space",
+    "validate_action_space",
+    "get_observation_shape",
+    "get_action_nvec",
+    # Supporting classes
+    "GameMap",
+    "ReplayWriter",
+    "StatsWriter",
+    # Optional visualization module
+    "mettascope",
+]
+
+
+def __getattr__(name: str) -> Any:  # pragma: no cover - thin import wrapper
+    """Dynamically import public attributes on first access."""
+
+    if name == "mettascope":
+        try:
+            module = importlib.import_module("mettagrid.mettascope")
+        except (ImportError, OSError):
+            module = None
+        globals()[name] = module
+        return module
+
+    try:
+        module_name, attribute_name = _LAZY_ATTRS[name]
+    except KeyError as exc:  # pragma: no cover - mirrors built-in behaviour
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}") from exc
+
+    module = importlib.import_module(module_name)
+    attr = getattr(module, attribute_name)
+    globals()[name] = attr
+    return attr
+
+
+def __dir__() -> list[str]:
+    """Expose lazy attributes in dir() results."""
+
+    return sorted({*__all__, *globals().keys()})

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,205 @@
+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,
+            ),
+        )
+    ],
+)
+
+assembler_lab = AssemblerConfig(
+    type_id=17,
+    recipes=[
+        (
+            ["Any"],
+            RecipeConfig(
+                input_resources={"battery_red": 1, "ore_red": 2},
+                output_resources={"laser": 1},
+                cooldown=10,
+            ),
+        )
+    ],
+)
+
+assembler_factory = AssemblerConfig(
+    type_id=18,
+    recipes=[
+        (
+            ["Any"],
+            RecipeConfig(
+                input_resources={"battery_red": 1, "ore_red": 2},
+                output_resources={"laser": 1},
+                cooldown=10,
+            ),
+        )
+    ],
+)
+
+assembler_temple = AssemblerConfig(
+    type_id=19,
+    recipes=[
+        (
+            ["Any"],
+            RecipeConfig(
+                input_resources={"battery_red": 1, "ore_red": 2},
+                output_resources={"laser": 1},
+                cooldown=10,
+            ),
+        )
+    ],
+)
+
+
+# Chest building definitions. Maybe not needed beyond the raw config?
+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 = []  # Default to no deposit positions
+    if withdrawal_positions is None:
+        withdrawal_positions = []  # Default to no withdrawal positions
+
+    return ChestConfig(
+        type_id=type_id,
+        resource_type=resource_type,
+        deposit_positions=deposit_positions,
+        withdrawal_positions=withdrawal_positions,
+    )
+
+
+# Example chest configurations
+chest_heart = make_chest("heart", 20, 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,
+)