PushPullEnvironment.py

import os
import sys
import cv2
import pygame
import numpy as np
import threading
import queue
import gymnasium as gym
from gymnasium import spaces

WIDTH, HEIGHT = 1000, 500
FPS = 60
BAR_MAX, BAR_MIN = 100.0, 0.0
STARTING_VALUE = 50.0
DECREASE_RATE, INCREASE_BUMP = 12.0, 12.0
WHITE, BLACK, RED, GRAY = (255, 255, 255), (20, 20, 20), (231, 76, 60), (149, 165, 166)
VIDEO_FILENAME = "hentai.mp4"

class PushPullEnv(gym.Env):
    def __init__(self, render_mode=None):
        super().__init__()
        self.render_mode = render_mode
        self.action_space = spaces.Discrete(2)
        self.observation_space = spaces.Box(
            low=np.array([BAR_MIN, 0.0, 0.0], dtype=np.float32),
            high=np.array([BAR_MAX, 1.0, 100.0], dtype=np.float32),
            dtype=np.float32
        )
        
        pygame.init()
        self.screen = pygame.display.set_mode((WIDTH, HEIGHT))
        self.clock = pygame.time.Clock()
        
        self.hand_sheet = pygame.image.load("hand.png").convert_alpha()
        sprite_w, sprite_h = self.hand_sheet.get_width() // 2, self.hand_sheet.get_height()
        self.sprite_w, self.sprite_h = sprite_w * 9, sprite_h * 9
        
        self.push_img = pygame.transform.scale(pygame.transform.flip(self.hand_sheet.subsurface((0, 0, sprite_w, sprite_h)), True, False), (self.sprite_w, self.sprite_h))
        self.pull_img = pygame.transform.scale(pygame.transform.flip(self.hand_sheet.subsurface((sprite_w, 0, sprite_w, sprite_h)), True, False), (self.sprite_w, self.sprite_h))

        self.video_path = os.path.join(os.getcwd(), VIDEO_FILENAME)
        self.frame_queue = queue.Queue(maxsize=30)
        self.stop_video, self.current_bg = False, None
        
        if os.path.exists(self.video_path):
            self.video_thread = threading.Thread(target=self._video_worker, daemon=True)
            self.video_thread.start()

        self.reset()

    def _video_worker(self):
        cap = cv2.VideoCapture(self.video_path)
        while not self.stop_video:
            if not self.frame_queue.full():
                ret, frame = cap.read()
                if not ret:
                    cap.set(cv2.CAP_PROP_POS_FRAMES, 0)
                    continue
                frame = cv2.resize(frame, (WIDTH, HEIGHT))
                frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB).swapaxes(0, 1)
                self.frame_queue.put(pygame.surfarray.make_surface(frame))
            else:
                pygame.time.wait(10)
        cap.release()

    def reset(self, *, seed=None, options=None):
        super().reset(seed=seed, options=options)
        self.bar_value, self.current_state = STARTING_VALUE, 0
        self.time_survived, self.time_since_last_switch = 0.0, 0.0
        return np.array([self.bar_value, self.current_state, self.time_since_last_switch], dtype=np.float32), {}

    def step(self, action):
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                self.stop_video = True
                pygame.quit()
                sys.exit()

        dt = self.clock.tick(FPS) / 1000.0 if self.render_mode == "human" else 1.0 / FPS
        self.time_survived += dt
        self.time_since_last_switch += dt
        reward = (self.bar_value - 50.0) / 50.0

        if action != self.current_state:
            if self.time_since_last_switch < 0.7:
                reward -= 1.0
                self.bar_value -= DECREASE_RATE * dt
            else:
                self.bar_value += INCREASE_BUMP
                self.current_state = action
                self.time_since_last_switch = 0.0
        else:
            self.bar_value -= DECREASE_RATE * dt

        terminated = self.bar_value <= BAR_MIN or self.bar_value >= BAR_MAX
        self.bar_value = max(BAR_MIN, min(BAR_MAX, self.bar_value))

        obs = np.array([self.bar_value, self.current_state, self.time_since_last_switch], dtype=np.float32)

        return obs, reward, terminated, False, {"time_survived": self.time_survived}

    def render(self):
        if not self.frame_queue.empty():
            self.current_bg = self.frame_queue.get()
            
        if self.current_bg:
            self.screen.blit(self.current_bg, (0, 0))
        else:
            self.screen.fill(BLACK)

        pygame.draw.rect(self.screen, GRAY, (50, 50, 60, HEIGHT - 100), 3)
        fill_h = (self.bar_value / BAR_MAX) * (HEIGHT - 100)
        pygame.draw.rect(self.screen, RED, (50, 50 + (HEIGHT - 100) - fill_h, 60, fill_h))

        img = self.push_img if self.current_state == 1 else self.pull_img
        self.screen.blit(img, (WIDTH - self.sprite_w, HEIGHT - self.sprite_h))

        pygame.display.flip()