Зацикливание алгоритма q-learning python

Знакомлюсь с reinforsment-learning на примере q-learning. Читаю эту статью на хабре, но код там мне совсем не понравился, поэтому я решил написать свой -- на его основе. Дело в том, что судя по выводу мой алгоритм где-то зацикливается (потому что в один момент программа просто перестает выводить новые данные, но не останавливается), но я не могу понять где именно. Код приведен ниже.

import random
import time
import os


class Q:
    def __init__(self):
        self.gamma = 0.95
        self.alpha = 0.05

        self.agent = None

        self.q_table = {}

    def set_agent_object(self, agent):
        self.agent = agent

    def teaching(self):

        self.agent.previous_action = self.agent.current_action
        self.agent.previous_state = self.agent.current_state

        self.agent.current_action = self.agent.select_action()
        self.agent.current_state = self.agent.get_state(self.agent.x, self.agent.y)

        if self.agent.previous_state not in self.q_table:
            self.q_table[self.agent.previous_state] = [0 for _ in self.agent.actions]

        if self.agent.current_state not in self.q_table:
            self.q_table[self.agent.current_state] = [0 for _ in self.agent.actions]

        q_max = max(self.q_table[self.agent.current_state])

        self.q_table[self.agent.previous_state][self.agent.previous_action] += \
            self.alpha * (self.agent.reward + self.gamma * q_max -
                          self.q_table[self.agent.previous_state][self.agent.previous_action]
                          )


class Environment:

    def __init__(self, dim, q_function):
        self.dim = dim

        self.q_model = q_function
        self.enemies = [Enemy(3, 3, dim), Enemy(4, 4, dim), Enemy(5, 5, dim)]
        self.agent = Agent(self.q_model, 1, 1, dim,  self.enemies)
        self.q_model.set_agent_object(self.agent)
        self.map = list([['=' for _ in range(self.dim)] for _ in range(self.dim)])

    def step(self):

        for enemy in self.enemies:
            enemy.move()

        self.agent.move()

    def get_reward(self, end_bool):
        if end_bool:
            self.agent.reward = 10
        else:
            self.agent.reward = -100

    def visualise(self):

        # os.system('cls')
        print(100*'\n')
        self.map = list([['=' for _ in range(self.dim)] for _ in range(self.dim)])

        agent_x, agent_y = self.agent.get_coordinates()
        self.map[agent_x][agent_y] = 'A'

        for enemy in self.enemies:
            enemy_x, enemy_y = enemy.get_coordinates()
            self.map[enemy_x][enemy_y] = 'E'

        for pixel in self.map:
            print(*pixel)

    def is_finished(self):
        finished = True

        agent_x, agent_y = self.agent.get_coordinates()

        for enemy in self.enemies:
            enemy_x, enemy_y = enemy.get_coordinates()

            finished = finished and ((agent_x, agent_y) == (enemy_x, enemy_y))

        return finished

    def play(self, visualise=True):

        finished = self.is_finished()

        iteration_ = 0

        while not finished:

            if visualise:
                time.sleep(1)
                self.visualise()

            self.step()

            finished = self.is_finished()

            self.get_reward(finished)

            self.q_model.teaching()

            iteration_ = iteration_ + 1


        return iteration_


class Unit:

    def __init__(self, x, y):
        self.x = x
        self.y = y

        self.actions = [(-1, -1), (-1, 0), (-1, 1), (0, -1), (0, 0),
                        (0, 1), (1, -1), (1, 0), (1, 1)]

    def get_coordinates(self):
        return self.x, self.y


class Agent(Unit):

    def __init__(self, q_function, x, y, dim, enemies):

        super().__init__(x, y)

        self.dim = dim

        self.enemies = enemies
        self.q_model = q_function

        self.dx = 0
        self.dy = 0

        self.epsilon = 0.7

        self.reward = 0

        self.current_state = self.get_state(x, y)
        self.current_action = self.select_action()

        self.previous_state = self.get_state(x, y)
        self.previous_action = self.select_action()

    def move(self):
        self.dx, self.dy = self.actions[self.select_action()]

        new_x = self.x + self.dx
        new_y = self.y + self.dy

        if (0 < new_x < self.dim) and (0 < new_y < self.dim):
            self.x = new_x
            self.y = new_y

    def select_action(self):

        if random.random() < self.epsilon:
            action = random.choice([i for i in range(len(self.actions))])
        else:
            if self.current_state not in self.q_model.q_table:
                self.q_model.q_table[self.current_state] = [0 for _ in self.actions]

            action = max(list(enumerate(self.q_model.q_table[self.current_state])), key=lambda x: x[1])[0]
        # print(self.q_model.q_table.get(self.current_state))
        # print(action)

        return action

    def get_state(self, x, y):
        #  состояние -- координаты всех врагов и величины,
        #  на которые хочет сдвинуться агент

        features = []

        for enemy in self.enemies:
            enemy_x, enemy_y = enemy.get_coordinates()

            features.append(enemy_x)
            features.append(enemy_y)

        features.append(x)
        features.append(y)

        features.append(self.dx)
        features.append(self.dy)

        state = tuple(features)

        return state


class Enemy(Unit):
    def __init__(self, x, y, dim):
        super().__init__(x, y)
        self.dim = dim

    def move(self):
        expr = False

        while not expr:

            new_x = self.x + random.choice(self.actions[0])
            new_y = self.y + random.choice(self.actions[1])

            expr = ((0 <= new_x < self.dim) and (0 <= new_y < self.dim))

            if expr:
                self.x = new_x
                self.y = new_y


if __name__ == "__main__":
    q_model = Q()
    iteration = 0

    for epoch in range(5):
        print(epoch)
        environment = Environment(7, q_model)
        iteration = environment.play()
        environment.visualise()

P.S 5 итераций я поставил только для отладки, по-хорошему нужно порядка 1000-5000

P.P.S Маленькая статья на вики про q-обучение для тех, кто хочет помочь, но не знаком с ним


Ответы (1 шт):

Автор решения: DKay

Ошибок было несколько: в формуле обновления q-функции и в методе передвижения агента.

→ Ссылка