Spaces:
Running on Zero
Running on Zero
| import gymnasium as gym | |
| import numpy as np | |
| import pygame | |
| class BirdAttackEnv(gym.Env): | |
| metadata = {"render_modes": ["rgb_array"], "render_fps": 30} | |
| def __init__(self, max_episode_steps=500): | |
| super().__init__() | |
| self.max_episode_steps = max_episode_steps | |
| self.width, self.height = 224, 224 | |
| self.area_size = 184 | |
| self.area_x = (self.width - self.area_size) // 2 | |
| self.area_y = (self.height - self.area_size) // 2 | |
| # agent | |
| self.agent_w, self.agent_h = 14, 8 | |
| self.agent_x = self.width // 2 - self.agent_w // 2 | |
| self.agent_y = self.area_y + self.area_size - 20 | |
| self.agent_speed = 6.0 | |
| self.reload_timer = 8 | |
| # birds | |
| self.bird_w, self.bird_h = 14, 12 | |
| self.bird_rows, self.bird_cols = 2, 8 | |
| self.birds = [] | |
| self.bird_dir = 1 # 1=right, -1=left | |
| self.bird_speed = 1.0 | |
| self.bird_drop = 4 | |
| # boss bird | |
| self.boss_spawn_step = 250 | |
| self.boss_active = False | |
| self.boss_pos = None # [x, y] | |
| self.boss_bullet = None # [x, y, vx, vy] | |
| self.boss_fire_timer = 0 | |
| self.boss_size = 28 | |
| self.boss_bullet_size = 10 | |
| self.boss_bullet_speed = 5.0 | |
| # Boss geometry + fade | |
| self.boss_span = 64 # wingtip-to-wingtip (px) | |
| self.boss_height = 16 # vertical amplitude of the "M" | |
| self.boss_fadein = 0 # frames remaining while fading in | |
| # projectiles | |
| self.agent_bullets = [] | |
| self.bird_bullets = [] | |
| self.bullet_speed = 7 | |
| self.max_agent_bullets = 5 | |
| self.max_bird_bullets = 10 | |
| max_birds = self.bird_rows * self.bird_cols | |
| obs_dim = ( | |
| 3 # [agent_x, reload, num_birds] | |
| + max_birds # bird x's | |
| + 1 # lowest bird y | |
| + self.max_agent_bullets*2 # agent bullets (x,y)×5 | |
| + self.max_bird_bullets*2 # bird bullets (x,y)×10 | |
| + 1 # stun | |
| + 4 # boss bullet (x,y,vx,vy) | |
| ) | |
| self.observation_space = gym.spaces.Box( | |
| low=-np.inf, high=np.inf, shape=(obs_dim,), dtype=np.float32 | |
| ) | |
| # continuous horizontal thrust in [-1,1] + fire button [-1,1] | |
| self.action_space = gym.spaces.Box( | |
| low=-np.ones(2, dtype=np.float32), high=np.ones(2, dtype=np.float32) | |
| ) | |
| self.score = 0 | |
| self.steps = 0 | |
| self.stun_timer = 0 | |
| self._surface = None | |
| self.font = None | |
| self.clock = None | |
| def reset(self, *, seed=None, options=None): | |
| super().reset(seed=seed) | |
| self.agent_x = self.width // 2 - self.agent_w // 2 | |
| self.reload_timer = 0 | |
| self.agent_bullets.clear() | |
| self.bird_bullets.clear() | |
| self.score = 0 | |
| self.steps = 0 | |
| self.stun_timer = 0 | |
| self.boss_active = False | |
| self.boss_pos = None | |
| self.boss_bullet = None | |
| self.boss_fire_timer = 0 | |
| # init birds in grid | |
| self.birds = [] | |
| start_x = self.area_x + 20 | |
| start_y = self.area_y + 20 | |
| spacing_x = 20 | |
| spacing_y = 20 | |
| for r in range(self.bird_rows): | |
| for c in range(self.bird_cols): | |
| self.birds.append([start_x + c*spacing_x, | |
| start_y + r*spacing_y, | |
| 2]) # hp = 2 | |
| self.bird_dir = 1 | |
| return self._get_obs(), {} | |
| def step(self, action): | |
| action = np.clip(action, -1, 1) | |
| move, fire = float(action[0]), float(action[1]) | |
| reward = 0.0 | |
| terminated, truncated = False, False | |
| # move agent | |
| speed = self.agent_speed | |
| if self.stun_timer > 0: | |
| speed *= 0.2 # lower speed when stunned | |
| self.agent_x += move * speed | |
| self.agent_x = np.clip(self.agent_x, self.area_x, | |
| self.area_x + self.area_size - self.agent_w) | |
| if self.stun_timer > 0: | |
| self.stun_timer -= 1 | |
| # fire | |
| if fire > 0 and self.reload_timer == 0: | |
| muzzle_x = self.agent_x + self.agent_w // 2 | |
| muzzle_y = self.agent_y - 6 | |
| self.agent_bullets.append([muzzle_x, muzzle_y]) | |
| self.reload_timer = 8 | |
| if self.reload_timer > 0: | |
| self.reload_timer -= 1 | |
| # move agent bullets | |
| for b in self.agent_bullets: | |
| b[1] -= self.bullet_speed | |
| self.agent_bullets = [b for b in self.agent_bullets if b[1] > self.area_y] | |
| # bird movement | |
| shift = self.bird_speed * self.bird_dir | |
| xs = [bx for bx, by, hp in self.birds] | |
| if xs: | |
| if max(xs) + self.bird_w + shift > self.area_x + self.area_size or min(xs) + shift < self.area_x: | |
| self.bird_dir *= -1 | |
| for b in self.birds: | |
| b[1] += self.bird_drop | |
| else: | |
| for b in self.birds: | |
| b[0] += shift | |
| # birds shoot randomly | |
| for bx, by, hp in self.birds: | |
| if self.np_random.random() < 0.0075: | |
| self.bird_bullets.append([bx + self.bird_w//2, by + self.bird_h]) | |
| # move bird bullets | |
| for b in self.bird_bullets: | |
| b[1] += self.bullet_speed | |
| self.bird_bullets = [b for b in self.bird_bullets if b[1] < self.area_y + self.area_size] | |
| new_birds = [] | |
| for bx, by, hp in self.birds: | |
| for bullet in self.agent_bullets: | |
| if abs(bullet[0] - (bx+self.bird_w/2)) < self.bird_w/2 and \ | |
| abs(bullet[1] - (by+self.bird_h/2)) < self.bird_h/2: | |
| hp -= 1 | |
| self.agent_bullets.remove(bullet) | |
| if hp == 0: | |
| reward += 1.0 | |
| self.score += 1 | |
| break | |
| if hp > 0: | |
| new_birds.append([bx, by, hp]) | |
| self.birds = new_birds | |
| # collisions: bird bullets vs agent | |
| ax, ay = self.agent_x, self.agent_y | |
| for bullet in self.bird_bullets: | |
| if abs(bullet[0] - (ax+self.agent_w/2)) < self.agent_w/2 and \ | |
| abs(bullet[1] - (ay+self.agent_h/2)) < self.agent_h/2: | |
| reward -= 2.0 | |
| self.bird_bullets.remove(bullet) | |
| self.stun_timer = 16 # 16-frame stun | |
| self.reload_timer = 16 # reset reload | |
| if not self.boss_active and self.steps >= self.boss_spawn_step: | |
| self.boss_active = True | |
| cx = self.area_x + self.area_size // 2 | |
| cy = self.area_y + 20 | |
| self.boss_pos = [cx, cy] # store as CENTER (cx, cy) | |
| self.boss_fire_timer = 50 | |
| self.boss_fadein = 5 # 5 frames: no move/fire | |
| if self.boss_active: | |
| if self.boss_fadein > 0: | |
| self.boss_fadein -= 1 | |
| else: | |
| # jitter within ±10% of area, but keep full M inside bounds | |
| half_w = self.boss_span / 2 | |
| half_h = self.boss_height / 2 | |
| spawn_cx = self.area_x + self.area_size // 2 | |
| spawn_cy = self.area_y + 20 | |
| jx = self.np_random.uniform(-0.03, 0.03) * self.area_size | |
| jy = self.np_random.uniform(-0.01, 0.01) * self.area_size | |
| cx = np.clip(spawn_cx + jx, self.area_x + half_w, self.area_x + self.area_size - half_w) | |
| cy = np.clip(spawn_cy + jy, self.area_y + half_h, self.area_y + self.area_size // 2 - half_h) | |
| self.boss_pos[0], self.boss_pos[1] = cx, cy | |
| # fire | |
| self.boss_fire_timer -= 1 | |
| if self.boss_fire_timer <= 0 and self.boss_bullet is None: | |
| cx, cy = self.boss_pos | |
| # fire from the center dip of the M (lower y) | |
| origin_x = cx | |
| origin_y = cy + (self.boss_height / 2) | |
| ax = self.agent_x + self.agent_w / 2 | |
| ay = self.agent_y + self.agent_h / 2 | |
| dx, dy = ax - origin_x, ay - origin_y | |
| n = max(1e-6, np.hypot(dx, dy)) | |
| vx = self.boss_bullet_speed * dx / n | |
| vy = self.boss_bullet_speed * dy / n | |
| self.boss_bullet = [origin_x, origin_y, vx, vy] | |
| self.boss_fire_timer = 50 | |
| # boss bullet move + bounce + collision | |
| if self.boss_bullet is not None: | |
| self.boss_bullet[0] += self.boss_bullet[2] | |
| self.boss_bullet[1] += self.boss_bullet[3] | |
| # bounce on vertical walls | |
| if (self.boss_bullet[0] - self.boss_bullet_size < self.area_x) or \ | |
| (self.boss_bullet[0] + self.boss_bullet_size > self.area_x + self.area_size): | |
| self.boss_bullet[2] *= -1 | |
| # despawn at bottom border (stop rendering past agent/border) | |
| if self.boss_bullet[1] - self.boss_bullet_size >= self.area_y + self.area_size: | |
| self.boss_bullet = None | |
| # collision | |
| if self.boss_bullet is not None: | |
| bx, by, vx, vy = self.boss_bullet | |
| if abs(bx - (ax+self.agent_w/2)) < (self.agent_w/2 + self.boss_bullet_size) and \ | |
| abs(by - (ay+self.agent_h/2)) < (self.agent_h/2 + self.boss_bullet_size): | |
| reward -= 2.0 | |
| self.stun_timer = 16 | |
| self.reload_timer = 16 | |
| self.boss_bullet = None # bullet consumed | |
| # Bonus treat for when there are no birds left | |
| if len(self.birds) == 0: | |
| reward += 1.0 | |
| self.steps += 1 | |
| if self.steps >= self.max_episode_steps: | |
| truncated = True | |
| return self._get_obs(), reward, terminated, truncated, {} | |
| def _get_obs(self): | |
| max_birds = self.bird_rows * self.bird_cols | |
| bird_xs = [bx / self.width for bx, by, hp in self.birds] | |
| bird_xs += [0.0] * (max_birds - len(bird_xs)) | |
| lowest_y = max([by for bx, by, hp in self.birds], default=self.area_y) / self.height | |
| agent_bullets = [(x/self.width, y/self.height) for x,y in self.agent_bullets[:self.max_agent_bullets]] | |
| agent_bullets += [(0.0, 0.0)] * (self.max_agent_bullets - len(agent_bullets)) | |
| agent_bullets_flat = [v for xy in agent_bullets for v in xy] | |
| bird_bullets = [(x/self.width, y/self.height) for x,y in self.bird_bullets[:self.max_bird_bullets]] | |
| bird_bullets += [(0.0, 0.0)] * (self.max_bird_bullets - len(bird_bullets)) | |
| bird_bullets_flat = [v for xy in bird_bullets for v in xy] | |
| # Boss bullet (x,y,vx,vy) | |
| if self.boss_bullet is not None: | |
| bx, by, vx, vy = self.boss_bullet | |
| boss_xy = (np.clip(bx/self.width, 0.0, 1.0), | |
| np.clip(by/self.height, 0.0, 1.0)) | |
| s = max(1e-6, self.boss_bullet_speed) | |
| boss_v = (np.clip(vx/s, -1.0, 1.0), | |
| np.clip(vy/s, -1.0, 1.0)) | |
| else: | |
| boss_xy = (0.0, 0.0) | |
| boss_v = (0.0, 0.0) | |
| return np.array( | |
| [self.agent_x / self.width, | |
| self.reload_timer / 5.0, | |
| len(self.birds) / max_birds, | |
| *bird_xs, | |
| lowest_y, | |
| *agent_bullets_flat, | |
| *bird_bullets_flat, | |
| self.stun_timer / 10.0, | |
| *boss_xy, | |
| *boss_v], | |
| dtype=np.float32 | |
| ) | |
| def render(self): | |
| if self._surface is None: | |
| pygame.init() | |
| self._surface = pygame.Surface((self.width, self.height)) | |
| self.clock = pygame.time.Clock() | |
| if self.font is None: | |
| self.font = pygame.font.SysFont("Arial", 18, bold=True) | |
| self._surface.fill((0, 0, 0)) | |
| # border | |
| pygame.draw.rect(self._surface, (255,255,255), | |
| pygame.Rect(self.area_x, self.area_y, self.area_size, self.area_size), 2) | |
| # agent body | |
| pygame.draw.rect(self._surface, (255,140,0), | |
| pygame.Rect(int(self.agent_x), int(self.agent_y), | |
| self.agent_w, self.agent_h)) | |
| # small triangle muzzle on top-center | |
| mx = self.agent_x + self.agent_w // 2 | |
| my = self.agent_y | |
| pygame.draw.polygon(self._surface, (255,140,0), | |
| [(mx-4, my), (mx+4, my), (mx, my-6)]) | |
| # birds (draw as "V") | |
| for bx, by, hp in self.birds: | |
| points = [(bx, by), | |
| (bx+self.bird_w//2, by+self.bird_h), | |
| (bx+self.bird_w, by)] | |
| if hp == 2: | |
| pygame.draw.polygon(self._surface, (200,200,200), points) | |
| else: | |
| pygame.draw.polygon(self._surface, (200,200,200), points, 2) | |
| # bullets | |
| for x,y in self.agent_bullets: | |
| pygame.draw.rect(self._surface, (255,140,0), | |
| pygame.Rect(int(x-2), int(y-6), 4, 6)) | |
| for x,y in self.bird_bullets: | |
| pygame.draw.rect(self._surface, (200,200,200), | |
| pygame.Rect(int(x-2), int(y), 4, 6)) | |
| # boss | |
| if self.boss_active and self.boss_pos is not None: | |
| cx, cy = self.boss_pos | |
| half_w = self.boss_span / 2 | |
| half_h = self.boss_height / 2 | |
| # "M" polyline (left tip → left peak → center dip → right peak → right tip) | |
| points = [ | |
| (int(cx - half_w), int(cy + half_h)), # left tip (lower) | |
| (int(cx - half_w * 0.5), int(cy - half_h)), # left peak (upper) | |
| (int(cx), int(cy + half_h)), # center dip (lower) | |
| (int(cx + half_w * 0.5), int(cy - half_h)), # right peak (upper) | |
| (int(cx + half_w), int(cy + half_h)), # right tip (lower) | |
| ] | |
| # Fade alpha: 0→255 over 5 frames | |
| alpha = 255 if self.boss_fadein == 0 else int(255 * (1.0 - self.boss_fadein / 5.0)) | |
| overlay = pygame.Surface((self.width, self.height), pygame.SRCALPHA) | |
| pygame.draw.lines(overlay, (200, 200, 200, alpha), False, points, 3) # outline width=3 | |
| self._surface.blit(overlay, (0, 0)) | |
| # boss bullet | |
| if self.boss_bullet is not None: | |
| bx, by, vx, vy = self.boss_bullet | |
| pygame.draw.circle(self._surface, (200,200,200), | |
| (int(bx), int(by)), self.boss_bullet_size) | |
| # score | |
| max_birds = self.bird_rows * self.bird_cols | |
| text_color = (80, 200, 110) if self.score >= max_birds else (255, 255, 255) | |
| text = self.font.render(str(self.score), True, text_color) | |
| self._surface.blit(text, (self.area_x+4, self.area_y-20)) | |
| # stun | |
| if self.stun_timer > 0: | |
| factor = self.stun_timer / 16.0 # 1.0 → 0.0 | |
| else: | |
| factor = 0.0 | |
| # red flash overlay (game area only) | |
| if self.stun_timer > 0: | |
| alpha = int(140 * factor) # max opacity ~140 (out of 255) | |
| overlay = pygame.Surface((self.area_size, self.area_size)) | |
| overlay.set_alpha(alpha) | |
| overlay.fill((220, 30, 10)) | |
| self._surface.blit(overlay, (self.area_x, self.area_y)) | |
| return np.transpose(np.array(pygame.surfarray.pixels3d(self._surface)), (1,0,2)).copy() | |
| def close(self): | |
| if self._surface is not None: | |
| pygame.quit() | |
| self._surface = None | |