Update

[ghstack-poisoned]

Author: Vincent Moens

torchrl/envs/libs/brax.py

@@ -7,7 +7,6 @@
 import importlib.util
 import warnings
 
-
 import torch
 from packaging import version
 from tensordict import TensorDict, TensorDictBase

torchrl/envs/libs/jax_utils.py

@@ -41,9 +41,7 @@ def _tree_flatten(x, batch_size: torch.Size):
 }
 
 
-def _ndarray_to_tensor(
-    value: jnp.ndarray | np.ndarray  # noqa: F821
-) -> torch.Tensor:
+def _ndarray_to_tensor(value: jnp.ndarray | np.ndarray) -> torch.Tensor:  # noqa: F821
     from jax import dlpack as jax_dlpack, numpy as jnp
 
     # JAX arrays generated by jax.vmap would have Numpy dtypes.

torchrl/envs/model_based/dreamer.py

@@ -4,7 +4,6 @@
 # LICENSE file in the root directory of this source tree.
 from __future__ import annotations
 
-
 import torch
 from tensordict import TensorDict
 from tensordict.nn import TensorDictModule

torchrl/envs/transforms/functional.py

@@ -4,7 +4,6 @@
 # LICENSE file in the root directory of this source tree.
 from __future__ import annotations
 
-
 from torch import Tensor
 
 

torchrl/modules/distributions/discrete.py

@@ -128,9 +128,7 @@ def entropy(self):
         return -p_log_p.sum(-1)
 
     @_one_hot_wrapper(D.Categorical)
-    def sample(
-        self, sample_shape: torch.Size | Sequence | None = None
-    ) -> torch.Tensor:
+    def sample(self, sample_shape: torch.Size | Sequence | None = None) -> torch.Tensor:
         ...
 
     def rsample(self, sample_shape: torch.Size | Sequence = None) -> torch.Tensor:

torchrl/modules/distributions/utils.py

@@ -4,7 +4,6 @@
 # LICENSE file in the root directory of this source tree.
 from __future__ import annotations
 
-
 import torch
 from torch import autograd, distributions as d
 from torch.distributions import Independent, Transform, TransformedDistribution

torchrl/objectives/value/utils.py

@@ -4,7 +4,6 @@
 # LICENSE file in the root directory of this source tree.
 from __future__ import annotations
 
-
 import torch
 
 from tensordict import TensorDictBase

torchrl/trainers/helpers/collectors.py

@@ -250,9 +250,7 @@ def _make_collector(
 
 def make_collector_offpolicy(
     make_env: Callable[[], EnvBase],
-    actor_model_explore: (
-        TensorDictModuleWrapper | ProbabilisticTensorDictSequential
-    ),
+    actor_model_explore: (TensorDictModuleWrapper | ProbabilisticTensorDictSequential),
     cfg: DictConfig,  # noqa: F821
     make_env_kwargs: dict | None = None,
 ) -> DataCollectorBase:
@@ -314,9 +312,7 @@ def make_collector_offpolicy(
 
 def make_collector_onpolicy(
     make_env: Callable[[], EnvBase],
-    actor_model_explore: (
-        TensorDictModuleWrapper | ProbabilisticTensorDictSequential
-    ),
+    actor_model_explore: (TensorDictModuleWrapper | ProbabilisticTensorDictSequential),
     cfg: DictConfig,  # noqa: F821
     make_env_kwargs: dict | None = None,
 ) -> DataCollectorBase:

torchrl/trainers/helpers/trainers.py

@@ -82,9 +82,7 @@ def make_trainer(
     loss_module: LossModule,
     recorder: EnvBase | None = None,
     target_net_updater: TargetNetUpdater | None = None,
-    policy_exploration: None | (
-        TensorDictModuleWrapper | TensorDictModule
-    ) = None,
+    policy_exploration: None | (TensorDictModuleWrapper | TensorDictModule) = None,
     replay_buffer: ReplayBuffer | None = None,
     logger: Logger | None = None,
     cfg: DictConfig = None,  # noqa: F821