Here is the code for a clone of Tetris that we made together! You did all the coding. There were lots and lots of errors and it took us a while to get there, but we did finally get there!



import pygame
import sys
import random
import collections
from tetris_agent import TetrisAgent, TetrisNN  # Import the agent and neural network
import torch.optim as optim
import torch.nn as nn
import torch
import torch.multiprocessing as mp
import numpy as np
import multiprocessing
import os

def save_training_state(agent, optimizer, generation, base_save_path, save_lock):
    save_state = {
        'model_state_dict': agent.model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
        'generation': generation,
    }
    
    if save_lock:
        with save_lock:
            torch.save(save_state, base_save_path)
            print(f"Training state saved at generation {generation}")

def load_training_state(agent, optimizer, base_save_path):
    if os.path.exists(base_save_path):
        checkpoint = torch.load(base_save_path)
        agent.model.load_state_dict(checkpoint['model_state_dict'])
        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
        generation = checkpoint['generation']
        print(f"Resumed training from generation {generation}")
        return generation
    else:
        print("No saved state found. Starting from scratch.")
        return 0

def get_state_feature_vector(state, device):
    try:
        grid = np.array(state['board'])  # Convert the board (a 2D list) to a numpy array
    except Exception as e:
        raise e

    # Flatten the grid to create a feature vector (1D list)
    flattened_state = grid.flatten()

    # Normalize the state (optional)
    max_value = np.max(flattened_state)
    normalized_state = flattened_state / max_value if max_value > 0 else flattened_state

    # Convert to a torch tensor and move to the specified device (GPU/CPU)
    return torch.tensor(normalized_state, dtype=torch.float32).to(device)

def draw_grid(grid, surface):
    reverse_color_mapping = {v: k for k, v in color_mapping.items()}  # Reverse the mapping

    for i, row in enumerate(grid):
        for j, block in enumerate(row):
            if block:
                color = reverse_color_mapping.get(block, WHITE)  # Map back to color
                pygame.draw.rect(surface, color, (j * BLOCK_SIZE, i * BLOCK_SIZE, BLOCK_SIZE, BLOCK_SIZE))
                pygame.draw.rect(surface, WHITE, (j * BLOCK_SIZE, i * BLOCK_SIZE, BLOCK_SIZE, BLOCK_SIZE), 1)

def draw_score(surface, score):
    font = pygame.font.SysFont(None, 36)
    score_text = font.render(f'Score: {score}', True, BLACK)
    surface.blit(score_text, (10, 10))

def draw_level(surface, level):
    font = pygame.font.SysFont(None, 36)
    level_text = font.render(f'Level: {level}', True, BLACK)
    surface.blit(level_text, (10, 40))

def draw_next_tetrimino(surface, next_tetrimino_key):
    tetrimino_info = TETRIMINOS[next_tetrimino_key]  # Access the shape and color from the dictionary
    draw_x = GRID_WIDTH * BLOCK_SIZE - 2 * BLOCK_SIZE  # Position to draw next Tetrimino
    draw_y = 0  # Adjust as needed for spacing

    draw_tetrimino(surface, tetrimino_info['shape'], draw_x, draw_y, tetrimino_info['color'], scale=0.5, is_ghost=False)

def get_game_state(grid):
    """
    Returns the current state of the game.
    This function will return the grid as the game state.
    
    Args:
    - grid: The current grid of the game.

    Returns:
    - state: A dictionary representing the current state of the game.
    """
    state = {
        'board': grid
    }
    return state

