Logging¶

def reset(self):
    """Clear all buffered statistics."""
    self.stats = {}

def add(self, new: Any):
    """Ingest a new object into the logger.

    Concrete subclasses define the supported input types.

    Args:
        new: Object to ingest.

    Raises:
        NotImplementedError: Always, in the base class.
    """
    raise NotImplementedError

def log(self, step: int):
    """Emit logs for a given global step.

    Args:
        step: Global step index used for TensorBoard summary steps.

    Raises:
        NotImplementedError: Always, in the base class.
    """
    raise NotImplementedError

def reset(self):
    """Clear all buffered scalar and distribution values."""
    super().reset()
    self.dists = {}
    self.stats_to_log = {}
    self.dists_to_log = {}

def add(self, new: Mapping[str, Union["Stats", "Dist", float, int, np.floating]]):
    """Add a batch of metrics to the logger.

    Args:
        new: Mapping from metric name to a supported metric object.

    Raises:
        NotImplementedError: If a value type is unsupported.
    """

    for key, val in new.items():
        if isinstance(val, Stats):
            met = val.get()
            for k, v in met.items():
                if k in self.stats:
                    self.stats[k].append(float(v))
                else:
                    maxlen = self.stats_window_size if k not in self.stats_window_overrides else self.stats_window_overrides[k]
                    self.stats[k] = deque([float(v)], maxlen=maxlen)
        elif isinstance(val, Dist):
            # Store a snapshot of the reservoir for later histogram logging.
            self.dists[key] = val.get()
        elif isinstance(val, (float, int, np.floating)):
            if key in self.stats:
                self.stats[key].append(float(val))
            else:
                maxlen = self.stats_window_size if key not in self.stats_window_overrides else self.stats_window_overrides[key]
                self.stats[key] = deque([float(val)], maxlen=maxlen)
        else:
            raise NotImplementedError(
                "StatsLogger.add() only supports types: Stats, Dist, and numeric scalars"
            )

def log(self, step: int):
    """Aggregate buffered values and emit logs.

    Args:
        step: Global step index used for TensorBoard summary steps.
    """
    self._create_logs()
    if self.tensorboard:
        self._log_to_tensorboard(step)
    if self.stdout:
        self._log_to_stdout(step)

def _create_logs(self):
    """Create :attr:`stats_to_log` and :attr:`dists_to_log` from buffers."""
    self._create_stats_to_log()
    self._create_dists_to_log()

def _create_stats_to_log(self):
    """Compute smoothed scalar values to emit."""
    self.stats_to_log = {}
    for key, val in self.stats.items():
        if len(val) > 0:
            if key.endswith("max") or key.endswith("mag"):
                agg = np.max(val)
            elif key.endswith("min"):
                agg = np.min(val)
            else:
                agg = np.mean(val)
        else:
            agg = last
        self.stats_to_log[key] = float(agg)

def _create_dists_to_log(self):
    """Collect distributions to emit as histograms."""
    self.dists_to_log = {}
    for key, val in self.dists.items():
        if len(val) > 0:
            self.dists_to_log[key] = val

def _log_to_tensorboard(self, step: int):
    """Write scalars and histograms to TensorBoard.

    Args:
        step: Global step index used for TensorBoard summary steps.
    """
    assert self.summary_writer is not None, (
        "You're trying to log to tensorboard without a summary writer setup!"
    )
    with self.summary_writer.as_default():
        for key, value in self.stats_to_log.items():
            tf.summary.scalar(self.prefix + key, value, step=step)
        for key, values in self.dists_to_log.items():
            tf.summary.histogram(self.prefix + key, data=values, step=step)

def _get_wandb_payload(self) -> Dict[str, Any]:
    """Return the current W&B payload dict without logging it.

    Used by :class:`TrainLogger` to batch all sub-logger payloads into a
    single :func:`wandb.log` call.

    Returns:
        Mapping from metric name to scalar or :class:`wandb.Histogram`.
    """
    import wandb as _wandb
    payload: Dict[str, Any] = {self.prefix + k: v for k, v in self.stats_to_log.items()}
    for k, v in self.dists_to_log.items():
        payload[self.prefix + k] = _wandb.Histogram(v)
    return payload

