[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]

.github/workflows/lint.yml

@@ -35,7 +35,7 @@ jobs:
         echo '::endgroup::'
 
         echo '::group::Install lint tools'
-        pip install --progress-bar=off pre-commit
+        pip install --progress-bar=off pre-commit autoflake
         echo '::endgroup::'
 
         echo '::group::Lint Python source and configs'

.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

@@ -3,4 +3,6 @@
 # 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 CMakeBuild, get_ext_modules  # noqa
+from .extension import CMakeBuild, get_ext_modules
+
+__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)

docs/source/reference/envs.rst

@@ -1220,7 +1220,7 @@ Recorders are transforms that register data as they come in, for logging purpose
 
 Helpers
 -------
-.. currentmodule:: torchrl.envs.utils
+.. currentmodule:: torchrl.envs
 
 .. autosummary::
     :toctree: generated/

docs/source/reference/objectives.rst

@@ -111,6 +111,7 @@ auto-completion to make their choice.
     :template: rl_template_noinherit.rst
 
     LossModule
+    add_random_module
 
 DQN
 ---

examples/rlhf/models/actor_critic.py

@@ -2,6 +2,8 @@
 #
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
+from __future__ import annotations
+
 from torchrl.modules.tensordict_module.actors import LMHeadActorValueOperator
 from torchrl.modules.tensordict_module.common import VmapModule
 

setup.cfg

@@ -45,3 +45,7 @@ ignore-decorators =
     test_*
 ;    test/*.py
 ;    .circleci/*
+
+[autoflake]
+per-file-ignores =
+    torchrl/trainers/helpers/envs.py *

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)