Source code for rsl_rl.utils.neptune_log_writer

# Copyright (c) 2021-2026, ETH Zurich and NVIDIA CORPORATION
# All rights reserved.
#
# SPDX-License-Identifier: BSD-3-Clause


from __future__ import annotations

import os
from dataclasses import asdict
from torch.utils.tensorboard import SummaryWriter

from rsl_rl.utils.log_writer import LogWriter

try:
    import neptune  # type: ignore
except ModuleNotFoundError:
    neptune = None




class NeptuneLogWriter(SummaryWriter, LogWriter):
    """Summary writer for Neptune."""



    def __init__(self, log_dir: str, project_name: str) -> None:
        """Initialize a Neptune run for logging."""
        if neptune is None:
            raise ModuleNotFoundError("neptune-client is required to log to Neptune.")
        super().__init__(log_dir, flush_secs=10)

        # Get the run name
        run_name = os.path.split(log_dir)[-1]

        try:
            token = os.environ["NEPTUNE_API_TOKEN"]
        except KeyError:
            raise KeyError(
                "Neptune api token not found. Please run or add to ~/.bashrc: export NEPTUNE_API_TOKEN=YOUR_API_TOKEN"
            ) from None
        try:
            entity = os.environ["NEPTUNE_USERNAME"]
        except KeyError:
            raise KeyError(
                "Neptune username not found. Please run or add to ~/.bashrc: export NEPTUNE_USERNAME=YOUR_USERNAME"
            ) from None

        # Initialize neptune
        neptune_project = entity + "/" + project_name
        self.run = neptune.init_run(project=neptune_project, api_token=token)
        self.run["log_dir"].log(run_name)

        # Name mapping for incompatible characters
        self.name_map = {
            "Train/mean_reward/time": "Train/mean_reward_time",
            "Train/mean_episode_length/time": "Train/mean_episode_length_time",
        }




    def add_scalar(
        self,
        tag: str,
        scalar_value: float,
        global_step: int | None = None,
        walltime: float | None = None,
        new_style: bool = False,
    ) -> None:
        """Log a scalar to both TensorBoard and Neptune."""
        super().add_scalar(tag, scalar_value, global_step=global_step, walltime=walltime, new_style=new_style)
        self.run[self._map_path(tag)].log(scalar_value, step=global_step)




    def store_config(self, env_cfg: dict | object, train_cfg: dict) -> None:
        """Upload environment and training configuration to Neptune."""
        self.run["train_cfg"] = train_cfg
        try:
            self.run["env_cfg"] = env_cfg.to_dict()  # type: ignore
        except Exception:
            self.run["env_cfg"] = asdict(env_cfg)  # type: ignore




    def save_model(self, model_path: str, it: int) -> None:
        """Upload a model checkpoint artifact to Neptune."""
        self.run["model/saved_model_" + str(it)].upload(model_path)




    def save_file(self, path: str) -> None:
        """Upload an arbitrary file artifact to Neptune."""
        name = path.rsplit("/", 1)[-1].split(".")[0]
        self.run["git_diff/" + name].upload(path)




    def stop(self) -> None:
        """Finish the active Neptune run."""
        self.run.stop()


    def _map_path(self, path: str) -> str:
        """Map metric names to Neptune-compatible keys."""
        if path in self.name_map:
            return self.name_map[path]
        else:
            return path