trackio histograms and str

Author: Xmaster6y

test/test_loggers.py

@@ -481,6 +481,15 @@ def test_log_scalar(self, steps, trackio_logger):
                 step=steps[i] if steps else None,
             )
 
+    @pytest.mark.parametrize("steps", [None, [1, 10, 11]])
+    def test_log_str(self, steps, trackio_logger):
+        for i in range(3):
+            trackio_logger.log_str(
+                name="foo",
+                value="bar",
+                step=steps[i] if steps else None,
+            )
+
     def test_log_video(self, trackio_logger):
         torch.manual_seed(0)
 
@@ -509,10 +518,14 @@ def test_log_hparams(self, trackio_logger, config):
         for key, value in config.items():
             assert trackio_logger.experiment.config[key] == value
 
-    def test_log_histogram(self, trackio_logger):
-        with pytest.raises(NotImplementedError):
-            data = torch.randn(10)
-            trackio_logger.log_histogram("hist", data, step=0, bins=2)
+    @pytest.mark.parametrize("steps", [None, [1, 10, 11]])
+    def test_log_histogram(self, steps, trackio_logger):
+        torch.manual_seed(0)
+        for i in range(3):
+            data = torch.randn(100)
+            trackio_logger.log_histogram(
+                "hist", data, step=steps[i] if steps else None, bins=10
+            )
 
 
 if __name__ == "__main__":

torchrl/record/loggers/trackio.py

@@ -128,7 +128,24 @@ def __repr__(self) -> str:
         return f"TrackioLogger(experiment={self.experiment.__repr__()})"
 
     def log_histogram(self, name: str, data: Sequence, **kwargs):
-        raise NotImplementedError("Logging histograms in trackio is not permitted.")
+        """Add histogram to log.
+
+        Args:
+            name (str): Data identifier
+            data (torch.Tensor, numpy.ndarray): Values to build histogram
+
+        Keyword Args:
+            step (int): Global step value to record
+            bins (int): Number of bins to use for the histogram
+
+        """
+        import trackio
+
+        num_bins = kwargs.pop("bins", None)
+        step = kwargs.pop("step", None)
+        self.experiment.log(
+            {name: trackio.Histogram(data, num_bins=num_bins)}, step=step
+        )
 
     def log_str(self, name: str, value: str, step: int | None = None) -> None:
         """Logs a string value to trackio using a table format for better visualization.
@@ -143,8 +160,4 @@ def log_str(self, name: str, value: str, step: int | None = None) -> None:
 
         # Create a table with a single row
         table = trackio.Table(columns=["text"], data=[[value]])
-
-        if step is not None:
-            self.experiment.log({name: value}, step=step)
-        else:
-            self.experiment.log({name: table})
+        self.experiment.log({name: table}, step=step)