mettagrid 0.0.1__tar.gz

mettagrid-0.0.1/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 David Bloomin
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

mettagrid-0.0.1/PKG-INFO

@@ -0,0 +1,23 @@
+Metadata-Version: 2.1
+Name: mettagrid
+Version: 0.0.1
+Summary: A fast grid-based open-ended MARL environment
+Home-page: https://daveey.github.io
+License: MIT
+Keywords: puffergrid,gridworld,minigrid,rl,reinforcement-learning,environment,gym
+Author: David Bloomin
+Author-email: daveey@gmail.com
+Requires-Python: >=3.10,<4.0
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Requires-Dist: cython (>=3.0.11,<4.0.0)
+Requires-Dist: numpy (>=1.21.0,<2.0.0)
+Project-URL: Repository, https://github.com/Metta-AI/mettagrid
+Description-Content-Type: text/markdown
+
+# mettagrid
+A fast grid-based open-ended MARL environment
+

mettagrid-0.0.1/README.md

@@ -0,0 +1,2 @@
+# mettagrid
+A fast grid-based open-ended MARL environment

mettagrid-0.0.1/mettagrid/actions/actions.pxd

@@ -0,0 +1,31 @@
+
+from libc.stdio cimport printf
+from libcpp.string cimport string
+from libcpp.map cimport map
+
+from puffergrid.grid_object cimport TypeId, GridObjectId
+from puffergrid.action cimport ActionHandler, ActionArg
+from mettagrid.objects cimport Agent
+
+cdef struct StatNames:
+    string action
+    string action_energy
+    map[TypeId, string] target
+    map[TypeId, string] target_energy
+
+cdef class MettaActionHandler(ActionHandler):
+    cdef StatNames _stats
+    cdef string action_name
+    cdef int action_cost
+
+    cdef char handle_action(
+        self,
+        unsigned int actor_id,
+        GridObjectId actor_object_id,
+        ActionArg arg)
+
+    cdef char _handle_action(
+        self,
+        unsigned int actor_id,
+        Agent * actor,
+        ActionArg arg)

mettagrid-0.0.1/mettagrid/actions/actions.pyx

@@ -0,0 +1,59 @@
+
+from libc.stdio cimport printf
+
+from omegaconf import OmegaConf
+
+from puffergrid.grid_object cimport GridObjectId
+from puffergrid.action cimport ActionHandler, ActionArg
+from mettagrid.objects cimport Agent, ObjectTypeNames
+
+cdef extern from "<string>" namespace "std":
+    string to_string(int val)
+
+cdef class MettaActionHandler(ActionHandler):
+    def __init__(self, cfg: OmegaConf, action_name, action_cost=0):
+        self.action_name = action_name
+
+        self._stats.action = "action." + action_name
+        self._stats.action_energy = "action." + action_name + ".energy"
+
+        for t, n in enumerate(ObjectTypeNames):
+            self._stats.target[t] = self._stats.action + "." + n
+            self._stats.target_energy[t] = self._stats.action_energy + "." + n
+
+        self.action_cost = cfg.cost
+
+    cdef char handle_action(
+        self,
+        unsigned int actor_id,
+        GridObjectId actor_object_id,
+        ActionArg arg):
+
+        cdef Agent *actor = <Agent*>self.env._grid.object(actor_object_id)
+
+        if actor.frozen:
+            return False
+
+        if actor.energy < self.action_cost:
+            return False
+
+        actor.energy -= self.action_cost
+        self.env._stats.agent_add(actor_id, self._stats.action_energy.c_str(), self.action_cost)
+
+        cdef char result = self._handle_action(actor_id, actor, arg)
+
+        if result:
+            self.env._stats.agent_incr(actor_id, self._stats.action.c_str())
+
+        return result
+
+    cdef char _handle_action(
+        self,
+        unsigned int actor_id,
+        Agent * actor,
+        ActionArg arg):
+        return False
+
+
+
+

mettagrid-0.0.1/mettagrid/actions/attack.pxd

@@ -0,0 +1,6 @@
+from mettagrid.actions.actions cimport MettaActionHandler
+
+cdef class Attack(MettaActionHandler):
+    cdef public int damage
+
+    pass

mettagrid-0.0.1/mettagrid/actions/attack.pyx

@@ -0,0 +1,61 @@
+
+from libc.stdio cimport printf
+from omegaconf import OmegaConf
+
+from puffergrid.grid_object cimport GridLocation, GridObjectId, Orientation, GridObject
+from puffergrid.action cimport ActionHandler, ActionArg
+from mettagrid.objects cimport MettaObject, ObjectType, Usable, Altar, Agent, Events, GridLayer
+from mettagrid.objects cimport Generator, Converter, InventoryItem, ObjectTypeNames, InventoryItemNames
+from mettagrid.actions.actions cimport MettaActionHandler
+
+cdef class Attack(MettaActionHandler):
+    def __init__(self, cfg: OmegaConf):
+        MettaActionHandler.__init__(self, cfg, "attack")
+
+    cdef char _handle_action(
+        self,
+        unsigned int actor_id,
+        Agent * actor,
+        ActionArg arg):
+
+        if arg > 9 or arg < 1:
+            return False
+
+        cdef short distance = 0
+        cdef short offset = 0
+        distance = 1 + (arg - 1) // 3
+        offset = (arg - 1) % 3 - 1
+
+        cdef GridLocation target_loc = self.env._grid.relative_location(
+            actor.location,
+            <Orientation>actor.orientation,
+            distance, offset)
+
+        target_loc.layer = GridLayer.Agent_Layer
+        cdef Agent * agent_target = <Agent *>self.env._grid.object_at(target_loc)
+        if agent_target:
+            self.env._stats.agent_incr(actor_id, self._stats.target[agent_target._type_id].c_str())
+            if agent_target.shield and agent_target.energy >= self.damage:
+                    agent_target.energy -= self.damage
+                    self.env._stats.agent_add(actor_id, "shield_damage", self.damage)
+            else:
+                self.env._stats.agent_add(actor_id, "shield_damage", agent_target.energy)
+                agent_target.energy = 0
+                agent_target.shield = False
+                agent_target.frozen = True
+                self.env._stats.agent_incr(actor_id, "attack.frozen")
+            return True
+
+        target_loc.layer = GridLayer.Object_Layer
+        cdef MettaObject * object_target = <MettaObject *>self.env._grid.object_at(target_loc)
+        if object_target:
+            self.env._stats.agent_incr(actor_id, self._stats.target[object_target._type_id].c_str())
+            object_target.hp -= 1
+            self.env._stats.agent_incr(actor_id, "damage." + ObjectTypeNames[object_target._type_id])
+            if object_target.hp <= 0:
+                self.env._grid.remove_object(object_target)
+                self.env._stats.agent_incr(actor_id, "destroyed." + ObjectTypeNames[object_target._type_id])
+
+            return True
+
+        return False

