pre-commit fixes

vwxyzjn · Sep 17, 2024 · c25ba62 · c25ba62
1 parent 2545864
commit c25ba62
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 10 deletions.
diff --git a/cleanrl/ppo_trxl/pom_env.py b/cleanrl/ppo_trxl/pom_env.py
@@ -1,7 +1,7 @@
 import gymnasium as gym
 import numpy as np
-from gymnasium import spaces
 import pygame
+from gymnasium import spaces
 
 gym.register(
     id="ProofofMemory-v0",
@@ -142,10 +142,14 @@ def render(self):
             left_goal_color = (0, 255, 0) if self._goals[0] > 0 else (255, 0, 0)
             pygame.draw.rect(canvas, left_goal_color, pygame.Rect(0, 0, self.cell_width, self.height))
             right_goal_color = (0, 255, 0) if self._goals[1] > 0 else (255, 0, 0)
-            pygame.draw.rect(canvas, right_goal_color, pygame.Rect(self.width - self.cell_width, 0, self.cell_width, self.height))
+            pygame.draw.rect(
+                canvas, right_goal_color, pygame.Rect(self.width - self.cell_width, 0, self.cell_width, self.height)
+            )
         else:
             pygame.draw.rect(canvas, (200, 200, 200), pygame.Rect(0, 0, self.cell_width, self.height))
-            pygame.draw.rect(canvas, (200, 200, 200), pygame.Rect(self.width - self.cell_width, 0, self.cell_width, self.height))
+            pygame.draw.rect(
+                canvas, (200, 200, 200), pygame.Rect(self.width - self.cell_width, 0, self.cell_width, self.height)
+            )
 
         # Render text information
         font = pygame.font.SysFont(None, 24)
@@ -157,9 +161,7 @@ def render(self):
             pygame.display.flip()
             self.clock.tick(self.metadata["render_fps"])
         elif self.render_mode in ["rgb_array", "debug_rgb_array"]:
-            return np.transpose(
-                np.array(pygame.surfarray.pixels3d(canvas)), axes=(1, 0, 2)
-            )
+            return np.transpose(np.array(pygame.surfarray.pixels3d(canvas)), axes=(1, 0, 2))
 
     def close(self):
         if self.window is not None:

diff --git a/cleanrl/ppo_trxl/ppo_trxl.py b/cleanrl/ppo_trxl/ppo_trxl.py
@@ -104,12 +104,12 @@ class Args:
 
 def make_env(env_id, idx, capture_video, run_name, render_mode="debug_rgb_array"):
     if "MiniGrid" in env_id:
-            if render_mode == "debug_rgb_array":
-                render_mode = "rgb_array"
+        if render_mode == "debug_rgb_array":
+            render_mode = "rgb_array"
+
     def thunk():
         if "MiniGrid" in env_id:
-            env = gym.make(env_id, agent_view_size=3, tile_size=28, 
-                           render_mode="rgb_array" if render_mode == "debug_rgb_array" else render_mode)
+            env = gym.make(env_id, agent_view_size=3, tile_size=28, render_mode=render_mode)
             env = ImgObsWrapper(RGBImgPartialObsWrapper(env, tile_size=28))
             env = gym.wrappers.TimeLimit(env, 96)
         else: