torchrl 0.11.0__cp314-cp314-macosx_11_0_arm64.whl

torchrl/record/loggers/mlflow.py

@@ -0,0 +1,142 @@
+# 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
+
+import os
+from collections.abc import Sequence
+from tempfile import TemporaryDirectory
+from typing import Any
+
+from torch import Tensor
+
+from torchrl.record.loggers.common import Logger
+
+_has_tv = importlib.util.find_spec("torchvision") is not None
+
+_has_mlflow = importlib.util.find_spec("mlflow") is not None
+_has_omegaconf = importlib.util.find_spec("omegaconf") is not None
+
+
+class MLFlowLogger(Logger):
+    """Wrapper for the mlflow logger.
+
+    Args:
+        exp_name (str): The name of the experiment.
+        tracking_uri (str): A tracking URI to a datastore that supports MLFlow or a local directory.
+
+    Keyword Args:
+        fps (int, optional): Number of frames per second when recording videos. Defaults to ``30``.
+
+    """
+
+    def __init__(
+        self,
+        exp_name: str,
+        tracking_uri: str,
+        tags: dict[str, Any] | None = None,
+        *,
+        video_fps: int = 30,
+        **kwargs,
+    ) -> None:
+        import mlflow
+
+        self._mlflow_kwargs = {
+            "name": exp_name,
+            "artifact_location": tracking_uri,
+            "tags": tags,
+        }
+        mlflow.set_tracking_uri(tracking_uri)
+        super().__init__(exp_name=exp_name, log_dir=tracking_uri)
+        self.video_log_counter = 0
+        self.video_fps = video_fps
+
+    def _create_experiment(self) -> mlflow.ActiveRun:  # noqa
+        import mlflow
+
+        """Creates an mlflow experiment.
+
+        Returns:
+            mlflow.ActiveRun: The mlflow experiment object.
+        """
+        if not _has_mlflow:
+            raise ImportError("MLFlow is not installed")
+
+        # Only create experiment if it doesnt exist
+        experiment = mlflow.get_experiment_by_name(self._mlflow_kwargs["name"])
+        if experiment is None:
+            self.id = mlflow.create_experiment(**self._mlflow_kwargs)
+        else:
+            self.id = experiment.experiment_id
+        return mlflow.start_run(experiment_id=self.id)
+
+    def log_scalar(self, name: str, value: float, step: int | None = None) -> None:
+        """Logs a scalar value to mlflow.
+
+        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.
+        """
+        import mlflow
+
+        mlflow.set_experiment(experiment_id=self.id)
+        mlflow.log_metric(key=name, value=value, step=step)
+
+    def log_video(self, name: str, video: Tensor, **kwargs) -> None:
+        """Log video inputs to mlflow.
+
+        Args:
+            name (str): The name of the video.
+            video (Tensor): The video to be logged, expected to be in (T, C, H, W) format
+                for consistency with other loggers.
+            **kwargs: Other keyword arguments. By construction, log_video
+                supports 'step' (integer indicating the step index) and 'fps' (defaults to ``self.video_fps``).
+        """
+        import mlflow
+        import torchvision
+
+        if not _has_tv:
+            raise ImportError(
+                "Logging a video with MLFlow requires torchvision to be installed."
+            )
+        mlflow.set_experiment(experiment_id=self.id)
+        if video.ndim == 5:
+            video = video[-1]  # N T C H W -> T C H W
+        video = video.permute(0, 2, 3, 1)  # T C H W -> T H W C
+        if video.size(dim=-1) != 3:
+            raise ValueError(
+                "The MLFlow logger only supports videos with 3 color channels."
+            )
+        self.video_log_counter += 1
+        fps = kwargs.pop("fps", self.video_fps)
+        step = kwargs.pop("step", None)
+        with TemporaryDirectory() as temp_dir:
+            video_name = f"{name}_step_{step:04}.mp4" if step else f"{name}.mp4"
+            with open(os.path.join(temp_dir, video_name), "wb") as f:
+                torchvision.io.write_video(filename=f.name, video_array=video, fps=fps)
+                mlflow.log_artifact(f.name, "videos")
+
+    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.
+        """
+        import mlflow
+        from omegaconf import OmegaConf
+
+        mlflow.set_experiment(experiment_id=self.id)
+        if type(cfg) is not dict and _has_omegaconf:
+            cfg = OmegaConf.to_container(cfg, resolve=True)
+        mlflow.log_params(cfg)
+
+    def __repr__(self) -> str:
+        return f"MLFlowLogger(experiment={self.experiment.__repr__()})"
+
+    def log_histogram(self, name: str, data: Sequence, **kwargs):
+        raise NotImplementedError("Logging histograms in cvs is not permitted.")

torchrl/record/loggers/tensorboard.py

@@ -0,0 +1,139 @@
+# 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
+
+import os
+from collections.abc import Sequence
+
+from torch import Tensor
+
+from .common import Logger
+
+_has_tb = importlib.util.find_spec("tensorboard") is not None
+_has_omgaconf = importlib.util.find_spec("omegaconf") is not None
+
+
+class TensorboardLogger(Logger):
+    """Wrapper for the Tensoarboard logger.
+
+    Args:
+        exp_name (str): The name of the experiment.
+        log_dir (str): the tensorboard log_dir. Defaults to ``td_logs``.
+
+    """
+
+    def __init__(self, exp_name: str, log_dir: str = "tb_logs") -> None:
+        super().__init__(exp_name=exp_name, log_dir=log_dir)
+        # re-write log_dir
+        self.log_dir = self.experiment.log_dir
+
+        self._has_imported_moviepy = False
+
+    def _create_experiment(self) -> SummaryWriter:  # noqa
+        """Creates a tensorboard experiment.
+
+        Args:
+            exp_name (str): The name of the experiment.
+
+        Returns:
+            SummaryWriter: The tensorboard experiment.
+
+        """
+        if not _has_tb:
+            raise ImportError("torch.utils.tensorboard could not be imported")
+
+        from torch.utils.tensorboard import SummaryWriter
+
+        log_dir = str(os.path.join(self.log_dir, self.exp_name))
+        return SummaryWriter(log_dir=log_dir)
+
+    def log_scalar(self, name: str, value: float, step: int | None = None) -> None:
+        """Logs a scalar value to the tensorboard.
+
+        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.add_scalar(name, value, global_step=step)
+
+    def log_video(
+        self, name: str, video: Tensor, step: int | None = None, **kwargs
+    ) -> None:
+        """Log videos inputs to the tensorboard.
+
+        Args:
+            name (str): The name of the video.
+            video (Tensor): The video to be logged.
+            step (int, optional): The step at which the video is logged. Defaults to None.
+
+        """
+        # check for correct format of the video tensor ((N), T, C, H, W)
+        # check that the color channel (C) is either 1 or 3
+        if video.dim() != 5 or video.size(dim=2) not in {1, 3}:
+            raise Exception(
+                "Wrong format of the video tensor. Should be ((N), T, C, H, W)"
+            )
+        if not self._has_imported_moviepy:
+            try:
+                import moviepy  # noqa
+
+                self._has_imported_moviepy = True
+            except ImportError:
+                raise Exception(
+                    "moviepy not found, videos cannot be logged with TensorboardLogger"
+                )
+        self.experiment.add_video(
+            tag=name,
+            vid_tensor=video,
+            global_step=step,
+            **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_omgaconf:
+            if not _has_omgaconf:
+                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.add_hparams(cfg, metric_dict={})
+
+    def __repr__(self) -> str:
+        return f"TensorboardLogger(experiment={self.experiment.__repr__()})"
+
+    def log_histogram(self, name: str, data: Sequence, **kwargs):
+        """Add histogram to summary.
+
+        Args:
+            name (str): Data identifier
+            data (torch.Tensor, numpy.ndarray, or string/blobname): Values to build histogram
+
+        Keyword Args:
+            step (int): Global step value to record
+            bins (str): One of {‘tensorflow’,’auto’, ‘fd’, …}. This determines how the bins are made. You can find other options in: https://docs.scipy.org/doc/numpy/reference/generated/numpy.histogram.html
+            walltime (:obj:`float`): Optional override default walltime (time.time()) seconds after epoch of event
+
+        """
+        global_step = kwargs.pop("step", None)
+        bins = kwargs.pop("bins")
+        walltime = kwargs.pop("walltime", None)
+        if len(kwargs):
+            raise TypeError(f"Unrecognised arguments {kwargs}.")
+        self.experiment.add_histogram(
+            tag=name, values=data, global_step=global_step, bins=bins, walltime=walltime
+        )

torchrl/record/loggers/trackio.py

@@ -0,0 +1,163 @@
+# 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
+
+import numpy as np
+
+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.numpy().astype(np.uint8), 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):
+        """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.
+
+        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]])
+        self.experiment.log({name: table}, step=step)

torchrl/record/loggers/utils.py

@@ -0,0 +1,78 @@
+# 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 os
+import pathlib
+import uuid
+from datetime import datetime
+
+from torchrl.record.loggers.common import Logger
+
+
+def generate_exp_name(model_name: str, experiment_name: str) -> str:
+    """Generates an ID (str) for the described experiment using UUID and current date."""
+    exp_name = "_".join(
+        (
+            model_name,
+            experiment_name,
+            str(uuid.uuid4())[:8],
+            datetime.now().strftime("%y_%m_%d-%H_%M_%S"),
+        )
+    )
+    return exp_name
+
+
+def get_logger(
+    logger_type: str, logger_name: str, experiment_name: str, **kwargs
+) -> Logger:
+    """Get a logger instance of the provided `logger_type`.
+
+    Args:
+        logger_type (str): One of tensorboard / csv / wandb / mlflow.
+            If empty, ``None`` is returned.
+        logger_name (str): Name to be used as a log_dir
+        experiment_name (str): Name of the experiment
+        kwargs (dict[str]): might contain either `wandb_kwargs`, `mlflow_kwargs` or `trackio_kwargs`
+    """
+    if logger_type == "tensorboard":
+        from torchrl.record.loggers.tensorboard import TensorboardLogger
+
+        logger = TensorboardLogger(log_dir=logger_name, exp_name=experiment_name)
+    elif logger_type == "csv":
+        from torchrl.record.loggers.csv import CSVLogger
+
+        logger = CSVLogger(
+            log_dir=logger_name, exp_name=experiment_name, video_format="mp4"
+        )
+    elif logger_type == "wandb":
+        from torchrl.record.loggers.wandb import WandbLogger
+
+        wandb_kwargs = kwargs.get("wandb_kwargs", {})
+        logger = WandbLogger(
+            log_dir=logger_name, exp_name=experiment_name, **wandb_kwargs
+        )
+    elif logger_type == "mlflow":
+        from torchrl.record.loggers.mlflow import MLFlowLogger
+
+        mlflow_kwargs = kwargs.get("mlflow_kwargs", {})
+        logger = MLFlowLogger(
+            tracking_uri=pathlib.Path(os.path.abspath(logger_name)).as_uri(),
+            exp_name=experiment_name,
+            **mlflow_kwargs,
+        )
+    elif logger_type == "trackio":
+        from torchrl.record.loggers.trackio import TrackioLogger
+
+        trackio_kwargs = kwargs.get("trackio_kwargs", {})
+        project = trackio_kwargs.pop("project", "torchrl")
+        logger = TrackioLogger(
+            project=project, exp_name=experiment_name, **trackio_kwargs
+        )
+    elif logger_type in ("", None):
+        return None
+    else:
+        raise NotImplementedError(f"Unsupported logger_type: '{logger_type}'")
+    return logger