def gauss_selfKL_firstfixed(mu, logstd):
def gauss_log_prob(mu, logstd, x):
def gauss_KL(mu1, logstd1, mu2, logstd2):
def gauss_ent(mu, logstd):
def discount(x, gamma):
def cat_sample(prob_nk, seed=1234):
def slice_2d(x, inds0, inds1):
def var_shape(x):
def numel(x):
def flatgrad(loss, var_list):
def conjugate_gradient(f_Ax, b, cg_iters=10, residual_tol=1e-10):
def linesearch(f, x, fullstep, expected_improve_rate):
  def __init__(self, session, var_list):
  def __call__(self, theta):
  def __init__(self, session, var_list):
  def __call__(self):
  def __init__(self, session, var_list):
  def __call__(self):
  def __init__(self, session, var_list):
  def __call__(self, weights):
def xavier_initializer(self, shape):
def fully_connected(input_layer, input_size, output_size, weight_init,
                    bias_init, scope, trainable):
  N = prob_nk.shape[0]
class SetFromFlat(object):
class GetFlat(object):
class GetPolicyWeights(object):
class SetPolicyWeights(object): class _MyFormatter(logging.Formatter):
    def format(self, record):
def set_file_handler(path=None, prefix='', time_str=''):
_LOGGING_METHOD = ['info', 'warning', 'error', 'critical',
                   'warn', 'exception', 'debug'] START_SIGNAL = 1
