Stack frame wrapper

component/atari_wrapper.py

@@ -1,4 +1,4 @@
-# This file is copied/apdated from
+# This file is apdated from
 # https://raw.githubusercontent.com/transedward/pytorch-dqn/master/utils/atari_wrapper.py
 
 import numpy as np
@@ -189,3 +189,20 @@ class NormalizeFrame(gym.Wrapper):
 
     def _reset(self):
         return self._normalize(self.env.reset())
+
+class StackFrame(gym.Wrapper):
+    def __init__(self, env=None, history_length=1):
+        super(StackFrame, self).__init__(env)
+        self.history_length = history_length
+        self.buffer = None
+
+    def _reset(self):
+        state = self.env.reset()
+        self.buffer = [state] * self.history_length
+        return np.vstack(self.buffer)
+
+    def _step(self, action):
+        state, reward, done, info = self.env.step(action)
+        self.buffer.pop(0)
+        self.buffer.append(state)
+        return np.vstack(self.buffer), reward, done, info

component/task.py

@@ -58,7 +58,7 @@ class LunarLander(BasicTask):
 
 class PixelAtari(BasicTask):
     def __init__(self, name, no_op, frame_skip, normalized_state=True,
-                 frame_size=84, max_steps=10000):
+                 frame_size=84, max_steps=10000, history_length=1):
         BasicTask.__init__(self, max_steps)
         self.normalized_state = normalized_state
         self.name = name
@@ -69,7 +69,8 @@ class PixelAtari(BasicTask):
         env = MaxAndSkipEnv(env, skip=frame_skip)
         if 'FIRE' in env.unwrapped.get_action_meanings():
             env = FireResetEnv(env)
-        self.env = ProcessFrame(env, frame_size)
+        env = ProcessFrame(env, frame_size)
+        self.env = StackFrame(env, history_length)
         self.action_dim = self.env.action_space.n
 
     def normalize_state(self, state):