mettagrid 0.0.1__py3-none-any.whl

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/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/actions/attack.pxd

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

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/actions/gift.pxd

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

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/actions/move.pxd

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

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/actions/rotate.pxd

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

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/actions/shield.pxd

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

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/actions/use.pxd

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

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/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/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/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)