import numpy as np
import tensorflow as tf
import gym
import time
import safe_rl.pg.trust_region as tro
from safe_rl.pg.agents import PPOAgent, TRPOAgent, CPOAgent, PPO_Agent_with_Mu
from safe_rl.pg.buffer import CPOBuffer, CPOBuffer_with_vio
from safe_rl.pg.network import count_vars, \
                               get_vars, \
                               mlp_actor_critic, mlp_actor_critic_with_lam,\
                               placeholders, \
                               placeholders_from_spaces
from safe_rl.pg.utils import values_as_sorted_list
from safe_rl.utils.logx import EpochLogger
from safe_rl.utils.mpi_tf import MpiAdamOptimizer, sync_all_params
from safe_rl.utils.mpi_tools import mpi_fork, proc_id, num_procs, mpi_sum, mpi_statistics_scalar

# Multi-purpose agent runner for policy optimization algos 
# (PPO, TRPO, their primal-dual equivalents, CPO)
def run_polopt_agent_da(env_fn,
                        agent=PPOAgent(),
                        actor_critic=mlp_actor_critic_with_lam,
                        ac_kwargs=dict(),
                        seed=0,
                        render=False,
                        rew_scale=0.1,
                        # Experience collection:
                        steps_per_epoch=4000,
                        epochs=50,
                        max_ep_len=1000,
                        # Discount factors:
                        gamma=0.99,
                        lam=0.95,
                        cost_gamma=0.99,
                        cost_lam=0.95,
                        # Policy learning:
                        ent_reg=0.,
                        # Cost constraints / penalties:
                        cost_lim=25,
                        penalty_init=1.,
                        penalty_lr=5e-2,
                        # KL divergence:
                        target_kl=0.01,
                        # Value learning:
                        vf_lr=3e-4,
                        vf_iters=80,
                        # Mu learning:
                        dual_ascent_interval=1,
                        # Logging:
                        logger=None,
                        logger_kwargs=dict(),
                        save_freq=1,
                        lr_decay=True
                        ):


    #=========================================================================#
    #  Prepare logger, seed, and environment in this process                  #
    #=========================================================================#

    logger = EpochLogger(**logger_kwargs) if logger is None else logger
    logger.save_config(locals())

    seed += 10000 * proc_id()
    tf.set_random_seed(seed)
    np.random.seed(seed)

    # Create summary writer
    logs_dir = logger_kwargs.get('output_dir') + '/logs'
    summary_writer = tf.summary.FileWriter(logs_dir, tf.get_default_graph())

    env = env_fn()

    agent.set_logger(logger)

    #=========================================================================#
    #  Create computation graph for actor and critic (not training routine)   #
    #=========================================================================#

    # Share information about action space with policy architecture
    ac_kwargs['action_space'] = env.action_space

    # Inputs to computation graph from environment spaces
    x_ph, a_ph = placeholders_from_spaces(env.observation_space, env.action_space)

    # Inputs to computation graph for batch data
    adv_ph, cadv_ph, ret_ph, cret_ph, logp_old_ph, vio_ph = placeholders(*(None for _ in range(6)))

    # Inputs to computation graph for special purposes
    surr_cost_rescale_ph = tf.placeholder(tf.float32, shape=())
    cur_cost_ph = tf.placeholder(tf.float32, shape=())

    # Outputs from actor critic
    ac_outs = actor_critic(x_ph, a_ph, **ac_kwargs)
    pi, logp, logp_pi, pi_info, pi_info_phs, d_kl, ent, v, vc, multiplier = ac_outs
    tf.summary.histogram('Optimizers/VCVals', vc)
    merged_summary = tf.summary.merge_all()

    # Organize placeholders for zipping with data from buffer on updates
    buf_phs = [x_ph, a_ph, adv_ph, cadv_ph, ret_ph, cret_ph, logp_old_ph, vio_ph]
    buf_phs += values_as_sorted_list(pi_info_phs)

    # Organize symbols we have to compute at each step of acting in env
    get_action_ops = dict(pi=pi, 
                          v=v, 
                          logp_pi=logp_pi,
                          pi_info=pi_info)

    # If agent is reward penalized, it doesn't use a separate value function
    # for costs and we don't need to include it in get_action_ops; otherwise we do.
    if not(agent.reward_penalized):
        get_action_ops['vc'] = vc

    # Count variables
    var_counts = tuple(count_vars(scope) for scope in ['pi', 'vf', 'vc'])
    logger.log('\nNumber of parameters: \t pi: %d, \t v: %d, \t vc: %d\n'%var_counts)

    # Make a sample estimate for entropy to use as sanity check
    approx_ent = tf.reduce_mean(-logp)


    #=========================================================================#
    #  Create replay buffer                                                   #
    #=========================================================================#

    # Obs/act shapes
    obs_shape = env.observation_space.shape
    act_shape = env.action_space.shape

    # Experience buffer
    local_steps_per_epoch = int(steps_per_epoch / num_procs())
    pi_info_shapes = {k: v.shape.as_list()[1:] for k,v in pi_info_phs.items()}
    value_cost_lim = cost_lim / float(max_ep_len) / (1-cost_lam)
    buf = CPOBuffer_with_vio(local_steps_per_epoch,
                            obs_shape,
                            act_shape,
                            pi_info_shapes,
                            gamma,
                            lam,
                            cost_gamma,
                            cost_lam,
                            cost_lim)


    #=========================================================================#
    #  Create computation graph for penalty learning, if applicable           #
    #=========================================================================#

    # if agent.use_penalty:
    #     with tf.variable_scope('penalty'):
    #         # param_init = np.log(penalty_init)
    #         param_init = np.log(max(np.exp(penalty_init)-1, 1e-8))
    #         penalty_param = tf.get_variable('penalty_param',
    #                                       initializer=float(param_init),
    #                                       trainable=agent.learn_penalty,
    #                                       dtype=tf.float32)
    #     # penalty = tf.exp(penalty_param)
    #     penalty = tf.nn.softplus(penalty_param)
    #
    # if agent.learn_penalty:
    #     if agent.penalty_param_loss:
    #         penalty_loss = -penalty_param * (cur_cost_ph - cost_lim)
    #     else:
    #         penalty_loss = -penalty * (cur_cost_ph - cost_lim) # penalty is an adaptive coefficient
    #     train_penalty = MpiAdamOptimizer(learning_rate=penalty_lr).minimize(penalty_loss)

    # if agent.dual_ascent:




    #=========================================================================#
    #  Create computation graph for policy learning                           #
    #=========================================================================#

    # Likelihood ratio
    ratio = tf.exp(logp - logp_old_ph)

    # Surrogate advantage / clipped surrogate advantage
    if agent.clipped_adv:
        min_adv = tf.where(adv_ph>0, 
                           (1+agent.clip_ratio)*adv_ph, 
                           (1-agent.clip_ratio)*adv_ph
                           )
        surr_adv = tf.reduce_mean(tf.minimum(ratio * adv_ph, min_adv))
    else:
        surr_adv = tf.reduce_mean(ratio * adv_ph)

    # old: Surrogate cost now: penalty
    surr_cadv = ratio * cadv_ph
    penalty = tf.reduce_mean(tf.multiply(tf.stop_gradient(multiplier), surr_cadv)) #todo:why use cadv here?
                                                                                   # because the grad is from IS ratio
    # complementary slackness cost
    clipped_vio = tf.clip_by_value(vio_ph, -0.1, 100)
    cs_cost = tf.reduce_mean(tf.multiply(multiplier, tf.stop_gradient(clipped_vio)))

    # Create policy objective function, including entropy regularization
    pi_objective = surr_adv + ent_reg * ent
    # lam_objective = tf.reduce_mean(tf.multiply(mu, tf.stop_gradient(vc)))

    # Possibly include surr_cost in pi_objective
    if agent.objective_penalized or agent.dual_ascent:
        lagrangian = pi_objective - penalty # max min [-(f(x)-mu * g(x))]

    # Loss function for pi is negative of pi_objective
    pi_loss = - lagrangian # todo refine if

    # Loss function for mu
    mu_loss = -cs_cost

    # Optimizer-specific symbols
    if agent.trust_region:

        # Symbols needed for CG solver for any trust region method
        pi_params = get_vars('pi')
        flat_g = tro.flat_grad(pi_loss, pi_params)
        v_ph, hvp = tro.hessian_vector_product(d_kl, pi_params)
        if agent.damping_coeff > 0:
            hvp += agent.damping_coeff * v_ph

        # Symbols needed for CG solver for CPO only
        flat_b = tro.flat_grad(penalty, pi_params) # todo: change to penalty, not confirm right

        # Symbols for getting and setting params
        get_pi_params = tro.flat_concat(pi_params)
        set_pi_params = tro.assign_params_from_flat(v_ph, pi_params)

        training_package = dict(flat_g=flat_g,
                                flat_b=flat_b,
                                v_ph=v_ph,
                                hvp=hvp,
                                get_pi_params=get_pi_params,
                                set_pi_params=set_pi_params)

    elif agent.first_order:

        if lr_decay:
            local_step = tf.Variable(0.0, trainable=False)
            pi_lr = tf.train.polynomial_decay(learning_rate=agent.pi_lr, global_step=local_step,
                                              decay_steps=epochs * agent.pi_iters, end_learning_rate=1e-6)
        else:
            pi_lr = agent.pi_lr
            # Optimizer for first-order policy optimization
        train_pi = MpiAdamOptimizer(learning_rate=pi_lr).minimize(pi_loss)

        if lr_decay:
            local_step = tf.Variable(0.0, trainable=False)
            mu_lr = tf.train.polynomial_decay(learning_rate=agent.mu_lr, global_step=local_step,
                                              decay_steps=int(epochs * agent.pi_iters / agent.ascent_interval),
                                              end_learning_rate=1e-6)
        else:
            mu_lr = agent.mu_lr
            # Optimizer for first-order policy optimization

        train_mu = MpiAdamOptimizer(learning_rate=mu_lr).minimize(mu_loss)

        # Prepare training package for agent
        training_package = dict(train_pi=train_pi, train_mu=train_mu)

    else:
        raise NotImplementedError

    # Provide training package to agent
    training_package.update(dict(pi_loss=pi_loss,
                                 mu_loss=mu_loss,
                                 penalty=penalty,
                                 d_kl=d_kl, 
                                 target_kl=target_kl,
                                 cost_lim=cost_lim))
    agent.prepare_update(training_package)

    #=========================================================================#
    #  Create computation graph for value learning                            #
    #=========================================================================#

    # Value losses
    v_loss = tf.reduce_mean((ret_ph - v)**2)
    vc_loss = tf.reduce_mean((cret_ph - vc)**2)

    # If agent uses penalty directly in reward function, don't train a separate
    # value function for predicting cost returns. (Only use one vf for r - p*c.)
    if agent.reward_penalized:
        total_value_loss = v_loss
    else:
        total_value_loss = v_loss + vc_loss

    # Optimizer for value learning
    if lr_decay:
        local_step_vf = tf.Variable(0.0, trainable=False)
        vf_lr = tf.train.polynomial_decay(learning_rate=vf_lr, global_step=local_step_vf,
                                          decay_steps=epochs * vf_iters, end_learning_rate=1e-6)

    train_vf = MpiAdamOptimizer(learning_rate=vf_lr).minimize(total_value_loss)


    #=========================================================================#
    #  Create session, sync across procs, and set up saver                    #
    #=========================================================================#

    sess = tf.Session()
    sess.run(tf.global_variables_initializer())

    # Sync params across processes
    sess.run(sync_all_params())

    # Setup model saving
    logger.setup_tf_saver(sess, inputs={'x': x_ph}, outputs={'pi': pi, 'v': v, 'vc': vc, 'multiplier': multiplier})


    #=========================================================================#
    #  Provide session to agent                                               #
    #=========================================================================#
    agent.prepare_session(sess)


    #=========================================================================#
    #  Create function for running update (called at end of each epoch)       #
    #=========================================================================#

    def update():
        cur_cost = logger.get_stats('EpCost')[0]
        c = cur_cost - cost_lim
        if c > 0 and agent.cares_about_cost:
            logger.log('Warning! Safety constraint is already violated.', 'red')

        #=====================================================================#
        #  Prepare feed dict                                                  #
        #=====================================================================#

        inputs = {k:v for k,v in zip(buf_phs, buf.get())}
        inputs[surr_cost_rescale_ph] = logger.get_stats('EpLen')[0]
        inputs[cur_cost_ph] = cur_cost

        #=====================================================================#
        #  Make some measurements before updating                             #
        #=====================================================================#

        measures = dict(LossPi=pi_loss,
                        LossMu=mu_loss,
                        Penalty=penalty,
                        Violation=clipped_vio,
                        Multiplier=multiplier,
                        LossV=v_loss,
                        Entropy=ent)
        if not(agent.reward_penalized):
            measures['LossVC'] = vc_loss
        if agent.use_penalty:
            measures['Penalty'] = penalty

        pre_update_measures = sess.run(measures, feed_dict=inputs)
        summary = sess.run(merged_summary, feed_dict)
        summary_writer.add_summary(summary, global_step=epoch * steps_per_epoch)
        logger.store(**pre_update_measures)

        #=====================================================================#
        #  Update penalty if learning penalty                                 #
        #=====================================================================#
        # if agent.learn_penalty:
        #     sess.run(train_penalty, feed_dict={cur_cost_ph: cur_cost})

        #=====================================================================#
        #  Update policy and mu                                               #
        #=====================================================================#
        agent.update_pi(inputs)

        #=====================================================================#
        #  Update value function                                              #
        #=====================================================================#
        for _ in range(vf_iters):
            sess.run(train_vf, feed_dict=inputs)

        #=====================================================================#
        #  Make some measurements after updating                              #
        #=====================================================================#

        del measures['Entropy']
        measures['KL'] = d_kl

        post_update_measures = sess.run(measures, feed_dict=inputs)
        deltas = dict()
        for k in post_update_measures:
            if k in pre_update_measures:
                deltas['Delta'+k] = post_update_measures[k] - pre_update_measures[k]
        logger.store(KL=post_update_measures['KL'], **deltas)




    #=========================================================================#
    #  Run main environment interaction loop                                  #
    #=========================================================================#

    start_time = time.time()
    o, r, d, c, ep_ret, ep_cost, ep_len = env.reset(), 0, False, 0, 0, 0, 0
    cur_penalty = 0
    cum_cost = 0

    for epoch in range(epochs):

        # if agent.use_penalty:
        #     cur_penalty = sess.run(penalty)
        cur_penalty = 0 # todo: remove relevant with cur_penalty

        for t in range(local_steps_per_epoch):

            # Possibly render
            if render and proc_id()==0 and t < 1000:
                env.render()
            
            # Get outputs from policy
            get_action_outs = sess.run(get_action_ops, 
                                       feed_dict={x_ph: o[np.newaxis]})
            a = get_action_outs['pi']
            v_t = get_action_outs['v']
            vc_t = get_action_outs.get('vc', 0)  # Agent may not use cost value func
            logp_t = get_action_outs['logp_pi']
            pi_info_t = get_action_outs['pi_info']

            # Step in environment
            o2, r, d, info = env.step(a)

            # Include penalty on cost
            c = info.get('cost', 0)

            # Track cumulative cost over training
            cum_cost += c

            # save and log
            if agent.reward_penalized:
                r_total = r - cur_penalty * c
                r_total = r_total / (1 + cur_penalty)
                buf.store(o, a, r_total, v_t, 0, 0, logp_t, pi_info_t)
            else:
                buf.store(o, a, r, v_t, c, vc_t, logp_t, pi_info_t)
            logger.store(VVals=v_t, CostVVals=vc_t)

            o = o2
            ep_ret += r
            ep_cost += c
            ep_len += 1

            terminal = d or (ep_len == max_ep_len)
            if terminal or (t==local_steps_per_epoch-1):

                # If trajectory didn't reach terminal state, bootstrap value target(s)
                if d and not(ep_len == max_ep_len):
                    # Note: we do not count env time out as true terminal state
                    last_val, last_cval = 0, 0
                else:
                    feed_dict={x_ph: o[np.newaxis]}
                    if agent.reward_penalized:
                        last_val = sess.run(v, feed_dict=feed_dict)
                        last_cval = 0
                    else:
                        last_val, last_cval = sess.run([v, vc], feed_dict=feed_dict)
                buf.finish_path(last_val, last_cval)

                # Only save EpRet / EpLen if trajectory finished
                if terminal:
                    logger.store(EpRet=ep_ret, EpLen=ep_len, EpCost=ep_cost) # , EpVelocity=velo
                else:
                    print('Warning: trajectory cut off by epoch at %d steps.'%ep_len)

                # Reset environment
                o, r, d, c, ep_ret, ep_len, ep_cost = env.reset(), 0, False, 0, 0, 0, 0
                velo = 0

        # Save model
        if (epoch % save_freq == 0) or (epoch == epochs-1):
            logger.save_state({'env': env}, None)

        #=====================================================================#
        #  Run RL update                                                      #
        #=====================================================================#
        update()

        #=====================================================================#
        #  Cumulative cost calculations                                       #
        #=====================================================================#
        cumulative_cost = mpi_sum(cum_cost)
        cost_rate = cumulative_cost / ((epoch+1)*steps_per_epoch)

        #=====================================================================#
        #  Log performance and stats                                          #
        #=====================================================================#

        logger.log_tabular('Epoch', epoch)

        # Performance stats
        logger.log_tabular('EpRet', with_min_and_max=True)
        logger.log_tabular('EpCost', with_min_and_max=True)
        logger.log_tabular('EpLen', average_only=True)
        logger.log_tabular('CumulativeCost', cumulative_cost)
        logger.log_tabular('CostRate', cost_rate)

        # Value function values
        logger.log_tabular('VVals', with_min_and_max=True)
        logger.log_tabular('CostVVals', with_min_and_max=True)

        # Pi loss and change
        logger.log_tabular('LossPi', average_only=True)
        logger.log_tabular('DeltaLossPi', average_only=True)

        # Surr cost and change
        logger.log_tabular('LossMu', average_only=True)
        logger.log_tabular('DeltaLossMu', average_only=True)
        logger.log_tabular('Violation', with_min_and_max=True)
        logger.log_tabular('Multiplier', with_min_and_max=True)


        # V loss and change
        logger.log_tabular('LossV', average_only=True)
        logger.log_tabular('DeltaLossV', average_only=True)

        # Vc loss and change, if applicable (reward_penalized agents don't use vc)
        if not(agent.reward_penalized):
            logger.log_tabular('LossVC', average_only=True)
            logger.log_tabular('DeltaLossVC', average_only=True)

        if agent.use_penalty or agent.save_penalty:
            logger.log_tabular('Penalty', average_only=True)
            logger.log_tabular('DeltaPenalty', average_only=True)
        else:
            logger.log_tabular('Penalty', 0)
            logger.log_tabular('DeltaPenalty', 0)

        # Anything from the agent?
        agent.log()

        # Policy stats
        logger.log_tabular('Entropy', average_only=True)
        logger.log_tabular('KL', average_only=True)

        # Time and steps elapsed
        logger.log_tabular('TotalEnvInteracts', (epoch+1)*steps_per_epoch)
        logger.log_tabular('Time', time.time()-start_time)

        # Show results!
        logger.dump_tabular()

