From ad8c8ad8538afbcee281e5f2309cd9727a939bc0 Mon Sep 17 00:00:00 2001 From: Shangtong Zhang Date: Sat, 20 May 2017 22:09:18 -0600 Subject: [PATCH] Pixel atari games with DQN --- README.md | 2 +- async_agent.py | 2 +- dqn_agent.py | 49 ++++++++++++++++++++++++++++++++++++------------ main.py | 47 +++++++++++++++++++++++++++------------------- network.py | 51 +++++++++++++++++++++++++++++++++++++++++++++----- policy.py | 10 +++++----- replay.py | 21 ++++++++++++++++----- task.py | 19 +++++++++++++++++-- 8 files changed, 151 insertions(+), 50 deletions(-) diff --git a/README.md b/README.md index fca92f0..d3d01d5 100644 --- a/README.md +++ b/README.md @@ -6,6 +6,6 @@ * Asynchronous N-Step Q-Learning * Asynchronous Advantage Actor Critic (A3C) ->Benchmarked by classical control tasks (CartPole, LunarLander). Atari games will make it difficult to replicate in a regular laptop without a good GPU. However it's fairly easy to adapt the components to fit Atari games. +>Tested with both classical control tasks (CartPole, LunarLander) and Atari games. >Try it out from ```main.py```! \ No newline at end of file diff --git a/async_agent.py b/async_agent.py index 4e7ed94..592e929 100644 --- a/async_agent.py +++ b/async_agent.py @@ -79,7 +79,6 @@ class AsyncAgent: episode_steps = 0 episode_return = 0 episode += 1 - policy.update_epsilon() terminal = False state = task.reset() state = state.reshape([1, -1]) @@ -97,6 +96,7 @@ class AsyncAgent: state = state.reshape([1, -1]) value = worker_network.predict(state) action = policy.sample(value.flatten()) + policy.update_epsilon() batch_rewards = self.bootstrap_fn(batch_states, batch_actions, batch_rewards, state, action, terminal, self) diff --git a/dqn_agent.py b/dqn_agent.py index 5b3829b..d5c9ee8 100644 --- a/dqn_agent.py +++ b/dqn_agent.py @@ -10,7 +10,18 @@ from policy import * import numpy as np class DQNAgent: - def __init__(self, task_fn, network_fn, optimizer_fn, policy_fn, replay_fn, discount, step_limit, target_network_update_freq): + def __init__(self, + task_fn, + network_fn, + optimizer_fn, + policy_fn, + replay_fn, + discount, + step_limit, + target_network_update_freq, + explore_steps, + history_length, + logger): self.learning_network = network_fn(optimizer_fn) self.target_network = network_fn(optimizer_fn) self.target_network.load_state_dict(self.learning_network.state_dict()) @@ -21,24 +32,37 @@ class DQNAgent: self.target_network_update_freq = target_network_update_freq self.policy = policy_fn() self.total_steps = 0 + self.explore_steps = explore_steps + self.history_length = history_length + self.logger = logger + + def get_state(self, history_buffer): + if self.history_length > 1: + return np.vstack(history_buffer) + return history_buffer[0] def episode(self): state = self.task.reset() + history_buffer = [state] * self.history_length total_reward = 0.0 steps = 0 while not self.step_limit or steps < self.step_limit: - value = self.learning_network.predict(np.reshape(state, (1, -1))) + state = self.get_state(history_buffer) + 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) + 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 - state = next_state + self.logger.debug('steps %d, reward %f, action %d' % (steps, reward, action)) if done: break - experiences = self.replay.sample() - if experiences is not None: + if self.total_steps > self.explore_steps: + experiences = self.replay.sample() states, actions, rewards, next_states, terminals = experiences targets = self.learning_network.predict(states) q_next = self.target_network.predict(next_states) @@ -46,10 +70,13 @@ class DQNAgent: q_next = np.where(terminals, 0, q_next) q_next = rewards + self.discount * q_next targets[np.arange(len(actions)), actions] = q_next + self.logger.debug('start minibatch') self.learning_network.learn(states, targets) + self.logger.debug('minibatch ended') if self.total_steps % self.target_network_update_freq == 0: self.target_network.load_state_dict(self.learning_network.state_dict()) - self.policy.update_epsilon() + if self.total_steps > self.explore_steps: + self.policy.update_epsilon() return total_reward def run(self): @@ -60,10 +87,8 @@ class DQNAgent: ep += 1 reward = self.episode() rewards.append(reward) - if len(rewards) > window_size: - reward = np.mean(rewards[-window_size:]) - print 'episode %d, epsilon %f, reward %f' % ( - ep, self.policy.epsilon, reward) - if reward > self.task.success_threshold: + avg_reward = np.mean(rewards[-window_size:]) + self.logger.info('episode %d, epsilon %f, reward %f, avg reward %f, total steps %d' % ( + ep, self.policy.epsilon, reward, avg_reward, self.total_steps)) + if avg_reward > self.task.success_threshold: break - diff --git a/main.py b/main.py index 9824bc2..ff8a896 100644 --- a/main.py +++ b/main.py @@ -1,12 +1,13 @@ from async_agent import * from dqn_agent import * +import logging def async_cart_pole(): config = dict() config['task_fn'] = lambda: CartPole() config['optimizer_fn'] = lambda params: torch.optim.SGD(params, 0.001) config['network_fn'] = lambda: FullyConnectedNet([4, 50, 200, 2]) - config['policy_fn'] = lambda: GreedyPolicy(epsilon=1.0, end_episode=500, min_epsilon=0.1) + config['policy_fn'] = lambda: GreedyPolicy(epsilon=1.0, final_step=5000, min_epsilon=0.1) # config['bootstrap_fn'] = OneStepQLearning # config['bootstrap_fn'] = NStepQLearning config['bootstrap_fn'] = OneStepSarsa @@ -25,7 +26,7 @@ def async_lunar_lander(): config['task_fn'] = lambda: LunarLander() config['optimizer_fn'] = lambda params: torch.optim.Adam(params, 0.001) config['network_fn'] = lambda: FullyConnectedNet([8, 50, 200, 4]) - config['policy_fn'] = lambda: GreedyPolicy(epsilon=1.0, end_episode=2000, min_epsilon=0.05) + config['policy_fn'] = lambda: GreedyPolicy(epsilon=1.0, final_step=40000, min_epsilon=0.05) config['bootstrap_fn'] = OneStepQLearning config['discount'] = 0.99 config['target_network_update_freq'] = 200 @@ -37,30 +38,19 @@ def async_lunar_lander(): agent = AsyncAgent(**config) agent.run() -# Mountain Car is fairly unstable -def dqn_mountain_car(): - config = dict() - config['task_fn'] = lambda: MountainCar() - config['optimizer_fn'] = lambda params: torch.optim.SGD(params, 0.001) - config['network_fn'] = lambda optimizer_fn: FullyConnectedNet([2, 50, 200, 3], optimizer_fn) - config['policy_fn'] = lambda: GreedyPolicy(epsilon=0.5, end_episode=500, min_epsilon=0.1) - config['replay_fn'] = lambda: Replay(memory_size=10000, batch_size=10) - config['discount'] = 0.99 - config['target_network_update_freq'] = 1000 - config['step_limit'] = 5000 - agent = DQNAgent(**config) - agent.run() - def dqn_cart_pole(): config = dict() config['task_fn'] = lambda: CartPole() config['optimizer_fn'] = lambda params: torch.optim.SGD(params, 0.001) config['network_fn'] = lambda optimizer_fn: FullyConnectedNet([4, 50, 200, 2], optimizer_fn) - config['policy_fn'] = lambda: GreedyPolicy(epsilon=1.0, end_episode=500, min_epsilon=0.1) + config['policy_fn'] = lambda: GreedyPolicy(epsilon=1.0, final_step=10000, min_epsilon=0.1) config['replay_fn'] = lambda: Replay(memory_size=10000, batch_size=10) config['discount'] = 0.99 config['target_network_update_freq'] = 200 - config['step_limit'] = 300 + config['step_limit'] = 0 + config['explore_steps'] = 1000 + config['logger'] = gym.logger + config['history_length'] = 1 agent = DQNAgent(**config) agent.run() @@ -81,9 +71,28 @@ def actor_critic_cart_pole(): 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) + config['policy_fn'] = lambda: GreedyPolicy(epsilon=1.0, final_step=1000000, min_epsilon=0.1) + config['replay_fn'] = lambda: Replay(memory_size=1000000, batch_size=32) + 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 + agent = DQNAgent(**config) + agent.run() + if __name__ == '__main__': + # gym.logger.setLevel(logging.DEBUG) + gym.logger.setLevel(logging.INFO) # async_cart_pole() # async_lunar_lander() # dqn_cart_pole() # dqn_mountain_car() - actor_critic_cart_pole() + # actor_critic_cart_pole() + dqn_pixel_atari('Breakout-v0') diff --git a/network.py b/network.py index d511123..0b48991 100644 --- a/network.py +++ b/network.py @@ -36,10 +36,6 @@ class FullyConnectedNet(nn.Module): y = self.fc3(y) return y - def sync_with(self, src_net): - for param_dst, param_src in zip(self.parameters(), src_net.parameters()): - param_dst.data.copy_(param_src.data) - def predict(self, x): return self.forward(x).cpu().data.numpy() @@ -103,4 +99,49 @@ class ActorCriticNet(nn.Module): def critic(self, x): phi = self.forward(x) - return self.fc_critic(phi).cpu().data.numpy() \ No newline at end of file + return self.fc_critic(phi).cpu().data.numpy() + +class ConvNet(nn.Module): + def __init__(self, in_channels, n_actions, optimizer_fn=None, gpu=True): + super(ConvNet, self).__init__() + self.conv1 = nn.Conv2d(in_channels, 32, kernel_size=8, stride=4) + self.conv2 = nn.Conv2d(32, 64, kernel_size=4, stride=2) + self.conv3 = nn.Conv2d(64, 64, kernel_size=3, stride=1) + self.fc4 = nn.Linear(7 * 7 * 64, 512) + self.fc5 = nn.Linear(512, n_actions) + + self.criterion = nn.MSELoss() + if optimizer_fn is not None: + self.optimizer = optimizer_fn(self.parameters()) + + self.gpu = gpu and torch.cuda.is_available() + if self.gpu: + print 'Transferring network to GPU...' + self.cuda() + print 'Network transferred.' + + def to_torch_variable(self, x): + x = torch.from_numpy(np.asarray(x, dtype='float32')) + if self.gpu: + x = x.cuda() + return Variable(x) + + def forward(self, x): + y = F.relu(self.conv1(x)) + y = F.relu(self.conv2(y)) + y = F.relu(self.conv3(y)) + y = y.view(y.size(0), -1) + y = F.relu(self.fc4(y)) + return self.fc5(y) + + def predict(self, x): + return self.forward(self.to_torch_variable(x)).cpu().data.numpy() + + def learn(self, x, target): + x = self.to_torch_variable(x) + target = self.to_torch_variable(target) + y = self.forward(x) + loss = self.criterion(y, target) + self.zero_grad() + loss.backward() + self.optimizer.step() diff --git a/policy.py b/policy.py index c30c91f..71cd9d0 100644 --- a/policy.py +++ b/policy.py @@ -7,11 +7,11 @@ import numpy as np class GreedyPolicy: - def __init__(self, epsilon, end_episode, min_epsilon): + def __init__(self, epsilon, final_step, min_epsilon): self.init_epsilon = self.epsilon = epsilon - self.current_episode = 0 + self.current_steps = 0 self.min_epsilon = min_epsilon - self.end_episode = end_episode + self.final_step = final_step def sample(self, action_value): if np.random.rand() < self.epsilon: @@ -19,9 +19,9 @@ class GreedyPolicy: return np.argmax(action_value) def update_epsilon(self): - self.epsilon = self.init_epsilon - float(self.current_episode) / self.end_episode * (self.init_epsilon - self.min_epsilon) + self.epsilon = self.init_epsilon - float(self.current_steps) / self.final_step * (self.init_epsilon - self.min_epsilon) self.epsilon = max(self.epsilon, self.min_epsilon) - self.current_episode += 1 + self.current_steps += 1 class SamplePolicy: def sample(self, action_value): diff --git a/replay.py b/replay.py index 69234e5..1e7cb89 100644 --- a/replay.py +++ b/replay.py @@ -37,10 +37,21 @@ class Replay: sampled_indices = np.arange(len(self.terminals)) np.random.shuffle(sampled_indices) sampled_indices = sampled_indices[: self.batch_size] - return [np.asarray(self.states)[sampled_indices], - np.asarray(self.actions)[sampled_indices], - np.asarray(self.rewards)[sampled_indices], - np.asarray(self.next_states)[sampled_indices], - np.asarray(self.terminals)[sampled_indices]] + sampled_states = [] + sampled_actions = [] + sampled_rewards = [] + sampled_next_states = [] + sampled_terminals = [] + for ind in sampled_indices: + sampled_states.append(self.states[ind]) + sampled_actions.append(self.actions[ind]) + sampled_rewards.append(self.rewards[ind]) + sampled_next_states.append(self.next_states[ind]) + sampled_terminals.append(self.terminals[ind]) + return [np.asarray(sampled_states), + np.asarray(sampled_actions), + np.asarray(sampled_rewards), + np.asarray(sampled_next_states), + np.asarray(sampled_terminals)] return None diff --git a/task.py b/task.py index 25efd48..4c2cddb 100644 --- a/task.py +++ b/task.py @@ -5,6 +5,8 @@ ####################################################################### import gym import sys +import numpy as np +import cv2 class BasicTask: def transfer_state(self, state): @@ -16,7 +18,7 @@ class BasicTask: def step(self, action): next_state, reward, done, info = self.env.step(action) next_state = self.transfer_state(next_state) - return next_state, reward, done, info + return next_state, np.sign(reward), done, info class MountainCar(BasicTask): name = 'MountainCar-v0' @@ -38,4 +40,17 @@ class LunarLander(BasicTask): success_threshold = 200 def __init__(self): - self.env = gym.make(self.name) \ No newline at end of file + self.env = gym.make(self.name) + +class PixelAtari(BasicTask): + width = 84 + height = 84 + success_threshold = 1000 + + def __init__(self, name): + 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 = cv2.resize(img, (self.width, self.height)) + return np.reshape(img, (1, self.width, self.height))