class QNetwork(Network):

    def __init__(self, conf):
        super(QNetwork, self).__init__(conf)
                
        with tf.variable_scope(self.name):
            self.target_ph = tf.placeholder('float32', [None], name='target')
            encoded_state = self._build_encoder()

            self.loss = self._build_q_head(encoded_state)
            self._build_gradient_ops(self.loss)


    def _build_q_head(self, input_state):
        self.w_out, self.b_out, self.output_layer = layers.fc('fc_out', input_state, self.num_actions, activation="linear")
        self.q_selected_action = tf.reduce_sum(self.output_layer * self.selected_action_ph, axis=1)

        diff = tf.subtract(self.target_ph, self.q_selected_action)
        return self._value_function_loss(diff)
 class DuelingNetwork(QNetwork):
 
    def _build_q_head(self, input_state):
        self.w_value, self.b_value, self.value = layers.fc('fc_value', input_state, 1, activation='linear')
        self.w_adv, self.b_adv, self.advantage = layers.fc('fc_advantage', input_state, self.num_actions, activation='linear')

        self.output_layer = (
            self.value + self.advantage
            - tf.reduce_mean(
                self.advantage,
                axis=1,
                keep_dims=True
            )
        )

        q_selected_action = tf.reduce_sum(self.output_layer * self.selected_action_ph, axis=1)
        diff = tf.subtract(self.target_ph, q_selected_action)
        return self._value_function_loss(diff) class ContinuousPolicyNetwork(ContinuousPolicyValueNetwork):
    def __init__(self, conf):
        super(ContinuousPolicyNetwork, self).__init__(conf, use_value_head=False) class ContinuousPolicyValueNetwork(PolicyValueNetwork):
    '''
    Shared policy-value network with polciy head parametrizing
    multivariate normal with diagonal covariance
    '''
    def __init__(self, conf, **kwargs):
        self.action_space = conf['args'].action_space
        self.use_state_dependent_std = False
        super(ContinuousPolicyValueNetwork, self).__init__(conf, **kwargs)

    def _build_policy_head(self, input_state):
        self.adv_actor_ph = tf.placeholder("float", [None], name='advantage')       
        self.w_mu, self.b_mu, self.mu = layers.fc(
            'mean', input_state, self.num_actions, activation='linear')
        self.sigma, dist_params = self._build_sigma(input_state)

        self.dist = DiagNormal(dist_params)
        self.log_output_selected_action = self.dist.log_likelihood(self.selected_action_ph)
        
        self.output_layer_entropy = self.dist.entropy()
        self.entropy = tf.reduce_sum(self.output_layer_entropy)

        self.actor_objective = -tf.reduce_sum(
            self.log_output_selected_action * self.adv_actor_ph
            + self.beta * self.output_layer_entropy
        )
        self.sample_action = self.dist.sample()

        return self.actor_objective

    def _build_sigma(self, input_state):
        if self.use_state_dependent_std:
            self.w_sigma2, self.b_sigma2, self.sigma_hat = layers.fc(
                'std2', input_state, self.num_actions, activation='linear')
            self.sigma2 = tf.log(1+tf.exp(self.sigma_hat))
            sigma = tf.sqrt(self.sigma2 + 1e-8)
            return sigma, tf.concat([self.mu, sigma], 1)
        else:
            self.log_sigma = tf.get_variable('log_sigma', self.mu.get_shape().as_list()[1],
                dtype=tf.float32, initializer=tf.random_uniform_initializer(-2, -1))
            sigma = tf.expand_dims(tf.exp(self.log_sigma), 0)
            tiled_sigma = tf.tile(sigma, [tf.shape(self.mu)[0], 1])
            return sigma, tf.concat([self.mu, tiled_sigma], 1)

    def get_action(self, session, state, lstm_state=None):
        feed_dict = {self.input_ph: [state]}
        if lstm_state is not None:
            feed_dict[self.step_size] = [1]
            feed_dict[self.initial_lstm_state] = lstm_state

            action, lstm_state, params = session.run([
                self.sample_action,
                self.lstm_state,
                self.dist.params()], feed_dict=feed_dict)

            return action[0], params, lstm_state
        else:
            action, params = session.run([
                self.sample_action,
                self.dist.params()], feed_dict=feed_dict)

            return action[0], params[0]

    def get_action_and_value(self, session, state, lstm_state=None):
        feed_dict = {self.input_ph: [state]}
        if lstm_state is not None:
            feed_dict[self.step_size] = [1]
            feed_dict[self.initial_lstm_state] = lstm_state

            action, v, lstm_state, params = session.run([
                self.sample_action,
                self.output_layer_v,
                self.lstm_state,
                self.dist.params()], feed_dict=feed_dict)

            return action[0], v[0, 0], params[0], lstm_state
        else:
            action, v, params = session.run([
                self.sample_action,
                self.output_layer_v,
                self.dist.params()], feed_dict=feed_dict)

            return action[0], v[0, 0], params[0] class PolicyNetwork(PolicyValueNetwork):
    def __init__(self, conf,):
        super(PolicyNetwork, self).__init__(conf, use_value_head=False) class PolicyValueNetwork(Network):
 
    def __init__(self, conf, use_policy_head=True, use_value_head=True):
        super(PolicyValueNetwork, self).__init__(conf)
        self.beta = conf['args'].entropy_regularisation_strength
        self.share_encoder_weights = conf['args'].share_encoder_weights

        encoded_state = self._build_encoder()
                
        with tf.variable_scope(self.name):
            self.loss = 0.0
            if use_policy_head:
                self.loss += self._build_policy_head(encoded_state)

            if use_value_head:
                if not self.share_encoder_weights:
                    with tf.variable_scope('value_encoder'):
                        encoded_state = self._build_encoder()

                self.loss += 0.5 * self._build_value_head(encoded_state)
            self.loss *= tf.cast(tf.shape(self.input_ph)[0], tf.float32) / self.max_local_steps
            self._build_gradient_ops(self.loss)

    def _build_policy_head(self, input_state):
        self.adv_actor_ph = tf.placeholder("float", [None], name='advantage')       
        self.wpi, self.bpi, self.logits = layers.fc(
            'logits', input_state, self.num_actions, activation='linear')
        self.dist = Discrete(self.logits)

        self.output_layer_entropy = self.dist.entropy()
        self.entropy = tf.reduce_sum(self.output_layer_entropy)
        
        self.output_layer_pi = self.dist.probs
        self.log_output_layer_pi = self.dist.log_probs
        self.log_output_selected_action = self.dist.log_likelihood(self.selected_action_ph)
        self.sample_action = self.dist.sample()

        self.actor_objective = -tf.reduce_sum(
            self.log_output_selected_action * self.adv_actor_ph
            + self.beta * self.output_layer_entropy
        )
        return self.actor_objective

    def _build_value_head(self, input_state):
        self.critic_target_ph = tf.placeholder('float32', [None], name='target')
        self.wv, self.bv, self.output_layer_v = layers.fc(
            'fc_value4', input_state, 1, activation='linear')
        # Advantage critic
        self.adv_critic = tf.subtract(self.critic_target_ph, tf.reshape(self.output_layer_v, [-1]))
        # Critic loss
        self.critic_loss = self._value_function_loss(self.adv_critic)
        return self.critic_loss

    def get_action(self, session, state, lstm_state=None):
        feed_dict = {self.input_ph: [state]}
        if lstm_state is not None:
            feed_dict[self.step_size] = [1]
            feed_dict[self.initial_lstm_state] = lstm_state

            action, logits, lstm_state = session.run([
                self.sample_action,
                self.logits,
                self.lstm_state], feed_dict=feed_dict)

            return action, logits[0], lstm_state
        else:
            action, logits = session.run([
                self.sample_action,
                self.logits], feed_dict=feed_dict)

            return action, logits[0]


    def get_action_and_value(self, session, state, lstm_state=None):
        feed_dict = {self.input_ph: [state]}
        if lstm_state is not None:
            feed_dict[self.step_size] = [1]
            feed_dict[self.initial_lstm_state] = lstm_state

            action, logits, v, lstm_state = session.run([
                self.sample_action,
                self.logits,
                self.output_layer_v,
                self.lstm_state], feed_dict=feed_dict)

            return action, v[0, 0], logits[0], lstm_state

        else:
            action, logits, v = session.run([
                self.sample_action,
                self.logits,
                self.output_layer_v], feed_dict=feed_dict)

            return action, v[0, 0], logits[0] class PolicyRepeatNetwork(PolicyValueNetwork):
 
    def __init__(self, conf):
        '''
        Extends action space to parametrize a discrete distribution over repetion lengths
        for each original action
        '''
        self.max_repeat = 5
        super(PolicyRepeatNetwork, self).__init__(conf)
        

    def _build_policy_head(self, input_state):
        self.adv_actor_ph = tf.placeholder("float", [None], name='advantage')

        self.wpi, self.bpi, self.output_layer_pi, self.log_output_layer_pi = layers.softmax_and_log_softmax(
            'action_policy', input_state, self.num_actions)

        self.w_ar, self.b_ar, self.action_repeat_probs, self.log_action_repeat_probs = layers.softmax_and_log_softmax(
            'repeat_policy', input_state, self.max_repeat)

        self.selected_repeat = tf.placeholder(tf.int32, [None], name='selected_repeat_placeholder')
        self.selected_repeat_onehot = tf.one_hot(self.selected_repeat, self.max_repeat)

        self.selected_repeat_prob = tf.reduce_sum(self.action_repeat_probs * self.selected_repeat_onehot, 1)
        self.log_selected_repeat_prob = tf.reduce_sum(self.log_action_repeat_probs * self.selected_repeat_onehot, 1)

        # Entropy: ∑_a[-p_a ln p_a]
        self.output_layer_entropy = tf.reduce_sum(
            - 1.0 * tf.multiply(
                tf.expand_dims(self.output_layer_pi, 2) * tf.expand_dims(self.action_repeat_probs, 1),
                tf.expand_dims(self.log_output_layer_pi, 2) + tf.expand_dims(self.log_action_repeat_probs, 1)
            ), axis=[1, 2])

        self.entropy = tf.reduce_sum(self.output_layer_entropy)

        self.log_output_selected_action = tf.reduce_sum(
            self.log_output_layer_pi*self.selected_action_ph,
            axis=1) + self.log_selected_repeat_prob

        self.actor_objective = -tf.reduce_sum(
            self.log_output_selected_action * self.adv_actor_ph
            + self.beta * self.output_layer_entropy)

        return self.actor_objective class SequencePolicyVNetwork(PolicyValueNetwork):

    def __init__(self, conf):
        '''Uses lstm decoder to output action sequences'''
        self.max_decoder_steps = conf['args'].max_decoder_steps
        self.max_local_steps = conf['args'].max_local_steps
        super(SequencePolicyVNetwork, self).__init__(conf)

        
    def _build_policy_head(self, input_state):
        self.adv_actor_ph = tf.placeholder("float", [self.batch_size], name='advantage')

        with tf.variable_scope(self.name+'/lstm_decoder') as vs:
            self.action_outputs = tf.placeholder(tf.float32, [self.batch_size, None, self.num_actions+1], name='action_outputs')
            self.action_inputs = tf.placeholder(tf.float32, [self.batch_size, None, self.num_actions+1], name='action_inputs')
                
            self.decoder_seq_lengths = tf.placeholder(tf.int32, [self.batch_size], name='decoder_seq_lengths')
            self.allowed_actions = tf.placeholder(tf.float32, [self.batch_size, None, self.num_actions+1], name='allowed_actions')
            self.use_fixed_action = tf.placeholder(tf.bool, name='use_fixed_action')
            self.temperature = tf.placeholder(tf.float32, name='temperature')

            self.decoder_hidden_state_size = input_state.get_shape().as_list()[-1]
            self.decoder_lstm_cell = CustomBasicLSTMCell(self.decoder_hidden_state_size, forget_bias=1.0)
            self.decoder_initial_state = tf.placeholder(tf.float32, [self.batch_size, 2*self.decoder_hidden_state_size], name='decoder_initial_state')

            self.network_state = tf.concat(axis=1, values=[
                tf.zeros_like(input_state), input_state
                # input_state, tf.zeros_like(input_state)
            ])

            self.W_actions = tf.get_variable('W_actions', shape=[self.decoder_hidden_state_size, self.num_actions+1], dtype='float32', initializer=tf.contrib.layers.xavier_initializer())
            self.b_actions = tf.get_variable('b_actions', shape=[self.num_actions+1], dtype='float32', initializer=tf.zeros_initializer())


            self.decoder_state, self.logits, self.actions = decoder(
                self.action_inputs,
                self.network_state,
                self.decoder_lstm_cell,
                self.decoder_seq_lengths,
                self.W_actions,
                self.b_actions,
                self.max_decoder_steps,
                vs,
                self.use_fixed_action,
                self.action_outputs,
                loop_function=loop_gumbel_softmax(self.temperature),
            )

            self.decoder_trainable_variables = [
                v for v in tf.trainable_variables()
                if v.name.startswith(vs.name)
            ]

        print 'Decoder out: s,l,a=', self.decoder_state.get_shape(), self.logits.get_shape(), self.actions.get_shape()


        #mask softmax by allowed actions
        exp_logits = tf.exp(self.logits) * self.allowed_actions
        Z = tf.expand_dims(tf.reduce_sum(exp_logits, 2), 2)
        self.action_probs = exp_logits / Z
        log_action_probs = self.logits - tf.log(Z)

        sequence_probs = tf.reduce_prod(tf.reduce_sum(self.action_probs * self.action_outputs, 2), 1)
        log_sequence_probs = tf.reduce_sum(tf.reduce_sum(log_action_probs * self.action_outputs, 2), 1)

        # ∏a_i * ∑ log a_i
        self.output_layer_entropy = - tf.reduce_sum(tf.stop_gradient(1 + log_sequence_probs) * log_sequence_probs)
        self.entropy = - tf.reduce_sum(log_sequence_probs)

        print 'sp, lsp:', sequence_probs.get_shape(), log_sequence_probs.get_shape()


        self.actor_advantage_term = tf.reduce_sum(log_sequence_probs[:self.max_local_steps] * self.adv_actor_ph)
        self.actor_entropy_term = self.beta * self.output_layer_entropy
        self.actor_objective = - (
            self.actor_advantage_term
            + self.actor_entropy_term
        )

        return self.actor_objective class SharedCounter(object):
    def __init__(self, initval=0):
        self.val = RawValue('i', initval)
        self.last_step_update_target = RawValue('i', initval)
        self.lock = Lock()

    def increment(self, elapsed_steps=None):
        self.val.value += 1
        if ((elapsed_steps is not None) 
            and ((self.val.value - self.last_step_update_target.value) 
                >= elapsed_steps)):
            self.last_step_update_target.value = self.val.value
            return self.val.value, True
        else:
            return self.val.value, False

    def set_value(self, value):
        self.lock.acquire()
        self.val.value = value
        self.lock.release()

    def value(self):
        return self.val.value
 class SharedVars(object):
    def __init__(self, params, opt_type=None, lr=0, step=0):
        self.var_shapes = [
            var.get_shape().as_list()
            for var in params]
        self.size = sum([np.prod(shape) for shape in self.var_shapes])
        self.step = RawValue(ctypes.c_int, step)

        if opt_type == 'adam':
            self.ms = self.malloc_contiguous(self.size)
            self.vs = self.malloc_contiguous(self.size)
            self.lr = RawValue(ctypes.c_float, lr)
        elif opt_type == 'adamax':
            self.ms = self.malloc_contiguous(self.size)
            self.vs = self.malloc_contiguous(self.size)
            self.lr = RawValue(ctypes.c_float, lr)
        elif opt_type == 'rmsprop':
            self.vars = self.malloc_contiguous(self.size, np.ones(self.size, dtype=np.float))
        elif opt_type == 'momentum':
            self.vars = self.malloc_contiguous(self.size)
        else:
            self.vars = self.malloc_contiguous(self.size)

            
    def malloc_contiguous(self, size, initial_val=None):
        if initial_val is None:
            return RawArray(ctypes.c_float, size)
        else:
            return RawArray(ctypes.c_float, initial_val)
 class SharedFlags(object):
    def __init__(self, num_actors):
        self.updated = RawArray(ctypes.c_int, num_actors)
 class Barrier:
    def __init__(self, n):
        self.n = n
        self.counter = SharedCounter(0)
        self.barrier = Semaphore(0)

    def wait(self):
        with self.counter.lock:
            self.counter.val.value += 1
            if self.counter.val.value == self.n:
                self.barrier.release()

        self.barrier.acquire()
        self.barrier.release()
 class A3CLearner(BaseA3CLearner):
    def __init__(self, args):
        super(A3CLearner, self).__init__(args)

        conf_learning = {'name': 'local_learning_{}'.format(self.actor_id),
                         'input_shape': self.input_shape,
                         'num_act': self.num_actions,
                         'args': args}

        self.local_network = args.network(conf_learning)
        self.reset_hidden_state()

        if self.is_master():
            var_list = self.local_network.params
            self.saver = tf.train.Saver(var_list=var_list, max_to_keep=3,
                                        keep_checkpoint_every_n_hours=2)


    def choose_next_action(self, state):
        return self.local_network.get_action_and_value(self.session, state) class A3CLSTMLearner(BaseA3CLearner):
    def __init__(self, args):
        super(A3CLSTMLearner, self).__init__(args)

        conf_learning = {'name': 'local_learning_{}'.format(self.actor_id),
                         'input_shape': self.input_shape,
                         'num_act': self.num_actions,
                         'args': args}

        self.local_network = args.network(conf_learning)
        self.reset_hidden_state()

        if self.is_master():
            var_list = self.local_network.params
            self.saver = tf.train.Saver(var_list=var_list, max_to_keep=3,
                                        keep_checkpoint_every_n_hours=2)


    def reset_hidden_state(self):
        self.lstm_state_out = np.zeros([1, 2*self.local_network.hidden_state_size])
        # self.lstm_state_out = tf.contrib.rnn.LSTMStateTuple(np.zeros([1, 256]),
        #                                                     np.zeros([1, 256]))


    def set_local_lstm_state(self):
        self.local_lstm_state = np.copy(self.lstm_state_out)


    def choose_next_action(self, state):
        action, v, dist, self.lstm_state_out = self.local_network.get_action_and_value(
            self.session, state, lstm_state=self.lstm_state_out)
        return action, v, dist


    def bootstrap_value(self, state, episode_over):
        if episode_over:
            R = 0
        else:
            R = self.session.run(
                self.local_network.output_layer_v,
                feed_dict={
                    self.local_network.input_ph:[state],
                    self.local_network.step_size: [1],
                    self.local_network.initial_lstm_state: self.lstm_state_out,
                }
            )[0][0]

        return R


    def apply_update(self, states, actions, targets, advantages):
        feed_dict={
            self.local_network.input_ph: states,
            self.local_network.selected_action_ph: actions,
            self.local_network.critic_target_ph: targets,
            self.local_network.adv_actor_ph: advantages,
            self.local_network.step_size : [len(states)],
            self.local_network.initial_lstm_state: self.local_lstm_state,
        }
        grads, entropy = self.session.run(
            [self.local_network.get_gradients, self.local_network.entropy],
            feed_dict=feed_dict)

        self.apply_gradients_to_shared_memory_vars(grads)
        return entropy class ActionSequenceA3CLearner(BaseA3CLearner):
    def __init__(self, args):

        super(ActionSequenceA3CLearner, self).__init__(args)
        
        # Shared mem vars
        self.learning_vars = args.learning_vars

        conf_learning = {'name': 'local_learning_{}'.format(self.actor_id),
                         'input_shape': args.input_shape,
                         'num_act': self.num_actions,
                         'args': args}
        
        self.local_network = SequencePolicyVNetwork(conf_learning)
        self.reset_hidden_state()
            
        if self.is_master():
            var_list = self.local_network.params
            self.saver = tf.train.Saver(var_list=var_list, max_to_keep=3, 
                                        keep_checkpoint_every_n_hours=2)


    def sample_action_sequence(self, state):
        allowed_actions = np.ones((self.local_network.max_decoder_steps, self.local_network.num_actions+1))
        allowed_actions[0, -1] = 0

        action_inputs = np.zeros((1,self.local_network.max_decoder_steps,self.num_actions+1))
        action_inputs[0, 0, -1] = 1

        actions, value = self.session.run(
            [
                self.local_network.actions,
                self.local_network.output_layer_v,
            ],
            feed_dict={
                self.local_network.input_ph:              [state],
                self.local_network.decoder_seq_lengths:   [self.local_network.max_decoder_steps],
                self.local_network.allowed_actions:       [allowed_actions],
                self.local_network.use_fixed_action:      False,
                self.local_network.temperature:           1.0,
                self.local_network.action_outputs:        np.zeros((1,self.local_network.max_decoder_steps,self.num_actions+1)),
                self.local_network.action_inputs:         action_inputs,
                self.local_network.decoder_initial_state: np.zeros((1, self.local_network.decoder_hidden_state_size*2)),
            }
        )

        return actions[0], value[0, 0]


    def train(self):
        """ Main actor learner loop for advantage actor critic learning. """
        logger.debug("Actor {} resuming at Step {}".format(self.actor_id, 
            self.global_step.value()))

        s = self.emulator.get_initial_state()
        total_episode_reward = 0

        s_batch = []
        a_batch = []
        y_batch = []
        adv_batch = []
        seq_len_batch = []
        
        reset_game = False
        episode_over = False
        start_time = time.time()
        steps_at_last_reward = self.local_step
        
        while (self.global_step.value() < self.max_global_steps):
            # Sync local learning net with shared mem
            self.sync_net_with_shared_memory(self.local_network, self.learning_vars)
            self.save_vars()

            local_step_start = self.local_step 
            
            rewards = []
            states = []
            actions = []
            values = []
            seq_lengths = []
            
            while not (episode_over 
                or (self.local_step - local_step_start 
                    == self.max_local_steps)):
                
                # Choose next action and execute it
                action_sequence, readout_v_t = self.sample_action_sequence(s)
                # if self.is_master() and (self.local_step % 100 == 0):
                #     logger.debug("pi={}, V={}".format(readout_pi_t, readout_v_t))
                
                acc_reward = 0.0
                length = 0

                for action in action_sequence:
                    length += 1
                    a = np.argmax(action)
                    if a == self.num_actions or episode_over:
                        break

                    new_s, reward, episode_over = self.emulator.next(action[:self.num_actions])
                    acc_reward += reward

                reward = acc_reward
                if reward != 0.0:
                    steps_at_last_reward = self.local_step


                total_episode_reward += reward
                # Rescale or clip immediate reward
                reward = self.rescale_reward(reward)
                
                rewards.append(reward)
                seq_lengths.append(length)
                states.append(s)
                actions.append(action_sequence)
                values.append(readout_v_t)
                

                s = new_s
                self.local_step += 1
                self.global_step.increment()

                if self.local_step % 1000 == 0:
                    pass
                    # Pdb().set_trace()
                
            
            # Calculate the value offered by critic in the new state.
            if episode_over:
                R = 0
            else:
                R = self.session.run(
                    self.local_network.output_layer_v,
                    feed_dict={self.local_network.input_ph:[new_s]})[0][0]
                            
             
            sel_actions = []
            for i in reversed(xrange(len(states))):
                R = rewards[i] + self.gamma * R

                y_batch.append(R)
                a_batch.append(actions[i])
                s_batch.append(states[i])
                adv_batch.append(R - values[i])
                seq_len_batch.append(seq_lengths[i])
                
                sel_actions.append(np.argmax(actions[i]))
            
            padded_output_sequences = np.array([
                np.vstack([seq[:length, :], np.zeros((max(seq_len_batch)-length, self.num_actions+1))])
                for length, seq in zip(seq_len_batch, a_batch)
            ])

            go_input = np.zeros((len(s_batch), 1, self.num_actions+1))
            go_input[:,:,self.num_actions] = 1
            padded_input_sequences = np.hstack([go_input, padded_output_sequences[:,:-1,:]])

            print 'Sequence lengths:', seq_lengths
            print 'Actions:', [np.argmax(a) for a in a_batch[0]]

            allowed_actions = np.ones((len(s_batch), max(seq_len_batch), self.num_actions+1))
            allowed_actions[:, 0, -1] = 0 #empty sequence is not a valid action

            feed_dict={
                self.local_network.input_ph:              s_batch, 
                self.local_network.critic_target_ph:      y_batch,
                self.local_network.adv_actor_ph:          adv_batch,
                self.local_network.decoder_initial_state: np.zeros((len(s_batch), self.local_network.decoder_hidden_state_size*2)),
                self.local_network.action_inputs:         padded_input_sequences,
                self.local_network.action_outputs:        padded_output_sequences,
                self.local_network.allowed_actions:       allowed_actions,
                self.local_network.use_fixed_action:      True,
                self.local_network.decoder_seq_lengths:   seq_lengths,
                self.local_network.temperature:           1.0,
            }
            entropy, grads = self.session.run(
                [
                    self.local_network.entropy,
                    # self.local_network.adv_critic,
                    # self.local_network.output_layer_v,
                    self.local_network.get_gradients
                ],
                feed_dict=feed_dict)

            print 'Entropy:', entropy #, 'Adv:', advantage #, 'Value:', value
            self.apply_gradients_to_shared_memory_vars(grads)     
            
            s_batch = []
            a_batch = []
            y_batch = []          
            adv_batch = []
            seq_len_batch = []

            
            # prevent the agent from getting stuck
            if (self.local_step - steps_at_last_reward > 5000
                or (self.emulator.get_lives() == 0
                    and self.emulator.game not in ONE_LIFE_GAMES)):

                steps_at_last_reward = self.local_step
                episode_over = True
                reset_game = True


            # Start a new game on reaching terminal state
            if episode_over:
                elapsed_time = time.time() - start_time
                global_t = self.global_step.value()
                steps_per_sec = global_t / elapsed_time
                perf = "{:.0f}".format(steps_per_sec)
                logger.info("T{} / STEP {} / REWARD {} / {} STEPS/s, Actions {}".format(self.actor_id, global_t, total_episode_reward, perf, sel_actions))
                
                self.log_summary(total_episode_reward, entropy)

                episode_over = False
                total_episode_reward = 0
                steps_at_last_reward = self.local_step

                if reset_game or self.emulator.game in ONE_LIFE_GAMES:
                    s = self.emulator.get_initial_state()
                    reset_game = False class ARA3CLearner(BaseA3CLearner):
    def __init__(self, args):
        super(ARA3CLearner, self).__init__(args)

        conf_learning = {'name': 'local_learning_{}'.format(self.actor_id),
                         'input_shape': self.input_shape,
                         'num_act': self.num_actions,
                         'args': args}

        self.local_network = PolicyRepeatNetwork(conf_learning)
        self.reset_hidden_state()

        if self.is_master():
            var_list = self.local_network.params
            self.saver = tf.train.Saver(var_list=var_list, max_to_keep=3,
                                        keep_checkpoint_every_n_hours=2)


    def choose_next_action(self, state):
        network_output_v, network_output_pi, action_repeat_probs = self.session.run(
            [
                self.local_network.output_layer_v,
                self.local_network.output_layer_pi,
                self.local_network.action_repeat_probs,
            ],
            feed_dict={
                self.local_network.input_ph: [state],
            })

        network_output_pi = network_output_pi.reshape(-1)
        network_output_v = np.asscalar(network_output_v)

        action_index = self.sample_policy_action(network_output_pi)
        new_action = np.zeros([self.num_actions])
        new_action[action_index] = 1

        action_repeat = 1 + self.sample_policy_action(action_repeat_probs[0])

        return new_action, network_output_v, network_output_pi, action_repeat


    def train(self):
        """ Main actor learner loop for advantage actor critic learning. """
        logger.debug("Actor {} resuming at Step {}".format(self.actor_id, 
            self.global_step.value()))

        s = self.emulator.get_initial_state()
        steps_at_last_reward = self.local_step
        total_episode_reward = 0.0
        mean_entropy = 0.0
        episode_start_step = 0
        
        while (self.global_step.value() < self.max_global_steps):
            # Sync local learning net with shared mem
            self.sync_net_with_shared_memory(self.local_network, self.learning_vars)
            self.save_vars()

            local_step_start = self.local_step 
            
            reset_game = False
            episode_over = False

            rewards        = list()
            states         = list()
            actions        = list()
            values         = list()
            s_batch        = list()
            a_batch        = list()
            y_batch        = list()
            ar_batch       = list()
            adv_batch      = list()
            action_repeats = list()
            
            while not (episode_over 
                or (self.local_step - local_step_start 
                    == self.max_local_steps)):
                
                # Choose next action and execute it
                a, readout_v_t, readout_pi_t, action_repeat = self.choose_next_action(s)
                
                if self.is_master() and (self.local_step % 100 == 0):
                    logger.debug("π_a={:.4f} / V={:.4f} repeat={}".format(
                        readout_pi_t[a.argmax()], readout_v_t, action_repeat))

                reward = 0.0
                for _ in range(action_repeat):
                    new_s, reward_i, episode_over = self.emulator.next(a)
                    reward += reward_i

                if reward != 0.0:
                    steps_at_last_reward = self.local_step


                total_episode_reward += reward
                # Rescale or clip immediate reward
                reward = self.rescale_reward(reward)
                
                rewards.append(reward)
                states.append(s)
                actions.append(a)
                values.append(readout_v_t)
                action_repeats.append(action_repeat)
                
                s = new_s
                self.local_step += 1
                self.global_step.increment()
                
            
            # Calculate the value offered by critic in the new state.
            if episode_over:
                R = 0
            else:
                R = self.session.run(
                    self.local_network.output_layer_v,
                    feed_dict={self.local_network.input_ph:[new_s]})[0][0]
                            
             
            sel_actions = []
            for i in reversed(xrange(len(states))):
                R = rewards[i] + self.gamma * R

                y_batch.append(R)
                a_batch.append(actions[i])
                s_batch.append(states[i])
                adv_batch.append(R - values[i])
                ar_batch.append(action_repeats[i])
                
                sel_actions.append(np.argmax(actions[i]))
                

            # Compute gradients on the local policy/V network and apply them to shared memory  
            feed_dict={
                self.local_network.input_ph: s_batch, 
                self.local_network.critic_target_ph: y_batch,
                self.local_network.selected_action_ph: a_batch,
                self.local_network.adv_actor_ph: adv_batch,
                self.local_network.selected_repeat: ar_batch,
            }
            grads, entropy = self.session.run(
                [self.local_network.get_gradients, self.local_network.entropy],
                feed_dict=feed_dict)

            self.apply_gradients_to_shared_memory_vars(grads)     

            delta_old = local_step_start - episode_start_step
            delta_new = self.local_step -  local_step_start
            mean_entropy = (mean_entropy*delta_old + entropy*delta_new) / (delta_old + delta_new)

            s, mean_entropy, episode_start_step, total_episode_reward, steps_at_last_reward = self.prepare_state(
                s, mean_entropy, episode_start_step, total_episode_reward, steps_at_last_reward, sel_actions, episode_over) class NStepQLearner(ValueBasedLearner):

    def train(self):
        """ Main actor learner loop for n-step Q learning. """
        logger.debug("Actor {} resuming at Step {}, {}".format(self.actor_id,
            self.global_step.value(), time.ctime()))

        episode_rewards = deque(maxlen=100)
        episode_max_q = deque(maxlen=100)

        while (self.global_step.value() < self.max_global_steps):

            s = self.emulator.get_initial_state()
            total_episode_reward = 0.0
            episode_over = False
            episode_start_step = self.local_step
            self.local_episode += 1
            exec_update_target = False

            while not episode_over:


                # Sync local learning net with shared mem
                self.sync_net_with_shared_memory(self.local_network, self.learning_vars)
                self.save_vars()

                rewards = list()
                states =  list()
                actions = list()
                max_q = list()
                local_step_start = self.local_step

                while self.local_step - local_step_start < self.max_local_steps and not episode_over:

                    # Choose next action and execute it
                    a, readout_t = self.choose_next_action(s)

                    new_s, reward, episode_over = self.emulator.next(a)

                    #if reward != 0.0:
                    #    steps_at_last_reward = self.local_step

                    total_episode_reward += reward
                    # Rescale or clip immediate reward
                    reward = self.rescale_reward(reward)

                    rewards.append(reward)
                    states.append(s)
                    actions.append(a)

                    s = new_s
                    self.local_step += 1
                    self.global_step.increment()
                    max_q.append(np.max(readout_t))

                    global_step, update_target = self.global_step.increment(
                        self.q_target_update_steps)

                    if update_target:
                        update_target = False
                        exec_update_target = True

                    self.local_network.global_step = global_step

                R = self.bootstrap_value(s, episode_over)
                targets = self.compute_targets(rewards, R)
                # Compute gradients on the local Q network
                self.apply_update(states, actions, targets)

            episode_max_q.append(np.max(max_q))
            episode_rewards.append(total_episode_reward)
            elapsed_time = time.time() - self.start_time
            steps_per_sec = self.global_step.value() / elapsed_time
            perf = "{:.0f}".format(steps_per_sec)

            logger.info("T{} / EPISODE {} / STEP {}k / MEAN REWARD {:.1f} / MEAN MAX Q {:.1f} / {} STEPS/s".format(
                self.actor_id,
                self.local_episode,
                self.global_step.value()/1000,
                np.array(episode_rewards).mean(),
                np.array(episode_max_q).mean(),
                perf))

            self.log_summary(
                total_episode_reward,
                np.array(max_q).mean(),
                self.epsilon)

            if exec_update_target:
                self.update_target()
                exec_update_target = False

            # Sync local tensorflow target network params with shared target network params
            if self.target_update_flags.updated[self.actor_id] == 1:
                self.sync_net_with_shared_memory(self.target_network, self.target_vars)
                self.target_update_flags.updated[self.actor_id] = 0
 class DuelingLearner(NStepQLearner):
    def __init__(self, args):
        super(DuelingLearner, self).__init__(args, network_type=DuelingNetwork)
 class OneStepSARSALearner(ValueBasedLearner):

    def generate_final_epsilon(self):
        return 0.1

    def train(self):
        """ Main actor learner loop for 1-step SARSA learning. """
        logger.debug("Actor {} resuming at Step {}, {}".format(self.actor_id,
            self.global_step.value(), time.ctime()))

        s = self.emulator.get_initial_state()

        states =  list()
        actions = list()
        targets = list()

        steps_at_last_reward = self.local_step
        exec_update_target = False
        total_episode_reward = 0
        episode_ave_max_q = 0
        episode_over = False
        qmax_down = 0
        qmax_up = 0
        prev_qmax = -10*6
        low_qmax = 0
        ep_t = 0

        # Choose initial action
        a, readout_t = self.choose_next_action(s)

        while (self.global_step.value() < self.max_global_steps):
            s_prime, reward, episode_over = self.emulator.next(a)
            if reward != 0.0:
                steps_at_last_reward = self.local_step

            ep_t += 1
            episode_ave_max_q += np.max(readout_t)

            actions.append(a)
            states.append(s)
            s = s_prime

            total_episode_reward += reward
            reward = self.rescale_reward(reward)
            if episode_over:
                y = reward
            else:
                # Choose action that we will execute in the next step
                a_prime, readout_t = self.choose_next_action(s_prime)
                q_prime = readout_t[a_prime.argmax()]
                # Q_target in the new state for the next step action
                q_prime = self.session.run(
                    self.target_network.output_layer,
                    feed_dict={self.target_network.input_ph: [s_prime]}
                )[0][a_prime.argmax()]

                y = reward + self.gamma * q_prime
                a = a_prime

            targets.append(y)

            self.local_step += 1
            global_step, update_target = self.global_step.increment(
                self.q_target_update_steps)

            # Compute grads and asynchronously apply them to shared memory
            if ((self.local_step % self.grads_update_steps == 0)
                or episode_over):

                # Compute gradients on the local Q network
                self.apply_update(states, actions, targets)
                self.sync_net_with_shared_memory(self.local_network, self.learning_vars)
                self.save_vars()

                states =  list()
                actions = list()
                targets = list()

            # Copy shared learning network params to shared target network params
            if update_target:
                self.update_target()

            # Sync local tensorflow target network params with shared target network params
            if self.target_update_flags.updated[self.actor_id] == 1:
                self.sync_net_with_shared_memory(self.target_network, self.target_vars)
                self.target_update_flags.updated[self.actor_id] = 0

            s, total_episode_reward, steps_at_last_reward, ep_t, episode_ave_max_q, episode_over = \
                self.prepare_state(s, total_episode_reward, steps_at_last_reward, ep_t, episode_ave_max_q, episode_over)
 class PseudoCountA3CLearner(A3CLearner, A3CDensityModelMixin):
    """
    Attempt at replicating the A3C+ model from the paper 'Unifying Count-Based Exploration and Intrinsic Motivation' (https://arxiv.org/abs/1606.01868)
    """
    def __init__(self, args):
        super(PseudoCountA3CLearner, self).__init__(args)
        self._init_density_model(args)

    def train(self):
        self._train() class PseudoCountA3CLSTMLearner(A3CLSTMLearner, A3CDensityModelMixin):
    def __init__(self, args):
        super(PseudoCountA3CLSTMLearner, self).__init__(args)
        self._init_density_model(args)

    def train(self):
        self._train() class PseudoCountQLearner(ValueBasedLearner, DensityModelMixin):
    """
    Based on DQN+CTS model from the paper 'Unifying Count-Based Exploration and Intrinsic Motivation' (https://arxiv.org/abs/1606.01868)
    Presently the implementation differs from the paper in that the novelty bonuses are computed online rather than by computing the
    prediction gains after the model has been updated with all frames from the episode. Async training with different final epsilon values
    tends to produce better results than just using a single actor-learner.
    """
    def __init__(self, args):
        self.args = args
        super(PseudoCountQLearner, self).__init__(args)

        self.cts_eta = args.cts_eta
        self.cts_beta = args.cts_beta
        self.batch_size = args.batch_update_size
        self.replay_memory = ReplayMemory(
            args.replay_size,
            self.local_network.get_input_shape(),
            self.num_actions)

        self._init_density_model(args)
        self._double_dqn_op()


    def generate_final_epsilon(self):
        if self.num_actor_learners == 1:
            return self.args.final_epsilon
        else:
            return super(PseudoCountQLearner, self).generate_final_epsilon()


    def _get_summary_vars(self):
        q_vars = super(PseudoCountQLearner, self)._get_summary_vars()

        bonus_q05 = tf.Variable(0., name='novelty_bonus_q05')
        s1 = tf.summary.scalar('Novelty_Bonus_q05_{}'.format(self.actor_id), bonus_q05)
        bonus_q50 = tf.Variable(0., name='novelty_bonus_q50')
        s2 = tf.summary.scalar('Novelty_Bonus_q50_{}'.format(self.actor_id), bonus_q50)
        bonus_q95 = tf.Variable(0., name='novelty_bonus_q95')
        s3 = tf.summary.scalar('Novelty_Bonus_q95_{}'.format(self.actor_id), bonus_q95)

        augmented_reward = tf.Variable(0., name='augmented_episode_reward')
        s4 = tf.summary.scalar('Augmented_Episode_Reward_{}'.format(self.actor_id), augmented_reward)       

        return q_vars + [bonus_q05, bonus_q50, bonus_q95, augmented_reward]


    #TODO: refactor to make this cleaner
    def prepare_state(self, state, total_episode_reward, steps_at_last_reward,
                      ep_t, episode_ave_max_q, episode_over, bonuses, total_augmented_reward):
        # Start a new game on reaching terminal state
        if episode_over:
            T = self.global_step.value() * self.max_local_steps
            t = self.local_step
            e_prog = float(t)/self.epsilon_annealing_steps
            episode_ave_max_q = episode_ave_max_q/float(ep_t)
            s1 = "Q_MAX {0:.4f}".format(episode_ave_max_q)
            s2 = "EPS {0:.4f}".format(self.epsilon)

            self.scores.insert(0, total_episode_reward)
            if len(self.scores) > 100:
                self.scores.pop()

            logger.info('T{0} / STEP {1} / REWARD {2} / {3} / {4}'.format(
                self.actor_id, T, total_episode_reward, s1, s2))
            logger.info('ID: {0} -- RUNNING AVG: {1:.0f} ± {2:.0f} -- BEST: {3:.0f}'.format(
                self.actor_id,
                np.array(self.scores).mean(),
                2*np.array(self.scores).std(),
                max(self.scores),
            ))

            self.log_summary(
                total_episode_reward,
                episode_ave_max_q,
                self.epsilon,
                np.percentile(bonuses, 5),
                np.percentile(bonuses, 50),
                np.percentile(bonuses, 95),
                total_augmented_reward,
            )

            state = self.emulator.get_initial_state()
            ep_t = 0
            total_episode_reward = 0
            episode_ave_max_q = 0
            episode_over = False

        return (
            state,
            total_episode_reward,
            steps_at_last_reward,
            ep_t,
            episode_ave_max_q,
            episode_over
        )


    def _double_dqn_op(self):
        q_local_action = tf.cast(tf.argmax(
            self.local_network.output_layer, axis=1), tf.int32)
        q_target_max = utils.ops.slice_2d(
            self.target_network.output_layer,
            tf.range(0, self.batch_size),
            q_local_action,
        )
        self.one_step_reward = tf.placeholder(tf.float32, self.batch_size, name='one_step_reward')
        self.is_terminal = tf.placeholder(tf.bool, self.batch_size, name='is_terminal')

        self.y_target = self.one_step_reward + self.cts_eta*self.gamma*q_target_max \
            * (1 - tf.cast(self.is_terminal, tf.float32))

        self.double_dqn_loss = self.local_network._value_function_loss(
            self.local_network.q_selected_action
            - tf.stop_gradient(self.y_target))

        self.double_dqn_grads = tf.gradients(self.double_dqn_loss, self.local_network.params)


    # def batch_update(self):
    #     if len(self.replay_memory) < self.replay_memory.maxlen//10:
    #         return

    #     s_i, a_i, r_i, s_f, is_terminal = self.replay_memory.sample_batch(self.batch_size)

    #     feed_dict={
    #         self.one_step_reward: r_i,
    #         self.target_network.input_ph: s_f,
    #         self.local_network.input_ph: np.vstack([s_i, s_f]),
    #         self.local_network.selected_action_ph: np.vstack([a_i, a_i]),
    #         self.is_terminal: is_terminal
    #     }
    #     grads = self.session.run(self.double_dqn_grads, feed_dict=feed_dict)
    #     self.apply_gradients_to_shared_memory_vars(grads)


    def batch_update(self):
        if len(self.replay_memory) < self.replay_memory.maxlen//10:
            return

        s_i, a_i, r_i, s_f, is_terminal = self.replay_memory.sample_batch(self.batch_size)

        feed_dict={
            self.local_network.input_ph: s_f,
            self.target_network.input_ph: s_f,
            self.is_terminal: is_terminal,
            self.one_step_reward: r_i,
        }
        y_target = self.session.run(self.y_target, feed_dict=feed_dict)

        feed_dict={
            self.local_network.input_ph: s_i,
            self.local_network.target_ph: y_target,
            self.local_network.selected_action_ph: a_i
        }
        grads = self.session.run(self.local_network.get_gradients,
                                 feed_dict=feed_dict)
        self.apply_gradients_to_shared_memory_vars(grads)


    def train(self):
        """ Main actor learner loop for n-step Q learning. """
        logger.debug("Actor {} resuming at Step {}, {}".format(self.actor_id, 
            self.global_step.value(), time.ctime()))

        s = self.emulator.get_initial_state()
        
        s_batch = list()
        a_batch = list()
        y_batch = list()
        bonuses = deque(maxlen=1000)
        episode_over = False
        
        t0 = time.time()
        global_steps_at_last_record = self.global_step.value()
        while (self.global_step.value() < self.max_global_steps):
            # # Sync local learning net with shared mem
            # self.sync_net_with_shared_memory(self.local_network, self.learning_vars)
            # self.save_vars()
            rewards =      list()
            states =       list()
            actions =      list()
            max_q_values = list()
            local_step_start = self.local_step
            total_episode_reward = 0
            total_augmented_reward = 0
            episode_ave_max_q = 0
            ep_t = 0

            while not episode_over:
                # Sync local learning net with shared mem
                self.sync_net_with_shared_memory(self.local_network, self.learning_vars)
                self.save_vars()

                # Choose next action and execute it
                a, q_values = self.choose_next_action(s)

                new_s, reward, episode_over = self.emulator.next(a)
                total_episode_reward += reward
                max_q = np.max(q_values)

                current_frame = new_s[...,-1]
                bonus = self.density_model.update(current_frame)
                bonuses.append(bonus)

                # Rescale or clip immediate reward
                reward = self.rescale_reward(self.rescale_reward(reward) + bonus)
                total_augmented_reward += reward
                ep_t += 1
                
                rewards.append(reward)
                states.append(s)
                actions.append(a)
                max_q_values.append(max_q)
                
                s = new_s
                self.local_step += 1
                episode_ave_max_q += max_q
                
                global_step, _ = self.global_step.increment()

                if global_step % self.q_target_update_steps == 0:
                    self.update_target()
                if global_step % self.density_model_update_steps == 0:
                    self.write_density_model()

                # Sync local tensorflow target network params with shared target network params
                if self.target_update_flags.updated[self.actor_id] == 1:
                    self.sync_net_with_shared_memory(self.target_network, self.target_vars)
                    self.target_update_flags.updated[self.actor_id] = 0
                if self.density_model_update_flags.updated[self.actor_id] == 1:
                    self.read_density_model()
                    self.density_model_update_flags.updated[self.actor_id] = 0

                if self.local_step % self.q_update_interval == 0:
                    self.batch_update()

                if self.is_master() and (self.local_step % 500 == 0):
                    bonus_array = np.array(bonuses)
                    steps = global_step - global_steps_at_last_record
                    global_steps_at_last_record = global_step

                    logger.debug('Mean Bonus={:.4f} / Max Bonus={:.4f} / STEPS/s={}'.format(
                        bonus_array.mean(), bonus_array.max(), steps/float(time.time()-t0)))
                    t0 = time.time()


            else:
                #compute monte carlo return
                mc_returns = np.zeros((len(rewards),), dtype=np.float32)
                running_total = 0.0
                for i, r in enumerate(reversed(rewards)):
                    running_total = r + self.gamma*running_total
                    mc_returns[len(rewards)-i-1] = running_total

                mixed_returns = self.cts_eta*np.asarray(rewards) + (1-self.cts_eta)*mc_returns

                #update replay memory
                states.append(new_s)
                episode_length = len(rewards)
                for i in range(episode_length):
                    self.replay_memory.append(
                        states[i],
                        actions[i],
                        mixed_returns[i],
                        i+1 == episode_length)

            s, total_episode_reward, _, ep_t, episode_ave_max_q, episode_over = \
                self.prepare_state(s, total_episode_reward, self.local_step, ep_t, episode_ave_max_q, episode_over, bonuses, total_augmented_reward) class TRPOLearner(BaseA3CLearner):
	'''
	Implementation of Trust Region Policy Optimization + Generalized Advantage Estimation 
	as described in https://arxiv.org/pdf/1506.02438.pdf

	∂'π = F^-1 ∂π where F is the Fischer Information Matrix
	We can't compute F^-1 directly except for very small networks
	so we'll use either conjugate gradient descent to approximate F^-1 ∂π
	'''

	def __init__(self, args):
		args.entropy_regularisation_strength = 0.0
		super(TRPOLearner, self).__init__(args)

		self.batch_size = 512
		self.max_cg_iters = 20
		self.num_epochs = args.num_epochs
		self.cg_damping = args.cg_damping
		self.cg_subsample = args.cg_subsample
		self.max_kl = args.max_kl
		self.max_rollout = args.max_rollout
		self.episodes_per_batch = args.episodes_per_batch
		self.baseline_vars = args.baseline_vars
		self.experience_queue = args.experience_queue
		self.task_queue = args.task_queue
		self.history_length = args.history_length
		self.append_timestep = args.arch == 'FC'


		policy_conf = {'name': 'policy_network_{}'.format(self.actor_id),
					   'input_shape': self.input_shape,
					   'num_act': self.num_actions,
					   'args': args}
		value_conf = policy_conf.copy()
		value_conf['name'] = 'value_network_{}'.format(self.actor_id)
		value_conf['input_shape'] = args.vf_input_shape

		self.device = '/gpu:0' if self.is_master() else '/cpu:0'
		with tf.device(self.device):
			#we use separate networks as in the paper since so we don't do damage to the trust region updates
			self.policy_network = args.network(policy_conf)
			self.value_network = PolicyValueNetwork(value_conf, use_policy_head=False)
			self.local_network = self.policy_network
			self._build_ops()

		if self.is_master():
			var_list = self.policy_network.params + self.value_network.params
			self.saver = tf.train.Saver(var_list=var_list, max_to_keep=3,
                                        keep_checkpoint_every_n_hours=2)


	def _build_ops(self):
		eps = 1e-10
		self.dist_params = self.policy_network.dist.params()
		num_params = self.dist_params.get_shape().as_list()[1]
		self.old_params = tf.placeholder(tf.float32, shape=[None, num_params], name='old_params')

		selected_prob = tf.exp(self.policy_network.log_output_selected_action)
		old_dist = self.policy_network.dist.__class__(self.old_params)
		old_selected_prob = tf.exp(old_dist.log_likelihood(self.policy_network.selected_action_ph))

		self.theta = self.policy_network.flat_vars
		self.policy_loss = -tf.reduce_mean(tf.multiply(
			self.policy_network.adv_actor_ph,
			selected_prob / old_selected_prob
		))
		self.pg = utils.ops.flatten_vars(
			tf.gradients(self.policy_loss, self.policy_network.params))

		self.kl = tf.reduce_mean(self.policy_network.dist.kl_divergence(self.old_params))
		self.kl_firstfixed = tf.reduce_mean(self.policy_network.dist.kl_divergence(
			tf.stop_gradient(self.dist_params)))


		kl_grads = tf.gradients(self.kl_firstfixed, self.policy_network.params)
		flat_kl_grads = utils.ops.flatten_vars(kl_grads)

		self.pg_placeholder = tf.placeholder(tf.float32, shape=self.pg.get_shape().as_list(), name='pg_placeholder')
		self.fullstep, self.neggdotstepdir = self._conjugate_gradient_ops(
			-self.pg_placeholder, flat_kl_grads, max_iterations=self.max_cg_iters)


	def _conjugate_gradient_ops(self, pg_grads, kl_grads, max_iterations=20, residual_tol=1e-10):
		'''
		Construct conjugate gradient descent algorithm inside computation graph for improved efficiency
		'''
		i0 = tf.constant(0, dtype=tf.int32)
		loop_condition = lambda i, r, p, x, rdotr: tf.logical_and(
			tf.greater(rdotr, residual_tol), tf.less(i, max_iterations))


		def body(i, r, p, x, rdotr):
			fvp = utils.ops.flatten_vars(tf.gradients(
				tf.reduce_sum(tf.stop_gradient(p)*kl_grads),
				self.policy_network.params))

			z = fvp + self.cg_damping * p

			alpha = rdotr / (tf.reduce_sum(p*z) + 1e-8)
			x += alpha * p
			r -= alpha * z

			new_rdotr = tf.reduce_sum(r*r)
			beta = new_rdotr / (rdotr + 1e-8)
			p = r + beta * p

			new_rdotr = tf.Print(new_rdotr, [i, new_rdotr], 'Iteration / Residual: ')

			return i+1, r, p, x, new_rdotr

		_, r, p, stepdir, rdotr = tf.while_loop(
			loop_condition,
			body,
			loop_vars=[i0,
					   pg_grads,
					   pg_grads,
					   tf.zeros_like(pg_grads),
					   tf.reduce_sum(pg_grads*pg_grads)])

		fvp = utils.ops.flatten_vars(tf.gradients(
			tf.reduce_sum(tf.stop_gradient(stepdir)*kl_grads),
			self.policy_network.params))

		shs = 0.5 * tf.reduce_sum(stepdir*fvp)
		lm = tf.sqrt((shs + 1e-8) / self.max_kl)
		fullstep = stepdir / lm
		neggdotstepdir = tf.reduce_sum(pg_grads*stepdir) / lm

		return fullstep, neggdotstepdir


	def choose_next_action(self, state):
		return self.policy_network.get_action(self.session, state)


	def run_minibatches(self, data, *ops):
		outputs = [np.zeros(op.get_shape().as_list(), dtype=np.float32) for op in ops]

		data_size = len(data['state'])
		for start in range(0, data_size, self.batch_size):
			end = start + np.minimum(self.batch_size, data_size-start)
			feed_dict={
				self.policy_network.input_ph:           data['state'][start:end],
				self.policy_network.selected_action_ph: data['action'][start:end],
				self.policy_network.adv_actor_ph:       data['reward'][start:end],
				self.old_params:                        data['pi'][start:end]
			}
			for i, output_i in enumerate(self.session.run(ops, feed_dict=feed_dict)):
				outputs[i] += output_i * (end-start)/float(data_size)

		return outputs


	def linesearch(self, data, x, fullstep, expected_improve_rate):
		accept_ratio = .1
		backtrack_ratio = .7
		max_backtracks = 15
    
		fval = self.run_minibatches(data, self.policy_loss)

		for (_n_backtracks, stepfrac) in enumerate(backtrack_ratio**np.arange(max_backtracks)):
		    xnew = x + stepfrac * fullstep
		    self.assign_vars(self.policy_network, xnew)
		    newfval, kl = self.run_minibatches(data, self.policy_loss, self.kl)

		    improvement = fval - newfval
		    logger.debug('Improvement {} / Mean KL {}'.format(improvement, kl))

		    expected_improve = expected_improve_rate * stepfrac
		    ratio = improvement / expected_improve
		    # if ratio > accept_ratio and improvement > 0:
		    if kl < self.max_kl and improvement > 0:
		        return xnew

		logger.debug('No update')
		return x


	def fit_baseline(self, data, mix_old=.9):
		data_size = len(data['state'])
		proc_state = self.preprocess_value_state(data)
		print 'diffs', (data['mc_return'] - data['values']).mean()
		target = (1-mix_old)*data['mc_return'] + mix_old*data['values']
		grads = [np.zeros(g.get_shape().as_list(), dtype=np.float32) for g in self.value_network.get_gradients]

		#permute data in minibatches so we don't introduce bias
		perm = np.random.permutation(data_size)
		for start in range(0, data_size, self.batch_size):
			end = start + np.minimum(self.batch_size, data_size-start)
			batch_idx = perm[start:end]
			feed_dict={
				self.value_network.input_ph:         proc_state[batch_idx],
				self.value_network.critic_target_ph: target[batch_idx]
			}
			output_i = self.session.run(self.value_network.get_gradients, feed_dict=feed_dict)
			
			for i, g in enumerate(output_i):
				grads[i] += g * (end-start)/float(data_size)

			self._apply_gradients_to_shared_memory_vars(output_i, self.baseline_vars)
			self.sync_net_with_shared_memory(self.value_network, self.baseline_vars)


	def preprocess_value_state(self, data):
		if self.append_timestep: #this is particularly helpful on MuJoCo environments
			return np.hstack([data['state'], np.tile(data['timestep'].reshape(-1, 1, 1), self.history_length)])
		else:
			return data['state']


	def predict_values(self, data):
		state = self.preprocess_value_state({
			'state': np.array(data['state']),
			'timestep': np.array(data['timestep'])})
		return self.session.run(
			self.value_network.output_layer_v,
			feed_dict={self.value_network.input_ph: state})[:, 0]


	def update_grads(self, data):
		#we need to compute the policy gradient in minibatches to avoid GPU OOM errors on Atari
		print 'fitting baseline...'
		self.fit_baseline(data)

		normalized_advantage = (data['advantage'] - data['advantage'].mean())/(data['advantage'].std() + 1e-8)
		data['reward'] = normalized_advantage

		print 'running policy gradient...'
		pg = self.run_minibatches(data, self.pg)[0]

		data_size = len(data['state'])
		subsample = np.random.choice(data_size, int(data_size*self.cg_subsample), replace=False)
		feed_dict={
			self.policy_network.input_ph:           data['state'][subsample],
			self.policy_network.selected_action_ph: data['action'][subsample],
			self.policy_network.adv_actor_ph:       data['reward'][subsample],
			self.old_params:                        data['pi'][subsample],
			self.pg_placeholder:                    pg
		}

		print 'running conjugate gradient descent...'
		theta_prev, fullstep, neggdotstepdir = self.session.run(
			[self.theta, self.fullstep, self.neggdotstepdir], feed_dict=feed_dict)

		print 'running linesearch...'
		new_theta = self.linesearch(data, theta_prev, fullstep, neggdotstepdir)
		self.assign_vars(self.policy_network, new_theta)

		return self.session.run(self.kl, feed_dict)


	def _run_worker(self):
		while True:
			signal = self.task_queue.get()
			if signal == 'EXIT':
				break

			self.sync_net_with_shared_memory(self.local_network, self.learning_vars)
			self.sync_net_with_shared_memory(self.value_network, self.baseline_vars)
			s = self.emulator.get_initial_state()

			data = {
				'state':     list(),
				'pi':        list(),
				'action':    list(),
				'reward':    list(),
			}
			episode_over = False
			accumulated_rewards = list()
			while not episode_over and len(accumulated_rewards) < self.max_rollout:
				a, pi = self.choose_next_action(s)
				new_s, reward, episode_over = self.emulator.next(a)
				
				accumulated_rewards.append(reward)

				data['state'].append(s)
				data['pi'].append(pi)
				data['action'].append(a)
				data['reward'].append(reward)

				s = new_s

			mc_returns = list()
			running_total = 0.0
			for r in reversed(accumulated_rewards):
				running_total = r + self.gamma*running_total
				mc_returns.insert(0, running_total)

			timestep = np.arange(len(mc_returns), dtype=np.float32)/self.emulator.env.spec.timestep_limit
			data['timestep'] = timestep
			data['mc_return'] = mc_returns
			episode_reward = sum(accumulated_rewards)
			logger.debug('T{} / Episode Reward {}'.format(
				self.actor_id, episode_reward))

			self.experience_queue.put((data, episode_reward))
			

	def _run_master(self):
		for epoch in range(self.num_epochs):
			data = {
				'state':     list(),
				'pi':        list(),
				'action':    list(),
				'reward':    list(),
				'advantage': list(),
				'mc_return': list(),
				'timestep':  list(),
				'values':    list(),
			}
			#launch worker tasks
			for i in xrange(self.episodes_per_batch):
				self.task_queue.put(i)

			#collect worker experience
			episode_rewards = list()
			t0 = time.time()
			for _ in xrange(self.episodes_per_batch):
				worker_data, reward = self.experience_queue.get()
				episode_rewards.append(reward)

				values = self.predict_values(worker_data)
				advantages = self.compute_gae(worker_data['reward'], values.tolist(), 0)
				# advantages = worker_data['mc_return'] - values
				worker_data['values'] = values
				worker_data['advantage'] = advantages
				for key, value in worker_data.items():
					data[key].extend(value)
					
			t1 = time.time()
			kl = self.update_grads({
				k: np.array(v) for k, v in data.items()})
			self.update_shared_memory()
			t2 = time.time()

			mean_episode_reward = np.array(episode_rewards).mean()
			logger.info('Epoch {} / Mean KL Divergence {} / Mean Reward {} / Experience Time {:.2f}s / Training Time {:.2f}s'.format(
				epoch+1, kl, mean_episode_reward, t1-t0, t2-t1))


	def train(self):
		if self.is_master():
			self._run_master()
			for _ in xrange(self.num_actor_learners):
				self.task_queue.put('EXIT')
		else:
			self._run_worker() class PGQLearner(BasePGQLearner):
    def choose_next_action(self, state):
        network_output_v, network_output_pi = self.session.run(
                [self.local_network.output_layer_v,
                 self.local_network.output_layer_pi], 
                feed_dict={self.local_network.input_ph: [state]})

        network_output_pi = network_output_pi.reshape(-1)
        network_output_v = np.asscalar(network_output_v)

        action_index = self.sample_policy_action(network_output_pi)
        new_action = np.zeros([self.num_actions])
        new_action[action_index] = 1

        return new_action, network_output_v, network_output_pi


    def apply_update(self, states, actions, targets, advantages):
        feed_dict={
            self.local_network.input_ph: states,
            self.local_network.selected_action_ph: actions,
            self.local_network.critic_target_ph: targets,
            self.local_network.adv_actor_ph: advantages,
        }
        grads, entropy = self.session.run(
            [self.local_network.get_gradients, self.local_network.entropy],
            feed_dict=feed_dict)
        self.apply_gradients_to_shared_memory_vars(grads)

        self.q_update_counter += 1
        if self.q_update_counter % self.q_update_interval == 0:
            self.batch_q_update()

        return entropy class CEMLearner(BaseA3CLearner):
	'''
	Implementation of Cross-Entropy Method, Useful as a baseline for simple environments
	'''
	def __init__(self, args):
		super(CEMLearner, self).__init__(args)

		policy_conf = {'name': 'local_learning_{}'.format(self.actor_id),
					   'input_shape': self.input_shape,
					   'num_act': self.num_actions,
					   'args': args}

		self.local_network = args.network(policy_conf)
		self.num_params = np.sum([
			np.prod(v.get_shape().as_list())
			for v in self.local_network.params])

		logger.info('Parameter count: {}'.format(self.num_params))
		self.mu = np.zeros(self.num_params)
		self.sigma = np.ones(self.num_params)
		self.num_samples = args.episodes_per_batch
		self.num_epochs = args.num_epochs

		if self.is_master():
			var_list = self.local_network.params
			self.saver = tf.train.Saver(var_list=var_list, max_to_keep=3,
                                        keep_checkpoint_every_n_hours=2)


	def choose_next_action(self, state):
		return self.local_network.get_action(self.session, state)


	def sample_theta(self, N):
		return self.mu + np.random.randn(N, self.num_params)*self.sigma


	def update_sample_distribution(self, population, rewards, keep_ratio=0.25, noise=0.01):
		num_to_keep = int(keep_ratio * len(population))
		elite = np.array(rewards).argsort()[-num_to_keep:]

		self.mu = population[elite].mean(axis=0)
		self.sigma = population[elite].std(axis=0) + noise

		return np.array(rewards)[elite].mean(), population[elite]


	def train(self):
		consecutive_successes = 0

		for epoch in xrange(self.num_epochs):
			episode_rewards = list()
			population = self.sample_theta(self.num_samples)

			for theta in population:
				self.assign_vars(self.local_network, theta)
				s = self.emulator.get_initial_state()

				total_episode_reward = 0.0
				episode_over = False

				while not episode_over:
					a, pi = self.choose_next_action(s)
					s, reward, episode_over = self.emulator.next(a)
					total_episode_reward += reward

				episode_rewards.append(total_episode_reward)

			population_mean_reward = np.array(episode_rewards).mean()
			elite_mean_reward, elite_set = self.update_sample_distribution(
				population, episode_rewards, noise=.1/(1+epoch))
			logger.info('Epoch {} / Population Mean {} / Elite Mean {}'.format(
				epoch+1, population_mean_reward, elite_mean_reward))

			if self.emulator.env.spec.reward_threshold and elite_mean_reward > self.emulator.env.spec.reward_threshold:
				consecutive_successes += 1
			else:
				consecutive_successes = 0

			#if the elite set is consistently good enough we'll stop updating the parameters
			if consecutive_successes > 10:
				logger.info('switching to testing mode...')
				episode_rewards = list()

				for _ in xrange(200):
					i = np.random.choice(elite_set.shape[0])
					self.assign_vars(self.local_network, elite_set[i])
					s = self.emulator.get_initial_state()
					total_episode_reward = 0.0
					episode_over = False

					while not episode_over:
						a, pi = self.choose_next_action(s)
						s, reward, episode_over = self.emulator.next(a)
						total_episode_reward += reward

					episode_rewards.append(total_episode_reward)
				
				logger.info('Average Test Reward: {}'.format(np.array(episode_rewards).mean()))
				return
