3
\$\begingroup\$

This is a very bare bones program where a population of blocks learn to reach a target on a 2d screen (My attempt at making a smart rockets program)

I've only used python for 2 weeks so i'm positive certain aspects of the program have extraneous lines of code and unnecessary junk, so what are some ways I can optimize the program/algorithm and make it more readable?

import random
import math
import pygame
import time

WINDOW_SIZE = (600, 600)
COLOR_KEY = {"RED": [255, 0, 0, 10],
             "GREEN": [0, 255, 0, 128],
             "BLUE": [0, 0, 255, 128],
             "WHITE": [255, 255, 255],
             "BLACK": [0, 0, 0]}

target = [500, 500]
pop_size = 20
max_moves = 1500
mutation_rate = 30


class Block:
    def __init__(self):
        self.color = COLOR_KEY["GREEN"]
        self.size = (40, 40)
        self.move_set = []
        self.position = [0, 0]
        self.fitness = -1
        self.target_reached = False

    def evaluate(self):
        dx = target[0] - self.position[0]
        dy = target[1] - self.position[1]
        if dx == 0 and dy == 0:
            self.fitness = 1.0
        else:
            self.fitness = 1 / math.sqrt((dx * dx) + (dy * dy))

    def move(self, frame_count):
        if not self.target_reached:
            self.position[0] += self.move_set[frame_count][0]
            self.position[1] += self.move_set[frame_count][1]

            if self.position[0] == target[0] and self.position[1] == target[1]:
                self.target_reached = True
                self.color = COLOR_KEY["BLUE"]


def create_pop():
    pop = []
    for block in range(pop_size):
        pop.append(Block())

    return pop


def generate_moves(population):
    for block in population:
        for _ in range(max_moves+1):
            rand_x = random.randint(-1, 1)
            rand_y = random.randint(-1, 1)
            block.move_set.append([rand_x, rand_y])

    return population


def fitness(population):
    for block in population:
        block.evaluate()


def selection(population):
    population = sorted(population, key=lambda block: block.fitness, reverse=True)
    best_fit = round(0.1 * len(population))
    population = population[:best_fit]
    return population


def cross_over(population):
    offspring = []

    for _ in range(int(pop_size/2)):
        parents = random.sample(population, 2)
        child1 = Block()
        child2 = Block()

        split = random.randint(0, max_moves)
        child1.move_set = parents[0].move_set[0:split] + parents[1].move_set[split:max_moves]
        child2.move_set = parents[1].move_set[0:split] + parents[0].move_set[split:max_moves]

        offspring.append(child1)
        offspring.append(child2)

    return offspring


def mutation(population):
    chance = random.randint(0, 100)
    num_mutated = random.randint(0, pop_size)

    if chance >= 100 - mutation_rate:
        for _ in range(num_mutated):
            mutated_block = population[random.randint(0, len(population) - 1)]
            for _ in range(50):
                if chance >= 100 - mutation_rate/2:
                    rand_x = random.randint(0, 1)
                    rand_y = random.randint(0, 1)
                else:
                    rand_x = random.randint(-1, 1)
                    rand_y = random.randint(-1, 1)

                mutated_block.move_set[random.randint(0, max_moves - 1)] = [rand_x, rand_y]

    return population


def calc_avg_fit(population):
    avg_sum = sum(block.fitness for block in population)
    return avg_sum/pop_size


def ga(population):
    fitness(population)
    avg_fit = calc_avg_fit(population)
    population = selection(population)
    population = cross_over(population)
    population = mutation(population)

    returning = (avg_fit, population)
    return returning


def main():
    pygame.init()
    window = pygame.display.set_mode(WINDOW_SIZE)
    pygame.display.set_caption("AI Algorithm")
    population = create_pop()
    population = generate_moves(population)

    my_font = pygame.font.SysFont("Arial", 16)
    frame_count = 0
    frame_rate = 0
    t0 = time.process_time()
    gen = 0
    avg_fit = 0

    while True:
        event = pygame.event.poll()
        if event.type == pygame.QUIT:
            break

        frame_count += 1
        if frame_count % max_moves == 0:
            t1 = time.process_time()
            frame_rate = 500 / (t1-t0)
            t0 = t1
            frame_count = 0
            data = ga(population)
            avg_fit = data[0]
            population = data[1]
            gen += 1

        window.fill(COLOR_KEY["BLACK"], (0, 0, WINDOW_SIZE[0], WINDOW_SIZE[1]))

        for block in population:
            block.move(frame_count)

        for block in population:
            window.fill(block.color, (block.position[0], block.position[1], block.size[0], block.size[1]))

        window.fill(COLOR_KEY["RED"], (target[0] + 10, target[1] + 10, 20, 20), 1)

        frame_rate_text = my_font.render("Frame = {0} rate = {1:.2f} fps Generation: {2}"
                                         .format(frame_count, frame_rate, gen), True, COLOR_KEY["WHITE"])

        fitness_text = my_font.render("Average Fitness: {0}".format(avg_fit), True, COLOR_KEY["WHITE"])

        window.blit(fitness_text, (WINDOW_SIZE[0] - 300, 40))
        window.blit(frame_rate_text, (WINDOW_SIZE[0] - 300, 10))

        pygame.display.flip()

    pygame.quit()


main()
\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Browse other questions tagged or ask your own question.