Плохо работает нейросеть

import pygame
import sys
import random
import numpy as np
import tensorflow as tf
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense
from tensorflow.keras.layers import Conv2D, Flatten
 
# Глобальные переменные
WIDTH = 480
HEIGHT = 480
BLOCK_SIZE = 20
FPS = 10  # Замедление игры
BUFFER_SIZE = 10000
UPDATE_FREQ = 10
 
# Определение цветов
WHITE = (255, 255, 255)
GREEN = (0, 128, 0)
RED = (255, 0, 0)
 
 
class Snake:
    def __init__(self):
        self.positions = [(WIDTH // 2, HEIGHT // 2)]
        self.direction = (0, -BLOCK_SIZE)
 
    def move(self):
        head_x, head_y = self.positions[0]
        dir_x, dir_y = self.direction
        new_position = (head_x + dir_x, head_y + dir_y)
        self.positions.insert(0, new_position)
        self.positions.pop()
 
    def change_direction(self, direction):
        if self.direction[0] == -direction[0] and self.direction[1] == -direction[1]:
            return
        self.direction = direction
 
    def grow(self):
        last_x, last_y = self.positions[-1]
        dx, dy = self.direction
        new_position = (last_x + dx, last_y + dy)
        self.positions.append(new_position)
 
    def collided_with_wall(self):
        head_x, head_y = self.positions[0]
        return head_x < 0 or head_x >= WIDTH or head_y < 0 or head_y >= HEIGHT    
 
    def collided_with_itself(self):
        return self.positions[0] in self.positions[1:]
 
    def draw(self, screen):
        for position in self.positions:
            pygame.draw.rect(screen, GREEN, (*position, BLOCK_SIZE, BLOCK_SIZE))
 
class Food:
    def __init__(self):
        self.position = self.generate_random_position()
 
    def generate_random_position(self):
        return (random.randint(0, (WIDTH - BLOCK_SIZE) // BLOCK_SIZE) * BLOCK_SIZE,
                random.randint(0, (HEIGHT - BLOCK_SIZE) // BLOCK_SIZE) * BLOCK_SIZE)
 
    def regenerate(self):
        self.position = self.generate_random_position()
 
    def draw(self, screen):
        pygame.draw.rect(screen, RED, (*self.position, BLOCK_SIZE, BLOCK_SIZE))
 
class ReplayBuffer:
    def __init__(self, buffer_size):
        self.buffer = []
        self.buffer_size = buffer_size
 
    def add(self, state, action, reward, next_state, done):
        if len(self.buffer) >= self.buffer_size:
            self.buffer.pop(0)
        self.buffer.append((state, action, reward, next_state, done))
 
    def sample(self, batch_size):
        return random.sample(self.buffer, batch_size)
        
 
def get_game_state(snake, food):
    food_x, food_y = food.position
    head_x, head_y = snake.positions[0]
    apple_x_rel, apple_y_rel = food_x - head_x, food_y - head_y
    dir_x, dir_y = snake.direction
 
    if dir_y == -BLOCK_SIZE:  # движение вверх
        apple_x_rel, apple_y_rel = -apple_y_rel, apple_x_rel
    elif dir_y == BLOCK_SIZE:  # движение вниз
        apple_x_rel, apple_y_rel = apple_y_rel, -apple_x_rel
    elif dir_x == -BLOCK_SIZE:  # движение влево
        apple_x_rel, apple_y_rel = -apple_x_rel, -apple_y_rel
 
 
    dir_idx = {(-BLOCK_SIZE, 0): 0, (BLOCK_SIZE, 0): 1, (0, -BLOCK_SIZE): 2, (0, BLOCK_SIZE): 3}[snake.direction]
 
    wall_distances = []
 
    for direction in [(0, -BLOCK_SIZE), (0, BLOCK_SIZE), (-BLOCK_SIZE, 0), (BLOCK_SIZE, 0)]:
        distance = 0
        temp_x, temp_y = head_x + direction[0], head_y + direction[1]
        while not (temp_x < 0 or temp_x >= WIDTH or temp_y < 0 or temp_y >= HEIGHT):
            distance += 1
            if (temp_x, temp_y) in snake.positions:
                break
            temp_x += direction[0]
            temp_y += direction[1]
 
        wall_distances.append(distance)
 
    state = np.zeros((5,))
    state[0] = wall_distances[(dir_idx - 1) % 4]
    state[1] = wall_distances[dir_idx]
    state[2] = wall_distances[(dir_idx + 1) % 4]
    state[3] = apple_x_rel
    state[4] = apple_y_rel
 
    return state
 
 
def create_q_network(input_shape, output_size):
    model = Sequential()
    model.add(Dense(128, activation='relu', input_shape=input_shape))
    model.add(Dense(64, activation='relu'))
    model.add(Dense(32, activation='relu'))
    model.add(Dense(output_size, activation='linear'))
    model.compile(optimizer='rmsprop', loss='mse', metrics=['accuracy'])
    return model
 
 
 
 
 
def action_to_index(action):
    return {(0, -BLOCK_SIZE): 0, (0, BLOCK_SIZE): 1, (-BLOCK_SIZE, 0): 2, (BLOCK_SIZE, 0): 3}[action]
 
 
def update_q_network(q_network, replay_buffer, batch_size, gamma=0.99):
    batch = replay_buffer.sample(batch_size)
    states, actions, rewards, next_states, dones = zip(*batch)
    states = np.array(states)
    actions = np.array(actions)
    rewards = np.array(rewards)
    next_states = np.array(next_states)
    dones = np.array(dones, dtype=np.float32)
 
    q_values = q_network.predict(states)
    q_values_next = q_network.predict(next_states)
 
    for i, (state, action, reward, next_state, done) in enumerate(batch):
        action_index = action_to_index(action)
        q_values[i, action_index] = reward + gamma * np.max(q_values_next[i]) * (1 - done)
 
 
    q_network.fit(states, q_values, verbose=0)
 
 
def choose_action(q_network, state, epsilon):
    if np.random.rand() < epsilon:
        return random.choice([(0, -BLOCK_SIZE), (0, BLOCK_SIZE), (-BLOCK_SIZE, 0), (BLOCK_SIZE, 0)])
    predictions = q_network.predict(state.reshape(1, 3))
    probabilities = tf.nn.softmax(predictions).numpy()[0]
    return np.random.choice([(0, -BLOCK_SIZE), (0, BLOCK_SIZE), (-BLOCK_SIZE, 0), (BLOCK_SIZE, 0)], p=probabilities)
 
 
 
 
 
 
def main():
    pygame.init()
    screen = pygame.display.set_mode((WIDTH, HEIGHT))
    clock = pygame.time.Clock()
    snake = Snake()
    food = Food()
    q_network = create_q_network((3,), 4)
    epsilon = 1.0
    replay_buffer = ReplayBuffer(BUFFER_SIZE)
    episode_counter = 0
    target_update_freq = 200
    save_counter = 0
    SAVE_MODEL_FREQ = 1000
    MODEL_SAVE_PATH = 'saved_model.h5'
 
 
    while True:
        state = get_game_state(snake, food)
        action = choose_action(q_network, state, epsilon)
        snake.change_direction(action)
        snake.move()
 
        if snake.collided_with_itself() or snake.collided_with_wall():
            snake = Snake()
            food.regenerate()
            if len(snake.positions) < 15:
                reward = -100
            else:
                reward = -10
            done = True
            episode_counter += 1
        elif snake.positions[0] == food.position:
            snake.grow()
            food.regenerate()
            eaten_apples = len(snake.positions) - 1
            reward = np.sqrt(eaten_apples) * 3.5
            done = False
        else:
            reward = -0.25
            done = False
 
 
        next_state = get_game_state(snake, food)
        replay_buffer.add(state, action, reward, next_state, done)
 
        if len(replay_buffer.buffer) >= BUFFER_SIZE:
            update_q_network(q_network, replay_buffer, batch_size=64)
            epsilon = max(epsilon * 0.995, 0.1)
            save_counter += 1
 
            if save_counter % SAVE_MODEL_FREQ == 0:
                save_model(q_network, MODEL_SAVE_PATH)
 
 
        screen.fill(WHITE)
        snake.draw(screen)
        food.draw(screen)
        pygame.display.flip()
 
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                sys.exit()
 
        clock.tick(FPS)
 
 
if __name__ == '__main__':
    main()