How do I make it so this saves frames in a batch so it doesnt lag the game, in order also to improve training

import re
import os
import cv2
import logging
import pydirectinput
import torch
import torch.nn as nn
import torch.optim as optim
import random
import pygetwindow as gw
import numpy as np
import subprocess
import pyautogui
import torchvision.models as models
import time
from collections import namedtuple
from pathlib import Path
from PIL import Image

# Disable failsafe feature of PyAutoGUI
pyautogui.FAILSAFE = False

# Configure logging
LOG_FILE = r"C:\Users\Ninja.000\Downloads\SFD_AI\training_log.txt"
logging.basicConfig(filename=LOG_FILE, level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Define the DQN model
class DQN(nn.Module):
    def __init__(self, input_size, output_size, use_pretrained=False):
        super(DQN, self).__init__()

        # Use a pre-trained model (e.g., ResNet18) as the feature extractor
        if use_pretrained:
            self.feature_extractor = models.resnet18(pretrained=True)
            self.feature_extractor.fc = nn.Identity()  # Remove the fully connected layer
        else:
            # Define your own convolutional layers for feature extraction
            self.conv1 = nn.Conv2d(in_channels=input_size[0], out_channels=32, kernel_size=8, stride=4)
            self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=4, stride=2)
            self.conv3 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=1)

            # Common layers for both cases
            self.fc1 = nn.Linear(64 * 7 * 7, 512)  # Set the correct input size
            self.fc2 = nn.Linear(512, output_size)

    def forward(self, x):
        if hasattr(self, 'feature_extractor'):  # Check if using a pre-trained model
            x = self.feature_extractor(x)
            x = x.view(x.size(0), -1)  # Flatten the output
        else:
            x = x / 255.0  # Normalize pixel values
            x = torch.as_tensor(x, dtype=torch.float32)  # Remove clone().detach()
            x = x.permute(0, 3, 1, 2)  # Change the order of dimensions (N, H, W, C) to (N, C, H, W)
            x = self.conv1(x)
            x = torch.relu(x)
            x = self.conv2(x)
            x = torch.relu(x)
            x = self.conv3(x)
            x = torch.relu(x)

            # Print the shape of the tensor at this point
            print("Shape after convolutional layers:", x.shape)

            x = torch.flatten(x, start_dim=1)  # Flatten the output

            # Update the fully connected layers with the correct input size
            self.fc1 = nn.Linear(x.size(1), 512)
            self.fc2 = nn.Linear(512, output_size)

            x = self.fc1(x)
            x = torch.relu(x)
            x = self.fc2(x)
        return x

# Define the process frame
def process_frame(frame):
    try:
        # Implement any necessary preprocessing here (e.g., resizing, cropping, converting to grayscale)
        processed_frame = cv2.resize(frame, (800, 600))
        processed_frame = cv2.cvtColor(processed_frame, cv2.COLOR_RGB2GRAY)  # Convert to grayscale
        return processed_frame
    except Exception as e:
        logger.error(f"Error processing frame: {str(e)}")
        raise

# Define the replay buffer
class ReplayBuffer:
    def __init__(self, capacity):
        self.capacity = capacity
        self.memory = []
        self.position = 0

    def push(self, transition):
        if len(self.memory) < self.capacity:
            self.memory.append(None)
        self.memory[self.position] = transition
        self.position = (self.position + 1) % self.capacity

    def sample(self, batch_size):
        return random.sample(self.memory, batch_size)

    def __len__(self):
        return len(self.memory)

# Define a named tuple for transitions
Transition = namedtuple('Transition', ('state', 'action', 'reward', 'next_state', 'done'))

# Define the DQN agent
class DQNAgent:
    def __init__(self, input_size, output_size, gamma=0.99, epsilon=1.0, epsilon_decay=0.995, epsilon_min=0.01):
        self.model = DQN(input_size, output_size)
        self.target_model = DQN(input_size, output_size)
        self.optimizer = optim.Adam(self.model.parameters(), lr=0.001)
        self.gamma = gamma
        self.epsilon = epsilon
        self.epsilon_decay = epsilon_decay
        self.epsilon_min = epsilon_min
        self.replay_buffer = ReplayBuffer(10000)

    def select_action(self, state):
        if np.random.rand() < self.epsilon:
            return np.random.randint(len(self.model.fc2.weight))
        else:
            with torch.no_grad():
                q_values = self.model(torch.tensor(state).float())
                return torch.argmax(q_values).item()

    def train_behavior_cloning(self, state_batch, action_batch):
        predicted_actions = self.model(state_batch)
        loss = nn.CrossEntropyLoss()(predicted_actions, action_batch)
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()

    def train(self, batch_size):
        if len(self.replay_buffer) < batch_size:
            return

        transitions = self.replay_buffer.sample(batch_size)
        batch = Transition(*zip(*transitions))

        state_batch = torch.tensor(np.stack(batch.state)).float()
        action_batch = torch.tensor(batch.action).long()
        reward_batch = torch.tensor(batch.reward).float()
        next_state_batch = torch.tensor(np.stack(batch.next_state)).float()
        done_batch = torch.tensor(batch.done).float()

        q_values = self.model(state_batch).gather(1, action_batch.unsqueeze(1))
        next_q_values = self.target_model(next_state_batch).max(1)[0].detach()
        expected_q_values = reward_batch + (1 - done_batch) * self.gamma * next_q_values

        loss = nn.MSELoss()(q_values, expected_q_values.unsqueeze(1))
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()

        self.epsilon = max(self.epsilon * self.epsilon_decay, self.epsilon_min)

    def update_target_model(self):
        self.target_model.load_state_dict(self.model.state_dict())

# Function to capture the game state from the screen
def get_state():
    try:
        game_window = gw.getWindowsWithTitle('Superfighters Deluxe')[0]
        if game_window:
            game_window.restore()
            game_window.activate()
            # Make sure the game window is in the foreground before capturing
            pyautogui.sleep(0.2)  # A short delay to ensure the window comes to the front
            x, y, width, height = game_window.left, game_window.top, game_window.width, game_window.height
            screenshot = pyautogui.screenshot(region=(x, y, width, height))
            state = np.array(screenshot)
            return state
        else:
            raise Exception("Game window not found")
    except Exception as e:
        logger.error(f"Error getting game state: {str(e)}")
        raise

# Function to check if the game is open
def is_game_open():
    game_window = gw.getWindowsWithTitle('Superfighters Deluxe')
    return bool(game_window)

# Function to simulate actions based on the agent's output
def simulate_action(action):
    if action == 0:
        # Randomly choose between Jump / ClimbLadder and Move to the Left or Move to the Right
        if np.random.rand() < 0.5:
            pydirectinput.keyDown('up')  
            pydirectinput.keyDown('left')
            logger.info("Simulating action: Jump / ClimbLadder and Move to the Left")
        else:
            pydirectinput.keyDown('up')
            pydirectinput.keyDown('right')
            logger.info("Simulating action: Jump / ClimbLadder and Move to the Right")
    elif action == 1:
        pydirectinput.keyDown('down')  # Crouch / SlideDownLadder
        logger.info("Simulating action: Crouch / SlideDownLadder")
    elif action == 2:
        # Move to the left
        pydirectinput.keyDown('left')
        logger.info("Simulating action: Move to the left")
    elif action == 3:
        # Move to the right
        pydirectinput.keyDown('right')
        logger.info("Simulating action: Move to the right")

    time.sleep(0.1)

    keys_to_release = ['up', 'down', 'left', 'right', 'd', 'a', 's', 'shift']
    keys_to_release.extend(map(str, range(1, 6)))
    keys_to_release.extend(['left', 'right', 'down', 'f'])

    for key in keys_to_release:
        pydirectinput.keyUp(key)

def expert_action(state):
    # This is a placeholder example for a simple expert action function.
    # The actual logic will depend on the specifics of your game.

    # Suppose the state is an image represented as a NumPy array
    # You might use basic image processing or heuristics to make decisions

    # Example: If the average pixel intensity in the top region of the screen is high, jump.
    top_region_intensity = np.mean(state[:100, :, :])
    
    # Example: If the player is close to the right edge, move left; otherwise, move right.
    player_position = np.argmax(state[:, :, 0])  # Assuming the player is represented by channel 0
    right_edge_threshold = 600  # Adjust based on your game's resolution
    if player_position > right_edge_threshold:
        action = 2  # Move left
    else:
        action = 3  # Move right

    return action

def main():
    global output_size  # Declare output_size as global

    input_size = (3, 650, 650)
    output_size = 4
    agent = DQNAgent(input_size, output_size)
    agent.model.to(device)
    agent.target_model.to(device)

    NUM_DEMONSTRATIONS = 10  # Number of expert demonstrations
    DEMONSTRATION_DELAY = 5  # Delay between expert demonstrations
    BATCH_SIZE = 32
    EPISODE_DELAY = 10
    LOG_FILE = r"C:\Users\Ninja.000\Downloads\SFD_AI\training_log.txt"
    MODEL_SAVE_PATH = r"C:\Users\Ninja.000\Downloads\SFD_AI\dqn_model.pth"
    BEST_MODEL_SAVE_PATH = r"C:\Users\Ninja.000\Downloads\SFD_AI\best_dqn_model.pth"

    # Flag to indicate expert demonstration phase
    expert_demonstration_phase = True
    train_agent = False  # Set to False for demonstration phase

    # Add a console handler to the logger
    console_handler = logging.StreamHandler()
    console_handler.setLevel(logging.INFO)
    logger.addHandler(console_handler)

    MAX_EPISODE_LENGTH = 100  # Set an appropriate value based on your application
    TUTORIAL_FOLDER = r"C:\Users\Ninja.000\Downloads\SFD_AI\gameplay\Tutorial"

    try:
        # Load the model and its parameters if it exists
        if os.path.exists(MODEL_SAVE_PATH):
            agent.model.load_state_dict(torch.load(MODEL_SAVE_PATH))

            # Extract episode number from the model path
            match = re.search(r'\d+', os.path.basename(MODEL_SAVE_PATH))
            episode_start = int(match.group()) if match else 0
            logger.info(f"Resuming training from Episode {episode_start + 1}")
        else:
            episode_start = 0
            logger.info("Starting training from Episode 1")

        # Load the best model parameters if it exists
        if os.path.exists(BEST_MODEL_SAVE_PATH):
            agent.best_model.load_state_dict(torch.load(BEST_MODEL_SAVE_PATH))
            best_reward = max(np.load(LOG_FILE), default=float('-inf'))
            logger.info(f"Resuming training with the best reward: {best_reward}")
        else:
            best_reward = float('-inf')
            logger.info("Starting training with no best reward information")

        episode_rewards = []

        # Expert demonstration phase
        if expert_demonstration_phase:
            for demonstration in range(NUM_DEMONSTRATIONS):
                logger.info(f"Expert Demonstration {demonstration + 1}")

                frames = []  # Collect frames for processing
                for t in range(1, MAX_EPISODE_LENGTH, 1):
                    state = get_state()
                    frames.append(state)
                    time.sleep(0.1)

                    # Save the processed frame to the tutorial folder with an available frame number
                    frame_number = 1
                    while True:
                        frame_path = os.path.join(TUTORIAL_FOLDER, f"frame_{frame_number}.jpg")
                        if not os.path.exists(frame_path):
                            break
                        frame_number += 1

                    processed_frame = process_frame(state)  # Apply processing to the frame
                    cv2.imwrite(frame_path, processed_frame)

                for t in range(len(frames) - 1):
                    next_state = frames[t + 1]
                    action = expert_action(next_state)  # Modified to take next_state as input

                    # During expert demonstration phase, actions are not simulated
                    state = frames[t]

                    # Save experiences for later training only if train_agent is True
                    if train_agent:
                        agent.replay_buffer.push(Transition(state, action, 0, next_state, False))  # No reward during demonstration

                time.sleep(DEMONSTRATION_DELAY)

        # Training phase (Skipped if expert_demonstration_phase is True)
        if not expert_demonstration_phase and train_agent:
            for episode in range(episode_start, episode_start + NUM_EPISODES):
                logger.info(f"Episode {episode + 1}")

                total_reward = 0
                experiences = []  # Collect experiences for training
                action = None  # Initialize action variable

                frames = []  # Collect frames for processing
                for t in range(1, MAX_EPISODE_LENGTH, 1):
                    frame_path = f"C:\\Users\\Ninja.000\\Downloads\\SFD_AI\\gameplay\\Tutorial\\frame_{t}.jpg"
                    logger.info(f"Observing frame: {frame_path}")

                    # Process the frame
                    state = process_frame(frame_path)
                    frames.append(state)

                for t in range(len(frames) - 1):
                    next_state = frames[t + 1]

                    if expert_demonstration_phase:
                        # During expert demonstration phase, use expert actions
                        action = expert_action(next_state)
                    else:
                        # During training phase, use DQN actions
                        action = agent.select_action(state)

                        # Simulate action only during the training phase
                        simulate_action(action)

                    state = frames[t]
                    reward, done = calculate_reward_and_check_termination(state, next_state)
                    agent.replay_buffer.push(Transition(state, action, reward, next_state, done))

                    # Process the frame
                    state = process_frame(frame_path)
                    frames.append(state)

                for t in range(len(frames) - 1):
                    # Capture the next frame for the experience
                    next_state = frames[t + 1]

                    # DQN action
                    action = agent.select_action(state)
                    simulate_action(action)

                    # Replace the placeholder reward and termination condition with actual logic
                    reward, done = calculate_reward_and_check_termination(state, next_state)
                    total_reward += reward
                    experiences.append(Transition(state, action, reward, next_state, done))

                episode_rewards.append(total_reward)
                logger.info(f"Episode {episode + 1}, Total Reward: {total_reward}")

                # Train the agent using collected experiences
                for _ in range(BATCH_SIZE):
                    agent.replay_buffer.push(random.choice(experiences))
                    agent.train(BATCH_SIZE)
                    agent.update_target_model()

                # Save the model parameters at the end of each episode
                torch.save(agent.model.state_dict(), MODEL_SAVE_PATH)

                # Check if this episode achieved the best reward
                if total_reward > best_reward:
                    best_reward = total_reward
                    torch.save(agent.model.state_dict(), BEST_MODEL_SAVE_PATH)

                # Save the best reward information
                np.save(LOG_FILE, [best_reward])

                # Clear experiences for the next episode
                experiences.clear()

                time.sleep(EPISODE_DELAY)

    except KeyboardInterrupt:
        logger.info("Training interrupted by user. Saving the model before exiting...")
        torch.save(agent.model.state_dict(), MODEL_SAVE_PATH)
        np.save(LOG_FILE, [best_reward])  # Save the best reward information
    except Exception as e:
        logger.error(f"An error occurred: {str(e)}")
    finally:
        exit()

if __name__ == "__main__":
    main()