fix lint errors

Author: smanolloff

ci/clean_notebooks.py

@@ -63,7 +63,6 @@ def clean_notebook(file: pathlib.Path, check_only=False) -> None:
         print(f"Checking {file}")
 
     for cell in nb.cells:
-
         # Remove empty cells
         if cell["cell_type"] == "code" and not cell["source"]:
             if check_only:

src/imitation/algorithms/adversarial/common.py

@@ -9,8 +9,8 @@
 import torch.utils.tensorboard as thboard
 import tqdm
 from stable_baselines3.common import base_class, on_policy_algorithm, policies, vec_env
-from stable_baselines3.common.type_aliases import MaybeCallback
 from stable_baselines3.common.callbacks import BaseCallback, ConvertCallback
+from stable_baselines3.common.type_aliases import MaybeCallback
 from stable_baselines3.sac import policies as sac_policies
 from torch.nn import functional as F
 

src/imitation/algorithms/bc.py

@@ -57,7 +57,6 @@ def __post_init__(self) -> None:
 
     def __iter__(self) -> Iterator[types.TransitionMapping]:
         def batch_iterator() -> Iterator[types.TransitionMapping]:
-
             # Note: the islice here ensures we do not exceed self.n_epochs
             for epoch_num in itertools.islice(itertools.count(), self.n_epochs):
                 some_batch_was_yielded = False

src/imitation/data/huggingface_utils.py

@@ -16,6 +16,7 @@ class TrajectoryDatasetSequence(Sequence[types.Trajectory]):
 
     def __init__(self, dataset: datasets.Dataset):
         """Construct a TrajectoryDatasetSequence."""
+
         # TODO: this is just a temporary workaround for
         #  https://github.com/huggingface/datasets/issues/5517
         #  switch to .with_format("numpy") once it's fixed
@@ -31,7 +32,6 @@ def __len__(self) -> int:
         return len(self._dataset)
 
     def __getitem__(self, idx):
-
         if isinstance(idx, slice):
             # Note: we could use self._dataset[idx] here and then convert the result of
             #   that to a series of trajectories, but if we do that, we run into trouble

src/imitation/scripts/train_adversarial.py

@@ -24,12 +24,15 @@
 
 
 class CheckpointCallback(BaseCallback):
+    """A callback for calling `save` at regular intervals."""
+
     def __init__(
         self,
         trainer: common.AdversarialTrainer,
         log_dir: pathlib.Path,
-        interval: int
+        interval: int,
     ):
+        """Creates new Checkpoint callback."""
         super().__init__(self)
         self.trainer = trainer
         self.log_dir = log_dir

tests/algorithms/test_adversarial.py

@@ -468,11 +468,10 @@ def test_regression_gail_with_sac(
 
 
 def test_gen_callback(trainer: common.AdversarialTrainer):
-    learner = stable_baselines3.PPO("MlpPolicy", env=trainer.venv)
-
     def make_fn_callback(calls, key):
         def cb(_a, _b):
             calls[key] += 1
+
         return cb
 
     class SB3Callback(BaseCallback):
@@ -490,10 +489,10 @@ def _on_step(self):
 
     trainer.train(n_steps, callback=make_fn_callback(calls, "fn"))
     trainer.train(n_steps, callback=SB3Callback(calls, "sb3"))
-    trainer.train(n_steps, callback=[
-        SB3Callback(calls, "list.0"),
-        SB3Callback(calls, "list.1")
-    ])
+    trainer.train(
+        n_steps,
+        callback=[SB3Callback(calls, "list.0"), SB3Callback(calls, "list.1")],
+    )
 
     # Env steps for off-plicy algos (DQN) may exceed `total_timesteps`,
     # so we check if the callback was called *at least* that many times.

tests/data/test_huggingface_utils.py

@@ -71,7 +71,6 @@ def test_sliced_access(data: st.DataObject, trajectories: Sequence[types.Traject
 
     # Note: we test if for 10 slices at a time because creating the dataset is slow
     for _ in range(10):
-
         # GIVEN
         the_slice = data.draw(slices_strategy)
         indices_of_slice = list(range(*the_slice.indices(len(trajectories))))
@@ -113,7 +112,6 @@ def test_sliced_info_dict_access(
 
     # Note: we test if for 10 slices at a time because creating the dataset is slow
     for _ in range(10):
-
         # GIVEN
         the_slice = data.draw(slices_strategy)
         indices_of_slice = list(range(*the_slice.indices(len(wrapped_info_dicts))))

tests/util/test_wb_logger.py

@@ -62,7 +62,6 @@ def log(
         commit: bool = False,
         sync: bool = False,
     ):
-
         assert self._initialized
         if sync:
             raise NotImplementedError("usage of sync to MockWandb.log not implemented")