Align Ruff config and enable docstring rules

Author: Trinkle23897

benchmark/atari_wrappers.py

@@ -33,30 +33,34 @@ class NoopResetEnv(gym.Wrapper):
     """
 
     def __init__(self, env, noop_max=30):
+        """Initialize the no-op reset wrapper."""
         super().__init__(env)
         self.noop_max = noop_max
         self.noop_action = 0
         assert env.unwrapped.get_action_meanings()[0] == "NOOP"
 
     def reset(self, *, seed=None, options=None):
+        """Reset the environment after a random number of no-op actions."""
         obs, info = self.env.reset(seed=seed, options=options)
         noops = self.unwrapped.np_random.integers(1, self.noop_max + 1)
         for _ in range(noops):
-            obs, _, terminated, truncated, info = self.env.step(self.noop_action)
+            obs, _, terminated, truncated, info = self.env.step(
+                self.noop_action
+            )
             if terminated or truncated:
                 obs, info = self.env.reset()
         return obs, info
 
 
 class MaxAndSkipEnv(gym.Wrapper):
-    """Return only every `skip`-th frame (frameskipping) using most recent raw
-    observations (for max pooling across time steps)
+    """Return every `skip`-th frame with max pooling over recent frames.
 
     :param gym.Env env: the environment to wrap.
     :param int skip: number of `skip`-th frame.
     """
 
     def __init__(self, env, skip=4):
+        """Initialize the instance."""
         super().__init__(env)
         self._skip = skip
 
@@ -85,11 +89,13 @@ class EpisodicLifeEnv(gym.Wrapper):
     """
 
     def __init__(self, env):
+        """Initialize episodic-life tracking."""
         super().__init__(env)
         self.lives = 0
         self.was_real_done = True
 
     def step(self, action):
+        """Step the environment and treat life loss as episode end."""
         obs, reward, terminated, truncated, info = self.env.step(action)
         self.was_real_done = terminated or truncated
         # check current lives, make loss of life terminal, then update lives to
@@ -129,11 +135,13 @@ class FireResetEnv(gym.Wrapper):
     """
 
     def __init__(self, env):
+        """Initialize the fire-reset wrapper."""
         super().__init__(env)
         assert env.unwrapped.get_action_meanings()[1] == "FIRE"
         assert len(env.unwrapped.get_action_meanings()) >= 3
 
     def reset(self, *, seed=None, options=None):
+        """Reset the environment and apply the fire action."""
         self.env.reset(seed=seed, options=options)
         obs, _, terminated, truncated, info = self.env.step(1)
         if terminated or truncated:
@@ -148,6 +156,7 @@ class WarpFrame(gym.ObservationWrapper):
     """
 
     def __init__(self, env):
+        """Initialize grayscale frame warping."""
         super().__init__(env)
         self.size = 84
         self.observation_space = gym.spaces.Box(
@@ -158,9 +167,11 @@ def __init__(self, env):
         )
 
     def observation(self, frame):
-        """Returns the current observation from a frame"""
+        """Returns the current observation from a frame."""
         frame = cv2.cvtColor(frame, cv2.COLOR_RGB2GRAY)
-        return cv2.resize(frame, (self.size, self.size), interpolation=cv2.INTER_AREA)
+        return cv2.resize(
+            frame, (self.size, self.size), interpolation=cv2.INTER_AREA
+        )
 
 
 class ScaledFloatFrame(gym.ObservationWrapper):
@@ -170,6 +181,7 @@ class ScaledFloatFrame(gym.ObservationWrapper):
     """
 
     def __init__(self, env):
+        """Initialize the instance."""
         super().__init__(env)
         low = np.min(env.observation_space.low)
         high = np.max(env.observation_space.high)
@@ -183,6 +195,7 @@ def __init__(self, env):
         )
 
     def observation(self, observation):
+        """Scale an observation to the [0, 1] range."""
         return (observation - self.bias) / self.scale
 
 
