yuh

tetris_env.py

@@ -91,45 +91,41 @@ class TetrisEnv(gym.Env):
         return self._get_obs(), reward, done, {}
 
     def _get_obs(self):
-        if self.game.field is None:
-            raise ValueError("The field attribute in self.game is None.")
         """
         Get the current state of the game as an observation
         """
+        if self.game.field is None:
+            raise ValueError("The field attribute in self.game is None.")
         return np.array(self.game.field)
 
     def render(self, mode='human'):
-    """
-    Render the game state as an RGB array or display it.
+        """
+        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
@@ -137,51 +133,46 @@ 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))
-    
+
             # 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:
+
+        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.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))
-    
-        pygame.display.flip()
+
+            pygame.display.flip()
 
     def close(self):
         """