diff --git a/.gitignore b/.gitignore index 220aa96..0347215 100644 --- a/.gitignore +++ b/.gitignore @@ -10,6 +10,7 @@ upload.py data draw_* log +figure # C extensions *.so diff --git a/dqn_agent.py b/dqn_agent.py index 06ed10b..42161a7 100644 --- a/dqn_agent.py +++ b/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' % ( diff --git a/main.py b/main.py index 214c8eb..8e28acd 100644 --- a/main.py +++ b/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() diff --git a/replay.py b/replay.py index 6dfafb2..d6a578a 100644 --- a/replay.py +++ b/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)] diff --git a/task.py b/task.py index 4c2cddb..0d460da 100644 --- a/task.py +++ b/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