MP

async_agent.py

@@ -0,0 +1,161 @@
+#######################################################################
+# Copyright (C) 2017 Shangtong Zhang(zhangshangtong.cpp@gmail.com)    #
+# Permission given to modify the code as long as you keep this        #
+# declaration at the top                                              #
+#######################################################################
+
+from network import *
+from policy import *
+import numpy as np
+import torch.multiprocessing as mp
+import time
+from task import *
+from network import *
+from torch.autograd import Variable
+import thread
+
+class AsyncAgent:
+    def __init__(self, task_fn, network_fn, optimizer_fn, policy_fn, discount, step_limit,
+                 target_network_update_freq, n_workers):
+        self.network_fn = network_fn
+        self.learning_network = network_fn()
+        self.learning_network.share_memory()
+        self.target_network = network_fn()
+        self.target_network.share_memory()
+        self.target_network.load_state_dict(self.learning_network.state_dict())
+
+        self.optimizer_fn = optimizer_fn
+        # self.optimizer = optimizer_fn(self.learning_network.parameters())
+        self.task_fn = task_fn
+        self.step_limit = step_limit
+        self.discount = discount
+        self.target_network_update_freq = target_network_update_freq
+        self.policy = policy_fn()
+        self.total_steps = mp.Value('i', 0)
+        self.lock = mp.Lock()
+        self.n_workers = n_workers
+
+    def async_update(self, worker_network, optimizer):
+        with self.lock:
+            optimizer.zero_grad()
+            for param, worker_param in zip(self.learning_network.parameters(), worker_network.parameters()):
+                param._grad = worker_param.grad
+            # print list(self.learning_network.parameters())[1].data
+            # print list(worker_network.parameters())[1].grad.data
+            optimizer.step()
+            # print list(self.learning_network.parameters())[1].data
+
+    def worker(self, id):
+        worker_network = self.network_fn()
+        task = self.task_fn()
+        episode = 0
+        optimizer = self.optimizer_fn(self.learning_network.parameters())
+        while True:
+            worker_network.load_state_dict(self.learning_network.state_dict())
+            worker_network.zero_grad()
+            state = np.reshape(task.reset(), (1, -1))
+            steps = 0
+            total_reward = 0
+            while not self.step_limit or steps < self.step_limit:
+                value = worker_network.predict(state)
+                action = self.policy.sample(value.flatten())
+                next_state, reward, done, info = task.step(action)
+                next_state = np.reshape(next_state, (1, -1))
+                q_next = self.learning_network.predict(next_state)
+                # q_next = self.target_network.predict(next_state)
+                q_next = np.max(q_next, axis=1)
+                if done:
+                    q_next = 0
+                q_next = reward + self.discount * q_next
+                value[0, action] = q_next
+                # if (steps > 0):
+                #     print list(worker_network.parameters())[1].grad.data
+                # worker_network.zero_grad()
+                worker_network.gradient(state, value)
+                # print list(worker_network.parameters())[1].grad
+                # worker_network.zero_grad()
+                # worker_network.gradient(state, value)
+                # print list(worker_network.parameters())[1].grad
+                # worker_network.gradient(state, value)
+                # print list(worker_network.parameters())[1].grad
+
+                steps += 1
+                total_reward += reward
+                with self.lock:
+                    self.total_steps.value += 1
+                    if self.total_steps.value % self.target_network_update_freq == 0:
+                        self.target_network.load_state_dict(self.learning_network.state_dict())
+                if done:
+                    break
+                state = next_state
+            print 'worker %d, episode %d, rewards %d' % (id, episode, total_reward)
+            episode += 1
+            self.async_update(worker_network, optimizer)
+            self.policy.update_epsilon()
+
+    def run(self):
+        procs = [mp.Process(target=self.worker, args=(i, )) for i in range(self.n_workers)]
+        for p in procs: p.start()
+        for p in procs: p.join()
+        # print list(self.learning_network.parameters())[1].data
+
+class Test:
+    def __init__(self):
+        # self.data = np.zeros((2, 3))
+        self.data = torch.zeros((2, 3))
+        self.data.share_memory_()
+        print self.data
+        self.val = mp.Value('i', 0)
+        self.lock = mp.Lock()
+
+    def fun(self):
+        # self.data += rank
+        for i in range(50):
+            time.sleep(0.01)
+            with self.lock:
+                self.val.value += 1
+
+    def run(self):
+        processes = []
+        for i in range(3):
+            processes.append(mp.Process(target=self.fun))
+        for p in processes:
+            p.start()
+        for p in processes:
+            p.join()
+
+class TestNet(nn.Module):
+    def __init__(self):
+        super(TestNet, self).__init__()
+        self.fc = nn.Linear(2, 1, bias=False)
+        for param in self.parameters():
+            param.data.copy_(torch.from_numpy(np.array([[1.0, 2.0]], dtype='float32').T))
+        self.criterion = nn.MSELoss()
+    def gradient(self, x, target):
+        y = self.fc(x)
+        loss = self.criterion(y, target)
+        loss.backward()
+
+# t = Test()
+# t.run()
+# print t.val.value
+# print t.data
+
+if __name__ == '__main__':
+    task_fn = lambda: CartPole()
+    network_fn = lambda: FullyConnectedNet([4, 50, 200, 2])
+    optimizer_fn = lambda params: torch.optim.SGD(params, 0.01)
+    policy_fn = lambda: GreedyPolicy(epsilon=0.5, decay_factor=0.99, min_epsilon=0.01)
+    # config = {'discount': 0.99, 'step_limit': 5000, 'target_network_update_freq': 200}
+    agent = AsyncAgent(task_fn, network_fn, optimizer_fn, policy_fn, 0.99, 0, 200, 6)
+    agent.run()
+
+    # t = TestNet()
+    # x = Variable(torch.from_numpy(np.array([[0.1, 0.2]], dtype='float32')))
+    # target = Variable(torch.from_numpy(np.array([1], dtype='float32')))
+    # t.gradient(x, target)
+    # for p in t.parameters(): print p.grad
+    # t.gradient(x, target)
+    # for p in t.parameters(): print p.grad
+    # t.gradient(x, target)
+    # for p in t.parameters(): print p.grad

dqn_agent.py

@@ -13,6 +13,7 @@ class DQNAgent:
     def __init__(self, task, network_fn, optimizer_fn, policy_fn, replay_fn, discount, step_limit, target_network_update_freq):
         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())
         self.task = task
         self.step_limit = step_limit
         self.replay = replay_fn()
@@ -26,7 +27,7 @@ class DQNAgent:
         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)) )
+            value = self.learning_network.predict(np.reshape(state, (1, -1)))
             action = self.policy.sample(value.flatten())
             next_state, reward, done, info = self.task.step(action)
             total_reward += reward
@@ -42,8 +43,8 @@ class DQNAgent:
                 targets = self.learning_network.predict(states)
                 q_next = self.target_network.predict(next_states)
                 q_next = np.max(q_next, axis=1)
-                q_next = rewards + self.discount * q_next
                 q_next = np.where(terminals, 0, q_next)
+                q_next = rewards + self.discount * q_next
                 targets[np.arange(len(actions)), actions] = q_next
                 self.learning_network.learn(states, targets)
             if self.total_steps % self.target_network_update_freq == 0:

network.py

@@ -11,13 +11,14 @@ import torch.nn.functional as F
 import numpy as np
 
 class FullyConnectedNet(nn.Module):
-    def __init__(self, dims, optimizer_fn, gpu=True):
+    def __init__(self, dims, optimizer_fn=None, gpu=True):
         super(FullyConnectedNet, self).__init__()
         self.fc1 = nn.Linear(dims[0], dims[1])
         self.fc2 = nn.Linear(dims[1], dims[2])
         self.fc3 = nn.Linear(dims[2], dims[3])
         self.criterion = nn.MSELoss()
-        self.optimizer = optimizer_fn(self.parameters())
+        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...'
@@ -54,5 +55,11 @@ class FullyConnectedNet(nn.Module):
         loss.backward()
         self.optimizer.step()
 
+    def gradient(self, x, target):
+        y = self.forward(x)
+        target = Variable(torch.from_numpy(target))
+        loss = self.criterion(y, target)
+        loss.backward()
+
     def output_transfer(self, y):
         return y