Atari wrapper

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

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

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

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')

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()

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