@@ -193,6 +206,7 @@ class ClipRewardEnv(gym.RewardWrapper):
     """
 
     def __init__(self, env):
+        """Initialize reward clipping."""
         super().__init__(env)
         self.reward_range = (-1, 1)
 
@@ -209,6 +223,7 @@ class FrameStack(gym.Wrapper):
     """
 
     def __init__(self, env, n_frames):
+        """Initialize the frame stack buffer."""
         super().__init__(env)
         self.n_frames = n_frames
         self.frames = deque([], maxlen=n_frames)
@@ -221,12 +236,14 @@ def __init__(self, env, n_frames):
         )
 
     def reset(self, *, seed=None, options=None):
+        """Reset the environment and refill the frame stack."""
         obs, info = self.env.reset(seed=seed, options=options)
         for _ in range(self.n_frames):
             self.frames.append(obs)
         return self._get_ob(), info
 
     def step(self, action):
+        """Step the environment and append the latest frame."""
         obs, reward, terminated, truncated, info = self.env.step(action)
         self.frames.append(obs)
         return self._get_ob(), reward, terminated, truncated, info

benchmark/plot.py

@@ -13,6 +13,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+"""Plot benchmark results."""
+
 import argparse
 
 import matplotlib.ticker as ticker
@@ -23,6 +25,7 @@
 
 
 def reset_data() -> None:
