cp

Author: daveey

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

@@ -4,6 +4,7 @@
 #include <algorithm>
 #include <cassert>
 #include <cmath>
+#include <iostream>
 #include <stdexcept>
 #include <string>
 #include <unordered_map>
@@ -250,6 +251,12 @@ class Assembler : public GridObject, public Usable {
   // Chest search distance - if > 0, assembler can use inventories from chests within this distance
   unsigned int chest_search_distance;
 
+  // Per-agent cooldown duration - number of timesteps before an agent can use this assembler again
+  unsigned int agent_cooldown;
+
+  // Per-agent cooldown tracking - maps agent_id to timestep when cooldown ends
+  std::vector<unsigned int> agent_cooldown_ends;
+
   Assembler(GridCoord r, GridCoord c, const AssemblerConfig& cfg, StatsTracker* stats)
       : protocols(build_protocol_map(cfg.protocols)),
         unclip_protocols(),
@@ -266,6 +273,7 @@ class Assembler : public GridObject, public Usable {
         obs_encoder(nullptr),
         allow_partial_usage(cfg.allow_partial_usage),
         chest_search_distance(cfg.chest_search_distance),
+        agent_cooldown(cfg.agent_cooldown),
         clipper_ptr(nullptr) {
     GridObject::init(cfg.type_id, cfg.type_name, GridLocation(r, c), cfg.tag_ids, cfg.initial_vibe);
   }
@@ -287,9 +295,36 @@ class Assembler : public GridObject, public Usable {
   }
 
   // Initialize the per-agent tracking array (call after knowing num_agents)
-  // Note: This is a no-op for now since we removed per-agent cooldown, but kept for API compat
-  void init_agent_tracking(unsigned int /*num_agents*/) {
-    // No-op - per-agent cooldown was removed
+  void init_agent_tracking(unsigned int num_agents) {
+    if (agent_cooldown > 0) {
+      agent_cooldown_ends.resize(num_agents, 0);
+    }
+  }
+
+  // Check if a specific agent is on cooldown
+  bool is_agent_on_cooldown(unsigned int agent_id) const {
+    if (agent_cooldown == 0 || agent_id >= agent_cooldown_ends.size()) {
+      return false;
+    }
+    return current_timestep_ptr && agent_cooldown_ends[agent_id] > *current_timestep_ptr;
+  }
+
+  // Get remaining cooldown for a specific agent
+  unsigned int get_agent_cooldown_remaining(unsigned int agent_id) const {
+    if (agent_cooldown == 0 || agent_id >= agent_cooldown_ends.size() || !current_timestep_ptr) {
+      return 0;
+    }
+    if (agent_cooldown_ends[agent_id] <= *current_timestep_ptr) {
+      return 0;
+    }
+    return agent_cooldown_ends[agent_id] - *current_timestep_ptr;
+  }
+
+  // Set cooldown for a specific agent
+  void set_agent_cooldown(unsigned int agent_id) {
+    if (agent_cooldown > 0 && agent_id < agent_cooldown_ends.size() && current_timestep_ptr) {
+      agent_cooldown_ends[agent_id] = *current_timestep_ptr + agent_cooldown;
+    }
   }
 
   // Get the remaining cooldown duration in ticks (0 when ready for use)
@@ -439,6 +474,11 @@ class Assembler : public GridObject, public Usable {
       return false;
     }
 
+    // Check per-agent cooldown
+    if (is_agent_on_cooldown(actor.agent_id)) {
+      return false;
+    }
+
     // Check if on cooldown and whether partial usage is allowed
     unsigned int remaining = cooldown_remaining();
     if (remaining > 0 && !allow_partial_usage) {
@@ -497,6 +537,9 @@ class Assembler : public GridObject, public Usable {
     cooldown_duration = static_cast<unsigned int>(protocol_to_use.cooldown);
     cooldown_end_timestep = *current_timestep_ptr + cooldown_duration;
 
+    // Set per-agent cooldown
+    set_agent_cooldown(actor.agent_id);
+
     // If we were clipped and successfully used an unclip protocol, become unclipped. Also, don't count this as a use.
     if (is_clipped) {
       become_unclipped();
@@ -515,6 +558,14 @@ class Assembler : public GridObject, public Usable {
       features.push_back({ObservationFeature::CooldownRemaining, static_cast<ObservationType>(remaining)});
     }
 
+    // Add per-agent cooldown remaining if the observer is a valid agent
+    if (observer_agent_id != UINT_MAX && agent_cooldown > 0) {
+      unsigned int agent_remaining = std::min(get_agent_cooldown_remaining(observer_agent_id), 255u);
+      if (agent_remaining > 0) {
+        features.push_back({ObservationFeature::AgentCooldownRemaining, static_cast<ObservationType>(agent_remaining)});
+      }
+    }
+
     // Add clipped status to observations if clipped
     if (is_clipped) {
       features.push_back({ObservationFeature::Clipped, static_cast<ObservationType>(1)});

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

@@ -17,10 +17,11 @@ struct AssemblerConfig : public GridObjectConfig {
   AssemblerConfig(TypeId type_id, const std::string& type_name, ObservationType initial_vibe = 0)
       : GridObjectConfig(type_id, type_name, initial_vibe),
         allow_partial_usage(false),
-        max_uses(0),                 // 0 means unlimited uses
-        clip_immune(false),          // Not immune by default
-        start_clipped(false),        // Not clipped at start by default
-        chest_search_distance(0) {}  // 0 means chests are not searched
+        max_uses(0),               // 0 means unlimited uses
+        clip_immune(false),        // Not immune by default
+        start_clipped(false),      // Not clipped at start by default
+        chest_search_distance(0),  // 0 means chests are not searched
+        agent_cooldown(0) {}       // 0 means no per-agent cooldown
 
   // List of protocols - GroupVibe keys will be calculated from each protocol's vibes vector
   std::vector<std::shared_ptr<Protocol>> protocols;
@@ -40,6 +41,10 @@ struct AssemblerConfig : public GridObjectConfig {
   // Distance is measured as Chebyshev distance (max of row and column differences)
   // 0 means chests are not searched
   unsigned int chest_search_distance;
+
+  // Per-agent cooldown duration - number of timesteps before an agent can use this assembler again
+  // 0 means no per-agent cooldown
+  unsigned int agent_cooldown;
 };
 
 namespace py = pybind11;
@@ -59,6 +64,7 @@ inline void bind_assembler_config(py::module& m) {
       .def_readwrite("clip_immune", &AssemblerConfig::clip_immune)
       .def_readwrite("start_clipped", &AssemblerConfig::start_clipped)
       .def_readwrite("chest_search_distance", &AssemblerConfig::chest_search_distance)
+      .def_readwrite("agent_cooldown", &AssemblerConfig::agent_cooldown)
       .def_readwrite("initial_vibe", &AssemblerConfig::initial_vibe);
 }
 

packages/mettagrid/python/src/mettagrid/config/mettagrid_c_config.py

@@ -292,6 +292,7 @@ def convert_to_cpp_game_config(game_config: GameConfig):
             cpp_assembler_config.clip_immune = object_config.clip_immune
             cpp_assembler_config.start_clipped = object_config.start_clipped
             cpp_assembler_config.chest_search_distance = object_config.chest_search_distance
+            cpp_assembler_config.agent_cooldown = object_config.agent_cooldown
             # Key by map_name so map grid (which uses map_name) resolves directly.
             objects_cpp_params[object_config.map_name or object_type] = cpp_assembler_config
         elif isinstance(object_config, ChestConfig):

packages/mettagrid/python/src/mettagrid/config/mettagrid_config.py

@@ -386,6 +386,11 @@ class AssemblerConfig(GridObjectConfig):
         ge=0,
         description="Distance within which assembler can use inventories from chests",
     )
+    agent_cooldown: int = Field(
+        default=0,
+        ge=0,
+        description="Per-agent cooldown duration in timesteps before they can use this assembler again",
+    )
 
 
 class ChestConfig(GridObjectConfig):

packages/mettagrid/python/src/mettagrid/simulator/multi_episode/rollout.py

@@ -31,6 +31,8 @@ class EpisodeRolloutResult(BaseModel):
     replay_path: str | None
     steps: int
     max_steps: int
+    agent_inventories: list[dict[str, int]] | None = None  # agent_id -> {resource_name -> amount}
+    resource_names: list[str] | None = None  # list of resource names
 
     def print_agent_stats(self, title: str = "Agent Summary") -> None:
         """Print agent inventories and stats as a table."""
@@ -167,6 +169,12 @@ def multi_episode_rollout(
             all_replay_paths = episode_replay_writer.get_written_replay_urls()
             replay_path = None if not all_replay_paths else list(all_replay_paths.values())[0]
 
+        # Collect agent inventories
+        agent_inventories = []
+        for agent_id in range(env_cfg.game.num_agents):
+            agent = rollout._sim.agent(agent_id)
+            agent_inventories.append(dict(agent.inventory))
+
         result = EpisodeRolloutResult(
             assignments=assignments.copy(),
             rewards=np.array(rollout._sim.episode_rewards, dtype=float),
@@ -175,6 +183,8 @@ def multi_episode_rollout(
             replay_path=replay_path,
             steps=rollout._sim.current_step,
             max_steps=rollout._sim.config.game.max_steps,
+            agent_inventories=agent_inventories,
+            resource_names=env_cfg.game.resource_names,
         )
 
         episode_results.append(result)