[Refactor] Move loggers to torchrl.record (#854)

Author: vmoens

docs/source/reference/trainers.rst

@@ -193,7 +193,7 @@ Utils
 Loggers
 -------
 
-.. currentmodule:: torchrl.trainers.loggers
+.. currentmodule:: torchrl.recorder.loggers
 
 .. autosummary::
     :toctree: generated/

examples/a2c/a2c.py

@@ -11,6 +11,7 @@
 from torchrl.envs.transforms import RewardScaling
 from torchrl.envs.utils import set_exploration_mode
 from torchrl.objectives.value import TDEstimate
+from torchrl.record.loggers import generate_exp_name, get_logger
 from torchrl.trainers.helpers.collectors import (
     make_collector_onpolicy,
     OnPolicyCollectorConfig,
@@ -27,7 +28,6 @@
 from torchrl.trainers.helpers.losses import A2CLossConfig, make_a2c_loss
 from torchrl.trainers.helpers.models import A2CModelConfig, make_a2c_model
 from torchrl.trainers.helpers.trainers import make_trainer, TrainerConfig
-from torchrl.trainers.loggers.utils import generate_exp_name, get_logger
 
 config_fields = [
     (config_field.name, config_field.type, config_field)

examples/ddpg/ddpg.py

@@ -13,6 +13,7 @@
 from torchrl.envs.utils import set_exploration_mode
 from torchrl.modules import OrnsteinUhlenbeckProcessWrapper
 from torchrl.record import VideoRecorder
+from torchrl.record.loggers import generate_exp_name, get_logger
 from torchrl.trainers.helpers.collectors import (
     make_collector_offpolicy,
     OffPolicyCollectorConfig,
@@ -30,7 +31,6 @@
 from torchrl.trainers.helpers.models import DDPGModelConfig, make_ddpg_actor
 from torchrl.trainers.helpers.replay_buffer import make_replay_buffer, ReplayArgsConfig
 from torchrl.trainers.helpers.trainers import make_trainer, TrainerConfig
-from torchrl.trainers.loggers.utils import generate_exp_name, get_logger
 
 config_fields = [
     (config_field.name, config_field.type, config_field)

examples/dqn/dqn.py

@@ -12,6 +12,7 @@
 from torchrl.envs.transforms import RewardScaling, TransformedEnv
 from torchrl.modules import EGreedyWrapper
 from torchrl.record import VideoRecorder
+from torchrl.record.loggers import generate_exp_name, get_logger
 from torchrl.trainers.helpers.collectors import (
     make_collector_offpolicy,
     OffPolicyCollectorConfig,
@@ -29,7 +30,6 @@
 from torchrl.trainers.helpers.models import DiscreteModelConfig, make_dqn_actor
 from torchrl.trainers.helpers.replay_buffer import make_replay_buffer, ReplayArgsConfig
 from torchrl.trainers.helpers.trainers import make_trainer, TrainerConfig
-from torchrl.trainers.loggers.utils import generate_exp_name, get_logger
 
 config_fields = [
     (config_field.name, config_field.type, config_field)

examples/dreamer/dreamer.py

@@ -27,6 +27,7 @@
     DreamerModelLoss,
     DreamerValueLoss,
 )
+from torchrl.record.loggers import generate_exp_name, get_logger
 from torchrl.trainers.helpers.collectors import (
     make_collector_offpolicy,
     OffPolicyCollectorConfig,
@@ -40,7 +41,6 @@
 from torchrl.trainers.helpers.models import DreamerConfig, make_dreamer
 from torchrl.trainers.helpers.replay_buffer import make_replay_buffer, ReplayArgsConfig
 from torchrl.trainers.helpers.trainers import TrainerConfig
-from torchrl.trainers.loggers.utils import generate_exp_name, get_logger
 from torchrl.trainers.trainers import Recorder, RewardNormalizer
 
 config_fields = [

examples/dreamer/dreamer_utils.py

@@ -24,8 +24,8 @@
     TransformedEnv,
 )
 from torchrl.envs.transforms.transforms import FlattenObservation, TensorDictPrimer
+from torchrl.record.loggers import Logger
 from torchrl.record.recorder import VideoRecorder
-from torchrl.trainers.loggers import Logger
 
 __all__ = [
     "transformed_env_constructor",

examples/ppo/ppo.py

@@ -13,6 +13,7 @@
 from torchrl.envs.utils import set_exploration_mode
 from torchrl.objectives.value import GAE
 from torchrl.record import VideoRecorder
+from torchrl.record.loggers import generate_exp_name, get_logger
 from torchrl.trainers.helpers.collectors import (
     make_collector_onpolicy,
     OnPolicyCollectorConfig,
@@ -29,7 +30,6 @@
 from torchrl.trainers.helpers.losses import make_ppo_loss, PPOLossConfig
 from torchrl.trainers.helpers.models import make_ppo_model, PPOModelConfig
 from torchrl.trainers.helpers.trainers import make_trainer, TrainerConfig
-from torchrl.trainers.loggers.utils import generate_exp_name, get_logger
 
 config_fields = [
     (config_field.name, config_field.type, config_field)

examples/redq/redq.py

@@ -15,6 +15,7 @@
 from torchrl.envs.utils import set_exploration_mode
 from torchrl.modules import OrnsteinUhlenbeckProcessWrapper
 from torchrl.record import VideoRecorder
+from torchrl.record.loggers import generate_exp_name, get_logger
 from torchrl.trainers.helpers.collectors import (
     make_collector_offpolicy,
     OffPolicyCollectorConfig,
@@ -32,7 +33,6 @@
 from torchrl.trainers.helpers.models import make_redq_model, REDQModelConfig
 from torchrl.trainers.helpers.replay_buffer import make_replay_buffer, ReplayArgsConfig
 from torchrl.trainers.helpers.trainers import make_trainer, TrainerConfig
-from torchrl.trainers.loggers.utils import generate_exp_name, get_logger
 
 config_fields = [
     (config_field.name, config_field.type, config_field)

examples/sac/sac.py

@@ -15,6 +15,7 @@
 from torchrl.envs.utils import set_exploration_mode
 from torchrl.modules import OrnsteinUhlenbeckProcessWrapper
 from torchrl.record import VideoRecorder
+from torchrl.record.loggers import generate_exp_name, get_logger
 from torchrl.trainers.helpers.collectors import (
     make_collector_offpolicy,
     OffPolicyCollectorConfig,
@@ -32,7 +33,6 @@
 from torchrl.trainers.helpers.models import make_sac_model, SACModelConfig
 from torchrl.trainers.helpers.replay_buffer import make_replay_buffer, ReplayArgsConfig
 from torchrl.trainers.helpers.trainers import make_trainer, TrainerConfig
-from torchrl.trainers.loggers.utils import generate_exp_name, get_logger
 
 config_fields = [
     (config_field.name, config_field.type, config_field)

examples/td3/td3.py

@@ -38,7 +38,7 @@
 
 from torchrl.objectives import SoftUpdate
 from torchrl.objectives.td3 import TD3Loss
-from torchrl.trainers.loggers.utils import generate_exp_name, get_logger
+from torchrl.record.loggers import generate_exp_name, get_logger
 
 
 def env_maker(task, frame_skip=1, device="cpu", from_pixels=False):

test/test_loggers.py

@@ -12,10 +12,15 @@
 
 import pytest
 import torch
-from torchrl.trainers.loggers.csv import CSVLogger
-from torchrl.trainers.loggers.mlflow import _has_mlflow, _has_tv, MLFlowLogger
-from torchrl.trainers.loggers.tensorboard import _has_tb, TensorboardLogger
-from torchrl.trainers.loggers.wandb import _has_wandb, WandbLogger
+from torchrl.record.loggers import (
+    CSVLogger,
+    MLFlowLogger,
+    TensorboardLogger,
+    WandbLogger,
+)
+from torchrl.record.loggers.mlflow import _has_mlflow, _has_tv
+from torchrl.record.loggers.tensorboard import _has_tb
+from torchrl.record.loggers.wandb import _has_wandb
 
 if _has_tv:
     import torchvision

test/test_trainer.py

@@ -16,7 +16,7 @@
 
 try:
     from tensorboard.backend.event_processing import event_accumulator
-    from torchrl.trainers.loggers.tensorboard import TensorboardLogger
+    from torchrl.record.loggers import TensorboardLogger
 
     _has_tb = True
 except ImportError:

torchrl/trainers/loggers/__init__.py renamed to torchrl/record/loggers/__init__.py

@@ -4,7 +4,6 @@
 # LICENSE file in the root directory of this source tree.
 
 import os
-import warnings
 from tempfile import TemporaryDirectory
 from typing import Any, Dict, Optional
 
@@ -18,22 +17,21 @@
 
 from .common import Logger
 
-_has_mlflow = False
+MLFLOW_ERR = None
 try:
     import mlflow
 
     _has_mlflow = True
-except ImportError:
-    warnings.warn("mlflow could not be imported")
-_has_omgaconf = False
+except ImportError as err:
+    _has_mlflow = False
+    MLFLOW_ERR = err
+
 try:
     from omegaconf import OmegaConf
 
     _has_omgaconf = True
 except ImportError:
-    warnings.warn(
-        "OmegaConf could not be imported. Cannot log hydra configs without OmegaConf"
-    )
+    _has_omgaconf = False
 
 
 class MLFlowLogger(Logger):
@@ -67,7 +65,7 @@ def _create_experiment(self) -> "mlflow.ActiveRun":
             mlflow.ActiveRun: The mlflow experiment object.
         """
         if not _has_mlflow:
-            raise ImportError("MLFlow is not installed")
+            raise ImportError("MLFlow is not installed") from MLFLOW_ERR
         self.id = mlflow.create_experiment(**self._mlflow_kwargs)
         return mlflow.start_run(experiment_id=self.id)
 

torchrl/trainers/loggers/common.py renamed to torchrl/record/loggers/common.py

@@ -9,7 +9,7 @@
 import uuid
 from datetime import datetime
 
-from torchrl.trainers.loggers.common import Logger
+from torchrl.record.loggers.common import Logger
 
 
 def generate_exp_name(model_name: str, experiment_name: str) -> str:
@@ -37,22 +37,22 @@ def get_logger(
         kwargs (dict[str]): might contain either `wandb_kwargs` or `mlflow_kwargs`
     """
     if logger_type == "tensorboard":
-        from torchrl.trainers.loggers.tensorboard import TensorboardLogger
+        from torchrl.record.loggers.tensorboard import TensorboardLogger
 
         logger = TensorboardLogger(log_dir=logger_name, exp_name=experiment_name)
     elif logger_type == "csv":
-        from torchrl.trainers.loggers.csv import CSVLogger
+        from torchrl.record.loggers.csv import CSVLogger
 
         logger = CSVLogger(log_dir=logger_name, exp_name=experiment_name)
     elif logger_type == "wandb":
-        from torchrl.trainers.loggers.wandb import WandbLogger
+        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.trainers.loggers.mlflow import MLFlowLogger
+        from torchrl.record.loggers.mlflow import MLFlowLogger
 
         mlflow_kwargs = kwargs.get("mlflow_kwargs", {})
         logger = MLFlowLogger(

torchrl/trainers/loggers/csv.py renamed to torchrl/record/loggers/csv.py

@@ -16,7 +16,7 @@
 from tensordict.tensordict import TensorDictBase
 
 from torchrl.envs.transforms import ObservationTransform, Transform
-from torchrl.trainers.loggers import Logger
+from torchrl.record.loggers import Logger
 
 
 class VideoRecorder(ObservationTransform):

torchrl/trainers/loggers/mlflow.py renamed to torchrl/record/loggers/mlflow.py

@@ -30,8 +30,8 @@
     VecNorm,
 )
 from torchrl.envs.transforms.transforms import FlattenObservation, gSDENoise
+from torchrl.record.loggers import Logger
 from torchrl.record.recorder import VideoRecorder
-from torchrl.trainers.loggers import Logger
 
 LIBS = {
     "gym": GymEnv,

torchrl/trainers/loggers/tensorboard.py renamed to torchrl/record/loggers/tensorboard.py

@@ -18,7 +18,7 @@
 from torchrl.modules import reset_noise, SafeModule
 from torchrl.objectives.common import LossModule
 from torchrl.objectives.utils import TargetNetUpdater
-from torchrl.trainers.loggers import Logger
+from torchrl.record.loggers import Logger
 from torchrl.trainers.trainers import (
     BatchSubSampler,
     ClearCudaCache,

torchrl/trainers/loggers/utils.py renamed to torchrl/record/loggers/utils.py

@@ -27,7 +27,7 @@
 from torchrl.envs.utils import set_exploration_mode
 from torchrl.modules import SafeModule
 from torchrl.objectives.common import LossModule
-from torchrl.trainers.loggers import Logger
+from torchrl.record.loggers import Logger
 
 try:
     from tqdm import tqdm