diff --git a/pokemonred_puffer/environment.py b/pokemonred_puffer/environment.py index 555c25f..0ca28cb 100644 --- a/pokemonred_puffer/environment.py +++ b/pokemonred_puffer/environment.py @@ -682,12 +682,8 @@ def step(self, action): self.step_count += 1 reset = ( - self.step_count - >= min( - self.max_steps, - self.max_steps - * (len(self.required_events) + len(self.required_items) * self.max_steps_scaling), - ) # or + self.step_count >= self.get_max_steps() + # or # self.caught_pokemon[6] == 1 # squirtle ) @@ -1294,9 +1290,7 @@ def agent_stats(self, action): "pokecenter_heal": self.pokecenter_heal, "in_battle": self.read_m("wIsInBattle") > 0, "event": self.progress_reward["event"], - "max_steps": self.max_steps - * (len(self.required_events) + len(self.required_items)) - * self.max_steps_scaling, + "max_steps": self.get_max_steps(), } | { "exploration": { @@ -1359,6 +1353,13 @@ def add_video_frame(self): def get_game_coords(self): return (self.read_m(0xD362), self.read_m(0xD361), self.read_m(0xD35E)) + def get_max_steps(self): + return min( + self.max_steps, + self.max_steps + * (len(self.required_events) + len(self.required_items) * self.max_steps_scaling), + ) + def update_seen_coords(self): inc = 0.5 if (self.read_m("wd736") & 0b1000_0000) else self.exploration_inc