diff --git a/CONTRIBUTORS.md b/CONTRIBUTORS.md index bafc69ed..13858e25 100644 --- a/CONTRIBUTORS.md +++ b/CONTRIBUTORS.md @@ -33,6 +33,7 @@ Please keep the lists sorted alphabetically. * Eric Vollenweider * Fabian Jenelten * Lorenzo Terenzi +* Maciej Aleksandrowicz * Marko Bjelonic * Matthijs van der Boon * Özhan Özen diff --git a/README.md b/README.md index c8bae884..abfbcbee 100644 --- a/README.md +++ b/README.md @@ -44,6 +44,7 @@ The package supports the following logging frameworks which can be configured th * Tensorboard: https://www.tensorflow.org/tensorboard/ * Weights & Biases: https://wandb.ai/site * Neptune: https://docs.neptune.ai/ +* ClearML: https://clear.ml/docs/latest/docs/ For a demo configuration of PPO, please check the [example_config.yaml](config/example_config.yaml) file. diff --git a/config/example_config.yaml b/config/example_config.yaml index 00a9483f..e1a211e4 100644 --- a/config/example_config.yaml +++ b/config/example_config.yaml @@ -11,9 +11,10 @@ runner: experiment_name: walking_experiment run_name: "" # Logging writer - logger: tensorboard # tensorboard, neptune, wandb + logger: tensorboard # tensorboard, neptune, wandb, clearml neptune_project: legged_gym wandb_project: legged_gym + clearml_project: legged_gym # Policy policy: diff --git a/rsl_rl/utils/clearml_utils.py b/rsl_rl/utils/clearml_utils.py new file mode 100644 index 00000000..8686f7cc --- /dev/null +++ b/rsl_rl/utils/clearml_utils.py @@ -0,0 +1,77 @@ +# 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 + +try: + from clearml import Task +except ModuleNotFoundError: + raise ModuleNotFoundError("clearml package is required to log to ClearML.") from None + + +class ClearmlSummaryWriter(SummaryWriter): + """Summary writer for ClearML.""" + + def __init__(self, log_dir: str, flush_secs: int, cfg: dict) -> None: + super().__init__(log_dir, flush_secs) + + # Get the run name + run_name = os.path.split(log_dir)[-1] + + # Get ClearML task + try: + project_name = cfg["clearml_project"] + except KeyError: + raise KeyError("Please specify clearml_project in the runner config, e.g. `legged_gym`.") from None + + # Initialize ClearML Task + self.task = Task.init( + project_name=project_name, task_name=run_name, auto_connect_frameworks={"tensorboard": False} + ) + + def store_config(self, env_cfg: dict | object, train_cfg: dict) -> None: + runner_cfg = dict(train_cfg) + runner_cfg.pop("policy", None) + runner_cfg.pop("algorithm", None) + + if isinstance(env_cfg, dict): + env_dict = env_cfg + else: + env_dict = env_cfg.to_dict() if hasattr(env_cfg, "to_dict") else asdict(env_cfg) + + self.task.connect(runner_cfg, name="runner_cfg") + self.task.connect(train_cfg.get("policy", {}), name="policy_cfg") + self.task.connect(train_cfg.get("algorithm", {}), name="alg_cfg") + self.task.connect(env_dict, name="env_cfg") + + def add_scalar( + self, + tag: str, + scalar_value: float, + global_step: int | None = None, + walltime: float | None = None, + new_style: bool = False, + ) -> None: + super().add_scalar( + tag, + scalar_value, + global_step=global_step, + walltime=walltime, + new_style=new_style, + ) + self.task.get_logger().report_scalar(tag, "series", scalar_value, iteration=global_step) + + def stop(self) -> None: + self.task.close() + + def save_model(self, model_path: str, it: int) -> None: + self.task.upload_artifact(name=f"model_{it}", artifact_object=model_path) + + def save_file(self, path: str) -> None: + self.task.upload_artifact(name=os.path.basename(path), artifact_object=path) diff --git a/rsl_rl/utils/logger.py b/rsl_rl/utils/logger.py index 22838dfa..713a9676 100644 --- a/rsl_rl/utils/logger.py +++ b/rsl_rl/utils/logger.py @@ -19,6 +19,8 @@ class Logger: """Logger to save the learning metrics to different logging services.""" + LOGGER_TYPES = ("wandb", "neptune", "clearml") + def __init__( self, log_dir: str | None, @@ -64,7 +66,7 @@ def __init__( self._store_code_state() # Log configuration - if self.writer and not self.disable_logs and self.logger_type in ["wandb", "neptune"]: + if self.writer and not self.disable_logs and self.logger_type in Logger.LOGGER_TYPES: self.writer.store_config(env_cfg, self.cfg) def process_env_step( @@ -230,11 +232,11 @@ def log( def save_model(self, path: str, it: int) -> None: """Save the model to external logging services if specified.""" - if self.writer and not self.disable_logs and self.logger_type in ["neptune", "wandb"]: + if self.writer and not self.disable_logs and self.logger_type in Logger.LOGGER_TYPES: self.writer.save_model(path, it) def _prepare_logging_writer(self) -> None: - """Prepare the logging writer, which can be either Tensorboard, W&B or Neptune.""" + """Prepare the logging writer, which can be either Tensorboard, W&B, Neptune or ClearML.""" if self.log_dir is not None and not self.disable_logs: self.logger_type = self.cfg.get("logger", "tensorboard") self.logger_type = self.logger_type.lower() @@ -247,12 +249,16 @@ def _prepare_logging_writer(self) -> None: from rsl_rl.utils.wandb_utils import WandbSummaryWriter self.writer = WandbSummaryWriter(log_dir=self.log_dir, flush_secs=10, cfg=self.cfg) + elif self.logger_type == "clearml": + from rsl_rl.utils.clearml_utils import ClearmlSummaryWriter + + self.writer = ClearmlSummaryWriter(log_dir=self.log_dir, flush_secs=10, cfg=self.cfg) elif self.logger_type == "tensorboard": from torch.utils.tensorboard import SummaryWriter self.writer = SummaryWriter(log_dir=self.log_dir, flush_secs=10) else: - raise ValueError("Logger type not found. Please choose 'wandb', 'neptune', or 'tensorboard'.") + raise ValueError(f"Logger type not found. Please choose one of: {', '.join(Logger.LOGGER_TYPES)}.") else: self.writer = None @@ -285,6 +291,6 @@ def _store_code_state(self) -> None: file_paths.append(diff_file_name) # Upload diff files to external logging services - if self.writer and self.logger_type in ["wandb", "neptune"] and file_paths: + if self.writer and self.logger_type in Logger.LOGGER_TYPES and file_paths: for path in file_paths: self.writer.save_file(path)