mettagrid-0.0.1/mettagrid/actions/gift.pxd

@@ -0,0 +1,4 @@
+from mettagrid.actions.actions cimport MettaActionHandler
+
+cdef class Gift(MettaActionHandler):
+    pass

mettagrid-0.0.1/mettagrid/actions/gift.pyx

@@ -0,0 +1,24 @@
+
+from libc.stdio cimport printf
+from omegaconf import OmegaConf
+
+from puffergrid.grid_object cimport GridLocation, GridObjectId, Orientation, GridObject
+from puffergrid.action cimport ActionHandler, ActionArg
+from mettagrid.objects cimport MettaObject, ObjectType, Usable, Altar, Agent, Events, GridLayer
+from mettagrid.objects cimport Generator, Converter, InventoryItem, ObjectTypeNames, InventoryItemNames
+from mettagrid.actions.actions cimport MettaActionHandler
+
+
+cdef class Gift(MettaActionHandler):
+    def __init__(self, cfg: OmegaConf):
+        MettaActionHandler.__init__(self, cfg, "gift")
+
+    cdef char _handle_action(
+        self,
+        unsigned int actor_id,
+        Agent * actor,
+        ActionArg arg):
+        return False
+
+
+

mettagrid-0.0.1/mettagrid/actions/move.pxd

@@ -0,0 +1,4 @@
+from mettagrid.actions.actions cimport MettaActionHandler
+
+cdef class Move(MettaActionHandler):
+    pass

mettagrid-0.0.1/mettagrid/actions/move.pyx

@@ -0,0 +1,31 @@
+
+from libc.stdio cimport printf
+
+from omegaconf import OmegaConf
+
+from puffergrid.grid_object cimport GridLocation, GridObjectId, GridObject, Orientation
+from puffergrid.action cimport ActionHandler, ActionArg
+from mettagrid.objects cimport MettaObject, ObjectType, Usable, Altar, Agent, Events, GridLayer
+from mettagrid.objects cimport Generator, Converter, InventoryItem, ObjectTypeNames, InventoryItemNames
+from mettagrid.actions.actions cimport MettaActionHandler
+
+cdef class Move(MettaActionHandler):
+    def __init__(self, cfg: OmegaConf):
+        MettaActionHandler.__init__(self, cfg, "move")
+
+    cdef char _handle_action(
+        self,
+        unsigned int actor_id,
+        Agent * actor,
+        ActionArg arg):
+
+        cdef unsigned short direction = arg
+        if direction >= 2:
+            return False
+
+        cdef Orientation orientation = <Orientation>((actor.orientation + 2*(direction)) % 4)
+        cdef GridLocation old_loc = actor.location
+        cdef GridLocation new_loc = self.env._grid.relative_location(old_loc, orientation)
+        if not self.env._grid.is_empty(new_loc.r, new_loc.c):
+            return False
+        return self.env._grid.move_object(actor.id, new_loc)

mettagrid-0.0.1/mettagrid/actions/rotate.pxd

@@ -0,0 +1,4 @@
+from mettagrid.actions.actions cimport MettaActionHandler
+
+cdef class Rotate(MettaActionHandler):
+    pass

mettagrid-0.0.1/mettagrid/actions/rotate.pyx

@@ -0,0 +1,28 @@
+
+from libc.stdio cimport printf
+
+from omegaconf import OmegaConf
+
+from puffergrid.grid_object cimport GridLocation, GridObjectId, Orientation, GridObject
+from puffergrid.action cimport ActionHandler, ActionArg
+from mettagrid.objects cimport MettaObject, ObjectType, Usable, Altar, Agent, Events, GridLayer
+from mettagrid.objects cimport Generator, Converter, InventoryItem, ObjectTypeNames, InventoryItemNames
+from mettagrid.actions.actions cimport MettaActionHandler
+
+
+cdef class Rotate(MettaActionHandler):
+    def __init__(self, cfg: OmegaConf):
+        MettaActionHandler.__init__(self, cfg, "rotate")
+
+    cdef char _handle_action(
+        self,
+        unsigned int actor_id,
+        Agent * actor,
+        ActionArg arg):
+
+        cdef unsigned short orientation = arg
+        if orientation >= 4:
+            return False
+
+        actor.orientation = orientation
+        return True

mettagrid-0.0.1/mettagrid/actions/shield.pxd

@@ -0,0 +1,4 @@
+from mettagrid.actions.actions cimport MettaActionHandler
+
+cdef class Shield(MettaActionHandler):
+    pass

mettagrid-0.0.1/mettagrid/actions/shield.pyx

