# 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.metadata as metadata
import importlib.util
import json
import os
import platform
import sys
from collections.abc import Mapping, Sequence
from typing import Any
from tensordict import TensorDictBase
from torch import Tensor
from .common import _make_metrics_safe, Logger
_has_wandb = importlib.util.find_spec("wandb") is not None
_has_omegaconf = importlib.util.find_spec("omegaconf") is not None
_has_moviepy = importlib.util.find_spec("moviepy") is not None
def _collect_env_metadata() -> dict[str, Any]:
packages = {}
editable_sources = {}
for distribution in metadata.distributions():
name = distribution.metadata.get("Name")
if name is None:
continue
packages[name] = distribution.version
direct_url = distribution.read_text("direct_url.json")
if direct_url is not None:
direct_url = json.loads(direct_url)
dir_info = direct_url.get("dir_info")
if dir_info is not None and dir_info.get("editable"):
editable_sources[name] = direct_url.get("url")
return {
"python": {
"version": sys.version,
"executable": sys.executable,
"prefix": sys.prefix,
"base_prefix": sys.base_prefix,
"platform": platform.platform(),
},
"packages": dict(sorted(packages.items())),
"editable_sources": dict(sorted(editable_sources.items())),
}
[docs]
class WandbLogger(Logger):
"""Wrapper for the wandb logger.
See also :class:`~torchrl.trainers.algorithms.configs.WandbLoggerConfig`.
The keyword arguments are mainly based on the :func:`wandb.init` kwargs.
See the doc `here <https://docs.wandb.ai/ref/python/init>`__.
Args:
exp_name (str): The name of the experiment.
offline (bool, optional): if ``True``, the logs will be stored locally
only. Defaults to ``False``.
save_dir (path, optional): the directory where to save data. Exclusive with
``log_dir``.
log_dir (path, optional): the directory where to save data. Exclusive with
``save_dir``.
id (str, optional): A unique ID for this run, used for resuming.
It must be unique in the project, and if you delete a run you can't reuse the ID.
project (str, optional): The name of the project where you're sending
the new run. If the project is not specified, the run is put in
an ``"Uncategorized"`` project.
base_url (str, optional): The W&B server URL used for authentication and
data synchronization. Construct the logger before other code imports
or calls W&B, because the W&B client reads ``WANDB_BASE_URL`` when it
is initialized.
log_env_packages (bool, optional): if ``True``, logs the Python runtime,
installed package versions, and editable source locations under
``wandb.config["env"]``. Defaults to ``True``.
Keyword Args:
fps (int, optional): Number of frames per second when recording videos. Defaults to ``30``.
**kwargs: Extra keyword arguments for ``wandb.init``. See relevant page for
more info.
"""
@classmethod
def __new__(cls, *args, **kwargs):
return super().__new__(cls)
def __init__(
self,
exp_name: str,
offline: bool = False,
save_dir: str | os.PathLike[str] | None = None,
id: str | None = None,
project: str | None = None,
base_url: str | None = None,
*,
video_fps: int = 32,
log_env_packages: bool = True,
**kwargs,
) -> None:
if not _has_wandb:
raise ImportError("wandb could not be imported")
log_dir: str | os.PathLike[str] | None = kwargs.pop("log_dir", None)
self.offline = offline
save_dir = save_dir or None
log_dir = log_dir or None
if save_dir is not None and log_dir is not None:
raise ValueError(
"log_dir and save_dir point to the same value in "
"WandbLogger. Both cannot be specified."
)
save_dir = save_dir if save_dir is not None else log_dir
if save_dir is not None:
save_dir = os.path.expanduser(os.fspath(save_dir))
os.makedirs(save_dir, exist_ok=True)
self.save_dir = save_dir
self.id = id
self.project = project
self.base_url = base_url
self.video_fps = video_fps
self.log_env_packages = log_env_packages
self._step_registry: dict[str, int] = {}
self._defined_step_metrics: set[str] = set()
self._defined_metrics: set[str] = set()
self._wandb_kwargs = {
"name": exp_name,
"dir": save_dir,
"id": id,
"project": project,
"resume": "allow",
**kwargs,
}
super().__init__(exp_name=exp_name, log_dir=save_dir)
if self.log_env_packages:
try:
self.experiment.config.update(
{"env": _collect_env_metadata()}, allow_val_change=True
)
except TypeError:
self.experiment.config.update({"env": _collect_env_metadata()})
if self.offline:
os.environ["WANDB_MODE"] = "dryrun"
def _checkpoint_state(self) -> dict[str, Any]:
return {
"id": getattr(self.experiment, "id", self.id),
"step_registry": dict(self._step_registry),
"defined_step_metrics": sorted(self._defined_step_metrics),
"defined_metrics": sorted(self._defined_metrics),
}
def _load_checkpoint_state(self, state_dict: Mapping[str, Any]) -> None:
saved_id = state_dict.get("id")
live_id = getattr(self.experiment, "id", self.id)
if saved_id and live_id and saved_id != live_id:
raise RuntimeError(
f"This WandbLogger logs to run {live_id!r} but the checkpoint was "
f"written by run {saved_id!r}. Construct the logger for the saved "
"run before loading, e.g. WandbLogger(..., id=saved_id, "
"resume='must') or get_logger(..., state_dict=saved_state)."
)
if saved_id:
self.id = saved_id
self._step_registry.clear()
self._step_registry.update(state_dict.get("step_registry", {}))
self._defined_step_metrics = set(state_dict.get("defined_step_metrics", ()))
self._defined_metrics = set(state_dict.get("defined_metrics", ()))
def _create_experiment(self):
"""Creates a wandb experiment.
Args:
exp_name (str): The name of the experiment.
Returns:
A wandb.Experiment object.
"""
if not _has_wandb:
raise ImportError("Wandb is not installed")
import wandb
if self.offline:
os.environ["WANDB_MODE"] = "dryrun"
if self.base_url is not None:
settings = self._wandb_kwargs.get("settings")
if settings is None:
settings = wandb.Settings(base_url=self.base_url)
elif isinstance(settings, Mapping):
settings = wandb.Settings(**{**settings, "base_url": self.base_url})
else:
settings.base_url = self.base_url
self._wandb_kwargs["settings"] = settings
if not self.offline:
wandb.login(host=self.base_url, key=getattr(settings, "api_key", None))
return wandb.init(**self._wandb_kwargs)
def log_scalar(
self,
name: str,
value: float,
step: int | None = None,
commit: bool = True,
*,
override_global_step: bool = False,
) -> None:
"""Logs a scalar value to wandb.
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.
commit (bool, optional): If ``True``, data for the current step is
assumed to be final (and no further data for this step should
be logged). Set to ``False`` to batch multiple calls into the
same W&B history row. Defaults to ``True``.
override_global_step (bool, optional): If ``True``, bypasses
per-group step injection and forwards ``step`` to wandb's
global ``step`` argument. Defaults to ``False``.
"""
self._log_payload(
{name: value},
step=step,
commit=commit,
override_global_step=override_global_step,
)
def log_video(self, name: str, video: Tensor, **kwargs) -> None:
"""Log videos inputs to wandb.
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.
Raises:
ImportError: If moviepy is not installed (required by wandb for video encoding).
"""
if not _has_moviepy:
raise ImportError(
"Video logging with wandb requires moviepy. "
"Install with: pip install moviepy\n"
"Or install wandb with media support: pip install 'wandb[media]'"
)
import wandb
fps = kwargs.pop("fps", self.video_fps)
format = kwargs.pop("format", "mp4")
self._log_payload(
{name: wandb.Video(video, fps=fps, format=format)},
step=kwargs.pop("step", None),
override_global_step=kwargs.pop("override_global_step", False),
**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, allow_val_change=True)
def __repr__(self) -> str:
return f"WandbLogger(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, 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
"""
import wandb
num_bins = kwargs.pop("bins", None)
step = kwargs.pop("step", None)
self._log_payload(
{name: wandb.Histogram(data, num_bins=num_bins)},
step=step,
override_global_step=kwargs.pop("override_global_step", False),
**kwargs,
)
def log_str(
self,
name: str,
value: str,
step: int | None = None,
*,
override_global_step: bool = False,
) -> None:
"""Logs a string value to wandb 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.
override_global_step (bool, optional): If ``True``, bypasses
per-group step injection and forwards ``step`` to wandb's
global ``step`` argument. Defaults to ``False``.
"""
import wandb
# Create a table with a single row
table = wandb.Table(columns=["text"], data=[[value]])
self._log_payload(
{name: value if step is not None else table},
step=step,
override_global_step=override_global_step,
)
def log_metrics(
self,
metrics: dict[str, Any] | TensorDictBase,
step: int | None = None,
*,
keys_sep: str = "/",
override_global_step: bool = False,
) -> dict[str, Any]:
"""Log multiple scalar metrics at once to wandb.
This method efficiently handles tensor values by batching CUDA->CPU
transfers and performing a single synchronization, then logs all
metrics in a single wandb API call.
Args:
metrics: Dictionary or TensorDict mapping metric names to values.
Tensor values are automatically converted to Python scalars/lists.
For TensorDict inputs, nested keys are flattened using ``keys_sep``.
step: Optional step value for all metrics.
keys_sep: Separator used to flatten nested TensorDict keys into strings.
Defaults to "/". Only used for TensorDict inputs.
override_global_step: If ``True``, bypasses per-group step injection
and forwards ``step`` to wandb's global ``step`` argument.
Defaults to ``False``.
Returns:
The converted metrics dictionary (with tensors converted to Python types).
"""
safe_metrics = _make_metrics_safe(metrics, keys_sep=keys_sep)
self._log_payload(
safe_metrics, step=step, override_global_step=override_global_step
)
return safe_metrics
@staticmethod
def _is_step_key(name: str) -> bool:
return name == "step" or name.endswith("/step")
@staticmethod
def _step_key(name: str) -> str:
if WandbLogger._is_step_key(name):
return name
prefix, sep, _ = name.rpartition("/")
return f"{prefix}{sep}step" if sep else "step"
def _consume_step(self, step_key: str, step: int | None) -> int:
last_step = self._step_registry.get(step_key, -1)
if step is None:
step = last_step + 1
self._step_registry[step_key] = max(last_step, step)
return step
def _define_metric(self, name: str, *, step_metric: str | None = None) -> None:
if step_metric is None:
if name in self._defined_step_metrics:
return
self._defined_step_metrics.add(name)
self.experiment.define_metric(name)
return
if name in self._defined_metrics:
return
self._defined_metrics.add(name)
self.experiment.define_metric(name, step_metric=step_metric)
def _prepare_payload(
self, payload: dict[str, Any], step: int | None
) -> dict[str, Any]:
prepared = dict(payload)
for key, value in list(prepared.items()):
if self._is_step_key(key):
self._consume_step(key, value)
for key in list(prepared):
if self._is_step_key(key):
continue
step_key = self._step_key(key)
if step_key not in prepared:
prepared[step_key] = self._consume_step(step_key, step)
return prepared
def _register_metrics(self, payload: dict[str, Any]) -> None:
for key in payload:
if self._is_step_key(key):
self._define_metric(key)
for key in payload:
if self._is_step_key(key):
continue
self._define_metric(key, step_metric=self._step_key(key))
def _log_payload(
self,
payload: dict[str, Any],
*,
step: int | None = None,
override_global_step: bool = False,
**kwargs,
) -> None:
if override_global_step:
self.experiment.log(payload, step=step, **kwargs)
return
payload = self._prepare_payload(payload, step)
self._register_metrics(payload)
self.experiment.log(payload, **kwargs)