if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--agent', type=str, default='ppo_dual_ascent')
    parser.add_argument('--env', type=str, default='Safexp-PointGoal1-v0')
    parser.add_argument('--hid', type=int, default=64)
    parser.add_argument('--l', type=int, default=2)
    parser.add_argument('--gamma', type=float, default=0.99)
    parser.add_argument('--cost_gamma', type=float, default=0.99)
    parser.add_argument('--seed', '-s', type=int, default=0)
    parser.add_argument('--cpu', type=int, default=4)
    parser.add_argument('--steps', type=int, default=4000)
    parser.add_argument('--epochs', type=int, default=50)
    parser.add_argument('--len', type=int, default=1000)
    parser.add_argument('--cost_lim', type=float, default=10)
    parser.add_argument('--exp_name', type=str, default='runagent')
    parser.add_argument('--kl', type=float, default=0.01)
    parser.add_argument('--render', action='store_true')
    parser.add_argument('--reward_penalized', action='store_true')
    parser.add_argument('--objective_penalized', action='store_true')
    parser.add_argument('--learn_penalty', action='store_true')
    parser.add_argument('--penalty_param_loss', action='store_true')
    parser.add_argument('--entreg', type=float, default=0.)
    args = parser.parse_args()

    try:
        import safety_gym
    except:
        print('Make sure to install Safety Gym to use constrained RL environments.')

    # mpi_fork(args.cpu)  # run parallel code with mpi

    # Prepare logger
    from safe_rl.utils.run_utils import setup_logger_kwargs
    logger_kwargs = setup_logger_kwargs(args.exp_name, args.seed)

    # Prepare agent
    agent_kwargs = dict(reward_penalized=args.reward_penalized,
                        objective_penalized=args.objective_penalized,
                        learn_penalty=args.learn_penalty,
                        penalty_param_loss=args.penalty_param_loss)
    dual_ascent_kwargs = dict(reward_penalized=False,
                        objective_penalized=True,
                        learn_penalty=False,
                        penalty_param_loss=False,
                        dual_ascent=True)
    if args.agent=='ppo':
        agent = PPOAgent(**agent_kwargs)
    elif args.agent=='trpo':
        agent = TRPOAgent(**agent_kwargs)
    elif args.agent=='cpo':
        agent = CPOAgent(**agent_kwargs)
    elif args.agent=='ppo_dual_ascent':
        agent = PPO_Agent_with_Mu(**dual_ascent_kwargs)

    run_polopt_agent_da(lambda : gym.make(args.env),
                        agent=agent,
                        actor_critic=mlp_actor_critic_with_lam,
                        ac_kwargs=dict(hidden_sizes=[args.hid]*args.l),
                        seed=args.seed,
                        render=args.render,
                        # Experience collection:
                        steps_per_epoch=args.steps,
                        epochs=args.epochs,
                        max_ep_len=args.len,
                        # Discount factors:
                        gamma=args.gamma,
                        cost_gamma=args.cost_gamma,
                        # Policy learning:
                        ent_reg=args.entreg,
                        # KL Divergence:
                        target_kl=args.kl,
                        cost_lim=args.cost_lim,
                        # Logging:
                        logger_kwargs=logger_kwargs,
                        save_freq=1
                        )

