mettagrid 0.2__cp311-none-macosx_11_0_arm64.whl

metta/mettagrid/__init__.py

@@ -0,0 +1,61 @@
+"""
+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 environment classes
+from metta.mettagrid.core import MettaGridCore
+
+# Import other commonly used classes
+from metta.mettagrid.gym_env import MettaGridGymEnv
+from metta.mettagrid.map_builder.map_builder import GameMap
+
+# Import data types from C++ module (source of truth)
+from metta.mettagrid.mettagrid_c import (
+    dtype_actions,
+    dtype_masks,
+    dtype_observations,
+    dtype_rewards,
+    dtype_success,
+    dtype_terminals,
+    dtype_truncations,
+)
+from metta.mettagrid.mettagrid_config import MettaGridConfig
+from metta.mettagrid.mettagrid_env import MettaGridEnv
+from metta.mettagrid.pettingzoo_env import MettaGridPettingZooEnv
+from metta.mettagrid.replay_writer import ReplayWriter
+from metta.mettagrid.stats_writer import StatsWriter
+
+__all__ = [
+    # Config
+    "MettaGridConfig",
+    # Core classes
+    "MettaGridCore",
+    # Main environment (backward compatible)
+    "MettaGridEnv",
+    # Environment adapters
+    "MettaGridGymEnv",
+    "MettaGridPettingZooEnv",
+    # Data types
+    "dtype_actions",
+    "dtype_observations",
+    "dtype_rewards",
+    "dtype_terminals",
+    "dtype_truncations",
+    "dtype_masks",
+    "dtype_success",
+    # Supporting classes
+    "GameMap",
+    "ReplayWriter",
+    "StatsWriter",
+]

metta/mettagrid/action_handler.hpp

@@ -0,0 +1,137 @@
+// action_handler.hpp
+#ifndef ACTION_HANDLER_HPP_
+#define ACTION_HANDLER_HPP_
+
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+
+#include <map>
+#include <string>
+#include <vector>
+
+#include "grid.hpp"
+#include "grid_object.hpp"
+#include "objects/agent.hpp"
+#include "objects/constants.hpp"
+#include "types.hpp"
+
+struct ActionConfig {
+  std::map<InventoryItem, InventoryQuantity> required_resources;
+  std::map<InventoryItem, InventoryQuantity> consumed_resources;
+
+  ActionConfig(const std::map<InventoryItem, InventoryQuantity>& required_resources = {},
+               const std::map<InventoryItem, InventoryQuantity>& consumed_resources = {})
+      : required_resources(required_resources), consumed_resources(consumed_resources) {}
+
+  virtual ~ActionConfig() {}
+};
+
+class ActionHandler {
+public:
+  unsigned char priority;
+  Grid* _grid{};
+
+  ActionHandler(const ActionConfig& cfg, const std::string& action_name)
+      : priority(0),
+        _action_name(action_name),
+        _required_resources(cfg.required_resources),
+        _consumed_resources(cfg.consumed_resources) {
+    for (const auto& [item, amount] : _required_resources) {
+      if (amount < _consumed_resources[item]) {
+        throw std::runtime_error("Required resources must be greater than or equal to consumed resources");
+      }
+    }
+  }
+
+  virtual ~ActionHandler() {}
+
+  void init(Grid* grid) {
+    this->_grid = grid;
+  }
+
+  bool handle_action(GridObjectId actor_object_id, ActionArg arg) {
+    Agent* actor = static_cast<Agent*>(_grid->object(actor_object_id));
+
+    // Handle frozen status
+    if (actor->frozen != 0) {
+      actor->stats.incr("status.frozen.ticks");
+      actor->stats.incr("status.frozen.ticks." + actor->group_name);
+      if (actor->frozen > 0) {
+        actor->frozen -= 1;
+      }
+      return false;
+    }
+
+    bool has_needed_resources = true;
+    for (const auto& [item, amount] : _required_resources) {
+      if (actor->inventory[item] < amount) {
+        has_needed_resources = false;
+        break;
+      }
+    }
+
+    // Execute the action
+    bool success = has_needed_resources && _handle_action(actor, arg);
+
+    // The intention here is to provide a metric that reports when an agent has stayed in one location for a long
+    // period, perhaps spinning in circles. We think this could be a good indicator that a policy has collapsed.
+    if (actor->location == actor->prev_location) {
+      actor->steps_without_motion += 1;
+      if (actor->steps_without_motion > actor->stats.get("status.max_steps_without_motion")) {
+        actor->stats.set("status.max_steps_without_motion", actor->steps_without_motion);
+      }
+    } else {
+      actor->steps_without_motion = 0;
+    }
+
+    // Update tracking for this agent
+    actor->prev_action_name = _action_name;
+    actor->prev_location = actor->location;
+
+    // Track success/failure
+    if (success) {
+      actor->stats.incr("action." + _action_name + ".success");
+      for (const auto& [item, amount] : _consumed_resources) {
+        [[maybe_unused]] InventoryDelta delta = actor->update_inventory(item, -amount);
+        // We consume resources after the action succeeds, but in the future we might have an action that uses the
+        // resource. This check will catch that.
+        assert(delta == -amount);
+      }
+    } else {
+      actor->stats.incr("action." + _action_name + ".failed");
+      actor->stats.incr("action.failure_penalty");
+      *actor->reward -= actor->action_failure_penalty;
+    }
+
+    return success;
+  }
+
+  virtual unsigned char max_arg() const {
+    return 0;
+  }
+
+  std::string action_name() const {
+    return _action_name;
+  }
+
+protected:
+  virtual bool _handle_action(Agent* actor, ActionArg arg) = 0;
+
+  std::string _action_name;
+  std::map<InventoryItem, InventoryQuantity> _required_resources;
+  std::map<InventoryItem, InventoryQuantity> _consumed_resources;
+};
+
+namespace py = pybind11;
+
+inline void bind_action_config(py::module& m) {
+  py::class_<ActionConfig, std::shared_ptr<ActionConfig>>(m, "ActionConfig")
+      .def(py::init<const std::map<InventoryItem, InventoryQuantity>&,
+                    const std::map<InventoryItem, InventoryQuantity>&>(),
+           py::arg("required_resources") = std::map<InventoryItem, InventoryQuantity>(),
+           py::arg("consumed_resources") = std::map<InventoryItem, InventoryQuantity>())
+      .def_readwrite("required_resources", &ActionConfig::required_resources)
+      .def_readwrite("consumed_resources", &ActionConfig::consumed_resources);
+}
+
+#endif  // ACTION_HANDLER_HPP_

metta/mettagrid/actions/attack.hpp

@@ -0,0 +1,235 @@
+#ifndef ACTIONS_ATTACK_HPP_
+#define ACTIONS_ATTACK_HPP_
+
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+
+#include <map>
+#include <string>
+#include <vector>
+
+#include "action_handler.hpp"
+#include "grid_object.hpp"
+#include "mettagrid_config.hpp"
+#include "objects/agent.hpp"
+#include "objects/constants.hpp"
+#include "types.hpp"
+
+// Attack takes an argument 0-8, which is the index of the target agent to attack.
+// Target agents are those found in a 3x3 grid in front of the agent, indexed in scan order.
+// If the argument (agent index to attack) > num_agents, the last agent is attacked.
+struct AttackActionConfig : public ActionConfig {
+  std::map<InventoryItem, InventoryQuantity> defense_resources;
+
+  AttackActionConfig(const std::map<InventoryItem, InventoryQuantity>& required_resources,
+                     const std::map<InventoryItem, InventoryQuantity>& consumed_resources,
+                     const std::map<InventoryItem, InventoryQuantity>& defense_resources)
+      : ActionConfig(required_resources, consumed_resources), defense_resources(defense_resources) {}
+};
+
+class Attack : public ActionHandler {
+public:
+  explicit Attack(const AttackActionConfig& cfg,
+                  const GameConfig* game_config,
+                  const std::string& action_name = "attack")
+      : ActionHandler(cfg, action_name), _defense_resources(cfg.defense_resources), _game_config(game_config) {
+    priority = 1;
+  }
+
+  unsigned char max_arg() const override {
+    return 8;
+  }
+
+protected:
+  std::map<InventoryItem, InventoryQuantity> _defense_resources;
+  const GameConfig* _game_config;
+
+  bool _handle_action(Agent* actor, ActionArg arg) override {
+    Agent* last_agent = nullptr;
+    short num_skipped = 0;
+
+    // Attack pattern depends on diagonal support
+    if (!_game_config->allow_diagonals) {
+      // Original 3x3 grid pattern for cardinal-only movement
+      // 7 6 8  (3 cells forward)
+      // 4 3 5  (2 cells forward)
+      // 1 0 2  (1 cell forward)
+      // . A .  (Agent position)
+
+      static constexpr short COL_OFFSETS[3] = {0, -1, 1};
+
+      for (short distance = 1; distance <= 3; distance++) {
+        for (short offset : COL_OFFSETS) {
+          GridLocation target_loc = _grid->relative_location(actor->location, actor->orientation, distance, offset);
+          target_loc.layer = GridLayer::AgentLayer;
+
+          Agent* target_agent = static_cast<Agent*>(_grid->object_at(target_loc));
+          if (target_agent) {
+            last_agent = target_agent;
+            if (num_skipped == arg) {
+              return _handle_target(*actor, *target_agent);
+            }
+            num_skipped++;
+          }
+        }
+      }
+    } else {
+      // Diagonal attack pattern when diagonals are enabled
+      // Agent facing NE example:
+      // . 4 6 8
+      // . 1 3 7
+      // . 0 2 5
+      // A . . .
+
+      // Define the 9 positions in order (0-8)
+      static constexpr struct {
+        short forward;
+        short lateral;
+      } DIAGONAL_POSITIONS[9] = {
+          {1, 0},   // 0: directly forward
+          {2, -1},  // 1: forward-left
+          {1, 1},   // 2: forward-right
+          {2, 0},   // 3: 2 forward
+          {3, -2},  // 4: far forward-left
+          {1, 2},   // 5: right-forward
+          {3, -1},  // 6: far forward-left-center
+          {2, 1},   // 7: forward-right-center
+          {3, 0}    // 8: 3 forward
+      };
+
+      for (const auto& pos : DIAGONAL_POSITIONS) {
+        GridLocation target_loc =
+            _grid->relative_location(actor->location, actor->orientation, pos.forward, pos.lateral);
+        target_loc.layer = GridLayer::AgentLayer;
+
+        Agent* target_agent = static_cast<Agent*>(_grid->object_at(target_loc));
+        if (target_agent) {
+          last_agent = target_agent;
+          if (num_skipped == arg) {
+            return _handle_target(*actor, *target_agent);
+          }
+          num_skipped++;
+        }
+      }
+    }
+
+    // If we got here, it means we skipped over all the targets. Attack the last one.
+    if (last_agent) {
+      return _handle_target(*actor, *last_agent);
+    }
+
+    return false;
+  }
+
+  bool _handle_target(Agent& actor, Agent& target) {
+    bool was_already_frozen = target.frozen > 0;
+
+    // Check if target can defend
+    if (!_defense_resources.empty()) {
+      bool target_can_defend = _check_defense_capability(target);
+
+      if (target_can_defend) {
+        _consume_defense_resources(target);
+        _log_blocked_attack(actor, target);
+        return true;
+      }
+    }
+
+    // Attack succeeds
+    target.frozen = target.freeze_duration;
+
+    if (!was_already_frozen) {
+      _steal_resources(actor, target);
+      _log_successful_attack(actor, target);
+    } else {
+      // Track wasted attacks on already-frozen targets
+      const std::string& actor_group = actor.group_name;
+      actor.stats.incr(_action_prefix(actor_group) + "wasted_on_frozen");
+    }
+    return true;
+  }
+
+private:
+  bool _check_defense_capability(const Agent& target) const {
+    for (const auto& [item, amount] : _defense_resources) {
+      auto it = target.inventory.find(item);
+      if (it == target.inventory.end() || it->second < amount) {
+        return false;
+      }
+    }
+    return true;
+  }
+
+  void _consume_defense_resources(Agent& target) {
+    for (const auto& [item, amount] : _defense_resources) {
+      [[maybe_unused]] InventoryDelta delta = target.update_inventory(item, -amount);
+      assert(delta == -amount);
+    }
+  }
+
+  void _steal_resources(Agent& actor, Agent& target) {
+    // Create snapshot to avoid iterator invalidation
+    std::vector<std::pair<InventoryItem, InventoryQuantity>> snapshot;
+    snapshot.reserve(target.inventory.size());
+    for (const auto& [item, amount] : target.inventory) {
+      snapshot.emplace_back(item, amount);
+    }
+
+    // Transfer resources
+    for (const auto& [item, amount] : snapshot) {
+      InventoryDelta stolen = actor.update_inventory(item, amount);
+      target.update_inventory(item, -stolen);
+
+      if (stolen > 0) {
+        _log_resource_theft(actor, target, item, stolen);
+      }
+    }
+  }
+
+  std::string _action_prefix(const std::string& group) const {
+    return "action." + _action_name + "." + group + ".";
+  }
+
+  void _log_blocked_attack(Agent& actor, const Agent& target) const {
+    const std::string& actor_group = actor.group_name;
+    const std::string& target_group = target.group_name;
+
+    actor.stats.incr(_action_prefix(actor_group) + "blocked_by." + target_group);
+  }
+
+  void _log_successful_attack(Agent& actor, Agent& target) const {
+    const std::string& actor_group = actor.group_name;
+    const std::string& target_group = target.group_name;
+    bool same_team = (actor_group == target_group);
+
+    if (same_team) {
+      actor.stats.incr(_action_prefix(actor_group) + "friendly_fire");
+    } else {
+      actor.stats.incr(_action_prefix(actor_group) + "hit." + target_group);
+      target.stats.incr(_action_prefix(target_group) + "hit_by." + actor_group);
+    }
+  }
+
+  void _log_resource_theft(Agent& actor, Agent& target, InventoryItem item, InventoryDelta amount) const {
+    const std::string& actor_group = actor.group_name;
+    const std::string& target_group = target.group_name;
+    const std::string item_name = actor.stats.resource_name(item);
+
+    actor.stats.add(_action_prefix(actor_group) + "steals." + item_name + ".from." + target_group, amount);
+  }
+};
+
+namespace py = pybind11;
+
+inline void bind_attack_action_config(py::module& m) {
+  py::class_<AttackActionConfig, ActionConfig, std::shared_ptr<AttackActionConfig>>(m, "AttackActionConfig")
+      .def(py::init<const std::map<InventoryItem, InventoryQuantity>&,
+                    const std::map<InventoryItem, InventoryQuantity>&,
+                    const std::map<InventoryItem, InventoryQuantity>&>(),
+           py::arg("required_resources") = std::map<InventoryItem, InventoryQuantity>(),
+           py::arg("consumed_resources") = std::map<InventoryItem, InventoryQuantity>(),
+           py::arg("defense_resources") = std::map<InventoryItem, InventoryQuantity>())
+      .def_readwrite("defense_resources", &AttackActionConfig::defense_resources);
+}
+
+#endif  // ACTIONS_ATTACK_HPP_

metta/mettagrid/actions/change_color.hpp

@@ -0,0 +1,40 @@
+#ifndef ACTIONS_CHANGE_COLOR_HPP_
+#define ACTIONS_CHANGE_COLOR_HPP_
+
+#include <string>
+
+#include "action_handler.hpp"
+#include "objects/agent.hpp"
+#include "types.hpp"
+
+class ChangeColor : public ActionHandler {
+public:
+  explicit ChangeColor(const ActionConfig& cfg) : ActionHandler(cfg, "change_color") {}
+
+  unsigned char max_arg() const override {
+    return 3;  // support fine and coarse adjustment
+  }
+
+protected:
+  bool _handle_action(Agent* actor, ActionArg arg) override {
+    // Note: 'color' is uint8_t which naturally wraps at 256.
+    // This could be interpreted as circular hue behavior (red -> orange -> ... -> violet -> red)
+
+    // Calculate step size once (integer division is intentional)
+    const uint8_t step_size = static_cast<uint8_t>(255 / (max_arg() + 1));
+
+    if (arg == 0) {  // Increment
+      actor->color = static_cast<uint8_t>(actor->color + 1);
+    } else if (arg == 1) {  // Decrement
+      actor->color = static_cast<uint8_t>(actor->color - 1);
+    } else if (arg == 2) {  // Large increment
+      actor->color = static_cast<uint8_t>(actor->color + step_size);
+    } else if (arg == 3) {  // Large decrement
+      actor->color = static_cast<uint8_t>(actor->color - step_size);
+    }
+
+    return true;
+  }
+};
+
+#endif  // ACTIONS_CHANGE_COLOR_HPP_

metta/mettagrid/actions/change_glyph.hpp

@@ -0,0 +1,54 @@
+#ifndef ACTIONS_CHANGE_GLYPH_HPP_
+#define ACTIONS_CHANGE_GLYPH_HPP_
+
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+
+#include <string>
+
+#include "action_handler.hpp"
+#include "objects/agent.hpp"
+#include "types.hpp"
+
+struct ChangeGlyphActionConfig : public ActionConfig {
+  const ObservationType number_of_glyphs;
+
+  ChangeGlyphActionConfig(const std::map<InventoryItem, InventoryQuantity>& required_resources,
+                          const std::map<InventoryItem, InventoryQuantity>& consumed_resources,
+                          const ObservationType number_of_glyphs)
+      : ActionConfig(required_resources, consumed_resources), number_of_glyphs(number_of_glyphs) {}
+};
+
+class ChangeGlyph : public ActionHandler {
+public:
+  explicit ChangeGlyph(const ChangeGlyphActionConfig& cfg)
+      : ActionHandler(cfg, "change_glyph"), _number_of_glyphs(cfg.number_of_glyphs) {}
+
+  unsigned char max_arg() const override {
+    // Return number_of_glyphs - 1 since args are 0-indexed
+    return _number_of_glyphs > 0 ? _number_of_glyphs - 1 : 0;
+  }
+
+protected:
+  const ObservationType _number_of_glyphs;
+
+  bool _handle_action(Agent* actor, ActionArg arg) override {
+    actor->glyph = static_cast<ObservationType>(arg);  // ActionArg is int32 for puffer compatibility
+    return true;
+  }
+};
+
+namespace py = pybind11;
+
+inline void bind_change_glyph_action_config(py::module& m) {
+  py::class_<ChangeGlyphActionConfig, ActionConfig, std::shared_ptr<ChangeGlyphActionConfig>>(m,
+                                                                                              "ChangeGlyphActionConfig")
+      .def(py::init<const std::map<InventoryItem, InventoryQuantity>&,
+                    const std::map<InventoryItem, InventoryQuantity>&,
+                    const int>(),
+           py::arg("required_resources") = std::map<InventoryItem, InventoryQuantity>(),
+           py::arg("consumed_resources") = std::map<InventoryItem, InventoryQuantity>(),
+           py::arg("number_of_glyphs"))
+      .def_readonly("number_of_glyphs", &ChangeGlyphActionConfig::number_of_glyphs);
+}
+#endif  // ACTIONS_CHANGE_GLYPH_HPP_

metta/mettagrid/actions/get_output.hpp

@@ -0,0 +1,88 @@
+#ifndef ACTIONS_GET_OUTPUT_HPP_
+#define ACTIONS_GET_OUTPUT_HPP_
+
+#include <string>
+
+#include "action_handler.hpp"
+#include "grid.hpp"
+#include "grid_object.hpp"
+#include "objects/agent.hpp"
+#include "objects/converter.hpp"
+#include "types.hpp"
+
+class GetOutput : public ActionHandler {
+public:
+  explicit GetOutput(const ActionConfig& cfg) : ActionHandler(cfg, "get_items") {}
+
+  unsigned char max_arg() const override {
+    return 0;
+  }
+
+protected:
+  bool _handle_action(Agent* actor, ActionArg /*arg*/) override {
+    GridLocation target_loc = _grid->relative_location(actor->location, static_cast<Orientation>(actor->orientation));
+    target_loc.layer = GridLayer::ObjectLayer;
+    // get_output only works on Converters, since only Converters have an output.
+    // Once we generalize this to `get`, we should be able to get from any HasInventory object, which
+    // should include agents. That's (e.g.) why we're checking inventory_is_accessible.
+    Converter* converter = dynamic_cast<Converter*>(_grid->object_at(target_loc));
+    Box* box = dynamic_cast<Box*>(_grid->object_at(target_loc));
+    if (!converter && !box) {
+      return false;
+    }
+    // If converter, get output from converter
+    if (converter) {
+      if (!converter->inventory_is_accessible()) {
+        return false;
+      }
+
+      // Actions is only successful if we take at least one item.
+      bool resources_taken = false;
+
+      for (const auto& [item, _] : converter->output_resources) {
+        if (converter->inventory.count(item) == 0) {
+          continue;
+        }
+        InventoryDelta resources_available = converter->inventory[item];
+
+        InventoryDelta taken = actor->update_inventory(item, resources_available);
+
+        if (taken > 0) {
+          actor->stats.add(actor->stats.resource_name(item) + ".get", taken);
+          converter->update_inventory(item, -taken);
+          resources_taken = true;
+        }
+      }
+      return resources_taken;
+    }
+
+    // If box, pick up box if allowed
+    if (box) {
+      if (actor->agent_id == box->creator_agent_id) {
+        // Creator cannot open their own box
+        return false;
+      }
+      // If the creator of the box is an agent, return blue battery to creator and penalize creator
+      if (box->creator_agent_id != 255) {
+        Agent* creator = dynamic_cast<Agent*>(_grid->object(box->creator_agent_object_id));
+        if (!creator) {
+          return false;
+        }
+        // Return required resources to create box to creator inventory
+        for (const auto& [item, amount] : box->returned_resources) {
+          if (amount > 0) {
+            creator->update_inventory(item, amount);
+          }
+        }
+      }
+
+      // Reward the agent for opening the box and teleport back to top-left corner
+      _grid->ghost_move_object(box->id, GridLocation(0, 0, GridLayer::ObjectLayer));
+      actor->stats.add("box.opened", 1.0f);
+      return true;
+    }
+    return false;
+  }
+};
+
+#endif  // ACTIONS_GET_OUTPUT_HPP_