diff --git a/async_agent.py b/async_agent.py index f29bf4d..299b1d4 100644 --- a/async_agent.py +++ b/async_agent.py @@ -10,9 +10,11 @@ import numpy as np import torch.multiprocessing as mp from task import * from network import * -from bootstrap import * +from worker import * import pickle import os +import traceback +import time class AsyncAgent: def __init__(self, @@ -20,7 +22,7 @@ class AsyncAgent: network_fn, optimizer_fn, policy_fn, - bootstrap, + worker_fn, discount, step_limit, target_network_update_freq, @@ -33,13 +35,10 @@ class AsyncAgent: self.network_fn = network_fn self.learning_network = network_fn() self.learning_network.share_memory() - if bootstrap != AdvantageActorCritic: - self.target_network = network_fn() - self.target_network.share_memory() - self.target_network.load_state_dict(self.learning_network.state_dict()) - else: - self.target_network = None - self.bootstrap = bootstrap + self.target_network = network_fn() + self.target_network.share_memory() + self.target_network.load_state_dict(self.learning_network.state_dict()) + self.worker_fn = worker_fn self.optimizer_fn = optimizer_fn self.task_fn = task_fn @@ -66,86 +65,41 @@ class AsyncAgent: total_rewards = 0 steps = 0 network.reset(True) - bootstrap = self.bootstrap(self) while not self.step_limit or steps < self.step_limit: - action = np.argmax(bootstrap.process_state(network, state)) + action_value = network.predict(np.stack([state])) + if self.worker_fn == AdvantageActorCritic: + action_value = action_value[0] + action = np.argmax(action_value.data.numpy().flatten()) state, reward, terminal, _ = task.step(action) steps += 1 total_rewards += reward if terminal: break - bootstrap.reset() return total_rewards - def worker(self, id): - optimizer = self.optimizer_fn(self.learning_network.parameters()) - worker_network = self.network_fn() - worker_network.load_state_dict(self.learning_network.state_dict()) - - bootstrap = self.bootstrap(self) - task = self.task_fn() - policy = self.policy_fn() + def train(self, id): + worker = self.worker_fn(self) episode = 0 - episode_steps = 0 - episode_returns = [0] - state = task.reset() - pending_steps = 0 + rewards = [] while True and not self.stop_signal.value: - action = policy.sample(bootstrap.process_state(worker_network, state)) - next_state, reward, terminal, _ = task.step(action) - bootstrap.process_interaction(action, reward, next_state) - - episode_returns[-1] += reward - episode_steps += 1 - if self.step_limit and episode_steps > self.step_limit: - terminal = True - with self.steps_lock: - self.total_steps.value += 1 - pending_steps += 1 - - if terminal or pending_steps >= self.update_interval: - loss = bootstrap.compute_loss(worker_network, terminal) - pending_steps = 0 - worker_network.zero_grad() - loss.backward() - nn.utils.clip_grad_norm(worker_network.parameters(), 40) - optimizer.zero_grad() - for param, worker_param in zip(self.learning_network.parameters(), worker_network.parameters()): - param._grad = worker_param.grad.clone().cpu() - optimizer.step() - worker_network.load_state_dict(self.learning_network.state_dict()) - worker_network.reset(terminal) - - if terminal: - state = task.reset() - episode += 1 - if id == 0: - self.logger.info('episode %d, return %f, avg return %f, episode steps %d, total steps %d' % ( - episode, episode_returns[-1], np.mean(episode_returns[-100:]), episode_steps, self.total_steps.value)) - episode_returns.append(0) - episode_steps = 0 - else: - state = next_state - - if self.target_network and self.total_steps.value % self.target_network_update_freq == 0: - self.target_network.load_state_dict(self.learning_network.state_dict()) + steps, reward = worker.episode() + rewards.append(reward) + self.logger.debug('worker %d, episode %d, return %f, avg return %f, episode steps %d, total steps %d' % ( + id, episode, rewards[-1], np.mean(rewards[-100:]), steps, self.total_steps.value)) def save(self, file_name): with open(file_name, 'wb') as f: pickle.dump(self.learning_network.state_dict(), f) - def run(self): - os.environ['OMP_NUM_THREADS'] = '1' - procs = [mp.Process(target=self.worker, args=(i, )) for i in range(self.n_workers)] - for p in procs: p.start() + def evaluate(self, id): test_rewards = [] test_points = [] test_network = self.network_fn() while True: - steps = self.total_steps.value + 1 + steps = self.total_steps.value if steps % self.test_interval == 0: test_network.load_state_dict(self.learning_network.state_dict()) - self.save('data/%s%s-model-%s.bin' % (self.tag, self.bootstrap.__name__, self.task.name)) + self.save('data/%s%s-model-%s.bin' % (self.tag, self.worker_fn.__name__, self.task.name)) rewards = np.zeros(self.test_repetitions) for i in range(self.test_repetitions): rewards[i] = self.deterministic_episode(self.task, test_network) @@ -154,10 +108,28 @@ class AsyncAgent: test_rewards.append(np.mean(rewards)) test_points.append(steps) with open('data/%s%s-statistics-%s.bin' % ( - self.tag, self.bootstrap.__name__, self.task.name + self.tag, self.worker_fn.__name__, self.task.name ), 'wb') as f: pickle.dump([test_points, test_rewards], f) if np.mean(rewards) > self.task.success_threshold: self.stop_signal.value = True break + + def run(self): + os.environ['OMP_NUM_THREADS'] = '1' + procs = [mp.Process(target=self.train, args=(i, )) for i in range(self.n_workers)] + procs.append(mp.Process(target=self.evaluate, args=(self.n_workers, ))) + for p in procs: p.start() + while True: + time.sleep(1) + for i, p in enumerate(procs): + if not p.is_alive() and not self.stop_signal.value: + self.logger.warning('Worker %d exited unexpectedly.' % i) + p.terminate() + procs[i] = mp.Process(target=self.train, args=(i, )) + procs[i].start() + self.logger.warning('Worker %d restarted.' % i) + break + if self.stop_signal.value: + break for p in procs: p.join() diff --git a/bootstrap.py b/bootstrap.py deleted file mode 100644 index ab99240..0000000 --- a/bootstrap.py +++ /dev/null @@ -1,139 +0,0 @@ -####################################################################### -# 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 # -####################################################################### -import numpy as np -import torch -from torch.autograd import Variable - -class OneStepSarsa: - def __init__(self, agent): - self.agent = agent - self.reset() - - def reset(self): - self.pending = [] - - def process_state(self, network, state): - q = network.predict(np.stack([state])) - self.pending.append([q]) - return q.data.numpy().flatten() - - def process_interaction(self, action, reward, next_state): - self.pending[-1].extend([action, reward, next_state]) - - def compute_loss(self, network, terminal): - loss = 0 - valid_length = len(self.pending) - if not terminal: - valid_length -= 1 - for i in range(valid_length): - q, action, reward, next_state = self.pending[i] - q_next = self.agent.target_network.predict(np.stack([next_state])).data - if i < len(self.pending) - 1: - next_action = self.pending[i + 1][1] - q_next = q_next.gather(1, torch.LongTensor([[next_action]])) - else: - q_next = torch.FloatTensor([[0]]) - q_next = self.agent.discount * q_next + reward - q = q.gather(1, Variable(torch.LongTensor([[action]]))) - loss += 0.5 * (q - Variable(q_next)).pow(2) - self.reset() - return loss - - -class OneStepQLearning: - def __init__(self, agent): - self.agent = agent - self.reset() - - def reset(self): - self.pending = [] - - def process_state(self, network, state): - q = network.predict(np.stack([state])) - self.pending.append([q]) - return q.data.numpy().flatten() - - def process_interaction(self, action, reward, next_state): - self.pending[-1].extend([action, reward, next_state]) - - def compute_loss(self, network, terminal): - loss = 0 - for i in range(len(self.pending)): - q, action, reward, next_state = self.pending[i] - q_next, _ = self.agent.target_network.predict(np.stack([next_state])).data.max(1) - if terminal and i == len(self.pending) - 1: - q_next = torch.FloatTensor([[0]]) - q_next = self.agent.discount * q_next + reward - q = q.gather(1, Variable(torch.LongTensor([[action]]))) - loss += 0.5 * (q - Variable(q_next)).pow(2) - self.reset() - return loss - -class NStepQLearning: - def __init__(self, agent): - self.agent = agent - self.reset() - - def reset(self): - self.pending = [] - - def process_state(self, network, state): - q = network.predict(np.stack([state])) - self.pending.append([q]) - return q.data.numpy().flatten() - - def process_interaction(self, action, reward, next_state): - self.pending[-1].extend([action, reward]) - self.tailing_state = next_state - - def compute_loss(self, network, terminal): - loss = 0 - if terminal: - R = torch.FloatTensor([[0]]) - else: - R, _ = self.agent.target_network.predict( - np.stack([self.tailing_state])).data.max(1) - - for i in reversed(range(len(self.pending))): - q, action, reward = self.pending[i] - R = reward + self.agent.discount * R - loss += 0.5 * (Variable(R) - q.gather(1, Variable(torch.LongTensor([[action]])))).pow(2) - self.reset() - return loss - -class AdvantageActorCritic: - def __init__(self, agent): - self.agent = agent - self.reset() - - def reset(self): - self.pending = [] - - def process_state(self, network, state): - prob, log_prob, value = network.predict(np.stack([state])) - self.pending.append([prob, log_prob, value]) - return prob.data.numpy().flatten() - - def process_interaction(self, action, reward, next_state): - self.pending[-1].extend([action, reward]) - self.tailing_state = next_state - - def compute_loss(self, network, terminal): - loss = 0 - if terminal: - R = torch.FloatTensor([[0]]) - else: - R = network.critic(np.stack([self.tailing_state])).data - for i in reversed(range(len(self.pending))): - prob, log_prob, value, action, reward = self.pending[i] - R = reward + self.agent.discount * R - advantage = Variable(R) - value - loss += 0.5 * advantage.pow(2) - loss += -log_prob.gather(1, Variable(torch.LongTensor([[action]]))) * Variable(advantage.data) - loss += 0.01 * torch.sum(torch.mul(prob, log_prob)) - self.reset() - return loss - diff --git a/main.py b/main.py index e62d4ec..7b07d27 100644 --- a/main.py +++ b/main.py @@ -1,6 +1,7 @@ from async_agent import * from dqn_agent import * import logging +import traceback def dqn_cart_pole(): config = dict() @@ -29,9 +30,9 @@ def async_cart_pole(): config['optimizer_fn'] = lambda params: torch.optim.Adam(params, 0.001) config['network_fn'] = lambda: FCNet([4, 50, 200, 2]) config['policy_fn'] = lambda: GreedyPolicy(epsilon=0.5, final_step=5000, min_epsilon=0.1) - config['bootstrap'] = OneStepQLearning - # config['bootstrap'] = NStepQLearning - # config['bootstrap'] = OneStepSarsa + config['worker_fn'] = OneStepQLearning + # config['worker_fn'] = NStepQLearning + # config['worker_fn'] = OneStepSarsa config['discount'] = 0.99 config['target_network_update_freq'] = 200 config['step_limit'] = 0 @@ -51,7 +52,7 @@ def a3c_cart_pole(): config['optimizer_fn'] = lambda params: torch.optim.Adam(params, 0.001) config['network_fn'] = lambda: ActorCriticFCNet([4, 200, 2]) config['policy_fn'] = SamplePolicy - config['bootstrap'] = AdvantageActorCritic + config['worker_fn'] = AdvantageActorCritic config['discount'] = 0.99 config['target_network_update_freq'] = 200 config['step_limit'] = 0 @@ -96,13 +97,13 @@ def async_pixel_atari(name): config['optimizer_fn'] = lambda params: torch.optim.Adam(params, lr=0.0001) config['network_fn'] = lambda: OpenAIConvNet(history_length, n_actions) - config['policy_fn'] = lambda: StochasticGreedyPolicy(epsilons=[0.5, 0.5, 0.5], + config['policy_fn'] = lambda: StochasticGreedyPolicy(epsilons=[0.7, 0.7, 0.7], final_step=2000000, - min_epsilons=[0.1, 0.01, 0.2], + min_epsilons=[0.1, 0.01, 0.5], probs=[0.4, 0.3, 0.3]) - # config['bootstrap'] = OneStepQLearning - # config['bootstrap'] = NStepQLearning - config['bootstrap'] = OneStepSarsa + # config['worker_fn'] = OneStepQLearning + # config['worker_fn'] = NStepQLearning + config['worker_fn'] = OneStepSarsa config['discount'] = 0.99 config['target_network_update_freq'] = 10000 config['step_limit'] = 10000 @@ -113,6 +114,7 @@ def async_pixel_atari(name): config['history_length'] = history_length config['logger'] = gym.logger agent = AsyncAgent(**config) + agent.tag = 'Centered-target-network-' agent.run() def a3c_pixel_atari(name): @@ -125,7 +127,7 @@ def a3c_pixel_atari(name): n_actions, LSTM=False) config['policy_fn'] = SamplePolicy - config['bootstrap'] = AdvantageActorCritic + config['worker_fn'] = AdvantageActorCritic config['discount'] = 0.99 config['target_network_update_freq'] = 0 config['step_limit'] = 10000 @@ -140,8 +142,8 @@ def a3c_pixel_atari(name): 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) # dqn_cart_pole() # async_cart_pole() diff --git a/worker.py b/worker.py new file mode 100644 index 0000000..d5271a4 --- /dev/null +++ b/worker.py @@ -0,0 +1,251 @@ +####################################################################### +# 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 # +####################################################################### +import numpy as np +import torch +from torch.autograd import Variable +import torch.nn as nn + +class AdvantageActorCritic: + def __init__(self, agent): + self.agent = agent + self.optimizer = agent.optimizer_fn(agent.learning_network.parameters()) + self.worker_network = agent.network_fn() + self.worker_network.load_state_dict(agent.learning_network.state_dict()) + self.task = agent.task_fn() + self.policy = agent.policy_fn() + + def episode(self): + state = self.task.reset() + steps = 0 + total_reward = 0 + pending = [] + while True and not self.agent.stop_signal.value: + prob, log_prob, value = self.worker_network.predict(np.stack([state])) + action = self.policy.sample(prob.data.numpy().flatten()) + next_state, reward, terminal, _ = self.task.step(action) + pending.append([prob, log_prob, value, action, reward]) + + steps += 1 + with self.agent.steps_lock: + self.agent.total_steps.value += 1 + total_reward += reward + + if terminal or len(pending) >= self.agent.update_interval: + loss = 0 + if terminal: + R = torch.FloatTensor([[0]]) + else: + R = self.worker_network.critic(np.stack([next_state])).data + for i in reversed(range(len(pending))): + prob, log_prob, value, action, reward = pending[i] + R = reward + self.agent.discount * R + advantage = Variable(R) - value + loss += 0.5 * advantage.pow(2) + loss += -log_prob.gather(1, Variable(torch.LongTensor([[action]]))) * Variable(advantage.data) + loss += 0.01 * torch.sum(torch.mul(prob, log_prob)) + + pending = [] + self.worker_network.zero_grad() + loss.backward() + nn.utils.clip_grad_norm(self.worker_network.parameters(), 40) + self.optimizer.zero_grad() + for param, worker_param in zip( + self.agent.learning_network.parameters(), self.worker_network.parameters()): + param._grad = worker_param.grad.clone() + self.optimizer.step() + self.worker_network.load_state_dict(self.agent.learning_network.state_dict()) + self.worker_network.reset(terminal) + + if terminal: + break + else: + state = next_state + + return steps, total_reward + +class NStepQLearning: + def __init__(self, agent): + self.agent = agent + self.optimizer = agent.optimizer_fn(agent.learning_network.parameters()) + self.worker_network = agent.network_fn() + self.worker_network.load_state_dict(agent.learning_network.state_dict()) + self.task = agent.task_fn() + self.policy = agent.policy_fn() + + def episode(self): + state = self.task.reset() + steps = 0 + total_reward = 0 + pending = [] + while True and not self.agent.stop_signal.value: + q = self.worker_network.predict(np.stack([state])) + action = self.policy.sample(q.data.numpy().flatten()) + next_state, reward, terminal, _ = self.task.step(action) + pending.append([q, action, reward]) + + steps += 1 + with self.agent.steps_lock: + self.agent.total_steps.value += 1 + total_reward += reward + + if terminal or len(pending) >= self.agent.update_interval: + loss = 0 + if terminal: + R = torch.FloatTensor([[0]]) + else: + R, _ = self.agent.target_network.predict( + np.stack([next_state])).data.max(1) + + for i in reversed(range(len(pending))): + q, action, reward = pending[i] + R = reward + self.agent.discount * R + loss += 0.5 * (Variable(R) - q.gather(1, Variable(torch.LongTensor([[action]])))).pow(2) + + pending = [] + self.worker_network.zero_grad() + loss.backward() + nn.utils.clip_grad_norm(self.worker_network.parameters(), 40) + self.optimizer.zero_grad() + for param, worker_param in zip( + self.agent.learning_network.parameters(), self.worker_network.parameters()): + param._grad = worker_param.grad.clone() + self.optimizer.step() + self.worker_network.load_state_dict(self.agent.learning_network.state_dict()) + self.worker_network.reset(terminal) + + if terminal: + break + else: + state = next_state + + if self.agent.total_steps.value % self.agent.target_network_update_freq == 0: + self.agent.target_network.load_state_dict( + self.agent.learning_network.state_dict()) + + return steps, total_reward + +class OneStepQLearning: + def __init__(self, agent): + self.agent = agent + self.optimizer = agent.optimizer_fn(agent.learning_network.parameters()) + self.worker_network = agent.network_fn() + self.worker_network.load_state_dict(agent.learning_network.state_dict()) + self.task = agent.task_fn() + self.policy = agent.policy_fn() + + def episode(self): + state = self.task.reset() + steps = 0 + total_reward = 0 + pending = [] + while True and not self.agent.stop_signal.value: + q = self.worker_network.predict(np.stack([state])) + action = self.policy.sample(q.data.numpy().flatten()) + next_state, reward, terminal, _ = self.task.step(action) + pending.append([q, action, reward, next_state]) + + steps += 1 + with self.agent.steps_lock: + self.agent.total_steps.value += 1 + total_reward += reward + + if terminal or len(pending) >= self.agent.update_interval: + loss = 0 + for i in range(len(pending)): + q, action, reward, next_state = pending[i] + q_next, _ = self.agent.target_network.predict(np.stack([next_state])).data.max(1) + if terminal and i == len(pending) - 1: + q_next = torch.FloatTensor([[0]]) + q_next = self.agent.discount * q_next + reward + q = q.gather(1, Variable(torch.LongTensor([[action]]))) + loss += 0.5 * (q - Variable(q_next)).pow(2) + + pending = [] + self.worker_network.zero_grad() + loss.backward() + nn.utils.clip_grad_norm(self.worker_network.parameters(), 40) + self.optimizer.zero_grad() + for param, worker_param in zip( + self.agent.learning_network.parameters(), self.worker_network.parameters()): + param._grad = worker_param.grad.clone() + self.optimizer.step() + self.worker_network.load_state_dict(self.agent.learning_network.state_dict()) + self.worker_network.reset(terminal) + + if terminal: + break + else: + state = next_state + + if self.agent.total_steps.value % self.agent.target_network_update_freq == 0: + self.agent.target_network.load_state_dict( + self.agent.learning_network.state_dict()) + + return steps, total_reward + +class OneStepSarsa: + def __init__(self, agent): + self.agent = agent + self.optimizer = agent.optimizer_fn(agent.learning_network.parameters()) + self.worker_network = agent.network_fn() + self.worker_network.load_state_dict(agent.learning_network.state_dict()) + self.task = agent.task_fn() + self.policy = agent.policy_fn() + + def episode(self): + state = self.task.reset() + q = self.worker_network.predict(np.stack([state])) + action = self.policy.sample(q.data.numpy().flatten()) + steps = 0 + total_reward = 0 + pending = [] + while True and not self.agent.stop_signal.value: + next_state, reward, terminal, _ = self.task.step(action) + next_q = self.worker_network.predict(np.stack([next_state])) + next_action = self.policy.sample(next_q.data.numpy().flatten()) + pending.append([q, action, reward, next_state, next_action]) + + steps += 1 + with self.agent.steps_lock: + self.agent.total_steps.value += 1 + total_reward += reward + + if terminal or len(pending) >= self.agent.update_interval: + loss = 0 + for i in range(len(pending)): + q, action, reward, next_state, next_action = pending[i] + q_next = self.agent.target_network.predict(np.stack([next_state])).data + if terminal and i == len(pending) - 1: + q_next = torch.FloatTensor([[0]]) + else: + q_next = q_next.gather(1, torch.LongTensor([[next_action]])) + q_next = self.agent.discount * q_next + reward + q = q.gather(1, Variable(torch.LongTensor([[action]]))) + loss += 0.5 * (q - Variable(q_next)).pow(2) + + pending = [] + self.worker_network.zero_grad() + loss.backward() + nn.utils.clip_grad_norm(self.worker_network.parameters(), 40) + self.optimizer.zero_grad() + for param, worker_param in zip( + self.agent.learning_network.parameters(), self.worker_network.parameters()): + param._grad = worker_param.grad.clone() + self.optimizer.step() + self.worker_network.load_state_dict(self.agent.learning_network.state_dict()) + self.worker_network.reset(terminal) + + if terminal: + break + else: + q = next_q + action = next_action + + if self.agent.total_steps.value % self.agent.target_network_update_freq == 0: + self.agent.target_network.load_state_dict( + self.agent.learning_network.state_dict()) + + return steps, total_reward