@@ -0,0 +1,26 @@
+
+from libc.stdio cimport printf
+
+from omegaconf import OmegaConf
+
+from puffergrid.grid_object cimport GridLocation, GridObjectId, Orientation, GridObject
+from puffergrid.action cimport ActionHandler, ActionArg
+from mettagrid.objects cimport MettaObject, ObjectType, Usable, Altar, Agent, Events, GridLayer
+from mettagrid.objects cimport Generator, Converter, InventoryItem, ObjectTypeNames, InventoryItemNames
+from mettagrid.actions.actions cimport MettaActionHandler
+
+cdef class Shield(MettaActionHandler):
+    def __init__(self, cfg: OmegaConf):
+        MettaActionHandler.__init__(self, cfg, "shield")
+
+    cdef char _handle_action(
+        self,
+        unsigned int actor_id,
+        Agent * actor,
+        ActionArg arg):
+
+        if actor.shield:
+            actor.shield = True
+        else:
+            actor.shield = False
+

mettagrid-0.0.1/mettagrid/actions/use.pxd

@@ -0,0 +1,4 @@
+from mettagrid.actions.actions cimport MettaActionHandler
+
+cdef class Use(MettaActionHandler):
+    pass

mettagrid-0.0.1/mettagrid/actions/use.pyx

@@ -0,0 +1,66 @@
+
+from libc.stdio cimport printf
+
+from omegaconf import OmegaConf
+
+from puffergrid.grid_object cimport GridLocation, GridObjectId, Orientation, GridObject
+from puffergrid.action cimport ActionHandler, ActionArg
+from mettagrid.objects cimport MettaObject, ObjectType, Usable, Altar, Agent, Events, GridLayer
+from mettagrid.objects cimport Generator, Converter, InventoryItem, ObjectTypeNames, InventoryItemNames
+from mettagrid.actions.actions cimport MettaActionHandler
+
+cdef class Use(MettaActionHandler):
+    def __init__(self, cfg: OmegaConf):
+        MettaActionHandler.__init__(self, cfg, "use")
+
+    cdef char _handle_action(
+        self,
+        unsigned int actor_id,
+        Agent * actor,
+        ActionArg arg):
+
+        cdef GridLocation target_loc = self.env._grid.relative_location(
+            actor.location,
+            <Orientation>actor.orientation
+        )
+        target_loc.layer = GridLayer.Object_Layer
+        cdef MettaObject *target = <MettaObject*>self.env._grid.object_at(target_loc)
+        if target == NULL:
+            return False
+
+        if not target.usable(actor):
+            return False
+
+        cdef Usable *usable = <Usable*> target
+        actor.energy -= usable.use_cost
+
+        usable.ready = 0
+        self.env._event_manager.schedule_event(Events.Reset, usable.cooldown, usable.id, 0)
+
+        self.env._stats.agent_incr(actor_id, self._stats.target[target._type_id].c_str())
+        self.env._stats.agent_add(actor_id, self._stats.target_energy[target._type_id].c_str(), usable.use_cost + self.action_cost)
+
+        if target._type_id == ObjectType.AltarT:
+            self.env._rewards[actor_id] += 1
+
+        cdef Generator *generator
+        if target._type_id == ObjectType.GeneratorT:
+            generator = <Generator*>target
+            generator.r1 -= 1
+            actor.update_inventory(InventoryItem.r1, 1)
+            self.env._stats.agent_incr(actor_id, "r1.gained")
+            self.env._stats.game_incr("r1.harvested")
+
+        cdef Converter *converter
+        if target._type_id == ObjectType.ConverterT:
+            converter = <Converter*>target
+            actor.update_inventory(converter.input_resource, -1)
+            self.env._stats.agent_incr(actor_id, InventoryItemNames[converter.input_resource] + ".used")
+
+            actor.update_inventory(converter.output_resource, 1)
+            self.env._stats.agent_incr(actor_id, InventoryItemNames[converter.input_resource] + ".gained")
+
+            actor.energy += converter.output_energy
+            self.env._stats.agent_add(actor_id, "energy.gained", converter.output_energy)
+
+        return True

mettagrid-0.0.1/mettagrid/config/game_builder.py

@@ -0,0 +1,101 @@
+from typing import Dict, List
+import numpy as np
+import yaml
+from omegaconf import OmegaConf
+
+class NoAliasDumper(yaml.Dumper):
+    def ignore_aliases(self, data):
+        return True
+
+class MettaGridGameBuilder():
+    def __init__(
+            self,
+            obs_width: int,
+            obs_height: int,
+            tile_size: int,
+            max_steps: int,
+            num_agents: int,
+            no_energy_steps: int,
+            objects,
+            actions,
+            map):
+
+        self.obs_width = obs_width
+        self.obs_height = obs_height
+        self.tile_size = tile_size
+        self.num_agents = num_agents
+        self.max_steps = max_steps
+
+        self._symbols = {
+            "agent": "A",
+            "altar": "a",
+            "converter": "c",
+            "generator": "g",
+            "wall": "w",
+            "empty": " ",
+        }
+
+        self.no_energy_steps = no_energy_steps
+        objects = OmegaConf.create(objects)
+        self.object_configs = objects
+        actions = OmegaConf.create(actions)
+        self.action_configs = actions
+        self.map_config = OmegaConf.create(map)
+
+
+    def level(self):
+        layout = self.map_config.layout
+
+        if "rooms" in layout:
+            return self.build_map(layout.rooms)
+        else:
+            return self.build_map(
+                [["room"] * layout.rooms_x] * layout.rooms_y)
+
+    def build_map(self, rooms):
+        num_agents = 0
+        layers = []
+        for layer in rooms:
+            rooms = []
+            for room_name in layer:
+                room_config = self.map_config[room_name]
+                rooms.append(self.build_room(room_config, num_agents + 1))
+                num_agents += room_config.objects.agent
+            layers.append(np.concatenate(rooms, axis=1))
+        level = np.concatenate(layers, axis=0)
+        assert num_agents == self.num_agents, f"Number of agents in map ({num_agents}) does not match num_agents ({self.num_agents})"
+
+        footer = np.full((1, level.shape[1]), "W", dtype="U6")
+        footer[0, 0] = "q"
+
+        level = np.concatenate([level, footer], axis=0)
+        return level
+
+    def build_room(self, room_config, starting_agent=1):
+        symbols = []
+        content_width = room_config.width - 2*room_config.border
+        content_height = room_config.height - 2*room_config.border
+        area = content_width * content_height
+
+        for obj_name, count in room_config.objects.items():
+            symbol = self._symbols[obj_name]
+            if obj_name == "agent":
+                symbols.extend([f"{symbol}{i+starting_agent}" for i in range(count)])
+            else:
+                symbols.extend([symbol] * count)
+
+        assert(len(symbols) <= area), f"Too many objects in room: {len(symbols)} > {area}"
+        symbols.extend(["."] * (area - len(symbols)))
+        symbols = np.array(symbols).astype("U8")
+        np.random.shuffle(symbols)
+        content = symbols.reshape(content_height, content_width)
+        room = np.full((room_config.height, room_config.width), "W", dtype="U6")
+        room[room_config.border:room_config.border+content_height,
+             room_config.border:room_config.border+content_width] = content
+
+        return room
+
+    def termination_conditions(self):
+        return {
+            "Win": [ {"lt": ["game:max_steps", "game:step"]} ],
+        }

mettagrid-0.0.1/mettagrid/config/sample_config.py

@@ -0,0 +1,27 @@
+
+
+import numpy as np
+from omegaconf import DictConfig, ListConfig
+
+
+def sample_config(value):
+    if isinstance(value, int):
+        return value
+    if isinstance(value, float):
+        return value
+    if isinstance(value, DictConfig):
+        return {
+            key: sample_config(value)
+            for key, value in value.items()
+        }
+    if isinstance(value, ListConfig):
+        if len(value) == 0:
+            return value
+        if isinstance(value[0], int):
+            assert len(value) == 2, f"Found a list with length != 2 {value}"
+            return np.random.randint(value[0], value[1])
+        if isinstance(value[0], float):
+            assert len(value) == 2, f"Found a list with length != 2 {value}"
+            return np.random.uniform(value[0], value[1])
+        return value
+    return value

mettagrid-0.0.1/mettagrid/mettagrid.pyx

@@ -0,0 +1,78 @@
+
+from libc.stdio cimport printf
+
+import numpy as np
+import gymnasium as gym
+from omegaconf import OmegaConf
+
+from puffergrid.grid_env cimport GridEnv
+
+from mettagrid.objects cimport ObjectLayers, Agent, ResetHandler, Wall, Generator, Converter, Altar
+from mettagrid.objects cimport MettaObservationEncoder
+from mettagrid.actions.move import Move
+from mettagrid.actions.rotate import Rotate
+from mettagrid.actions.use import Use
+from mettagrid.actions.attack import Attack
+from mettagrid.actions.shield import Shield
+from mettagrid.actions.gift import Gift
+
+cdef class MettaGrid(GridEnv):
+    cdef:
+        object _cfg
+
+    def __init__(self, cfg: OmegaConf, map: np.ndarray):
+        self._cfg = cfg
+
+        GridEnv.__init__(
+            self,
+            cfg.num_agents,
+            map.shape[1],
+            map.shape[0],
+            cfg.max_steps,
+            dict(ObjectLayers).values(),
+            cfg.obs_width, cfg.obs_height,
+            MettaObservationEncoder(),
+            [
+                Move(cfg.actions.move),
+                Rotate(cfg.actions.rotate),
+                Use(cfg.actions.use),
+                Attack(cfg.actions.attack),
+                Shield(cfg.actions.shield),
+                Gift(cfg.actions.gift),
+            ],
+            [
+                ResetHandler()
+            ]
+        )
+
+
+        cdef Agent *agent
+        for r in range(map.shape[0]):
+            for c in range(map.shape[1]):
+                if map[r,c] == "W":
+                    self._grid.add_object(new Wall(r, c, cfg.objects.wall))
+                    self._stats.game_incr("objects.wall")
+                elif map[r,c] == "g":
+                    self._grid.add_object(new Generator(r, c, cfg.objects.generator))
+                    self._stats.game_incr("objects.generator")
+                elif map[r,c] == "c":
+                    self._grid.add_object(new Converter(r, c, cfg.objects.converter))
+                    self._stats.game_incr("objects.converter")
+                elif map[r,c] == "a":
+                    self._grid.add_object(new Altar(r, c, cfg.objects.altar))
+                    self._stats.game_incr("objects.altar")
+                elif map[r,c][0] == "A":
+                    agent = new Agent(r, c, cfg.objects.agent)
+                    self._grid.add_object(agent)
+                    self.add_agent(agent)
+                    self._stats.game_incr("objects.agent")
+
+
+    def render(self):
+        grid = self.render_ascii(["A", "#", "g", "c", "a"])
+        for r in grid:
+            print("".join(r))
+
+    @property
+    def action_space(self):
+        return gym.spaces.MultiDiscrete((self.num_actions(), 10), dtype=np.uint32)

mettagrid-0.0.1/mettagrid/mettagrid_env.py

