import numpy as np
import cv2
import mss
import pyautogui
import time
import keyboard
from stable_baselines3 import DQN
from stable_baselines3.common.env_checker import check_env
import gymnasium as gym
from gymnasium import spaces


class GeometryDashEnv(gym.Env):
    def __init__(self):
        super(GeometryDashEnv, self).__init__()
        self.screen_width = 3840
        self.screen_height = 2160
        self.image_width = 128  # Размер изображения
        self.image_height = 128
        self.monitor = {
            "top": 0,
            "left": self.screen_width - self.screen_height,
            "width": self.screen_height,
            "height": self.screen_height,
        }
        self.action_space = spaces.Discrete(2)
        self.observation_space = spaces.Box(
            low=0,
            high=255,
            shape=(self.image_height, self.image_width, 1),
            dtype=np.uint8,
        )
        self.sct = mss.mss()
        self.done = False
        self.reward = 0
        self.total_reward = 0
        self.prev_obs = None
        self.no_change_frames = 0
        self.no_change_threshold = 10
        self.change_threshold = 0.001
        self.survival_time = 0  # Счетчик времени выживания

    def reset(self, *, seed=None, options=None):
        super().reset(seed=seed)
        self.restart_game()
        obs = self.get_observation()
        self.prev_obs = obs.copy()
        self.done = False
        self.total_reward = 0
        self.no_change_frames = 0
        self.survival_time = 0
        info = {}
        return obs, info

    def step(self, action):
        self.perform_action(action)
        obs = self.get_observation()
        self.done = self.is_game_over(obs)
        additional_reward = 0
        negative_reward = 0
        if self.done:
            self.reward = -10
        else:
            self.survival_time += 1
            self.reward = (
                np.log(self.survival_time + 1) + additional_reward + negative_reward
            )
        self.total_reward += self.reward
        truncated = False
        info = {}
        self.prev_obs = obs.copy()
        return obs, self.reward, self.done, truncated, info

    def get_observation(self):
        img = np.array(self.sct.grab(self.monitor))
        img = cv2.cvtColor(img, cv2.COLOR_BGRA2GRAY)
        img = cv2.resize(img, (self.image_width, self.image_height))
        img = np.expand_dims(img, axis=2)
        return img

    def perform_action(self, action):
        if action == 1:
            pyautogui.mouseDown(button="left")
            pyautogui.mouseUp(button="left")

    def is_game_over(self, obs):
        if self.prev_obs is not None:
            diff = np.mean((obs - self.prev_obs) ** 2) / 255.0
            if diff < self.change_threshold:
                self.no_change_frames += 1
            else:
                self.no_change_frames = 0
            if self.no_change_frames >= self.no_change_threshold:
                return True
            else:
                return False
        else:
            return False

    def restart_game(self):
        pyautogui.press("space")

    def render(self):
        pass

    def close(self):
        pass


time.sleep(10)
env = GeometryDashEnv()
check_env(env)
model = DQN("CnnPolicy", env, verbose=1, tensorboard_log="./geometry_dash_logs/")
model.learn(total_timesteps=100000)
model.save("geometry_dash_dqn_model")
model = DQN.load("geometry_dash_dqn_model", env=env)
obs, info = env.reset()

while True:
    action, _states = model.predict(obs)
    obs, reward, done, truncated, info = env.step(action)
    if done:
        obs, info = env.reset()

Я бы хотел сделать здесь новую систему поощрения
Я бы хотел сохранять время, которое выжил персонаж в игре - время от старта до ресета и хранить максимальное такое время. А дальше давать поощрение или наказание за каждые 200мс не пройденные или пройденные