in safe_rl/pg/run_agent.py [0:0]
def run_polopt_agent(env_fn,
agent=PPOAgent(),
actor_critic=mlp_actor_critic,
ac_kwargs=dict(),
seed=0,
render=False,
# Experience collection:
steps_per_epoch=4000,
epochs=50,
max_ep_len=1000,
# Discount factors:
gamma=0.99,
lam=0.97,
cost_gamma=0.99,
cost_lam=0.97,
# 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=1e-3,
vf_iters=80,
# Logging:
logger=None,
logger_kwargs=dict(),
save_freq=1
):
#=========================================================================#
# 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)
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 = placeholders(*(None for _ in range(5)))
# 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 = ac_outs
# 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]
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()}
buf = CPOBuffer(local_steps_per_epoch,
obs_shape,
act_shape,
pi_info_shapes,
gamma,
lam,
cost_gamma,
cost_lam)
#=========================================================================#
# 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)
train_penalty = MpiAdamOptimizer(learning_rate=penalty_lr).minimize(penalty_loss)
#=========================================================================#
# 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)
# Surrogate cost
surr_cost = tf.reduce_mean(ratio * cadv_ph)
# Create policy objective function, including entropy regularization
pi_objective = surr_adv + ent_reg * ent
# Possibly include surr_cost in pi_objective
if agent.objective_penalized:
pi_objective -= penalty * surr_cost
pi_objective /= (1 + penalty)
# Loss function for pi is negative of pi_objective
pi_loss = -pi_objective
# 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(surr_cost, pi_params)
# 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:
# Optimizer for first-order policy optimization
train_pi = MpiAdamOptimizer(learning_rate=agent.pi_lr).minimize(pi_loss)
# Prepare training package for agent
training_package = dict(train_pi=train_pi)
else:
raise NotImplementedError
# Provide training package to agent
training_package.update(dict(pi_loss=pi_loss,
surr_cost=surr_cost,
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
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})
#=========================================================================#
# 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,
SurrCost=surr_cost,
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)
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 #
#=====================================================================#
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)
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)
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
# 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('SurrCost', average_only=True)
logger.log_tabular('DeltaSurrCost', average_only=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()