Update tetris_env.py

tetris_env.py

@@ -101,75 +101,35 @@ class TetrisEnv(gym.Env):
     def render(self, mode='human'):
     """
     Render the game state as an RGB array or display it.
-    """
-    if mode == 'rgb_array':
-        if self.screen is None:
-            pygame.init()
-            size = (self.x * 2 + self.zoom * self.width, self.y * 2 + self.zoom * self.height)
-            self.screen = pygame.Surface(size)
-
-        self.screen.fill((173, 216, 230))  # WHITE background
-
-        # Check if the game field is initialized
-        if self.game.field is None:
-            raise ValueError("Game field is None.")
-
-        # Draw the game field
-        for i in range(self.game.height):
-            for j in range(self.game.width):
-                rect = pygame.Rect(self.x + self.zoom * j, self.y + self.zoom * i, self.zoom, self.zoom)
-                pygame.draw.rect(self.screen, (128, 128, 128), rect, 1)  # Grid lines
-                if self.game.field[i][j] > 0:
-                    if self.game.field[i][j] not in colors:
-                        raise ValueError(f"Color for field value {self.game.field[i][j]} is not defined.")
-                    pygame.draw.rect(self.screen, 
-                                     colors[self.game.field[i][j]],
-                                     rect.inflate(-2, -2))
-
-        # Draw the current figure
-        if self.game.figure is not None:
-            if self.game.figure.image() is None:
-                raise ValueError("Figure's image is None.")
-            for i in range(4):
-                for j in range(4):
-                    p = i * 4 + j
-                    if p in self.game.figure.image():
-                        rect = pygame.Rect(self.x + self.zoom * (j + self.game.figure.x),
-                                           self.y + self.zoom * (i + self.game.figure.y),
-                                           self.zoom, self.zoom)
-                        if self.game.figure.color not in colors:
-                            raise ValueError(f"Color for figure value {self.game.figure.color} is not defined.")
-                        pygame.draw.rect(self.screen, 
-                                         colors[self.game.figure.color],
-                                         rect.inflate(-2, -2))
-
-        # Convert Pygame surface to RGB array
-        if self.screen is None:
-            raise ValueError("Screen is None when trying to create RGB array.")
-        return pygame.surfarray.array3d(self.screen)
-
-
-        elif mode == 'human':
+        """
+        if mode == 'rgb_array':
             if self.screen is None:
                 pygame.init()
                 size = (self.x * 2 + self.zoom * self.width, self.y * 2 + self.zoom * self.height)
-                self.screen = pygame.display.set_mode(size)
-                pygame.display.set_caption("Tetris RL")
-
+                self.screen = pygame.Surface(size)
+    
             self.screen.fill((173, 216, 230))  # WHITE background
-
+    
+            # Check if the game field is initialized
+            if self.game.field is None:
+                raise ValueError("Game field is None.")
+    
             # Draw the game field
             for i in range(self.game.height):
                 for j in range(self.game.width):
                     rect = pygame.Rect(self.x + self.zoom * j, self.y + self.zoom * i, self.zoom, self.zoom)
                     pygame.draw.rect(self.screen, (128, 128, 128), rect, 1)  # Grid lines
                     if self.game.field[i][j] > 0:
+                        if self.game.field[i][j] not in colors:
+                            raise ValueError(f"Color for field value {self.game.field[i][j]} is not defined.")
                         pygame.draw.rect(self.screen, 
                                          colors[self.game.field[i][j]],
                                          rect.inflate(-2, -2))
-
+    
             # Draw the current figure
             if self.game.figure is not None:
+                if self.game.figure.image() is None:
+                    raise ValueError("Figure's image is None.")
                 for i in range(4):
                     for j in range(4):
                         p = i * 4 + j
@@ -177,11 +137,51 @@ class TetrisEnv(gym.Env):
                             rect = pygame.Rect(self.x + self.zoom * (j + self.game.figure.x),
                                                self.y + self.zoom * (i + self.game.figure.y),
                                                self.zoom, self.zoom)
+                            if self.game.figure.color not in colors:
+                                raise ValueError(f"Color for figure value {self.game.figure.color} is not defined.")
                             pygame.draw.rect(self.screen, 
                                              colors[self.game.figure.color],
                                              rect.inflate(-2, -2))
-
-            pygame.display.flip()
+    
+            # Convert Pygame surface to RGB array
+            if self.screen is None:
+                raise ValueError("Screen is None when trying to create RGB array.")
+            return pygame.surfarray.array3d(self.screen)
+    
+    
+            elif mode == 'human':
+                if self.screen is None:
+                    pygame.init()
+                    size = (self.x * 2 + self.zoom * self.width, self.y * 2 + self.zoom * self.height)
+                    self.screen = pygame.display.set_mode(size)
+                    pygame.display.set_caption("Tetris RL")
+    
+                self.screen.fill((173, 216, 230))  # WHITE background
+    
+                # Draw the game field
+                for i in range(self.game.height):
+                    for j in range(self.game.width):
+                        rect = pygame.Rect(self.x + self.zoom * j, self.y + self.zoom * i, self.zoom, self.zoom)
+                        pygame.draw.rect(self.screen, (128, 128, 128), rect, 1)  # Grid lines
+                        if self.game.field[i][j] > 0:
+                            pygame.draw.rect(self.screen, 
+                                             colors[self.game.field[i][j]],
+                                             rect.inflate(-2, -2))
+    
+                # Draw the current figure
+                if self.game.figure is not None:
+                    for i in range(4):
+                        for j in range(4):
+                            p = i * 4 + j
+                            if p in self.game.figure.image():
+                                rect = pygame.Rect(self.x + self.zoom * (j + self.game.figure.x),
+                                                   self.y + self.zoom * (i + self.game.figure.y),
+                                                   self.zoom, self.zoom)
+                                pygame.draw.rect(self.screen, 
+                                                 colors[self.game.figure.color],
+                                                 rect.inflate(-2, -2))
+    
+        pygame.display.flip()
 
     def close(self):
         """