[Lint] pyupgrade

.github/unittest/helpers/coverage_run_parallel.py

@@ -28,8 +28,8 @@ def write_config(config_path: Path, argv: List[str]) -> None:
         argv: Arguments passed to this script, which need to be converted to config file entries
     """
     assert not config_path.exists(), "Temporary coverage config exists already"
-    cmdline = " ".join(shlex.quote(arg) for arg in argv[1:])
-    with open(str(config_path), "wt", encoding="utf-8") as fh:
+    cmdline = shlex.join(argv[1:])
+    with open(str(config_path), "w", encoding="utf-8") as fh:
         fh.write(
             f"""# .coveragerc to control coverage.py
 [run]

.pre-commit-config.yaml

@@ -35,3 +35,17 @@ repos:
     hooks:
       - id: pydocstyle
         files: ^torchrl/
+
+  - repo: https://github.com/asottile/pyupgrade
+    rev: v3.9.0
+    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/__init__.py

@@ -4,3 +4,5 @@
 # LICENSE file in the root directory of this source tree.
 
 from .extension import CMakeBuild, get_ext_modules  # noqa
+
+__all__ = ["CMakeBuild", "get_ext_modules"]

build_tools/setup_helpers/extension.py

@@ -14,7 +14,6 @@
 from setuptools import Extension
 from setuptools.command.build_ext import build_ext
 
-
 _THIS_DIR = Path(__file__).parent.resolve()
 _ROOT_DIR = _THIS_DIR.parent.parent.resolve()
 _TORCHRL_DIR = _ROOT_DIR / "torchrl"
@@ -130,7 +129,7 @@ def build_extension(self, ext):
             # using -j in the build_ext call, not supported by pip or PyPA-build.
             if hasattr(self, "parallel") and self.parallel:
                 # CMake 3.12+ only.
-                build_args += ["-j{}".format(self.parallel)]
+                build_args += [f"-j{self.parallel}"]
 
         if not os.path.exists(self.build_temp):
             os.makedirs(self.build_temp)

setup.py

@@ -30,7 +30,7 @@
 
 def get_version():
     version_txt = os.path.join(cwd, "version.txt")
-    with open(version_txt, "r") as f:
+    with open(version_txt) as f:
         version = f.readline().strip()
     if os.getenv("TORCHRL_BUILD_VERSION"):
         version = os.getenv("TORCHRL_BUILD_VERSION")
@@ -64,8 +64,8 @@ def parse_args(argv: List[str]) -> argparse.Namespace:
 def write_version_file(version):
     version_path = os.path.join(cwd, "torchrl", "version.py")
     with open(version_path, "w") as f:
-        f.write("__version__ = '{}'\n".format(version))
-        f.write("git_version = {}\n".format(repr(sha)))
+        f.write(f"__version__ = '{version}'\n")
+        f.write(f"git_version = {repr(sha)}\n")
 
 
 def _get_pytorch_version(is_nightly, is_local):
@@ -185,7 +185,7 @@ def _main(argv):
         version = get_version()
         write_version_file(version)
     TORCHRL_BUILD_VERSION = os.getenv("TORCHRL_BUILD_VERSION")
-    logging.info("Building wheel {}-{}".format(package_name, version))
+    logging.info(f"Building wheel {package_name}-{version}")
     logging.info(f"TORCHRL_BUILD_VERSION is {TORCHRL_BUILD_VERSION}")
 
     is_local = TORCHRL_BUILD_VERSION is None

sota-implementations/a2c/a2c_atari.py

@@ -13,7 +13,7 @@
 
 
 @hydra.main(config_path="", config_name="config_atari", version_base="1.1")
-def main(cfg: "DictConfig"):  # noqa: F821
+def main(cfg: DictConfig):  # noqa: F821
 
     from copy import deepcopy
 

sota-implementations/a2c/a2c_mujoco.py

@@ -13,7 +13,7 @@
 
 
 @hydra.main(config_path="", config_name="config_mujoco", version_base="1.1")
-def main(cfg: "DictConfig"):  # noqa: F821
+def main(cfg: DictConfig):  # noqa: F821
 
     from copy import deepcopy
 

sota-implementations/cql/cql_offline.py

@@ -15,16 +15,13 @@
 
 import hydra
 import numpy as np
-
 import torch
 import tqdm
 from tensordict.nn import CudaGraphModule
-
 from torchrl._utils import timeit
 from torchrl.envs.utils import ExplorationType, set_exploration_type
 from torchrl.objectives import group_optimizers
 from torchrl.record.loggers import generate_exp_name, get_logger