import os
import sys
import time
import yaml
import argparse
import numpy as np
import utils.logger
import multiprocessing
import tensorflow as tf
from networks.q_network import QNetwork
from multiprocessing import Process, Queue
from networks.dueling_network import DuelingNetwork
from networks.continuous_actions import ContinuousPolicyNetwork, ContinuousPolicyValueNetwork
from networks.policy_v_network import PolicyNetwork, PolicyValueNetwork, PolicyRepeatNetwork, SequencePolicyVNetwork
from utils.shared_memory import SharedCounter, SharedVars, SharedFlags, Barrier
from algorithms.policy_based_actor_learner import A3CLearner, A3CLSTMLearner
from algorithms.sequence_decoder_actor_learner import ActionSequenceA3CLearner, ARA3CLearner
from algorithms.value_based_actor_learner import NStepQLearner, DuelingLearner, OneStepSARSALearner
from algorithms.intrinsic_motivation_actor_learner import PseudoCountA3CLearner, PseudoCountA3CLSTMLearner, PseudoCountQLearner
from algorithms.trpo_actor_learner import TRPOLearner
from algorithms.pgq_actor_learner import PGQLearner
from algorithms.cem_actor_learner import CEMLearner
    from ale_python_interface import ALEInterface
        from environments import atari_environment

# -*- coding: utf-8 -*-


logger = utils.logger.getLogger('main')


ALGORITHMS = {
    'q': (NStepQLearner, QNetwork),

    'sarsa': (OneStepSARSALearner, QNetwork),