@@ -0,0 +1,181 @@
+from typing import Any, Dict
+
+import pufferlib
+import numpy as np
+from omegaconf import OmegaConf
+
+from mettagrid.config.game_builder import MettaGridGameBuilder
+from mettagrid.renderer.raylib_client import MettaRaylibClient
+from mettagrid.config.sample_config import sample_config
+from mettagrid.mettagrid_c import MettaGrid
+from pufferlib.environments.ocean.render import GridRender
+
+class GridClient:
+    def __init__(self, width, height):
+        self._width = width
+        self._height = height
+
+class MettaGridEnv(pufferlib.PufferEnv):
+    def __init__(self, render_mode: str, **cfg):
+        super().__init__()
+
+        self._render_mode = render_mode
+        self._cfg = OmegaConf.create(cfg)
+        self.make_env()
+
+        if render_mode == "human":
+            self._renderer = MettaRaylibClient(
+                self._env.map_width(), self._env.map_height(),
+            )
+        elif render_mode == "raylib":
+            self._renderer = GridRender(
+                self._env.map_width(), self._env.map_height(),
+                fps=10
+            )
+
+
+    def make_env(self):
+        game_cfg = OmegaConf.create(sample_config(self._cfg.game))
+        self._game_builder = MettaGridGameBuilder(**game_cfg)
+        level = self._game_builder.level()
+        self._c_env = MettaGrid(game_cfg, level)
+        self._grid_env = self._c_env
+        self._num_agents = self._c_env.num_agents()
+
+        # self._grid_env = PufferGridEnv(self._c_env)
+        env = self._grid_env
+
+        self._env = env
+        #self._env = LastActionTracker(self._grid_env)
+        #self._env = Kinship(**sample_config(self._cfg.kinship), env=self._env)
+        #self._env = RewardTracker(self._env)
+        #self._env = FeatureMasker(self._env, self._cfg.hidden_features)
+        self.done = False
+
+    def reset(self, **kwargs):
+        self.make_env()
+        if hasattr(self, "buf"):
+            self._c_env.set_buffers(
+                self.buf.observations,
+                self.buf.terminals,
+                self.buf.truncations,
+                self.buf.rewards)
+
+        # obs, infos = self._env.reset(**kwargs)
+        # self._compute_max_energy()
+        # return obs, infos
+        obs, infos = self._c_env.reset()
+        return obs, infos
+
+    def step(self, actions):
+        obs, rewards, terminated, truncated, infos = self._c_env.step(actions.astype(np.int32))
+
+        rewards_sum = rewards.sum()
+        if rewards_sum != 0:
+            reward_mean = rewards_sum / self._num_agents
+            rewards -= reward_mean
+
+        if terminated.all() or truncated.all():
+            self.done = True
+
+            stats = self._c_env.get_episode_stats()
+            episode_rewards = self._c_env.get_episode_rewards()
+            episode_rewards_sum = episode_rewards.sum()
+            episode_rewards_mean = episode_rewards_sum / self._num_agents
+
+            infos = {
+                "episode/reward.sum": episode_rewards_sum,
+                "episode/reward.mean": episode_rewards_mean,
+                "episode/reward.min": episode_rewards.min(),
+                "episode/reward.max": episode_rewards.max(),
+                "episode_length": self._c_env.current_timestep(),
+            }
+
+            agent_stats = {}
+            for a_stats in stats["agent_stats"]:
+                for k, v in a_stats.items():
+                    if k not in agent_stats:
+                        agent_stats[k] = 0
+                    agent_stats[k] += v
+
+            for k, v in agent_stats.items():
+                infos[f"agent_stats/{k}"] = float(v) / self._num_agents
+
+        return obs, list(rewards), terminated.all(), truncated.all(), infos
+
+    def process_episode_stats(self, episode_stats: Dict[str, Any]):
+        for agent_stats in episode_stats["agent_stats"]:
+            extra_stats = {}
+            for stat_name in agent_stats.keys():
+                if stat_name.startswith("action_"):
+                    extra_stats[stat_name + "_pct"] = agent_stats[stat_name] / self._grid_env.current_timestep
+
+
+            #     for object in self._game_builder.object_configs.keys():
+            #         if stat_name.startswith(f"stats_{object}_") and object != "agent":
+            #             symbol = self._game_builder._objects[object].symbol
+            #             num_obj = self._griddly_yaml["Environment"]["Levels"][0].count(symbol)
+            #             if num_obj == 0:
+            #                 num_obj = 1
+            #             extra_stats[stat_name + "_pct"] = agent_stats[stat_name] / num_obj
+
+            agent_stats.update(extra_stats)
+            agent_stats.update(episode_stats["game_stats"])
+            # agent_stats["level_max_energy"] = self._max_level_energy
+            # agent_stats["level_max_energy_per_agent"] = self._max_level_energy_per_agent
+            # agent_stats["level_max_reward_per_agent"] = self._max_level_reward_per_agent
+
+    def _compute_max_energy(self):
+        pass
+        # num_generators = self._griddly_yaml["Environment"]["Levels"][0].count("g")
+        # num_converters = self._griddly_yaml["Environment"]["Levels"][0].count("c")
+        # max_resources = num_generators * min(
+        #     self._game_builder.object_configs.generator.initial_resources,
+        #     self._max_steps / self._game_builder.object_configs.generator.cooldown)
+
+        # max_conversions = num_converters * (
+        #     self._max_steps / self._game_builder.object_configs.converter.cooldown
+        # )
+        # max_conv_energy = min(max_resources, max_conversions) * \
+        #     np.mean(list(self._game_builder.object_configs.converter.energy_output.values()))
+
+        # initial_energy = self._game_builder.object_configs.agent.initial_energy * self._game_builder.num_agents
+
+        # self._max_level_energy = max_conv_energy + initial_energy
+        # self._max_level_energy_per_agent = self._max_level_energy / self._game_builder.num_agents
+
+        # self._max_level_reward_per_agent = self._max_level_energy_per_agent
+
+
+    @property
+    def _max_steps(self):
+        return self._game_builder.max_steps
+
+    @property
+    def observation_space(self):
+        return self._env.observation_space
+
+    @property
+    def action_space(self):
+        return self._env.action_space
+
+    @property
+    def player_count(self):
+        return self._num_agents
+
+    def render(self, *args, **kwargs):
+        return self._renderer.render(
+            self._c_env.grid_objects(),
+        )
+
+    @property
+    def grid_features(self):
+        return self._env.grid_features()
+
+    @property
+    def global_features(self):
+        return []
+
+    @property
+    def render_mode(self):
+        return self._render_mode

mettagrid-0.0.1/mettagrid/objects.pxd

