Simplified Smart Rockets using Genetic Algorithm

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