ant_fsm.py

import numpy as np
from scipy.spatial import KDTree
import pygame
from pygame.locals import *
from copy import deepcopy

np.random.seed(20)

STATES = 12
POP = 100
MOV_LIM = 500
GEN = 200


SCREEN_SIZE = (417, 417)
BLOCKS = 32
SIZE = 12
MARGIN = 1
FOOD = 80


WHITE = 0xffffff
BLACK = 0x000000
RED = 0xff0000
GREEN = 0x00ff00
BLUE = 0x0000ff
YELLOW = 0xffff00

FPS = 60


def gen_grid():
    global FOOD
    grid = [[WHITE for _ in range(BLOCKS)] for _ in range(BLOCKS)]
    pts = SCREEN_SIZE[0] * np.random.random((BLOCKS, 2))
    tree = KDTree(pts)
    pts = list(tree.sparse_distance_matrix(tree, FOOD).keys())
    FOOD = len(pts)
    for x, y in pts:
        grid[x][y] = GREEN
    return grid


def gen_pop():
    pop = []
    for _ in range(POP):
        fsm = []
        for _ in range(STATES):
            state = []
            for i in range(4):
                if i % 2 == 0:
                    x = np.random.randint(3)
                else:
                    x = np.random.randint(STATES)
                state.append(x)
            fsm.append(state)
        pop.append(fsm)
    return pop


grid = gen_grid()
pop = gen_pop()


LEFT = {(1, 0): (0, -1),
        (0, -1): (-1, 0),
        (-1, 0): (0, 1),
        (0, 1): (1, 0)}
RIGHT = {(1, 0): (0, 1),
         (0, 1): (-1, 0),
         (-1, 0): (0, -1),
         (0, -1): (1, 0)}


def cost(pop, p=0):
    costs = []

    for i in range(POP):
        vis = {(i, j): False for i in range(BLOCKS) for j in range(BLOCKS)}
        fsm = pop[i]
        state = 0
        eat = 0
        row, col = BLOCKS - 1, 0
        d = (1, 0)
        for g in range(MOV_LIM):
            if grid[row][col] == GREEN and not vis[(row, col)]:
                eat += 1
                vis[(row, col)] = True
                if eat == FOOD:
                    break
            x, y = ((row + d[0] + BLOCKS) % BLOCKS,
                    (col + d[1] + BLOCKS) % BLOCKS)
            if grid[x][y] == GREEN and not vis[(x, y)]:
                next_move, state = fsm[state][2:]
            else:
                next_move, state = fsm[state][:2]

            if next_move == 1:
                d = LEFT[d]
            elif next_move == 2:
                d = RIGHT[d]
            else:
                row, col = x, y

        eat = FOOD - eat
        if eat != 0:
            eat += MOV_LIM
        costs.append((i + p, eat))

    return costs


def mutate(pop):
    newpop = deepcopy(pop)
    for i in range(POP):
        fsm = newpop[i]
        for j in range(STATES):
            state = fsm[j]
            for k in range(4):
                r = np.random.standard_normal()
                if k % 2 == 0:
                    state[k] += 3 * r
                    state[k] = int(round(state[k]))
                    if state[k] >= 3 or state[k] < 0:
                        state[k] = np.random.randint(3)
                else:
                    state[k] += STATES * r
                    state[k] = int(round(state[k]))
                    if state[k] >= STATES or state[k] < 0:
                        state[k] = np.random.randint(STATES)

    return newpop


print(FOOD)
for g in range(GEN):
    c = cost(pop)
    if g % 10 == 0:
        print("Gen", g, "Min Cost =", min(c, key=lambda x: x[1])[1])
    # newpop = mutate(pop)
    newpop = gen_pop()
    newcost = cost(newpop, p=POP)
    c.extend(newcost)
    c.sort(key=lambda x: x[1])
    c = c[:POP]
    ind = [c[i][0] for i in range(POP)]
    tmp = []
    for i in ind:
        if i < 100:
            tmp.append(pop[i])
        else:
            tmp.append(newpop[i - POP])
    pop = tmp


def color_block(pos, color):
    block = [(MARGIN + SIZE) * pos[1] + MARGIN,
             (MARGIN + SIZE) * pos[0] + MARGIN, SIZE, SIZE]
    pygame.draw.rect(screen, color, block)


pygame.init()
pygame.display.set_caption("Artificial Ant Algorithm")
screen = pygame.display.set_mode(SCREEN_SIZE, 0, 32)
clock = pygame.time.Clock()

fsm = pop[0]
state = 0

itr = 0
row = BLOCKS - 1
col = 0
d = (1, 0)

count = 0

while True:
    for event in pygame.event.get():
        if event.type == QUIT:
            pygame.quit()
            exit()

    screen.fill(BLACK)

    for i in range(BLOCKS):
        for j in range(BLOCKS):
            color_block((i, j), grid[i][j])

    color_block((row, col), BLUE)
    if grid[row][col] == GREEN:
        grid[row][col] = RED
        color_block((row, col), RED)

    if count == MOV_LIM:
        pygame.time.delay(5000)
        exit()

    count += 1

    sense_food = False
    next_pos = ((row + d[0] + BLOCKS) % BLOCKS,
                (col + d[1] + BLOCKS) % BLOCKS)

    if grid[next_pos[0]][next_pos[1]] == GREEN:
        sense_food = True

    if sense_food:
        next_move, state = fsm[state][2:]
    else:
        next_move, state = fsm[state][:2]

    itr += 1
    if next_move == 1:
        d = LEFT[d]
    elif next_move == 2:
        d = RIGHT[d]
    else:
        if grid[row][col] != RED:
            grid[row][col] = BLACK
        row, col = next_pos

    clock.tick(FPS)
    pygame.display.update()


pygame.quit()