diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index 7cb344b..c686539 100644 --- a/docs/misc/changelog.rst +++ b/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 diff --git a/stable_baselines3/common/atari_wrappers.py b/stable_baselines3/common/atari_wrappers.py index 785d911..32c1bda 100644 --- a/stable_baselines3/common/atari_wrappers.py +++ b/stable_baselines3/common/atari_wrappers.py @@ -106,7 +106,13 @@ class EpisodicLifeEnv(gym.Wrapper): 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 @@ class MaxAndSkipEnv(gym.Wrapper): return max_frame, total_reward, done, info - def reset(self, **kwargs) -> GymObs: - return self.env.reset(**kwargs) - class ClipRewardEnv(gym.RewardWrapper): """