USER
I have this code:
```py
import pygame
import torch
import math
from sys import argv
from random import randint
from model import model, train_step, path
WIDTH = 800
HEIGHT = 800
FPS = 0
DEBUG = True
SHOULD_LEARN = False
NOMODEL = False
BACKGROUND_COLOR = (64, 128, 64)
FOREGROUND_COLOR = (0, 192, 0)
class Direction:
LEFT = 0
RIGHT = 1
UP = 2
DOWN = 3
def reverse(dir):
match dir:
case Direction.LEFT:
return Direction.RIGHT
case Direction.UP:
return Direction.DOWN
case Direction.DOWN:
return Direction.UP
case Direction.RIGHT:
return Direction.LEFT
def move(vec, dir):
match dir:
case Direction.RIGHT:
return [vec[0] + 1, vec[1]]
case Direction.LEFT:
return [vec[0] - 1, vec[1]]
case Direction.UP:
return [vec[0], vec[1] - 1]
case Direction.DOWN:
return [vec[0], vec[1] + 1]
def print(dir):
match dir:
case Direction.RIGHT:
return "Right"
case Direction.LEFT:
return "Left"
case Direction.UP:
return "Up"
case Direction.DOWN:
return "Down"
class Field:
width = WIDTH - 100
height = HEIGHT - 100
x = int((WIDTH - width) / 2)
y = int((HEIGHT - height) / 2)
cell_count = 20
cell_size = (width / cell_count, height / cell_count)
head_color = (0, 255, 255)
body_color = (0, 200, 200)
apple = [6, 6]
def __init__(self):
Field.apple = [randint(0, Field.cell_count - 1), randint(0, Field.cell_count - 1)]
self.snake = Field.Snake(Direction.RIGHT)
### Init
def move(self):
val = self.snake.move()
return val
### Move
def draw(self, screen):
pygame.draw.rect(screen, (0, 160, 0), (Field.x, Field.y, Field.width, Field.height))
self.snake.draw(screen)
pygame.draw.rect(screen, (255, 0, 0), (Field.x + Field.apple[0] * Field.cell_size[0], Field.y + Field.apple[1]* Field.cell_size[1], Field.cell_size[0], Field.cell_size[1]))
return
### Draw
class Snake:
class Body:
def __init__(self, x, y):
self.x = x
self.y = y
### Init
### Body
def __init__(self, direction):
self.body = []
self.direction = direction
base = [5, 5]
self.body.append(Field.Snake.Body(base[0], base[1]))
base = Direction.move(base, Direction.reverse(direction))
self.body.append(Field.Snake.Body(base[0], base[1]))
base = Direction.move(base, Direction.reverse(direction))
self.body.append(Field.Snake.Body(base[0], base[1]))
self.queued_move = -1
self.should_move = True
### Init
def draw(self, screen):
screen.fill(Field.head_color, (Field.x + self.body[0].x * Field.cell_size[0], Field.y + self.body[0].y * Field.cell_size[1], Field.cell_size[0], Field.cell_size[1]))
texture=pygame.image.load("pic.jpg")
screen.blit(texture, (Field.x + self.body[0].x * Field.cell_size[0], Field.y + self.body[0].y * Field.cell_size[1], Field.cell_size[0], Field.cell_size[1]))
for i in range(1, len(self.body)):
screen.fill(Field.body_color, (Field.x + self.body[i].x * Field.cell_size[0], Field.y + self.body[i].y * Field.cell_size[1], Field.cell_size[0], Field.cell_size[1]))
### draw
def move(self):
if self.queued_move != -1:
self.direction = self.queued_move
self.queued_move = -1
self.should_move = True
for i in range(0, len(self.body) - 1):
self.body[len(self.body) - i - 1].y = self.body[len(self.body) - i - 1 - 1].y
self.body[len(self.body) - i - 1].x = self.body[len(self.body) - i - 1 - 1].x
if self.body[i].x < 0 or self.body[i].y < 0 or self.body[i].x == Field.cell_count or self.body[i].y == Field.cell_count:
return -1
moved = Direction.move([self.body[0].x, self.body[1].y], self.direction)
self.body[0].x = moved[0]
self.body[0].y = moved[1]
if self.body[0].x < 0 or self.body[0].y < 0 or self.body[0].x == Field.cell_count or self.body[0].y == Field.cell_count:
return -1
for i in range(1, len(self.body)):
if self.body[i].x == self.body[0].x and self.body[i].y == self.body[0].y:
return -2
if self.body[0].x == Field.apple[0] and Field.apple[1] == self.body[0].y:
self.add_body()
Field.apple = [randint(0, Field.cell_count - 1), randint(0, Field.cell_count - 1)]
return 1
return 0
### move
def add_body(self):
pos = Direction.move([self.body[len(self.body) - 1].x, self.body[len(self.body) - 1].y], Direction.reverse(self.direction))
self.body.append(Field.Snake.Body(pos[0], pos[1]))
### add body
def change_direction(self, direction):
if self.should_move:
self.queued_move = direction
self.should_move = False
### change_direction
def get_direction(self):
return self.direction
### change_direction
def cast_rays(self):
res = []
res.append(self.cast_to(Direction.LEFT))
res.append(self.cast_to(Direction.RIGHT))
res.append(self.cast_to(Direction.UP))
res.append(self.cast_to(Direction.DOWN))
return res
def cast_to(self, direction):
res = 0
startpos = [ self.body[0].x, self.body[0].y]
while (startpos[0] < Field.cell_count and startpos[0] >= 0 and startpos[1] >= 0 and startpos[1] < Field.cell_count):
for i in range(1, len(self.body)):
if self.body[i].x == startpos[0] and self.body[i].y == startpos[1]:
return res
startpos = Direction.move(startpos, direction)
res += 1
return res
### cast
### Snake
### Field
def distance_head_apple(field):
return math.sqrt((Field.apple[0] - field.snake.body[0].x)**2 + (Field.apple[1] - field.snake.body[0].y)**2)
def main():
pygame.init()
screen = pygame.display.set_mode((WIDTH, HEIGHT))
clock = pygame.time.Clock()
SHOULD_LEARN = False
BASE_SPEED = 0
NOMODEL = False
print(argv)
for i in range(1, len(argv)):
match argv[i]:
case "--train":
SHOULD_LEARN = True
print("training flag")
case "--speed":
try:
BASE_SPEED = int(argv[i + 1])
print(f"speed flag. speed = {BASE_SPEED}")
except IndexError:
print("Did not recognize the speed")
case "--nomodel":
NOMODEL = True
case _:
print("Еблан???")
exit(-1)
field = Field()
counter = 0
speed = BASE_SPEED
if not NOMODEL:
try:
model.load_state_dict(torch.load(path))
except FileNotFoundError:
pass
rays = field.snake.cast_rays()
data = [Field.apple[0], Field.apple[1], field.snake.body[0].x, field.snake.body[0].y, rays[0], rays[1], rays[2], rays[3], Field.cell_count, Field.cell_count, len(field.snake.body), field.snake.direction, distance_head_apple(field)]
data = torch.tensor(data, dtype=torch.float32)
data = model(data)
print(data)
print(torch.argmax(data))
record = 0
collected = 0
death_counter = 0
death_from_body_collide_counter = 0
running = True
while running:
if speed == 1000000:
field = None
field = Field()
speed = BASE_SPEED
counter = 0
continue
got_apple = 0
died = 0
clock.tick(FPS)
for event in pygame.event.get():
if event.type == pygame.QUIT:
running = False
if event.type == pygame.KEYDOWN:
match event.key:
case pygame.K_RIGHT:
if field.snake.get_direction() != Direction.reverse(Direction.RIGHT):
field.snake.change_direction(Direction.RIGHT)
case pygame.K_LEFT:
if field.snake.get_direction() != Direction.reverse(Direction.LEFT):
field.snake.change_direction(Direction.LEFT)
case pygame.K_UP:
if field.snake.get_direction() != Direction.reverse(Direction.UP):
field.snake.change_direction(Direction.UP)
case pygame.K_DOWN:
if field.snake.get_direction() != Direction.reverse(Direction.DOWN):
field.snake.change_direction(Direction.DOWN)
case pygame.K_ESCAPE:
running = False
case pygame.K_SPACE:
field.snake.add_body()
case pygame.K_LCTRL:
print(field.snake.cast_rays())
if counter == speed:
if not NOMODEL:
distance_before_move = distance_head_apple(field)
rays = field.snake.cast_rays()
data = [Field.apple[0], Field.apple[1], field.snake.body[0].x, field.snake.body[0].y, rays[0], rays[1], rays[2], rays[3], Field.cell_count, Field.cell_count, len(field.snake.body), field.snake.direction, distance_before_move]
state = data
state = torch.tensor(state, dtype=torch.float32)
result = model(state)
#print(result)
move = torch.argmax(result)
match move:
case Direction.LEFT:
if field.snake.get_direction() != Direction.reverse(Direction.LEFT):
field.snake.change_direction(Direction.LEFT)
case Direction.RIGHT:
if field.snake.get_direction() != Direction.reverse(Direction.RIGHT):
field.snake.change_direction(Direction.RIGHT)
case Direction.UP:
if field.snake.get_direction() != Direction.reverse(Direction.UP):
field.snake.change_direction(Direction.UP)
case Direction.DOWN:
if field.snake.get_direction() != Direction.reverse(Direction.DOWN):
field.snake.change_direction(Direction.DOWN)
print(f"Record: {record}\t Deaths: {death_counter}\tDeath from self: {death_from_body_collide_counter}\t Current: {collected} \t argmax = {torch.argmax(result)} \t maxval = {torch.max(result)}")
res = field.move()
counter = 0
if res == -1:
speed = 1000000
counter = 0
died = -0.05
collected = 0
death_counter += 1
elif res == -2:
speed = 1000000
counter = 0
died = -0.05
collected = 0
death_counter += 1
death_from_body_collide_counter += 1
elif res == 1:
got_apple = 1
collected += 1
if SHOULD_LEARN and not NOMODEL:
distance_after_move = distance_head_apple(field)
rays = field.snake.cast_rays()
next_data = [Field.apple[0], Field.apple[1], field.snake.body[0].x, field.snake.body[0].y, rays[0], rays[1], rays[2], rays[3], Field.cell_count, Field.cell_count, len(field.snake.body), field.snake.direction, distance_after_move]
next_state = torch.tensor(next_data, dtype=torch.float32)
#print(f'Before = {distance_before_move} ; After = {distance_after_move}')
if got_apple == 1:
train_step(state, move, 1, next_state, True)
got_apple = 0
elif died != 0:
train_step(state, move, -1, next_state, False)
elif math.fabs(distance_before_move) < math.fabs(distance_after_move):
train_step(state, move, -0.1, next_state, False)
elif math.fabs(distance_before_move) > math.fabs(distance_after_move):
train_step(state, move, 0.1, next_state, False)
else:
train_step(state, move, 0.00001, next_state, False)
else:
counter += 1
#
if collected >= record:
record = collected
screen.fill(BACKGROUND_COLOR)
field.draw(screen)
pygame.display.flip()
pygame.quit()
if not NOMODEL:
torch.save(model.state_dict(), path)
if __name__ == "__main__":
main()
```
Make method in Field class, that will compose all field into list[list[int]] cell_count*cell_count where 0 - empty cell of field, 1 - body of snake, 2 - head of snake, 3 - appleASSISTANT
Certainly! You can add a method called `get_grid` to the `Field` class that generates a 2D list representing the current state of the game field. Here's how you can implement it:
### Step-by-Step Implementation:
1. **Initialize the Grid:**
- Create a 2D list filled with `0`s representing empty cells.
2. **Place the Apple:**
- Set the cell corresponding to the apple's position to `3`.
3. **Place the Snake:**
- Iterate through the snake's body segments.
- Set the head of the snake to `2`.
- Set the rest of the snake's body to `1`.
4. **Return the Grid:**
- The method returns the completed grid.
### Updated `Field` Class with `get_grid` Method:
```python
class Field:
width = WIDTH - 100
height = HEIGHT - 100
x = int((WIDTH - width) / 2)
y = int((HEIGHT - height) / 2)
cell_count = 20
cell_size = (width / cell_count, height / cell_count)
head_color = (0, 255, 255)
body_color = (0, 200, 200)
apple = [6, 6]
def __init__(self):
Field.apple = [randint(0, Field.cell_count - 1), randint(0, Field.cell_count - 1)]
self.snake = Field.Snake(Direction.RIGHT)
### Init
def move(self):
val = self.snake.move()
return val
### Move
def draw(self, screen):
pygame.draw.rect(screen, (0, 160, 0), (Field.x, Field.y, Field.width, Field.height))
self.snake.draw(screen)
pygame.draw.rect(screen, (255, 0, 0), (Field.x + Field.apple[0] * Field.cell_size[0], Field.y + Field.apple[1]* Field.cell_size[1], Field.cell_size[0], Field.cell_size[1]))
return
### Draw
def get_grid(self) -> list[list[int]]:
"""
Returns a 2D list representing the game field.
0 - empty cell
1 - snake body
2 - snake head
3 - apple
"""
# Initialize grid with 0s
grid = [[0 for _ in range(Field.cell_count)] for _ in range(Field.cell_count)]
# Place the apple
apple_x, apple_y = Field.apple
if 0 <= apple_x < Field.cell_count and 0 <= apple_y < Field.cell_count:
grid[apple_y][apple_x] = 3 # Note: grid[y][x]
# Place the snake
for idx, segment in enumerate(self.snake.body):
x, y = segment.x, segment.y
if 0 <= x < Field.cell_count and 0 <= y < Field.cell_count:
if idx == 0:
grid[y][x] = 2 # Head
else:
grid[y][x] = 1 # Body
return grid
### Get Grid
class Snake:
class Body:
def __init__(self, x, y):
self.x = x
self.y = y
### Init
### Body
def __init__(self, direction):
self.body = []
self.direction = direction
base = [5, 5]
self.body.append(Field.Snake.Body(base[0], base[1]))
base = Direction.move(base, Direction.reverse(direction))
self.body.append(Field.Snake.Body(base[0], base[1]))
base = Direction.move(base, Direction.reverse(direction))
self.body.append(Field.Snake.Body(base[0], base[1]))
self.queued_move = -1
self.should_move = True
### Init
def draw(self, screen):
screen.fill(Field.head_color, (Field.x + self.body[0].x * Field.cell_size[0], Field.y + self.body[0].y * Field.cell_size[1], Field.cell_size[0], Field.cell_size[1]))
texture=pygame.image.load("pic.jpg")
screen.blit(texture, (Field.x + self.body[0].x * Field.cell_size[0], Field.y + self.body[0].y * Field.cell_size[1], Field.cell_size[0], Field.cell_size[1]))
for i in range(1, len(self.body)):
screen.fill(Field.body_color, (Field.x + self.body[i].x * Field.cell_size[0], Field.y + self.body[i].y * Field.cell_size[1], Field.cell_size[0], Field.cell_size[1]))
### draw
def move(self):
if self.queued_move != -1:
self.direction = self.queued_move
self.queued_move = -1
self.should_move = True
for i in range(0, len(self.body) - 1):
self.body[len(self.body) - i - 1].y = self.body[len(self.body) - i - 1 - 1].y
self.body[len(self.body) - i - 1].x = self.body[len(self.body) - i - 1 - 1].x
if self.body[i].x < 0 or self.body[i].y < 0 or self.body[i].x == Field.cell_count or self.body[i].y == Field.cell_count:
return -1
moved = Direction.move([self.body[0].x, self.body[1].y], self.direction)
self.body[0].x = moved[0]
self.body[0].y = moved[1]
if self.body[0].x < 0 or self.body[0].y < 0 or self.body[0].x == Field.cell_count or self.body[0].y == Field.cell_count:
return -1
for i in range(1, len(self.body)):
if self.body[i].x == self.body[0].x and self.body[i].y == self.body[0].y:
return -2
if self.body[0].x == Field.apple[0] and Field.apple[1] == self.body[0].y:
self.add_body()
Field.apple = [randint(0, Field.cell_count - 1), randint(0, Field.cell_count - 1)]
return 1
return 0
### move
def add_body(self):
pos = Direction.move([self.body[len(self.body) - 1].x, self.body[len(self.body) - 1].y], Direction.reverse(self.direction))
self.body.append(Field.Snake.Body(pos[0], pos[1]))
### add body
def change_direction(self, direction):
if self.should_move:
self.queued_move = direction
self.should_move = False
### change_direction
def get_direction(self):
return self.direction
### get_direction
def cast_rays(self):
res = []
res.append(self.cast_to(Direction.LEFT))
res.append(self.cast_to(Direction.RIGHT))
res.append(self.cast_to(Direction.UP))
res.append(self.cast_to(Direction.DOWN))
return res
def cast_to(self, direction):
res = 0
startpos = [ self.body[0].x, self.body[0].y]
while (startpos[0] < Field.cell_count and startpos[0] >= 0 and startpos[1] >= 0 and startpos[1] < Field.cell_count):
for i in range(1, len(self.body)):
if self.body[i].x == startpos[0] and self.body[i].y == startpos[1]:
return res
startpos = Direction.move(startpos, direction)
res += 1
return res
### cast
### Snake
### Field
```
### Explanation:
1. **Grid Initialization:**
```python
grid = [[0 for _ in range(Field.cell_count)] for _ in range(Field.cell_count)]
```
- This creates a 2D list filled with `0`s, representing an empty field.
2. **Placing the Apple:**
```python
apple_x, apple_y = Field.apple
if 0 <= apple_x < Field.cell_count and 0 <= apple_y < Field.cell_count:
grid[apple_y][apple_x] = 3
```
- The apple's position is retrieved and set to `3` in the grid.
- Note that `grid[y][x]` is used because the first index represents the row (y-axis) and the second index represents the column (x-axis).
3. **Placing the Snake:**
```python
for idx, segment in enumerate(self.snake.body):
x, y = segment.x, segment.y
if 0 <= x < Field.cell_count and 0 <= y < Field.cell_count:
if idx == 0:
grid[y][x] = 2 # Head
else:
grid[y][x] = 1 # Body
```
- Iterates through each segment of the snake's body.
- The head (first segment) is marked as `2`, and the rest of the body is marked as `1`.
4. **Returning the Grid:**
```python
return grid
```
- The method returns the populated grid.
### Usage Example:
You can use the `get_grid` method whenever you need a representation of the current state of the game field. For example, you might use it for debugging, logging, or integrating with other systems (like AI agents).
```python
def main():
# ... existing code ...
field = Field()
# Example usage of get_grid
current_grid = field.get_grid()
for row in current_grid:
print(row) # This will print the grid to the console
# ... existing game loop ...
```
### Additional Considerations:
- **Boundary Checks:**
- The method includes boundary checks to ensure that only valid positions within the grid are marked. This prevents potential `IndexError` exceptions.
- **Grid Representation:**
- Depending on your needs, you might want to represent the grid differently (e.g., using NumPy arrays for efficiency). However, for simplicity and compatibility, a list of lists is used here.
- **Performance:**
- If your grid size is large or if you call `get_grid` frequently, you might need to optimize the method. For most standard Snake game implementations, this should be sufficiently performant.
### Complete Integration:
Ensure that you place the `get_grid` method within the `Field` class as shown above. This method seamlessly integrates with your existing `Field` and `Snake` classes, providing a clear representation of the game state.
Let me know if you need further assistance or additional features!