TrackioLogger

Author: Xmaster6y

torchrl/record/loggers/trackio.py

@@ -0,0 +1,146 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+from __future__ import annotations
+
+import importlib.util
+
+from collections.abc import Sequence
+
+from torch import Tensor
+
+from .common import Logger
+
+_has_trackio = importlib.util.find_spec("trackio") is not None
+_has_omegaconf = importlib.util.find_spec("omegaconf") is not None
+
+
+class TrackioLogger(Logger):
+    """Wrapper for the trackio logger.
+
+    Args:
+        exp_name (str): The name of the experiment.
+        project (str): The name of the project.
+
+    Keyword Args:
+        fps (int, optional): Number of frames per second when recording videos. Defaults to ``30``.
+        **kwargs: Extra keyword arguments for ``trackio.init``.
+
+    """
+
+    @classmethod
+    def __new__(cls, *args, **kwargs):
+        return super().__new__(cls)
+
+    def __init__(
+        self,
+        exp_name: str,
+        project: str,
+        *,
+        video_fps: int = 32,
+        **kwargs,
+    ) -> None:
+        if not _has_trackio:
+            raise ImportError("trackio could not be imported")
+
+        self.video_fps = video_fps
+        self._trackio_kwargs = {
+            "name": exp_name,
+            "project": project,
+            "resume": "allow",
+            **kwargs,
+        }
+
+        super().__init__(exp_name=exp_name, log_dir=project)
+
+    def _create_experiment(self):
+        """Creates a trackio experiment.
+
+        Args:
+            exp_name (str): The name of the experiment.
+
+        Returns:
+            A trackio.Experiment object.
+        """
+        if not _has_trackio:
+            raise ImportError("Trackio is not installed")
+        import trackio
+
+        return trackio.init(**self._trackio_kwargs)
+
+    def log_scalar(
+        self, name: str, value: float, step: int | None = None
+    ) -> None:
+        """Logs a scalar value to trackio.
+
+        Args:
+            name (str): The name of the scalar.
+            value (float): The value of the scalar.
+            step (int, optional): The step at which the scalar is logged.
+                Defaults to None.
+        """
+        self.experiment.log({name: value}, step=step)
+
+    def log_video(self, name: str, video: Tensor, **kwargs) -> None:
+        """Log videos inputs to trackio.
+
+        Args:
+            name (str): The name of the video.
+            video (Tensor): The video to be logged.
+            **kwargs: Other keyword arguments. By construction, log_video
+                supports 'step' (integer indicating the step index), 'format'
+                (default is 'mp4') and 'fps' (defaults to ``self.video_fps``). Other kwargs are
+                passed as-is to the :obj:`experiment.log` method.
+        """
+        import trackio
+
+        fps = kwargs.pop("fps", self.video_fps)
+        format = kwargs.pop("format", "mp4")
+        self.experiment.log(
+            {name: trackio.Video(video, fps=fps, format=format)},
+            **kwargs,
+        )
+
+    def log_hparams(self, cfg: DictConfig | dict) -> None:  # noqa: F821
+        """Logs the hyperparameters of the experiment.
+
+        Args:
+            cfg (DictConfig or dict): The configuration of the experiment.
+
+        """
+        if type(cfg) is not dict and _has_omegaconf:
+            if not _has_omegaconf:
+                raise ImportError(
+                    "OmegaConf could not be imported. "
+                    "Cannot log hydra configs without OmegaConf."
+                )
+            from omegaconf import OmegaConf
+
+            cfg = OmegaConf.to_container(cfg, resolve=True)
+        self.experiment.config.update(cfg)
+
+    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.")
+
+    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.
+
+        Args:
+            name (str): The name of the string data.
+            value (str): The string value to log.
+            step (int, optional): The step at which the string is logged.
+                Defaults to None.
+        """
+        import trackio
+
+        # 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})