@@ -0,0 +1,197 @@
+# distutils: language=c++
+# cython: warn.undeclared=False
+
+cimport cython
+
+from libcpp.vector cimport vector
+from libcpp.map cimport map
+from libcpp.string cimport string
+from puffergrid.grid_env import StatsTracker
+from libc.stdio cimport printf
+from puffergrid.observation_encoder cimport ObservationEncoder, ObsType
+from puffergrid.grid_object cimport GridObject, TypeId, GridCoord, GridLocation, GridObjectId
+from puffergrid.event cimport EventHandler, EventArg
+
+cdef enum GridLayer:
+    Agent_Layer = 0
+    Object_Layer = 1
+
+ctypedef map[string, int] ObjectConfig
+
+cdef cppclass MettaObject(GridObject):
+    unsigned int hp
+
+    inline void init_mo(ObjectConfig cfg):
+        this.hp = cfg[b"hp"]
+
+    inline char usable(const Agent *actor):
+        return False
+
+    inline char attackable():
+        return False
+
+cdef cppclass Usable(MettaObject):
+    unsigned int use_cost
+    unsigned int cooldown
+    unsigned char ready
+
+    inline void init_usable(ObjectConfig cfg):
+        this.use_cost = cfg[b"use_cost"]
+        this.cooldown = cfg[b"cooldown"]
+        this.ready = 1
+
+    inline char usable(const Agent *actor):
+        return this.ready and this.use_cost <= actor.energy
+
+cdef enum ObjectType:
+    AgentT = 0
+    WallT = 1
+    GeneratorT = 2
+    ConverterT = 3
+    AltarT = 4
+    Count = 5
+
+cdef vector[string] ObjectTypeNames # defined in objects.pyx
+
+cdef enum InventoryItem:
+    r1 = 0,
+    r2 = 1,
+    r3 = 2,
+    InventoryCount = 3
+
+cdef vector[string] InventoryItemNames # defined in objects.pyx
+
+
+cdef cppclass Agent(MettaObject):
+    char frozen
+    unsigned int energy
+    unsigned int orientation
+    char shield
+    vector[unsigned short] inventory
+
+    inline Agent(GridCoord r, GridCoord c, ObjectConfig cfg):
+        GridObject.init(ObjectType.AgentT, GridLocation(r, c, GridLayer.Agent_Layer))
+        MettaObject.init_mo(cfg)
+        this.frozen = False
+        this.energy = cfg[b"initial_energy"]
+        this.orientation = 0
+        this.inventory.resize(InventoryItem.InventoryCount)
+
+    inline void update_inventory(InventoryItem item, short amount):
+        this.inventory[<InventoryItem>item] += amount
+
+    inline void obs(ObsType[:] obs):
+        obs[0] = 1
+        obs[1] = this.hp
+        obs[2] = this.frozen
+        obs[3] = this.energy
+        obs[4] = this.orientation
+        obs[5] = this.shield
+
+        cdef unsigned short idx = 6
+        cdef unsigned short i
+        for i in range(InventoryItem.InventoryCount):
+            obs[idx + i] = this.inventory[i]
+
+    @staticmethod
+    inline vector[string] feature_names():
+        return [
+            "agent", "agent:hp", "agent:frozen", "agent:energy", "agent:orientation",
+            "agent:shield"
+        ] + [
+            "agent:inv:" + n for n in InventoryItemNames]
+
+cdef cppclass Wall(MettaObject):
+    inline Wall(GridCoord r, GridCoord c, ObjectConfig cfg):
+        GridObject.init(ObjectType.WallT, GridLocation(r, c, GridLayer.Object_Layer))
+        MettaObject.init_mo(cfg)
+
+    inline void obs(ObsType[:] obs):
+        obs[0] = 1
+        obs[1] = hp
+
+    @staticmethod
+    inline vector[string] feature_names():
+        return ["wall", "wall:hp"]
+
+cdef cppclass Generator(Usable):
+    unsigned int r1
+
+    inline Generator(GridCoord r, GridCoord c, ObjectConfig cfg):
+        GridObject.init(ObjectType.GeneratorT, GridLocation(r, c, GridLayer.Object_Layer))
+        MettaObject.init_mo(cfg)
+        Usable.init_usable(cfg)
+        this.r1 = cfg[b"initial_resources"]
+
+    inline char usable(const Agent *actor):
+        return Usable.usable(actor) and this.r1 > 0
+
+    inline void obs(ObsType[:] obs):
+        obs[0] = 1
+        obs[1] = this.hp
+        obs[2] = this.r1
+        obs[3] = this.ready
+
+
+    @staticmethod
+    inline vector[string] feature_names():
+        return ["generator", "generator:hp", "generator:r1", "generator:ready"]
+
+cdef cppclass Converter(Usable):
+    InventoryItem input_resource
+    InventoryItem output_resource
+    short output_energy
+
+    inline Converter(GridCoord r, GridCoord c, ObjectConfig cfg):
+        GridObject.init(ObjectType.ConverterT, GridLocation(r, c, GridLayer.Object_Layer))
+        MettaObject.init_mo(cfg)
+        Usable.init_usable(cfg)
+        this.input_resource = InventoryItem.r1
+        this.output_resource = InventoryItem.r2
+        this.output_energy = cfg[b"energy_output.r1"]
+
+    inline char usable(const Agent *actor):
+        return Usable.usable(actor) and actor.inventory[this.input_resource] > 0
+
+    inline obs(ObsType[:] obs):
+        obs[0] = 1
+        obs[1] = hp
+        obs[2] = input_resource
+        obs[3] = output_resource
+        obs[4] = output_energy
+        obs[5] = ready
+
+    @staticmethod
+    inline vector[string] feature_names():
+        return ["converter", "converter:hp", "converter:input_resource", "converter:output_resource", "converter:output_energy", "converter:ready"]
+
+cdef cppclass Altar(Usable):
+    inline Altar(GridCoord r, GridCoord c, ObjectConfig cfg):
+        GridObject.init(ObjectType.AltarT, GridLocation(r, c, GridLayer.Object_Layer))
+        MettaObject.init_mo(cfg)
+        Usable.init_usable(cfg)
+
+    inline void obs(ObsType[:] obs):
+        obs[0] = 1
+        obs[1] = hp
+        obs[2] = ready
+
+    @staticmethod
+    inline vector[string] feature_names():
+        return ["altar", "altar:hp", "altar:ready"]
+
+cdef map[TypeId, GridLayer] ObjectLayers
+
+cdef class ResetHandler(EventHandler):
+    cdef inline void handle_event(self, GridObjectId obj_id, EventArg arg):
+        cdef Usable *usable = <Usable*>self.env._grid.object(obj_id)
+        usable.ready = True
+        self.env._stats.game_incr("resets." + ObjectTypeNames[usable._type_id])
+
+cdef enum Events:
+    Reset = 0
+
+cdef class MettaObservationEncoder(ObservationEncoder):
+    cdef vector[short] _offsets
+    cdef vector[string] _feature_names
+