END_SIGNAL = -1
END_ROLLOUT_SIGNAL = -2
AGENT_COLLECT_FILTER_INFO = 23
AGENT_SYNCHRONIZE_FILTER = 233
AGENT_SET_POLICY_WEIGHTS = 2333
AGENT_EVOLUTION_TRAIN = 23333
AGENT_EVOLUTION_START = 233333 class normalizer(object):
    def __init__(self,
                 mean=0, variance=0, num_steps=0.,
                 filter_mean=True, clip_value=5.0):
    def raw_filter(self, o):
    def filter(self, o, update=False):
    def set_parameters(self, mean, v, n):
    def get_parameters(self): class base_agent(multiprocessing.Process):

    def __init__(self,
                 args,
                 observation_size,
                 action_size,
                 task_q,
                 result_q,
                 name_scope='trpo_agent',
                 is_rollout_agent=False):

        # the multiprocessing initialization
        multiprocessing.Process.__init__(self)
        self.task_q = task_q
        self.result_q = result_q
        self.gnn_placeholder_list = None
        self.obs_placeholder = None

        # the configurations for the agent
        self.args = args

        # the network parameters
        self.name_scope = name_scope
        self.observation_size = observation_size
        self.action_size = action_size
        self.is_rollout_agent = is_rollout_agent

        # the variables and networks to be used, init them before use them
        self.policy_network = None
        self.step_policy_network = None
        self.policy_var_list = None

        self.tf_var_list = None
        self.iteration = None

        # the gnn parameters
        if self.args.use_gnn_as_policy:
            self.gnn_parameter_initialization()

        self.base_path = init_path.get_base_dir()

    def build_session(self):
        if self.args.use_gpu:
            config = tf.ConfigProto(device_count={'GPU': 1})
        else:
            config = tf.ConfigProto(device_count={'GPU': 0})
        config.gpu_options.allow_growth = True  # don't take full gpu memory
        self.session = tf.Session(config=config)

    def fetch_policy_info(self):
        assert self.policy_network is not None, \
            logger.error('Init the policy network before using it')

        # input placeholders to the policy networks
        if self.args.use_gnn_as_policy:
            # index placeholders
            self.receive_idx_placeholder, self.send_idx_placeholder, \
                self.node_type_idx_placeholder, \
                self.inverse_node_type_idx_placeholder, \
                self.output_type_idx_placeholder, \
                self.inverse_output_type_idx_placeholder, \
                self.batch_size_int_placeholder = \
                self.policy_network.get_gnn_idx_placeholder()

            # the graph_obs placeholders and graph_parameters_placeholders
            self.graph_obs_placeholder = \
                self.policy_network.get_input_obs_placeholder()
            self.graph_parameters_placeholder = \
                self.policy_network.get_input_parameters_placeholder()

            self.gnn_placeholder_list = [
                self.receive_idx_placeholder,
                self.send_idx_placeholder,
                self.node_type_idx_placeholder,
                self.inverse_node_type_idx_placeholder,
                self.output_type_idx_placeholder,
                self.inverse_output_type_idx_placeholder,
                self.batch_size_int_placeholder,
                self.graph_obs_placeholder,
                self.graph_parameters_placeholder
            ]
            if self.args.nervenetplus:

                self.step_receive_idx_placeholder, \
                    self.step_send_idx_placeholder, \
                    self.step_node_type_idx_placeholder, \
                    self.step_inverse_node_type_idx_placeholder, \
                    self.step_output_type_idx_placeholder, \
                    self.step_inverse_output_type_idx_placeholder, \
                    self.step_batch_size_int_placeholder = \
                    self.step_policy_network.get_gnn_idx_placeholder()

                # the graph_obs placeholders and graph_parameters_placeholders
                self.step_graph_obs_placeholder = \
                    self.step_policy_network.get_input_obs_placeholder()
                self.step_graph_parameters_placeholder = \
                    self.step_policy_network.get_input_parameters_placeholder()
        else:
            self.obs_placeholder = self.policy_network.get_input_placeholder()

        self.raw_obs_placeholder = None

        # output from the policy networks means for each action
        self.action_dist_mu = self.policy_network.get_action_dist_mu()

        # log std parameters of actions (all the same)
        self.action_dist_logstd_param = \
            self.policy_network.get_action_dist_logstd_param()
        self.action_dist_logstd = self.action_dist_logstd_param

        if self.args.nervenetplus:
            self.step_action_dist_mu = \
                self.step_policy_network.get_action_dist_mu()

            # log std parameters of actions (all the same)
            self.step_action_dist_logstd_param = \
                self.step_policy_network.get_action_dist_logstd_param()
            self.step_action_dist_logstd = self.step_action_dist_logstd_param
        # "policy_var_list": to be passed to the rollout agents
        # "all_policy_var_list": to be saved into the checkpoint
        self.policy_var_list, self.all_policy_var_list = \
            self.policy_network.get_var_list()
        self.iteration = self.policy_network.get_iteration_var()
        self.iteration_add_op = self.iteration.assign_add(1)

    def build_policy_network(self, adj_matrix=None, node_attr=None):
        if self.args.nervenetplus:
            assert self.args.use_gnn_as_policy and self.args.use_nervenet

        if self.args.use_gnn_as_policy:
            if self.args.use_nervenet:
                if self.args.nervenetplus:
                    if self.args.tree_net:
                        self.policy_network = treenet_policy.nervenet(
                            session=self.session,
                            name_scope=self.name_scope + '_policy',
                            input_size=self.observation_size,
                            output_size=self.action_size,
                            adj_matrix=adj_matrix,
                            node_attr=node_attr,
                            args=self.args,
                            is_rollout_agent=self.is_rollout_agent
                        )
                        self.step_policy_network = treenet_policy.nervenet(
                            session=self.session,
                            name_scope=self.name_scope + 'step_policy',
                            input_size=self.observation_size,
                            output_size=self.action_size,
                            adj_matrix=self.adj_matrix,
                            node_attr=self.node_attr,
                            args=self.args,
                            is_rollout_agent=True
                        )
                        self.node_info = self.policy_network.get_node_info()

                        self.step_policy_var_list, _ = \
                            self.step_policy_network.get_var_list()

                        self.set_step_policy = \
                            utils.SetPolicyWeights(self.session,
                                                   self.step_policy_var_list)
                    else:
                        self.policy_network = nervenetplus_policy.nervenet(
                            session=self.session,
                            name_scope=self.name_scope + '_policy',
                            input_size=self.observation_size,
                            output_size=self.action_size,
                            adj_matrix=adj_matrix,
                            node_attr=node_attr,
                            args=self.args,
                            is_rollout_agent=self.is_rollout_agent
                        )
                        self.step_policy_network = nervenetplus_policy.nervenet(
                            session=self.session,
                            name_scope=self.name_scope + 'step_policy',
                            input_size=self.observation_size,
                            output_size=self.action_size,
                            adj_matrix=self.adj_matrix,
                            node_attr=self.node_attr,
                            args=self.args,
                            is_rollout_agent=True
                        )
                        self.node_info = self.policy_network.get_node_info()

                        self.step_policy_var_list, _ = \
                            self.step_policy_network.get_var_list()

                        self.set_step_policy = \
                            utils.SetPolicyWeights(self.session,
                                                   self.step_policy_var_list)
                else:
                    self.policy_network = nervenet_policy.nervenet(
                        session=self.session,
                        name_scope=self.name_scope + '_policy',
                        input_size=self.observation_size,
                        output_size=self.action_size,
                        adj_matrix=adj_matrix,
                        node_attr=node_attr,
                        args=self.args
                    )
            else:
                self.policy_network = gated_graph_policy_network.GGNN(
                    session=self.session,
                    name_scope=self.name_scope + '_policy',
                    input_size=self.observation_size,
                    output_size=self.action_size,
                    ob_placeholder=None,
                    trainable=True,
                    build_network_now=True,
                    is_baseline=False,
                    placeholder_list=None,
                    args=self.args
                )
            self.raw_obs_placeholder = None
            self.node_info = self.policy_network.get_node_info()
        else:
            self.policy_network = policy_network.policy_network(
                session=self.session,
                name_scope=self.name_scope + '_policy',
                input_size=self.observation_size,
                output_size=self.action_size,
                ob_placeholder=None,
                trainable=True,
                build_network_now=True,
                define_std=True,
                is_baseline=False,
                args=self.args
            )

        # if use the nervenetplus model
        # if self.args.nervenetplus:
        #     # build the action model
        #     with tf.variable_scope('', reuse=True):
        #         self.step_policy_network = nervenetplus_policy.nervenet(
        #             session=self.session,
        #             name_scope=self.name_scope + '_policy',
        #             input_size=self.observation_size,
        #             output_size=self.action_size,
        #             adj_matrix=self.adj_matrix,
        #             node_attr=self.node_attr,
        #             args=self.args,
        #             is_rollout_agent=True
        #         )
        #     self.node_info = self.policy_network.get_node_info()
        self.fetch_policy_info()

    def gnn_parameter_initialization(self):
        '''
            @brief:
                the parameters for the gnn, see the gated_graph_network_policy
                file for details what these variables mean.
        '''
        self.receive_idx = None
        self.send_idx = None
        self.node_type_idx = None
        self.inverse_node_type_idx = None
        self.output_type_idx = None
        self.inverse_output_type_idx = None
        self.last_batch_size = -1
        self.nervenetplus_batch_pos = None

    def prepared_policy_network_feeddict(self, obs_n, rollout_data=None,
                                         step_model=False):
        '''
            @brief: prepare the feed dict for the policy network part
        '''
        nervenetplus_batch_pos = None

        if self.args.use_gnn_as_policy:

            if not self.args.nervenetplus or obs_n.shape[0] == 1:
                graph_obs, graph_parameters, \
                    self.receive_idx, self.send_idx, \
                    self.node_type_idx, self.inverse_node_type_idx, \
                    self.output_type_idx, self.inverse_output_type_idx, \
                    self.last_batch_size = \
                    graph_data_util.construct_graph_input_feeddict(
                        self.node_info,
                        obs_n,
                        self.receive_idx,
                        self.send_idx,
                        self.node_type_idx,
                        self.inverse_node_type_idx,
                        self.output_type_idx,
                        self.inverse_output_type_idx,
                        self.last_batch_size,
                        request_data=['ob', 'idx']
                    )
            else:
                assert rollout_data is not None

                # preprocess the episodic information
                graph_obs, graph_parameters, _, _, _, _, _, _, _ = \
                    graph_data_util.construct_graph_input_feeddict(
                        self.node_info, obs_n,
                        -1, -1, -1, -1, -1, -1, -1,
                        request_data=['ob']
                    )
                nervenetplus_batch_pos, total_size = \
                    nervenetplus_util.nervenetplus_step_assign(
                        rollout_data, self.args.gnn_num_prop_steps
                    )
                _, _, self.receive_idx, self.send_idx, \
                    self.node_type_idx, self.inverse_node_type_idx, \
                    self.output_type_idx, self.inverse_output_type_idx, \
                    self.last_batch_size = \
                    graph_data_util.construct_graph_input_feeddict(
                        self.node_info,
                        np.empty(
                            [int(total_size / self.args.gnn_num_prop_steps)]
                        ),
                        self.receive_idx,
                        self.send_idx,
                        self.node_type_idx,
                        self.inverse_node_type_idx,
                        self.output_type_idx,
                        self.inverse_output_type_idx,
                        self.last_batch_size,
                        request_data=['idx']
                    )

            if step_model:
                feed_dict = {
                    self.step_batch_size_int_placeholder:
                        int(self.last_batch_size),
                    self.step_receive_idx_placeholder:
                        self.receive_idx,
                    self.step_inverse_node_type_idx_placeholder:
                        self.inverse_node_type_idx,
                    self.step_inverse_output_type_idx_placeholder:
                        self.inverse_output_type_idx
                }

                # append the input obs and parameters
                for i_node_type in self.node_info['node_type_dict']:
                    feed_dict[self.step_graph_obs_placeholder[i_node_type]] = \
                        graph_obs[i_node_type]
                    feed_dict[self.step_graph_parameters_placeholder[i_node_type]] = \
                        graph_parameters[i_node_type]

                # append the send idx
                for i_edge in self.node_info['edge_type_list']:
                    feed_dict[self.step_send_idx_placeholder[i_edge]] = \
                        self.send_idx[i_edge]

                # append the node type idx
                for i_node_type in self.node_info['node_type_dict']:
                    feed_dict[self.step_node_type_idx_placeholder[i_node_type]] \
                        = self.node_type_idx[i_node_type]

                # append the output type idx
                for i_output_type in self.node_info['output_type_dict']:
                    feed_dict[self.step_output_type_idx_placeholder[i_output_type]] \
                        = self.output_type_idx[i_output_type]

                # if the raw_obs is needed for the baseline
                if self.raw_obs_placeholder is not None:
                    feed_dict[self.raw_obs_placeholder] = obs_n
            else:
                feed_dict = {
                    self.batch_size_int_placeholder:
                        int(self.last_batch_size),
                    self.receive_idx_placeholder:
                        self.receive_idx,
                    self.inverse_node_type_idx_placeholder:
                        self.inverse_node_type_idx,
                    self.inverse_output_type_idx_placeholder:
                        self.inverse_output_type_idx
                }

                # append the input obs and parameters
                for i_node_type in self.node_info['node_type_dict']:
                    feed_dict[self.graph_obs_placeholder[i_node_type]] = \
                        graph_obs[i_node_type]
                    feed_dict[self.graph_parameters_placeholder[i_node_type]] = \
                        graph_parameters[i_node_type]

                # append the send idx
                for i_edge in self.node_info['edge_type_list']:
                    feed_dict[self.send_idx_placeholder[i_edge]] = \
                        self.send_idx[i_edge]

                # append the node type idx
                for i_node_type in self.node_info['node_type_dict']:
                    feed_dict[self.node_type_idx_placeholder[i_node_type]] \
                        = self.node_type_idx[i_node_type]

                # append the output type idx
                for i_output_type in self.node_info['output_type_dict']:
                    feed_dict[self.output_type_idx_placeholder[i_output_type]] \
                        = self.output_type_idx[i_output_type]

                # if the raw_obs is needed for the baseline
                if self.raw_obs_placeholder is not None:
                    feed_dict[self.raw_obs_placeholder] = obs_n
        else:
            # it is the most easy case, nice and easy
            feed_dict = {self.obs_placeholder: obs_n}
        self.nervenetplus_batch_pos = nervenetplus_batch_pos

        return feed_dict, nervenetplus_batch_pos

    def build_update_op_preprocess(self):
        '''
            @brief: The preprocess that is shared by trpo, ppo and vpg updates
        '''
        # the input placeholders for the input
        self.action_placeholder = tf.placeholder(
            tf.float32, [None, self.action_size],
            name='action_sampled_in_rollout'
        )
        self.advantage_placeholder = tf.placeholder(
            tf.float32, [None], name='advantage_value'
        )
        self.oldaction_dist_mu_placeholder = tf.placeholder(
            tf.float32, [None, self.action_size], name='old_act_dist_mu'
        )
        self.oldaction_dist_logstd_placeholder = tf.placeholder(
            tf.float32, [None, self.action_size], name='old_act_dist_logstd'
        )
        self.batch_size_float_placeholder = tf.placeholder(
            tf.float32, [], name='batch_size_float'
        )

        # the adaptive kl penalty
        if self.args.use_kl_penalty:
            self.kl_lambda_placeholder = tf.placeholder(tf.float32, [],
                                                        name='kl_lambda')

        # what are the probabilities of taking self.action, given new and old
        # distributions
        self.log_p_n = utils.gauss_log_prob(
            self.action_dist_mu,
            self.action_dist_logstd,
            self.action_placeholder
        )
        self.log_oldp_n = utils.gauss_log_prob(
            self.oldaction_dist_mu_placeholder,
            self.oldaction_dist_logstd_placeholder,
            self.action_placeholder
        )

        self.ratio = tf.exp(self.log_p_n - self.log_oldp_n)

        # the kl divergence between the old and new action
        self.kl = utils.gauss_KL(
            self.oldaction_dist_mu_placeholder,
            self.oldaction_dist_logstd_placeholder,
            self.action_dist_mu,
            self.action_dist_logstd
        )
        self.kl = self.kl / self.batch_size_float_placeholder

        # the entropy
        self.ent = utils.gauss_ent(
            self.action_dist_mu, self.action_dist_logstd
        )
        self.ent = self.ent / self.batch_size_float_placeholder

    def build_update_op_postprocess(self):
        '''
            @brief: The postprocess that is shared by trpo, ppo and vpg updates
        '''
        raise NotImplementedError

    def run(self):
        '''
            @brief:
                This is the standard function to be called by the
                "multiprocessing.Process"
        '''
        raise NotImplementedError

    def build_models(self):
        '''
            @brief:
                This is the function where the rollout agents and trpo agent
                build their networks, set up the placeholders, and gather the
                variable list.
        '''
        raise NotImplementedError

    def prepared_network_feeddict(self, data_dict):
        '''
            @brief:
                For the general policy network and graph network, we have
                different format, as we batch the network in the policy network
            @return:
                The feed_dict structure
        '''
        raise NotImplementedError

    def get_sess(self):
        return self.session

    def get_iteration_count(self):
        return self.session.run(self.iteration)

    def get_experiment_name(self):
        '''
            @brief:
                this is the unique id of the experiments. it might be useful if
                we are running several tasks on the server
        '''
        return self.args.task + '_' + self.args.time_id _ABS_BASE_PATH = init_path.get_abs_base_dir()
def construct_graph_input_feeddict(node_info,
                                   obs_n,
                                   receive_idx,
                                   send_idx,
                                   node_type_idx,
                                   inverse_node_type_idx,
                                   output_type_idx,
                                   inverse_output_type_idx,
                                   last_batch_size,
                                   request_data=['ob', 'idx']):
def _construct_index(node_info, batch_size):
def _get_obs(node_info, raw_ob, batch_size):
def _add_offset(idx, offset, batch_size, if_reshape):
def _add_inverse_offset(extype_offset, intype_offset, self_offset, batch_size): class parallel_rollout_master_agent():

    def __init__(self, args, observation_size, action_size):
        '''
            @brief:
                the master agent has several actors (or samplers) to do the
                sampling for it.
        '''
        self.args = args
        self.observation_size = observation_size
        self.action_size = action_size

        # init the running means
        self.reset_running_mean_info()
        self.load_running_means()

        # init the multiprocess actors
        self.tasks = multiprocessing.JoinableQueue()
        self.results = multiprocessing.Queue()
        self.init_actors()

        # we will start by running 20,000 / 1000 = 20 episodes for the first
        # iteration
        self.average_timesteps_in_episode = args.max_pathlength

    def reset_running_mean_info(self):
        self.running_mean_info = {
            'mean': 0.0,
            'variance': 1,
            'step': 0.01,
            'square_sum': 0.01,
            'sum': 0.0
        }

    def rollout(self):
        # step 1: ask the samplers to generate rollouts
        rollout_data = self.ask_for_rollouts()

        # step 2: update the running means
        self.update_running_means(rollout_data)

        # step 3: synchronize the filter statistics for all samplers
        for i_agent in xrange(self.args.num_threads + 1):
            self.tasks.put(
                (parallel_util.AGENT_SYNCHRONIZE_FILTER, self.running_mean_info)
            )
        self.tasks.join()

        # tell the trpo agent about the information of normalizer
        rollout_data.append(self.running_mean_info)
        return rollout_data

    def set_policy_weights(self, parameters):
        for i_agent in xrange(self.args.num_threads + 1):
            self.tasks.put((parallel_util.AGENT_SET_POLICY_WEIGHTS, parameters))
        self.tasks.join()

    def end(self):
        for i_agent in xrange(self.args.num_threads + 1):
            self.tasks.put((parallel_util.END_ROLLOUT_SIGNAL, None))

    def ask_for_rollouts(self):
        '''
            @brief:
                Run the experiments until a total of @timesteps_per_batch
                timesteps are collected.
        '''
        num_timesteps_received = 0
        timesteps_needed = self.args.timesteps_per_batch
        rollout_data = []

        while True:
            # calculate the expected number of episodes needed
            num_rollouts = int(
                np.ceil(timesteps_needed / self.average_timesteps_in_episode)
            )

            if self.args.test:
                num_rollouts = self.args.test

            # request episodes
            for i_agent in xrange(num_rollouts):
                self.tasks.put((parallel_util.START_SIGNAL, None))
            self.tasks.join()

            # collect episodes
            for _ in range(num_rollouts):
                traj_episode = self.results.get()
                rollout_data.append(traj_episode)
                num_timesteps_received += len(traj_episode['rewards'])

            # update average timesteps per episode
            self.average_timesteps_in_episode = \
                float(num_timesteps_received) / len(rollout_data)

            timesteps_needed = self.args.timesteps_per_batch - \
                num_timesteps_received

            if timesteps_needed <= 0 or self.args.test:
                break

        logger.info('Rollouts generating ...')
        logger.info('{} time steps from {} episodes collected'.format(
            num_timesteps_received, len(rollout_data)
        ))

        return rollout_data

    def update_running_means(self, paths):
        # collect the info
        new_sum = 0.0
        new_step_sum = 0.0
        new_sq_sum = 0.0
        # from util import fpdb; fpdb.fpdb().set_trace()
        for path in paths:
            raw_obs = path['raw_obs']  # shape: [timesteps, num_observation]
            new_sum += raw_obs.sum(axis=0)  # shape [num_observation]
            new_step_sum += raw_obs.shape[0]
            new_sq_sum += (np.square(raw_obs)).sum(axis=0)

        # update the parameters
        self.running_mean_info['sum'] += new_sum
        self.running_mean_info['square_sum'] += new_sq_sum
        self.running_mean_info['step'] += new_step_sum
        self.running_mean_info['mean'] = \
            self.running_mean_info['sum'] / self.running_mean_info['step']

        self.running_mean_info['variance'] = np.maximum(
            self.running_mean_info['square_sum'] /
            self.running_mean_info['step'] -
            np.square(self.running_mean_info['mean']),
            1e-2
        )

    def load_running_means(self):
        # load the observation running mean
        if self.args.ckpt_name is not None:
            base_path = os.path.join(init_path.get_base_dir(), 'checkpoint')
            logger.info('[LOAD_CKPT] loading observation normalizer info')
            self.running_mean_info = model_saver.load_numpy_model(
                os.path.join(base_path,
                             self.args.ckpt_name + '_normalizer.npy'),
                numpy_var_list=self.running_mean_info
            )
            self.running_mean_info['transfer_env'] = self.args.transfer_env

        if not self.args.transfer_env == 'Nothing2Nothing':
            if self.args.mlp_raw_transfer == 0:
                assert 'shared' in self.args.gnn_embedding_option

            ienv, oenv = [env + '-v1'
                          for env in self.args.transfer_env.split('2')]
            self.running_mean_info = \
                structure_mapper.map_transfer_env_running_mean(
                    ienv, oenv, self.running_mean_info,
                    self.observation_size,
                    self.args.gnn_node_option, self.args.root_connection_option,
                    self.args.gnn_output_option, self.args.gnn_embedding_option
                )

    def set_running_means(self, new_running_mean):
        self.running_mean_info = copy.deepcopy(new_running_mean)

    def init_actors(self):
        '''
            @brief: init the actors and start the multiprocessing
        '''

        self.actors = []
        # the main actor who can sampler video result
        self.actors.append(
            rollout_agent(self.args,
                          self.observation_size,
                          self.action_size,
                          self.tasks,
                          self.results,
                          0,
                          self.args.monitor,
                          init_filter_parameters=self.running_mean_info)
        )

        # the sub actor that only do the sampling
        for i in xrange(self.args.num_threads):
            self.actors.append(
                rollout_agent(
                    self.args,
                    self.observation_size,
                    self.action_size,
                    self.tasks,
                    self.results,
                    i + 1,
                    False,
                    init_filter_parameters=self.running_mean_info
                )
            )

        for i_actor in self.actors:
            i_actor.start() def build_baseline_network(args, session, name_scope,
                           observation_size, gnn_placeholder_list,
                           obs_placeholder): _BASE_PATH = init_path.get_abs_base_dir()
def prepare_network_feeddict(data_dict, is_nervenet, baseline_network,
                             node_info, current_idx_dict, nervenetplus=False,
                             gnn_num_prop_steps=-1):
def prepare_policy_network_feeddict(observations, is_nervenet,
                                    node_info, current_idx_dict,
                                    rollout_data, nervenetplus,
                                    gnn_num_prop_steps):
def generate_advantage(data_dict, baseline_network): ELLIPSOID_X_RATIO = 1
ELLIPSOID_Y_RATIO = 8
ELLIPSOID_Z_RATIO = 4
CYLINDER_R_RATIO = 1
CYLINDER_H_RATIO = 6
RAD_TO_DEG = 57.29577951308232
BASE_PATH = init_path.get_abs_base_dir()
                R = np.array([[1,   0,    0],
                              [0, cos, -sin],
                              [0, sin,  cos]])
                R = np.array([[ cos, 0, sin],
                              [   0, 1,   0],
                              [-sin, 0, cos]])
                R = np.array([[cos, -sin, 0],
                              [sin,  cos, 0],
                              [  0,    0, 1]])
            R = np.zeros((3, 3))
            R = np.matmul(R_z, R_y)
            R = np.matmul(R, R_z)
            R = H[0:3, 0:3]
            H = np.zeros((4, 4))
            H[0:3, 0:3] = R
            H[0:3, 3] = d
            H[3, 3] = 1
            P = np.hstack((p, np.array(1)))
        N, _ = adj_matrix.shape
        N, _ = adj_matrix.shape
        N, _ = adj_matrix.shape
        N, _ = adj_matrix.shape
    N, _ = adj_matrix.shape
def save_generated_xml(root, filename):
def fish_xml_generator(adj_matrix, node_attr, options=None, filename=None):
    def add_mujoco_header(root):
    def add_mujoco_options(root, options):
    def add_worldbody(root, adj_matrix, node_attr_list):
        def worldbody_preliminary(worldbody):
        def torso_preliminary(torso):
        def rotation_matrix(axis, angle):
        def euler_rotation(angle1, angle2, angle3):
        def ellipsoid_line_intersection(a, b, c, dx, dy, dz):
        def homogeneous_inverse(H):
        def homogeneous_transform(R, d):
        def homogeneous_representation(p):
        def ellipsoid_normal_vec(a, b, c, p):
        def angle_between(vec1, vec2):
        def ellipsoid_boundary_pt(a, b, c):
    def add_mujoco_actuator(root, adj_matrix):
    def add_mujoco_sensor(root):
def walker_xml_generator(adj_matrix, node_attr, options=None, filename=None):
    def add_mujoco_header(root):
    def add_mujoco_options(root, options):
    def add_worldbody(root, adj_matrix, node_attr_list):
        def worldbody_preliminary(worldbody):
        def torso_preliminary(torso):
    def add_mujoco_sensor(root, adj_matrix):
    def add_mujoco_actuator(root, adj_matrix, node_attr):
def hopper_xml_generator(adj_matrix, node_attr, options=None, filename=None):
    def add_mujoco_header(root):
    def add_mujoco_options(root, options):
    def add_worldbody(root, adj_matrix, node_attr_list):
        def worldbody_preliminary(worldbody):
        def torso_preliminary(torso):
    def add_mujoco_sensor(root, adj_matrix):
    def add_mujoco_actuator(root, adj_matrix):
def cheetah_xml_generator(adj_matrix, node_attr, options=None, filename=None):
    def add_mujoco_header(root):
    def add_mujoco_options(root, options):
    def add_worldbody(root, adj_matrix, node_attr_list):
        def worldbody_preliminary(worldbody):
        def torso_preliminary(torso):
    def add_mujoco_sensor(root, adj_matrix):
    def add_mujoco_actuator(root, adj_matrix, node_attr):
def gen_test_adj_mat(task='fish', hinge_opt=7, shape=(5, 5)):
def gen_test_node_attr(task='fish', node_num=5, discrete_rv=True):
    def root_default(task):
    def gen_one_attr(task, discrete_rv=True):
def get_initial_settings(): class Node:
class Tree:
class Species:
    def __init__(self, node_id, attr):
    def set_id(self, new_id):
    def get_all_node_id(self):
    def get_all_descendents(self):
    def get_node_id(self):
    def get_attr(self):
    def set_attr(self, new_attr):
    def get_child_list(self):
    def add_child(self, node):
    def is_same_size(self, other_node):
    def make_symmetric(self, task='fish', discrete=True):
    def remove_child(self, node_id):
    def printinfo(self):
    def __eq__(self, other):
def node_count(node, counter_start=0):
    def __init__(self, node=None):
    def get_root(self):
    def sample_node(self, p=0.5):
        def check_node(node, p=0.5):
    def sample_leaf(self, p=0.5):
        def find_leaf(node, total_res):
    def get_lvl_node(self, lvl=0):
    def get_all_nodes(self):
    def get_pc_relation(self):
    def add_sub_tree(self, parent, child):
    def remove_sub_tree(self, parent, child):
    def to_mujoco_format(self):
    def __init__(self, 
                 args,
                 body_num=3, 
                 discrete=True,
                 allow_hierarchy=True,
                 filter_ratio=0.5):
    def mutate(self):
    def generate_one_node(self, only_one=False):
    def perturb_add(self, no_selfcross=False):
    def perturb_remove(self):
    def perturb_attr(self):
    def get_gene(self):
    def get_xml(self):
        N = self.total_num
NEW_STRUCT_OPTS = ['l1-basic', 'l2-symmetry'] PATH = init_path.get_abs_base_dir()
WEIGHT_NAMES = ['policy_logstd']  # 'policy_0/w:0', 'policy_output'
INHERIT_LIST = [
    'policy_weights', 'baseline_weights', 'running_mean_info',
    'var_name', 'node_info', 'observation_size', 'action_size', 'lr'
]
GENERATED_LIST = \
    ['xml_str', 'adj_matrix', 'node_attr', 'debug_info', 'PrtID',
     'SpcID', 'node_info']
INVALID_LIST = [
    'stats', 'start_time', 'rollout_time', 'agent_id',
    'env_name', 'rank_info', 'LastRwd', 'AvgRwd',
]
ALL_LIST = INHERIT_LIST + GENERATED_LIST + INVALID_LIST
MUTATION_METHOD = ['add', 'delete', 'perturb']
def weight_s2s_update(p_species, c_species, node_order, is_nervenet):
def running_stats_s2s_update(p_species, c_species, node_order, is_nervenet):
def process_inherited_info(raw_species_info, current_species_format,
                           is_nervenet):
def mutate_species(p_species, mutation_method, new_spc_struct=False):
def node_order_assertion(c_species, p_species):
def parse_policy_logstd(c_species, p_species, c_var_id, p_var_id, node_order):
def parse_embedding(c_species, p_species, c_var_id, p_var_id, node_order):
import numpy as np
import tensorflow as tf
import init_path
import random
import os
import time
from util import utils
from util import logger
from util import parallel_util
from util import ob_normalizer
from .agent import base_agent
from graph_util import graph_data_util
from .rollout_master_agent import parallel_rollout_master_agent
from util import agent_util
from graph_util import gnn_io_util
from env import model_gen
from env import hierarchy_model
from lxml import etree
from evolution import species2species
            from env import fish_env_wrapper
            from env import walker_env_wrapper
        # samples the guassian distribution
        act = action_dist_mu + np.exp(action_dist_logstd) * \
            self._npr.randn(*action_dist_logstd.shape)
        act = act.ravel()
        return act, action_dist_mu, action_dist_logstd

    def build_env(self, received_data=None):
        assert 'evo' in self.args.task

        allow_monitor = 0
        if received_data is None:
            if self.args.new_species_struct:
                if self.args.more_body_nodes_at_start:
                    body_num = random.randint(3, 6)
                else:
                    body_num = 3
                species = hierarchy_model.Species(self.args, body_num=body_num)
                adj_matrix, node_attr = species.get_gene()
                xml_struct, xml_str = species.get_xml()
            else:
                adj_matrix, node_attr, xml_str = model_gen.get_initial_settings()
                xml_str = etree.tostring(xml_str, pretty_print=True)

        else:
            # if we set the visualization open
            if 'rank_info' in received_data:
                generation, rank = [int(info) for info
                                    in received_data['rank_info'].split('_')]
                if generation % self.args.species_visualize_freq == 0 and \
                        rank <= self.args.visualize_top_species:
                    allow_monitor = True

            if self.args.new_species_struct:
                species = received_data['species']

            adj_matrix = received_data['adj_matrix']
            node_attr = received_data['node_attr']
            xml_str = received_data['xml_str']
            # self.current_lr = received_data['lr']

        # create the environment for this agent
        if 'fish' in self.args.task:
            self.env = fish_env_wrapper.dm_evofish3d_wrapper(
                self.args, self._npr.randint(0, 99999), allow_monitor,
                adj_matrix=adj_matrix, xml_str=xml_str
            )
        elif 'walker' in self.args.task or \
             'cheetah' in self.args.task or \
             'hopper' in self.args.task:
            self.env = walker_env_wrapper.dm_evowalker_wrapper(
                self.args, self._npr.randint(0, 99999), allow_monitor,
                adj_matrix=adj_matrix, xml_str=xml_str
            )
        else:
            raise NotImplementedError

        if allow_monitor and received_data is not None:
            self.env.set_output_dir(received_data['video_save_path'])

        self.is_dm_env = True
        self.action_size = self.env.get_action_size()
        self.observation_size = self.env.get_observation_size()

        if self.args.new_species_struct:
            self.species = species

        self.adj_matrix = adj_matrix
        self.node_attr = node_attr
        self.xml_str = xml_str

    def run(self):
        '''
            @brief:
                this is the standard function to be called by the
                "multiprocessing.Process"

            @NOTE:
                check the parallel_util.py for definitions
        '''
        # the main training process
        received_signal, received_data = self.task_q.get()

        if received_signal == parallel_util.END_SIGNAL:
            # kill the learner
            self.task_q.task_done()

        elif received_signal == parallel_util.AGENT_EVOLUTION_START:
            self.build_env(received_data)
            self.build_models()

            training_stats = self.rollout_and_train()
            self.task_q.task_done()
            species_data = self.get_species_info(training_stats, received_data)
            self.result_q.put(species_data)

        elif received_signal == parallel_util.AGENT_EVOLUTION_TRAIN:
            self.build_env(received_data)
            self.build_models(received_data)

            training_stats = self.rollout_and_train()
            self.task_q.task_done()
            species_data = self.get_species_info(training_stats, received_data)
            self.result_q.put(species_data)

    def build_models(self, received_data=None):
        '''
            @brief:
                this is the function where the rollout agents and optimization
                agent build their networks, set up the placeholders, and gather
                the variable list.
        '''
        # make sure that the agent has a session
        self.start_time = time.time()
        self.build_session()

        self.build_policy_network(adj_matrix=self.adj_matrix,
                                  node_attr=self.node_attr)

        self.baseline_network, self.target_return_placeholder, \
            self.raw_obs_placeholder = \

            agent_util.build_baseline_network(