Added kwarg passthrough in FrameStack wrapper (#757)

* Update Framestack reset function to pass kwargs * Update changelog * Update changelog.rst Co-authored-by: Antonin RAFFIN <antonin.raffin@ensta.org>
Stable-Baselines-Team · Mar 22, 2020 · d2364c9 · d2364c9
1 parent ae4f6c5
commit d2364c9
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst
@@ -20,6 +20,7 @@ Bug Fixes:
 - Fixed DDPG sampling empty replay buffer when combined with HER  (@tirafesi)
 - Fixed a bug in ``HindsightExperienceReplayWrapper``, where the openai-gym signature for ``compute_reward`` was not matched correctly (@johannes-dornheim)
 - Fixed SAC/TD3 checking time to update on learn steps instead of total steps (@solliet)
+- Added ``**kwarg`` pass through for ``reset`` method in ``atari_wrappers.FrameStack`` (@solliet)
 
 Deprecations:
 ^^^^^^^^^^^^^

diff --git a/stable_baselines/common/atari_wrappers.py b/stable_baselines/common/atari_wrappers.py
@@ -214,8 +214,8 @@ def __init__(self, env, n_frames):
         self.observation_space = spaces.Box(low=0, high=255, shape=(shp[0], shp[1], shp[2] * n_frames),
                                             dtype=env.observation_space.dtype)
 
-    def reset(self):
-        obs = self.env.reset()
+    def reset(self, **kwargs):
+        obs = self.env.reset(**kwargs)
         for _ in range(self.n_frames):
             self.frames.append(obs)
         return self._get_ob()