def get_reward_and_next_state(grid, lines_cleared, game_over, rotations=0, moved_horizontally=False, locked=False, device=None):
    reward = SCORING_RULES.get(lines_cleared, lines_cleared * SCORE_INCREMENT)
    
    if locked:
        reward += 10
    
    if rotations > 5:
        reward -= (rotations - 5)
    
    # Penalize for the height of the tallest column
    max_height = max(GRID_HEIGHT - row for col in range(GRID_WIDTH) for row in range(GRID_HEIGHT) if grid[row][col])
    reward -= max_height * 0.5
    
    # Penalize for holes in the grid
    holes = sum(1 for col in range(GRID_WIDTH) for row in range(1, GRID_HEIGHT) if grid[row][col] == 0 and grid[row - 1][col] != 0)
    reward -= holes * 10
    
    # Reward for clearing multiple lines
    if lines_cleared > 1:
        reward += lines_cleared * 50
    if lines_cleared == 1:
        reward += 80
    
    # Penalize for an uneven surface
    surface_area = sum(abs(max(grid[row][col] for row in range(GRID_HEIGHT)) - max(grid[row][col + 1] for row in range(GRID_HEIGHT))) for col in range(GRID_WIDTH - 1))
    reward -= surface_area * 2
    
    # Penalize heavily for game over
    if game_over:
        reward -= 1000
    
    # Add an exploration bonus
    exploration_bonus = random.uniform(-2, 2)
    reward += exploration_bonus
    
    # Create the next state based on the current grid
    next_state_dict = {"board": grid}
    try:
        next_state = get_state_feature_vector(next_state_dict, device)  # Pass the device to the function
    except Exception as e:
        raise e

    return reward, next_state

# Initialize Pygame
pygame.init()
pygame.mixer.init()

# Constants
pygame.display.set_caption("Renis")  # Set window title to "Renis"
WINDOW_WIDTH = 300
WINDOW_HEIGHT = 600
BLOCK_SIZE = 30

GRID_WIDTH = WINDOW_WIDTH // BLOCK_SIZE
GRID_HEIGHT = WINDOW_HEIGHT // BLOCK_SIZE

LEVEL_UP_SCORE = 1000  # Score needed to level up
MIN_FALL_SPEED = 100   # Minimum fall speed in milliseconds
FALL_SPEED_DECREMENT = 50  # Speed increase per level

SCORING_RULES = {
    1: 100,
    2: 300,
    3: 600,
    4: 1000,
}

SCORE_INCREMENT = 100  # Kept for backward compatibility  # Score increase for each cleared line

# Colors
WHITE = (255, 255, 255)
BLUE = (0, 0, 255)
BLACK = (0, 0, 0)

# Sounds
rotate_sound = pygame.mixer.Sound("rotate.wav")
lock_sound = pygame.mixer.Sound("lock.wav")
hold_sound = pygame.mixer.Sound("hold.wav")
clear_sound = pygame.mixer.Sound("clear.wav")
renis_clear_sound = pygame.mixer.Sound("renisclear.wav")



# Tetrimino shapes
TETRIMINOS = {
    'O': {'shape': [[1, 1],
                    [1, 1]], 'color': (255, 255, 0)},  # Yellow
    'I': {'shape': [[1, 1, 1, 1]], 'color': (0, 255, 255)},  # Cyan
    'S': {'shape': [[0, 1, 1],
                    [1, 1, 0]], 'color': (0, 255, 0)},  # Green
    'Z': {'shape': [[1, 1, 0],
                    [0, 1, 1]], 'color': (255, 0, 0)},  # Red
    'L': {'shape': [[1, 1, 1],
                    [1, 0, 0]], 'color': (255, 165, 0)},  # Orange
    'J': {'shape': [[1, 1, 1],
                    [0, 0, 1]], 'color': (0, 0, 255)},  # Blue
    'T': {'shape': [[1, 1, 1],
                    [0, 1, 0]], 'color': (128, 0, 128)},  # Purple
}
color_mapping = {
    (0, 255, 255): 1,    # Cyan (I)
    (255, 255, 0): 2,    # Yellow (O)
    (0, 255, 0): 3,      # Green (S)
    (255, 0, 0): 4,      # Red (Z)
    (255, 165, 0): 5,    # Orange (L)
    (0, 0, 255): 6,      # Blue (J)
    (128, 0, 128): 7,    # Purple (T)
}

# To keep track of the last three generated pieces
last_three_pieces = collections.deque(maxlen=3)

# Initialize window
window = pygame.display.set_mode((WINDOW_WIDTH, WINDOW_HEIGHT))


