Fix Atari wrapper bug: tried to step environment that needs reset (#1297

)

* fix 1060

* update changelog

docs/misc/changelog.rst

@@ -23,6 +23,7 @@ New Features:
 
 Bug Fixes:
 ^^^^^^^^^^
+- Fixed Atari wrapper that missed the reset condition (@luizapozzobon)
 
 Deprecations:
 ^^^^^^^^^^^^^
@@ -1218,4 +1219,4 @@ And all the contributors:
 @Gregwar @ycheng517 @quantitative-technologies @bcollazo @git-thor @TibiGG @cool-RR @MWeltevrede
 @Melanol @qgallouedec @francescoluciano @jlp-ue @burakdmb @timothe-chaumont @honglu2875 @yuanmingqi
 @anand-bala @hughperkins @sidney-tio @AlexPasqua @dominicgkerr @Akhilez @Rocamonde @tobirohrer @ZikangXiong
-@DavyMorgan
+@DavyMorgan @luizapozzobon

stable_baselines3/common/atari_wrappers.py

@@ -106,7 +106,13 @@ def reset(self, **kwargs) -> np.ndarray:
             obs = self.env.reset(**kwargs)
         else:
             # no-op step to advance from terminal/lost life state
-            obs, _, _, _ = self.env.step(0)
+            obs, _, done, _ = self.env.step(0)
+
+            # The no-op step can lead to a game over, so we need to check it again
+            # to see if we should reset the environment and avoid the
+            # monitor.py `RuntimeError: Tried to step environment that needs reset`
+            if done:
+                obs = self.env.reset(**kwargs)
         self.lives = self.env.unwrapped.ale.lives()
         return obs
 
@@ -150,9 +156,6 @@ def step(self, action: int) -> GymStepReturn:
 
         return max_frame, total_reward, done, info
 
-    def reset(self, **kwargs) -> GymObs:
-        return self.env.reset(**kwargs)
-
 
 class ClipRewardEnv(gym.RewardWrapper):
     """