mettagrid-0.0.1/mettagrid/objects.pyx

@@ -0,0 +1,67 @@
+# distutils: language=c++
+
+from libc.stdio cimport printf
+from libcpp.string cimport string
+from libcpp.vector cimport vector
+from puffergrid.grid_object cimport GridObject, GridObjectId
+
+cdef vector[string] ObjectTypeNames = [
+    "agent",
+    "wall",
+    "generator",
+    "converter",
+    "altar"
+]
+
+cdef vector[string] InventoryItemNames = [
+    "r1",
+    "r2",
+    "r3"
+]
+
+ObjectLayers = {
+    ObjectType.AgentT: GridLayer.Agent_Layer,
+    ObjectType.WallT: GridLayer.Object_Layer,
+    ObjectType.GeneratorT: GridLayer.Object_Layer,
+    ObjectType.ConverterT: GridLayer.Object_Layer,
+    ObjectType.AltarT: GridLayer.Object_Layer,
+}
+
+cdef class MettaObservationEncoder(ObservationEncoder):
+    def __init__(self) -> None:
+        self._offsets.resize(ObjectType.Count)
+
+        features = []
+        self._offsets[ObjectType.AgentT] = 0
+        features.extend(Agent.feature_names())
+
+        self._offsets[ObjectType.WallT] = len(features)
+        features.extend(Wall.feature_names())
+
+        self._offsets[ObjectType.GeneratorT] = len(features)
+        features.extend(Generator.feature_names())
+
+        self._offsets[ObjectType.ConverterT] = len(features)
+        features.extend(Converter.feature_names())
+
+        self._offsets[ObjectType.AltarT] = len(features)
+        features.extend(Altar.feature_names())
+
+        self._feature_names = features
+
+    cdef encode(self, GridObject *obj, ObsType[:] obs):
+        if obj._type_id == ObjectType.AgentT:
+            (<Agent*>obj).obs(obs[self._offsets[ObjectType.AgentT]:])
+        elif obj._type_id == ObjectType.WallT:
+            (<Wall*>obj).obs(obs[self._offsets[ObjectType.WallT]:])
+        elif obj._type_id == ObjectType.GeneratorT:
+            (<Generator*>obj).obs(obs[self._offsets[ObjectType.GeneratorT]:])
+        elif obj._type_id == ObjectType.ConverterT:
+            (<Converter*>obj).obs(obs[self._offsets[ObjectType.ConverterT]:])
+        elif obj._type_id == ObjectType.AltarT:
+            (<Altar*>obj).obs(obs[self._offsets[ObjectType.AltarT]:])
+        else:
+            printf("Encoding object of unknown type: %d\n", obj._type_id)
+
+    cdef vector[string] feature_names(self):
+        return self._feature_names

mettagrid-0.0.1/mettagrid/renderer/raylib_client.py