def generate_tetrimino():
    while True:
        tetrimino_key = random.choice(list(TETRIMINOS.keys()))
        # Check if the new key is not repeating more than 3 times
        if last_three_pieces.count(tetrimino_key) < 3:
            break
    last_three_pieces.append(tetrimino_key)  # Add the new key to the queue
    
    tetrimino_shape = TETRIMINOS[tetrimino_key]['shape']
    x_position = random.randint(0, GRID_WIDTH - len(tetrimino_shape[0]))
    return {'shape': tetrimino_shape, 'position': (x_position, 0), 'color': TETRIMINOS[tetrimino_key]['color']}
    
def update_tetriminos(tetriminos):
    for tetrimino in tetriminos:
        x, y = tetrimino['position']
        tetrimino['position'] = (x, y + 1)
    tetriminos[:] = [t for t in tetriminos if t['position'][1] < GRID_HEIGHT]

def draw_tetrimino(surface, shape, x, y, color, scale=1):
    for i, row in enumerate(shape):
        for j, block in enumerate(row):
            if block:
                scaled_block_size = int(BLOCK_SIZE * scale)
                pygame.draw.rect(surface, color, (x + j*scaled_block_size, y + i*scaled_block_size, scaled_block_size, scaled_block_size))
                pygame.draw.rect(surface, WHITE, (x + j*scaled_block_size, y + i*scaled_block_size, scaled_block_size, scaled_block_size), 1)  # Always draw white border

