[BE] Avoid warning users who don't care about PRB (#3099)

Author: vmoens

.gitignore

@@ -14,7 +14,7 @@ __pycache__/
 .installed.cfg
 MANIFEST
 build/
-data/
+./data/
 develop-eggs/
 dist/
 downloads/

build_tools/__init__.py

@@ -11,7 +11,6 @@
 import sys
 import time
 import unittest
-import warnings
 from functools import wraps
 from typing import Callable
 
@@ -24,6 +23,7 @@
 
 from torchrl._utils import (
     implement_for,
+    logger,
     logger as torchrl_logger,
     RL_WARNINGS,
     seed_generator,
@@ -792,7 +792,7 @@ def _call_value_nets(
             )
         else:
             if RL_WARNINGS:
-                warnings.warn(
+                logger.warning(
                     "Got a tensordict without a time-marked dimension, assuming time is along the last dimension. "
                     "This warning can be turned off by setting the environment variable RL_WARNINGS to False."
                 )

build_tools/setup_helpers/__init__.py

@@ -48,7 +48,7 @@ def _is_nightly(version):
         " - make sure ninja and cmake were installed\n"
         " - make sure you ran `python setup.py clean && python setup.py develop` and that no error was raised\n"
         " - make sure the version of PyTorch you are using matches the one that was present in your virtual env during "
-        f"setup. This package was built with PyTorch {pytorch_version}."
+        f"setup. This package was built with PyTorch {pytorch_version}. You can deactivate this warning by setting the environment variable `RL_WARNINGS=0`."
     )
 
 else:
@@ -59,5 +59,6 @@ def _is_nightly(version):
         "prioritized replay buffers can only be used with the PyTorch version they were built against. "
         f"This package was built with PyTorch {pytorch_version}. "
         "Workarounds include: (1) upgrading/downgrading PyTorch or TorchRL to compatible versions, "
-        "or (2) making a local install using `pip install git+https://github.com/pytorch/rl.git@<version>`."
+        "or (2) making a local install using `pip install git+https://github.com/pytorch/rl.git@<version>`. "
+        "You can deactivate this warning by setting the environment variable `RL_WARNINGS=0`."
     )

build_tools/setup_helpers/extension.py

@@ -18,7 +18,6 @@
 import torch
 from tensordict import MemoryMappedTensor, TensorDict
 from tensordict.utils import NestedKey
-
 from torch.utils._pytree import tree_map
 from torchrl._extension import EXTENSION_WARNING
 from torchrl._utils import _replace_last, logger
@@ -33,7 +32,7 @@
         SumSegmentTreeFp64,
     )
 except ImportError:
-    warnings.warn(EXTENSION_WARNING)
+    logger.warning(EXTENSION_WARNING)
 
 _EMPTY_STORAGE_ERROR = "Cannot sample from an empty storage."
 

test/_utils_internal.py

@@ -27,7 +27,7 @@
 from tensordict.utils import NestedKey, unravel_key
 from torch import Tensor
 
-from torchrl._utils import RL_WARNINGS
+from torchrl._utils import logger, RL_WARNINGS
 from torchrl.envs.utils import step_mdp
 from torchrl.objectives.utils import (
     _maybe_get_or_select,
@@ -452,7 +452,7 @@ def _call_value_nets(
                 ndim = list(data.names).index("time") + 1
             except ValueError:
                 if RL_WARNINGS:
-                    warnings.warn(
+                    logger.warning(
                         "Got a tensordict without a time-marked dimension, assuming time is along the last dimension. "
                         "This warning can be turned off by setting the environment variable RL_WARNINGS to False."
                     )