mettagrid 0.0.1__py3-none-any.whl → 0.0.3__py3-none-any.whl

mettagrid/actions/actions.pxd

@@ -18,13 +18,13 @@ cdef class MettaActionHandler(ActionHandler):
     cdef string action_name
     cdef int action_cost
 
-    cdef char handle_action(
+    cdef bint handle_action(
         self,
         unsigned int actor_id,
         GridObjectId actor_object_id,
         ActionArg arg)
 
-    cdef char _handle_action(
+    cdef bint _handle_action(
         self,
         unsigned int actor_id,
         Agent * actor,

mettagrid/actions/actions.pyx

@@ -11,7 +11,7 @@ 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):
+    def __init__(self, cfg: OmegaConf, action_name):
         self.action_name = action_name
 
         self._stats.action = "action." + action_name
@@ -23,7 +23,7 @@ cdef class MettaActionHandler(ActionHandler):
 
         self.action_cost = cfg.cost
 
-    cdef char handle_action(
+    cdef bint handle_action(
         self,
         unsigned int actor_id,
         GridObjectId actor_object_id,
@@ -31,7 +31,17 @@ cdef class MettaActionHandler(ActionHandler):
 
         cdef Agent *actor = <Agent*>self.env._grid.object(actor_object_id)
 
-        if actor.frozen:
+        if actor.shield:
+            actor.energy -= actor.shield_upkeep
+            self.env._stats.agent_add(actor_id, "shield_upkeep", actor.shield_upkeep)
+            self.env._stats.agent_incr(actor_id, "status.shield.ticks")
+            if actor.energy <= 0:
+                actor.shield = False
+                actor.energy = 0
+
+        if actor.frozen > 0:
+            self.env._stats.agent_incr(actor_id, "status.frozen.ticks")
+            actor.frozen -= 1
             return False
 
         if actor.energy < self.action_cost:
@@ -40,14 +50,18 @@ cdef class MettaActionHandler(ActionHandler):
         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)
+        cdef bint result = self._handle_action(actor_id, actor, arg)
 
         if result:
             self.env._stats.agent_incr(actor_id, self._stats.action.c_str())
 
+        if actor.energy < 0:
+            actor.energy = 0
+            self.env._stats.agent_incr(actor_id, "error.negative_energy")
+
         return result
 
-    cdef char _handle_action(
+    cdef bint _handle_action(
         self,
         unsigned int actor_id,
         Agent * actor,

mettagrid/actions/attack.pyx

@@ -12,7 +12,10 @@ cdef class Attack(MettaActionHandler):
     def __init__(self, cfg: OmegaConf):
         MettaActionHandler.__init__(self, cfg, "attack")
 
-    cdef char _handle_action(
+    cdef unsigned char max_arg(self):
+        return 9
+
+    cdef bint _handle_action(
         self,
         unsigned int actor_id,
         Agent * actor,
@@ -24,7 +27,7 @@ cdef class Attack(MettaActionHandler):
         cdef short distance = 0
         cdef short offset = 0
         distance = 1 + (arg - 1) // 3
-        offset = (arg - 1) % 3 - 1
+        offset = -((arg - 1) % 3 - 1)
 
         cdef GridLocation target_loc = self.env._grid.relative_location(
             actor.location,
@@ -42,8 +45,14 @@ cdef class Attack(MettaActionHandler):
                 self.env._stats.agent_add(actor_id, "shield_damage", agent_target.energy)
                 agent_target.energy = 0
                 agent_target.shield = False
-                agent_target.frozen = True
+                agent_target.frozen = agent_target.freeze_duration
                 self.env._stats.agent_incr(actor_id, "attack.frozen")
+                for item in range(InventoryItem.InventoryCount):
+                    actor.update_inventory(item, agent_target.inventory[item])
+                    self.env._stats.agent_add(actor_id, InventoryItemNames[item] + ".stolen", agent_target.inventory[item])
+                    self.env._stats.agent_add(actor_id, InventoryItemNames[item] + ".gained", agent_target.inventory[item])
+                    agent_target.inventory[item] = 0
+
             return True
 
         target_loc.layer = GridLayer.Object_Layer

mettagrid/actions/gift.pyx

@@ -13,7 +13,10 @@ cdef class Gift(MettaActionHandler):
     def __init__(self, cfg: OmegaConf):
         MettaActionHandler.__init__(self, cfg, "gift")
 
-    cdef char _handle_action(
+    cdef unsigned char max_arg(self):
+        return 1
+
+    cdef bint _handle_action(
         self,
         unsigned int actor_id,
         Agent * actor,

mettagrid/actions/move.pyx

@@ -1,4 +1,3 @@
-
 from libc.stdio cimport printf
 
 from omegaconf import OmegaConf
@@ -13,19 +12,30 @@ cdef class Move(MettaActionHandler):
     def __init__(self, cfg: OmegaConf):
         MettaActionHandler.__init__(self, cfg, "move")
 
-    cdef char _handle_action(
+    cdef unsigned char max_arg(self):
+        return 1
+
+    cdef bint _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 Orientation orientation = <Orientation>(actor.orientation)
+        if direction == 1:
+            if orientation == Orientation.Up:
+                orientation = Orientation.Down
+            elif orientation == Orientation.Down:
+                orientation = Orientation.Up
+            elif orientation == Orientation.Left:
+                orientation = Orientation.Right
+            elif orientation == Orientation.Right:
+                orientation = Orientation.Left
+
         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 0
         return self.env._grid.move_object(actor.id, new_loc)

mettagrid/actions/noop.pxd

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

mettagrid/actions/noop.pyx

@@ -0,0 +1,25 @@
+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 Noop(MettaActionHandler):
+    def __init__(self, cfg: OmegaConf):
+        MettaActionHandler.__init__(self, cfg, "noop")
+
+    cdef unsigned char max_arg(self):
+        return 0
+
+    cdef bint _handle_action(
+        self,
+        unsigned int actor_id,
+        Agent * actor,
+        ActionArg arg):
+
+        return 1
+

mettagrid/actions/rotate.pyx

@@ -14,15 +14,17 @@ cdef class Rotate(MettaActionHandler):
     def __init__(self, cfg: OmegaConf):
         MettaActionHandler.__init__(self, cfg, "rotate")
 
-    cdef char _handle_action(
+    cdef unsigned char max_arg(self):
+        return 3
+
+
+    cdef bint _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.pyx

@@ -13,14 +13,21 @@ cdef class Shield(MettaActionHandler):
     def __init__(self, cfg: OmegaConf):
         MettaActionHandler.__init__(self, cfg, "shield")
 
-    cdef char _handle_action(
+    cdef unsigned char max_arg(self):
+        return 0
+
+
+    cdef bint _handle_action(
         self,
         unsigned int actor_id,
         Agent * actor,
         ActionArg arg):
 
         if actor.shield:
-            actor.shield = True
-        else:
             actor.shield = False
+        elif actor.energy >= actor.shield_upkeep:
+            actor.shield = True
+
+        return True
+
 

mettagrid/actions/use.pyx

@@ -13,7 +13,10 @@ cdef class Use(MettaActionHandler):
     def __init__(self, cfg: OmegaConf):
         MettaActionHandler.__init__(self, cfg, "use")
 
-    cdef char _handle_action(
+    cdef unsigned char max_arg(self):
+        return 0
+
+    cdef bint _handle_action(
         self,
         unsigned int actor_id,
         Agent * actor,
@@ -52,15 +55,17 @@ cdef class Use(MettaActionHandler):
             self.env._stats.game_incr("r1.harvested")
 
         cdef Converter *converter
+        cdef unsigned int energy_gain = 0
         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")
+            self.env._stats.agent_incr(actor_id, InventoryItemNames[converter.output_resource] + ".gained")
 
-            actor.energy += converter.output_energy
-            self.env._stats.agent_add(actor_id, "energy.gained", converter.output_energy)
+            energy_gain = min(actor.max_energy - actor.energy, converter.output_energy)
+            actor.energy += energy_gain
+            self.env._stats.agent_add(actor_id, "energy.gained", energy_gain)
 
         return True

mettagrid/mettagrid.pyx

@@ -1,11 +1,16 @@
 
 from libc.stdio cimport printf
 
+
 import numpy as np
+cimport numpy as cnp
 import gymnasium as gym
 from omegaconf import OmegaConf
+from types import SimpleNamespace
 
 from puffergrid.grid_env cimport GridEnv
+from puffergrid.grid_object cimport GridObject
+from puffergrid.observation_encoder cimport ObsType
 
 from mettagrid.objects cimport ObjectLayers, Agent, ResetHandler, Wall, Generator, Converter, Altar
 from mettagrid.objects cimport MettaObservationEncoder
@@ -15,6 +20,9 @@ from mettagrid.actions.use import Use
 from mettagrid.actions.attack import Attack
 from mettagrid.actions.shield import Shield
 from mettagrid.actions.gift import Gift
+from mettagrid.actions.noop import Noop
+
+obs_np_type = np.uint8
 
 cdef class MettaGrid(GridEnv):
     cdef:
@@ -33,6 +41,7 @@ cdef class MettaGrid(GridEnv):
             cfg.obs_width, cfg.obs_height,
             MettaObservationEncoder(),
             [
+                Noop(SimpleNamespace(cost=0)),
                 Move(cfg.actions.move),
                 Rotate(cfg.actions.rotate),
                 Use(cfg.actions.use),
@@ -45,7 +54,6 @@ cdef class MettaGrid(GridEnv):
             ]
         )
 
-
         cdef Agent *agent
         for r in range(map.shape[0]):
             for c in range(map.shape[1]):
@@ -76,3 +84,29 @@ cdef class MettaGrid(GridEnv):
     @property
     def action_space(self):
         return gym.spaces.MultiDiscrete((self.num_actions(), 10), dtype=np.uint32)
+
+    cpdef grid_objects(self):
+        cdef GridObject *obj
+        cdef ObsType[:] obj_data = np.zeros(len(self._obs_encoder.feature_names()), dtype=obs_np_type)
+        cdef unsigned int obj_id, i
+        cdef MettaObservationEncoder obs_encoder = <MettaObservationEncoder>self._obs_encoder
+        objects = {}
+        for obj_id in range(1, self._grid.objects.size()):
+            obj = self._grid.object(obj_id)
+            if obj == NULL:
+                continue
+            objects[obj_id] = {
+                "id": obj_id,
+                "type": obj._type_id,
+                "r": obj.location.r,
+                "c": obj.location.c,
+                "layer": obj.location.layer
+            }
+            obs_encoder._encode(obj, obj_data, 0)
+            for i, name in enumerate(obs_encoder._type_feature_names[obj._type_id]):
+                objects[obj_id][name] = obj_data[i]
+
+        for agent_idx in range(self._agents.size()):
+            objects[self._agents[agent_idx].id]["agent_id"] = agent_idx
+
+        return objects

mettagrid/mettagrid_env.py

@@ -1,19 +1,13 @@
 from typing import Any, Dict
 
-import pufferlib
 import numpy as np
+import pufferlib
 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
+from mettagrid.mettagrid_c import MettaGrid # pylint: disable=E0611
 
-class GridClient:
-    def __init__(self, width, height):
-        self._width = width
-        self._height = height
 
 class MettaGridEnv(pufferlib.PufferEnv):
     def __init__(self, render_mode: str, **cfg):
@@ -23,20 +17,17 @@ class MettaGridEnv(pufferlib.PufferEnv):
         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
-            )
+        self._renderer = None
+
+        self.done = False
+        self.buf = None
 
 
     def make_env(self):
-        game_cfg = OmegaConf.create(sample_config(self._cfg.game))
-        self._game_builder = MettaGridGameBuilder(**game_cfg)
+        scfg = sample_config(self._cfg.game)
+        assert isinstance(scfg, Dict)
+        game_cfg = OmegaConf.create(scfg)
+        self._game_builder = MettaGridGameBuilder(**scfg) # type: ignore
         level = self._game_builder.level()
         self._c_env = MettaGrid(game_cfg, level)
         self._grid_env = self._c_env
@@ -52,9 +43,9 @@ class MettaGridEnv(pufferlib.PufferEnv):
         #self._env = FeatureMasker(self._env, self._cfg.hidden_features)
         self.done = False
 
-    def reset(self, **kwargs):
+    def reset(self, seed=None):
         self.make_env()
-        if hasattr(self, "buf"):
+        if hasattr(self, "buf") and self.buf is not None:
             self._c_env.set_buffers(
                 self.buf.observations,
                 self.buf.terminals,
@@ -70,60 +61,37 @@ class MettaGridEnv(pufferlib.PufferEnv):
     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 self._cfg.normalize_rewards:
+            rewards -= rewards.mean()
 
+        infos = {}
         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
-
+            self.process_episode_stats(infos)
         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 process_episode_stats(self, infos: Dict[str, Any]):
+        episode_rewards = self._c_env.get_episode_rewards()
+        episode_rewards_sum = episode_rewards.sum()
+        episode_rewards_mean = episode_rewards_sum / self._num_agents
+        infos.update({
+            "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(),
+        })
+        stats = self._c_env.get_episode_stats()
+
+        infos["episode_rewards"] = episode_rewards
+        infos["agent_raw"] = stats["agent"]
+        infos["game"] = stats["game"]
+        infos["agent"] = {}
+        for agent_stats in stats["agent"]:
+            for n, v in agent_stats.items():
+                infos["agent"][n] = infos["agent"].get(n, 0) + v
+        for n, v in infos["agent"].items():
+            infos["agent"][n] = v / self._num_agents
 
     def _compute_max_energy(self):
         pass
@@ -163,9 +131,13 @@ class MettaGridEnv(pufferlib.PufferEnv):
     def player_count(self):
         return self._num_agents
 
-    def render(self, *args, **kwargs):
+    def render(self):
+        if self._renderer is None:
+            return None
+
         return self._renderer.render(
-            self._c_env.grid_objects(),
+            self._c_env.current_timestep(),
+            self._c_env.grid_objects()
         )
 
     @property

mettagrid/objects.pxd

@@ -1,5 +1,6 @@
 # distutils: language=c++
 # cython: warn.undeclared=False
+# cython: c_api_binop_methods=True
 
 cimport cython
 
@@ -24,10 +25,7 @@ cdef cppclass MettaObject(GridObject):
     inline void init_mo(ObjectConfig cfg):
         this.hp = cfg[b"hp"]
 
-    inline char usable(const Agent *actor):
-        return False
-
-    inline char attackable():
+    inline bint usable(const Agent *actor):
         return False
 
 cdef cppclass Usable(MettaObject):
@@ -40,7 +38,7 @@ cdef cppclass Usable(MettaObject):
         this.cooldown = cfg[b"cooldown"]
         this.ready = 1
 
-    inline char usable(const Agent *actor):
+    inline bint usable(const Agent *actor):
         return this.ready and this.use_cost <= actor.energy
 
 cdef enum ObjectType:
@@ -63,22 +61,32 @@ cdef vector[string] InventoryItemNames # defined in objects.pyx
 
 
 cdef cppclass Agent(MettaObject):
-    char frozen
-    unsigned int energy
+    unsigned int frozen
+    unsigned int freeze_duration
+    unsigned char energy
     unsigned int orientation
-    char shield
-    vector[unsigned short] inventory
+    unsigned char shield
+    unsigned char shield_upkeep
+    vector[unsigned char] inventory
+    unsigned char max_items
+    unsigned char max_energy
 
     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.frozen = 0
+        this.freeze_duration = cfg[b"freeze_duration"]
         this.energy = cfg[b"initial_energy"]
+        this.max_energy = cfg[b"max_energy"]
+        this.shield_upkeep = cfg[b"upkeep.shield"]
         this.orientation = 0
         this.inventory.resize(InventoryItem.InventoryCount)
+        this.max_items = cfg[b"max_inventory"]
 
     inline void update_inventory(InventoryItem item, short amount):
         this.inventory[<InventoryItem>item] += amount
+        if this.inventory[<InventoryItem>item] > this.max_items:
+            this.inventory[<InventoryItem>item] = this.max_items
 
     inline void obs(ObsType[:] obs):
         obs[0] = 1
@@ -123,7 +131,7 @@ cdef cppclass Generator(Usable):
         Usable.init_usable(cfg)
         this.r1 = cfg[b"initial_resources"]
 
-    inline char usable(const Agent *actor):
+    inline bint usable(const Agent *actor):
         return Usable.usable(actor) and this.r1 > 0
 
     inline void obs(ObsType[:] obs):
@@ -150,7 +158,7 @@ cdef cppclass Converter(Usable):
         this.output_resource = InventoryItem.r2
         this.output_energy = cfg[b"energy_output.r1"]
 
-    inline char usable(const Agent *actor):
+    inline bint usable(const Agent *actor):
         return Usable.usable(actor) and actor.inventory[this.input_resource] > 0
 
     inline obs(ObsType[:] obs):
@@ -185,6 +193,9 @@ 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)
+        if usable is NULL:
+            return
+
         usable.ready = True
         self.env._stats.game_incr("resets." + ObjectTypeNames[usable._type_id])
 
@@ -192,6 +203,8 @@ cdef enum Events:
     Reset = 0
 
 cdef class MettaObservationEncoder(ObservationEncoder):
+    cdef _encode(self, GridObject *obj, ObsType[:] obs, unsigned int offset)
     cdef vector[short] _offsets
     cdef vector[string] _feature_names
+    cdef vector[vector[string]] _type_feature_names