| | import pygame
|
| | import time
|
| | import numpy as np
|
| | from stable_baselines3 import PPO
|
| | from gymnasium import spaces
|
| | import gymnasium as gym
|
| |
|
| |
|
| |
|
| |
|
| | class PixelCopterCertEnv(gym.Env):
|
| | def __init__(self, screen_width=50, screen_height=10, gap_size=6):
|
| | super().__init__()
|
| | self.screen_width = screen_width
|
| | self.screen_height = screen_height
|
| | self.copter_y = self.screen_height // 2
|
| | self.copter_velocity = 0
|
| | self.gravity = 0.25
|
| | self.lift = -0.9
|
| | self.done = False
|
| | self.timestep = 0
|
| | self.max_timesteps = 500
|
| | self.gap_size = gap_size
|
| | self.wall_gap_positions = [np.random.randint(1, self.screen_height - self.gap_size -1)
|
| | for _ in range(screen_width)]
|
| | self.action_space = spaces.Discrete(2)
|
| | self.observation_space = spaces.Box(
|
| | low=0, high=self.screen_height, shape=(self.screen_width + 1,), dtype=np.float32
|
| | )
|
| |
|
| | def reset(self, seed=None, options=None):
|
| | self.copter_y = self.screen_height // 2
|
| | self.copter_velocity = 0
|
| | self.done = False
|
| | self.timestep = 0
|
| | self.wall_gap_positions = [np.random.randint(1, self.screen_height - self.gap_size -1)
|
| | for _ in range(self.screen_width)]
|
| | obs = np.array([self.copter_y] + self.wall_gap_positions, dtype=np.float32)
|
| | return obs, {}
|
| |
|
| | def step(self, action):
|
| | if action == 1:
|
| | self.copter_velocity += self.lift
|
| | self.copter_velocity += self.gravity
|
| | self.copter_y += self.copter_velocity
|
| | self.copter_y = np.clip(self.copter_y, 0, self.screen_height)
|
| |
|
| | self.wall_gap_positions = self.wall_gap_positions[1:]
|
| | last_gap = self.wall_gap_positions[-1]
|
| | new_gap = last_gap + np.random.choice([-1,0,1])
|
| | new_gap = np.clip(new_gap, 1, self.screen_height - self.gap_size -1)
|
| | self.wall_gap_positions.append(new_gap)
|
| |
|
| | gap_top = self.wall_gap_positions[0]
|
| | gap_bottom = gap_top + self.gap_size
|
| | if self.copter_y <= gap_top or self.copter_y >= gap_bottom:
|
| | self.done = True
|
| | reward = -5
|
| | else:
|
| | reward = 1
|
| |
|
| | self.timestep += 1
|
| | if self.timestep >= self.max_timesteps:
|
| | self.done = True
|
| |
|
| | obs = np.array([self.copter_y] + self.wall_gap_positions, dtype=np.float32)
|
| | return obs, reward, self.done, False, {}
|
| |
|
| |
|
| |
|
| |
|
| | pygame.init()
|
| | WINDOW_WIDTH = 800
|
| | WINDOW_HEIGHT = 400
|
| | WIN = pygame.display.set_mode((WINDOW_WIDTH, WINDOW_HEIGHT))
|
| | pygame.display.set_caption("PixelCopter Certification Test")
|
| | WHITE = (255,255,255)
|
| | BLUE = (0,102,255)
|
| | BLACK = (0,0,0)
|
| |
|
| | def draw_screen(obs, env):
|
| | WIN.fill(WHITE)
|
| | scale_y = WINDOW_HEIGHT / env.screen_height
|
| | scale_x = WINDOW_WIDTH / env.screen_width
|
| | copter_y = int(obs[0])
|
| | wall_gaps = [int(gap) for gap in obs[1:]]
|
| |
|
| |
|
| | for i, gap_top in enumerate(wall_gaps):
|
| | gap_bottom = gap_top + env.gap_size
|
| | pygame.draw.rect(WIN, BLACK, pygame.Rect(i*scale_x, 0, scale_x, gap_top*scale_y))
|
| | pygame.draw.rect(WIN, BLACK, pygame.Rect(i*scale_x, gap_bottom*scale_y, scale_x, WINDOW_HEIGHT - gap_bottom*scale_y))
|
| |
|
| |
|
| | pygame.draw.rect(WIN, BLUE, pygame.Rect(50, copter_y*scale_y - 10, 20, 20))
|
| | pygame.display.update()
|
| |
|
| |
|
| |
|
| |
|
| | env = PixelCopterCertEnv(screen_width=80, screen_height=10, gap_size=6)
|
| | model = PPO.load("ppo_pixelcopter_cert")
|
| |
|
| | episodes = 5
|
| | for ep in range(episodes):
|
| | obs, _ = env.reset()
|
| | done = False
|
| | total_reward = 0
|
| | while not done:
|
| | for event in pygame.event.get():
|
| | if event.type==pygame.QUIT:
|
| | done=True
|
| | pygame.quit()
|
| | action, _ = model.predict(obs)
|
| | obs, reward, done, truncated, info = env.step(action)
|
| | total_reward += reward
|
| | draw_screen(obs, env)
|
| | time.sleep(0.03)
|
| | print(f"Episode {ep+1} reward: {total_reward}")
|
| |
|
| | pygame.quit()
|
| |
|
| |
|
| |
|
| |
|
| |
|
| |
|
| |
|