import numpy as np
import tensorflow as tf
import gym
import time
import safe_rl.pg.trust_region as tro
from safe_rl.pg.agents import PPOAgent, TRPOAgent, CPOAgent, PPO_Agent_with_Mu
from safe_rl.pg.buffer import CPOBuffer, CPOBuffer_with_vio
from safe_rl.pg.network import count_vars, \
                               get_vars, \
                               mlp_actor_critic, mlp_actor_critic_with_lam,\
                               placeholders, \
                               placeholders_from_spaces
from safe_rl.pg.utils import values_as_sorted_list
from safe_rl.utils.logx import EpochLogger
from safe_rl.utils.mpi_tf import MpiAdamOptimizer, sync_all_params
from safe_rl.utils.mpi_tools import mpi_fork, proc_id, num_procs, mpi_sum, mpi_statistics_scalar

# Multi-purpose agent runner for policy optimization algos 
# (PPO, TRPO, their primal-dual equivalents, CPO)
def run_polopt_agent_da(env_fn,
                        agent=PPOAgent(),
                        actor_critic=mlp_actor_critic_with_lam,
                        ac_kwargs=dict(),
                        seed=0,
                        render=False,
                        rew_scale=0.1,
                        # Experience collection:
                        steps_per_epoch=4000,
                        epochs=50,
                        max_ep_len=1000,
                        # Discount factors:
                        gamma=0.99,
                        lam=0.95,
                        cost_gamma=0.99,
                        cost_lam=0.95,
                        # Policy learning:
                        ent_reg=0.,
                        # Cost constraints / penalties:
                        cost_lim=25,
                        penalty_init=1.,
                        penalty_lr=5e-2,
                        # KL divergence:
                        target_kl=0.01,
                        # Value learning:
                        vf_lr=3e-4,
                        vf_iters=80,
                        # Mu learning:
                        dual_ascent_interval=1,
                        # Logging:
                        logger=None,
                        logger_kwargs=dict(),
                        save_freq=1,
                        lr_decay=True
                        ):


    #=========================================================================#
    #  Prepare logger, seed, and environment in this process                  #
    #=========================================================================#

    logger = EpochLogger(**logger_kwargs) if logger is None else logger
    logger.save_config(locals())

    seed += 10000 * proc_id()
    tf.set_random_seed(seed)
    np.random.seed(seed)

    # Create summary writer
    logs_dir = logger_kwargs.get('output_dir') + '/logs'
    summary_writer = tf.summary.FileWriter(logs_dir, tf.get_default_graph())

    env = env_fn()

    agent.set_logger(logger)

    #=========================================================================#
    #  Create computation graph for actor and critic (not training routine)   #
    #=========================================================================#

    # Share information about action space with policy architecture
    ac_kwargs['action_space'] = env.action_space

    # Inputs to computation graph from environment spaces
    x_ph, a_ph = placeholders_from_spaces(env.observation_space, env.action_space)

    # Inputs to computation graph for batch data
    adv_ph, cadv_ph, ret_ph, cret_ph, logp_old_ph, vio_ph = placeholders(*(None for _ in range(6)))

    # Inputs to computation graph for special purposes
    surr_cost_rescale_ph = tf.placeholder(tf.float32, shape=())
    cur_cost_ph = tf.placeholder(tf.float32, shape=())

    # Outputs from actor critic
    ac_outs = actor_critic(x_ph, a_ph, **ac_kwargs)
    pi, logp, logp_pi, pi_info, pi_info_phs, d_kl, ent, v, vc, multiplier = ac_outs
    tf.summary.histogram('Optimizers/VCVals', vc)
    merged_summary = tf.summary.merge_all()

    # Organize placeholders for zipping with data from buffer on updates
    buf_phs = [x_ph, a_ph, adv_ph, cadv_ph, ret_ph, cret_ph, logp_old_ph, vio_ph]
    buf_phs += values_as_sorted_list(pi_info_phs)

    # Organize symbols we have to compute at each step of acting in env
    get_action_ops = dict(pi=pi, 
                          v=v, 
                          logp_pi=logp_pi,
                          pi_info=pi_info)

    # If agent is reward penalized, it doesn't use a separate value function
    # for costs and we don't need to include it in get_action_ops; otherwise we do.
    if not(agent.reward_penalized):
        get_action_ops['vc'] = vc

    # Count variables
    var_counts = tuple(count_vars(scope) for scope in ['pi', 'vf', 'vc'])
    logger.log('\nNumber of parameters: \t pi: %d, \t v: %d, \t vc: %d\n'%var_counts)

    # Make a sample estimate for entropy to use as sanity check
    approx_ent = tf.reduce_mean(-logp)


    #=========================================================================#
    #  Create replay buffer                                                   #
    #=========================================================================#

    # Obs/act shapes
    obs_shape = env.observation_space.shape
    act_shape = env.action_space.shape

    # Experience buffer
    local_steps_per_epoch = int(steps_per_epoch / num_procs())
    pi_info_shapes = {k: v.shape.as_list()[1:] for k,v in pi_info_phs.items()}
    value_cost_lim = cost_lim / float(max_ep_len) / (1-cost_lam)
    buf = CPOBuffer_with_vio(local_steps_per_epoch,
                            obs_shape,
                            act_shape,
                            pi_info_shapes,
                            gamma,
                            lam,
                            cost_gamma,
                            cost_lam,
                            cost_lim)


    #=========================================================================#
    #  Create computation graph for penalty learning, if applicable           #
    #=========================================================================#

    # if agent.use_penalty:
    #     with tf.variable_scope('penalty'):
    #         # param_init = np.log(penalty_init)
    #         param_init = np.log(max(np.exp(penalty_init)-1, 1e-8))
    #         penalty_param = tf.get_variable('penalty_param',
    #                                       initializer=float(param_init),
    #                                       trainable=agent.learn_penalty,
    #                                       dtype=tf.float32)
    #     # penalty = tf.exp(penalty_param)
    #     penalty = tf.nn.softplus(penalty_param)
    #
    # if agent.learn_penalty:
    #     if agent.penalty_param_loss:
    #         penalty_loss = -penalty_param * (cur_cost_ph - cost_lim)
    #     else:
    #         penalty_loss = -penalty * (cur_cost_ph - cost_lim) # penalty is an adaptive coefficient
    #     train_penalty = MpiAdamOptimizer(learning_rate=penalty_lr).minimize(penalty_loss)

    # if agent.dual_ascent:




    #=========================================================================#
    #  Create computation graph for policy learning                           #
    #=========================================================================#

    # Likelihood ratio
    ratio = tf.exp(logp - logp_old_ph)

    # Surrogate advantage / clipped surrogate advantage
    if agent.clipped_adv:
        min_adv = tf.where(adv_ph>0, 
                           (1+agent.clip_ratio)*adv_ph, 
                           (1-agent.clip_ratio)*adv_ph
                           )
        surr_adv = tf.reduce_mean(tf.minimum(ratio * adv_ph, min_adv))
    else:
        surr_adv = tf.reduce_mean(ratio * adv_ph)

    # old: Surrogate cost now: penalty
    surr_cadv = ratio * cadv_ph
    penalty = tf.reduce_mean(tf.multiply(tf.stop_gradient(multiplier), surr_cadv)) #todo:why use cadv here?
                                                                                   # because the grad is from IS ratio
    # complementary slackness cost
    clipped_vio = tf.clip_by_value(vio_ph, -0.1, 100)
    cs_cost = tf.reduce_mean(tf.multiply(multiplier, tf.stop_gradient(clipped_vio)))

    # Create policy objective function, including entropy regularization
    pi_objective = surr_adv + ent_reg * ent
    # lam_objective = tf.reduce_mean(tf.multiply(mu, tf.stop_gradient(vc)))

    # Possibly include surr_cost in pi_objective
    if agent.objective_penalized or agent.dual_ascent:
        lagrangian = pi_objective - penalty # max min [-(f(x)-mu * g(x))]

    # Loss function for pi is negative of pi_objective
    pi_loss = - lagrangian # todo refine if

    # Loss function for mu
    mu_loss = -cs_cost

    # Optimizer-specific symbols
    if agent.trust_region:

        # Symbols needed for CG solver for any trust region method
        pi_params = get_vars('pi')
        flat_g = tro.flat_grad(pi_loss, pi_params)
        v_ph, hvp = tro.hessian_vector_product(d_kl, pi_params)
        if agent.damping_coeff > 0:
            hvp += agent.damping_coeff * v_ph

        # Symbols needed for CG solver for CPO only
        flat_b = tro.flat_grad(penalty, pi_params) # todo: change to penalty, not confirm right

        # Symbols for getting and setting params
        get_pi_params = tro.flat_concat(pi_params)
        set_pi_params = tro.assign_params_from_flat(v_ph, pi_params)

        training_package = dict(flat_g=flat_g,
                                flat_b=flat_b,
                                v_ph=v_ph,
                                hvp=hvp,
                                get_pi_params=get_pi_params,
                                set_pi_params=set_pi_params)

    elif agent.first_order:

        if lr_decay:
            local_step = tf.Variable(0.0, trainable=False)
            pi_lr = tf.train.polynomial_decay(learning_rate=agent.pi_lr, global_step=local_step,
                                              decay_steps=epochs * agent.pi_iters, end_learning_rate=1e-6)
        else:
            pi_lr = agent.pi_lr
            # Optimizer for first-order policy optimization
        train_pi = MpiAdamOptimizer(learning_rate=pi_lr).minimize(pi_loss)

        if lr_decay:
            local_step = tf.Variable(0.0, trainable=False)
            mu_lr = tf.train.polynomial_decay(learning_rate=agent.mu_lr, global_step=local_step,
                                              decay_steps=int(epochs * agent.pi_iters / agent.ascent_interval),
                                              end_learning_rate=1e-6)
        else:
            mu_lr = agent.mu_lr
            # Optimizer for first-order policy optimization

        train_mu = MpiAdamOptimizer(learning_rate=mu_lr).minimize(mu_loss)

        # Prepare training package for agent
        training_package = dict(train_pi=train_pi, train_mu=train_mu)

    else:
        raise NotImplementedError

    # Provide training package to agent
    training_package.update(dict(pi_loss=pi_loss,
                                 mu_loss=mu_loss,
                                 penalty=penalty,
                                 d_kl=d_kl, 
                                 target_kl=target_kl,
                                 cost_lim=cost_lim))
    agent.prepare_update(training_package)

    #=========================================================================#
    #  Create computation graph for value learning                            #
    #=========================================================================#

    # Value losses
    v_loss = tf.reduce_mean((ret_ph - v)**2)
    vc_loss = tf.reduce_mean((cret_ph - vc)**2)

    # If agent uses penalty directly in reward function, don't train a separate
    # value function for predicting cost returns. (Only use one vf for r - p*c.)
    if agent.reward_penalized:
        total_value_loss = v_loss
    else:
        total_value_loss = v_loss + vc_loss

    # Optimizer for value learning
    if lr_decay:
        local_step_vf = tf.Variable(0.0, trainable=False)
        vf_lr = tf.train.polynomial_decay(learning_rate=vf_lr, global_step=local_step_vf,
                                          decay_steps=epochs * vf_iters, end_learning_rate=1e-6)

    train_vf = MpiAdamOptimizer(learning_rate=vf_lr).minimize(total_value_loss)


    #=========================================================================#
    #  Create session, sync across procs, and set up saver                    #
    #=========================================================================#

    sess = tf.Session()
    sess.run(tf.global_variables_initializer())

    # Sync params across processes
    sess.run(sync_all_params())

    # Setup model saving
    logger.setup_tf_saver(sess, inputs={'x': x_ph}, outputs={'pi': pi, 'v': v, 'vc': vc, 'multiplier': multiplier})


    #=========================================================================#
    #  Provide session to agent                                               #
    #=========================================================================#
    agent.prepare_session(sess)


    #=========================================================================#
    #  Create function for running update (called at end of each epoch)       #
    #=========================================================================#

    def update():
        cur_cost = logger.get_stats('EpCost')[0]
        c = cur_cost - cost_lim
        if c > 0 and agent.cares_about_cost:
            logger.log('Warning! Safety constraint is already violated.', 'red')

        #=====================================================================#
        #  Prepare feed dict                                                  #
        #=====================================================================#

        inputs = {k:v for k,v in zip(buf_phs, buf.get())}
        inputs[surr_cost_rescale_ph] = logger.get_stats('EpLen')[0]
        inputs[cur_cost_ph] = cur_cost

        #=====================================================================#
        #  Make some measurements before updating                             #
        #=====================================================================#

        measures = dict(LossPi=pi_loss,
                        LossMu=mu_loss,
                        Penalty=penalty,
                        Violation=clipped_vio,
                        Multiplier=multiplier,
                        LossV=v_loss,
                        Entropy=ent)
        if not(agent.reward_penalized):
            measures['LossVC'] = vc_loss
        if agent.use_penalty:
            measures['Penalty'] = penalty

        pre_update_measures = sess.run(measures, feed_dict=inputs)
        summary = sess.run(merged_summary, feed_dict)
        summary_writer.add_summary(summary, global_step=epoch * steps_per_epoch)
        logger.store(**pre_update_measures)

        #=====================================================================#
        #  Update penalty if learning penalty                                 #
        #=====================================================================#
        # if agent.learn_penalty:
        #     sess.run(train_penalty, feed_dict={cur_cost_ph: cur_cost})

        #=====================================================================#
        #  Update policy and mu                                               #
        #=====================================================================#
        agent.update_pi(inputs)

        #=====================================================================#
        #  Update value function                                              #
        #=====================================================================#
        for _ in range(vf_iters):
            sess.run(train_vf, feed_dict=inputs)

        #=====================================================================#
        #  Make some measurements after updating                              #
        #=====================================================================#

        del measures['Entropy']
        measures['KL'] = d_kl

        post_update_measures = sess.run(measures, feed_dict=inputs)
        deltas = dict()
        for k in post_update_measures:
            if k in pre_update_measures:
                deltas['Delta'+k] = post_update_measures[k] - pre_update_measures[k]
        logger.store(KL=post_update_measures['KL'], **deltas)




    #=========================================================================#
    #  Run main environment interaction loop                                  #
    #=========================================================================#

    start_time = time.time()
    o, r, d, c, ep_ret, ep_cost, ep_len = env.reset(), 0, False, 0, 0, 0, 0
    cur_penalty = 0
    cum_cost = 0

    for epoch in range(epochs):

        # if agent.use_penalty:
        #     cur_penalty = sess.run(penalty)
        cur_penalty = 0 # todo: remove relevant with cur_penalty

        for t in range(local_steps_per_epoch):

            # Possibly render
            if render and proc_id()==0 and t < 1000:
                env.render()
            
            # Get outputs from policy
            get_action_outs = sess.run(get_action_ops, 
                                       feed_dict={x_ph: o[np.newaxis]})
            a = get_action_outs['pi']
            v_t = get_action_outs['v']
            vc_t = get_action_outs.get('vc', 0)  # Agent may not use cost value func
            logp_t = get_action_outs['logp_pi']
            pi_info_t = get_action_outs['pi_info']

            # Step in environment
            o2, r, d, info = env.step(a)

            # Include penalty on cost
            c = info.get('cost', 0)

            # Track cumulative cost over training
            cum_cost += c

            # save and log
            if agent.reward_penalized:
                r_total = r - cur_penalty * c
                r_total = r_total / (1 + cur_penalty)
                buf.store(o, a, r_total, v_t, 0, 0, logp_t, pi_info_t)
            else:
                buf.store(o, a, r, v_t, c, vc_t, logp_t, pi_info_t)
            logger.store(VVals=v_t, CostVVals=vc_t)

            o = o2
            ep_ret += r
            ep_cost += c
            ep_len += 1

            terminal = d or (ep_len == max_ep_len)
            if terminal or (t==local_steps_per_epoch-1):

                # If trajectory didn't reach terminal state, bootstrap value target(s)
                if d and not(ep_len == max_ep_len):
                    # Note: we do not count env time out as true terminal state
                    last_val, last_cval = 0, 0
                else:
                    feed_dict={x_ph: o[np.newaxis]}
                    if agent.reward_penalized:
                        last_val = sess.run(v, feed_dict=feed_dict)
                        last_cval = 0
                    else:
                        last_val, last_cval = sess.run([v, vc], feed_dict=feed_dict)
                buf.finish_path(last_val, last_cval)

                # Only save EpRet / EpLen if trajectory finished
                if terminal:
                    logger.store(EpRet=ep_ret, EpLen=ep_len, EpCost=ep_cost) # , EpVelocity=velo
                else:
                    print('Warning: trajectory cut off by epoch at %d steps.'%ep_len)

                # Reset environment
                o, r, d, c, ep_ret, ep_len, ep_cost = env.reset(), 0, False, 0, 0, 0, 0
                velo = 0

        # Save model
        if (epoch % save_freq == 0) or (epoch == epochs-1):
            logger.save_state({'env': env}, None)

        #=====================================================================#
        #  Run RL update                                                      #
        #=====================================================================#
        update()

        #=====================================================================#
        #  Cumulative cost calculations                                       #
        #=====================================================================#
        cumulative_cost = mpi_sum(cum_cost)
        cost_rate = cumulative_cost / ((epoch+1)*steps_per_epoch)

        #=====================================================================#
        #  Log performance and stats                                          #
        #=====================================================================#

        logger.log_tabular('Epoch', epoch)

        # Performance stats
        logger.log_tabular('EpRet', with_min_and_max=True)
        logger.log_tabular('EpCost', with_min_and_max=True)
        logger.log_tabular('EpLen', average_only=True)
        logger.log_tabular('CumulativeCost', cumulative_cost)
        logger.log_tabular('CostRate', cost_rate)

        # Value function values
        logger.log_tabular('VVals', with_min_and_max=True)
        logger.log_tabular('CostVVals', with_min_and_max=True)

        # Pi loss and change
        logger.log_tabular('LossPi', average_only=True)
        logger.log_tabular('DeltaLossPi', average_only=True)

        # Surr cost and change
        logger.log_tabular('LossMu', average_only=True)
        logger.log_tabular('DeltaLossMu', average_only=True)
        logger.log_tabular('Violation', with_min_and_max=True)
        logger.log_tabular('Multiplier', with_min_and_max=True)


        # V loss and change
        logger.log_tabular('LossV', average_only=True)
        logger.log_tabular('DeltaLossV', average_only=True)

        # Vc loss and change, if applicable (reward_penalized agents don't use vc)
        if not(agent.reward_penalized):
            logger.log_tabular('LossVC', average_only=True)
            logger.log_tabular('DeltaLossVC', average_only=True)

        if agent.use_penalty or agent.save_penalty:
            logger.log_tabular('Penalty', average_only=True)
            logger.log_tabular('DeltaPenalty', average_only=True)
        else:
            logger.log_tabular('Penalty', 0)
            logger.log_tabular('DeltaPenalty', 0)

        # Anything from the agent?
        agent.log()

        # Policy stats
        logger.log_tabular('Entropy', average_only=True)
        logger.log_tabular('KL', average_only=True)

        # Time and steps elapsed
        logger.log_tabular('TotalEnvInteracts', (epoch+1)*steps_per_epoch)
        logger.log_tabular('Time', time.time()-start_time)

        # Show results!
        logger.dump_tabular()

