UD04 · EX6 — Aprendizaje por refuerzo con NEAT¶

%pip install aitk neat-python graphviz

import aitk.robots as bots
from aitk.utils import Grid
from aitk.algorithms import neat

import numpy as np
import cv2
import requests
import matplotlib.pyplot as plt

def neat_controller(robot):
    # Si ya nos hemos dado la vuelta, salimos
    #sequence = robot.state["grid"].sequence

    #if (4, 1) in sequence and (5, 1) in sequence:
    #    return True

    cam = robot['ground-camera']
    image = cam.get_image()

    # Convirtemos la imagen en una matriz numpy para trabajar en OpenCV
    im = np.array(image)

    # Nos quedamos con la parte superior de la imagen
    im_baix = im[00:20, 0:amplada_camera]

    # convertir a escala de grises
    im_bn = cv2.cvtColor(im_baix, cv2.COLOR_RGB2GRAY)

    # Aplicamos un filtro de ilusión
    im_blurred = cv2.GaussianBlur(im_bn, (5, 5), 0)

    # Binarizamos la imagen
    ret, im_thresh = cv2.threshold(im_blurred, 100, 200, cv2.THRESH_BINARY)

    # Buscamos los contornos
    contours, hierarchy = cv2.findContours(im_thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    # Dibujamos los contornos
    cv2.drawContours(im_baix, contours, -1, (0, 255, 0), 1)


    # Definir las variables que se devuelven
    centre = 0

    # Calculamos el centro del borde más grande
    if len(contours) >= 2:
        cnt = max(contours, key=cv2.contourArea)
        M = cv2.moments(cnt)
        # En m tenemos los momentos de la imagen
        M00 = M['m00'] # area
        M10 = M['m10'] # sum(x)

        # Si el área es diferente de 0, calculamos el centro de masa
        if M00 != 0:
            centre = int(M10/M00)
            inputs = [centre/amplada_camera]
            out = robot.state["net"].activate(inputs)

        # Aumentamos la distancia recorrida
        x, y, a = robot.get_pose()
        robot.state["grid"].update(x, y)
        robot.state["grid"].draw()

        robot.move(out[0], out[1])
    else:
        robot.state["out"] = True
        return True

    if robot.stalled:
        return True


def fitness(net, seconds=60, real_time=False):
    # Evalúa el cromosoma y devuelve su valor físico

    # Mantenemos la red en el robot para usarla en el controlador
    robot.state["net"] = net

    # Volvemos a poner el robot en la posición inicial
    world.reset()
    robot.state["net"] = net
    robot.state["grid"] = Grid((10, 10), robot.world)
    robot.state["out"] = False

    # Ejecutamos el robot durante 100 segundos como máximo
    world.seconds(
        seconds, [neat_controller],
        real_time=real_time
    )

    # Calculamos el valor de fitness
    score = 0

    if not robot.stalled and not robot.state["out"]:
        visites  = robot.state["grid"].analyze_visits()
        score = visites*10 - world.time / 200


    return score

def eval_population(pop, config):
    """Esta función recibirá una población de cromosomas y los evaluará"""
    for genome_id, genome in pop:
        net = neat.nn.FeedForwardNetwork.create(genome, config)
        genome.fitness = fitness(net, 60, real_time=False)

def run(config_file):
    # Load configuration
    config = neat.Config(neat.DefaultGenome, neat.DefaultReproduction,
                         neat.DefaultSpeciesSet, neat.DefaultStagnation,
                         config_file)

    # Create the population, which is the top-level object for a NEAT run.
    p = neat.Population(config)

    # Add a stdout reporter to show progress in the terminal.
    p.add_reporter(neat.StdOutReporter(True))
    stats = neat.StatisticsReporter()
    p.add_reporter(stats)
    # Uncomment this to save checkpoints every 5 generations
    # p.add_reporter(neat.Checkpointer(5))

    winner = p.run(eval_population, 50)

    # Display the winning genome.
    print('\nBest genome:\n{!s}'.format(winner))
    node_names = {-1:'d1', -2:'d2', -3:'d3', -4:'timer',
                  -5:'stall', -6:'bias',
                  0:'translate', 1: 'rotate'}
    neat.visualize.draw_net(config, winner, True, node_names=node_names)
    neat.visualize.plot_stats(stats, ylog=False, view=True)
    neat.visualize.plot_species(stats, view=True)

    return winner

with open("winner.pkl", "wb") as f:
    import pickle

    pickle.dump(winner, f)
    f.close()