Добавь изменение цвета по клавише enter...)))

import pygame
import random
import math

WIDTH = 500
HEIGHT = 500

pygame.init()
screen = pygame.display.set_mode((WIDTH, HEIGHT))
running = True
r = random.randint(0,255)
g = random.randint(0,255)
b = random.randint(0,255)
x = 250
y = 0
speed= 1
state = 0
while running:
    
    for event in pygame.event.get():
    
        
        if event.type == pygame.QUIT:
            running = False   
   

    screen.fill((0,0,0))
    
    for x in range(50,450,50):
        for y in range(50,450,50):
            if state == 0:
                state = 1
                speed = -speed
            else:
                state = 0
                speed = -speed
                x = x + speed
                pygame.draw.rect(screen,(r,g,b),(x, y, 50,50))
        pygame.display.flip()
pygame.quit()