Simplified Smart Rockets using Genetic Algorithm

Question

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. 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()

Answer 1

UX

When I run the code, I don't always see the value for "Generation" printed in the GUI. It is sometimes clipped at the right edge of the screen. It would be better to move it to its own line.

Also, "Average Fitness" shows a lot more digits than are needed. It would be better to show fewer digits.

Here are suggested changes:

frame_rate_text = my_font.render("Frame = {0} rate = {1:.2f} fps"
 .format(frame_count, frame_rate), True, COLOR_KEY["WHITE"])
gen_text = my_font.render(f"Generation: {gen}", True, COLOR_KEY["WHITE"])
fitness_text = my_font.render(f"Average Fitness: {avg_fit:.9f}", True, COLOR_KEY["WHITE"])
window.blit(gen_text, (WINDOW_SIZE[0] - 300, 70))

This code makes use of f-strings to simplify the formatting code:

f"Generation: {gen}"

Documentation

The PEP 8 style guide recommends adding docstrings for classes and functions.

For example, Block is a generic name for a class. The docstring should describe what you mean by a block:

class Block:
 """
 Rectangle on the 2D screen.
 Represents a member of the population.
 """
 def evaluate(self):
 """Evaluate fitness of population member"""

Naming

PEP 8 recommends all caps for constants.

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

Would become:

TARGET = [500, 500]
POP_SIZE = 20
MAX_MOVES = 1500
MUTATION_RATE = 30

The word "pop" has many meanings. It would be better to spell out "population". Instead of:

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

consider:

def create_population():
 population = []
 for _ in range(POPULATION_SIZE):
 population.append(Block())
 return population

Since the block variable was unused, it was replaced by _.

ga is not a very descriptive name for a function. Again, it would be better to spell it out: genetic_algo.

Simpler

Instead of sqrt:

self.fitness = 1 / math.sqrt((dx * dx) + (dy * dy))

consider hypot:

self.fitness = 1 / math.hypot(dx, dy)

In the ga function, you can replace these 2 lines:

returning = (avg_fit, population)
return returning

with 1 line:

return (avg_fit, population)

This eliminates the generically named intermediate variable: returning

Similarly, in the main function, replace:

population = create_pop()
population = generate_moves(population)

with:

population = generate_moves(create_pop())

Main

It is customary to add a main guard:

if __name__ == "__main__":
 main()