From 8168ba4002fe932d8f1c8a64d10067a991314c67 Mon Sep 17 00:00:00 2001
From: thatguy11325 <148832074+thatguy11325@users.noreply.github.com>
Date: Sat, 8 Jun 2024 23:13:07 -0400
Subject: [PATCH] Better debug mode, add no wild encounters mode

---
 config.yaml                      |  5 +++++
 pokemonred_puffer/environment.py | 18 ++++++++++++++++--
 2 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/config.yaml b/config.yaml
index dfeee75..5f87f58 100644
--- a/config.yaml
+++ b/config.yaml
@@ -9,6 +9,8 @@ debug:
     stream_wrapper: False
     init_state: cut
     max_steps: 1_000_000
+    disable_wild_encounters: True
+    disable_ai_actions: True
   train:
     device: cpu
     compile: False
@@ -48,6 +50,9 @@ env:
   two_bit: True
   log_frequency: 2000
   auto_flash: True
+  disable_wild_encounters: False
+  disable_ai_actions: False
+
 
 train:
   seed: 1
diff --git a/pokemonred_puffer/environment.py b/pokemonred_puffer/environment.py
index 3f4fa5e..d3bc01a 100644
--- a/pokemonred_puffer/environment.py
+++ b/pokemonred_puffer/environment.py
@@ -192,6 +192,8 @@ def __init__(self, env_config: pufferlib.namespace):
         self.log_frequency = env_config.log_frequency
         self.two_bit = env_config.two_bit
         self.auto_flash = env_config.auto_flash
+        self.disable_wild_encounters = env_config.disable_wild_encounters
+        self.disable_ai_actions = env_config.disable_ai_actions
         self.action_space = ACTION_SPACE
 
         # Obs space-related. TODO: avoid hardcoding?
@@ -303,6 +305,13 @@ def register_hooks(self):
         self.pyboy.hook_register(None, "SetLastBlackoutMap.done", self.blackout_update_hook, None)
         # self.pyboy.hook_register(None, "UsedCut.nothingToCut", self.cut_hook, context=True)
         # self.pyboy.hook_register(None, "UsedCut.canCut", self.cut_hook, context=False)
+        if self.disable_wild_encounters:
+            self.pyboy.hook_register(
+                None,
+                "TryDoWildEncounter.gotWildEncounterType",
+                self.disable_wild_encounter_hook,
+                None,
+            )
 
     def update_state(self, state: bytes):
         self.reset(seed=random.randint(0, 10), options={"state": state})
@@ -630,8 +639,9 @@ def run_action_on_emulator(self, action):
         self.action_hist[action] += 1
         # press button then release after some steps
         # TODO: Add video saving logic
-        self.pyboy.send_input(VALID_ACTIONS[action])
-        self.pyboy.send_input(VALID_RELEASE_ACTIONS[action], delay=8)
+        if not self.disable_ai_actions:
+            self.pyboy.send_input(VALID_ACTIONS[action])
+            self.pyboy.send_input(VALID_RELEASE_ACTIONS[action], delay=8)
         self.pyboy.tick(self.action_freq, render=True)
 
         if self.read_bit(0xD803, 0):
@@ -811,6 +821,10 @@ def cut_hook(self, context):
         self.cut_explore_map[local_to_global(y, x, map_id)] = 1
         self.cut_tiles[wTileInFrontOfPlayer] = 1
 
+    def disable_wild_encounter_hook(self, *args, **kwargs):
+        self.pyboy.memory[self.pyboy.symbol_lookup("wRepelRemainingSteps")[1]] = 100
+        self.pyboy.memory[self.pyboy.symbol_lookup("wCurEnemyLVL")[1]] = 1
+
     def agent_stats(self, action):
         levels = [self.read_m(f"wPartyMon{i+1}Level") for i in range(self.read_m("wPartyCount"))]
         return {