From a907dac0bb831f8f2b068b4820a89e355e8c4e0c Mon Sep 17 00:00:00 2001
From: Shangtong Zhang <zhangshangtong.cpp@icloud.com>
Date: Mon, 19 Feb 2018 13:55:57 -0700
Subject: [PATCH] Ram Atari

---
 component/atari_wrapper.py | 25 +++++++++++++++++++++++--
 component/task.py          | 17 +++++++++++++++++
 main.py                    | 27 +++++++++++++++++++++++++--
 3 files changed, 65 insertions(+), 4 deletions(-)

diff --git a/component/atari_wrapper.py b/component/atari_wrapper.py
index 63735c9..667964d 100644
--- a/component/atari_wrapper.py
+++ b/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
diff --git a/component/task.py b/component/task.py
index c292441..187231f 100644
--- a/component/task.py
+++ b/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
diff --git a/main.py b/main.py
index d156d3e..293a4dd 100644
--- a/main.py
+++ b/main.py
@@ -114,6 +114,27 @@ def dqn_pixel_atari(name):
     config.double_q = False
     run_episodes(DQNAgent(config))
 
+def dqn_ram_atari(name):
+    config = Config()
+    config.history_length = 1
+    config.task_fn = lambda: RamAtari(name, no_op=30, frame_skip=4)
+    action_dim = config.task_fn().action_dim
+    config.optimizer_fn = lambda params: torch.optim.RMSprop(params, lr=0.00025, alpha=0.95, eps=0.01)
+    config.network_fn = lambda: FCNet([128, 64, 64, action_dim], gpu=2)
+    config.policy_fn = lambda: GreedyPolicy(epsilon=0.1, final_step=1000000, min_epsilon=0.1)
+    config.replay_fn = lambda: Replay(memory_size=100000, batch_size=32, dtype=np.uint8)
+    config.reward_shift_fn = lambda r: np.sign(r)
+    config.discount = 0.99
+    config.target_network_update_freq = 10000
+    config.max_episode_length = 0
+    config.exploration_steps= 100
+    config.logger = Logger('./log', logger)
+    config.test_interval = 0
+    config.test_repetitions = 10
+    config.double_q = True
+    # config.double_q = False
+    run_episodes(DQNAgent(config))
+
 def async_pixel_atari(name):
     config = Config()
     config.history_length = 1
@@ -322,7 +343,7 @@ if __name__ == '__main__':
     # logger.setLevel(logging.DEBUG)
     logger.setLevel(logging.INFO)
 
-    # dqn_cart_pole()
+    dqn_cart_pole()
     # async_cart_pole()
     # a3c_cart_pole()
     # a2c_cart_pole()
@@ -334,11 +355,13 @@ if __name__ == '__main__':
     # dqn_pixel_atari('PongNoFrameskip-v4')
     # async_pixel_atari('PongNoFrameskip-v4')
     # a3c_pixel_atari('PongNoFrameskip-v4')
-    a2c_pixel_atari('PongNoFrameskip-v4')
+    # a2c_pixel_atari('PongNoFrameskip-v4')
 
     # dqn_pixel_atari('BreakoutNoFrameskip-v4')
     # async_pixel_atari('BreakoutNoFrameskip-v4')
     # a3c_pixel_atari('BreakoutNoFrameskip-v4')
 
+    dqn_ram_atari('Pong-ramNoFrameskip-v4')
+
     # acvp.train('PongNoFrameskip-v4')