def draw_title_and_instructions(surface):
    title_font = pygame.font.SysFont(None, 72)
    instr_font = pygame.font.SysFont(None, 36)

    title_text = title_font.render('Renis', True, WHITE)
    instr_text = instr_font.render('Press any key to start', True, WHITE)


    title_rect = title_text.get_rect(center=(WINDOW_WIDTH // 2, WINDOW_HEIGHT // 2 - 50))
    instr_rect = instr_text.get_rect(center=(WINDOW_WIDTH // 2, WINDOW_HEIGHT // 2 + 50))

    surface.blit(title_text, title_rect)
    surface.blit(instr_text, instr_rect)

def display_title_screen():
    # Load and play the music
    pygame.mixer.music.load('title.mp3')
    pygame.mixer.music.play(-1)  # The -1 will make the music loop indefinitely    running = True

    tetriminos = []
    last_generation_time = 0
    generation_interval = 1000

    while True:
        current_time = pygame.time.get_ticks()
        window.fill(BLACK)

        if current_time - last_generation_time > generation_interval:
            tetriminos.append(generate_tetrimino())
            last_generation_time = current_time

        update_tetriminos(tetriminos)
        draw_tetriminos(tetriminos, window)
        draw_title_and_instructions(window)

        pygame.display.flip()

        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                sys.exit()
            elif event.type == pygame.KEYDOWN:
                return

        pygame.time.delay(100)

    window.fill(BLACK)  # Clearing the screen
    title_font = pygame.font.SysFont(None, 72)  # Larger font for the title
    instr_font = pygame.font.SysFont(None, 36)  # Smaller font for instructions

    title_text = title_font.render('Renis', True, WHITE)
    instr_text = instr_font.render('Press any key to start', True, WHITE)

    title_rect = title_text.get_rect(center=(WINDOW_WIDTH // 2, WINDOW_HEIGHT // 2 - 50))
    instr_rect = instr_text.get_rect(center=(WINDOW_WIDTH // 2, WINDOW_HEIGHT // 2 + 50))

    window.blit(title_text, title_rect)
    window.blit(instr_text, instr_rect)

    pygame.display.flip()  # Update the display

def draw_transparent_block(surface, color, x, y, block_size, alpha=128):
    block = pygame.Surface((block_size, block_size), pygame.SRCALPHA)  
    block.fill(color)  
    block.set_alpha(alpha)  
    surface.blit(block, (x, y))

def draw_tetrimino(surface, shape, x, y, color, is_ghost=False, scale=1):
    for i, row in enumerate(shape):
        for j, block in enumerate(row):
            if block:
                scaled_block_size = int(BLOCK_SIZE * scale)
                if is_ghost:
                    draw_transparent_block(surface, color, x + j*scaled_block_size, y + i*scaled_block_size, scaled_block_size)
                else:
                    pygame.draw.rect(surface, color, (x + j*scaled_block_size, y + i*scaled_block_size, scaled_block_size, scaled_block_size))
                pygame.draw.rect(surface, WHITE, (x + j*scaled_block_size, y + i*scaled_block_size, scaled_block_size, scaled_block_size), 1)  # Always draw white border

#ghost tetrimino
def get_ghost_position(tetrimino, x, y, grid):
    shape = tetrimino['shape']
    ghost_y = y
    
    # TODO: Replace 'collision' with your actual collision check function
    while not collision(shape, (x, ghost_y + 1), grid):  
        ghost_y += 1
        
    return (x, ghost_y)

#collision
def collision(shape, position, grid):
    """
    Check if the space is occupied by the tetrimino.
    
    :param shape: 2D list representing the tetrimino shape.
    :param position: (x, y) position of the tetrimino.
    :param grid: 2D list representing the game grid.
    :return: Boolean value representing whether a collision occurred.
    """
    for i, row in enumerate(shape):
        for j, cell in enumerate(row):
            try:
                if cell and grid[i + position[1]][j + position[0]]:
                    return True  # Collision detected
            except IndexError:  # Trying to access out of bounds index
                return True
    return False

# Rotate the tetrimino
def rotate_tetrimino(tetrimino, grid, x, y):
    shape = tetrimino['shape']
    
    rotated_shape = [
        [shape[x][y] for x in reversed(range(len(shape)))]
        for y in range(len(shape[0]))
    ]
    
    rotated_tetrimino = {'shape': rotated_shape, 'color': tetrimino['color']}
    x_move = 0
    y_move = 0
    original_y = y

    found_valid_position = False

    while y >= 0 and not found_valid_position:
        for x_offset in [0, -1, 1]:
            if is_valid_move(grid, rotated_tetrimino['shape'], x + x_offset, y):
                x_move = x_offset
                y_move = y - original_y
                found_valid_position = True
                break
        y -= 1

    if not found_valid_position:
        return tetrimino, 0, 0

    # Play rotation sound if the tetrimino was rotated
    if (rotated_tetrimino, x_move, y_move) != (tetrimino, 0, 0):
        rotate_sound.play()

    return rotated_tetrimino, x_move, y_move
    
# Counter-clockwise rotation
def rotate_tetrimino_ccw(tetrimino, grid, x, y):
    shape = tetrimino['shape']
    
    rotated_shape = [
        [shape[x][y] for x in range(len(shape))]
        for y in reversed(range(len(shape[0])))
    ]
    
    rotated_tetrimino = {'shape': rotated_shape, 'color': tetrimino['color']}
    x_move, y_move = 0, 0

    if not is_valid_move(grid, rotated_tetrimino['shape'], x, y):
        if is_valid_move(grid, rotated_tetrimino['shape'], x + 1, y):
            x_move = 1
        elif is_valid_move(grid, rotated_tetrimino['shape'], x - 1, y):
            x_move = -1
        elif is_valid_move(grid, rotated_tetrimino['shape'], x, y - 1):
            y_move = -1
        else:
            return tetrimino, 0, 0

    # Play rotation sound if the tetrimino was rotated
    if (rotated_tetrimino, x_move, y_move) != (tetrimino, 0, 0):
        rotate_sound.play()

    return rotated_tetrimino, x_move, y_move    
#Drop the Tetrimino to the bottom
def hard_drop(current_tetrimino, current_x, current_y, current_grid):
    # Find how far the Tetrimino can drop
    while is_valid_move(current_grid, current_tetrimino, current_x, current_y + 1):
        current_y += 1
    # Return final position without altering the grid
    return current_x, current_y


# Check if the new position/rotation is valid
def is_valid_move(grid, shape, x, y):
    for i, row in enumerate(shape):
        for j, cell in enumerate(row):
            if cell:
                if y + i >= len(grid) or x + j < 0 or x + j >= len(grid[0]) or grid[y + i][x + j]:
                    return False
    return True

# Merge the tetrimino into the grid and return the number of cleared lines
def merge_tetrimino(grid, tetrimino, x, y):
    cleared_lines = 0
    for i, row in enumerate(tetrimino['shape']):
        for j, block in enumerate(row):
            if block:
                # Use an integer to represent the block in the grid
                grid[y + i][x + j] = color_mapping[tetrimino['color']]

    for i, row in enumerate(grid):
        if all(row):
            print(f"Clearing line {i}: {row}")  # Debug log
            del grid[i]
            grid.insert(0, [0] * GRID_WIDTH)
            cleared_lines += 1

    # Play clear sound effects here
    if cleared_lines == 4:  # If four lines are cleared, it's a Tetris!
        renis_clear_sound.play()
    elif cleared_lines > 0:  # If 1 to 3 lines are cleared
        clear_sound.play()

    return cleared_lines
    
# Main loop
def wait_for_key_press():
    while True:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                sys.exit()
            elif event.type == pygame.KEYDOWN:
                return

def draw_text(surface, text, size, x, y, color):
    pygame.font.init()  # You only need to call this once, so you might do it in the main part of your program
    font = pygame.font.Font(None, size)  # Choose the font and size
    text_surface = font.render(text, True, color)  # Create the text surface
    text_rect = text_surface.get_rect()
    text_rect.midtop = (x, y)
    surface.blit(text_surface, text_rect)

def display_game_over_screen(final_score):
    # Load and play the music
    pygame.mixer.music.load('gameover.mp3')
    pygame.mixer.music.play(-1)  # The -1 will make the music loop indefinitely    running = True

    flying_tetriminos = []
    for _ in range(5):
        tetrimino_key = random.choice(list(TETRIMINOS.keys()))
        tetrimino = TETRIMINOS[tetrimino_key]
        
        while True:
            x = random.randint(0, WINDOW_WIDTH - len(tetrimino['shape'][0]) * BLOCK_SIZE)
            y = random.randint(-WINDOW_HEIGHT, WINDOW_HEIGHT//2)
            
            is_too_close = False
            for x_existing, y_existing, _ in flying_tetriminos:
                if abs(x - x_existing) < BLOCK_SIZE and abs(y - y_existing) < BLOCK_SIZE:
                    is_too_close = True
                    break
            
            if not is_too_close:
                break
        
        flying_tetriminos.append((x, y, tetrimino))
    
    clock = pygame.time.Clock()
    elapsed_time = 0
    
    input_enabled = False  # Disable input initially
    input_enable_time = pygame.time.get_ticks() + 2000  # Enable input after 2000 ms (2 seconds)

    while True:
        window.fill(BLACK)
        elapsed_time += clock.tick(60)

        if elapsed_time >= 500:
            tetrimino_key = random.choice(list(TETRIMINOS.keys()))
            tetrimino = TETRIMINOS[tetrimino_key]
            x = random.randint(0, WINDOW_WIDTH - len(tetrimino['shape'][0]) * BLOCK_SIZE)
            y = -len(tetrimino['shape']) * BLOCK_SIZE
            flying_tetriminos.append((x, y, tetrimino))
            elapsed_time = 0

        for i, (x, y, tetrimino) in enumerate(flying_tetriminos):
            flying_tetriminos[i] = (x, y + 5, tetrimino)
            draw_tetrimino(window, tetrimino['shape'], x, y, tetrimino['color'])
        
        flying_tetriminos = [(x, y, tetrimino) for x, y, tetrimino in flying_tetriminos if y < WINDOW_HEIGHT]
        
        font = pygame.font.SysFont(None, 70)
        game_over_text = font.render('GAME OVER', True, WHITE)
        text_x = WINDOW_WIDTH // 2 - game_over_text.get_width() // 2
        text_y = WINDOW_HEIGHT // 2 - game_over_text.get_height() // 2 - 40
        window.blit(game_over_text, (text_x, text_y))

        score_label_font = pygame.font.SysFont(None, 50)
        score_label_text = score_label_font.render('Final Score:', True, WHITE)
        score_label_x = WINDOW_WIDTH // 2 - score_label_text.get_width() // 2
        score_label_y = WINDOW_HEIGHT // 2 - score_label_text.get_height() // 2
        window.blit(score_label_text, (score_label_x, score_label_y))

        score_value_font = pygame.font.SysFont(None, 50)
        score_value_text = score_value_font.render(str(final_score), True, WHITE)
        score_value_x = WINDOW_WIDTH // 2 - score_value_text.get_width() // 2
        score_value_y = score_label_y + score_label_text.get_height()
        window.blit(score_value_text, (score_value_x, score_value_y))
        
        pygame.display.flip()
        
        # Enable input after a delay
        if pygame.time.get_ticks() >= input_enable_time:
            input_enabled = True

        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                sys.exit()
            elif event.type == pygame.KEYDOWN and input_enabled:
                # Check if the key is not the print screen key
                if event.key not in [pygame.K_PRINT, pygame.K_PRINTSCREEN]:
                    return
        
        clock.tick(30)

    game_over_font = pygame.font.SysFont(None, 72)  
    game_over_text = game_over_font.render("Game Over", True, WHITE)
    window.blit(game_over_text, (WINDOW_WIDTH // 2 - game_over_text.get_width() // 2, WINDOW_HEIGHT // 2 - game_over_text.get_height() // 2))
    pygame.display.update()

def main_game(agent, instance_id, base_save_path="tetris_agent_model_shared.pth", counter=None, save_lock=None):
    device = agent.device
    optimizer = optim.Adam(agent.model.parameters())  # Define the optimizer
    generation = load_training_state(agent, optimizer, base_save_path)  # Load the state if exists
    
    # Initialize Pygame window at the default position (top-left corner)
    window = pygame.display.set_mode((WINDOW_WIDTH, WINDOW_HEIGHT))
    
    while True:  # Train indefinitely
        generation += 1
        print(f"\n=== Instance {instance_id}: Starting Generation {generation} ===")
        running = True
        grid = [[0] * GRID_WIDTH for _ in range(GRID_HEIGHT)]

        next_tetriminos = [random.choice(list(TETRIMINOS.keys())) for _ in range(3)]
        tetrimino_key = next_tetriminos.pop(0)
        tetrimino = TETRIMINOS[tetrimino_key]
        x, y = GRID_WIDTH // 2 - len(tetrimino['shape'][0]) // 2, 0        
        held_tetrimino = None
        can_swap = True

        clock = pygame.time.Clock()

        level = 1
        fall_speed = 500
        move_speed = 150
        fast_fall_speed = 50
        last_move_time = 0
        last_fast_fall_time = 0
        score = 0
        lock_timer = 0
        fall_time = 0
        lines_cleared = 0

        rotations = 0
        action_count = 0  # Counter to track the number of actions

        if not is_valid_move(grid, tetrimino['shape'], x, y):
            print("Game Over: Invalid initial position")
            running = False
            break

        state = get_state_feature_vector({"board": grid}, device)  # Move state to GPU

        while running:
            clock.tick(60)  # Ensure the game runs at 60 FPS

            moved_horizontally = False
            locked = False

            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    running = False  # End the current generation and move to the next
                    if counter is not None:
                        with counter.get_lock():
                            counter.value -= 1  # Decrement the counter when a window is closed
                    pygame.quit()
                    save_training_state(agent, optimizer, generation, base_save_path, save_lock)  # Save state on quit
                    return
                elif event.type == pygame.KEYDOWN:
                    if event.key == pygame.K_ESCAPE:  # Optional: Press ESC to end the current generation
                        running = False

            window.fill(WHITE)
            current_time = pygame.time.get_ticks()

            action = agent.select_action(state)

            if action == 0:  # Rotate clockwise
                rotated_tetrimino, x_move, y_move = rotate_tetrimino(tetrimino, grid, x, y)
                if is_valid_move(grid, rotated_tetrimino['shape'], x + x_move, y + y_move):
                    tetrimino = rotated_tetrimino
                    x += x_move
                    y += y_move
                    lock_timer = 0
                    rotations += 1  # Increment rotation count
            elif action == 1:  # Rotate counterclockwise
                rotated_tetrimino, x_move, y_move = rotate_tetrimino_ccw(tetrimino, grid, x, y)
                if is_valid_move(grid, rotated_tetrimino['shape'], x + x_move, y + y_move):
                    tetrimino = rotated_tetrimino
                    x += x_move
                    y += y_move
                    lock_timer = 0
                    rotations += 1  # Increment rotation count
            elif action == 2:  # Move left
                if is_valid_move(grid, tetrimino['shape'], x - 1, y):
                    x -= 1
                    lock_timer = 0
                    moved_horizontally = True
            elif action == 3:  # Move right
                if is_valid_move(grid, tetrimino['shape'], x + 1, y):
                    x += 1
                    lock_timer = 0
                    moved_horizontally = True
            elif action == 4:  # Soft drop
                if is_valid_move(grid, tetrimino['shape'], x, y + 1):
                    y += 1
                    lock_timer = 0
            elif action == 5:  # Hard drop
                while is_valid_move(grid, tetrimino['shape'], x, y + 1):
                    y += 1
                locked = True

            fall_time += clock.get_rawtime()

            if fall_time > fall_speed:
                if is_valid_move(grid, tetrimino['shape'], x, y + 1):
                    y += 1
                    lock_timer = 0
                else:
                    lock_timer += clock.get_rawtime()
                    if lock_timer >= 500:
                        lines_cleared = merge_tetrimino(grid, tetrimino, x, y)
                        score += SCORING_RULES.get(lines_cleared, lines_cleared * SCORE_INCREMENT)
                        if score >= level * 1000:
                            level += 1
                            fall_speed = max(MIN_FALL_SPEED, fall_speed - FALL_SPEED_DECREMENT)
                        tetrimino_key = next_tetriminos.pop(0)
                        tetrimino = TETRIMINOS[tetrimino_key]
                        next_tetriminos.append(random.choice(list(TETRIMINOS.keys())))
                        x, y = GRID_WIDTH // 2 - len(tetrimino['shape'][0]) // 2, 0
                        if not is_valid_move(grid, tetrimino['shape'], x, y):
                            running = False
                        can_swap = True
                        lock_timer = 0
                        rotations = 0  # Reset rotations after locking in place
                        locked = True
                fall_time = 0

            temp_grid = [row[:] for row in grid]
            for i, row in enumerate(tetrimino['shape']):
                for j, block in enumerate(row):
                    if block:
                        temp_grid[y + i][x + j] = color_mapping[tetrimino['color']]  # Use the integer value
            
            reward, next_state = get_reward_and_next_state(temp_grid, lines_cleared, not running, rotations, moved_horizontally, locked)
            next_state = torch.tensor(next_state, dtype=torch.float32).to(device)  # Move next state to GPU

            agent.update(state, action, reward, next_state, not running)
            state = next_state

            draw_grid(temp_grid, window)
            draw_score(window, score)
            draw_next_tetrimino(window, next_tetriminos[0])
            pygame.display.flip()

            action_count += 1
            if action_count % 50 == 0:  # Output information every 50 actions
                print(f"Instance {instance_id} | Action Count: {action_count} | Score: {score} | Reward: {reward} | Lines Cleared: {lines_cleared}")

        print(f"Instance {instance_id}: Generation {generation} ended with score {score}")

        # Save the model and training state safely using a multiprocessing lock
        save_training_state(agent, optimizer, generation, base_save_path, save_lock)

        # Additionally, save a snapshot every 10 generations
        if generation % 10 == 0:
            generation_save_path = f"{base_save_path.split('.pth')[0]}_gen{generation}.pth"
            agent.save_model(generation_save_path)
            print(f"Instance {instance_id}: Snapshot saved at {generation_save_path}")

        # Proceed to the next generation instead of quitting
        pygame.quit()
        pygame.init()
        window = pygame.display.set_mode((WINDOW_WIDTH, WINDOW_HEIGHT))

if __name__ == "__main__":
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {'GPU' if torch.cuda.is_available() else 'CPU'}")  # Output whether GPU or CPU is used
    
    num_instances = 10
    processes = []
    counter = mp.Value('i', num_instances)  # Shared counter initialized to the number of instances
    save_lock = mp.Lock()  # Lock to ensure safe saving of the model

    for i in range(num_instances):
        agent = TetrisAgent(state_size=GRID_WIDTH * GRID_HEIGHT, action_size=5, device=device)
        p = mp.Process(target=main_game, args=(agent, i, "tetris_agent_model_shared.pth", counter, save_lock))
        processes.append(p)
        p.start()

    while True:
        if counter.value <= 0:  # Check if all windows are closed
            print("All windows closed. Terminating training.")
            for p in processes:
                p.terminate()  # Terminate all processes
            break
        pygame.time.delay(1000)  # Check every second