From dde67bece1d4cf51f51cf79a04992710a32fc922 Mon Sep 17 00:00:00 2001 From: Shangtong Zhang Date: Wed, 10 May 2017 14:12:37 -0600 Subject: [PATCH] MP --- async_agent.py | 161 +++++++++++++++++++++++++++++++++++++++++++++++++ dqn_agent.py | 5 +- network.py | 11 +++- 3 files changed, 173 insertions(+), 4 deletions(-) create mode 100644 async_agent.py diff --git a/async_agent.py b/async_agent.py new file mode 100644 index 0000000..bb6ac5e --- /dev/null +++ b/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 diff --git a/dqn_agent.py b/dqn_agent.py index c2436e8..bfb8fe3 100644 --- a/dqn_agent.py +++ b/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: diff --git a/network.py b/network.py index 0480771..ee656b7 100644 --- a/network.py +++ b/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