Space_Invaders.py (4110B)
1 import math 2 from MAMEToolkit.emulator import Emulator 3 from MAMEToolkit.emulator import Address 4 from Actions import Actions 5 from Steps import * 6 7 def add_rewards(old_data, new_data): 8 for k in old_data.keys(): 9 if "rewards" in k: 10 for player in old_data[k]: 11 new_data[k][player] += old_data[k][player] 12 return new_data 13 14 def memory_addresses(): 15 return { 16 "score": Address('0x20F8', 's16'), 17 "player_x": Address('0x201B', 's16'), 18 "player_y": Address('0x201A', 's16'), 19 "shot_status": Address('0x2025', 's16'), 20 "num_aliens": Address('0x2082', 's16'), 21 "alien_shot_x": Address('0x207C', 's16'), 22 "alien_shot_y": Address('0x207B', 's16'), 23 "shot_collision": Address('0x2061', 's16'), 24 "player_alive": Address('0x2068', 's16'), 25 "has_lives": Address('0x20E7', 's16'), 26 "ships_remain": Address('0x21FF', 's16') 27 } 28 29 def index_to_action(action): 30 return { 31 0: [Actions.p1_shoot], 32 1: [Actions.p1_left], 33 2: [Actions.p1_right], 34 3: [] 35 }[action] 36 37 class Space_Invaders(object): 38 def __init__(self, env_id, roms_path, frame_ratio=3, frames_per_step=3): 39 self.frame_ratio = frame_ratio 40 self.frames_per_step = frames_per_step 41 self.emu = Emulator(env_id, roms_path, "invaders", memory_addresses(), frame_ratio=frame_ratio) 42 self.expected_score = {"score": 1000, "alive": 1} 43 self.game_over = False 44 45 def run_steps(self, steps): 46 for step in steps: 47 for i in range(step["wait"]): 48 self.emu.step([]) 49 self.emu.step([action.value for action in step["actions"]]) 50 51 def start(self): 52 self.run_steps(start_game(self.frame_ratio)) 53 self.started = True 54 55 def check_game_over(self, data): 56 if data["has_lives"] < 257: 57 self.game_over = True 58 59 return data 60 61 def new_game(self): 62 self.run_steps(new_game(self.frame_ratio)) 63 self.expected_score = {"scoreL": 1000, "alive": 1} 64 self.game_over = False 65 66 def gather_frames(self, actions): 67 data = self.sub_step(actions) 68 frames = [data["frame"]] 69 for i in range(self.frames_per_step - 1): 70 data = add_rewards(data, self.sub_step(actions)) 71 frames.append(data["frame"]) 72 data["frame"] = frames[0] if self.frames_per_step == 1 else frames 73 return data 74 75 def sub_step(self, actions): 76 data = self.emu.step([action.value for action in actions]) 77 78 score = data["score"] 79 aliens = data["num_aliens"] 80 alive = data["player_alive"] 81 x_diff = (data["player_x"] - data["alien_shot_x"]) 82 y_diff = (data["player_y"] - data["alien_shot_y"]) 83 alien_shot_x = data["alien_shot_x"] 84 alien_shot_y = data["alien_shot_y"] 85 shot_status = (data["shot_status"]) 86 has_lives = (data["has_lives"]) 87 distance = math.sqrt( ((data["player_x"] - data["alien_shot_x"])**2) + ((data["player_y"] - data["alien_shot_y"])**2)) 88 collide = data["shot_collision"] 89 90 rewards = { 91 "score": score, 92 "aliens": aliens, 93 "x_diff": x_diff, 94 "y_diff": y_diff, 95 "alien_shot_x": alien_shot_x, 96 "alien_shot_y": alien_shot_y, 97 "distance": distance, 98 "alive": alive, 99 "shot_status": shot_status, 100 "collide": collide, 101 "has_lives": has_lives 102 } 103 104 data["rewards"] = rewards 105 return data 106 107 def step(self, action): 108 if self.started: 109 if not self.game_over: 110 actions = [] 111 actions += index_to_action(action) 112 data = self.gather_frames(actions) 113 data = self.check_game_over(data) 114 return data["frame"], data["rewards"], self.game_over 115 else: 116 raise EnvironmentError("Attempted to step while game not playing") 117 else: 118 raise EnvironmentError("Start must be called before stepping") 119 120 def close(self): 121 self.emu.close()