env.py

import gymnasium as gym
from gymnasium import spaces
import numpy as np
from dino import Dino
import asyncio
import time
from scipy.ndimage import convolve
import numpy as np
import gymnasium as gym
from tensorboardX import SummaryWriter
from datetime import datetime
from PIL import Image
from datetime import datetime

class DinoEnv(gym.Env):
    def __init__(self):
        super(DinoEnv, self).__init__()

        # Define the action and observation space
        # Actions: 0=up, 1=down, 2 do nothing
        self.action_space = spaces.Discrete(3)
        
        self.sobel_kernel = np.array([
            [-1, 0, 1],
            [-2, 0, 2],
            [-1, 0, 1]
        ])

        self.im_size = 9000
        
        # Define the observation space: 11316 integer pixels
        self.observation_space = spaces.Box(low=0, high=255, shape=(self.im_size,), dtype=np.int32)

        # Initialize the state (e.g., all pixels set to 0)
        self.state = np.zeros((self.im_size,), dtype=np.int32)

        # Initialize other variables (e.g., a variable to keep track of the score)
        self.score = 0

        self.current_step = 0

    def step(self, action):
        print( str(action) + ' : ' + str(self.score))
        self.current_step += 1
        if action == 0:  # up
            #asyncio.get_event_loop().run_until_complete(self.dino.send_key_event2("down"))
            asyncio.get_event_loop().run_until_complete(self.dino.send_key_event("ArrowUp", "ArrowUp", 38))
            asyncio.get_event_loop().run_until_complete(self.dino.complete_action())
        elif action == 1:  # down
            #asyncio.get_event_loop().run_until_complete(self.dino.send_key_event2("down"))
            asyncio.get_event_loop().run_until_complete(self.dino.send_key_event("ArrowDown", "ArrowDown", 40))
            asyncio.get_event_loop().run_until_complete(self.dino.complete_action())
        elif action == 2:  # do nothing
            time.sleep(0.1)
            pass

        self.state = self.get_screenshot()

        status = asyncio.get_event_loop().run_until_complete(self.dino.check_status())
        
        if not status or status['crashed']:
            with open('scores.txt', 'a') as f:
               now = datetime.now() # current date and time
               date_time = now.strftime("%m/%d/%Y, %H:%M:%S")
               print("---------date and time:",date_time)
               f.write(date_time + ' - '+ str(self.score) + '\n')
            reward = -100
            done = True
            if not status:
                asyncio.get_event_loop().run_until_complete(self.dino.open_dino())
        else:
            reward = 2
            #if action == 1: reward = 3
            done = False

        self.score+=reward

        info = {}

        #observation, reward, terminated, truncated, info
        return self.state, reward, done, False, info
        

    def reset(self, seed=None, options=None):
        super().reset(seed=seed)

        self.dino = Dino('runner-canvas')
        
        asyncio.get_event_loop().run_until_complete(self.dino.start())
        
        asyncio.get_event_loop().run_until_complete(self.dino.open_dino())
        
        time.sleep(2)
        
        asyncio.get_event_loop().run_until_complete(self.dino.send_key_event(" ", "Space", 32))
        
        time.sleep(2)

        asyncio.get_event_loop().run_until_complete(self.dino.capture_screenshot2())

        self.state = self.get_screenshot()

        info = {}

        self.current_step = 0

        self.score = 0
        
        asyncio.get_event_loop().run_until_complete(self.dino.enable_all_obstacles())

        return self.state, info
    
    def get_screenshot(self):
        image = asyncio.get_event_loop().run_until_complete(self.dino.capture_screenshot2())

        gray_image = image.convert('L')

        image_array = np.array(gray_image)

        convolved_image = convolve(image_array, self.sobel_kernel)
        # Get the current date and time
        current_time = datetime.now()

        # Format the date and time as a string
        timestamp_string = current_time.strftime('%H:%M:%S')

        image = Image.fromarray(convolved_image)
    
        # Save the image
        #image.save('im/' + timestamp_string + 'conv_image.png')

        flattened_array = convolved_image.flatten()

        return flattened_array

    def render(self, mode='human'):
       pass
        

# Register the custom environment
gym.envs.registration.register(
    id='DinoEnv-v0',
    entry_point=__name__ + ':DinoEnv',
)