diff --git a/async_agent.py b/async_agent.py index 592e929..b43f532 100644 --- a/async_agent.py +++ b/async_agent.py @@ -13,8 +13,20 @@ from network import * from bootstrap import * class AsyncAgent: - def __init__(self, task_fn, network_fn, optimizer_fn, policy_fn, bootstrap_fn, discount, step_limit, - target_network_update_freq, n_workers, batch_size, test_interval, test_repeats): + def __init__(self, + task_fn, + network_fn, + optimizer_fn, + policy_fn, + bootstrap_fn, + discount, + step_limit, + target_network_update_freq, + n_workers, + batch_size, + test_interval, + test_repeats, + logger): self.network_fn = network_fn self.learning_network = network_fn() self.learning_network.share_memory() @@ -38,12 +50,14 @@ class AsyncAgent: self.batch_size = batch_size self.test_interval = test_interval self.test_repeats = test_repeats + self.logger = logger def deterministic_episode(self, task, network): state = np.asarray([task.reset()]) total_rewards = 0 steps = 0 - while True and steps < self.step_limit: + terminal = False + while not terminal and steps < self.step_limit: action_values = network.predict(state) steps += 1 action = np.argmax(action_values.flatten()) @@ -75,7 +89,7 @@ class AsyncAgent: batch_states, batch_actions, batch_rewards = [], [], [] if terminal: if id == 0: - print 'worker %d, episode %d, return %f' % (id, episode, episode_return) + self.logger.debug('worker %d, episode %d, return %f' % (id, episode, episode_return)) episode_steps = 0 episode_return = 0 episode += 1 @@ -86,8 +100,7 @@ class AsyncAgent: action = policy.sample(value.flatten()) while not terminal and len(batch_states) < self.batch_size: episode_steps += 1 - with self.steps_lock: - self.total_steps.value += 1 + self.total_steps.value += 1 batch_states.append(state) batch_actions.append(action) state, reward, terminal, _ = task.step(action) @@ -119,14 +132,15 @@ class AsyncAgent: task = self.task_fn() test_network = self.network_fn() while True: - if self.total_steps.value % self.test_interval == 0: + steps = self.total_steps.value + 1 + if steps % self.test_interval == 0: with self.network_lock: test_network.load_state_dict(self.learning_network.state_dict()) rewards = np.zeros(self.test_repeats) for i in range(self.test_repeats): rewards[i] = self.deterministic_episode(task, test_network) - print 'total steps: %d, averaged return per episode: %f' %\ - (self.total_steps.value, np.mean(rewards)) + self.logger.info('total steps: %d, averaged return per episode: %f' %\ + (steps, np.mean(rewards))) if np.mean(rewards) > task.success_threshold: self.stop_signal.value = True break diff --git a/atari_wrapper.py b/atari_wrapper.py new file mode 100644 index 0000000..a95b43b --- /dev/null +++ b/atari_wrapper.py @@ -0,0 +1,133 @@ +# This file is copied/apdated from +# https://raw.githubusercontent.com/transedward/pytorch-dqn/master/utils/atari_wrapper.py + +import numpy as np +from collections import deque +import gym +from gym import spaces +from PIL import Image + +class NoopResetEnv(gym.Wrapper): + def __init__(self, env=None, noop_max=30): + """Sample initial states by taking random number of no-ops on reset. + No-op is assumed to be action 0. + """ + super(NoopResetEnv, self).__init__(env) + self.noop_max = noop_max + assert env.unwrapped.get_action_meanings()[0] == 'NOOP' + + def _reset(self): + """ Do no-op action for a number of steps in [1, noop_max].""" + self.env.reset() + noops = np.random.randint(1, self.noop_max + 1) + for _ in range(noops): + obs, _, _, _ = self.env.step(0) + return obs + +class FireResetEnv(gym.Wrapper): + def __init__(self, env=None): + """Take action on reset for environments that are fixed until firing.""" + super(FireResetEnv, self).__init__(env) + assert env.unwrapped.get_action_meanings()[1] == 'FIRE' + assert len(env.unwrapped.get_action_meanings()) >= 3 + + def _reset(self): + self.env.reset() + obs, _, _, _ = self.env.step(1) + obs, _, _, _ = self.env.step(2) + return obs + +class EpisodicLifeEnv(gym.Wrapper): + def __init__(self, env=None): + """Make end-of-life == end-of-episode, but only reset on true game over. + Done by DeepMind for the DQN and co. since it helps value estimation. + """ + super(EpisodicLifeEnv, self).__init__(env) + self.lives = 0 + self.was_real_done = True + self.was_real_reset = False + + def _step(self, action): + obs, reward, done, info = self.env.step(action) + self.was_real_done = done + # check current lives, make loss of life terminal, + # then update lives to handle bonus lives + lives = self.env.unwrapped.ale.lives() + if lives < self.lives and lives > 0: + # for Qbert somtimes we stay in lives == 0 condtion for a few frames + # so its important to keep lives > 0, so that we only reset once + # the environment advertises done. + done = True + self.lives = lives + return obs, reward, done, info + + def _reset(self): + """Reset only when lives are exhausted. + This way all states are still reachable even though lives are episodic, + and the learner need not know about any of this behind-the-scenes. + """ + if self.was_real_done: + obs = self.env.reset() + self.was_real_reset = True + else: + # no-op step to advance from terminal/lost life state + obs, _, _, _ = self.env.step(0) + self.was_real_reset = False + self.lives = self.env.unwrapped.ale.lives() + return obs + +class MaxAndSkipEnv(gym.Wrapper): + def __init__(self, env=None, skip=4): + """Return only every `skip`-th frame""" + super(MaxAndSkipEnv, self).__init__(env) + # most recent raw observations (for max pooling across time steps) + self._obs_buffer = deque(maxlen=2) + 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) + self._obs_buffer.append(obs) + total_reward += reward + if done: + break + + max_frame = np.max(np.stack(self._obs_buffer), axis=0) + + return max_frame, total_reward, done, info + + def _reset(self): + """Clear past frame buffer and init. to first obs. from inner env.""" + self._obs_buffer.clear() + obs = self.env.reset() + self._obs_buffer.append(obs) + return obs + +def _process_frame84(frame): + img = np.reshape(frame, [210, 160, 3]).astype(np.float32) + img = img[:, :, 0] * 0.299 + img[:, :, 1] * 0.587 + img[:, :, 2] * 0.114 + img = Image.fromarray(img) + resized_screen = img.resize((84, 110), Image.BILINEAR) + resized_screen = np.array(resized_screen) + x_t = resized_screen[18:102, :] + x_t = np.reshape(x_t, [1, 84, 84]) + return x_t.astype(np.uint8) + +class ProcessFrame84(gym.Wrapper): + def __init__(self, env=None): + super(ProcessFrame84, self).__init__(env) + self.observation_space = spaces.Box(low=0, high=255, shape=(1, 84, 84)) + + def _step(self, action): + obs, reward, done, info = self.env.step(action) + return _process_frame84(obs), reward, done, info + + def _reset(self): + return _process_frame84(self.env.reset()) + +class ClippedRewardsWrapper(gym.Wrapper): + def _step(self, action): + obs, reward, done, info = self.env.step(action) + return obs, np.sign(reward), done, info diff --git a/dqn_agent.py b/dqn_agent.py index f717696..7f77442 100644 --- a/dqn_agent.py +++ b/dqn_agent.py @@ -25,6 +25,8 @@ class DQNAgent: target_network_update_freq, explore_steps, history_length, + test_interval, + test_repetitions, logger): self.learning_network = network_fn(optimizer_fn) self.target_network = network_fn(optimizer_fn) @@ -40,12 +42,13 @@ class DQNAgent: self.history_length = history_length self.logger = logger self.process = psutil.Process(os.getpid()) - self.report_interval = 1000 + self.test_interval = test_interval + self.test_repetitions = test_repetitions def get_state(self, history_buffer): return np.vstack(history_buffer) - def episode(self): + def episode(self, deterministic=False): episode_start_time = time.time() state = self.task.reset() history_buffer = [state] * self.history_length @@ -55,40 +58,34 @@ class DQNAgent: 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()) + if deterministic: + action = np.argmax(value.flatten()) + else: + 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) - self.replay.feed([state, action, reward, next_state, int(done)]) + if not deterministic: + next_state = self.get_state(history_buffer) + self.replay.feed([state, action, reward, next_state, int(done)]) + self.total_steps += 1 total_reward += reward steps += 1 - self.total_steps += 1 if done: break - if self.total_steps > self.explore_steps: - sample_start_time = time.time() + if not deterministic and self.total_steps > self.explore_steps: experiences = self.replay.sample() - 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() q_next = self.target_network.predict(next_states) - 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 - minibatch_start_time = time.time() self.learning_network.learn(states, actions, q_next) - # self.learning_network.clippedLearn(states, actions, q_next) - 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: + if not deterministic and 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: + if not deterministic and self.total_steps > self.explore_steps: self.policy.update_epsilon() episode_time = time.time() - episode_start_time info = self.process.memory_info() @@ -107,11 +104,19 @@ class DQNAgent: while True: ep += 1 reward = self.episode() - if ep % 1000 == 0: - self.save('data/dqn-episode-%d.bin' % (ep)) 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' % ( ep, self.policy.epsilon, reward, avg_reward, self.total_steps)) - if avg_reward > self.task.success_threshold: - break + + if ep % self.test_interval == 0: + self.logger.info('Testing...') + self.save('data/dqn-episode-%d.bin' % (ep)) + test_rewards = [] + for _ in range(self.test_repetitions): + test_rewards.append(self.episode(True)) + avg_reward = np.mean(test_rewards) + self.logger.info('Avg reward %f(%f)' % ( + avg_reward, np.std(test_rewards) / np.sqrt(self.test_repetitions))) + if avg_reward > self.task.success_threshold: + break \ No newline at end of file diff --git a/main.py b/main.py index 1ea6b24..4efa602 100644 --- a/main.py +++ b/main.py @@ -51,32 +51,36 @@ def dqn_cart_pole(): config['explore_steps'] = 1000 config['logger'] = gym.logger config['history_length'] = 2 + config['test_interval'] = 100 + config['test_repetitions'] = 50 agent = DQNAgent(**config) agent.run() def actor_critic_cart_pole(): config = dict() config['task_fn'] = lambda: CartPole() - config['optimizer_fn'] = lambda params: torch.optim.SGD(params, 0.001) + config['optimizer_fn'] = lambda params: torch.optim.Adam(params, 0.001) config['network_fn'] = lambda: ActorCriticNet([4, 200, 2]) config['policy_fn'] = SamplePolicy config['bootstrap_fn'] = AdvantageActorCritic config['discount'] = 0.99 config['target_network_update_freq'] = 200 - config['step_limit'] = 300 - config['n_workers'] = 8 + config['step_limit'] = 200 + config['n_workers'] = 10 config['batch_size'] = 5 config['test_interval'] = 50000 config['test_repeats'] = 5 + config['logger'] = gym.logger agent = AsyncAgent(**config) agent.run() def dqn_pixel_atari(name): config = dict() history_length = 4 - config['task_fn'] = lambda: PixelAtari(name, 30) + n_actions = 6 + config['task_fn'] = lambda: PixelAtari(name, 30, 4) 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['network_fn'] = lambda optimizer_fn: ConvNet(history_length, n_actions, 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, dtype=np.uint8) config['discount'] = 0.99 @@ -85,15 +89,18 @@ def dqn_pixel_atari(name): config['explore_steps'] = 50000 config['logger'] = gym.logger config['history_length'] = history_length + config['test_interval'] = 1000 + config['test_repetitions'] = 50 agent = DQNAgent(**config) agent.run() if __name__ == '__main__': - gym.logger.setLevel(logging.DEBUG) - # gym.logger.setLevel(logging.INFO) + # gym.logger.setLevel(logging.DEBUG) + gym.logger.setLevel(logging.INFO) + benchmark = gym.benchmark_spec('Atari40M') + # async_cart_pole() # async_lunar_lander() + # actor_critic_cart_pole() # dqn_cart_pole() - actor_critic_cart_pole() - # dqn_pixel_atari('Breakout-v0') - # dqn_pixel_atari('SpaceInvaders-v0') + dqn_pixel_atari('BreakoutNoFrameskip-v3') diff --git a/network.py b/network.py index e00212c..8965426 100644 --- a/network.py +++ b/network.py @@ -10,6 +10,7 @@ import torch.nn as nn import torch.nn.functional as F import numpy as np + class FullyConnectedNet(nn.Module): def __init__(self, dims, optimizer_fn=None, gpu=True): super(FullyConnectedNet, self).__init__() @@ -44,39 +45,33 @@ class FullyConnectedNet(nn.Module): return Variable(x) def learn(self, x, actions, targets): + self.zero_grad() + self.gradient(x, actions, targets) + self.optimizer.step() + + # def clippedLearn(self, x, actions, targets): + # y = self.forward(x) + # actions = self.to_torch_variable(actions, 'int64').unsqueeze(1) + # targets = self.to_torch_variable(targets).unsqueeze(1) + # y = y.gather(1, actions) + # bellman_error = targets - y + # bellman_error = bellman_error.clamp(-1, 1) * -1 + # self.zero_grad() + # y.backward(bellman_error.data) + # self.optimizer.step() + + def gradient(self, x, actions, targets): y = self.forward(x) actions = self.to_torch_variable(actions, 'int64').unsqueeze(1) targets = self.to_torch_variable(targets).unsqueeze(1) y = y.gather(1, actions) - error = -(targets - y) * 0.5 - self.zero_grad() - y.backward(error.data) - self.optimizer.step() - - - def clippedLearn(self, x, actions, targets): - y = self.forward(x) - actions = self.to_torch_variable(actions, 'int64').unsqueeze(1) - targets = self.to_torch_variable(targets).unsqueeze(1) - y = y.gather(1, actions) - bellman_error = targets - y - bellman_error = bellman_error.clamp(-1, 1) * -1 - self.zero_grad() - y.backward(bellman_error.data) - self.optimizer.step() - - - def gradient(self, x, actions, rewards): - y = self.forward(x) - target = np.copy(y.data.numpy()) - target[np.arange(target.shape[0]), actions] = np.asarray(rewards) - target = Variable(torch.from_numpy(target)) - loss = self.criterion(y, target) + loss = self.criterion(y, targets) loss.backward() def output_transfer(self, y): return y + class ActorCriticNet(nn.Module): def __init__(self, dims, gpu=True): super(ActorCriticNet, self).__init__() @@ -89,12 +84,14 @@ class ActorCriticNet(nn.Module): self.cuda() print 'Network transferred.' - def forward(self, x): - x = torch.from_numpy(np.asarray(x, dtype='float32')) + def to_torch_variable(self, x, dtype='float32'): + x = torch.from_numpy(np.asarray(x, dtype=dtype)) if self.gpu: x = x.cuda() - x = Variable(x) - phi = self.fc1(x) + return Variable(x) + + def forward(self, x): + phi = self.fc1(self.to_torch_variable(x)) return phi def predict(self, x): @@ -104,19 +101,23 @@ class ActorCriticNet(nn.Module): def gradient(self, x, actions, rewards): phi = self.forward(x) logit = self.fc_actor(phi) - log_prob = F.log_softmax(logit) + prob = F.softmax(logit) + log_prob_ = F.log_softmax(logit) state_value = self.fc_critic(phi) - log_prob = log_prob.gather(1, Variable(torch.from_numpy(np.asarray([actions]).reshape([-1, 1])))) + log_prob = log_prob_.gather(1, self.to_torch_variable(np.asarray([actions]).reshape([-1, 1]), 'int64')) advantage = np.asarray([rewards]).reshape([-1, 1]) - state_value.cpu().data.numpy() - policy_loss = -torch.sum(log_prob * Variable(torch.from_numpy(np.asarray(advantage, dtype='float32')))) - value_loss = 0.5 * torch.sum(torch.pow(state_value - Variable(torch.from_numpy(np.asarray(rewards, dtype='float32'))), 2)) - (policy_loss + value_loss).backward() + policy_loss = -torch.sum(log_prob * self.to_torch_variable(advantage)) + value_loss = 0.5 * torch.sum( + torch.pow(state_value - Variable(torch.from_numpy(np.asarray(rewards, dtype='float32'))), 2)) + entropy = -torch.sum(torch.mul(prob, log_prob_)) + (policy_loss + value_loss - 0.01 * entropy).backward() nn.utils.clip_grad_norm(self.parameters(), 40) def critic(self, x): phi = self.forward(x) 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__() @@ -154,22 +155,14 @@ class ConvNet(nn.Module): return self.forward(self.to_torch_variable(x)).cpu().data.numpy() def learn(self, x, actions, targets): - y = self.forward(self.to_torch_variable(x)) - actions = self.to_torch_variable(actions, 'int64').unsqueeze(1) - targets = self.to_torch_variable(targets).unsqueeze(1) - y = y.gather(1, actions) - error = -(targets - y) * 0.5 self.zero_grad() - y.backward(error.data) + self.gradient(x, actions, targets) self.optimizer.step() - def clippedLearn(self, x, actions, targets): - y = self.forward(self.to_torch_variable(x)) + def gradient(self, x, actions, targets): + y = self.forward(x) actions = self.to_torch_variable(actions, 'int64').unsqueeze(1) targets = self.to_torch_variable(targets).unsqueeze(1) y = y.gather(1, actions) - bellman_error = -(targets - y) - bellman_error = bellman_error.clamp(-1, 1) - self.zero_grad() - y.backward(bellman_error.data) - self.optimizer.step() + loss = self.criterion(y, targets) + loss.backward() diff --git a/task.py b/task.py index fbd0ff5..49629f1 100644 --- a/task.py +++ b/task.py @@ -6,11 +6,9 @@ import gym import sys import numpy as np -import cv2 +from atari_wrapper import * class BasicTask: - no_op = 0 - def transfer_state(self, state): return state @@ -19,9 +17,6 @@ class BasicTask: def reset(self): 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): @@ -52,35 +47,18 @@ class LunarLander(BasicTask): self.env = gym.make(self.name) class PixelAtari(BasicTask): - width = 84 - height = 84 success_threshold = 1000 - def __init__(self, name, no_op): - self.no_op = no_op - self.env = gym.make(name) - self.done = True - self.lives = 0 - - def reset(self): - if self.done: - return BasicTask.reset(self) - else: - state, _, _, _ = BasicTask.step(self, 0) - return state - - def step(self, action): - next_state, reward, done, info = BasicTask.step(self, action) - self.done = done - if self.lives > 0 and info['ale.lives'] < self.lives: - done = True - self.lives = info['ale.lives'] - return next_state, reward, done, info - - def transfer_state(self, state): - img = cv2.cvtColor(state, cv2.COLOR_RGB2GRAY) - img = cv2.resize(img, (self.width, self.height)) - return np.asarray(np.reshape(img, (1, self.width, self.height)), np.uint8) + def __init__(self, name, no_op, frame_skip): + env = gym.make(name) + assert 'NoFrameskip' in env.spec.id + env = EpisodicLifeEnv(env) + env = NoopResetEnv(env, noop_max=no_op) + env = MaxAndSkipEnv(env, skip=frame_skip) + if 'FIRE' in env.unwrapped.get_action_meanings(): + env = FireResetEnv(env) + env = ProcessFrame84(env) + self.env = ClippedRewardsWrapper(env) def normalize_state(self, state): return np.asarray(state, dtype=np.float32) / 255.0