@@ -0,0 +1,180 @@
+from pdb import set_trace as T
+import numpy as np
+import os
+
+import pettingzoo
+import gymnasium
+
+import pufferlib
+from pufferlib.environments.ocean import render
+
+class MettaRaylibClient:
+    def __init__(self, width, height, tile_size=32):
+        self.width = width
+        self.height = height
+        self.tile_size = tile_size
+
+
+        sprite_sheet_path = os.path.join(
+            *self.__module__.split('.')[:-1], './puffer_chars.png')
+        self.asset_map = {
+            1: (0, 0, 128, 128),
+            3: (128, 0, 128, 128),
+            4: (256, 0, 128, 128),
+            # 5: (384, 0, 128, 128),
+            5: (512, 0, 128, 128), #star
+        }
+
+        from raylib import rl, colors
+        rl.InitWindow(width*tile_size, height*tile_size,
+            "PufferLib Ray Grid".encode())
+        rl.SetTargetFPS(10)
+        self.puffer = rl.LoadTexture(sprite_sheet_path.encode())
+        self.rl = rl
+        self.colors = colors
+
+        import pyray as ray
+        camera = ray.Camera2D()
+        camera.target = ray.Vector2(0.0, 0.0)
+        camera.rotation = 0.0
+        camera.zoom = 1.0
+        self.camera = camera
+
+        from cffi import FFI
+        self.ffi = FFI()
+
+    def _cdata_to_numpy(self):
+        image = self.rl.LoadImageFromScreen()
+        width, height, channels = image.width, image.height, 4
+        cdata = self.ffi.buffer(image.data, width*height*channels)
+        return np.frombuffer(cdata, dtype=np.uint8
+            ).reshape((height, width, channels))[:, :, :3]
+
+    def render(self, grid):
+        rl = self.rl
+        colors = self.colors
+        ay, ax = None, None
+
+        ts = self.tile_size
+
+        pos = rl.GetMousePosition()
+        raw_mouse_x = pos.x
+        raw_mouse_y = pos.y
+        mouse_x = int(raw_mouse_x // ts)
+        mouse_y = int(raw_mouse_y // ts)
+        ay = int(np.clip((pos.y - ts*self.height//2) / 50, -3, 3)) + 3
+        ax = int(np.clip((pos.x - ts*self.width//2) / 50, -3, 3)) + 3
+
+        if rl.IsKeyDown(rl.KEY_ESCAPE):
+            exit(0)
+
+        action_id = 0
+        action_arg = 0
+
+        if rl.IsKeyDown(rl.KEY_E):
+            action_id = 0
+            action_arg = 0
+        elif rl.IsKeyDown(rl.KEY_Q):
+            action_id = 0
+            action_arg = 1
+
+        elif rl.IsKeyDown(rl.KEY_W):
+            action_id = 1
+            action_arg = 0
+        elif rl.IsKeyDown(rl.KEY_S):
+            action_id = 1
+            action_arg = 1
+        elif rl.IsKeyDown(rl.KEY_A):
+            action_id = 1
+            action_arg = 2
+        elif rl.IsKeyDown(rl.KEY_R):
+            action_id = 1
+            action_arg = 3
+
+        # if rl.IsKeyDown(rl.KEY_LEFT_SHIFT):
+        #     target_heros = 2
+
+        action = (action_id, action_arg)
+
+        rl.BeginDrawing()
+        rl.BeginMode2D(self.camera)
+        rl.ClearBackground([6, 24, 24, 255])
+        for y in range(self.height):
+            for x in range(self.width):
+                tile = grid[y, x]
+                tx = x*ts
+                ty = y*ts
+                if tile == 0:
+                    continue
+                elif tile == 2:
+                    # Wall
+                    rl.DrawRectangle(x*ts, y*ts, ts, ts, [0, 0, 0, 255])
+                    continue
+                else:
+                    # Player
+                    source_rect = self.asset_map[tile]
+                    dest_rect = (tx, ty, ts, ts)
+                    rl.DrawTexturePro(self.puffer, source_rect, dest_rect,
+                        (0, 0), 0, colors.WHITE)
+
+        # Draw circle at mouse x, y
+        rl.DrawCircle(ts*mouse_x + ts//2, ts*mouse_y + ts//8, ts//8, [255, 0, 0, 255])
+
+        rl.EndMode2D()
+
+        # Draw HUD
+        # player = entities[0]
+        # hud_y = self.height*ts - 2*ts
+        # draw_bars(rl, player, 2*ts, hud_y, 10*ts, 24, draw_text=True)
+
+        # off_color = [255, 255, 255, 255]
+        # on_color = [0, 255, 0, 255]
+
+        # q_color = on_color if skill_q else off_color
+        # w_color = on_color if skill_w else off_color
+        # e_color = on_color if skill_e else off_color
+
+        # q_cd = player.q_timer
+        # w_cd = player.w_timer
+        # e_cd = player.e_timer
+
+        # rl.DrawText(f'Q: {q_cd}'.encode(), 13*ts, hud_y - 20, 40, q_color)
+        # rl.DrawText(f'W: {w_cd}'.encode(), 17*ts, hud_y - 20, 40, w_color)
+        # rl.DrawText(f'E: {e_cd}'.encode(), 21*ts, hud_y - 20, 40, e_color)
+        # rl.DrawText(f'Stun: {player.stun_timer}'.encode(), 25*ts, hud_y - 20, 20, e_color)
+        # rl.DrawText(f'Move: {player.move_timer}'.encode(), 25*ts, hud_y, 20, e_color)
+
+        rl.EndDrawing()
+        return self._cdata_to_numpy(), action
+
+def draw_bars(rl, entity, x, y, width, height=4, draw_text=False):
+    health_bar = entity.health / entity.max_health
+    mana_bar = entity.mana / entity.max_mana
+    if entity.max_health == 0:
+        health_bar = 2
+    if entity.max_mana == 0:
+        mana_bar = 2
+    rl.DrawRectangle(x, y, width, height, [255, 0, 0, 255])
+    rl.DrawRectangle(x, y, int(width*health_bar), height, [0, 255, 0, 255])
+
+    if entity.entity_type == 0:
+        rl.DrawRectangle(x, y - height - 2, width, height, [255, 0, 0, 255])
+        rl.DrawRectangle(x, y - height - 2, int(width*mana_bar), height, [0, 255, 255, 255])
+
+    if draw_text:
+        health = int(entity.health)
+        mana = int(entity.mana)
+        max_health = int(entity.max_health)
+        max_mana = int(entity.max_mana)
+        rl.DrawText(f'Health: {health}/{max_health}'.encode(),
+            x+8, y+2, 20, [255, 255, 255, 255])
+        rl.DrawText(f'Mana: {mana}/{max_mana}'.encode(),
+            x+8, y+2 - height - 2, 20, [255, 255, 255, 255])
+
+        #rl.DrawRectangle(x, y - 2*height - 4, int(width*mana_bar), height, [255, 255, 0, 255])
+        rl.DrawText(f'Experience: {entity.xp}'.encode(),
+            x+8, y - 2*height - 4, 20, [255, 255, 255, 255])
+
+    elif entity.entity_type == 0:
+        rl.DrawText(f'Level: {entity.level}'.encode(),
+            x+4, y -2*height - 12, 12, [255, 255, 255, 255])

mettagrid-0.0.1/pyproject.toml

@@ -0,0 +1,22 @@
+[build-system]
+requires = ["poetry-core>=1.0.0"]
+build-backend = "poetry.core.masonry.api"
+
+[tool.poetry]
+name = "mettagrid"
+version = "0.0.1"
+description = "A fast grid-based open-ended MARL environment"
+authors = ["David Bloomin <daveey@gmail.com>"]
+license = "MIT"
+readme = "README.md"
+homepage = "https://daveey.github.io"
+repository = "https://github.com/Metta-AI/mettagrid"
+keywords = ["puffergrid", "gridworld", "minigrid", "rl", "reinforcement-learning", "environment", "gym"]
+
+[tool.poetry.dependencies]
+python = ">=3.10,<4.0"
+numpy = "^1.21.0"
+cython = "^3.0.11"
+
+[tool.poetry.dev-dependencies]
+pytest = "^6.2.4"