Update

[ghstack-poisoned]

Author: Vincent Moens

.pre-commit-config.yaml

@@ -41,3 +41,11 @@ repos:
     hooks:
       - id: pyupgrade
         args: [--py38-plus]
+
+  - repo: local
+    hooks:
+    - id: autoflake
+      name: autoflake
+      entry: autoflake --in-place --remove-unused-variables --remove-all-unused-imports
+      language: system
+      types: [python]

build_tools/setup_helpers/extension.py

@@ -8,7 +8,7 @@
 import platform
 import subprocess
 from pathlib import Path
-from subprocess import CalledProcessError, STDOUT, check_output
+from subprocess import CalledProcessError, check_output, STDOUT
 
 import torch
 from setuptools import Extension

sota-implementations/dreamer/dreamer.py

@@ -11,6 +11,7 @@
 import torch
 import torch.cuda
 import tqdm
+
 from dreamer_utils import (
     _default_device,
     dump_video,
@@ -20,6 +21,7 @@
     make_environments,
     make_replay_buffer,
 )
+
 # mixed precision training
 from torch.amp import GradScaler
 from torch.nn.utils import clip_grad_norm_

sota-implementations/redq/utils.py

@@ -17,7 +17,8 @@
 )
 from torch import distributions as d, nn, optim
 from torch.optim.lr_scheduler import CosineAnnealingLR
-from torchrl._utils import VERBOSE, logger as torchrl_logger
+
+from torchrl._utils import logger as torchrl_logger, VERBOSE
 from torchrl.collectors.collectors import DataCollectorBase
 from torchrl.data import (
     LazyMemmapStorage,
@@ -35,10 +36,12 @@
     Compose,
     DMControlEnv,
     DoubleToFloat,
+    env_creator,
     EnvBase,
     EnvCreator,
     FlattenObservation,
     GrayScale,
+    gSDENoise,
     GymEnv,
     InitTracker,
     NoopResetEnv,
@@ -50,8 +53,6 @@
     ToTensorImage,
     TransformedEnv,
     VecNorm,
-    env_creator,
-    gSDENoise,
 )
 from torchrl.envs.utils import ExplorationType, set_exploration_type
 from torchrl.modules import (

test/_utils_internal.py

@@ -17,15 +17,21 @@
 import pytest
 import torch
 import torch.cuda
-from tensordict import NestedKey, TensorDict, TensorDictBase, tensorclass
+from tensordict import NestedKey, tensorclass, TensorDict, TensorDictBase
 from tensordict.nn import TensorDictModuleBase
 from torch import nn, vmap
-from torchrl._utils import (RL_WARNINGS, implement_for, logger as torchrl_logger, seed_generator)
+
+from torchrl._utils import (
+    implement_for,
+    logger as torchrl_logger,
+    RL_WARNINGS,
+    seed_generator,
+)
 from torchrl.data.utils import CloudpickleWrapper
 from torchrl.envs import MultiThreadedEnv, ObservationNorm
 from torchrl.envs.batched_envs import ParallelEnv, SerialEnv
 from torchrl.envs.libs.envpool import _has_envpool
-from torchrl.envs.libs.gym import GymEnv, _has_gym, gym_backend
+from torchrl.envs.libs.gym import _has_gym, gym_backend, GymEnv
 from torchrl.envs.transforms import (
     Compose,
     RewardClipping,

test/mocking_classes.py

@@ -10,9 +10,10 @@
 import numpy as np
 import torch
 import torch.nn as nn
-from tensordict import TensorDict, TensorDictBase, tensorclass
+from tensordict import tensorclass, TensorDict, TensorDictBase
 from tensordict.nn import TensorDictModuleBase
-from tensordict.utils import NestedKey, expand_right
+from tensordict.utils import expand_right, NestedKey
+
 from torchrl.data import (
     Binary,
     Bounded,
@@ -28,7 +29,11 @@
 from torchrl.envs import Transform
 from torchrl.envs.common import EnvBase
 from torchrl.envs.model_based.common import ModelBasedEnvBase
-from torchrl.envs.utils import (MarlGroupMapType, _terminated_or_truncated, check_marl_grouping)
+from torchrl.envs.utils import (
+    _terminated_or_truncated,
+    check_marl_grouping,
+    MarlGroupMapType,
+)
 
 spec_dict = {
     "bounded": Bounded,

test/smoke_test_deps.py

@@ -9,8 +9,6 @@
 
 import pytest
 
-from torchrl.envs.libs.gym import gym_backend
-
 
 def test_dm_control():
     import dm_control  # noqa: F401

test/test_collector.py

@@ -17,23 +17,34 @@
 import pytest
 import torch
 from packaging import version
-from tensordict import (LazyStackedTensorDict, NonTensorData, TensorDict, TensorDictBase, assert_allclose_td)
+from tensordict import (
+    assert_allclose_td,
+    LazyStackedTensorDict,
+    NonTensorData,
+    TensorDict,
+    TensorDictBase,
+)
 from tensordict.nn import (
     CudaGraphModule,
     TensorDictModule,
     TensorDictModuleBase,
     TensorDictSequential,
 )
 from torch import nn
+
 from torchrl._utils import (
     _make_ordinal_device,
     _replace_last,
     logger as torchrl_logger,
     prod,
     seed_generator,
 )
-from torchrl.collectors import SyncDataCollector, aSyncDataCollector
-from torchrl.collectors.collectors import (MultiSyncDataCollector, MultiaSyncDataCollector, _Interruptor)
+from torchrl.collectors import aSyncDataCollector, SyncDataCollector
+from torchrl.collectors.collectors import (
+    _Interruptor,
+    MultiaSyncDataCollector,
+    MultiSyncDataCollector,
+)
 from torchrl.collectors.utils import split_trajectories
 from torchrl.data import (
     Composite,
@@ -54,9 +65,14 @@
     StepCounter,
     Transform,
 )
-from torchrl.envs.libs.gym import GymEnv, _has_gym, gym_backend, set_gym_backend
+from torchrl.envs.libs.gym import _has_gym, gym_backend, GymEnv, set_gym_backend
 from torchrl.envs.transforms import TransformedEnv, VecNorm
-from torchrl.envs.utils import (PARTIAL_MISSING_ERR, RandomPolicy, _aggregate_end_of_traj, check_env_specs)
+from torchrl.envs.utils import (
+    _aggregate_end_of_traj,
+    check_env_specs,
+    PARTIAL_MISSING_ERR,
+    RandomPolicy,
+)
 from torchrl.modules import Actor, OrnsteinUhlenbeckProcessModule, SafeModule
 
 if os.getenv("PYTORCH_TEST_FBCODE"):

test/test_cost.py

@@ -5859,7 +5859,6 @@ def test_crossq_tensordict_keys(self, td_est):
 
         actor = self._create_mock_actor()
         qvalue = self._create_mock_qvalue()
-        value = None
 
         loss_fn = CrossQLoss(
             actor_network=actor,

test/test_exploration.py

@@ -381,7 +381,7 @@ def test_nested(
         action_spec = env.action_spec
         d_act = action_spec.shape[-1]
 
-        net = nn.LazyLinear(d_act).to(device)
+        nn.LazyLinear(d_act).to(device)
         policy = TensorDictModule(
             CountingEnvCountModule(action_spec=action_spec),
             in_keys=[("data", "states") if nested_obs_action else "observation"],