+    """Reset the benchmark data accumulators."""
     global data
     data = {
         "Num. Workers": [],
@@ -34,6 +37,7 @@ def reset_data() -> None:
 
 
 def parse_table(env: str, system: str, suffix: str) -> None:
+    """Parse a benchmark table from a markdown report."""
     private_copy = {
         "Num. Workers": [],
         "FPS": [],
@@ -47,15 +51,17 @@ def parse_table(env: str, system: str, suffix: str) -> None:
     for line in raw[2:]:
         line = line.split("|")[1:-1]
         method = line.pop(0).strip()
-        for w, f in zip(worker_num, line):
+        for w, f in zip(worker_num, line, strict=False):
             for d in [data, private_copy]:
                 d["Num. Workers"].append(w)
                 d["FPS"].append(None if f.strip() == "/" else float(f))
                 d["Env"].append(env)
                 d["System"].append(system)
                 d["Method"].append(method)
     d = pd.DataFrame(private_copy)
-    plot = sns.lineplot(x="Num. Workers", y="FPS", hue="Method", data=d, marker="o")
+    plot = sns.lineplot(
+        x="Num. Workers", y="FPS", hue="Method", data=d, marker="o"
+    )
     plot.xaxis.set_major_formatter(ticker.EngFormatter())
     plot.yaxis.set_major_formatter(ticker.EngFormatter())
     plot.legend(fontsize=9)
@@ -68,6 +74,7 @@ def parse_table(env: str, system: str, suffix: str) -> None:
 
 
 def benchmark(suffix: str) -> None:
+    """Generate throughput plots for benchmark results."""
     global data
     reset_data()
     for env in ["Atari", "Mujoco"]:
@@ -95,7 +102,7 @@ def mapping(x, y, **kwargs):
     g.add_legend(bbox_to_anchor=(0.52, 1.02), ncol=6)
     axes = g.axes.flatten()
     alphabet = "abcdefgh"
-    for ax, i in zip(axes, alphabet):
+    for ax, i in zip(axes, alphabet, strict=False):
         env, system = ax.get_title().split("|")
         env = env.split("=")[-1].strip()
         system = system.split("=")[-1].strip()

benchmark/test_dmc.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+"""Benchmark EnvPool against DeepMind Control."""
+
 import argparse
 import time
 
@@ -23,6 +25,7 @@
 
 
 def run_dmc(env, action, frame_skip, total_step):
+    """Benchmark a DeepMind Control environment."""
     ts = env.reset()
     t = time.time()
     for i in tqdm.trange(total_step):
@@ -36,6 +39,7 @@ def run_dmc(env, action, frame_skip, total_step):
 
 
 def run_envpool(env, action, frame_skip, total_step):
+    """Benchmark an EnvPool DeepMind Control environment."""
     ts = env.reset()
     t = time.time()
     for i in tqdm.trange(total_step):
@@ -61,9 +65,10 @@ def run_envpool(env, action, frame_skip, total_step):
     env = suite.load(args.domain, args.task, {"random": args.seed})
     np.random.seed(args.seed)
     minimum, maximum = env.action_spec().minimum, env.action_spec().maximum
-    action = np.array(
-        [np.random.uniform(low=minimum, high=maximum) for _ in range(args.total_step)]
-    )
+    action = np.array([
+        np.random.uniform(low=minimum, high=maximum)
+        for _ in range(args.total_step)
+    ])
     frame_skip = env._n_sub_steps
 
     fps_dmc = run_dmc(env, action, frame_skip, args.total_step)

benchmark/test_envpool.py

@@ -72,7 +72,9 @@
     if args.batch_size is None:
         args.batch_size = min(248, args.num_envs)
     elif args.batch_size > args.num_envs:
-        raise ValueError("--batch-size must be less than or equal to --num-envs")
+        raise ValueError(
+            "--batch-size must be less than or equal to --num-envs"
+        )
     print(args)
     task_id = {
         "atari": "Pong-v5",
@@ -91,7 +93,9 @@
     env = envpool.make_gym(task_id, **kwargs)
     env.async_reset()
     env.action_space.seed(args.seed)
-    action = np.array([env.action_space.sample() for _ in range(args.batch_size)])
+    action = np.array([
+        env.action_space.sample() for _ in range(args.batch_size)
+    ])
     t = time.time()
     for _ in tqdm.trange(args.total_step):
         info = env.recv()[-1]

benchmark/test_gym.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+"""Benchmark EnvPool against Gym vector environments."""
+
 import argparse
 import time
 
@@ -22,6 +24,7 @@
 
 
 def make_vector_env(num_envs, async_, make_env):
+    """Create a Gym vector environment."""
     if async_:
         vector_env_cls = gym.vector.AsyncVectorEnv
     else:
@@ -30,6 +33,7 @@ def make_vector_env(num_envs, async_, make_env):
 
 
 def run(env, num_envs, total_step, async_):
+    """Benchmark a vectorized environment."""
     if env == "atari":
         gym.register_envs(ale_py)
         task_id = "ALE/Pong-v5"
@@ -88,7 +92,9 @@ def run(env, num_envs, total_step, async_):
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
-    parser.add_argument("--env", type=str, default="atari", choices=["atari", "mujoco", "box2d"])
+    parser.add_argument(
+        "--env", type=str, default="atari", choices=["atari", "mujoco", "box2d"]
+    )
     parser.add_argument("--async_", action="store_true")
     parser.add_argument("--num-envs", type=int, default=10)
     parser.add_argument("--total-step", type=int, default=5000)

docs/conf.py

@@ -4,6 +4,8 @@
 # list see the documentation:
 # https://www.sphinx-doc.org/en/master/usage/configuration.html
 
+"""Sphinx configuration for the EnvPool docs."""
+
 # -- Path setup --------------------------------------------------------------
 
 # If extensions (or modules to document with autodoc) are in another directory,
@@ -19,6 +21,7 @@
 
 def get_version() -> str:
     # https://packaging.python.org/guides/single-sourcing-package-version/
+    """Return the project version from the package metadata."""
     with open(os.path.join("..", "envpool", "__init__.py"), "r") as f:
         init = f.read().split()
     return init[init.index("__version__") + 2][1:-1]
@@ -71,6 +74,7 @@ def get_version() -> str:
 
 
 def setup(app):
+    """Register the Sphinx configuration hooks."""
     app.add_js_file("js/copybutton.js")
     app.add_css_file("css/style.css")