def _log_to_stdout(self, step: int):
    """Print the current scalar log table to stdout.

    Args:
        step: Global step index (unused; included for API symmetry).
    """
    stats_to_print = {key: "{0:.4g}".format(val) for key, val in self.stats_to_log.items()}
    if not stats_to_print:
        return

    max_key_len = max([len(key) for key in stats_to_print] + [max(0, len(self.prefix) - 2)])
    max_val_len = max([len(val) for val in stats_to_print.values()])

    stdout = ""
    max_len = 1 + 4 + max_key_len + 2 + 1 + 2 + max_val_len + 2 + 1
    stdout += ("-" * max_len + "\n")
    stdout += (
        "|  "
        + self.prefix
        + " " * (2 + max_key_len - len(self.prefix) + 2)
        + "|"
        + " " * (2 + max_val_len + 2)
        + "|\n"
    )
    for key, val in stats_to_print.items():
        stdout += (
            "|    "
            + key
            + " " * (max_key_len - len(key) + 2)
            + "|  "
            + val
            + " " * (max_val_len - len(val) + 2)
            + "|\n"
        )
    stdout += ("-" * max_len + "\n")

    if self.tqdm:
        tqdm.write(stdout)
    else:
        print(stdout)

def add(self, info: Mapping[str, Any], verbose: int = 0):
    """Ingest an ``info`` dict and extract episodic scalars.

    Args:
        info: Rollout ``info`` mapping (typically from environment step).
            If present, the logger reads:
            ``info["constraint"]["episode"]`` and/or
            ``info["metrics"]["episode"]``.
        verbose: Reserved for compatibility; currently unused.
    """
    if self.start_time is None:
        self.start_time = time.time()

    constraint = info.get("constraint", {})
    if isinstance(constraint, Mapping) and "episode" in constraint:
        ep_metrics = constraint["episode"]
        if isinstance(ep_metrics, Mapping):
            self._add_scalars(ep_metrics)

    metrics = info.get("metrics", {})
    if isinstance(metrics, Mapping) and "episode" in metrics:
        ep_metrics = metrics["episode"]
        if isinstance(ep_metrics, Mapping):
            self._add_scalars(ep_metrics)

def log(self, step: int):
    """Aggregate buffered episode metrics and emit logs.

    Args:
        step: Global step index used for TensorBoard summary steps.
    """
    self._create_logs()
    if self.tensorboard:
        self._log_to_tensorboard(step)
    if self.stdout:
        self._log_to_stdout(step)

def _create_logs(self):
    """Create :attr:`stats_to_log` from episode buffers."""
    self._create_stats_to_log()

def _add_scalars(self, scalars: Mapping[str, Union[float, int, np.floating]]):
    """Append scalar episode metrics into rolling windows.

    Args:
        scalars: Mapping from metric names to numeric values.
    """
    for k, v in scalars.items():
        if k in self.stats:
            self.stats[k].append(float(v))
        else:
            # +1 because we keep the most recent value as "current episode".
            maxlen = self.stats_window_size if k not in self.stats_window_overrides else self.stats_window_overrides[k]
            self.stats[k] = deque([float(v)], maxlen=maxlen + 1)

def _create_stats_to_log(self):
    """Compute per-metric mean over *completed* episodes."""
    self.stats_to_log = {}
    for key, val in self.stats.items():
        if len(val) > 1:
            # Temporarily drop last (current) value from the mean.
            last = val.pop()
            self.stats_to_log[key] = float(np.mean(val)) if len(val) > 0 else float(last)
            val.append(last)

def _log_to_tensorboard(self, step: int):
    """Write episode scalar summaries to TensorBoard.

    Args:
        step: Global step index used for TensorBoard summary steps.
    """
    assert self.summary_writer is not None, (
        "You're trying to log to tensorboard without a summary writer setup!"
    )
    with self.summary_writer.as_default():
        for key, value in self.stats_to_log.items():
            tf.summary.scalar(self.prefix + key, value, step=step)

def _get_wandb_payload(self) -> Dict[str, Any]:
    """Return the current W&B payload dict without logging it.

    Used by :class:`TrainLogger` to batch all sub-logger payloads into a
    single :func:`wandb.log` call.

    Returns:
        Mapping from metric name to scalar value.
    """
    return {self.prefix + k: v for k, v in self.stats_to_log.items()}