if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--agent', type=str, default='ppo_dual_ascent')
    parser.add_argument('--env', type=str, default='Safexp-PointGoal1-v0')
    parser.add_argument('--hid', type=int, default=64)
    parser.add_argument('--l', type=int, default=2)
    parser.add_argument('--gamma', type=float, default=0.99)
    parser.add_argument('--cost_gamma', type=float, default=0.99)
    parser.add_argument('--seed', '-s', type=int, default=0)
    parser.add_argument('--cpu', type=int, default=4)
    parser.add_argument('--steps', type=int, default=4000)
    parser.add_argument('--epochs', type=int, default=50)
    parser.add_argument('--len', type=int, default=1000)
    parser.add_argument('--cost_lim', type=float, default=10)
    parser.add_argument('--exp_name', type=str, default='runagent')
    parser.add_argument('--kl', type=float, default=0.01)
    parser.add_argument('--render', action='store_true')
    parser.add_argument('--reward_penalized', action='store_true')
    parser.add_argument('--objective_penalized', action='store_true')
    parser.add_argument('--learn_penalty', action='store_true')
    parser.add_argument('--penalty_param_loss', action='store_true')
    parser.add_argument('--entreg', type=float, default=0.)
    args = parser.parse_args()

    try:
        import safety_gym
    except:
        print('Make sure to install Safety Gym to use constrained RL environments.')

    # mpi_fork(args.cpu)  # run parallel code with mpi

    # Prepare logger
    from safe_rl.utils.run_utils import setup_logger_kwargs
    logger_kwargs = setup_logger_kwargs(args.exp_name, args.seed)

    # Prepare agent
    agent_kwargs = dict(reward_penalized=args.reward_penalized,
                        objective_penalized=args.objective_penalized,
                        learn_penalty=args.learn_penalty,
                        penalty_param_loss=args.penalty_param_loss)
    dual_ascent_kwargs = dict(reward_penalized=False,
                        objective_penalized=True,
                        learn_penalty=False,
                        penalty_param_loss=False,
                        dual_ascent=True)
    if args.agent=='ppo':
        agent = PPOAgent(**agent_kwargs)
    elif args.agent=='trpo':
        agent = TRPOAgent(**agent_kwargs)
    elif args.agent=='cpo':
        agent = CPOAgent(**agent_kwargs)
    elif args.agent=='ppo_dual_ascent':
        agent = PPO_Agent_with_Mu(**dual_ascent_kwargs)

    run_polopt_agent_da(lambda : gym.make(args.env),
                        agent=agent,
                        actor_critic=mlp_actor_critic_with_lam,
                        ac_kwargs=dict(hidden_sizes=[args.hid]*args.l),
                        seed=args.seed,
                        render=args.render,
                        # Experience collection:
                        steps_per_epoch=args.steps,
                        epochs=args.epochs,
                        max_ep_len=args.len,
                        # Discount factors:
                        gamma=args.gamma,
                        cost_gamma=args.cost_gamma,
                        # Policy learning:
                        ent_reg=args.entreg,
                        # KL Divergence:
                        target_kl=args.kl,
                        cost_lim=args.cost_lim,
                        # Logging:
                        logger_kwargs=logger_kwargs,
                        save_freq=1
                        )

