Add logging

emote/memory/memory.py

@@ -5,16 +5,20 @@
 of the sequence builder is to consume individual timesteps per agent and collate
 them into episodes before submission into the memory.
 """
+from __future__ import annotations
 
 import logging
 import os
 import time
 import warnings
 
 from collections import defaultdict
+from collections.abc import Iterable
 from dataclasses import dataclass, field
 from typing import Callable, Dict, List, Mapping, Optional, Tuple, Union
 
+import torch
+
 from torch.utils.tensorboard import SummaryWriter
 
 from emote.callback import Callback
@@ -191,31 +195,79 @@ def add(
         self.completed_episodes += len(observations) - len(responses)
 
         for obs in observations.values():
-            # todo: handle info lists how?
-
             if obs.metadata is None:
                 continue
 
-            # all infos for agents are windowed
-            for k, v in obs.metadata.info.items():
-                if k.startswith("histogram:"):
-                    continue
-
-                self.log_windowed_scalar(k, v)
+            self.report(obs.metadata.info, obs.metadata.info_lists)
 
         if (self._counter % self._log_interval) == 0:
             self._end_cycle()
             self._counter = 0
 
         return self._inner.add(observations, responses)
 
+    def report(
+        self,
+        metrics: dict[str, float],
+        metrics_lists: dict[str, list[float]],
+    ):
+
+        for key, value in metrics.items():
+            if key.startswith("histogram:"):
+                self.log_histogram(key[10:], value)
+            else:
+                self.log_windowed_scalar(key, value)
+
+        for key, value in metrics_lists.items():
+            if key.startswith("histogram:"):
+                self.log_histogram(key[10:], value)
+            else:
+                self.log_windowed_scalar(key, value)
+
+    def get_report(
+        self, keys: List[str]
+    ) -> Tuple[dict[str, Union[int, float, list[float]]], dict[str, list[float]]]:
+        keys = set(keys)
+        out = {}
+        out_lists = {}
+
+        for key in keys:
+            if key.startswith("histogram:") and key[10:] in self.hist_logs:
+                window = self.hist_logs[key[10:]]
+                out[key] = sum(window) / len(window)
+            elif key in self.windowed_scalar:
+                window = self.windowed_scalar[key]
+                out_lists[key] = list(window)
+                out[key] = sum(window) / len(window)
+                out[f"{key}/cumulative"] = self.windowed_scalar_cumulative[key]
+
+        return out, out_lists
+
     def _end_cycle(self):
         now_time = time.perf_counter()
         cycle_time = now_time - self._cycle_start_time
         cycle_infs = self.completed_inferences - self._cycle_start_infs
         inf_step = self.completed_inferences
         self.log_scalar("training/inf_per_sec", cycle_infs / cycle_time)
         self.log_scalar("episode/completed", self.completed_episodes)
+
+        if "episode/reward" in self.windowed_scalar:
+            rewards = self.windowed_scalar["episode/reward"]
+            average_reward = sum(rewards) / len(rewards)
+            rewards_tensor = torch.Tensor(rewards)
+
+            self._writer.add_scalar(
+                "env_vs_episode/reward", average_reward, self.completed_episodes
+            )
+            self._writer.add_histogram(
+                "episode/reward_distribution", rewards_tensor, inf_step
+            )
+            self._writer.add_histogram(
+                "env_vs_episode/reward_distribution",
+                rewards_tensor,
+                self.completed_episodes,
+            )
+
         suffix = "inf_step"
         for k, v in self.scalar_logs.items():
             if suffix:
@@ -254,6 +306,13 @@ def _end_cycle(self):
                 k = "/".join(k_split)
             self._writer.add_video(k, video_array, inf_step, fps=fps, walltime=None)
 
+        for k, v in self.hist_logs:
+            if suffix:
+                k_split = k.split("/")
+                k_split[0] = k_split[0] + "_" + suffix
+                k = "/".join(k_split)
+            self._writer.add_histogram(k, v, inf_step)
+
         time_since_start = time.monotonic() - self._start_time
 
         self._writer.add_scalar(

emote/mixins/logging.py

@@ -1,4 +1,5 @@
 from collections import deque
+from collections.abc import Iterable
 from typing import Any, Dict, Tuple, Union
 
 import numpy as np
@@ -33,7 +34,11 @@ def log_scalar(self, key: str, value: Union[float, torch.Tensor]):
         else:
             self.scalar_logs[key] = value
 
-    def log_windowed_scalar(self, key: str, value: Union[float, torch.Tensor]):
+    def log_windowed_scalar(
+        self,
+        key: str,
+        value: Union[float, torch.Tensor, Iterable[Union[torch.Tensor, float]]],
+    ):
         """Log scalars using a moving window average.
 
         By default this will use `default_window_length` from the constructor as the window
@@ -54,10 +59,14 @@ def log_windowed_scalar(self, key: str, value: Union[float, torch.Tensor]):
             self.windowed_scalar[key] = deque(maxlen=length)
             self.windowed_scalar_cumulative[key] = 0
 
-        if isinstance(value, torch.Tensor):
-            self.windowed_scalar[key].append(value.item())
+        if isinstance(value, Iterable):
+            val = value.numpy() if isinstance(value, torch.Tensor) else value
+            self.windowed_scalar[key].extend(val)
+            self.windowed_scalar_cumulative[key] += sum(val)
         else:
-            self.windowed_scalar[key].append(value)
+            val = value.item() if isinstance(value, torch.Tensor) else value
+            self.windowed_scalar[key].append(val)
+            self.windowed_scalar_cumulative[key] += val
 
     def log_image(self, key: str, value: torch.Tensor):
         """Use log_image to periodically log image data."""
@@ -68,8 +77,20 @@ def log_video(self, key: str, value: Tuple[np.ndarray, int]):
         """Use log_scalar to periodically log scalar data."""
         self.video_logs[key] = value
 
-    def log_histogram(self, key: str, value: torch.Tensor):
-        self.hist_logs[key] = value.detach()
+    def log_histogram(
+        self,
+        key: str,
+        value: Union[torch.Tensor, float, Iterable[Union[torch.Tensor, float]]],
+    ):
+        if isinstance(value, Iterable):
+            self.hist_logs[key] = (
+                value.detach() if isinstance(value, torch.Tensor) else value
+            )
+        else:
+            if key not in self.hist_logs:
+                self.hist_logs[key] = deque(maxlen=self._default_window_length)
+
+            self.hist_logs[key].append(value)
 
     def state_dict(self):
         state_dict = super().state_dict()