Optimize replay buffer

.gitignore

@@ -10,6 +10,7 @@ upload.py
 data
 draw_*
 log
+figure
 
 # C extensions
 *.so

dqn_agent.py

@@ -11,6 +11,7 @@ import numpy as np
 import time
 import psutil
 import os
+import pickle
 
 class DQNAgent:
     def __init__(self,
@@ -39,6 +40,7 @@ class DQNAgent:
         self.history_length = history_length
         self.logger = logger
         self.process = psutil.Process(os.getpid())
+        self.report_interval = 1000
 
     def get_state(self, history_buffer):
         if self.history_length > 1:
@@ -53,50 +55,53 @@ class DQNAgent:
         steps = 0
         while not self.step_limit or steps < self.step_limit:
             state = self.get_state(history_buffer)
+            state = self.task.normalize_state(state)
             value = self.learning_network.predict(np.reshape(state, (1, ) + state.shape))
             action = self.policy.sample(value.flatten())
             next_state, reward, done, info = self.task.step(action)
+            self.replay.feed([history_buffer[-1], action, reward, next_state, int(done)])
             history_buffer.pop(0)
             history_buffer.append(next_state)
-            next_state = self.get_state(history_buffer)
             total_reward += reward
-            self.replay.feed([state, action, reward, next_state, int(done)])
             steps += 1
             self.total_steps += 1
             if done:
                 break
             if self.total_steps > self.explore_steps:
                 sample_start_time = time.time()
-                experiences = self.replay.sample()
-                self.logger.debug('sample time %f' % (time.time() - sample_start_time))
+                experiences = self.replay.sample(self.history_length)
+                if self.total_steps % self.report_interval == 0:
+                    self.logger.debug('sample time %f' % (time.time() - sample_start_time))
                 states, actions, rewards, next_states, terminals = experiences
+                states = self.task.normalize_state(states)
+                next_states = self.task.normalize_state(next_states)
                 predict_start_time = time.time()
                 targets = self.learning_network.predict(states)
                 q_next = self.target_network.predict(next_states)
-                self.logger.debug('prediction time %f' % (time.time() - predict_start_time))
+                if self.total_steps % self.report_interval == 0:
+                    self.logger.debug('prediction time %f' % (time.time() - predict_start_time))
                 q_next = np.max(q_next, axis=1)
                 q_next = np.where(terminals, 0, q_next)
                 q_next = rewards + self.discount * q_next
                 targets[np.arange(len(actions)), actions] = q_next
                 minibatch_start_time = time.time()
                 self.learning_network.learn(states, targets)
-                self.logger.debug('minibatch time %f' % (time.time() - minibatch_start_time))
+                if self.total_steps % self.report_interval == 0:
+                    self.logger.debug('minibatch time %f' % (time.time() - minibatch_start_time))
             if self.total_steps % self.target_network_update_freq == 0:
                 self.target_network.load_state_dict(self.learning_network.state_dict())
             if self.total_steps > self.explore_steps:
                 self.policy.update_epsilon()
         episode_time = time.time() - episode_start_time
-        info = self.process.memory_full_info()
-        if hasattr(info, 'swap'):
-            info_stat = info.swap
-        elif hasattr(info, 'pfaults'):
-            info_stat = info.pfaults
-        else:
-            info_stat = -1
-        self.logger.debug('episode steps %d, episode time %f, time per step %f, memory_info %d' %
-                          (steps, episode_time, episode_time / float(steps), info_stat))
+        info = self.process.memory_info()
+        self.logger.debug('episode steps %d, episode time %f, time per step %f, rss %d, vms %d' %
+                          (steps, episode_time, episode_time / float(steps), info.rss, info.vms))
         return total_reward
 
+    def save(self, file_name):
+        with open(file_name, 'wb') as f:
+            pickle.dump(self.learning_network.state_dict(), f)
+
     def run(self):
         window_size = 100
         ep = 0
@@ -104,6 +109,8 @@ class DQNAgent:
         while True:
             ep += 1
             reward = self.episode()
+            if ep % 1000 == 0:
+                self.save('data/dqn-episode-%d.bin')
             rewards.append(reward)
             avg_reward = np.mean(rewards[-window_size:])
             self.logger.info('episode %d, epsilon %f, reward %f, avg reward %f, total steps %d' % (

main.py

@@ -66,24 +66,25 @@ def actor_critic_cart_pole():
     config['step_limit'] = 300
     config['n_workers'] = 8
     config['batch_size'] = 5
-    config['test_interval'] = 500
+    config['test_interval'] = 50000
     config['test_repeats'] = 5
     agent = AsyncAgent(**config)
     agent.run()
 
 def dqn_pixel_atari(name):
     config = dict()
-    config['task_fn'] = lambda: PixelAtari(name)
-    config['optimizer_fn'] = lambda params: torch.optim.RMSprop(params, lr=0.00025)
-    config['network_fn'] = lambda optimizer_fn: ConvNet(4, 6, optimizer_fn)
+    history_length = 4
+    config['task_fn'] = lambda: PixelAtari(name, 30)
+    config['optimizer_fn'] = lambda params: torch.optim.RMSprop(params, lr=0.00025, alpha=0.95, eps=0.01)
+    config['network_fn'] = lambda optimizer_fn: ConvNet(history_length, 6, optimizer_fn)
     config['policy_fn'] = lambda: GreedyPolicy(epsilon=1.0, final_step=1000000, min_epsilon=0.1)
-    config['replay_fn'] = lambda: Replay(memory_size=200000, batch_size=32)
+    config['replay_fn'] = lambda: Replay(memory_size=1000000, batch_size=32, dtype=np.uint8)
     config['discount'] = 0.99
     config['target_network_update_freq'] = 10000
     config['step_limit'] = 0
     config['explore_steps'] = 50000
     config['logger'] = gym.logger
-    config['history_length'] = 4
+    config['history_length'] = history_length
     agent = DQNAgent(**config)
     agent.run()
 

replay.py

@@ -7,9 +7,10 @@
 import numpy as np
 
 class Replay:
-    def __init__(self, memory_size, batch_size):
+    def __init__(self, memory_size, batch_size, dtype=np.float32):
         self.memory_size = memory_size
         self.batch_size = batch_size
+        self.dtype = dtype
 
         self.states = None
         self.actions = np.empty(self.memory_size, dtype=np.int8)
@@ -25,8 +26,8 @@ class Replay:
         state, action, reward, next_state, done = experience
 
         if self.states is None:
-            self.states = np.empty((self.memory_size, ) + state.shape)
-            self.next_states = np.empty((self.memory_size, ) + state.shape)
+            self.states = np.empty((self.memory_size, ) + state.shape, dtype=self.dtype)
+            self.next_states = np.empty((self.memory_size, ) + state.shape, dtype=self.dtype)
 
         self.states[self.pos][:] = state
         self.actions[self.pos] = action
@@ -39,16 +40,37 @@ class Replay:
             self.full = True
             self.pos = 0
 
-    def sample(self):
+    def sample(self, history_length):
         upper_bound = self.memory_size if self.full else self.pos
         sampled_indices = np.random.randint(0, upper_bound, size=self.batch_size)
-        sampled_states = self.states[sampled_indices]
-        sampled_actions = self.actions[sampled_indices]
-        sampled_rewards = self.rewards[sampled_indices]
-        sampled_next_states = self.next_states[sampled_indices]
-        sampled_terminals = self.terminals[sampled_indices]
-        return [sampled_states,
-                sampled_actions,
-                sampled_rewards,
-                sampled_next_states,
-                sampled_terminals]
+        sampled_states = []
+        sampled_actions = []
+        sampled_rewards = []
+        sampled_next_states = []
+        sampled_terminals = []
+        for index in sampled_indices:
+            if history_length == 1:
+                sampled_states.append(self.states[index])
+                sampled_next_states.append(self.next_states[index])
+            else:
+                full_indices = [(index - i + self.memory_size) % self.memory_size for i in range(history_length)]
+                if self.pos in full_indices:
+                    for i in range(full_indices.index(self.pos), len(full_indices)):
+                        full_indices[i] = self.pos
+                state = [self.states[i] for i in full_indices]
+                state = np.vstack(state)
+                sampled_states.append(state)
+
+                next_state = [self.next_states[i] for i in full_indices]
+                next_state = np.vstack(next_state)
+                sampled_next_states.append(next_state)
+
+            sampled_rewards.append(self.rewards[index])
+            sampled_actions.append(self.actions[index])
+            sampled_terminals.append(self.terminals[index])
+
+        return [np.asarray(sampled_states),
+                np.asarray(sampled_actions),
+                np.asarray(sampled_rewards),
+                np.asarray(sampled_next_states),
+                np.asarray(sampled_terminals)]

task.py

@@ -9,11 +9,20 @@ import numpy as np
 import cv2
 
 class BasicTask:
+    no_op = 0
+
     def transfer_state(self, state):
         return state
 
+    def normalize_state(self, state):
+        return state
+
     def reset(self):
-        return self.transfer_state(self.env.reset())
+        state = self.env.reset()
+        if self.no_op > 0:
+            for _ in range(np.random.randint(1, self.no_op + 1)):
+                state, _, _, _ = self.env.step(0)
+        return self.transfer_state(state)
 
     def step(self, action):
         next_state, reward, done, info = self.env.step(action)
@@ -47,10 +56,14 @@ class PixelAtari(BasicTask):
     height = 84
     success_threshold = 1000
 
-    def __init__(self, name):
+    def __init__(self, name, no_op):
+        self.no_op = no_op
         self.env = gym.make(name)
 
     def transfer_state(self, state):
-        img = (state[:, :, 0] * 0.299 + state[:, :, 1] * 0.587 + state[:, :, 2] * 0.114) / 255.0
+        img = (state[:, :, 0] * 0.299 + state[:, :, 1] * 0.587 + state[:, :, 2] * 0.114)
         img = cv2.resize(img, (self.width, self.height))
-        return np.reshape(img, (1, self.width, self.height))
+        return np.asarray(np.reshape(img, (1, self.width, self.height)), np.uint8)
+
+    def normalize_state(self, state):
+        return np.asarray(state, dtype=np.float32) / 255.0