Use correct function.

sobomax · Jan 16, 2024 · 3380986 · 3380986
1 parent 2323949
commit 3380986
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/python/DiggerGym.py b/python/DiggerGym.py
@@ -26,13 +26,13 @@ def step(self, action):
         reward = (newscore:=self.digger.getscore(0)) - self.lastscore
         assert reward >= 0
         self.lastscore = newscore
-        return self.screenshot(), reward, done, False, {}
+        return self.digger.getscreenrgb(), reward, done, False, {}
 
     def reset(self):
         self.digger.initgame()
         self.digger.startlevel()
         self.lastscore = self.digger.getscore(0)
-        return self.screenshot(), {}
+        return self.digger.getscreenrgb(), {}
 
     def screenshot(self):
         s = self.digger.screenshot()