cp

Author: daveey

packages/mettagrid/cpp/bindings/mettagrid_c.cpp

@@ -85,7 +85,7 @@ MettaGrid::MettaGrid(const GameConfig& game_config, const py::list map, unsigned
 
   init_action_handlers();
 
-  _init_grid(game_config, map);
+  _init_grid(_game_config, map);
 
   // Pre-compute goal_obs tokens for each agent
   if (_global_obs_config.goal_obs) {
@@ -136,6 +136,9 @@ void MettaGrid::_init_grid(const GameConfig& game_config, const py::list& map) {
     object_type_names[type_id] = object_cfg->type_name;
   }
 
+  // Collect assemblers to initialize agent tracking after all agents are created
+  std::vector<Assembler*> assemblers;
+
   // Initialize objects from map
   for (GridCoord r = 0; r < height; r++) {
     for (GridCoord c = 0; c < width; c++) {
@@ -184,6 +187,7 @@ void MettaGrid::_init_grid(const GameConfig& game_config, const py::list& map) {
         assembler->set_grid(_grid.get());
         assembler->set_current_timestep_ptr(&current_step);
         assembler->set_obs_encoder(_obs_encoder.get());
+        assemblers.push_back(assembler);
         continue;
       }
 
@@ -201,6 +205,7 @@ void MettaGrid::_init_grid(const GameConfig& game_config, const py::list& map) {
                                std::to_string(c) + ")");
     }
   }
+
 }
 
 void MettaGrid::_make_buffers(unsigned int num_agents) {
@@ -502,9 +507,10 @@ void MettaGrid::_compute_observation(GridCoord observer_row,
     int obs_r = r - static_cast<int>(observer_row) + static_cast<int>(obs_height_radius);
     int obs_c = c - static_cast<int>(observer_col) + static_cast<int>(obs_width_radius);
 
-    // Encode location and add tokens
+    // Encode location and add tokens (pass agent_idx for agent-specific observations like per-agent cooldown)
     uint8_t location = PackedCoordinate::pack(static_cast<uint8_t>(obs_r), static_cast<uint8_t>(obs_c));
-    attempted_tokens_written += _obs_encoder->encode_tokens(obj, obs_tokens, location);
+    attempted_tokens_written +=
+        _obs_encoder->encode_tokens(obj, obs_tokens, location, static_cast<unsigned int>(agent_idx));
     tokens_written = std::min(attempted_tokens_written, static_cast<size_t>(observation_view.shape(1)));
   }
 
@@ -598,9 +604,29 @@ void MettaGrid::_step() {
     }
   }
 
-  // Check and apply damage for all agents
+  // Apply cell effects to agents (AOE effects from nearby objects)
   for (auto* agent : _agents) {
-    agent->check_and_apply_damage(_rng);
+    const CellEffect& effect = _grid->effect_at(agent->location.r, agent->location.c);
+    for (const auto& [item, delta] : effect.resource_deltas) {
+      agent->inventory.update(item, delta);
+
+      // Track AOE stats
+      const std::string& resource_name = _stats->resource_name(item);
+      if (delta > 0) {
+        agent->stats.add("aoe." + resource_name + ".gained", static_cast<float>(delta));
+        _stats->add("aoe." + resource_name + ".gained", static_cast<float>(delta));
+      } else if (delta < 0) {
+        agent->stats.add("aoe." + resource_name + ".lost", static_cast<float>(-delta));
+        _stats->add("aoe." + resource_name + ".lost", static_cast<float>(-delta));
+      }
+      agent->stats.add("aoe." + resource_name + ".delta", static_cast<float>(delta));
+      _stats->add("aoe." + resource_name + ".delta", static_cast<float>(delta));
+    }
+  }
+
+  // Check and apply damage for all agents (randomized order for fairness)
+  for (const auto& agent_idx : agent_indices) {
+    _agents[agent_idx]->check_and_apply_damage(_rng);
   }
 
   // Apply global systems
@@ -1009,9 +1035,19 @@ PYBIND11_MODULE(mettagrid_c, m) {
       .def_readwrite("cost", &DemolishConfig::cost)
       .def_readwrite("scrap", &DemolishConfig::scrap);
 
+  // Bind AOEEffectConfig for AOE effects on any object
+  py::class_<AOEEffectConfig>(m, "AOEEffectConfig")
+      .def(py::init<>())
+      .def(py::init<unsigned int, const std::unordered_map<InventoryItem, InventoryDelta>&>(),
+           py::arg("range") = 1,
+           py::arg("resource_deltas") = std::unordered_map<InventoryItem, InventoryDelta>())
+      .def_readwrite("range", &AOEEffectConfig::range)
+      .def_readwrite("resource_deltas", &AOEEffectConfig::resource_deltas);
+
   // Expose this so we can cast python WallConfig / AgentConfig to a common GridConfig cpp object.
   py::class_<GridObjectConfig, std::shared_ptr<GridObjectConfig>>(m, "GridObjectConfig")
-      .def_readwrite("demolish", &GridObjectConfig::demolish);
+      .def_readwrite("demolish", &GridObjectConfig::demolish)
+      .def_readwrite("aoe", &GridObjectConfig::aoe);
 
   bind_wall_config(m);
 

packages/mettagrid/cpp/include/mettagrid/actions/attack.hpp

@@ -152,7 +152,7 @@ class Attack : public ActionHandler {
     actor.stats.incr(_action_prefix(actor_group) + "demolish." + target_type_name);
     actor.stats.incr("demolish." + target_type_name);
 
-    // Call on_demolish before removing
+    // Call on_demolish before removing (handles AOE cleanup)
     target->on_demolish();
 
     // Give scrap resources to actor

packages/mettagrid/cpp/include/mettagrid/core/grid.hpp

@@ -3,6 +3,7 @@
 
 #include <algorithm>
 #include <memory>
+#include <unordered_map>
 #include <vector>
 
 #include "core/grid_object.hpp"
@@ -13,6 +14,26 @@ using std::unique_ptr;
 using std::vector;
 using GridType = std::vector<std::vector<GridObject*>>;
 
+// Accumulated resource effects at a single grid cell
+struct CellEffect {
+  std::unordered_map<InventoryItem, InventoryDelta> resource_deltas;
+
+  void add(InventoryItem item, InventoryDelta delta) {
+    resource_deltas[item] += delta;
+    if (resource_deltas[item] == 0) {
+      resource_deltas.erase(item);
+    }
+  }
+
+  void subtract(InventoryItem item, InventoryDelta delta) {
+    add(item, -delta);
+  }
+
+  bool empty() const {
+    return resource_deltas.empty();
+  }
+};
+
 class Grid {
 public:
   const GridCoord height;
@@ -21,29 +42,20 @@ class Grid {
 
 private:
   GridType grid;
+  std::vector<std::vector<CellEffect>> _effects;
 
 public:
   Grid(GridCoord height, GridCoord width)
-      : height(height),
-        width(width),
-        objects(),  // Initialize objects in member init list
-        grid() {    // Initialize grid in member init list
+      : height(height), width(width), objects(), grid(), _effects(height, std::vector<CellEffect>(width)) {
     grid.resize(height, std::vector<GridObject*>(width, nullptr));
 
     // Reserve space for objects to avoid frequent reallocations
-    // Assume ~50% of grid cells will contain objects
     size_t estimated_objects = static_cast<size_t>(height) * width / 2;
-
-    // Cap preallocation at ~100MB of pointer memory
     constexpr size_t MAX_PREALLOCATED_OBJECTS = 12'500'000;
     size_t reserved_objects = std::min(estimated_objects, MAX_PREALLOCATED_OBJECTS);
-
     objects.reserve(reserved_objects);
 
     // GridObjectId "0" is reserved to mean empty space (GridObject pointer = nullptr).
-    // By pushing nullptr at index 0, we ensure that:
-    //   1. Grid initialization with zeros automatically represents empty spaces
-    //   2. Object IDs match their index in the objects vector (no off-by-one adjustments)
     objects.push_back(nullptr);
   }
 
@@ -64,6 +76,11 @@ class Grid {
     obj->id = static_cast<GridObjectId>(this->objects.size());
     this->objects.push_back(std::unique_ptr<GridObject>(obj));
     this->grid[obj->location.r][obj->location.c] = obj;
+
+    // Register AOE effects if configured
+    obj->aoe.init(this);
+    obj->aoe.register_effects(obj->location.r, obj->location.c);
+
     return true;
   }
 
@@ -104,6 +121,9 @@ class Grid {
       return false;  // Object not at expected location
     }
 
+    // Call on_demolish for cleanup (includes AOE unregistration)
+    obj.on_demolish();
+
     // Clear the grid cell
     grid[obj.location.r][obj.location.c] = nullptr;
 
@@ -129,6 +149,41 @@ class Grid {
   inline bool is_empty(GridCoord row, GridCoord col) const {
     return grid[row][col] == nullptr;
   }
+
+  // Get the effect at a cell
+  const CellEffect& effect_at(GridCoord r, GridCoord c) const {
+    return _effects[r][c];
+  }
+
+  // Apply AOE effects from a source at (center_r, center_c) with given radius
+  // If adding=true, adds effects; if false, removes them
+  void apply_aoe(GridCoord center_r,
+                 GridCoord center_c,
+                 unsigned int radius,
+                 const std::unordered_map<InventoryItem, InventoryDelta>& resource_deltas,
+                 bool adding) {
+    int r_start = std::max(0, static_cast<int>(center_r) - static_cast<int>(radius));
+    int r_end = std::min(static_cast<int>(height), static_cast<int>(center_r) + static_cast<int>(radius) + 1);
+    int c_start = std::max(0, static_cast<int>(center_c) - static_cast<int>(radius));
+    int c_end = std::min(static_cast<int>(width), static_cast<int>(center_c) + static_cast<int>(radius) + 1);
+
+    for (int r = r_start; r < r_end; ++r) {
+      for (int c = c_start; c < c_end; ++c) {
+        // Manhattan distance for AOE range
+        int dr = std::abs(r - static_cast<int>(center_r));
+        int dc = std::abs(c - static_cast<int>(center_c));
+        if (dr + dc <= static_cast<int>(radius)) {
+          for (const auto& [item, delta] : resource_deltas) {
+            if (adding) {
+              _effects[r][c].add(item, delta);
+            } else {
+              _effects[r][c].subtract(item, delta);
+            }
+          }
+        }
+      }
+    }
+  }
 };
 
 #endif  // PACKAGES_METTAGRID_CPP_INCLUDE_METTAGRID_CORE_GRID_HPP_

packages/mettagrid/cpp/include/mettagrid/core/grid_object.hpp

@@ -13,6 +13,9 @@
 #include "objects/constants.hpp"
 #include "objects/has_vibe.hpp"
 
+// Forward declaration
+class Grid;
+
 using TypeId = ObservationType;
 using ObservationCoord = ObservationType;
 using Vibe = ObservationType;
@@ -61,19 +64,74 @@ struct DemolishConfig {
       : cost(cost), scrap(scrap) {}
 };
 
+// Configuration for Area of Effect (AOE) resource effects
+struct AOEEffectConfig {
+  unsigned int range = 1;                                             // Radius of effect (Manhattan distance)
+  std::unordered_map<InventoryItem, InventoryDelta> resource_deltas;  // Per-tick resource changes
+
+  AOEEffectConfig() = default;
+  AOEEffectConfig(unsigned int range, const std::unordered_map<InventoryItem, InventoryDelta>& resource_deltas)
+      : range(range), resource_deltas(resource_deltas) {}
+};
+
 struct GridObjectConfig {
   TypeId type_id;
   std::string type_name;
   std::vector<int> tag_ids;
   ObservationType initial_vibe;
   std::optional<DemolishConfig> demolish;  // If set, object can be demolished
+  std::optional<AOEEffectConfig> aoe;      // If set, object emits AOE effects
 
   GridObjectConfig(TypeId type_id, const std::string& type_name, ObservationType initial_vibe = 0)
-      : type_id(type_id), type_name(type_name), tag_ids({}), initial_vibe(initial_vibe), demolish(std::nullopt) {}
+      : type_id(type_id),
+        type_name(type_name),
+        tag_ids({}),
+        initial_vibe(initial_vibe),
+        demolish(std::nullopt),
+        aoe(std::nullopt) {}
 
   virtual ~GridObjectConfig() = default;
 };
 
+// Helper class for managing AOE effects on grid objects
+class AOEHelper {
+public:
+  AOEHelper() = default;
+
+  // Initialize with grid reference
+  void init(Grid* grid) {
+    _grid = grid;
+  }
+
+  // Set the AOE config (call from object constructor)
+  void set_config(const AOEEffectConfig* config) {
+    _config = config;
+  }
+
+  // Check if this helper has AOE configured
+  bool has_aoe() const {
+    return _config != nullptr && _grid != nullptr;
+  }
+
+  // Register AOE effects at the given location
+  void register_effects(GridCoord r, GridCoord c);
+
+  // Unregister AOE effects (call on demolish or removal)
+  void unregister_effects();
+
+  // Get the config
+  const AOEEffectConfig* config() const {
+    return _config;
+  }
+
+private:
+  const AOEEffectConfig* _config = nullptr;
+  Grid* _grid = nullptr;
+  bool _registered = false;
+  GridCoord _location_r = 0;
+  GridCoord _location_c = 0;
+};
+
 class GridObject : public HasVibe {
 public:
   GridObjectId id{};
@@ -82,6 +140,7 @@ class GridObject : public HasVibe {
   std::string type_name;
   std::vector<int> tag_ids;
   const DemolishConfig* demolish_config = nullptr;  // Optional demolish config for buildings
+  AOEHelper aoe;                                    // AOE effect helper
 
   virtual ~GridObject() = default;
 
@@ -90,24 +149,34 @@ class GridObject : public HasVibe {
             const GridLocation& object_location,
             const std::vector<int>& tags,
             ObservationType object_vibe = 0,
-            const DemolishConfig* demolish = nullptr) {
+            const DemolishConfig* demolish = nullptr,
+            const std::optional<AOEEffectConfig>& aoe_config = std::nullopt) {
     this->type_id = object_type_id;
     this->type_name = object_type_name;
     this->location = object_location;
     this->tag_ids = tags;
     this->vibe = object_vibe;
     this->demolish_config = demolish;
+    if (aoe_config.has_value()) {
+      _aoe_config = aoe_config.value();
+      this->aoe.set_config(&_aoe_config.value());
+    }
   }
 
   // Called when this object is demolished. Override for cleanup.
-  virtual void on_demolish() {}
+  virtual void on_demolish() {
+    aoe.unregister_effects();
+  }
 
   // observer_agent_id: The agent observing this object (UINT_MAX means no specific observer)
   // Used by Assembler to report agent-specific cooldowns
   virtual std::vector<PartialObservationToken> obs_features(unsigned int observer_agent_id = UINT_MAX) const {
     (void)observer_agent_id;  // Unused in base class
     return {};                // Default: no observable features
   }
+
+private:
+  std::optional<AOEEffectConfig> _aoe_config;
 };
 
 #endif  // PACKAGES_METTAGRID_CPP_INCLUDE_METTAGRID_CORE_GRID_OBJECT_HPP_

packages/mettagrid/cpp/include/mettagrid/objects/assembler.hpp

@@ -276,7 +276,7 @@ class Assembler : public GridObject, public Usable {
         agent_cooldown(cfg.agent_cooldown),
         clipper_ptr(nullptr) {
     const DemolishConfig* demolish = cfg.demolish.has_value() ? &cfg.demolish.value() : nullptr;
-    GridObject::init(cfg.type_id, cfg.type_name, GridLocation(r, c), cfg.tag_ids, cfg.initial_vibe, demolish);
+    GridObject::init(cfg.type_id, cfg.type_name, GridLocation(r, c), cfg.tag_ids, cfg.initial_vibe, demolish, cfg.aoe);
   }
   virtual ~Assembler() = default;