Баг нейросети. Как решить?

Я сделал нейронную сеть, которая играет в змейку. На вход нейросети я даю координаты змейки и координаты еды, а на выходе должен получить направление. Т.е на выходе 2 нейрона, один отвечает за направление по координате y(0 - вниз, 1 - вверх), а второй, соответственно за направление по координате x(0 - влево, 1 - вправо). Но когда я запускаю игру, то нейросеть сначала нормально работает, но змейка начинает ползать по кругу. Как решить проблему?
Я новичок, так что не бейте сильно :)
Вот код:

import pygame
import random
import numpy as np

# AI
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
import matplotlib.pyplot as plt

# model = keras.models.load_model('aisnake.h5')

x_train = []
y_train = []

for i in range(1001):
    snake_x = random.randrange(10, 630, 10)
    food_x = random.randrange(10, 630, 10)
    snake_y = random.randrange(10, 470, 10)
    food_y = random.randrange(10, 470, 10)
    # right - 1 | left - 0
    # up -1 | down - 0
    if snake_x > food_x:
        if snake_y > food_y:
            result_x = 0
            result_y = 1
            x_train.append([1/snake_x, 1/food_x, 1/snake_y, 1/food_y])
            y_train.append([result_x, result_y])
    elif snake_x < food_x:
        if snake_y < food_y:
            result_x = 1
            result_y = 0
            x_train.append([1/snake_x, 1/food_x, 1/snake_y, 1/food_y])
            y_train.append([result_x, result_y])
    elif snake_x == food_x and snake_y == food_y:
        continue
    elif snake_x == food_x or snake_y == food_y:
        continue

x_train = np.array(x_train)
y_train = np.array(y_train)

model = keras.Sequential()
model.add(layers.Dense(2, input_shape=(4,), activation="sigmoid"))
model.add(layers.Dense(2, activation="sigmoid"))
model.add(layers.Dense(2))
opt = keras.optimizers.Adam(learning_rate=0.1)
 
model.compile(
    optimizer=opt,
    loss='mse',
    metrics=["accuracy"])
 
history = model.fit( x_train, y_train, batch_size=100, epochs=1000)

plt.plot(history.history['loss'], label='Ошибок')
plt.plot(history.history['accuracy'], label='Угадываний')
plt.xlabel('Эпоха обучения')
plt.ylabel('Значение')
plt.legend()
plt.show()
# model.save('aisnake.h5')

# initialize pygame
pygame.init()

# create display & run update
width = 640
height = 480
display = pygame.display.set_mode((width, height))

pygame.display.update()
pygame.display.set_caption("AIsnake by Ovsyanka")

# define colors
colors = {
    "snake_head": (0, 200, 0),
    "snake_tail": (102, 102, 51),
    "apple": (255, 0, 0)
}

# snake position with offsets
snake_pos = {
    "x": width/2-10,
    "y": height/2-10,
    "x_change": 0,
    "y_change": 0
}

# snake el size
snake_size = (10, 10)

# current snake movement speed
snake_speed = 10

# snake tails
snake_tails = []

for i in range(3):
    snake_tails.append([snake_pos["x"] + 10*i, snake_pos["y"]])

# food
food_pos = {
    "x": round(random.randrange(0, width - snake_size[0]) / 10) * 10,
    "y": round(random.randrange(0, height - snake_size[1]) / 10) * 10,
}

food_size = (10, 10)
food_eaten = 0

# mode
AI = 1

# score
score = 0
pygame.display.set_caption("AIsnake by Ovsyanka | SCORE = " + str(score))
pygame.display.update()

# start loop
game_end = False
clock = pygame.time.Clock()

