Ram Atari

component/atari_wrapper.py

@@ -108,6 +108,27 @@ class MaxAndSkipEnv(gym.Wrapper):
         self._obs_buffer.append(obs)
         return obs
 
+class SkipEnv(gym.Wrapper):
+    def __init__(self, env=None, skip=4):
+        """Return only every `skip`-th frame"""
+        super(SkipEnv, self).__init__(env)
+        self._skip       = skip
+
+    def step(self, action):
+        total_reward = 0.0
+        done = None
+        for _ in range(self._skip):
+            obs, reward, done, info = self.env.step(action)
+            total_reward += reward
+            if done:
+                break
+
+        return obs, total_reward, done, info
+
+    def reset(self):
+        obs = self.env.reset()
+        return obs
+
 class DatasetEnv(gym.Wrapper):
     def __init__(self, env=None):
         super(DatasetEnv, self).__init__(env)
@@ -174,10 +195,10 @@ class StackFrame(gym.Wrapper):
     def reset(self):
         state = self.env.reset()
         self.buffer = [state] * self.history_length
-        return np.vstack(self.buffer)
+        return np.asarray(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
+        return np.asarray(np.vstack(self.buffer)), reward, done, info

component/task.py

@@ -73,6 +73,23 @@ class PixelAtari(BasicTask):
     def normalize_state(self, state):
         return np.asarray(state) / 255.0
 
+class RamAtari(BasicTask):
+    def __init__(self, name, no_op, frame_skip, max_steps=10000):
+        BasicTask.__init__(self, max_steps)
+        self.name = name
+        env = gym.make(name)
+        assert 'NoFrameskip' in env.spec.id
+        env = EpisodicLifeEnv(env)
+        env = NoopResetEnv(env, noop_max=no_op)
+        env = SkipEnv(env, skip=frame_skip)
+        if 'FIRE' in env.unwrapped.get_action_meanings():
+            env = FireResetEnv(env)
+        self.env = env
+        self.action_dim = self.env.action_space.n
+
+    def normalize_state(self, state):
+        return np.asarray(state) / 255.0
+
 class ContinuousMountainCar(BasicTask):
     name = 'MountainCarContinuous-v0'
     success_threshold = 90