Update tetris_env.py

tetris_env.py

@@ -99,42 +99,55 @@ class TetrisEnv(gym.Env):
         return np.array(self.game.field)
 
     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)
+    """
+    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
+        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:
+        # 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.field[i][j]],
+                                         colors[self.game.figure.color],
                                          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))
+        # 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)
 
-            # Convert Pygame surface to RGB array
-            return pygame.surfarray.array3d(self.screen)
 
         elif mode == 'human':
             if self.screen is None: