# Path: util/logger.py
# 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']
# 
# Path: util/utils.py
# 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):
# 
# Path: graph_util/nervenetplus_util.py
# def nervenetplus_step_assign(rollout_data, nstep):
# 
# Path: graph_util/graph_data_util.py
# _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):


import init_path
import numpy as np
from util import logger
from util import utils
from graph_util import nervenetplus_util
from graph_util import graph_data_util
                             node_info, current_idx_dict, nervenetplus=False,
                             gnn_num_prop_steps=-1):
    # feed_dict = {}

    # step 1: For ggnn and fc policy, we have different obs in feed_dict
    observations = np.concatenate([path["obs"] for path in data_dict])
    feed_dict, nervenetplus_batch_pos = prepare_policy_network_feeddict(
        observations, is_nervenet, node_info, current_idx_dict,
        data_dict, nervenetplus, gnn_num_prop_steps
    )

    # step 2: prepare the advantage function, old action / obs needed for
    # trpo / ppo updates
    advant_n = generate_advantage(data_dict, baseline_network)
    action_dist_mu = np.concatenate(
        [path["action_dists_mu"] for path in data_dict]
    )
    action_dist_logstd = np.concatenate(
        [path["action_dists_logstd"] for path in data_dict]
    )
    action_n = np.concatenate([path["actions"] for path in data_dict])
    feed_dict.update(
        {'action_placeholder': action_n,
         'advantage_placeholder': advant_n,
         'oldaction_dist_mu_placeholder': action_dist_mu,
         'oldaction_dist_logstd_placeholder': action_dist_logstd,
         'batch_size_float_placeholder': np.array(float(len(observations)))}
    )

    # step 3: feed_dict to update value function

# Path: util/logger.py
# 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']
# 
# Path: util/utils.py
# 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):
# 
# Path: graph_util/nervenetplus_util.py
# def nervenetplus_step_assign(rollout_data, nstep):
# 
# Path: graph_util/graph_data_util.py
# _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):


# Path: graph_util/gnn_io_util.py
import init_path
import numpy as np
from util import logger
from util import utils
from graph_util import nervenetplus_util
from graph_util import graph_data_util

                             node_info, current_idx_dict, nervenetplus=False,
                             gnn_num_prop_steps=-1):
    # feed_dict = {}

    # step 1: For ggnn and fc policy, we have different obs in feed_dict
    observations = np.concatenate([path["obs"] for path in data_dict])
    feed_dict, nervenetplus_batch_pos = prepare_policy_network_feeddict(
        observations, is_nervenet, node_info, current_idx_dict,
        data_dict, nervenetplus, gnn_num_prop_steps
    )

    # step 2: prepare the advantage function, old action / obs needed for
    # trpo / ppo updates
    advant_n = generate_advantage(data_dict, baseline_network)
    action_dist_mu = np.concatenate(
        [path["action_dists_mu"] for path in data_dict]
    )
    action_dist_logstd = np.concatenate(
        [path["action_dists_logstd"] for path in data_dict]
    )
    action_n = np.concatenate([path["actions"] for path in data_dict])
    feed_dict.update(
        {'action_placeholder': action_n,
         'advantage_placeholder': advant_n,
         'oldaction_dist_mu_placeholder': action_dist_mu,
         'oldaction_dist_logstd_placeholder': action_dist_logstd,
         'batch_size_float_placeholder': np.array(float(len(observations)))}
    )

    # step 3: feed_dict to update value function

    target_return = \