From 9c73abfbff24a360dcec1d503e980bd50c0bde92 Mon Sep 17 00:00:00 2001 From: Shangtong Zhang Date: Thu, 8 Jun 2017 18:36:04 -0600 Subject: [PATCH] Major refactor to support LSTM layer --- async_agent.py | 127 +++++++++++++++--------------------- atari_wrapper.py | 28 ++++++-- bootstrap.py | 161 ++++++++++++++++++++++++++++++++------------- dqn_agent.py | 8 +-- main.py | 65 ++++++++++--------- network.py | 165 ++++++++++++++++++++++++++++++++++++----------- policy.py | 12 ++-- task.py | 8 +-- 8 files changed, 370 insertions(+), 204 deletions(-) diff --git a/async_agent.py b/async_agent.py index 09b7a34..6e4d346 100644 --- a/async_agent.py +++ b/async_agent.py @@ -20,12 +20,12 @@ class AsyncAgent: network_fn, optimizer_fn, policy_fn, - bootstrap_fn, + bootstrap, discount, step_limit, target_network_update_freq, n_workers, - batch_size, + update_interval, test_interval, test_repetitions, history_length, @@ -33,16 +33,17 @@ class AsyncAgent: self.network_fn = network_fn self.learning_network = network_fn() self.learning_network.share_memory() - if bootstrap_fn != AdvantageActorCritic: + 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_fn = bootstrap_fn + self.bootstrap = bootstrap self.optimizer_fn = optimizer_fn self.task_fn = task_fn + self.task = self.task_fn() self.step_limit = step_limit self.discount = discount self.optimizer_fn = optimizer_fn @@ -53,7 +54,7 @@ class AsyncAgent: self.total_steps = mp.Value('i', 0) self.stop_signal = mp.Value('i', False) self.n_workers = n_workers - self.batch_size = batch_size + self.update_interval = update_interval self.test_interval = test_interval self.test_repetitions = test_repetitions self.logger = logger @@ -63,92 +64,69 @@ class AsyncAgent: 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 = np.vstack(buffer) - action_values = network.predict(np.stack([state])) - steps += 1 - action = np.argmax(action_values.flatten()) + 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)) state, reward, terminal, _ = task.step(action) - buffer.pop(0) - buffer.append(state) + steps += 1 total_rewards += reward if terminal: break return total_rewards - def async_update(self, worker_network, optimizer): - 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() - 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() - terminal = True episode = 0 episode_steps = 0 - episode_return = 0 - episode_returns = [] - update_target_network = False + episode_returns = [0] + state = task.reset() + pending_steps = 0 while True and not self.stop_signal.value: - batch_states, batch_actions, batch_rewards = [], [], [] - if terminal: - if episode and id == 0: - episode_returns.append(episode_return) - self.logger.info('episode %d, return %f, avg return %f, episode steps %d, total steps %d' % ( - episode, episode_return, np.mean(episode_returns[-100: ]), episode_steps, - self.total_steps.value)) - episode_steps = 0 - episode_return = 0 - episode += 1 - terminal = False - state = task.reset() - buffer = [state] * self.history_length - state = np.vstack(buffer) - value = worker_network.predict(np.stack([state])) - action = policy.sample(value.flatten()) - while not terminal and len(batch_states) < self.batch_size: - episode_steps += 1 - with self.steps_lock: - self.total_steps.value += 1 - self.total_steps.value += 1 - if self.total_steps.value % self.target_network_update_freq == 0: - update_target_network = True - batch_states.append(state) - batch_actions.append(action) - state, reward, terminal, _ = task.step(action) - batch_rewards.append(reward) - episode_return += reward - buffer.pop(0) - buffer.append(state) - state = np.vstack(buffer) - value = worker_network.predict(np.stack([state])) - action = policy.sample(value.flatten()) - policy.update_epsilon() - - batch_rewards = self.bootstrap_fn(batch_states, batch_actions, batch_rewards, - state, action, terminal, worker_network, self.discount) + 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 - worker_network.zero_grad() - worker_network.gradient(np.asarray(batch_states), - worker_network.to_torch_variable(batch_actions, 'int64').unsqueeze(1), - worker_network.to_torch_variable(batch_rewards).unsqueeze(1)) - self.async_update(worker_network, optimizer) - worker_network.load_state_dict(self.learning_network.state_dict()) + 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 self.target_network is not None and update_target_network: + 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: with self.network_lock: self.target_network.load_state_dict(self.learning_network.state_dict()) - update_target_network = False def save(self, file_name): with open(file_name, 'wb') as f: @@ -158,28 +136,27 @@ class AsyncAgent: 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() - task = self.task_fn() - test_network = self.network_fn() test_rewards = [] test_points = [] + test_network = self.network_fn() while True: steps = self.total_steps.value + 1 if steps % self.test_interval == 0: 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)) + self.save('data/%s-model-%s.bin' % (self.bootstrap.__name__, self.task.name)) rewards = np.zeros(self.test_repetitions) for i in range(self.test_repetitions): - rewards[i] = self.deterministic_episode(task, test_network) + rewards[i] = self.deterministic_episode(self.task, test_network) 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)) test_points.append(steps) with open('data/%s-statistics-%s.bin' % ( - self.bootstrap_fn.__name__, task.name + self.bootstrap.__name__, self.task.name ), 'wb') as f: pickle.dump([test_points, test_rewards], f) - if np.mean(rewards) > task.success_threshold: + if np.mean(rewards) > self.task.success_threshold: self.stop_signal.value = True break for p in procs: p.join() diff --git a/atari_wrapper.py b/atari_wrapper.py index a95b43b..7f9ef10 100644 --- a/atari_wrapper.py +++ b/atari_wrapper.py @@ -115,17 +115,33 @@ def _process_frame84(frame): x_t = np.reshape(x_t, [1, 84, 84]) return x_t.astype(np.uint8) -class ProcessFrame84(gym.Wrapper): - def __init__(self, env=None): - super(ProcessFrame84, self).__init__(env) - self.observation_space = spaces.Box(low=0, high=255, shape=(1, 84, 84)) +def _process_frame42(frame): + img = np.reshape(frame, [210, 160, 3]).astype(np.float32) + img = img[:, :, 0] * 0.299 + img[:, :, 1] * 0.587 + img[:, :, 2] * 0.114 + img = img[34:34 + 160, :160] + img = Image.fromarray(img) + img = img.resize((80, 80), Image.BILINEAR) + img = img.resize((42, 42), Image.BILINEAR) + resized_screen = np.array(img).reshape(1, 42, 42) + return resized_screen.astype(np.uint8) + +class ProcessFrame(gym.Wrapper): + def __init__(self, env=None, frame_size=84): + super(ProcessFrame, self).__init__(env) + self.observation_space = spaces.Box(low=0, high=255, shape=(1, frame_size, frame_size)) + if frame_size == 84: + self.process_fn = _process_frame84 + elif frame_size == 42: + self.process_fn = _process_frame42 + else: + assert(False, "Unknown frame size") def _step(self, action): obs, reward, done, info = self.env.step(action) - return _process_frame84(obs), reward, done, info + return self.process_fn(obs), reward, done, info def _reset(self): - return _process_frame84(self.env.reset()) + return self.process_fn(self.env.reset()) class ClippedRewardsWrapper(gym.Wrapper): def _step(self, action): diff --git a/bootstrap.py b/bootstrap.py index 95964c4..0a730a3 100644 --- a/bootstrap.py +++ b/bootstrap.py @@ -4,51 +4,124 @@ # declaration at the top # ####################################################################### import numpy as np +import torch +from torch.autograd import Variable -def NStepQLearning(batch_states, batch_actions, batch_rewards, - tailing_state, tailing_action, terminal, network, discount): - if terminal: - reward = 0 - else: - reward = np.max(network.predict(np.stack([tailing_state])).flatten()) - rewards = [] - for r in reversed(batch_rewards): - reward = r + discount * reward - rewards.append(reward) - return rewards +class OneStepSarsa: + def __init__(self, agent): + self.agent = agent + self.pending = [] -def OneStepQLearning(batch_states, batch_actions, batch_rewards, - tailing_state, tailing_action, terminal, network, discount): - batch_states.append(tailing_state) - 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) + discount * q_next - return batch_rewards + def process_state(self, network, state): + q = network.predict(np.stack([state])) + self.pending.append([q]) + return q.data.numpy().flatten() -def OneStepSarsa(batch_states, batch_actions, batch_rewards, - tailing_state, tailing_action, terminal, network, discount): - batch_states.append(tailing_state) - batch_actions.append(tailing_action) - 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) + discount * q_next - return batch_rewards + 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.pending = [] + return loss + + +class OneStepQLearning: + def __init__(self, agent): + self.agent = agent + 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.pending = [] + return loss + +class NStepQLearning: + def __init__(self, agent): + self.agent = agent + 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.pending = [] + return loss + +class AdvantageActorCritic: + def __init__(self, agent): + self.agent = agent + 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.pending = [] + return loss -def AdvantageActorCritic(batch_states, batch_actions, batch_rewards, - tailing_state, tailing_action, terminal, network, discount): - if terminal: - reward = 0 - else: - reward = np.asscalar(network.critic(np.stack([tailing_state]))) - rewards = [] - for r in reversed(batch_rewards): - reward = r + discount * reward - rewards.append(reward) - return rewards diff --git a/dqn_agent.py b/dqn_agent.py index 13acc59..288d664 100644 --- a/dqn_agent.py +++ b/dqn_agent.py @@ -59,7 +59,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.stack([self.task.normalize_state(state)])) + value = self.learning_network.predict(np.stack([self.task.normalize_state(state)]), True) if deterministic: action = np.argmax(value.flatten()) else: @@ -81,9 +81,9 @@ class DQNAgent: states, actions, rewards, next_states, terminals = experiences states = self.task.normalize_state(states) next_states = self.task.normalize_state(next_states) - q_next = self.target_network.predict(next_states, False).detach() + q_next = self.target_network.predict(next_states).detach() if self.double_q: - _, best_actions = self.learning_network.predict(next_states, False).detach().max(1) + _, best_actions = self.learning_network.predict(next_states).detach().max(1) q_next = q_next.gather(1, best_actions) else: q_next, _ = q_next.max(1) @@ -92,7 +92,7 @@ class DQNAgent: q_next = q_next * (1 - terminals) q_next.add_(rewards) actions = self.learning_network.to_torch_variable(actions, 'int64').unsqueeze(1) - q = self.learning_network.predict(states, False) + q = self.learning_network.predict(states) q = q.gather(1, actions) loss = self.learning_network.criterion(q, q_next) self.learning_network.zero_grad() diff --git a/main.py b/main.py index b767792..d04ce61 100644 --- a/main.py +++ b/main.py @@ -6,8 +6,8 @@ def dqn_cart_pole(): config = dict() config['task_fn'] = lambda: CartPole() config['optimizer_fn'] = lambda params: torch.optim.RMSprop(params, 0.001) - # config['network_fn'] = lambda optimizer_fn: FullyConnectedNet([8, 50, 200, 2], optimizer_fn) - config['network_fn'] = lambda optimizer_fn: DuelingFullyConnectedNet([8, 50, 200, 2], optimizer_fn) + config['network_fn'] = lambda optimizer_fn: FullyConnectedNet([8, 50, 200, 2], optimizer_fn) + # config['network_fn'] = lambda optimizer_fn: DuelingFullyConnectedNet([8, 50, 200, 2], optimizer_fn) config['policy_fn'] = lambda: GreedyPolicy(epsilon=1.0, final_step=10000, min_epsilon=0.1) config['replay_fn'] = lambda: Replay(memory_size=10000, batch_size=10) config['discount'] = 0.99 @@ -27,16 +27,16 @@ def async_cart_pole(): config = dict() config['task_fn'] = lambda: CartPole() config['optimizer_fn'] = lambda params: torch.optim.Adam(params, 0.001) - config['network_fn'] = lambda: FullyConnectedNet([4, 50, 200, 2], gpu=False) - config['policy_fn'] = lambda: GreedyPolicy(epsilon=1.0, final_step=5000, min_epsilon=0.1) - config['bootstrap_fn'] = OneStepQLearning - # config['bootstrap_fn'] = NStepQLearning - # config['bootstrap_fn'] = OneStepSarsa + config['network_fn'] = lambda: FullyConnectedNet([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['discount'] = 0.99 config['target_network_update_freq'] = 200 config['step_limit'] = 0 config['n_workers'] = 16 - config['batch_size'] = 6 + config['update_interval'] = 6 config['test_interval'] = 4000 config['test_repetitions'] = 50 config['history_length'] = 1 @@ -45,19 +45,20 @@ def async_cart_pole(): agent.run() def a3c_cart_pole(): + update_interval = 6 config = dict() config['task_fn'] = lambda: CartPole() config['optimizer_fn'] = lambda params: torch.optim.Adam(params, 0.001) - config['network_fn'] = lambda: FCActorCriticNet([4, 200, 2], gpu=False) + config['network_fn'] = lambda: FCActorCriticNet([4, 200, 2]) config['policy_fn'] = SamplePolicy - config['bootstrap_fn'] = AdvantageActorCritic + config['bootstrap'] = AdvantageActorCritic config['discount'] = 0.99 - config['target_network_update_freq'] = 0 + config['target_network_update_freq'] = 200 config['step_limit'] = 0 config['n_workers'] = 16 - config['batch_size'] = 6 - config['test_interval'] = 4000 + config['update_interval'] = update_interval config['history_length'] = 1 + config['test_interval'] = 4000 config['test_repetitions'] = 50 config['logger'] = gym.logger agent = AsyncAgent(**config) @@ -89,23 +90,25 @@ def dqn_pixel_atari(name): def async_pixel_atari(name): config = dict() - history_length = 4 + history_length = 1 n_actions = 6 - config['task_fn'] = lambda: PixelAtari(name, no_op=30, frame_skip=4) + config['task_fn'] = lambda: PixelAtari(name, no_op=30, frame_skip=4, frame_size=42) config['optimizer_fn'] = lambda params: torch.optim.Adam(params, lr=0.0001) - config['network_fn'] = lambda: NipsConvNet(history_length, n_actions, gpu=False) + config['network_fn'] = lambda: OpenAIConvNet(history_length, + n_actions, + LSTM=False) config['policy_fn'] = lambda: StochasticGreedyPolicy(epsilons=[1.0, 1.0, 1.0], final_step=1000000, min_epsilons=[0.1, 0.01, 0.5], probs=[0.4, 0.3, 0.3]) - # config['bootstrap_fn'] = OneStepQLearning - # config['bootstrap_fn'] = NStepQLearning - config['bootstrap_fn'] = OneStepSarsa + # config['bootstrap'] = OneStepQLearning + config['bootstrap'] = NStepQLearning + # config['bootstrap'] = OneStepSarsa config['discount'] = 0.99 config['target_network_update_freq'] = 10000 config['step_limit'] = 10000 config['n_workers'] = 16 - config['batch_size'] = 32 + config['update_interval'] = 32 config['test_interval'] = 50000 config['test_repetitions'] = 1 config['history_length'] = history_length @@ -117,16 +120,18 @@ def a3c_pixel_atari(name): config = dict() history_length = 1 n_actions = 6 - config['task_fn'] = lambda: PixelAtari(name, no_op=30, frame_skip=4) + config['task_fn'] = lambda: PixelAtari(name, no_op=30, frame_skip=4, frame_size=42) config['optimizer_fn'] = lambda params: torch.optim.Adam(params, lr=0.0001) - config['network_fn'] = lambda: ConvActorCriticNet(history_length, n_actions, gpu=False) + config['network_fn'] = lambda: OpenAIConvActorCriticNet(history_length, + n_actions, + LSTM=True) config['policy_fn'] = SamplePolicy - config['bootstrap_fn'] = AdvantageActorCritic + config['bootstrap'] = AdvantageActorCritic config['discount'] = 0.99 config['target_network_update_freq'] = 0 config['step_limit'] = 10000 config['n_workers'] = 16 - config['batch_size'] = 20 + config['update_interval'] = 20 config['test_interval'] = 50000 config['test_repetitions'] = 1 config['history_length'] = history_length @@ -138,12 +143,14 @@ if __name__ == '__main__': # gym.logger.setLevel(logging.DEBUG) gym.logger.setLevel(logging.INFO) - # async_cart_pole() # dqn_cart_pole() - # dqn_pixel_atari('BreakoutNoFrameskip-v3') - # async_pixel_atari('BreakoutNoFrameskip-v3') - # a3c_pixel_atari('BreakoutNoFrameskip-v3') + # async_cart_pole() # a3c_cart_pole() + # dqn_pixel_atari('PongNoFrameskip-v3') - async_pixel_atari('PongNoFrameskip-v3') + # async_pixel_atari('PongNoFrameskip-v3') # a3c_pixel_atari('PongNoFrameskip-v3') + + # dqn_pixel_atari('BreakoutNoFrameskip-v3') + async_pixel_atari('BreakoutNoFrameskip-v3') + # a3c_pixel_atari('BreakoutNoFrameskip-v3') diff --git a/network.py b/network.py index 631d3ce..659303b 100644 --- a/network.py +++ b/network.py @@ -12,63 +12,57 @@ import numpy as np # Base class for all kinds of network class BasicNet: - def __init__(self, optimizer_fn, gpu): + def __init__(self, optimizer_fn, gpu, LSTM=False): if optimizer_fn is not None: self.optimizer = optimizer_fn(self.parameters()) self.gpu = gpu and torch.cuda.is_available() + self.LSTM = LSTM if self.gpu: print 'Transferring network to GPU...' self.cuda() print 'Network transferred.' def to_torch_variable(self, x, dtype='float32'): - x = torch.from_numpy(np.asarray(x, dtype=dtype)) + if not isinstance(x, torch.FloatTensor): + x = torch.from_numpy(np.asarray(x, dtype=dtype)) if self.gpu: x = x.cuda() return Variable(x) + def reset(self, terminal): + if not self.LSTM: + return + if terminal: + self.h.data.zero_() + self.c.data.zero_() + self.h = Variable(self.h.data) + self.c = Variable(self.c.data) + # Base class for value based methods class VanillaNet(BasicNet): - def predict(self, x, to_numpy=True): + def predict(self, x, to_numpy=False): y = self.forward(x) if to_numpy: y = y.cpu().data.numpy() return y - def gradient(self, x, actions, targets): - y = self.forward(x) - y = y.gather(1, actions) - loss = self.criterion(y, targets) - loss.backward() - # Base class for actor critic method class ActorCriticNet(BasicNet): def predict(self, x): - phi = self.forward(x) - return F.softmax(self.fc_actor(phi)).cpu().data.numpy() - - def gradient(self, x, actions, rewards): - phi = self.forward(x) - logit = self.fc_actor(phi) - prob = F.softmax(logit) - log_prob_ = F.log_softmax(logit) - state_value = self.fc_critic(phi) - log_prob = log_prob_.gather(1, actions) - advantage = (rewards - state_value).detach() - policy_loss = -torch.sum(log_prob * advantage) - value_loss = 0.5 * torch.sum(torch.pow(rewards - state_value, 2)) - entropy = -torch.sum(torch.mul(prob, log_prob_)) - loss = policy_loss + value_loss - self.xentropy_weight * entropy - loss.backward() - nn.utils.clip_grad_norm(self.parameters(), self.grad_threshold) + phi = self.forward(x, True) + pre_prob = self.fc_actor(phi) + prob = F.softmax(pre_prob) + log_prob = F.log_softmax(pre_prob) + value = self.fc_critic(phi) + return prob, log_prob, value def critic(self, x): - phi = self.forward(x) - return self.fc_critic(phi).cpu().data.numpy() + phi = self.forward(x, False) + return self.fc_critic(phi) # Base class for dueling architecture class DuelingNet(BasicNet): - def predict(self, x, to_numpy=True): + def predict(self, x, to_numpy=False): phi = self.forward(x) value = self.fc_value(phi) advantange = self.fc_advantage(phi) @@ -77,6 +71,8 @@ class DuelingNet(BasicNet): return q.cpu().data.numpy() return q +# Starting of several network instances + # Network for CartPole with value based methods class FullyConnectedNet(nn.Module, VanillaNet): def __init__(self, dims, optimizer_fn=None, gpu=True): @@ -178,21 +174,30 @@ class DuelingConvNet(nn.Module, DuelingNet): class FCActorCriticNet(nn.Module, ActorCriticNet): def __init__(self, dims, - xentropy_weight=0.01, - grad_threshold=40, - gpu=True): + LSTM=False): super(FCActorCriticNet, self).__init__() - self.fc1 = nn.Linear(dims[0], dims[1]) + if LSTM: + self.layer1 = nn.LSTMCell(dims[0], dims[1]) + else: + self.layer1 = 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) + BasicNet.__init__(self, optimizer_fn=None, gpu=False, LSTM=LSTM) + if LSTM: + self.h = self.to_torch_variable(np.zeros((1, dims[1]))) + self.c = self.to_torch_variable(np.zeros((1, dims[1]))) - def forward(self, x): + def forward(self, x, update_LSTM=True): x = self.to_torch_variable(x) x = x.view(x.size(0), -1) - phi = self.fc1(x) + if self.LSTM: + h, c = self.layer1(x, (self.h, self.c)) + if update_LSTM: + self.h = h + self.c = c + phi = h + else: + phi = self.layer1(x) return phi # Network for pixel Atari game with actor critic @@ -222,3 +227,87 @@ class ConvActorCriticNet(nn.Module, ActorCriticNet): y = y.view(y.size(0), -1) return F.elu(self.fc4(y)) +class OpenAIConvActorCriticNet(nn.Module, ActorCriticNet): + def __init__(self, + in_channels, + n_actions, + LSTM=False): + super(OpenAIConvActorCriticNet, self).__init__() + self.conv1 = nn.Conv2d(in_channels, 32, 3, stride=2, padding=1) + self.conv2 = nn.Conv2d(32, 32, 3, stride=2, padding=1) + self.conv3 = nn.Conv2d(32, 32, 3, stride=2, padding=1) + self.conv4 = nn.Conv2d(32, 32, 3, stride=2, padding=1) + + self.LSTM = LSTM + hidden_units = 256 + + if LSTM: + self.layer5 = nn.LSTMCell(32 * 3 * 3, hidden_units) + else: + self.layer5 = nn.Linear(32 * 3 * 3, hidden_units) + + self.fc_actor = nn.Linear(hidden_units, n_actions) + self.fc_critic = nn.Linear(hidden_units, 1) + BasicNet.__init__(self, optimizer_fn=None, gpu=False, LSTM=LSTM) + if LSTM: + self.h = self.to_torch_variable(np.zeros((1, hidden_units))) + self.c = self.to_torch_variable(np.zeros((1, hidden_units))) + + def forward(self, x, update_LSTM=True): + x = self.to_torch_variable(x) + y = F.elu(self.conv1(x)) + y = F.elu(self.conv2(y)) + y = F.elu(self.conv3(y)) + y = F.elu(self.conv4(y)) + y = y.view(y.size(0), -1) + if self.LSTM: + h, c = self.layer5(y, (self.h, self.c)) + if update_LSTM: + self.h = h + self.c = c + phi = h + else: + phi = F.elu(self.layer5(y)) + return phi + +class OpenAIConvNet(nn.Module, VanillaNet): + def __init__(self, + in_channels, + n_actions, + LSTM=False): + super(OpenAIConvNet, self).__init__() + self.conv1 = nn.Conv2d(in_channels, 32, 3, stride=2, padding=1) + self.conv2 = nn.Conv2d(32, 32, 3, stride=2, padding=1) + self.conv3 = nn.Conv2d(32, 32, 3, stride=2, padding=1) + self.conv4 = nn.Conv2d(32, 32, 3, stride=2, padding=1) + + self.LSTM = LSTM + hidden_units = 256 + + if LSTM: + self.layer5 = nn.LSTMCell(32 * 3 * 3, hidden_units) + else: + self.layer5 = nn.Linear(32 * 3 * 3, hidden_units) + + self.fc6 = nn.Linear(hidden_units, n_actions) + BasicNet.__init__(self, optimizer_fn=None, gpu=False, LSTM=LSTM) + if LSTM: + self.h = self.to_torch_variable(np.zeros((1, hidden_units))) + self.c = self.to_torch_variable(np.zeros((1, hidden_units))) + + def forward(self, x, update_LSTM=True): + x = self.to_torch_variable(x) + y = F.elu(self.conv1(x)) + y = F.elu(self.conv2(y)) + y = F.elu(self.conv3(y)) + y = F.elu(self.conv4(y)) + y = y.view(y.size(0), -1) + if self.LSTM: + h, c = self.layer5(y, (self.h, self.c)) + if update_LSTM: + self.h = h + self.c = c + phi = h + else: + phi = F.elu(self.layer5(y)) + return self.fc6(phi) diff --git a/policy.py b/policy.py index 3e4c13e..84b2c53 100644 --- a/policy.py +++ b/policy.py @@ -13,7 +13,9 @@ class GreedyPolicy: self.min_epsilon = min_epsilon self.final_step = final_step - def sample(self, action_value): + def sample(self, action_value, deterministic=False): + if deterministic: + return np.argmax(action_value) if np.random.rand() < self.epsilon: return np.random.randint(0, len(action_value)) return np.argmax(action_value) @@ -30,15 +32,17 @@ class StochasticGreedyPolicy: for epsilon, min_epsilon in zip(epsilons, min_epsilons): self.policies.append(GreedyPolicy(epsilon, final_step, min_epsilon)) - def sample(self, action_value): - return np.random.choice(self.policies, p=self.probs).sample(action_value) + def sample(self, action_value, deterministic=False): + return np.random.choice(self.policies, p=self.probs).sample(action_value, deterministic) def update_epsilon(self): for policy in self.policies: policy.update_epsilon() class SamplePolicy: - def sample(self, action_value): + def sample(self, action_value, deterministic=False): + if deterministic: + return np.argmax(action_value) return np.random.choice(np.arange(len(action_value)), p=action_value) def update_epsilon(self): pass \ No newline at end of file diff --git a/task.py b/task.py index 7026c23..f6cf517 100644 --- a/task.py +++ b/task.py @@ -53,12 +53,12 @@ class LunarLander(BasicTask): self.env = gym.make(self.name) class PixelAtari(BasicTask): - success_threshold = 1000 - - def __init__(self, name, no_op, frame_skip, normalized_state=True): + def __init__(self, name, no_op, frame_skip, normalized_state=True, + frame_size=84, success_threshold=1000): BasicTask.__init__(self) self.normalized_state = normalized_state self.name = name + self.success_threshold = success_threshold env = gym.make(name) assert 'NoFrameskip' in env.spec.id env = EpisodicLifeEnv(env) @@ -66,7 +66,7 @@ class PixelAtari(BasicTask): env = MaxAndSkipEnv(env, skip=frame_skip) if 'FIRE' in env.unwrapped.get_action_meanings(): env = FireResetEnv(env) - env = ProcessFrame84(env) + env = ProcessFrame(env, frame_size) self.env = ClippedRewardsWrapper(env) def normalize_state(self, state):