[Feature] Remove wild imports in the library (#642)

* init

* missing docstrings

* Add explicit imports for subpackages

* fix lint

* Revert "fix lint"

This reverts commit 040e156.

* fix lint

Co-authored-by: vmoens <vincentmoens@gmail.com>

Author: sosmond

build_tools/setup_helpers/__init__.py

@@ -3,4 +3,4 @@
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
-from .extension import *  # noqa
+from .extension import get_ext_modules, CMakeBuild  # noqa

build_tools/setup_helpers/extension.py

@@ -14,10 +14,6 @@
 from setuptools import Extension
 from setuptools.command.build_ext import build_ext
 
-__all__ = [
-    "get_ext_modules",
-    "CMakeBuild",
-]
 
 _THIS_DIR = Path(__file__).parent.resolve()
 _ROOT_DIR = _THIS_DIR.parent.parent.resolve()

torchrl/collectors/__init__.py

@@ -3,4 +3,9 @@
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
-from .collectors import *
+from .collectors import (
+    SyncDataCollector,
+    aSyncDataCollector,
+    MultiaSyncDataCollector,
+    MultiSyncDataCollector,
+)

torchrl/collectors/collectors.py

@@ -20,24 +20,16 @@
 from torch import multiprocessing as mp
 from torch.utils.data import IterableDataset
 
+from torchrl.envs.transforms import TransformedEnv
 from torchrl.envs.utils import set_exploration_mode, step_mdp
 from .._utils import _check_for_faulty_process, prod
-from ..modules.tensordict_module import ProbabilisticTensorDictModule, TensorDictModule
-from .utils import split_trajectories
-
-__all__ = [
-    "SyncDataCollector",
-    "aSyncDataCollector",
-    "MultiaSyncDataCollector",
-    "MultiSyncDataCollector",
-]
-
-from torchrl.envs.transforms import TransformedEnv
 from ..data import TensorSpec
 from ..data.tensordict.tensordict import TensorDictBase, TensorDict
 from ..data.utils import CloudpickleWrapper, DEVICE_TYPING
 from ..envs.common import EnvBase
 from ..envs.vec_env import _BatchedEnv
+from ..modules.tensordict_module import ProbabilisticTensorDictModule, TensorDictModule
+from .utils import split_trajectories
 
 _TIMEOUT = 1.0
 _MIN_TIMEOUT = 1e-3  # should be several orders of magnitude inferior wrt time spent collecting a trajectory
@@ -47,6 +39,8 @@
 
 
 class RandomPolicy:
+    """A random policy for data collectors."""
+
     def __init__(self, action_spec: TensorSpec):
         """Random policy for a given action_spec.
 
@@ -73,6 +67,7 @@ def __call__(self, td: TensorDictBase) -> TensorDictBase:
 
 
 def recursive_map_to_cpu(dictionary: OrderedDict) -> OrderedDict:
+    """Maps the tensors to CPU through a nested dictionary."""
     return OrderedDict(
         **{
             k: recursive_map_to_cpu(item)

torchrl/data/__init__.py

@@ -3,7 +3,38 @@
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
-from .postprocs import *
-from .replay_buffers import *
-from .tensor_specs import *
-from .tensordict import *
+from .postprocs import MultiStep
+from .replay_buffers import (
+    ReplayBuffer,
+    PrioritizedReplayBuffer,
+    TensorDictReplayBuffer,
+    TensorDictPrioritizedReplayBuffer,
+    Storage,
+    ListStorage,
+    LazyMemmapStorage,
+    LazyTensorStorage,
+)
+from .tensor_specs import (
+    TensorSpec,
+    BoundedTensorSpec,
+    OneHotDiscreteTensorSpec,
+    UnboundedContinuousTensorSpec,
+    UnboundedDiscreteTensorSpec,
+    NdBoundedTensorSpec,
+    NdUnboundedContinuousTensorSpec,
+    NdUnboundedDiscreteTensorSpec,
+    BinaryDiscreteTensorSpec,
+    MultOneHotDiscreteTensorSpec,
+    DiscreteTensorSpec,
+    CompositeSpec,
+    DEVICE_TYPING,
+)
+from .tensordict import (
+    MemmapTensor,
+    set_transfer_ownership,
+    TensorDict,
+    SubTensorDict,
+    merge_tensordicts,
+    LazyStackedTensorDict,
+    SavedTensorDict,
+)

torchrl/data/postprocs/__init__.py

@@ -3,4 +3,4 @@
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
-from .postprocs import *
+from .postprocs import MultiStep

torchrl/data/postprocs/postprocs.py

@@ -14,8 +14,6 @@
 from torchrl.data.tensordict.tensordict import TensorDictBase
 from torchrl.data.utils import expand_as_right
 
-__all__ = ["MultiStep"]
-
 
 def _conv1d_reward(
     reward: torch.Tensor, gammas: torch.Tensor, n_steps_max: int
@@ -76,7 +74,7 @@ def _get_steps_to_next_obs(nonterminal: torch.Tensor, n_steps_max: int) -> torch
     return steps_to_next_obs
 
 
-def select_and_repeat(
+def _select_and_repeat(
     tensor: torch.Tensor,
     terminal: torch.Tensor,
     post_terminal: torch.Tensor,
@@ -203,7 +201,7 @@ def forward(self, tensordict: TensorDictBase) -> TensorDictBase:
         for key, item in selected_td.items():
             tensordict.set_(
                 key,
-                select_and_repeat(
+                _select_and_repeat(
                     item,
                     terminal,
                     post_terminal,

torchrl/data/replay_buffers/__init__.py

@@ -3,5 +3,10 @@
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
-from .replay_buffers import *
-from .storages import *
+from .replay_buffers import (
+    ReplayBuffer,
+    PrioritizedReplayBuffer,
+    TensorDictReplayBuffer,
+    TensorDictPrioritizedReplayBuffer,
+)
+from .storages import Storage, ListStorage, LazyMemmapStorage, LazyTensorStorage

torchrl/data/replay_buffers/replay_buffers.py

@@ -32,13 +32,6 @@
 )
 from torchrl.data.utils import DEVICE_TYPING
 
-__all__ = [
-    "ReplayBuffer",
-    "PrioritizedReplayBuffer",
-    "TensorDictReplayBuffer",
-    "TensorDictPrioritizedReplayBuffer",
-]
-
 
 def stack_tensors(list_of_tensor_iterators: List) -> Tuple[torch.Tensor]:
     """Zips a list of iterables containing tensor-like objects and stacks the resulting lists of tensors together.
@@ -730,6 +723,13 @@ def sample(self, size: int, return_weight: bool = False) -> TensorDictBase:
 
 
 class InPlaceSampler:
+    """A sampler to write tennsordicts in-place.
+
+    To be used cautiously as this may lead to unexpected behaviour (i.e. tensordicts
+    overwritten during execution).
+
+    """
+
     def __init__(self, device: Optional[DEVICE_TYPING] = None):
         self.out = None
         if device is None:

torchrl/data/replay_buffers/storages.py

@@ -23,8 +23,6 @@
 except ImportError:
     _has_ts = False
 
-__all__ = ["Storage", "ListStorage", "LazyMemmapStorage", "LazyTensorStorage"]
-
 
 class Storage:
     """A Storage is the container of a replay buffer.