def _log_to_stdout(self, step: int):
    """Print episode summaries and runtime diagnostics to stdout.

    Args:
        step: Global step index used for fps and total timestep display.
    """
    stats_to_print: Dict[str, str] = {
        key: "{0:.4g}".format(val) for key, val in self.stats_to_log.items()
    }
    if self.start_time is not None:
        current_time = time.time()
        elapsed = current_time - self.start_time
        if elapsed > 0:
            stats_to_print["fps"] = "{0:.4g}".format(step / elapsed)
        stats_to_print["time_elapsed"] = "{0:.4g}".format(elapsed)
    stats_to_print["total_timesteps"] = str(step)

    if not stats_to_print:
        return

    max_key_len = max([len(key) for key in stats_to_print] + [max(0, len(self.prefix) - 2)])
    max_val_len = max([len(val) for val in stats_to_print.values()])

    stdout = ""
    max_len = 1 + 4 + max_key_len + 2 + 1 + 2 + max_val_len + 2 + 1
    stdout += ("-" * max_len + "\n")
    stdout += (
        "|  "
        + self.prefix
        + " " * (2 + max_key_len - len(self.prefix) + 2)
        + "|"
        + " " * (2 + max_val_len + 2)
        + "|\n"
    )
    for key, val in stats_to_print.items():
        stdout += (
            "|    "
            + key
            + " " * (max_key_len - len(key) + 2)
            + "|  "
            + val
            + " " * (max_val_len - len(val) + 2)
            + "|\n"
        )
    stdout += ("-" * max_len + "\n")

    if self.tqdm:
        tqdm.write(stdout)
    else:
        print(stdout)

def add(self, key: str, obj: Any):
    """Add an object to a named sub-logger.

    Args:
        key: The sub-logger key as provided in ``loggers`` during
            construction.
        obj: The object to forward to ``self.loggers[key].add(...)``.

    Raises:
        KeyError: If ``key`` is not a configured sub-logger.
    """
    if self.start_time is None:
        self.start_time = time.time()
    self.loggers[key].add(obj)

def log(self, step: int):
    """Emit TensorBoard logs (per sub-logger) and a combined stdout table.

    Args:
        step: Global step index used for TensorBoard summary steps.
    """
    for key in self.loggers.keys():
        # Rely on sub-logger internal API (common to StatsLogger/RolloutLogger).
        self.loggers[key]._create_logs()  # type: ignore[attr-defined]
        if self.tensorboard:
            self.loggers[key]._log_to_tensorboard(step)  # type: ignore[attr-defined]

    if self.wandb:
        # Batch all sub-logger payloads into a single wandb.log() call to
        # keep all metrics on the same step without step-alignment warnings.
        import wandb as _wandb
        payload: Dict[str, Any] = {}
        for logger in self.loggers.values():
            payload.update(logger._get_wandb_payload())  # type: ignore[attr-defined]
        if payload:
            _wandb.log(payload, step=step)

    if self.stdout:
        self._log_to_stdout(step)

def _log_to_stdout(self, step: int) -> None:
    """Print a combined stdout table for all sub-loggers.

    Args:
        step: Global step index used for runtime diagnostics.
    """
    stats_to_print: Dict[str, Dict[str, str]] = {}
    stats_to_print["run"] = {}

    if self.start_time is not None:
        current_time = time.time()
        elapsed = current_time - self.start_time
        if elapsed > 0:
            stats_to_print["run"]["fps"] = "{0:.4g}".format(step / elapsed)
        stats_to_print["run"]["time_elapsed"] = "{0:.4g}".format(elapsed)
    stats_to_print["run"]["total_timesteps"] = str(step)

    for key, logger in self.loggers.items():
        # Sub-loggers populate stats_to_log in _create_logs().
        stats = getattr(logger, "stats_to_log", {})
        stats_to_print[key] = {k: "{0:.4g}".format(v) for k, v in stats.items()}

    max_key_len = 0
    max_val_len = 0
    for group in stats_to_print.values():
        if not group:
            continue
        max_key_len = max(max_key_len, max([len(k) for k in group.keys()] + [max(0, len(self.prefix) - 2)]))
        max_val_len = max(max_val_len, max([len(v) for v in group.values()]))

    stdout = ""
    max_len = 1 + 4 + max_key_len + 2 + 1 + 2 + max_val_len + 2 + 1
    stdout += ("-" * max_len + "\n")
    if self.prefix:
        stdout += (
            "|  "
            + self.prefix
            + " " * (2 + max_key_len - len(self.prefix) + 2)
            + "|"
            + " " * (2 + max_val_len + 2)
            + "|\n"
        )
        stdout += ("-" * max_len + "\n")

    for group_key, group in stats_to_print.items():
        if not group:
            continue
        group_prefix = group_key + "/"
        stdout += (
            "|  "
            + group_prefix
            + " " * (2 + max_key_len - len(group_prefix) + 2)
            + "|"
            + " " * (2 + max_val_len + 2)
            + "|\n"
        )
        for k, v in group.items():
            stdout += (
                "|    "
                + k
                + " " * (max_key_len - len(k) + 2)
                + "|  "
                + v
                + " " * (max_val_len - len(v) + 2)
                + "|\n"
            )
        stdout += ("-" * max_len + "\n")

    if self.tqdm:
        tqdm.write(stdout)
    else:
        print(stdout)