Game2.py

import pygame
import numpy as np
import Main2
import importlib
import time


importlib.reload(Main2)
class RocketLandingGameRunner:
    def __init__(self, width=800, height=600, render=False):
        """Initialize the game runner with screen dimensions and rendering option."""
        self.width = width
        self.height = height
        self.step_duration = 0.1  # Time step for each action
        self.frame_time = 1.0 / 60  # Frame time for rendering (60 FPS)
        self.white = (255, 255, 255)
        self.render_enabled = render
        self.game = Main2.RocketLandingGame(self.render_enabled)  # Initialize the RocketLandingGame instance

        # Define state and action spaces
        self.state_space = {
            'x': (0, self.width),  # Position x
            'y': (0, self.height),  # Position y
            'velocity_x': (-50, 50),  # Velocity bounds (adjust as needed)
            'velocity_y': (-50, 50),
            'angle': (-180, 180)  # Angle in degrees
        }
        self.action_space = [0, 1, 2, 3]  # 0: Nothing, 1: Thrust, 2: Rotate left, 3: Rotate right

    def reset(self):
        """Reset the game to initial state and return initial state."""
        self.game.rocket = Main2.Rocket()
        self.game.game_over = False
        self.game.win = False
        self.game.running = True
        self.game.action_log = []
        return self._get_state()

    def _get_state(self):
        """Return the current state as a numpy array without exponent format."""
        state = np.array([
            self.game.rocket.x,
            self.game.rocket.y,
            self.game.rocket.velocity_x,
            self.game.rocket.velocity_y,
            self.game.rocket.angle
        ], dtype=np.float32)
        
        return np.round(state, 1)


    def step(self, action):
        """Take an action, update the game, and return (next_state, reward, done, info)."""
        # Map action to game inputs
        thrusting = (action == 1)
        rotate_left = (action == 2)
        rotate_right = (action == 3)

        # Log the action
        if action == 1:
            self.game.add_to_action_log("Thrust activated.")
        elif action == 2:
            self.game.add_to_action_log("Rotation left.")
        elif action == 3:
            self.game.add_to_action_log("Rotation right.")
        else:
            self.game.add_to_action_log("Doing nothing.")

        # Update game state for one step
        reward, game_over, win = self.game.game_step(thrusting, rotate_left, rotate_right, time_step=self.step_duration)

        # Get next state
        next_state = self._get_state()

        # Determine done flag
        done=game_over
        # Info dictionary for additional data
        info = {'win': win, 'action_log': self.game.action_log[-1]}

        # Render if enabled
        if self.render_enabled:
            self._render()

        return next_state, reward, done

    def _render(self):
        """Render the game state."""
        self.game.screen.fill(self.white)

        self.game.draw_terrain()
        self.game.rocket.draw(self.game.screen)

        self.game.draw_info()

        if self.game.game_over:
            if self.game.win:
                win_text = pygame.font.SysFont('Arial', 40).render(
                    "You Landed Safely!", True, (0, 128, 0)
                )
                self.game.screen.blit(
                    win_text, (self.width // 2 - win_text.get_width() // 2, self.height // 2 - 100)
                )
            else:
                crash_text = pygame.font.SysFont('Arial', 40).render(
                    "Crash! Try Again.", True, (255, 0, 0)
                )
                self.game.screen.blit(
                    crash_text, (self.width // 2 - crash_text.get_width() // 2, self.height // 2 - 100)
                )
            self.game.draw_restart_button()
        if self.render_enabled:
            pygame.display.update()
            self.game.clock.tick(30)
            pygame.event.pump()

    def run(self, agent=None):
        if agent is None:
            # Original human-controlled game loop
            while self.game.running:
                step_start_time = time.time()
                thrusting = False
                rotate_left = False
                rotate_right = False
                reset_game = False
                action_logged = False

                for event in pygame.event.get():
                    if event.type == pygame.QUIT:
                        self.game.running = False
                    if event.type == pygame.MOUSEBUTTONDOWN and self.game.game_over:
                        mouse_pos = pygame.mouse.get_pos()
                        if self.game.draw_restart_button().collidepoint(mouse_pos):
                            reset_game = True
                            self.game.add_to_action_log("Game reset by user.")
                            action_logged = True

                if not self.game.game_over and not reset_game:
                    keys = pygame.key.get_pressed()
                    if keys[pygame.K_w] and not action_logged:
                        thrusting = True
                        self.game.add_to_action_log("Thrust activated.")
                        action_logged = True
                    elif keys[pygame.K_a] and not action_logged:
                        rotate_left = True
                        self.game.add_to_action_log("Rotation left.")
                        action_logged = True
                    elif keys[pygame.K_d] and not action_logged:
                        rotate_right = True
                        self.game.add_to_action_log("Rotation right.")
                        action_logged = True
                    elif keys[pygame.K_r] and not action_logged:
                        reset_game = True
                        self.game.add_to_action_log("Game reset by 'R' key.")
                        action_logged = True

                if not action_logged:
                    self.game.add_to_action_log("Doing nothing.")

                if reset_game:
                    self.reset()

                while time.time() - step_start_time < self.step_duration:
                    if not reset_game and not self.game.game_over:
                        reward, game_over, win = self.game.game_step(
                            thrusting, rotate_left, rotate_right, time_step=self.frame_time
                        )
                        print(reward)
                    else:
                        reward, game_over, win = 0, self.game.game_over, self.game.win
                    if self.render_enabled:
                        self._render()

            print("\nAction Log:")
            for entry in self.game.action_log:
                print(entry)
            pygame.quit()

        else:
            # RL agent-controlled loop
            state = self.reset()
            total_reward = 0
            np.set_printoptions(suppress=True)
            while self.game.running:
                # if self.render_enabled:
                state_input = np.expand_dims(state, axis=0)  # make it shape (1,5)
                action_values = agent(state_input)            # get q_values
                action = np.argmax(action_values.numpy()[0])  # pick best action
                # action = np.random.choice(self.action_space)
                next_state, reward, done = self.step(action)
                total_reward += reward
                # agent.update(state, action, reward, next_state, done)  # Update agent's policy
                state = next_state
                if done:
                    break
            start_time = time.time()
            button_rect = self.game.draw_restart_button()

            # Keep checking for mouse events during the 5-second window
            while time.time() - start_time < 5:
                for event in pygame.event.get():
                    if event.type == pygame.MOUSEBUTTONDOWN:
                        mouse_pos = event.pos
                        if button_rect.collidepoint(mouse_pos):       
                            # If the restart button is clicked, restart the game
                            game_runner2 = RocketLandingGameRunner(width=800, height=600, render=True)
                            game_runner2.run(agent)
                            return  # Exit to avoid pygame.quit() being called twice
                    elif event.type == pygame.QUIT:
                        pygame.quit()
                        return

                pygame.display.update()  # Make sure the display is updated

            # If no click within 5 seconds, quit the game
            pygame.quit()

    def get_action_log(self):
        """Return the current action log."""
        return self.game.action_log