From 3659cbed3a2170a95e09f165c473f991ee44ef42 Mon Sep 17 00:00:00 2001 From: Shangtong Zhang Date: Wed, 31 May 2017 15:43:17 -0600 Subject: [PATCH] Fix bug for async agents --- async_agent.py | 41 +++++++++++++++++++++---------- bootstrap.py | 29 ++++++++++------------ main.py | 30 +++++++++++++++++++---- network.py | 65 ++++++++++++++++++++++++++++++++++---------------- 4 files changed, 112 insertions(+), 53 deletions(-) diff --git a/async_agent.py b/async_agent.py index be0fd7a..fef1eb4 100644 --- a/async_agent.py +++ b/async_agent.py @@ -11,6 +11,7 @@ import torch.multiprocessing as mp from task import * from network import * from bootstrap import * +import pickle class AsyncAgent: def __init__(self, @@ -29,11 +30,12 @@ class AsyncAgent: history_length, logger): self.network_fn = network_fn - self.learning_network = network_fn() + self.learning_network = network_fn(False) 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()) + if bootstrap_fn != AdvantageActorCritic: + self.target_network = network_fn(False) + self.target_network.share_memory() + self.target_network.load_state_dict(self.learning_network.state_dict()) self.bootstrap_fn = bootstrap_fn self.optimizer_fn = optimizer_fn @@ -55,14 +57,14 @@ class AsyncAgent: self.history_length = history_length def deterministic_episode(self, task, network): - state = np.asarray([task.reset()]) + state = task.reset() total_rewards = 0 steps = 0 terminal = False buffer = [state] * self.history_length while not terminal and (not self.step_limit or steps < self.step_limit): state = task.normalize_state(np.vstack(buffer)) - action_values = network.predict(np.reshape(state, (1, ) + state.shape)) + action_values = network.predict(np.stack([state])) steps += 1 action = np.argmax(action_values.flatten()) state, reward, terminal, _ = task.step(action) @@ -77,7 +79,7 @@ class AsyncAgent: 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.clone() + param._grad = worker_param.grad.clone().cpu() optimizer.step() def worker(self, id): @@ -125,7 +127,7 @@ class AsyncAgent: policy.update_epsilon() batch_rewards = self.bootstrap_fn(batch_states, batch_actions, batch_rewards, - state, action, terminal, self) + state, action, terminal, worker_network, self.discount) if self.step_limit and episode_steps > self.step_limit: terminal = True @@ -137,25 +139,40 @@ class AsyncAgent: self.async_update(worker_network, optimizer) worker_network.load_state_dict(self.learning_network.state_dict()) - if self.total_steps.value % self.target_network_update_freq == 0: + if self.target_network_update_freq and \ + 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 save(self, file_name): + with open(file_name, 'wb') as f: + pickle.dump(self.learning_network.state_dict(), f) + def run(self): procs = [mp.Process(target=self.worker, args=(i, )) for i in range(self.n_workers)] for p in procs: p.start() task = self.task_fn() test_network = self.network_fn() + test_rewards = [0] + test_points = [0] while True: steps = self.total_steps.value + 1 - if steps % self.test_interval == 0: + if steps >= test_points[-1] + self.test_interval: + test_points.append(steps) + self.logger.info('Testing...') with self.network_lock: test_network.load_state_dict(self.learning_network.state_dict()) + self.save('data/%s-model-%s.bin' % (self.bootstrap_fn.__name__, task.name)) rewards = np.zeros(self.test_repetitions) for i in range(self.test_repetitions): rewards[i] = self.deterministic_episode(task, test_network) - self.logger.info('total steps: %d, averaged return per episode: %f' %\ - (steps, np.mean(rewards))) + self.logger.info('total steps: %d, averaged return per episode: %f(%f)' %\ + (steps, np.mean(rewards), np.std(rewards) / np.sqrt(self.test_repetitions))) + test_rewards.append(np.mean(rewards)) + with open('data/%s-statistics-%s.bin' % ( + self.bootstrap_fn.__name__, task.name + ), 'wb') as f: + pickle.dump([test_points, test_rewards], f) if np.mean(rewards) > task.success_threshold: self.stop_signal.value = True break diff --git a/bootstrap.py b/bootstrap.py index 0f52a2e..95964c4 100644 --- a/bootstrap.py +++ b/bootstrap.py @@ -6,54 +6,49 @@ import numpy as np def NStepQLearning(batch_states, batch_actions, batch_rewards, - tailing_state, tailing_action, terminal, agent): + tailing_state, tailing_action, terminal, network, discount): if terminal: reward = 0 else: - with agent.network_lock: - reward = np.max(agent.target_network.predict( - np.reshape(tailing_state, (1, ) + tailing_state.shape))) + reward = np.max(network.predict(np.stack([tailing_state])).flatten()) rewards = [] for r in reversed(batch_rewards): - reward = r + agent.discount * reward + reward = r + discount * reward rewards.append(reward) return rewards def OneStepQLearning(batch_states, batch_actions, batch_rewards, - tailing_state, tailing_action, terminal, agent): + tailing_state, tailing_action, terminal, network, discount): batch_states.append(tailing_state) - with agent.network_lock: - q_next = agent.target_network.predict(np.asarray(batch_states[1:])) + q_next = network.predict(np.asarray(batch_states[1:])) q_next = np.max(q_next, axis=1) if terminal: q_next[-1] = 0 batch_states.pop(-1) - batch_rewards = np.asarray(batch_rewards) + agent.discount * q_next + batch_rewards = np.asarray(batch_rewards) + discount * q_next return batch_rewards def OneStepSarsa(batch_states, batch_actions, batch_rewards, - tailing_state, tailing_action, terminal, agent): + tailing_state, tailing_action, terminal, network, discount): batch_states.append(tailing_state) batch_actions.append(tailing_action) - with agent.network_lock: - q_next = agent.target_network.predict(np.asarray(batch_states[1:])) + q_next = network.predict(np.asarray(batch_states[1:])) q_next = q_next[np.arange(len(batch_actions[1:])), batch_actions[1:]] if terminal: q_next[-1] = 0 batch_states.pop(-1) batch_actions.pop(-1) - batch_rewards = np.asarray(batch_rewards) + agent.discount * q_next + batch_rewards = np.asarray(batch_rewards) + discount * q_next return batch_rewards def AdvantageActorCritic(batch_states, batch_actions, batch_rewards, - tailing_state, tailing_action, terminal, agent): + tailing_state, tailing_action, terminal, network, discount): if terminal: reward = 0 else: - with agent.network_lock: - reward = np.asscalar(agent.learning_network.critic(np.stack([tailing_state]))) + reward = np.asscalar(network.critic(np.stack([tailing_state]))) rewards = [] for r in reversed(batch_rewards): - reward = r + agent.discount * reward + reward = r + discount * reward rewards.append(reward) return rewards diff --git a/main.py b/main.py index 49c4cdf..eb9b7a8 100644 --- a/main.py +++ b/main.py @@ -58,7 +58,7 @@ def dqn_cart_pole(): agent = DQNAgent(**config) agent.run() -def actor_critic_cart_pole(): +def a3c_cart_pole(): config = dict() config['task_fn'] = lambda: CartPole() config['optimizer_fn'] = lambda params: torch.optim.RMSprop(params, 0.001) @@ -66,7 +66,7 @@ def actor_critic_cart_pole(): config['policy_fn'] = SamplePolicy config['bootstrap_fn'] = AdvantageActorCritic config['discount'] = 0.99 - config['target_network_update_freq'] = 200 + config['target_network_update_freq'] = 0 config['step_limit'] = 0 config['n_workers'] = 16 config['batch_size'] = 6 @@ -120,14 +120,36 @@ def async_pixel_atari(name): agent = AsyncAgent(**config) agent.run() +def a3c_pixel_atari(name): + config = dict() + history_length = 4 + n_actions = 6 + config['task_fn'] = lambda: PixelAtari(name, 30, 4) + config['optimizer_fn'] = lambda params: torch.optim.RMSprop(params, lr=0.0001, alpha=0.99, eps=0.01) + config['network_fn'] = lambda gpu=True: ConvActorCriticNet(history_length, n_actions, gpu=gpu) + config['policy_fn'] = SamplePolicy + config['bootstrap_fn'] = AdvantageActorCritic + config['discount'] = 0.99 + config['target_network_update_freq'] = 0 + config['step_limit'] = 0 + config['n_workers'] = 16 + config['batch_size'] = 10 + config['test_interval'] = 10000 + config['test_repetitions'] = 20 + config['history_length'] = 4 + config['logger'] = gym.logger + agent = AsyncAgent(**config) + agent.run() + if __name__ == '__main__': # gym.logger.setLevel(logging.DEBUG) gym.logger.setLevel(logging.INFO) benchmark = gym.benchmark_spec('Atari40M') # async_cart_pole() - # actor_critic_cart_pole() + # a3c_cart_pole() # async_lunar_lander() - dqn_cart_pole() + # dqn_cart_pole() # dqn_pixel_atari('BreakoutNoFrameskip-v3') # async_pixel_atari('BreakoutNoFrameskip-v3') + a3c_pixel_atari('BreakoutNoFrameskip-v3') diff --git a/network.py b/network.py index 236fa06..50103cf 100644 --- a/network.py +++ b/network.py @@ -80,26 +80,6 @@ class FullyConnectedNet(nn.Module, VanillaNet): y = self.fc3(y) return y -class FCActorCriticNet(nn.Module, ActorCriticNet): - def __init__(self, - dims, - xentropy_weight=0.01, - grad_threshold=40, - gpu=True): - super(FCActorCriticNet, self).__init__() - self.fc1 = nn.Linear(dims[0], dims[1]) - self.fc_actor = nn.Linear(dims[1], dims[2]) - self.fc_critic = nn.Linear(dims[1], 1) - self.xentropy_weight = xentropy_weight - self.grad_threshold = grad_threshold - BasicNet.__init__(self, optimizer_fn=None, gpu=gpu) - - def forward(self, x): - x = self.to_torch_variable(x) - x = x.view(x.size(0), -1) - phi = self.fc1(x) - return phi - class ConvNet(nn.Module, VanillaNet): def __init__(self, in_channels, n_actions, optimizer_fn=None, gpu=True): super(ConvNet, self).__init__() @@ -120,4 +100,49 @@ class ConvNet(nn.Module, VanillaNet): y = F.relu(self.fc4(y)) return self.fc5(y) +class FCActorCriticNet(nn.Module, ActorCriticNet): + def __init__(self, + dims, + xentropy_weight=0.01, + grad_threshold=40, + gpu=True): + super(FCActorCriticNet, self).__init__() + self.fc1 = nn.Linear(dims[0], dims[1]) + self.fc_actor = nn.Linear(dims[1], dims[2]) + self.fc_critic = nn.Linear(dims[1], 1) + self.xentropy_weight = xentropy_weight + self.grad_threshold = grad_threshold + BasicNet.__init__(self, optimizer_fn=None, gpu=gpu) + + def forward(self, x): + x = self.to_torch_variable(x) + x = x.view(x.size(0), -1) + phi = self.fc1(x) + return phi + +class ConvActorCriticNet(nn.Module, ActorCriticNet): + def __init__(self, + in_channels, + n_actions, + xentropy_weight=0.01, + grad_threshold=40, + gpu=True): + super(ConvActorCriticNet, self).__init__() + self.conv1 = nn.Conv2d(in_channels, 32, kernel_size=8, stride=4) + self.conv2 = nn.Conv2d(32, 64, kernel_size=4, stride=2) + self.conv3 = nn.Conv2d(64, 64, kernel_size=3, stride=1) + self.fc4 = nn.Linear(7 * 7 * 64, 512) + self.fc_actor = nn.Linear(512, n_actions) + self.fc_critic = nn.Linear(512, 1) + self.xentropy_weight = xentropy_weight + self.grad_threshold = grad_threshold + BasicNet.__init__(self, optimizer_fn=None, gpu=gpu) + + def forward(self, x): + x = self.to_torch_variable(x) + y = F.relu(self.conv1(x)) + y = F.relu(self.conv2(y)) + y = F.relu(self.conv3(y)) + y = y.view(y.size(0), -1) + return F.relu(self.fc4(y))