viberl.envs.grid_world.snake_env

Classes:

Name	Description
Direction
SnakeGameEnv

Direction

Bases: Enum

Attributes:

Name	Type	Description
UP
RIGHT
DOWN
LEFT

UP class-attribute instance-attribute

UP = 0

RIGHT class-attribute instance-attribute

RIGHT = 1

DOWN class-attribute instance-attribute

DOWN = 2

LEFT class-attribute instance-attribute

LEFT = 3

SnakeGameEnv

SnakeGameEnv(render_mode: str | None = None, grid_size: int = 20)

Bases: Env

Methods:

Name	Description
reset
step
render
close

Attributes:

Name	Type	Description
metadata	dict[str, int | list[str]]
grid_size
render_mode
action_space
observation_space
window
clock
cell_size
window_size

Source code in viberl/envs/grid_world/snake_env.py

def __init__(self, render_mode: str | None = None, grid_size: int = 20):
    super().__init__()

    self.grid_size = grid_size
    self.render_mode = render_mode

    # Action space: 0=UP, 1=RIGHT, 2=DOWN, 3=LEFT
    self.action_space = spaces.Discrete(4)

    # Observation space: grid with snake body, food, and empty spaces
    self.observation_space = spaces.Box(
        low=0, high=3, shape=(grid_size, grid_size), dtype=np.uint8
    )

    # Initialize game state
    self.reset()

    # Pygame setup for rendering
    self.window = None
    self.clock = None
    self.cell_size = 20
    self.window_size = self.grid_size * self.cell_size

metadata class-attribute

metadata: dict[str, int | list[str]] = {'render_modes': ['human', 'rgb_array'], 'render_fps': 10}

grid_size instance-attribute

grid_size = grid_size

render_mode instance-attribute

render_mode = render_mode

action_space instance-attribute

action_space = Discrete(4)

observation_space instance-attribute

observation_space = Box(low=0, high=3, shape=(grid_size, grid_size), dtype=uint8)

window instance-attribute

window = None

clock instance-attribute

clock = None

cell_size instance-attribute

cell_size = 20

window_size instance-attribute

window_size = grid_size * cell_size

reset

reset(seed: int | None = None, options: dict | None = None) -> tuple[ndarray, dict[str, Any]]

Source code in viberl/envs/grid_world/snake_env.py

def reset(
    self, seed: int | None = None, options: dict | None = None
) -> tuple[np.ndarray, dict[str, Any]]:
    super().reset(seed=seed)

    if seed is not None:
        np.random.seed(seed)

    # Initialize snake at center of grid
    center = self.grid_size // 2
    self.snake = [(center - 1, center), (center, center), (center + 1, center)]
    self.direction = Direction.RIGHT

    # Place food
    self.food = self._place_food()

    # Game state
    self.game_over = False
    self.score = 0
    self.steps = 0
    self.max_steps = self.grid_size * self.grid_size * 4

    return self._get_observation(), self._get_info()

step

step(action: int) -> tuple[ndarray, float, bool, bool, dict[str, Any]]

Source code in viberl/envs/grid_world/snake_env.py

def step(self, action: int) -> tuple[np.ndarray, float, bool, bool, dict[str, Any]]:
    if self.game_over:
        return self._get_observation(), 0.0, True, False, self._get_info()

    self.steps += 1

    # Convert action to direction (ensure snake doesn't reverse into itself)
    new_direction = Direction(action)

    # Prevent moving directly opposite to current direction
    if (
        (new_direction == Direction.UP and self.direction == Direction.DOWN)
        or (new_direction == Direction.DOWN and self.direction == Direction.UP)
        or (new_direction == Direction.LEFT and self.direction == Direction.RIGHT)
        or (new_direction == Direction.RIGHT and self.direction == Direction.LEFT)
    ):
        new_direction = self.direction

    self.direction = new_direction

    # Move snake
    head_x, head_y = self.snake[-1]

    if self.direction == Direction.UP:
        new_head = (head_x - 1, head_y)
    elif self.direction == Direction.RIGHT:
        new_head = (head_x, head_y + 1)
    elif self.direction == Direction.DOWN:
        new_head = (head_x + 1, head_y)
    elif self.direction == Direction.LEFT:
        new_head = (head_x, head_y - 1)

    # Check collision with walls
    if (
        new_head[0] < 0
        or new_head[0] >= self.grid_size
        or new_head[1] < 0
        or new_head[1] >= self.grid_size
    ):
        self.game_over = True
        return self._get_observation(), -10.0, True, False, self._get_info()

    # Check collision with self
    if new_head in self.snake[:-1]:
        self.game_over = True
        return self._get_observation(), -10.0, True, False, self._get_info()

    # Move snake
    self.snake.append(new_head)

    reward = 0.0

    # Check if food eaten
    if new_head == self.food:
        self.score += 1
        reward = 10.0
        self.food = self._place_food()
    else:
        # Remove tail if no food eaten
        self.snake.pop(0)

    # Give negative reward to encourage faster completion
    reward += -0.1

    # Check if maximum steps reached
    if self.steps >= self.max_steps:
        self.game_over = True
        return self._get_observation(), reward, True, False, self._get_info()

    return self._get_observation(), reward, False, False, self._get_info()

render

render() -> None | ndarray

Source code in viberl/envs/grid_world/snake_env.py

def render(self) -> None | np.ndarray:
    if self.render_mode is None:
        return None

    if self.render_mode == 'rgb_array':
        return self._render_frame()
    if self.render_mode == 'human':
        self._render_frame()

close

close()

Source code in viberl/envs/grid_world/snake_env.py

def close(self):
    if self.window is not None:
        pygame.display.quit()
        pygame.quit()