Refactor atari task

dqn_agent.py

@@ -75,16 +75,15 @@ class DQNAgent:
                 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)
                 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.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:

main.py

@@ -93,7 +93,7 @@ if __name__ == '__main__':
     # gym.logger.setLevel(logging.INFO)
     # async_cart_pole()
     # async_lunar_lander()
-    dqn_cart_pole()
-    # actor_critic_cart_pole()
+    # dqn_cart_pole()
+    actor_critic_cart_pole()
     # dqn_pixel_atari('Breakout-v0')
     # dqn_pixel_atari('SpaceInvaders-v0')

network.py

@@ -27,10 +27,7 @@ class FullyConnectedNet(nn.Module):
 
     def forward(self, x):
         x = x.reshape((x.shape[0], -1))
-        x = torch.from_numpy(np.asarray(x, dtype='float32'))
-        if self.gpu:
-            x = x.cuda()
-        x = Variable(x)
+        x = self.to_torch_variable(x)
 
         y = F.relu(self.fc1(x))
         y = F.relu(self.fc2(y))
@@ -40,17 +37,35 @@ class FullyConnectedNet(nn.Module):
     def predict(self, x):
         return self.forward(x).cpu().data.numpy()
 
-    def learn(self, x, target):
-        target = torch.from_numpy(target)
+    def to_torch_variable(self, x, dtype='float32'):
+        x = torch.from_numpy(np.asarray(x, dtype=dtype))
         if self.gpu:
-            target = target.cuda()
-        target = Variable(target)
+            x = x.cuda()
+        return Variable(x)
+
+    def learn(self, x, actions, targets):
         y = self.forward(x)
-        loss = self.criterion(y, target)
+        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()
-        loss.backward()
+        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())
@@ -121,8 +136,8 @@ class ConvNet(nn.Module):
             self.cuda()
             print 'Network transferred.'
 
-    def to_torch_variable(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()
         return Variable(x)
@@ -138,11 +153,23 @@ class ConvNet(nn.Module):
     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)
+    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()
-        loss.backward()
+        y.backward(error.data)
+        self.optimizer.step()
+
+    def clippedLearn(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)
+        bellman_error = -(targets - y)
+        bellman_error = bellman_error.clamp(-1, 1)
+        self.zero_grad()
+        y.backward(bellman_error.data)
         self.optimizer.step()

task.py

@@ -59,9 +59,26 @@ class PixelAtari(BasicTask):
     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 = (state[:, :, 0] * 0.299 + state[:, :, 1] * 0.587 + state[:, :, 2] * 0.114)
+        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)