Reformatting due to switch to ruff

Author: opcode81

docs/nbstripout.py

@@ -1,4 +1,5 @@
 """Implements a platform-independent way of calling nbstripout (used in pyproject.toml)."""
+
 import glob
 import os
 from pathlib import Path

examples/inverse/irl_gail.py

@@ -247,7 +247,7 @@ def dist(loc_scale: tuple[torch.Tensor, torch.Tensor]) -> Distribution:
     test_collector = Collector[CollectStats](algorithm, test_envs)
     # log
     t0 = datetime.datetime.now().strftime("%m%d_%H%M%S")
-    log_file = f'seed_{args.seed}_{t0}-{args.task.replace("-", "_")}_gail'
+    log_file = f"seed_{args.seed}_{t0}-{args.task.replace('-', '_')}_gail"
     log_path = os.path.join(args.logdir, args.task, "gail", log_file)
     writer = SummaryWriter(log_path)
     writer.add_text("args", str(args))

examples/offline/convert_rl_unplugged_atari.py

@@ -27,6 +27,7 @@
   * episode_return: Total episode return computed using per-step [-1, 1]
       clipping.
 """
+
 import os
 from argparse import ArgumentParser, Namespace
 
@@ -179,13 +180,16 @@ def download(url: str, fname: str, chunk_size: int | None = 1024) -> None:
     if os.path.exists(fname):
         print(f"Found cached file at {fname}.")
         return
-    with open(fname, "wb") as ofile, tqdm(
-        desc=fname,
-        total=total,
-        unit="iB",
-        unit_scale=True,
-        unit_divisor=1024,
-    ) as bar:
+    with (
+        open(fname, "wb") as ofile,
+        tqdm(
+            desc=fname,
+            total=total,
+            unit="iB",
+            unit_scale=True,
+            unit_divisor=1024,
+        ) as bar,
+    ):
         for data in resp.iter_content(chunk_size=chunk_size):
             size = ofile.write(data)
             bar.update(size)

test/base/env.py

@@ -34,9 +34,9 @@ def __init__(
         random_sleep: bool = False,
         array_state: bool = False,
     ) -> None:
-        assert (
-            dict_state + recurse_state + array_state <= 1
-        ), "dict_state / recurse_state / array_state can be only one true"
+        assert dict_state + recurse_state + array_state <= 1, (
+            "dict_state / recurse_state / array_state can be only one true"
+        )
         self.size = size
         self.sleep = sleep
         self.random_sleep = random_sleep
@@ -208,9 +208,9 @@ def step(
 
 class MyGoalEnv(MoveToRightEnv):
     def __init__(self, *args: Any, **kwargs: Any) -> None:
-        assert (
-            kwargs.get("dict_state", 0) + kwargs.get("recurse_state", 0) == 0
-        ), "dict_state / recurse_state not supported"
+        assert kwargs.get("dict_state", 0) + kwargs.get("recurse_state", 0) == 0, (
+            "dict_state / recurse_state not supported"
+        )
         super().__init__(*args, **kwargs)
         super().reset(options={"state": 0})
 

test/base/test_buffer.py

@@ -1,7 +1,6 @@
 import os
 import pickle
 import tempfile
-from test.base.env import MoveToRightEnv, MyGoalEnv
 from typing import cast
 
 import h5py
@@ -10,6 +9,7 @@
 import pytest
 import torch
 
+from test.base.env import MoveToRightEnv, MyGoalEnv
 from tianshou.data import (
     Batch,
     CachedReplayBuffer,
@@ -1451,10 +1451,7 @@ def test_custom_key() -> None:
     # Check if they have the same keys
     assert set(batch.get_keys()) == set(
         sampled_batch.get_keys(),
-    ), "Batches have different keys: {} and {}".format(
-        set(batch.get_keys()),
-        set(sampled_batch.get_keys()),
-    )
+    ), f"Batches have different keys: {set(batch.get_keys())} and {set(sampled_batch.get_keys())}"
     # Compare the values for each key
     for key in batch.get_keys():
         if isinstance(batch.__dict__[key], np.ndarray) and isinstance(

test/base/test_collector.py

@@ -1,12 +1,12 @@
 from collections.abc import Callable, Sequence
-from test.base.env import MoveToRightEnv, NXEnv
 from typing import Any
 
 import gymnasium as gym
 import numpy as np
 import pytest
 import tqdm
 
+from test.base.env import MoveToRightEnv, NXEnv
 from tianshou.algorithm.algorithm_base import Policy, episode_mc_return_to_go
 from tianshou.data import (
     AsyncCollector,
@@ -410,11 +410,8 @@ def test_collector_with_dict_state() -> None:
     result = c1.collect(n_episode=8)
     assert result.n_collected_episodes == 8
     lens = np.bincount(result.lens)
-    assert (
-        result.n_collected_steps == 21
-        and np.all(lens == [0, 0, 2, 2, 2, 2])
-        or result.n_collected_steps == 20
-        and np.all(lens == [0, 0, 3, 1, 2, 2])
+    assert (result.n_collected_steps == 21 and np.all(lens == [0, 0, 2, 2, 2, 2])) or (
+        result.n_collected_steps == 20 and np.all(lens == [0, 0, 3, 1, 2, 2])
     )
     batch, _ = c1.buffer.sample(10)
     c0.buffer.update(c1.buffer)

test/base/test_env.py

@@ -1,14 +1,14 @@
 import sys
 import time
 from collections.abc import Callable
-from test.base.env import MoveToRightEnv, NXEnv
 from typing import Any, Literal
 
 import gymnasium as gym
 import numpy as np
 import pytest
 from gymnasium.spaces.discrete import Discrete
 
+from test.base.env import MoveToRightEnv, NXEnv
 from tianshou.data import Batch
 from tianshou.env import (
     ContinuousToDiscrete,

test/continuous/test_ddpg.py

@@ -1,12 +1,12 @@
 import argparse
 import os
-from test.determinism_test import AlgorithmDeterminismTest
 
 import gymnasium as gym
 import numpy as np
 import torch
 from torch.utils.tensorboard import SummaryWriter
 
+from test.determinism_test import AlgorithmDeterminismTest
 from tianshou.algorithm import DDPG
 from tianshou.algorithm.algorithm_base import Algorithm
 from tianshou.algorithm.modelfree.ddpg import ContinuousDeterministicPolicy

test/continuous/test_npg.py

@@ -1,6 +1,5 @@
 import argparse
 import os
-from test.determinism_test import AlgorithmDeterminismTest
 
 import gymnasium as gym
 import numpy as np
@@ -9,6 +8,7 @@
 from torch.distributions import Distribution, Independent, Normal
 from torch.utils.tensorboard import SummaryWriter
 
+from test.determinism_test import AlgorithmDeterminismTest
 from tianshou.algorithm import NPG
 from tianshou.algorithm.algorithm_base import Algorithm
 from tianshou.algorithm.modelfree.reinforce import ProbabilisticActorPolicy

test/continuous/test_ppo.py

@@ -1,13 +1,13 @@
 import argparse
 import os
-from test.determinism_test import AlgorithmDeterminismTest
 
 import gymnasium as gym
 import numpy as np
 import torch
 from torch.distributions import Distribution, Independent, Normal
 from torch.utils.tensorboard import SummaryWriter
 
+from test.determinism_test import AlgorithmDeterminismTest
 from tianshou.algorithm import PPO
 from tianshou.algorithm.algorithm_base import Algorithm
 from tianshou.algorithm.modelfree.reinforce import ProbabilisticActorPolicy