diff --git a/agents/models.py b/agents/models.py new file mode 100644 index 0000000..9377959 --- /dev/null +++ b/agents/models.py @@ -0,0 +1,385 @@ +""" +A2C, IA2C, MA2C models +@author: Tianshu Chu +""" + +import os +from agents.utils import * +from agents.policies import * +import logging +import multiprocessing as mp +import numpy as np +import tensorflow.compat.v1 as tf + + +class A2C: + def __init__(self, n_s, n_a, total_step, model_config, seed=0, n_f=None): + # load parameters + self.name = 'a2c' + self.n_agent = 1 + # init reward norm/clip + self.reward_clip = model_config.getfloat('reward_clip') + self.reward_norm = model_config.getfloat('reward_norm') + self.n_s = n_s + self.n_a = n_a + self.n_step = model_config.getint('batch_size') + # init tf + tf.reset_default_graph() + tf.set_random_seed(seed) + config = tf.ConfigProto(allow_soft_placement=True) + self.sess = tf.Session(config=config) + self.policy = self._init_policy(n_s, n_a, n_f, model_config) + self.saver = tf.train.Saver(max_to_keep=5) + if total_step: + # training + self.total_step = total_step + self._init_scheduler(model_config) + self._init_train(model_config) + self.sess.run(tf.global_variables_initializer()) + + def _init_policy(self, n_s, n_a, n_w, n_f, model_config, agent_name=None): + n_fw = model_config.getint('num_fw') + n_ft = model_config.getint('num_ft') + n_lstm = model_config.getint('num_lstm') + n_fp = model_config.getint('num_fp') + policy = FPLstmACPolicy(n_s, n_a, n_w, n_f, self.n_step, n_fc_wave=n_fw, + n_fc_wait=n_ft, n_fc_fp=n_fp, n_lstm=n_lstm, name=agent_name) + return policy + + def _init_scheduler(self, model_config): + lr_init = model_config.getfloat('lr_init') + lr_decay = model_config.get('lr_decay') + beta_init = model_config.getfloat('entropy_coef_init') + beta_decay = model_config.get('entropy_decay') + if lr_decay == 'constant': + self.lr_scheduler = Scheduler(lr_init, decay=lr_decay) + else: + lr_min = model_config.getfloat('LR_MIN') + self.lr_scheduler = Scheduler( + lr_init, lr_min, self.total_step, decay=lr_decay) + if beta_decay == 'constant': + self.beta_scheduler = Scheduler(beta_init, decay=beta_decay) + else: + beta_min = model_config.getfloat('ENTROPY_COEF_MIN') + beta_ratio = model_config.getfloat('ENTROPY_RATIO') + self.beta_scheduler = Scheduler(beta_init, beta_min, self.total_step * beta_ratio, + decay=beta_decay) + + def _init_train(self, model_config): + # init loss + v_coef = model_config.getfloat('value_coef') + max_grad_norm = model_config.getfloat('max_grad_norm') + alpha = model_config.getfloat('rmsp_alpha') + epsilon = model_config.getfloat('rmsp_epsilon') + self.policy.prepare_loss(v_coef, max_grad_norm, alpha, epsilon) + + # init replay buffer + gamma = model_config.getfloat('gamma') + self.trans_buffer = OnPolicyBuffer(gamma) + + def save(self, model_dir, global_step): + self.saver.save(self.sess, model_dir + 'checkpoint', + global_step=global_step) + + def load(self, model_dir, checkpoint=None): + save_file = None + save_step = 0 + if os.path.exists(model_dir): + if checkpoint is None: + for file in os.listdir(model_dir): + if file.startswith('checkpoint'): + prefix = file.split('.')[0] + tokens = prefix.split('-') + if len(tokens) != 2: + continue + cur_step = int(tokens[1]) + if cur_step > save_step: + save_file = prefix + save_step = cur_step + else: + save_file = 'checkpoint-' + str(int(checkpoint)) + if save_file is not None: + self.saver.restore(self.sess, model_dir + save_file) + logging.info('Checkpoint loaded: %s' % save_file) + return True + logging.error('Can not find old checkpoint for %s' % model_dir) + return False + + def reset(self): + self.policy._reset() + + def backward(self, R, summary_writer=None, global_step=None): + cur_lr = self.lr_scheduler.get(self.n_step) + cur_beta = self.beta_scheduler.get(self.n_step) + obs, acts, dones, Rs, Advs = self.trans_buffer.sample_transition(R) + self.policy.backward(self.sess, obs, acts, dones, Rs, Advs, cur_lr, cur_beta, + summary_writer=summary_writer, global_step=global_step) + + def forward(self, ob, done, out_type='pv'): + return self.policy.forward(self.sess, ob, done, out_type) + + def add_transition(self, ob, action, reward, value, done): + # Hard code the reward norm for negative reward only + if (self.reward_norm): + reward /= self.reward_norm + if self.reward_clip: + reward = np.clip(reward, -self.reward_clip, self.reward_clip) + self.trans_buffer.add_transition(ob, action, reward, value, done) + + +class IA2C(A2C): + def __init__(self, n_s_ls, n_a_ls, n_w_ls, total_step, + model_config, seed=0): + self.name = 'ia2c' + self.agents = [] + self.n_agent = len(n_s_ls) + self.reward_clip = model_config.getfloat('reward_clip') + self.reward_norm = model_config.getfloat('reward_norm') + self.n_s_ls = n_s_ls + self.n_a_ls = n_a_ls + self.n_w_ls = n_w_ls + self.n_step = model_config.getint('batch_size') + # init tf + tf.reset_default_graph() + tf.set_random_seed(seed) + config = tf.ConfigProto(allow_soft_placement=True) + self.sess = tf.Session(config=config) + self.policy_ls = [] + for i, (n_s, n_w, n_a) in enumerate(zip(self.n_s_ls, self.n_w_ls, self.n_a_ls)): + # agent_name is needed to differentiate multi-agents + self.policy_ls.append(self._init_policy(n_s - n_w, n_a, n_w, 0, model_config, + agent_name='{:d}a'.format(i))) + self.saver = tf.train.Saver(max_to_keep=5) + if total_step: + # training + self.total_step = total_step + self._init_scheduler(model_config) + self._init_train(model_config) + self.sess.run(tf.global_variables_initializer()) + + def _init_train(self, model_config): + # init loss + v_coef = model_config.getfloat('value_coef') + max_grad_norm = model_config.getfloat('max_grad_norm') + alpha = model_config.getfloat('rmsp_alpha') + epsilon = model_config.getfloat('rmsp_epsilon') + gamma = model_config.getfloat('gamma') + self.trans_buffer_ls = [] + for i in range(self.n_agent): + self.policy_ls[i].prepare_loss( + v_coef, max_grad_norm, alpha, epsilon) + self.trans_buffer_ls.append(OnPolicyBuffer(gamma)) + + def backward(self, R_ls, summary_writer=None, global_step=None): + cur_lr = self.lr_scheduler.get(self.n_step) + cur_beta = self.beta_scheduler.get(self.n_step) + for i in range(self.n_agent): + obs, acts, dones, Rs, Advs = self.trans_buffer_ls[i].sample_transition( + R_ls[i]) + # Check if len(mini_batch) = batch_size or not + if len(obs) == self.n_step: + if i == 0: + self.policy_ls[i].backward(self.sess, obs, acts, dones, Rs, Advs, cur_lr, cur_beta, + summary_writer=summary_writer, global_step=global_step) + else: + self.policy_ls[i].backward( + self.sess, obs, acts, dones, Rs, Advs, cur_lr, cur_beta) + + def forward(self, obs, done, out_type='pv'): + if len(out_type) == 1: + out = [] + elif len(out_type) == 2: + out1, out2 = [], [] + for i in range(self.n_agent): + cur_out = self.policy_ls[i].forward( + self.sess, obs[i], done, out_type) + if len(out_type) == 1: + out.append(cur_out) + else: + out1.append(cur_out[0]) + out2.append(cur_out[1]) + if len(out_type) == 1: + return out + else: + return out1, out2 + + def backward_mp(self, R_ls, summary_writer=None, global_step=None): + cur_lr = self.lr_scheduler.get(self.n_step) + cur_beta = self.beta_scheduler.get(self.n_step) + + def worker(i): + obs, acts, dones, Rs, Advs = self.trans_buffer_ls[i].sample_transition( + R_ls[i]) + self.policy_ls[i].backward(self.sess, obs, acts, dones, Rs, Advs, cur_lr, cur_beta, + summary_writer=summary_writer, global_step=global_step) + mps = [] + for i in range(self.n_agent): + p = mp.Process(target=worker, args=(i)) + p.start() + mps.append(p) + for p in mps: + p.join() + + def reset(self): + for policy in self.policy_ls: + policy._reset() + + def add_transition(self, obs, actions, rewards, values, done): + if (self.reward_norm): + rewards = rewards / self.reward_norm + if self.reward_clip: + rewards = np.clip(rewards, -self.reward_clip, self.reward_clip) + for i in range(self.n_agent): + self.trans_buffer_ls[i].add_transition(obs[i], actions[i], + rewards[i], values[i], done) + + +class MA2C(IA2C): + def __init__(self, n_s_ls, n_a_ls, n_w_ls, n_f_ls, total_step, + model_config, seed=0): + self.name = 'ma2c' + self.agents = [] + self.n_agent = len(n_s_ls) + self.reward_clip = model_config.getfloat('reward_clip') + self.reward_norm = model_config.getfloat('reward_norm') + self.n_s_ls = n_s_ls + self.n_a_ls = n_a_ls + self.n_f_ls = n_f_ls + self.n_w_ls = n_w_ls + self.n_step = model_config.getint('batch_size') + + tf.reset_default_graph() + tf.set_random_seed(seed) + config = tf.ConfigProto(allow_soft_placement=True) + self.sess = tf.Session(config=config) + self.policy_ls = [] + for i, (n_s, n_a, n_w, n_f) in enumerate(zip(self.n_s_ls, self.n_a_ls, self.n_w_ls, self.n_f_ls)): + # agent_name is needed to differentiate multi-agents + self.policy_ls.append(self._init_policy(n_s - n_f - n_w, n_a, n_w, n_f, model_config, + agent_name='{:d}a'.format(i))) + self.saver = tf.train.Saver(max_to_keep=5) + if total_step: + # training + self.total_step = total_step + self._init_scheduler(model_config) + self._init_train(model_config) + self.sess.run(tf.global_variables_initializer()) + + +class IQL(A2C): + def __init__(self, n_s_ls, n_a_ls, n_w_ls, total_step, model_config, seed=0, model_type='dqn'): + self.name = 'iql' + self.model_type = model_type + self.agents = [] + self.n_agent = len(n_s_ls) + self.reward_clip = model_config.getfloat('reward_clip') + self.reward_norm = model_config.getfloat('reward_norm') + self.n_s_ls = n_s_ls + self.n_a_ls = n_a_ls + self.n_w_ls = n_w_ls + self.n_step = model_config.getint('batch_size') + # init tf + tf.reset_default_graph() + tf.set_random_seed(seed) + config = tf.ConfigProto(allow_soft_placement=True) + self.sess = tf.Session(config=config) + self.policy_ls = [] + for i, (n_s, n_a, n_w) in enumerate(zip(self.n_s_ls, self.n_a_ls, self.n_w_ls)): + # agent_name is needed to differentiate multi-agents + self.policy_ls.append(self._init_policy(n_s, n_a, n_w, model_config, + agent_name='{:d}a'.format(i))) + self.saver = tf.train.Saver(max_to_keep=5) + if total_step: + # training + self.total_step = total_step + self._init_scheduler(model_config) + self._init_train(model_config) + self.cur_step = 0 + self.sess.run(tf.global_variables_initializer()) + + def _init_policy(self, n_s, n_a, n_w, model_config, agent_name=None): + if self.model_type == 'dqn': + n_h = model_config.getint('num_h') + n_fc = model_config.getint('num_fc') + policy = DeepQPolicy(n_s - n_w, n_a, n_w, self.n_step, n_fc0=n_fc, n_fc=n_h, + name=agent_name) + else: + policy = LRQPolicy(n_s, n_a, self.n_step, name=agent_name) + return policy + + def _init_scheduler(self, model_config): + lr_init = model_config.getfloat('lr_init') + lr_decay = model_config.get('lr_decay') + eps_init = model_config.getfloat('epsilon_init') + eps_decay = model_config.get('epsilon_decay') + if lr_decay == 'constant': + self.lr_scheduler = Scheduler(lr_init, decay=lr_decay) + else: + lr_min = model_config.getfloat('lr_min') + self.lr_scheduler = Scheduler( + lr_init, lr_min, self.total_step, decay=lr_decay) + if eps_decay == 'constant': + self.eps_scheduler = Scheduler(eps_init, decay=eps_decay) + else: + eps_min = model_config.getfloat('epsilon_min') + eps_ratio = model_config.getfloat('epsilon_ratio') + self.eps_scheduler = Scheduler(eps_init, eps_min, self.total_step * eps_ratio, + decay=eps_decay) + + def _init_train(self, model_config): + # init loss + max_grad_norm = model_config.getfloat('max_grad_norm') + gamma = model_config.getfloat('gamma') + buffer_size = model_config.getfloat('buffer_size') + self.trans_buffer_ls = [] + for i in range(self.n_agent): + self.policy_ls[i].prepare_loss(max_grad_norm, gamma) + self.trans_buffer_ls.append(ReplayBuffer(buffer_size, self.n_step)) + + def backward(self, summary_writer=None, global_step=None): + cur_lr = self.lr_scheduler.get(self.n_step) + if self.trans_buffer_ls[0].size < self.trans_buffer_ls[0].batch_size: + return + for i in range(self.n_agent): + for k in range(10): + obs, acts, next_obs, rs, dones = self.trans_buffer_ls[i].sample_transition( + ) + if i == 0: + self.policy_ls[i].backward(self.sess, obs, acts, next_obs, dones, rs, cur_lr, + summary_writer=summary_writer, + global_step=global_step + k) + else: + self.policy_ls[i].backward( + self.sess, obs, acts, next_obs, dones, rs, cur_lr) + + def forward(self, obs, mode='act', stochastic=False): + if mode == 'explore': + eps = self.eps_scheduler.get(1) + action = [] + qs_ls = [] + for i in range(self.n_agent): + qs = self.policy_ls[i].forward(self.sess, obs[i]) + if (mode == 'explore') and (np.random.random() < eps): + action.append(np.random.randint(self.n_a_ls[i])) + else: + if not stochastic: + action.append(np.argmax(qs)) + else: + qs = qs / np.sum(qs) + action.append(np.random.choice(np.arange(len(qs)), p=qs)) + qs_ls.append(qs) + return action, qs_ls + + def reset(self): + # do nothing + return + + def add_transition(self, obs, actions, rewards, next_obs, done): + if (self.reward_norm): + rewards = rewards / self.reward_norm + if self.reward_clip: + rewards = np.clip(rewards, -self.reward_clip, self.reward_clip) + for i in range(self.n_agent): + + self.trans_buffer_ls[i].add_transition(obs[i], actions[i], + rewards[i], next_obs[i], done) diff --git a/agents/policies.py b/agents/policies.py new file mode 100644 index 0000000..6f602e8 --- /dev/null +++ b/agents/policies.py @@ -0,0 +1,410 @@ +from agents.utils import * +import numpy as np +# import tensorflow +# import tensorflow.compat.v1 as tf +# tf.disable_v2_behavior() +import tensorflow as tf +tf.compat.v1.disable_v2_behavior() + + +class ACPolicy: + def __init__(self, n_a, n_s, n_step, policy_name, agent_name): + self.name = policy_name + if agent_name is not None: + # for multi-agent system + self.name += '_' + str(agent_name) + self.n_a = n_a + self.n_s = n_s + self.n_step = n_step + + def forward(self, ob, *_args, **_kwargs): + raise NotImplementedError() + + def _build_out_net(self, h, out_type): + if out_type == 'pi': + pi = fc(h, out_type, self.n_a, act=tf.nn.softmax) + return tf.squeeze(pi) + else: + v = fc(h, out_type, 1, act=lambda x: x) + return tf.squeeze(v) + + def _get_forward_outs(self, out_type): + outs = [] + if 'p' in out_type: + outs.append(self.pi) + if 'v' in out_type: + outs.append(self.v) + return outs + + def _return_forward_outs(self, out_values): + if len(out_values) == 1: + return out_values[0] + return out_values + + def prepare_loss(self, v_coef, max_grad_norm, alpha, epsilon): + self.A = tf.compat.v1.placeholder(tf.int32, [self.n_step]) + self.ADV = tf.compat.v1.placeholder(tf.float32, [self.n_step]) + self.R = tf.compat.v1.placeholder(tf.float32, [self.n_step]) + self.entropy_coef = tf.compat.v1.placeholder(tf.float32, []) + A_sparse = tf.one_hot(self.A, self.n_a) + log_pi = tf.compat.v1.log(tf.clip_by_value(self.pi, 1e-10, 1.0)) + entropy = -tf.reduce_sum(self.pi * log_pi, axis=1) + entropy_loss = -tf.reduce_mean(entropy) * self.entropy_coef + policy_loss = - \ + tf.reduce_mean(tf.reduce_sum(log_pi * A_sparse, axis=1) * self.ADV) + value_loss = tf.reduce_mean(tf.square(self.R - self.v)) * 0.5 * v_coef + self.loss = policy_loss + value_loss + entropy_loss + + wts = tf.compat.v1.trainable_variables(scope=self.name) + grads = tf.compat.v1.gradients(self.loss, wts) + if max_grad_norm > 0: + grads, self.grad_norm = tf.clip_by_global_norm( + grads, max_grad_norm) + self.lr = tf.compat.v1.placeholder(tf.float32, []) + self.optimizer = tf.compat.v1.train.RMSPropOptimizer(learning_rate=self.lr, decay=alpha, + epsilon=epsilon) + self._train = self.optimizer.apply_gradients(list(zip(grads, wts))) + # monitor training + if self.name.endswith('_0a'): + summaries = [] + # summaries.append(tf.summary.scalar('loss/%s_entropy_loss' % self.name, entropy_loss)) + summaries.append(tf.compat.v1.summary.scalar( + 'loss/%s_policy_loss' % self.name, policy_loss)) + summaries.append(tf.compat.v1.summary.scalar( + 'loss/%s_value_loss' % self.name, value_loss)) + summaries.append(tf.compat.v1.summary.scalar( + 'loss/%s_total_loss' % self.name, self.loss)) + # summaries.append(tf.summary.scalar('train/%s_lr' % self.name, self.lr)) + # summaries.append(tf.summary.scalar('train/%s_entropy_beta' % self.name, self.entropy_coef)) + summaries.append(tf.compat.v1.summary.scalar( + 'train/%s_gradnorm' % self.name, self.grad_norm)) + self.summary = tf.compat.v1.summary.merge(summaries) + + +class LstmACPolicy(ACPolicy): + def __init__(self, n_s, n_a, n_w, n_step, n_fc_wave=128, n_fc_wait=32, n_lstm=64, name=None): + super().__init__(n_a, n_s, n_step, 'lstm', name) + self.n_lstm = n_lstm + self.n_fc_wait = n_fc_wait + self.n_fc_wave = n_fc_wave + self.n_w = n_w + self.ob_fw = tf.compat.v1.placeholder( + tf.float32, [1, n_s + n_w]) # forward 1-step + self.done_fw = tf.compat.v1.placeholder(tf.float32, [1]) + self.ob_bw = tf.compat.v1.placeholder( + tf.float32, [n_step, n_s + n_w]) # backward n-step + self.done_bw = tf.compat.v1.placeholder(tf.float32, [n_step]) + self.states = tf.compat.v1.placeholder(tf.float32, [2, n_lstm * 2]) + with tf.variable_scope(self.name): + # pi and v use separate nets + self.pi_fw, pi_state = self._build_net('forward', 'pi') + self.v_fw, v_state = self._build_net('forward', 'v') + pi_state = tf.expand_dims(pi_state, 0) + v_state = tf.expand_dims(v_state, 0) + self.new_states = tf.concat([pi_state, v_state], 0) + with tf.variable_scope(self.name, reuse=True): + self.pi, _ = self._build_net('backward', 'pi') + self.v, _ = self._build_net('backward', 'v') + self._reset() + + def _build_net(self, in_type, out_type): + if in_type == 'forward': + ob = self.ob_fw + done = self.done_fw + else: + ob = self.ob_bw + done = self.done_bw + if out_type == 'pi': + states = self.states[0] + else: + states = self.states[1] + if self.n_w == 0: + h = fc(ob, out_type + '_fcw', self.n_fc_wave) + else: + h0 = fc(ob[:, :self.n_s], out_type + '_fcw', self.n_fc_wave) + h1 = fc(ob[:, self.n_s:], out_type + '_fct', self.n_fc_wait) + h = tf.concat([h0, h1], 1) + h, new_states = lstm(h, done, states, out_type + '_lstm') + out_val = self._build_out_net(h, out_type) + return out_val, new_states + + def _reset(self): + # forget the cumulative states every cum_step + self.states_fw = np.zeros((2, self.n_lstm * 2), dtype=np.float32) + self.states_bw = np.zeros((2, self.n_lstm * 2), dtype=np.float32) + + def forward(self, sess, ob, done, out_type='pv'): + outs = self._get_forward_outs(out_type) + # update state only when p is called + if 'p' in out_type: + outs.append(self.new_states) + out_values = sess.run(outs, {self.ob_fw: np.array([ob]), + self.done_fw: np.array([done]), + self.states: self.states_fw}) + if 'p' in out_type: + self.states_fw = out_values[-1] + out_values = out_values[:-1] + return self._return_forward_outs(out_values) + + def backward(self, sess, obs, acts, dones, Rs, Advs, cur_lr, cur_beta, + summary_writer=None, global_step=None): + if summary_writer is None: + ops = self._train + else: + ops = [self.summary, self._train] + outs = sess.run(ops, + {self.ob_bw: obs, + self.done_bw: dones, + self.states: self.states_bw, + self.A: acts, + self.ADV: Advs, + self.R: Rs, + self.lr: cur_lr, + self.entropy_coef: cur_beta}) + self.states_bw = np.copy(self.states_fw) + if summary_writer is not None: + summary_writer.add_summary(outs[0], global_step=global_step) + + def _get_forward_outs(self, out_type): + outs = [] + if 'p' in out_type: + outs.append(self.pi_fw) + if 'v' in out_type: + outs.append(self.v_fw) + return outs + + +class FPLstmACPolicy(LstmACPolicy): + def __init__(self, n_s, n_a, n_w, n_f, n_step, n_fc_wave=128, n_fc_wait=32, n_fc_fp=32, n_lstm=64, name=None): + ACPolicy.__init__(self, n_a, n_s, n_step, 'fplstm', name) + self.n_lstm = n_lstm + self.n_fc_wave = n_fc_wave + self.n_fc_wait = n_fc_wait + self.n_fc_fp = n_fc_fp + self.n_w = n_w + self.ob_fw = tf.compat.v1.placeholder( + tf.float32, [1, n_s + n_w + n_f]) # forward 1-step + self.done_fw = tf.compat.v1.placeholder(tf.float32, [1]) + self.ob_bw = tf.compat.v1.placeholder( + tf.float32, [n_step, n_s + n_w + n_f]) # backward n-step + self.done_bw = tf.compat.v1.placeholder(tf.float32, [n_step]) + self.states = tf.compat.v1.placeholder(tf.float32, [2, n_lstm * 2]) + with tf.compat.v1.variable_scope(self.name): + # pi and v use separate nets + self.pi_fw, pi_state = self._build_net('forward', 'pi') + self.v_fw, v_state = self._build_net('forward', 'v') + pi_state = tf.expand_dims(pi_state, 0) + v_state = tf.expand_dims(v_state, 0) + self.new_states = tf.concat([pi_state, v_state], 0) + with tf.compat.v1.variable_scope(self.name, reuse=True): + self.pi, _ = self._build_net('backward', 'pi') + self.v, _ = self._build_net('backward', 'v') + self._reset() + + def _build_net(self, in_type, out_type): + if in_type == 'forward': + ob = self.ob_fw + done = self.done_fw + else: + ob = self.ob_bw + done = self.done_bw + if out_type == 'pi': + states = self.states[0] + else: + states = self.states[1] + h0 = fc(ob[:, :self.n_s], out_type + '_fcw', self.n_fc_wave) + h1 = fc(ob[:, (self.n_s + self.n_w):], out_type + '_fcf', self.n_fc_fp) + if self.n_w == 0: + h = tf.concat([h0, h1], 1) + else: + h2 = fc(ob[:, self.n_s: (self.n_s + self.n_w)], + out_type + '_fct', self.n_fc_wait) + h = tf.concat([h0, h1, h2], 1) + h, new_states = lstm(h, done, states, out_type + '_lstm') + out_val = self._build_out_net(h, out_type) + return out_val, new_states + + +class FcACPolicy(ACPolicy): + def __init__(self, n_s, n_a, n_w, n_step, n_fc_wave=128, n_fc_wait=32, n_lstm=64, name=None): + super().__init__(n_a, n_s, n_step, 'fc', name) + self.n_fc_wave = n_fc_wave + self.n_fc_wait = n_fc_wait + self.n_fc = n_lstm + self.n_w = n_w + self.obs = tf.placeholder(tf.float32, [None, n_s + n_w]) + with tf.variable_scope(self.name): + # pi and v use separate nets + self.pi = self._build_net('pi') + self.v = self._build_net('v') + + def _build_net(self, out_type): + if self.n_w == 0: + h = fc(self.obs, out_type + '_fcw', self.n_fc_wave) + else: + h0 = fc(self.obs[:, :self.n_s], out_type + '_fcw', self.n_fc_wave) + h1 = fc(self.obs[:, self.n_s:], out_type + '_fct', self.n_fc_wait) + h = tf.concat([h0, h1], 1) + h = fc(h, out_type + '_fc', self.n_fc) + return self._build_out_net(h, out_type) + + def forward(self, sess, ob, done, out_type='pv'): + outs = self._get_forward_outs(out_type) + out_values = sess.run(outs, {self.obs: np.array([ob])}) + return self._return_forward_outs(out_values) + + def backward(self, sess, obs, acts, dones, Rs, Advs, cur_lr, cur_beta, + summary_writer=None, global_step=None): + if summary_writer is None: + ops = self._train + else: + ops = [self.summary, self._train] + outs = sess.run(ops, + {self.obs: obs, + self.A: acts, + self.ADV: Advs, + self.R: Rs, + self.lr: cur_lr, + self.entropy_coef: cur_beta}) + if summary_writer is not None: + summary_writer.add_summary(outs[0], global_step=global_step) + + +class FPFcACPolicy(FcACPolicy): + def __init__(self, n_s, n_a, n_w, n_f, n_step, n_fc_wave=128, n_fc_wait=32, n_fc_fp=32, n_lstm=64, name=None): + ACPolicy.__init__(self, n_a, n_s, n_step, 'fpfc', name) + self.n_fc_wave = n_fc_wave + self.n_fc_wait = n_fc_wait + self.n_fc_fp = n_fc_fp + self.n_fc = n_lstm + self.n_w = n_w + self.obs = tf.placeholder(tf.float32, [None, n_s + n_w + n_f]) + with tf.variable_scope(self.name): + # pi and v use separate nets + self.pi = self._build_net('pi') + self.v = self._build_net('v') + + def _build_net(self, out_type): + h0 = fc(ob[:, :self.n_s], out_type + '_fcw', self.n_fc_wave) + h1 = fc(ob[:, (self.n_s + self.n_w):], out_type + '_fcf', self.n_fc_fp) + if self.n_w == 0: + h = tf.concat([h0, h1], 1) + else: + h2 = fc(ob[:, self.n_s: (self.n_s + self.n_w)], + out_type + '_fct', self.n_fc_wait) + h = tf.concat([h0, h1, h2], 1) + h = fc(h, out_type + '_fc', self.n_fc) + return self._build_out_net(h, out_type) + + +class QPolicy: + def __init__(self, n_a, n_s, n_step, policy_name, agent_name): + self.name = policy_name + if agent_name is not None: + # for multi-agent system + self.name += '_' + str(agent_name) + self.n_a = n_a + self.n_s = n_s + self.n_step = n_step + + def forward(self, ob, *_args, **_kwargs): + raise NotImplementedError() + + def _build_fc_net(self, h, n_fc_ls): + for i, n_fc in enumerate(n_fc_ls): + h = fc(h, 'q_fc_%d' % i, n_fc) + q = fc(h, 'q', self.n_a, act=lambda x: x) + return tf.squeeze(q) + + def _build_net(self): + raise NotImplementedError() + + def prepare_loss(self, max_grad_norm, gamma): + self.A = tf.placeholder(tf.int32, [self.n_step]) + self.S1 = tf.placeholder( + tf.float32, [self.n_step, self.n_s + self.n_w]) + self.R = tf.placeholder(tf.float32, [self.n_step]) + self.DONE = tf.placeholder(tf.bool, [self.n_step]) + A_sparse = tf.one_hot(self.A, self.n_a) + + # backward + with tf.variable_scope(self.name + '_q', reuse=True): + q0s = self._build_net(self.S) + q0 = tf.reduce_sum(q0s * A_sparse, axis=1) + with tf.variable_scope(self.name + '_q', reuse=True): + q1s = self._build_net(self.S1) + q1 = tf.reduce_max(q1s, axis=1) + tq = tf.stop_gradient(tf.where(self.DONE, self.R, self.R + gamma * q1)) + self.loss = tf.reduce_mean(tf.square(q0 - tq)) + + wts = tf.trainable_variables(scope=self.name) + grads = tf.gradients(self.loss, wts) + if max_grad_norm > 0: + grads, self.grad_norm = tf.clip_by_global_norm( + grads, max_grad_norm) + self.lr = tf.placeholder(tf.float32, []) + self.optimizer = tf.train.AdamOptimizer(learning_rate=self.lr) + self._train = self.optimizer.apply_gradients(list(zip(grads, wts))) + # monitor training + if self.name.endswith('_0a'): + summaries = [] + summaries.append(tf.summary.scalar( + 'train/%s_loss' % self.name, self.loss)) + summaries.append(tf.summary.scalar('train/%s_q' % + self.name, tf.reduce_mean(q0))) + summaries.append(tf.summary.scalar('train/%s_tq' % + self.name, tf.reduce_mean(tq))) + summaries.append(tf.summary.scalar( + 'train/%s_gradnorm' % self.name, self.grad_norm)) + self.summary = tf.summary.merge(summaries) + + +class DeepQPolicy(QPolicy): + def __init__(self, n_s, n_a, n_w, n_step, n_fc0=128, n_fc=64, name=None): + super().__init__(n_a, n_s, n_step, 'dqn', name) + self.n_fc = n_fc + self.n_fc0 = n_fc0 + self.n_w = n_w + self.S = tf.placeholder(tf.float32, [None, n_s + n_w]) + with tf.variable_scope(self.name + '_q'): + self.qvalues = self._build_net(self.S) + + def _build_net(self, S): + if self.n_w == 0: + h = fc(S, 'q_fcw', self.n_fc0) + else: + h0 = fc(S[:, :self.n_s], 'q_fcw', self.n_fc0) + h1 = fc(S[:, self.n_s:], 'q_fct', self.n_fc0 / 4) + h = tf.concat([h0, h1], 1) + return self._build_fc_net(h, [self.n_fc]) + + def forward(self, sess, ob): + return sess.run(self.qvalues, {self.S: np.array([ob])}) + + def backward(self, sess, obs, acts, next_obs, dones, rs, cur_lr, + summary_writer=None, global_step=None): + if summary_writer is None: + ops = self._train + else: + ops = [self.summary, self._train] + outs = sess.run(ops, + {self.S: obs, + self.A: acts, + self.S1: next_obs, + self.DONE: dones, + self.R: rs, + self.lr: cur_lr}) + if summary_writer is not None: + summary_writer.add_summary(outs[0], global_step=global_step) + + +class LRQPolicy(DeepQPolicy): + def __init__(self, n_s, n_a, n_step, name=None): + QPolicy.__init__(self, n_a, n_s, n_step, 'lr', name) + self.S = tf.compat.v1.placeholder(tf.float32, [None, n_s]) + self.n_w = 0 + with tf.compat.v1.variable_scope(self.name + '_q'): + self.qvalues = self._build_net(self.S) + + def _build_net(self, S): + return self._build_fc_net(S, []) diff --git a/agents/utils.py b/agents/utils.py new file mode 100644 index 0000000..af6b3b9 --- /dev/null +++ b/agents/utils.py @@ -0,0 +1,300 @@ +import numpy as np +import random +import tensorflow as tf + +""" +initializers +""" +DEFAULT_SCALE = np.sqrt(2) +DEFAULT_MODE = 'fan_in' + + +def ortho_init(scale=DEFAULT_SCALE, mode=None): + def _ortho_init(shape, dtype, partition_info=None): + # lasagne ortho init for tf + shape = tuple(shape) + if len(shape) == 2: # fc: in, out + flat_shape = shape + elif (len(shape) == 3) or (len(shape) == 4): # 1d/2dcnn: (in_h), in_w, in_c, out + flat_shape = (np.prod(shape[:-1]), shape[-1]) + a = np.random.standard_normal(flat_shape) + u, _, v = np.linalg.svd(a, full_matrices=False) + q = u if u.shape == flat_shape else v # pick the one with the correct shape + q = q.reshape(shape) + return (scale * q).astype(np.float32) + return _ortho_init + + +def norm_init(scale=DEFAULT_SCALE, mode=DEFAULT_MODE): + def _norm_init(shape, dtype, partition_info=None): + shape = tuple(shape) + if len(shape) == 2: + n_in = shape[0] + elif (len(shape) == 3) or (len(shape) == 4): + n_in = np.prod(shape[:-1]) + a = np.random.standard_normal(shape) + if mode == 'fan_in': + n = n_in + elif mode == 'fan_out': + n = shape[-1] + elif mode == 'fan_avg': + n = 0.5 * (n_in + shape[-1]) + return (scale * a / np.sqrt(n)).astype(np.float32) + + +DEFAULT_METHOD = ortho_init +""" +layers +""" + + +def conv(x, scope, n_out, f_size, stride=1, pad='VALID', f_size_w=None, act=tf.nn.relu, + conv_dim=1, init_scale=DEFAULT_SCALE, init_mode=None, init_method=DEFAULT_METHOD): + with tf.variable_scope(scope): + b = tf.get_variable( + "b", [n_out], initializer=tf.constant_initializer(0.0)) + if conv_dim == 1: + n_c = x.shape[2].value + w = tf.get_variable("w", [f_size, n_c, n_out], + initializer=init_method(init_scale, init_mode)) + z = tf.nn.conv1d(x, w, stride=stride, padding=pad) + b + elif conv_dim == 2: + n_c = x.shape[3].value + if f_size_w is None: + f_size_w = f_size + w = tf.get_variable("w", [f_size, f_size_w, n_c, n_out], + initializer=init_method(init_scale, init_mode)) + z = tf.nn.conv2d( + x, w, strides=[1, stride, stride, 1], padding=pad) + b + return act(z) + + +def fc(x, scope, n_out, act=tf.nn.relu, init_scale=DEFAULT_SCALE, + init_mode=DEFAULT_MODE, init_method=DEFAULT_METHOD): + with tf.compat.v1.variable_scope(scope): + n_in = x.shape[1].value + w = tf.compat.v1.get_variable("w", [n_in, n_out], + initializer=init_method(init_scale, init_mode)) + b = tf.compat.v1.get_variable( + "b", [n_out], initializer=tf.constant_initializer(0.0)) + z = tf.matmul(x, w) + b + return act(z) + + +def batch_to_seq(x): + n_step = x.shape[0].value + if len(x.shape) == 1: + x = tf.expand_dims(x, -1) + return tf.split(axis=0, num_or_size_splits=n_step, value=x) + + +def seq_to_batch(x): + return tf.concat(axis=0, values=x) + + +def lstm(xs, dones, s, scope, init_scale=DEFAULT_SCALE, init_mode=DEFAULT_MODE, + init_method=DEFAULT_METHOD): + xs = batch_to_seq(xs) + # need dones to reset states + dones = batch_to_seq(dones) + n_in = xs[0].shape[1].value + n_out = s.shape[0] // 2 + with tf.compat.v1.variable_scope(scope): + wx = tf.compat.v1.get_variable("wx", [n_in, n_out*4], + initializer=init_method(init_scale, init_mode)) + wh = tf.compat.v1.get_variable("wh", [n_out, n_out*4], + initializer=init_method(init_scale, init_mode)) + b = tf.compat.v1.get_variable( + "b", [n_out*4], initializer=tf.constant_initializer(0.0)) + s = tf.expand_dims(s, 0) + c, h = tf.split(axis=1, num_or_size_splits=2, value=s) + for ind, (x, done) in enumerate(zip(xs, dones)): + c = c * (1-done) + h = h * (1-done) + z = tf.matmul(x, wx) + tf.matmul(h, wh) + b + i, f, o, u = tf.split(axis=1, num_or_size_splits=4, value=z) + i = tf.nn.sigmoid(i) + f = tf.nn.sigmoid(f) + o = tf.nn.sigmoid(o) + u = tf.tanh(u) + c = f*c + i*u + h = o*tf.tanh(c) + xs[ind] = h + s = tf.concat(axis=1, values=[c, h]) + return seq_to_batch(xs), tf.squeeze(s) + + +def test_layers(): + print(tf.__version__) + tf.reset_default_graph() + sess = tf.Session() + n_step = 5 + fc_x = tf.placeholder(tf.float32, [None, 10]) + lstm_x = tf.placeholder(tf.float32, [n_step, 2]) + lstm_done = tf.placeholder(tf.float32, [n_step]) + lstm_s = tf.placeholder(tf.float32, [20]) + conv1_x = tf.placeholder(tf.float32, [None, 8, 1]) + conv2_x = tf.placeholder(tf.float32, [None, 8, 8, 1]) + fc_out = fc(fc_x, 'fc', 10) + lstm_out, lstm_ns = lstm(lstm_x, lstm_done, lstm_s, 'lstm') + conv1_out = conv(conv1_x, 'conv1', 10, 4, conv_dim=1) + conv2_out = conv(conv2_x, 'conv2', 10, 4, conv_dim=2) + sess.run(tf.global_variables_initializer()) + inputs = {'fc': {fc_x: np.random.randn(n_step, 10)}, + 'lstm_done': {lstm_x: np.zeros((n_step, 2)), + lstm_done: np.ones(n_step), + lstm_s: np.random.randn(20)}, + 'lstm': {lstm_x: np.random.randn(n_step, 2), + lstm_done: np.zeros(n_step), + lstm_s: np.random.randn(20)}, + 'conv1': {conv1_x: np.random.randn(n_step, 8, 1)}, + 'conv2': {conv2_x: np.random.randn(n_step, 8, 8, 1)}} + outputs = {'fc': [fc_out], 'lstm_done': [lstm_out, lstm_ns], + 'conv1': [conv1_out], 'conv2': [conv2_out], + 'lstm': [lstm_out, lstm_ns]} + for scope in ['fc', 'lstm', 'conv1', 'conv2']: + print(scope) + wts = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope=scope) + for wt in wts: + wt_val = wt.eval(sess) + print(wt_val.shape) + print(np.mean(wt_val), np.std(wt_val), + np.min(wt_val), np.max(wt_val)) + print('=====================================') + for x_name in inputs: + print(x_name) + out = sess.run(outputs[x_name], inputs[x_name]) + if x_name.startswith('lstm'): + print(out[0]) + print(out[1]) + else: + print(out[0].shape) + + +""" +buffers +""" + + +class TransBuffer: + def reset(self): + self.buffer = [] + + @property + def size(self): + return len(self.buffer) + + def add_transition(self, ob, a, r, *_args, **_kwargs): + raise NotImplementedError() + + def sample_transition(self, *_args, **_kwargs): + raise NotImplementedError() + + +class OnPolicyBuffer(TransBuffer): + def __init__(self, gamma): + self.gamma = gamma + self.reset() + + def reset(self, done=False): + # the done before each step is required + self.obs = [] + self.acts = [] + self.rs = [] + self.vs = [] + self.dones = [done] + + def add_transition(self, ob, a, r, v, done): + self.obs.append(ob) + self.acts.append(a) + self.rs.append(r) + self.vs.append(v) + self.dones.append(done) + + def _add_R_Adv(self, R): + Rs = [] + Advs = [] + # use post-step dones here + for r, v, done in zip(self.rs[::-1], self.vs[::-1], self.dones[:0:-1]): + R = r + self.gamma * R * (1.-done) + Adv = R - v + Rs.append(R) + Advs.append(Adv) + Rs.reverse() + Advs.reverse() + self.Rs = Rs + self.Advs = Advs + + def sample_transition(self, R, discrete=True): + self._add_R_Adv(R) + obs = np.array(self.obs, dtype=np.float32) + if discrete: + acts = np.array(self.acts, dtype=np.int32) + else: + acts = np.array(self.acts, dtype=np.float32) + Rs = np.array(self.Rs, dtype=np.float32) + Advs = np.array(self.Advs, dtype=np.float32) + # use pre-step dones here + dones = np.array(self.dones[:-1], dtype=np.bool) + self.reset(self.dones[-1]) + return obs, acts, dones, Rs, Advs + + +class ReplayBuffer(TransBuffer): + def __init__(self, buffer_size, batch_size): + self.buffer_size = buffer_size + self.batch_size = batch_size + self.cum_size = 0 + self.buffer = [] + + def add_transition(self, ob, a, r, next_ob, done): + experience = (ob, a, r, next_ob, done) + if self.cum_size < self.buffer_size: + self.buffer.append(experience) + else: + ind = int(self.cum_size % self.buffer_size) + self.buffer[ind] = experience + self.cum_size += 1 + + def reset(self): + self.buffer = [] + self.cum_size = 0 + + def sample_transition(self): + # Randomly sample batch_size examples + minibatch = random.sample(self.buffer, self.batch_size) + state_batch = np.asarray([data[0] for data in minibatch]) + action_batch = np.asarray([data[1] for data in minibatch]) + next_state_batch = np.asarray([data[3] for data in minibatch]) + reward_batch = np.asarray([data[2] for data in minibatch]) + done_batch = np.asarray([data[4] for data in minibatch]) + return state_batch, action_batch, next_state_batch, reward_batch, done_batch + + @property + def size(self): + return min(self.buffer_size, self.cum_size) + + +""" +util functions +""" + + +class Scheduler: + def __init__(self, val_init, val_min=0, total_step=0, decay='linear'): + self.val = val_init + self.N = float(total_step) + self.val_min = val_min + self.decay = decay + self.n = 0 + + def get(self, n_step): + self.n += n_step + if self.decay == 'linear': + return max(self.val_min, self.val * (1 - self.n / self.N)) + else: + return self.val + + +if __name__ == '__main__': + test_layers()