-
 from utils import (
     dump_video,
     log_metrics,
@@ -39,7 +36,7 @@
 
 
 @hydra.main(config_path="", config_name="offline_config", version_base="1.1")
-def main(cfg: "DictConfig"):  # noqa: F821
+def main(cfg: DictConfig):  # noqa: F821
     # Create logger
     exp_name = generate_exp_name("CQL-offline", cfg.logger.exp_name)
     logger = None

sota-implementations/cql/cql_online.py

@@ -21,12 +21,10 @@
 import tqdm
 from tensordict import TensorDict
 from tensordict.nn import CudaGraphModule
-
 from torchrl._utils import timeit
 from torchrl.envs.utils import ExplorationType, set_exploration_type
 from torchrl.objectives import group_optimizers
 from torchrl.record.loggers import generate_exp_name, get_logger
-
 from utils import (
     dump_video,
     log_metrics,
@@ -42,7 +40,7 @@
 
 
 @hydra.main(version_base="1.1", config_path="", config_name="online_config")
-def main(cfg: "DictConfig"):  # noqa: F821
+def main(cfg: DictConfig):  # noqa: F821
     # Create logger
     exp_name = generate_exp_name("CQL-online", cfg.logger.exp_name)
     logger = None

sota-implementations/cql/discrete_cql_online.py

@@ -16,16 +16,12 @@
 
 import hydra
 import numpy as np
-
 import torch
 import torch.cuda
 import tqdm
 from tensordict.nn import CudaGraphModule
-
 from torchrl._utils import timeit
-
 from torchrl.envs.utils import ExplorationType, set_exploration_type
-
 from torchrl.record.loggers import generate_exp_name, get_logger
 from utils import (
     log_metrics,
@@ -41,7 +37,7 @@
 
 
 @hydra.main(version_base="1.1", config_path="", config_name="discrete_cql_config")
-def main(cfg: "DictConfig"):  # noqa: F821
+def main(cfg: DictConfig):  # noqa: F821
     device = cfg.optim.device
     if device in ("", None):
         if torch.cuda.is_available():

sota-implementations/crossq/crossq.py

@@ -15,19 +15,15 @@
 import warnings
 
 import hydra
-
 import numpy as np
-
 import torch
 import torch.cuda
 import tqdm
 from tensordict import TensorDict
 from tensordict.nn import CudaGraphModule
-
 from torchrl._utils import timeit
 from torchrl.envs.utils import ExplorationType, set_exploration_type
 from torchrl.objectives import group_optimizers
-
 from torchrl.record.loggers import generate_exp_name, get_logger
 from utils import (
     log_metrics,
@@ -43,7 +39,7 @@
 
 
 @hydra.main(version_base="1.1", config_path=".", config_name="config")
-def main(cfg: "DictConfig"):  # noqa: F821
+def main(cfg: DictConfig):  # noqa: F821
     device = cfg.network.device
     if device in ("", None):
         if torch.cuda.is_available():

sota-implementations/ddpg/ddpg.py

@@ -15,16 +15,13 @@
 import warnings
 
 import hydra
-
 import numpy as np
 import torch
 import torch.cuda
 import tqdm
 from tensordict import TensorDict
 from tensordict.nn import CudaGraphModule
-
 from torchrl._utils import timeit
-
 from torchrl.envs.utils import ExplorationType, set_exploration_type
 from torchrl.objectives import group_optimizers
 from torchrl.record.loggers import generate_exp_name, get_logger
@@ -41,7 +38,7 @@
 
 
 @hydra.main(version_base="1.1", config_path="", config_name="config")
-def main(cfg: "DictConfig"):  # noqa: F821
+def main(cfg: DictConfig):  # noqa: F821
     device = cfg.optim.device
     if device in ("", None):
         if torch.cuda.is_available():

sota-implementations/decision_transformer/dt.py

@@ -19,11 +19,9 @@
 from tensordict.nn import CudaGraphModule
 from torchrl._utils import logger as torchrl_logger, timeit
 from torchrl.envs.libs.gym import set_gym_backend
-
 from torchrl.envs.utils import ExplorationType, set_exploration_type
 from torchrl.modules.tensordict_module import DecisionTransformerInferenceWrapper
 from torchrl.record import VideoRecorder
-
 from utils import (
     dump_video,
     log_metrics,
@@ -37,7 +35,7 @@
 
 
 @hydra.main(config_path="", config_name="dt_config", version_base="1.1")
-def main(cfg: "DictConfig"):  # noqa: F821
+def main(cfg: DictConfig):  # noqa: F821
     set_gym_backend(cfg.env.backend).set()
 
     model_device = cfg.optim.device

sota-implementations/decision_transformer/online_dt.py

@@ -20,7 +20,6 @@
 from torchrl.envs.utils import ExplorationType, set_exploration_type
 from torchrl.modules.tensordict_module import DecisionTransformerInferenceWrapper
 from torchrl.record import VideoRecorder
-
 from utils import (
     dump_video,
     log_metrics,
@@ -34,7 +33,7 @@
 
 
 @hydra.main(config_path="", config_name="odt_config", version_base="1.1")
-def main(cfg: "DictConfig"):  # noqa: F821
+def main(cfg: DictConfig):  # noqa: F821
     set_gym_backend(cfg.env.backend).set()
 
     model_device = cfg.optim.device

sota-implementations/discrete_sac/discrete_sac.py

@@ -38,7 +38,7 @@
 
 
 @hydra.main(version_base="1.1", config_path="", config_name="config")
-def main(cfg: "DictConfig"):  # noqa: F821
+def main(cfg: DictConfig):  # noqa: F821
     device = cfg.network.device
     if device in ("", None):
         if torch.cuda.is_available():

sota-implementations/dqn/dqn_atari.py

@@ -18,7 +18,6 @@
 import tqdm
 from tensordict.nn import CudaGraphModule, TensorDictSequential
 from torchrl._utils import timeit
-
 from torchrl.collectors import SyncDataCollector
 from torchrl.data import LazyMemmapStorage, TensorDictReplayBuffer
 from torchrl.envs import ExplorationType, set_exploration_type
@@ -32,7 +31,7 @@
 
 
 @hydra.main(config_path="", config_name="config_atari", version_base="1.1")
-def main(cfg: "DictConfig"):  # noqa: F821
+def main(cfg: DictConfig):  # noqa: F821
 
     device = cfg.device
     if device in ("", None):

sota-implementations/dqn/dqn_cartpole.py

@@ -11,7 +11,6 @@
 import torch.nn
 import torch.optim
 import tqdm
-
 from tensordict.nn import CudaGraphModule, TensorDictSequential
 from torchrl._utils import timeit
 from torchrl.collectors import SyncDataCollector
@@ -27,7 +26,7 @@
 
 
 @hydra.main(config_path="", config_name="config_cartpole", version_base="1.1")
-def main(cfg: "DictConfig"):  # noqa: F821
+def main(cfg: DictConfig):  # noqa: F821
 
     device = cfg.device
     if device in ("", None):

sota-implementations/dreamer/dreamer.py

@@ -11,6 +11,7 @@
 import torch
 import torch.cuda
 import tqdm
+
 from dreamer_utils import (
     _default_device,
     dump_video,
@@ -27,7 +28,6 @@
 from torchrl._utils import logger as torchrl_logger, timeit
 from torchrl.envs.utils import ExplorationType, set_exploration_type
 from torchrl.modules import RSSMRollout
-
 from torchrl.objectives.dreamer import (
     DreamerActorLoss,
     DreamerModelLoss,
@@ -37,7 +37,7 @@
 
 
 @hydra.main(version_base="1.1", config_path="", config_name="config")
-def main(cfg: "DictConfig"):  # noqa: F821
+def main(cfg: DictConfig):  # noqa: F821
     # cfg = correct_for_frame_skip(cfg)
 
     device = _default_device(cfg.networks.device)

sota-implementations/gail/gail.py

@@ -17,29 +17,25 @@
 import numpy as np
 import torch
 import tqdm
-
 from gail_utils import log_metrics, make_gail_discriminator, make_offline_replay_buffer
 from ppo_utils import eval_model, make_env, make_ppo_models
 from tensordict.nn import CudaGraphModule
-
 from torchrl._utils import compile_with_warmup, timeit
 from torchrl.collectors import SyncDataCollector
 from torchrl.data import LazyTensorStorage, TensorDictReplayBuffer
 from torchrl.data.replay_buffers.samplers import SamplerWithoutReplacement
-
 from torchrl.envs import set_gym_backend
 from torchrl.envs.utils import ExplorationType, set_exploration_type
 from torchrl.objectives import ClipPPOLoss, GAILLoss, group_optimizers
 from torchrl.objectives.value.advantages import GAE
 from torchrl.record import VideoRecorder
 from torchrl.record.loggers import generate_exp_name, get_logger
 
-
 torch.set_float32_matmul_precision("high")
 
 
 @hydra.main(config_path="", config_name="config")
-def main(cfg: "DictConfig"):  # noqa: F821
+def main(cfg: DictConfig):  # noqa: F821
     set_gym_backend(cfg.env.backend).set()
 
     device = cfg.gail.device

sota-implementations/gail/gail_utils.py

@@ -6,11 +6,9 @@
 
 import torch.nn as nn
 import torch.optim
-
 from torchrl.data.datasets.d4rl import D4RLExperienceReplay
 from torchrl.data.replay_buffers import SamplerWithoutReplacement
 from torchrl.envs import DoubleToFloat
-
 from torchrl.modules import SafeModule
 
 
@@ -45,7 +43,7 @@ def make_gail_discriminator(cfg, train_env, device="cpu"):
     # Define Discriminator Network
     class Discriminator(nn.Module):
         def __init__(self, state_dim, action_dim):
-            super(Discriminator, self).__init__()
+            super().__init__()
             self.fc1 = nn.Linear(state_dim + action_dim, hidden_dim)
             self.fc2 = nn.Linear(hidden_dim, hidden_dim)
             self.fc3 = nn.Linear(hidden_dim, 1)

sota-implementations/impala/impala_multi_node_ray.py

@@ -14,7 +14,7 @@
 
 
 @hydra.main(config_path="", config_name="config_multi_node_ray", version_base="1.1")
-def main(cfg: "DictConfig"):  # noqa: F821
+def main(cfg: DictConfig):  # noqa: F821
 
     import time