update mario

RLE-Foundation · May 25, 2024 · 0940ab4 · 0940ab4
1 parent 3615d8b
commit 0940ab4
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/rllte/env/mario/wrappers.py b/rllte/env/mario/wrappers.py
@@ -22,12 +22,12 @@ def step(self, action):
 
     def reset(self, **kwargs):
         if self.was_real_done:
-            obs = self.env.reset(**kwargs)
+            obs, infos = self.env.reset(**kwargs)
         else:
             # no-op step to advance from terminal/lost life state
-            obs, _, _, _, _ = self.env.step(0)
+            obs, _, _, _, infos = self.env.step(0)
         self.lives = self.env.unwrapped.env._life
-        return obs
+        return obs, infos
 
 class SkipFrame(gym.Wrapper):
     def __init__(self, env, skip):