Use `pygame` to render custom pong env. Goal is to have a controllable env to train RL agents.

octopong.py

#!/usr/bin/env python
# coding: utf-8
import pygame
import time
import random
from joblib import load

pygame.init()

# Run on either ubuntu or Powershell, don't use wsl2 as it's not supported
print(pygame.display.list_modes())

# TODO Refactor to have a sep class object for each player
class OctopusPong(object):
    def __init__(self, 
        width = 1000, 
        height = 650, 
        paddle_width = 40, 
        paddle_height  = 100, 
        width_offset = 50, 
        ball_radius = 50, 
        ball_x_vel = -0.5, 
        ball_y_vel = -0.25, 
        background=(0,0,175),
        octopus_path = "sprites/octopus.png",
        squid_path = "sprites/squid.png",
        squid2_path = "sprites/squid.png",
        bg_path = "sprites/bg.jpg"):
        pygame.display.set_caption('Octopus Pong')

        self.screen_width = width
        self.screen_height = height
        self.paddle_width = paddle_width
        self.paddle_height = paddle_height

        # position of player paddle
        self.player_x = width_offset
        self.player_y = (height / 2) - (paddle_height / 2)

        # position of the AI's paddle
        self.ai_x = width - 2 * width_offset
        self.ai_y = (height / 2) - (paddle_height / 2)

        # position of ball & velocity of the ball
        self.ball_x = width / 2
        self.ball_y = height / 2
        self.ball_x_vel = ball_x_vel
        self.ball_y_vel = ball_y_vel

        self.game_over = False

        self.background = background

        self.screen = pygame.display.set_mode((width, height))
        self.screen.fill(background)

        self.ball_radius = ball_radius
        
        # ensure the octopus.png or whatever else is in the same folder as this code
        if octopus_path:
            self.octopus = self.get_image(octopus_path, (ball_radius, ball_radius))
        else:
            self.octopus = None

        if squid_path:
            self.squid = self.get_image(squid_path, (paddle_width, paddle_height))
        else:
            self.squid = None

        if squid2_path:
            self.squid2 = self.get_image(squid2_path, (paddle_width, paddle_height))
        else:
            self.squid2 = None

        if bg_path:
            self.bg = self.get_image(bg_path, (width, height))
        else:
            self.bg = None

        # load model here when ready
        self.model = None

    def get_image(self, path, scale):
        image = pygame.image.load(path)
        if 'png' in path:
            image = image.convert_alpha()
        else:
            image = image.convert()
        return pygame.transform.scale(image, scale)

    def update_player_and_ball(self):
        # Draw background
        if self.bg:
            self.screen.blit(self.bg, (0, 0))
        else:
            self.blackout_screen()
        # draw the player's paddle
        if self.squid:
            self.screen.blit(self.squid, (self.player_x, self.player_y))
        else:
            # use dummy rectangle to illustrate player
            pygame.draw.rect(
                self.screen, 
                (255, 255, 255), 
                (self.player_x, self.player_y, self.paddle_width, self.paddle_height)
            )

        # draw the AI's paddle
        if self.squid2:
            self.screen.blit(self.squid2, (self.ai_x, self.ai_y))
        else:
            # use dummy rectangle to illustrate AI
            pygame.draw.rect(
                self.screen, 
                (255, 255, 255), 
                (self.ai_x, self.ai_y, self.paddle_width, self.paddle_height)
            )

        # draw the octopus
        if self.octopus:
            self.screen.blit(self.octopus, (self.ball_x, self.ball_y))
        else:
            # use dummy circle to illustrate octopus
            pygame.draw.circle(
                self.screen,
                (255, 255, 255),
                (self.ball_x, self.ball_y),
                self.ball_radius
            )

    def update_ai_pos(self):
        # TODO Replace this with prediction of model
        prediction = self.get_prediction()

        # if the predicted paddle position is less than the current position move up by 1
        # if it's greater than the current position, move down by 1, otherwise stay in place
        if prediction < self.ai_y:
            self.ai_y -= 1
        elif prediction > self.ai_y:
            self.ai_y += 1

    def update_ball_pos(self):
        self.ball_x += self.ball_x_vel
        self.ball_y += self.ball_y_vel
        
        # check if ball collided with top or bottom wall
        if self.ball_y <= 0 or self.ball_y >= self.screen_height:
            self.ball_y_vel *= -1 # Reverse y velocity

        # check if ball collided with a paddle
        if (self.ball_x <= self.player_x + self.paddle_width and self.ball_y >= self.player_y and self.ball_y <= self.player_y + self.paddle_height) \
                or (self.ball_x >= self.ai_x and self.ball_y >= self.ai_y and self.ball_y <= self.ai_y + self.paddle_height):
            self.ball_x_vel *= -1 # Reverse x velocity

        # check if ball went off the screen
        if self.ball_x <= 0 or self.ball_x >= self.screen_width:
            self.game_over = True

    def reverse_ball_x_direction(self):
        self.ball_x_vel *= -1

    def get_prediction(self):
        if self.model == None:
            return random.randint(0, self.screen_height)
        else:
            # TODO implement model predicting behavior here
        
    def reverse_ball_y_direction(self):
        self.ball_y_vel *= -1

    def blackout_screen(self):
        self.screen.fill(self.background)

    def update_display(self):
        pygame.display.update()   

    def move_player_up(self):
        self.player_y -= 1

    def move_player_down(self):
        self.player_y += 1

    def end_game(self):
        self.game_over = True

    def is_end_game(self):
        return self.game_over

def main_loop():
    """
    main loop for the game
    """
    pong = OctopusPong()
    # start the ball in a random direction
    if int(random.random() * 2) == 1:
        pong.reverse_ball_x_direction()

    if int(random.random() * 2) == 1:
        pong.reverse_ball_y_direction()

    while not pong.is_end_game():
        pong.update_player_and_ball()
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                sys.exit()

        keys = pygame.key.get_pressed()

        if keys[pygame.K_w]:
            pong.move_player_up()
        elif keys[pygame.K_s]:
            pong.move_player_down()
        elif keys[pygame.K_ESCAPE]: # escape to quite the game
            pong.end_game()

        pong.update_ai_pos()
        pong.update_ball_pos()

        pong.update_display()

if __name__ == '__main__':
    main_loop()
    pygame.quit()
    quit()