From e5e2a0623c34807d62da2160050f246ebbda54f2 Mon Sep 17 00:00:00 2001
From: aboundy <andrew@boundy.com>
Date: Sat, 28 Oct 2023 21:45:30 +0200
Subject: [PATCH] first go

---
 .gitignore                                    |  1 +
 pufferlib/environments/pokemon_red_updated.py | 12 +++++++-----
 tests/test_pokemon_red_step.py                | 19 +++++++++++++++++++
 3 files changed, 27 insertions(+), 5 deletions(-)
 create mode 100644 tests/test_pokemon_red_step.py

diff --git a/.gitignore b/.gitignore
index 07d65164..730fb29f 100644
--- a/.gitignore
+++ b/.gitignore
@@ -5,6 +5,7 @@ __pycache__/
 
 # C extensions
 *.so
+*.c
 
 # Distribution / packaging
 .Python
diff --git a/pufferlib/environments/pokemon_red_updated.py b/pufferlib/environments/pokemon_red_updated.py
index 3cd3b50b..782160ba 100644
--- a/pufferlib/environments/pokemon_red_updated.py
+++ b/pufferlib/environments/pokemon_red_updated.py
@@ -1,8 +1,7 @@
 
 import sys
-import uuid 
-import os
-from math import floor, sqrt
+import uuid
+from math import floor
 import json
 from pathlib import Path
 
@@ -38,7 +37,7 @@ def __init__(
             extra_buttons=False,
             explore_weight=3 # 2.5
             ):
-        self.s_path = Path(f'session_{str(uuid.uuid4())[:8]}')
+        self.s_path = Path(f'runs/session_{str(uuid.uuid4())[:8]}')
         self.gb_path=str(Path(__file__).parent / 'pokemon_red.gb')
         self.init_state=str(Path(__file__).parent / 'has_pokedex_nballs.state')
         
@@ -63,7 +62,7 @@ def __init__(
         self.extra_buttons = extra_buttons
         self.instance_id = str(uuid.uuid4())[:8]
 
-        self.s_path.mkdir(exist_ok=True)
+        self.s_path.mkdir(parents=True, exist_ok=True)
         self.reset_count = 0
         self.all_runs = []
 
@@ -133,7 +132,10 @@ def __init__(
         self.reset()
 
     def reset(self, seed=None):
+        # Gym v26 favours `Env.reset(seed=seed)` this allows seeding to only be changed on environment reset.
+        # https://gymnasium.farama.org/content/migration-guide/#seed-and-random-number-generator
         self.seed = seed
+
         # restart game, skipping credits
         self.initial_state.seek(0)
         self.pyboy.load_state(self.initial_state)
diff --git a/tests/test_pokemon_red_step.py b/tests/test_pokemon_red_step.py
new file mode 100644
index 00000000..0ef07642
--- /dev/null
+++ b/tests/test_pokemon_red_step.py
@@ -0,0 +1,19 @@
+from pdb import set_trace as T
+import numpy as np
+from pufferlib.environments import PokemonRed
+
+env = PokemonRed()
+r_ob, r_info = env.reset()
+isinstance(r_ob, np.ndarray)
+isinstance(r_info, dict)
+
+for i in range(100):
+    action = env.action_space.sample()
+    ob, reward, terminal, truncated, info = env.step(action)
+    print(f'Step: {i}, Info: {info}')
+
+    # check datatype output
+    isinstance(ob, np.ndarray)
+    isinstance(reward, float)
+    isinstance(terminal, bool)
+    isinstance(truncated, bool)
\ No newline at end of file