while not game_end:
    # game loop
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            game_end = True
        elif AI == 0:
            if event.type == pygame.KEYDOWN:
                if event.key == pygame.K_LEFT and snake_pos["x_change"] == 0:
                    # move left
                    snake_pos["x_change"] = -snake_speed
                    snake_pos["y_change"] = 0

                elif event.key == pygame.K_RIGHT and snake_pos["x_change"] == 0:
                    # move right
                    snake_pos["x_change"] = snake_speed
                    snake_pos["y_change"] = 0

                elif event.key == pygame.K_UP and snake_pos["y_change"] == 0:
                    # move up
                    snake_pos["x_change"] = 0
                    snake_pos["y_change"] = -snake_speed

                elif event.key == pygame.K_DOWN and snake_pos["y_change"] == 0:
                    # move down
                    snake_pos["x_change"] = 0
                    snake_pos["y_change"] = snake_speed
    # clear screen
    display.fill((51,51,51))

    # move snake tails
    ltx = snake_pos["x"]
    lty = snake_pos["y"]

    for i,v in enumerate(snake_tails):
        _ltx = snake_tails[i][0]
        _lty = snake_tails[i][1]

        snake_tails[i][0] = ltx
        snake_tails[i][1] = lty

        ltx = _ltx
        lty = _lty

    # draw snake tails
    for t in snake_tails:
        pygame.draw.rect(display, colors["snake_tail"], [
            t[0],
            t[1],
            snake_size[0],
            snake_size[1]])

    # draw snake
    snake_pos["x"] += snake_pos["x_change"]
    snake_pos["y"] += snake_pos["y_change"]

    # barrier
    if(snake_pos["x"] < 0 or snake_pos["x"] > 630):
        print('Вы врезались в стену!')
        print('Максимальный рекорд: ' + str(score))
        quit = True
    
    if(snake_pos["y"] < 0 or snake_pos["y"] > 470):
        print('Вы врезались в стену!')
        print('Максимальный рекорд: ' + str(score))
        quit = True
    
    pygame.draw.rect(display, colors["snake_head"], [
        snake_pos["x"],
        snake_pos["y"],
        snake_size[0],
        snake_size[1]])

    # draw food
    for i, v in enumerate(snake_tails):
        if(snake_pos["x"]+snake_pos["x_change"] == snake_tails[i][0]
                and snake_pos["y"]+snake_pos["y_change"] == snake_tails[i][1]):
                    food_pos.clear()
                    food_pos = {
                        "x": round(random.randrange(0, width - snake_size[0]) / 10) * 10,
                        "y": round(random.randrange(0, height - snake_size[1]) / 10) * 10,
                    }
    pygame.draw.rect(display, colors["apple"], [
        food_pos["x"],
        food_pos["y"],
        food_size[0],
        food_size[1]])

    if AI == 1:
        pred = model.predict(np.array([[1/snake_pos['x'], 1/food_pos['x'], 1/snake_pos['y'], 1/food_pos['y']]]))
        for i in pred:
            pred_x = i[0]
            print(pred_x)
            pred_y = i[1]
            print(pred_y)
            if pred_x > 0.5:
                pred_x = 1
            elif pred_x < 0.5:
                pred_x = 0
            if pred_y > 0.5:
                pred_y = 1
            elif pred_y < 0.5:
                pred_y = 0
        if pred_x == 1 and snake_pos["x_change"] == 0:
            # move right
            snake_pos["x_change"] = snake_speed
            snake_pos["y_change"] = 0
        elif pred_x == 0 and snake_pos["x_change"] == 0:
            # move left
            snake_pos["x_change"] = -snake_speed
            snake_pos["y_change"] = 0
        elif pred_y == 0 and snake_pos["y_change"] == 0:
            # move down
            snake_pos["x_change"] = 0
            snake_pos["y_change"] = snake_speed
        elif pred_y == 1 and snake_pos["y_change"] == 0:
            # move up
            snake_pos["x_change"] = 0
            snake_pos["y_change"] = -snake_speed

    # detect collision with food
    if(snake_pos["x"] == food_pos["x"]
        and snake_pos["y"] == food_pos["y"]):
        food_eaten += 1
        snake_tails.append([food_pos["x"], food_pos["y"]])

        food_pos = {
            "x": round(random.randrange(0, width - snake_size[0]) / 10) * 10,
            "y": round(random.randrange(0, height - snake_size[1]) / 10) * 10,
        }
        score += 1
        pygame.display.set_caption("AIsnake by Ovsyanka | SCORE = " + str(score))
        pygame.display.update()

    # detect collision with tail
    for i,v in enumerate(snake_tails):
        if(snake_pos["x"]+snake_pos["x_change"] == snake_tails[i][0]
            and snake_pos["y"]+snake_pos["y_change"] == snake_tails[i][1]):
            snake_tails = snake_tails[:i]
            print('Вы врезались в себя!')
            break

    pygame.display.update()
    
    # set FPS
    clock.tick(30)

# close app, if required
pygame.quit()
quit()