# Multi-purpose agent runner for policy optimization algos # (PPO, TRPO, their primal-dual equivalents, CPO) # Experience collection: # Discount factors: # Policy learning: # Cost constraints / penalties: # KL divergence: # Value learning: # Mu learning: # Logging: #=========================================================================# #  Prepare logger, seed, and environment in this process                  # #=========================================================================# # Create summary writer #=========================================================================# #  Create computation graph for actor and critic (not training routine)   # #=========================================================================# # Share information about action space with policy architecture # Inputs to computation graph from environment spaces # Inputs to computation graph for batch data # Inputs to computation graph for special purposes # Outputs from actor critic # Organize placeholders for zipping with data from buffer on updates # Organize symbols we have to compute at each step of acting in env # If agent is reward penalized, it doesn't use a separate value function # for costs and we don't need to include it in get_action_ops; otherwise we do. # Count variables # Make a sample estimate for entropy to use as sanity check #=========================================================================# #  Create replay buffer                                                   # #=========================================================================# # Obs/act shapes # Experience buffer #=========================================================================# #  Create computation graph for penalty learning, if applicable           # #=========================================================================# # if agent.use_penalty: #     with tf.variable_scope('penalty'): #         # param_init = np.log(penalty_init) #         param_init = np.log(max(np.exp(penalty_init)-1, 1e-8)) #         penalty_param = tf.get_variable('penalty_param', #                                       initializer=float(param_init), #                                       trainable=agent.learn_penalty, #                                       dtype=tf.float32) #     # penalty = tf.exp(penalty_param) #     penalty = tf.nn.softplus(penalty_param) # # if agent.learn_penalty: #     if agent.penalty_param_loss: #         penalty_loss = -penalty_param * (cur_cost_ph - cost_lim) #     else: #         penalty_loss = -penalty * (cur_cost_ph - cost_lim) # penalty is an adaptive coefficient #     train_penalty = MpiAdamOptimizer(learning_rate=penalty_lr).minimize(penalty_loss) # if agent.dual_ascent: #=========================================================================# #  Create computation graph for policy learning                           # #=========================================================================# # Likelihood ratio # Surrogate advantage / clipped surrogate advantage # old: Surrogate cost now: penalty #todo:why use cadv here? # because the grad is from IS ratio # complementary slackness cost # Create policy objective function, including entropy regularization # lam_objective = tf.reduce_mean(tf.multiply(mu, tf.stop_gradient(vc))) # Possibly include surr_cost in pi_objective # max min [-(f(x)-mu * g(x))] # Loss function for pi is negative of pi_objective # todo refine if # Loss function for mu # Optimizer-specific symbols # Symbols needed for CG solver for any trust region method # Symbols needed for CG solver for CPO only # todo: change to penalty, not confirm right # Symbols for getting and setting params # Optimizer for first-order policy optimization # Optimizer for first-order policy optimization # Prepare training package for agent # Provide training package to agent #=========================================================================# #  Create computation graph for value learning                            # #=========================================================================# # Value losses # If agent uses penalty directly in reward function, don't train a separate # value function for predicting cost returns. (Only use one vf for r - p*c.) # Optimizer for value learning #=========================================================================# #  Create session, sync across procs, and set up saver                    # #=========================================================================# # Sync params across processes # Setup model saving #=========================================================================# #  Provide session to agent                                               # #=========================================================================# #=========================================================================# #  Create function for running update (called at end of each epoch)       # #=========================================================================# #=====================================================================# #  Prepare feed dict                                                  # #=====================================================================# #=====================================================================# #  Make some measurements before updating                             # #=====================================================================# #=====================================================================# #  Update penalty if learning penalty                                 # #=====================================================================# # if agent.learn_penalty: #     sess.run(train_penalty, feed_dict={cur_cost_ph: cur_cost}) #=====================================================================# #  Update policy and mu                                               # #=====================================================================# #=====================================================================# #  Update value function                                              # #=====================================================================# #=====================================================================# #  Make some measurements after updating                              # #=====================================================================# #=========================================================================# #  Run main environment interaction loop                                  # #=========================================================================# # if agent.use_penalty: #     cur_penalty = sess.run(penalty) # todo: remove relevant with cur_penalty # Possibly render # Get outputs from policy # Agent may not use cost value func # Step in environment # Include penalty on cost # Track cumulative cost over training # save and log # If trajectory didn't reach terminal state, bootstrap value target(s) # Note: we do not count env time out as true terminal state # Only save EpRet / EpLen if trajectory finished # , EpVelocity=velo # Reset environment # Save model #=====================================================================# #  Run RL update                                                      # #=====================================================================# #=====================================================================# #  Cumulative cost calculations                                       # #=====================================================================# #=====================================================================# #  Log performance and stats                                          # #=====================================================================# # Performance stats # Value function values # Pi loss and change # Surr cost and change # V loss and change # Vc loss and change, if applicable (reward_penalized agents don't use vc) # Anything from the agent? # Policy stats # Time and steps elapsed # Show results! # mpi_fork(args.cpu)  # run parallel code with mpi # Prepare logger # Prepare agent # Experience collection: # Discount factors: # Policy learning: # KL Divergence: # Logging: