[Refactor] A less verbose torchrl (#973)

Author: vmoens

torchrl/_utils.py

@@ -1,13 +1,17 @@
 import collections
+
 import math
 import os
 import time
 import warnings
+from distutils.util import strtobool
 from functools import wraps
 from importlib import import_module
 
 import numpy as np
 
+VERBOSE = strtobool(os.environ.get("VERBOSE", "0"))
+
 
 class timeit:
     """A dirty but easy to use decorator for profiling code."""

torchrl/collectors/collectors.py

@@ -23,7 +23,7 @@
 from torch import multiprocessing as mp
 from torch.utils.data import IterableDataset
 
-from torchrl._utils import _check_for_faulty_process, prod
+from torchrl._utils import _check_for_faulty_process, prod, VERBOSE
 from torchrl.collectors.utils import split_trajectories
 from torchrl.data import TensorSpec
 from torchrl.data.utils import CloudpickleWrapper, DEVICE_TYPING
@@ -1662,7 +1662,7 @@ def _main_async_collector(
     idx: int = 0,
     exploration_mode: str = DEFAULT_EXPLORATION_MODE,
     reset_when_done: bool = True,
-    verbose: bool = False,
+    verbose: bool = VERBOSE,
 ) -> None:
     pipe_parent.close()
     #  init variables that will be cleared when closing

torchrl/data/replay_buffers/storages.py

@@ -15,7 +15,7 @@
 from tensordict.prototype import is_tensorclass
 from tensordict.tensordict import is_tensor_collection, TensorDict, TensorDictBase
 
-from torchrl._utils import _CKPT_BACKEND
+from torchrl._utils import _CKPT_BACKEND, VERBOSE
 from torchrl.data.replay_buffers.utils import INT_CLASSES
 
 try:
@@ -224,7 +224,8 @@ def load_state_dict(self, state_dict):
         self._len = state_dict["_len"]
 
     def _init(self, data: Union[TensorDictBase, torch.Tensor]) -> None:
-        print("Creating a TensorStorage...")
+        if VERBOSE:
+            print("Creating a TensorStorage...")
         if isinstance(data, torch.Tensor):
             # if Tensor, we just create a MemmapTensor of the desired shape, device and dtype
             out = torch.empty(
@@ -353,16 +354,18 @@ def load_state_dict(self, state_dict):
         self._len = state_dict["_len"]
 
     def _init(self, data: Union[TensorDictBase, torch.Tensor]) -> None:
-        print("Creating a MemmapStorage...")
+        if VERBOSE:
+            print("Creating a MemmapStorage...")
         if isinstance(data, torch.Tensor):
             # if Tensor, we just create a MemmapTensor of the desired shape, device and dtype
             out = MemmapTensor(
                 self.max_size, *data.shape, device=self.device, dtype=data.dtype
             )
             filesize = os.path.getsize(out.filename) / 1024 / 1024
-            print(
-                f"The storage was created in {out.filename} and occupies {filesize} Mb of storage."
-            )
+            if VERBOSE:
+                print(
+                    f"The storage was created in {out.filename} and occupies {filesize} Mb of storage."
+                )
         elif is_tensorclass(data):
             out = (
                 data.clone()
@@ -374,12 +377,13 @@ def _init(self, data: Union[TensorDictBase, torch.Tensor]) -> None:
                 out.items(include_nested=True, leaves_only=True), key=str
             ):
                 filesize = os.path.getsize(tensor.filename) / 1024 / 1024
-                print(
-                    f"\t{key}: {tensor.filename}, {filesize} Mb of storage (size: {tensor.shape})."
-                )
+                if VERBOSE:
+                    print(
+                        f"\t{key}: {tensor.filename}, {filesize} Mb of storage (size: {tensor.shape})."
+                    )
         else:
-            # out = TensorDict({}, [self.max_size, *data.shape])
-            print("The storage is being created: ")
+            if VERBOSE:
+                print("The storage is being created: ")
             out = (
                 data.clone()
                 .expand(self.max_size, *data.shape)
@@ -390,9 +394,10 @@ def _init(self, data: Union[TensorDictBase, torch.Tensor]) -> None:
                 out.items(include_nested=True, leaves_only=True), key=str
             ):
                 filesize = os.path.getsize(tensor.filename) / 1024 / 1024
-                print(
-                    f"\t{key}: {tensor.filename}, {filesize} Mb of storage (size: {tensor.shape})."
-                )
+                if VERBOSE:
+                    print(
+                        f"\t{key}: {tensor.filename}, {filesize} Mb of storage (size: {tensor.shape})."
+                    )
         self._storage = out
         self.initialized = True
 

torchrl/envs/libs/dm_control.py

@@ -19,13 +19,16 @@
     UnboundedDiscreteTensorSpec,
 )
 
+from ..._utils import VERBOSE
+
 from ...data.utils import DEVICE_TYPING, numpy_to_torch_dtype_dict
 from ..gym_like import GymLikeEnv
 
 if torch.has_cuda and torch.cuda.device_count() > 1:
     n = torch.cuda.device_count() - 1
     os.environ["EGL_DEVICE_ID"] = str(1 + (os.getpid() % n))
-    print("EGL_DEVICE_ID: ", os.environ["EGL_DEVICE_ID"])
+    if VERBOSE:
+        print("EGL_DEVICE_ID: ", os.environ["EGL_DEVICE_ID"])
 
 try:
 

torchrl/envs/vec_env.py

@@ -21,7 +21,7 @@
 from tensordict import TensorDict
 from tensordict.tensordict import LazyStackedTensorDict, TensorDictBase
 from torch import multiprocessing as mp
-from torchrl._utils import _check_for_faulty_process
+from torchrl._utils import _check_for_faulty_process, VERBOSE
 from torchrl.data import (
     CompositeSpec,
     DiscreteTensorSpec,
@@ -126,7 +126,7 @@ class _BatchedEnv(EnvBase):
 
     """
 
-    _verbose: bool = False
+    _verbose: bool = VERBOSE
     _excluded_wrapped_keys = [
         "is_closed",
         "parent_channels",

torchrl/modules/models/exploration.py

@@ -2,8 +2,8 @@
 #
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
-
 import math
+import warnings
 from typing import Optional, Sequence, Union
 
 import torch
@@ -390,7 +390,7 @@ def forward(self, mu, state, _eps_gSDE):
 
         sigma = (sigma * state.unsqueeze(-2)).pow(2).sum(-1).clamp_min(1e-5).sqrt()
         if not torch.isfinite(sigma).all():
-            print("inf sigma")
+            warnings.warn("inf sigma")
 
         if self.transform is not None:
             action = self.transform(action)

torchrl/modules/tensordict_module/common.py

@@ -20,7 +20,7 @@
 
     _has_functorch = True
 except ImportError:
-    print(
+    warnings.warn(
         "failed to import functorch. TorchRL's features that do not require "
         "functional programming should work, but functionality and performance "
         "may be affected. Consider installing functorch and/or upgrating pytorch."

torchrl/trainers/helpers/envs.py

@@ -8,6 +8,7 @@
 
 import torch
 
+from torchrl._utils import VERBOSE
 from torchrl.envs import ParallelEnv
 from torchrl.envs.common import EnvBase
 from torchrl.envs.env_creator import env_creator, EnvCreator
@@ -394,7 +395,8 @@ def get_stats_random_rollout(
             cfg=cfg, use_env_creator=False, stats={"loc": 0.0, "scale": 1.0}
         )()
 
-    print("computing state stats")
+    if VERBOSE:
+        print("computing state stats")
     if not hasattr(cfg, "init_env_steps"):
         raise AttributeError("init_env_steps missing from arguments.")
 
@@ -426,11 +428,12 @@ def get_stats_random_rollout(
     m[s == 0] = 0.0
     s[s == 0] = 1.0
 
-    print(
-        f"stats computed for {val_stats.numel()} steps. Got: \n"
-        f"loc = {m}, \n"
-        f"scale = {s}"
-    )
+    if VERBOSE:
+        print(
+            f"stats computed for {val_stats.numel()} steps. Got: \n"
+            f"loc = {m}, \n"
+            f"scale = {s}"
+        )
     if not torch.isfinite(m).all():
         raise RuntimeError("non-finite values found in mean")
     if not torch.isfinite(s).all():

torchrl/trainers/helpers/trainers.py

@@ -12,6 +12,7 @@
 from torch import optim
 from torch.optim.lr_scheduler import CosineAnnealingLR
 
+from torchrl._utils import VERBOSE
 from torchrl.collectors.collectors import _DataCollector
 from torchrl.data import ReplayBuffer
 from torchrl.envs.common import EnvBase
@@ -168,16 +169,17 @@ def make_trainer(
     else:
         raise NotImplementedError(f"lr scheduler {cfg.lr_scheduler}")
 
-    print(
-        f"collector = {collector}; \n"
-        f"loss_module = {loss_module}; \n"
-        f"recorder = {recorder}; \n"
-        f"target_net_updater = {target_net_updater}; \n"
-        f"policy_exploration = {policy_exploration}; \n"
-        f"replay_buffer = {replay_buffer}; \n"
-        f"logger = {logger}; \n"
-        f"cfg = {cfg}; \n"
-    )
+    if VERBOSE:
+        print(
+            f"collector = {collector}; \n"
+            f"loss_module = {loss_module}; \n"
+            f"recorder = {recorder}; \n"
+            f"target_net_updater = {target_net_updater}; \n"
+            f"policy_exploration = {policy_exploration}; \n"
+            f"replay_buffer = {replay_buffer}; \n"
+            f"logger = {logger}; \n"
+            f"cfg = {cfg}; \n"
+        )
 
     if logger is not None:
         # log hyperparams

torchrl/trainers/trainers.py

@@ -19,7 +19,7 @@
 from tensordict.utils import expand_right
 from torch import nn, optim
 
-from torchrl._utils import _CKPT_BACKEND, KeyDependentDefaultDict
+from torchrl._utils import _CKPT_BACKEND, KeyDependentDefaultDict, VERBOSE
 from torchrl.collectors.collectors import _DataCollector
 from torchrl.data import TensorDictPrioritizedReplayBuffer, TensorDictReplayBuffer
 from torchrl.data.utils import DEVICE_TYPING
@@ -451,7 +451,8 @@ def __del__(self):
         self.collector.shutdown()
 
     def shutdown(self):
-        print("shutting down collector")
+        if VERBOSE:
+            print("shutting down collector")
         self.collector.shutdown()
 
     def optim_steps(self, batch: TensorDictBase) -> None: