environment.py

from nes_py.wrappers import JoypadSpace
from gym_super_mario_bros import make
from gym_super_mario_bros.actions import SIMPLE_MOVEMENT
import cv2
import numpy as np

def create_environment():
    """
    Initialize the Super Mario Bros environment.
    """
    env = make('SuperMarioBros-1-1-v0')
    env = JoypadSpace(env, SIMPLE_MOVEMENT)
    return env

def preprocess_observation(observation):
    """
    Preprocess the observation (game screen) for input to the neural network.
    """
    gray_obs = cv2.cvtColor(observation, cv2.COLOR_BGR2GRAY)  # Convert to grayscale
    resized_obs = cv2.resize(gray_obs, (84, 84))  # Resize to 84x84
    normalized_obs = resized_obs / 255.0  # Normalize pixel values to [0, 1]
    return np.expand_dims(normalized_obs, axis=0)  # Add channel dimension