Fix PR #77 review findings: revert transition() override break, ndarray edge cases

Agent.transition()'s widened return type broke its override of
Entity.transition() (reportIncompatibleMethodOverride), so it's reverted back
to -> None; collect_turn_stats goes back to reconstructing agent stats from
agent.model.memory tail-index arithmetic instead. Also hardens
TensorboardLogger.record_step's np.ndarray branch against empty arrays
(previously fell through to a crashing float()) and caps per-element scalar
explosion for large arrays.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

Author: bing-j

sorrel/agents/agent.py

@@ -152,7 +152,7 @@ def model_take_action(self, state: np.ndarray):
             return threadsafe_model.threadsafe_take_action(state)
         return self.model.take_action(state)
 
-    def transition(self, world: W) -> tuple[int, float, bool]:
+    def transition(self, world: W) -> None:
         """Processes a full transition step for the agent.
 
         This function does the following:
@@ -163,10 +163,6 @@ def transition(self, world: W) -> tuple[int, float, bool]:
 
         Args:
             env (Gridworld): the environment that this agent is acting in.
-
-        Returns:
-            tuple[int, float, bool]: the ``(action, reward, done)`` values computed
-                for this transition.
         """
         state = self.pov(world)
         action = self.get_action(state)
@@ -175,7 +171,6 @@ def transition(self, world: W) -> tuple[int, float, bool]:
 
         world.total_reward += reward
         self.add_memory(state, action, reward, done)
-        return action, reward, done
 
 
 class MovingAgent[W: Gridworld](Agent[W]):

sorrel/environment.py

@@ -102,11 +102,10 @@ def take_turn(self, epoch: int = 0) -> None:
             x: Entity
             if x.has_transitions and not isinstance(x, Agent):
                 x.transition(self.world)
-        agent_results: list[tuple[int, float, bool]] = [
-            agent.transition(self.world) for agent in self.agents
-        ]
+        for agent in self.agents:
+            agent.transition(self.world)
         if self._active_logger is not None:
-            stats = self.collect_turn_stats(epoch, agent_results)
+            stats = self.collect_turn_stats(epoch)
             self.on_turn_end(stats)
 
     def _model_start_epoch_action(self, agent: Agent[W], epoch: int) -> None:
@@ -121,9 +120,7 @@ def _model_train_step(self, agent: Agent[W]) -> np.ndarray:
         """Run model train step."""
         return agent.model.train_step()
 
-    def collect_turn_stats(
-        self, epoch: int, agent_results: list[tuple[int, float, bool]]
-    ) -> TurnStats:
+    def collect_turn_stats(self, epoch: int) -> TurnStats:
         """Collect per-turn statistics after all transitions have run.
 
         Called automatically by :meth:`take_turn` (when an active logger is
@@ -134,25 +131,25 @@ def collect_turn_stats(
 
         Args:
             epoch: Current epoch index.
-            agent_results: The ``(action, reward, done)`` tuples returned by
-                each agent's :meth:`~sorrel.agents.agent.Agent.transition`
-                call this turn, in the same order as :attr:`agents`.
 
         Returns:
             A :class:`~sorrel.utils.turn_stats.TurnStats` snapshot for this turn.
         """
-        agent_stats: list[AgentTurnStats] = [
-            AgentTurnStats(
-                agent_id=i,
-                location=tuple(agent.location),
-                last_action=action,
-                last_reward=reward,
-                last_done=done,
-            )
-            for i, (agent, (action, reward, done)) in enumerate(
-                zip(self.agents, agent_results)
+        agent_stats: list[AgentTurnStats] = []
+        for i, agent in enumerate(self.agents):
+            mem = agent.model.memory
+            if mem.size == 0:
+                continue
+            tail = (mem.idx - 1) % mem.capacity
+            agent_stats.append(
+                AgentTurnStats(
+                    agent_id=i,
+                    location=tuple(agent.location),
+                    last_action=int(mem.actions[tail]),
+                    last_reward=float(mem.rewards[tail]),
+                    last_done=bool(mem.dones[tail]),
+                )
             )
-        ]
         return TurnStats(
             epoch=epoch,
             turn=self.turn,

sorrel/utils/logging.py

@@ -189,6 +189,10 @@ class TensorboardLogger(Logger):
         additional_values: A dictionary of optional values to be stored.
     """
 
+    # Above this many elements, a multi-element np.ndarray extra is logged as a
+    # single mean scalar instead of one TensorBoard series per element.
+    _MAX_EXTRA_ARRAY_ELEMENTS = 32
+
     def __init__(self, max_epochs: int, log_dir: str | os.PathLike, *args):
         """Initialize a Tensorboard log.
 
@@ -245,7 +249,14 @@ def record_step(self, stats: TurnStats) -> None:
         for key, value in stats.extra.items():
             if isinstance(value, dict):
                 self.writer.add_scalars(f"turn/{key}", value, step)
-            elif isinstance(value, np.ndarray) and value.size > 1:
+            elif isinstance(value, np.ndarray) and value.size != 1:
+                if value.size == 0:
+                    continue
+                if value.size > self._MAX_EXTRA_ARRAY_ELEMENTS:
+                    self.writer.add_scalar(
+                        f"turn/{key}_mean", float(value.mean()), step
+                    )
+                    continue
                 per_element = {str(i): float(v) for i, v in enumerate(value.ravel())}
                 self.writer.add_scalars(f"turn/{key}", per_element, step)
             else: