Refactor APIs

async_agent.py

@@ -8,161 +8,113 @@ 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 threading
 
 class AsyncAgent:
     def __init__(self, task_fn, network_fn, optimizer_fn, policy_fn, discount, step_limit,
-                 target_network_update_freq, n_workers):
+                 target_network_update_freq, n_workers, batch_size, test_interval):
         self.network_fn = network_fn
         self.learning_network = network_fn()
-        # self.learning_network.share_memory()
+        self.learning_network.share_memory()
         self.target_network = network_fn()
-        # self.target_network.share_memory()
+        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.optimizer_fn = optimizer_fn
         self.task_fn = task_fn
         self.step_limit = step_limit
         self.discount = discount
+        self.optimizer_fn = optimizer_fn
         self.target_network_update_freq = target_network_update_freq
-        # self.policy = policy_fn()
         self.policy_fn = policy_fn
-        # self.total_steps = mp.Value('i', 0)
-        self.total_steps = 0
-        self.lock = threading.Lock()
-        # self.lock = mp.Lock()
+        self.steps_lock = mp.Lock()
+        self.network_lock = mp.Lock()
+        self.total_steps = mp.Value('i', 0)
+        self.stop_signal = mp.Value('i', False)
         self.n_workers = n_workers
-        self.batch_size = 5
+        self.batch_size = batch_size
+        self.test_interval = test_interval
+
+    def deterministic_episode(self, task):
+        state = np.asarray([task.reset()])
+        total_rewards = 0
+        steps = 0
+        while True and steps < self.step_limit:
+            with self.network_lock:
+                action_values = self.learning_network.predict(state)
+            steps += 1
+            action = np.argmax(action_values.flatten())
+            state, reward, terminal, _ = task.step(action)
+            total_rewards += reward
+            if terminal:
+                break
+            state = state.reshape([1, -1])
+        return total_rewards
 
     def async_update(self, worker_network, optimizer):
-        with self.lock:
+        with self.network_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
+                param._grad = worker_param.grad.clone()
             optimizer.step()
-            # print list(self.learning_network.parameters())[1].data
 
     def worker(self, id):
-        # worker_network = self.network_fn()
+        optimizer = self.optimizer_fn(self.learning_network.parameters())
+        worker_network = self.network_fn()
+        worker_network.load_state_dict(self.learning_network.state_dict())
         task = self.task_fn()
-        episode = 0
-        # optimizer = self.optimizer_fn(self.learning_network.parameters())
         policy = self.policy_fn()
         terminal = True
+        episode = 0
         episode_steps = 0
-        while True:
+        while True and not self.stop_signal.value:
             batch_states, batch_actions, batch_rewards = [], [], []
             if terminal:
-                if id == 0:
-                    print 'episode %d, epsilon %f, steps: %d' % \
-                          (episode, policy.epsilon, episode_steps)
                 episode_steps = 0
                 episode += 1
                 policy.update_epsilon()
                 terminal = False
                 state = task.reset()
-                state = np.reshape(state, (1, -1))
+                state = state.reshape([1, -1])
             while not terminal and len(batch_states) < self.batch_size:
                 episode_steps += 1
-                with self.lock:
-                    self.total_steps += 1
+                with self.steps_lock:
+                    self.total_steps.value += 1
                 batch_states.append(state)
-                value = self.learning_network.predict(state)
+                value = worker_network.predict(state)
                 action = policy.sample(value.flatten())
                 batch_actions.append(action)
                 state, reward, terminal, _ = task.step(action)
-                state = np.reshape(state, (1, -1))
+                state = state.reshape([1, -1])
                 if not terminal:
-                    q_next = np.max(self.target_network.predict(state))
+                    with self.network_lock:
+                        q_next = np.max(self.target_network.predict(state))
                     reward += self.discount * q_next
                 batch_rewards.append(reward)
 
-            self.learning_network.learn_from_raw(np.vstack(batch_states), batch_actions, batch_rewards)
+            worker_network.zero_grad()
+            worker_network.gradient(np.vstack(batch_states), batch_actions, batch_rewards)
+            self.async_update(worker_network, optimizer)
+            worker_network.load_state_dict(self.learning_network.state_dict())
 
-            if self.total_steps % self.target_network_update_freq == 0:
-                with self.lock:
+            if self.total_steps.value % self.target_network_update_freq == 0:
+                with self.network_lock:
                     self.target_network.load_state_dict(self.learning_network.state_dict())
 
     def run(self):
-        procs = [threading.Thread(target=self.worker, args=(i, )) for i in range(self.n_workers)]
-        # procs = [mp.Process(target=self.worker, args=(i, )) for i in range(self.n_workers)]
+        procs = [mp.Process(target=self.worker, args=(i, )) for i in range(self.n_workers)]
         for p in procs: p.start()
-        # while True:
-        #     time.sleep(0.01)
+        task = self.task_fn()
+        while True:
+            if self.total_steps.value % self.test_interval == 0:
+                test_repeats = 5
+                rewards = np.zeros(test_repeats)
+                for i in range(test_repeats):
+                    rewards[i] = self.deterministic_episode(task)
+                print 'total stpes: %d, test process epsidoe reward: %f' %\
+                      (self.total_steps.value, np.mean(rewards))
+                if np.mean(rewards) > task.success_threshold:
+                    self.stop_signal.value = True
+                    break
         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()
-        self.lock = threading.Lock()
-
-    def fun(self):
-        # self.data += rank
-        for i in range(50):
-            time.sleep(0.01)
-            # with self.lock:
-            #     self.val.value += 1
-            self.data += 1
-
-    def run(self):
-        processes = []
-        for i in range(3):
-            # processes.append(mp.Process(target=self.fun))
-            processes.append(threading.Thread(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()
-    optimizer_fn = lambda params: torch.optim.SGD(params, 0.001)
-    network_fn = lambda: FullyConnectedNet([4, 50, 200, 2], optimizer_fn = optimizer_fn)
-    policy_fn = lambda: GreedyPolicy(epsilon=1.0, end_episode=500, min_epsilon=0.1)
-    # 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, 8)
-    agent.run()
-
-    # t = Test()
-    # t.run()
-    # print t.data
-
-    # 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

@@ -10,11 +10,11 @@ from policy import *
 import numpy as np
 
 class DQNAgent:
-    def __init__(self, task, 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):
         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.task = task_fn()
         self.step_limit = step_limit
         self.replay = replay_fn()
         self.discount = discount
@@ -48,6 +48,21 @@ class DQNAgent:
                 targets[np.arange(len(actions)), actions] = q_next
                 self.learning_network.learn(states, targets)
             if self.total_steps % self.target_network_update_freq == 0:
-                self.target_network.sync_with(self.learning_network)
+                self.target_network.load_state_dict(self.learning_network.state_dict())
         self.policy.update_epsilon()
         return total_reward
+
+    def run(self):
+        window_size = 100
+        ep = 0
+        rewards = []
+        while True:
+            ep += 1
+            reward = self.episode()
+            rewards.append(reward)
+            if len(rewards) > window_size:
+                reward = np.mean(rewards[-window_size:])
+            print 'episode %d: %f' % (ep, reward)
+            if reward > self.task.success_threshold:
+                break
+

main.py

@@ -0,0 +1,34 @@
+from async_agent import *
+from dqn_agent import *
+
+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['discount'] = 0.99
+    config['target_network_update_freq'] = 200
+    config['step_limit'] = 300
+    config['n_workers'] = 8
+    config['batch_size'] = 5
+    config['test_interval'] = 500
+    agent = AsyncAgent(**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['replay_fn'] = lambda: Replay(memory_size=10000, batch_size=10)
+    config['discount'] = 0.99
+    config['target_network_update_freq'] = 200
+    config['step_limit'] = 300
+    agent = DQNAgent(**config)
+    agent.run()
+
+if __name__ == '__main__':
+    async_cart_pole()
+    # dqn_cart_pole()

network.py

@@ -43,16 +43,6 @@ class FullyConnectedNet(nn.Module):
     def predict(self, x):
         return self.forward(x).cpu().data.numpy()
 
-    def learn_from_raw(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)
-        self.zero_grad()
-        loss.backward()
-        self.optimizer.step()
-
     def learn(self, x, target):
         target = torch.from_numpy(target)
         if self.gpu:
@@ -64,8 +54,10 @@ class FullyConnectedNet(nn.Module):
         loss.backward()
         self.optimizer.step()
 
-    def gradient(self, x, target):
+    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.backward()

task.py

@@ -39,35 +39,8 @@ class MountainCar(BasicTask):
         self.replay_fn = lambda: Replay(memory_size=10000, batch_size=10)
 
 class CartPole(BasicTask):
-    state_space_size = 4
-    action_space_size = 2
     name = 'CartPole-v0'
     success_threshold = 195
-    discount = 0.99
-    step_limit = 5000
-    target_network_update_freq = 200
 
     def __init__(self):
         self.env = gym.make(self.name)
-        self.optimizer_fn = lambda params: torch.optim.SGD(params, 0.001)
-        self.network_fn = lambda optimizer_fn: FullyConnectedNet([self.state_space_size, 50, 200, self.action_space_size], optimizer_fn)
-        self.policy_fn = lambda: GreedyPolicy(epsilon=0.5, decay_factor=0.99, min_epsilon=0.01)
-        self.replay_fn = lambda: Replay(memory_size=10000, batch_size=10)
-
-if __name__ == '__main__':
-    task = CartPole()
-    optimizer_fn = lambda params: torch.optim.SGD(params, 0.001)
-    agent = DQNAgent(task, task.network_fn, optimizer_fn, task.policy_fn, task.replay_fn,
-                     task.discount, task.step_limit, task.target_network_update_freq)
-    window_size = 100
-    ep = 0
-    rewards = []
-    while True:
-        ep += 1
-        reward = agent.episode()
-        rewards.append(reward)
-        if len(rewards) > window_size:
-            reward = np.mean(rewards[-window_size:])
-        print 'episode %d: %f' % (ep, reward)
-        if reward > task.success_threshold:
-            break