in safe_rl/sac/sac.py [0:0]
def sac(env_fn, actor_fn=mlp_actor, critic_fn=mlp_critic, ac_kwargs=dict(), seed=0,
steps_per_epoch=1000, epochs=100, replay_size=int(1e6), gamma=0.99,
polyak=0.995, lr=1e-4, batch_size=1024, local_start_steps=int(1e3),
max_ep_len=1000, logger_kwargs=dict(), save_freq=10, local_update_after=int(1e3),
update_freq=1, render=False,
fixed_entropy_bonus=None, entropy_constraint=-1.0,
fixed_cost_penalty=None, cost_constraint=None, cost_lim=None,
reward_scale=1,
):
"""
Args:
env_fn : A function which creates a copy of the environment.
The environment must satisfy the OpenAI Gym API.
actor_fn: A function which takes in placeholder symbols
for state, ``x_ph``, and action, ``a_ph``, and returns the actor
outputs from the agent's Tensorflow computation graph:
=========== ================ ======================================
Symbol Shape Description
=========== ================ ======================================
``mu`` (batch, act_dim) | Computes mean actions from policy
| given states.
``pi`` (batch, act_dim) | Samples actions from policy given
| states.
``logp_pi`` (batch,) | Gives log probability, according to
| the policy, of the action sampled by
| ``pi``. Critical: must be differentiable
| with respect to policy parameters all
| the way through action sampling.
=========== ================ ======================================
critic_fn: A function which takes in placeholder symbols
for state, ``x_ph``, action, ``a_ph``, and policy ``pi``,
and returns the critic outputs from the agent's Tensorflow computation graph:
=========== ================ ======================================
Symbol Shape Description
=========== ================ ======================================
``critic`` (batch,) | Gives one estimate of Q* for
| states in ``x_ph`` and actions in
| ``a_ph``.
``critic_pi`` (batch,) | Gives another estimate of Q* for
| states in ``x_ph`` and actions in
| ``a_ph``.
=========== ================ ======================================
ac_kwargs (dict): Any kwargs appropriate for the actor_fn / critic_fn
function you provided to SAC.
seed (int): Seed for random number generators.
steps_per_epoch (int): Number of steps of interaction (state-action pairs)
for the agent and the environment in each epoch.
epochs (int): Number of epochs to run and train agent.
replay_size (int): Maximum length of replay buffer.
gamma (float): Discount factor. (Always between 0 and 1.)
polyak (float): Interpolation factor in polyak averaging for target
networks. Target networks are updated towards main networks
according to:
.. math:: \\theta_{\\text{targ}} \\leftarrow
\\rho \\theta_{\\text{targ}} + (1-\\rho) \\theta
where :math:`\\rho` is polyak. (Always between 0 and 1, usually
close to 1.)
lr (float): Learning rate (used for both policy and value learning).
batch_size (int): Minibatch size for SGD.
local_start_steps (int): Number of steps for uniform-random action selection,
before running real policy. Helps exploration.
max_ep_len (int): Maximum length of trajectory / episode / rollout.
logger_kwargs (dict): Keyword args for EpochLogger.
save_freq (int): How often (in terms of gap between epochs) to save
the current policy and value function.
fixed_entropy_bonus (float or None): Fixed bonus to reward for entropy.
Units are (points of discounted sum of future reward) / (nats of policy entropy).
If None, use ``entropy_constraint`` to set bonus value instead.
entropy_constraint (float): If ``fixed_entropy_bonus`` is None,
Adjust entropy bonus to maintain at least this much entropy.
Actual constraint value is multiplied by the dimensions of the action space.
Units are (nats of policy entropy) / (action dimenson).
fixed_cost_penalty (float or None): Fixed penalty to reward for cost.
Units are (points of discounted sum of future reward) / (points of discounted sum of future costs).
If None, use ``cost_constraint`` to set penalty value instead.
cost_constraint (float or None): If ``fixed_cost_penalty`` is None,
Adjust cost penalty to maintain at most this much cost.
Units are (points of discounted sum of future costs).
Note: to get an approximate cost_constraint from a cost_lim (undiscounted sum of costs),
multiply cost_lim by (1 - gamma ** episode_len) / (1 - gamma).
If None, use cost_lim to calculate constraint.
cost_lim (float or None): If ``cost_constraint`` is None,
calculate an approximate constraint cost from this cost limit.
Units are (expectation of undiscounted sum of costs in a single episode).
If None, cost_lim is not used, and if no cost constraints are used, do naive optimization.
"""
use_costs = fixed_cost_penalty or cost_constraint or cost_lim
logger = EpochLogger(**logger_kwargs)
logger.save_config(locals())
# Env instantiation
env, test_env = env_fn(), env_fn()
obs_dim = env.observation_space.shape[0]
act_dim = env.action_space.shape[0]
# Setting seeds
tf.set_random_seed(seed)
np.random.seed(seed)
env.seed(seed)
test_env.seed(seed)
# Action limit for clamping: critically, assumes all dimensions share the same bound!
act_limit = env.action_space.high[0]
# Share information about action space with policy architecture
ac_kwargs['action_space'] = env.action_space
# Inputs to computation graph
x_ph, a_ph, x2_ph, r_ph, d_ph, c_ph = placeholders(obs_dim, act_dim, obs_dim, None, None, None)
# Main outputs from computation graph
with tf.variable_scope('main'):
mu, pi, logp_pi = actor_fn(x_ph, a_ph, **ac_kwargs)
qr1, qr1_pi = critic_fn(x_ph, a_ph, pi, name='qr1', **ac_kwargs)
qr2, qr2_pi = critic_fn(x_ph, a_ph, pi, name='qr2', **ac_kwargs)
qc, qc_pi = critic_fn(x_ph, a_ph, pi, name='qc', **ac_kwargs)
with tf.variable_scope('main', reuse=True):
# Additional policy output from a different observation placeholder
# This lets us do separate optimization updates (actor, critics, etc)
# in a single tensorflow op.
_, pi2, logp_pi2 = actor_fn(x2_ph, a_ph, **ac_kwargs)
# Target value network
with tf.variable_scope('target'):
_, qr1_pi_targ = critic_fn(x2_ph, a_ph, pi2, name='qr1', **ac_kwargs)
_, qr2_pi_targ = critic_fn(x2_ph, a_ph, pi2, name='qr2', **ac_kwargs)
_, qc_pi_targ = critic_fn(x2_ph, a_ph, pi2, name='qc', **ac_kwargs)
# Entropy bonus
if fixed_entropy_bonus is None:
with tf.variable_scope('entreg'):
soft_alpha = tf.get_variable('soft_alpha',
initializer=0.0,
trainable=True,
dtype=tf.float32)
alpha = tf.nn.softplus(soft_alpha)
else:
alpha = tf.constant(fixed_entropy_bonus)
log_alpha = tf.log(alpha)
# Cost penalty
if use_costs:
if fixed_cost_penalty is None:
with tf.variable_scope('costpen'):
soft_beta = tf.get_variable('soft_beta',
initializer=0.0,
trainable=True,
dtype=tf.float32)
beta = tf.nn.softplus(soft_beta)
log_beta = tf.log(beta)
else:
beta = tf.constant(fixed_cost_penalty)
log_beta = tf.log(beta)
else:
beta = 0.0 # costs do not contribute to policy optimization
print('Not using costs')
# Experience buffer
replay_buffer = ReplayBuffer(obs_dim=obs_dim, act_dim=act_dim, size=replay_size)
# Count variables
if proc_id()==0:
var_counts = tuple(count_vars(scope) for scope in
['main/pi', 'main/qr1', 'main/qr2', 'main/qc', 'main'])
print(('\nNumber of parameters: \t pi: %d, \t qr1: %d, \t qr2: %d, \t qc: %d, \t total: %d\n')%var_counts)
# Min Double-Q:
min_q_pi = tf.minimum(qr1_pi, qr2_pi)
min_q_pi_targ = tf.minimum(qr1_pi_targ, qr2_pi_targ)
# Targets for Q and V regression
q_backup = tf.stop_gradient(r_ph + gamma*(1-d_ph)*(min_q_pi_targ - alpha * logp_pi2))
qc_backup = tf.stop_gradient(c_ph + gamma*(1-d_ph)*qc_pi_targ)
# Soft actor-critic losses
pi_loss = tf.reduce_mean(alpha * logp_pi - min_q_pi + beta * qc_pi)
qr1_loss = 0.5 * tf.reduce_mean((q_backup - qr1)**2)
qr2_loss = 0.5 * tf.reduce_mean((q_backup - qr2)**2)
qc_loss = 0.5 * tf.reduce_mean((qc_backup - qc)**2)
q_loss = qr1_loss + qr2_loss + qc_loss
# Loss for alpha
entropy_constraint *= act_dim
pi_entropy = -tf.reduce_mean(logp_pi)
# alpha_loss = - soft_alpha * (entropy_constraint - pi_entropy)
alpha_loss = - alpha * (entropy_constraint - pi_entropy)
print('using entropy constraint', entropy_constraint)
# Loss for beta
if use_costs:
if cost_constraint is None:
# Convert assuming equal cost accumulated each step
# Note this isn't the case, since the early in episode doesn't usually have cost,
# but since our algorithm optimizes the discounted infinite horizon from each entry
# in the replay buffer, we should be approximately correct here.
# It's worth checking empirical total undiscounted costs to see if they match.
cost_constraint = cost_lim * (1 - gamma ** max_ep_len) / (1 - gamma) / max_ep_len
print('using cost constraint', cost_constraint)
beta_loss = beta * (cost_constraint - qc)
# Policy train op
# (has to be separate from value train op, because qr1_pi appears in pi_loss)
train_pi_op = MpiAdamOptimizer(learning_rate=lr).minimize(pi_loss, var_list=get_vars('main/pi'), name='train_pi')
# Value train op
with tf.control_dependencies([train_pi_op]):
train_q_op = MpiAdamOptimizer(learning_rate=lr).minimize(q_loss, var_list=get_vars('main/q'), name='train_q')
if fixed_entropy_bonus is None:
entreg_optimizer = MpiAdamOptimizer(learning_rate=lr)
with tf.control_dependencies([train_q_op]):
train_entreg_op = entreg_optimizer.minimize(alpha_loss, var_list=get_vars('entreg'))
if use_costs and fixed_cost_penalty is None:
costpen_optimizer = MpiAdamOptimizer(learning_rate=lr)
with tf.control_dependencies([train_entreg_op]):
train_costpen_op = costpen_optimizer.minimize(beta_loss, var_list=get_vars('costpen'))
# Polyak averaging for target variables
target_update = get_target_update('main', 'target', polyak)
# Single monolithic update with explicit control dependencies
with tf.control_dependencies([train_pi_op]):
with tf.control_dependencies([train_q_op]):
grouped_update = tf.group([target_update])
if fixed_entropy_bonus is None:
grouped_update = tf.group([grouped_update, train_entreg_op])
if use_costs and fixed_cost_penalty is None:
grouped_update = tf.group([grouped_update, train_costpen_op])
# Initializing targets to match main variables
# As a shortcut, use our exponential moving average update w/ coefficient zero
target_init = get_target_update('main', 'target', 0.0)
sess = tf.Session()
sess.run(tf.global_variables_initializer())
sess.run(target_init)
# Sync params across processes
sess.run(sync_all_params())
# Setup model saving
logger.setup_tf_saver(sess, inputs={'x': x_ph, 'a': a_ph},
outputs={'mu': mu, 'pi': pi, 'qr1': qr1, 'qr2': qr2, 'qc': qc})
def get_action(o, deterministic=False):
act_op = mu if deterministic else pi
return sess.run(act_op, feed_dict={x_ph: o.reshape(1,-1)})[0]
def test_agent(n=10):
for j in range(n):
o, r, d, ep_ret, ep_cost, ep_len, ep_goals, = test_env.reset(), 0, False, 0, 0, 0, 0
while not(d or (ep_len == max_ep_len)):
# Take deterministic actions at test time
o, r, d, info = test_env.step(get_action(o, True))
if render and proc_id() == 0 and j == 0:
test_env.render()
ep_ret += r
ep_cost += info.get('cost', 0)
ep_len += 1
ep_goals += 1 if info.get('goal_met', False) else 0
logger.store(TestEpRet=ep_ret, TestEpCost=ep_cost, TestEpLen=ep_len, TestEpGoals=ep_goals)
start_time = time.time()
o, r, d, ep_ret, ep_cost, ep_len, ep_goals = env.reset(), 0, False, 0, 0, 0, 0
total_steps = steps_per_epoch * epochs
# variables to measure in an update
vars_to_get = dict(LossPi=pi_loss, LossQR1=qr1_loss, LossQR2=qr2_loss, LossQC=qc_loss,
QR1Vals=qr1, QR2Vals=qr2, QCVals=qc, LogPi=logp_pi, PiEntropy=pi_entropy,
Alpha=alpha, LogAlpha=log_alpha, LossAlpha=alpha_loss)
if use_costs:
vars_to_get.update(dict(Beta=beta, LogBeta=log_beta, LossBeta=beta_loss))
print('starting training', proc_id())
# Main loop: collect experience in env and update/log each epoch
local_steps = 0
local_steps_per_epoch = steps_per_epoch // num_procs()
local_batch_size = batch_size // num_procs()
epoch_start_time = time.time()
for t in range(total_steps // num_procs()):
"""
Until local_start_steps have elapsed, randomly sample actions
from a uniform distribution for better exploration. Afterwards,
use the learned policy.
"""
if t > local_start_steps:
a = get_action(o)
else:
a = env.action_space.sample()
# Step the env
o2, r, d, info = env.step(a)
r *= reward_scale # yee-haw
c = info.get('cost', 0)
ep_ret += r
ep_cost += c
ep_len += 1
ep_goals += 1 if info.get('goal_met', False) else 0
local_steps += 1
# Ignore the "done" signal if it comes from hitting the time
# horizon (that is, when it's an artificial terminal signal
# that isn't based on the agent's state)
d = False if ep_len==max_ep_len else d
# Store experience to replay buffer
replay_buffer.store(o, a, r, o2, d, c)
# Super critical, easy to overlook step: make sure to update
# most recent observation!
o = o2
if d or (ep_len == max_ep_len):
logger.store(EpRet=ep_ret, EpCost=ep_cost, EpLen=ep_len, EpGoals=ep_goals)
o, r, d, ep_ret, ep_cost, ep_len, ep_goals = env.reset(), 0, False, 0, 0, 0, 0
if t > 0 and t % update_freq == 0:
for j in range(update_freq):
batch = replay_buffer.sample_batch(local_batch_size)
feed_dict = {x_ph: batch['obs1'],
x2_ph: batch['obs2'],
a_ph: batch['acts'],
r_ph: batch['rews'],
c_ph: batch['costs'],
d_ph: batch['done'],
}
if t < local_update_after:
logger.store(**sess.run(vars_to_get, feed_dict))
else:
values, _ = sess.run([vars_to_get, grouped_update], feed_dict)
logger.store(**values)
# End of epoch wrap-up
if t > 0 and t % local_steps_per_epoch == 0:
epoch = t // local_steps_per_epoch
# Save model
if (epoch % save_freq == 0) or (epoch == epochs-1):
logger.save_state({'env': env}, None)
# Test the performance of the deterministic version of the agent.
test_start_time = time.time()
test_agent()
logger.store(TestTime=time.time() - test_start_time)
logger.store(EpochTime=time.time() - epoch_start_time)
epoch_start_time = time.time()
# Log info about epoch
logger.log_tabular('Epoch', epoch)
logger.log_tabular('EpRet', with_min_and_max=True)
logger.log_tabular('TestEpRet', with_min_and_max=True)
logger.log_tabular('EpCost', with_min_and_max=True)
logger.log_tabular('TestEpCost', with_min_and_max=True)
logger.log_tabular('EpLen', average_only=True)
logger.log_tabular('TestEpLen', average_only=True)
logger.log_tabular('EpGoals', average_only=True)
logger.log_tabular('TestEpGoals', average_only=True)
logger.log_tabular('TotalEnvInteracts', mpi_sum(local_steps))
logger.log_tabular('QR1Vals', with_min_and_max=True)
logger.log_tabular('QR2Vals', with_min_and_max=True)
logger.log_tabular('QCVals', with_min_and_max=True)
logger.log_tabular('LogPi', with_min_and_max=True)
logger.log_tabular('LossPi', average_only=True)
logger.log_tabular('LossQR1', average_only=True)
logger.log_tabular('LossQR2', average_only=True)
logger.log_tabular('LossQC', average_only=True)
logger.log_tabular('LossAlpha', average_only=True)
logger.log_tabular('LogAlpha', average_only=True)
logger.log_tabular('Alpha', average_only=True)
if use_costs:
logger.log_tabular('LossBeta', average_only=True)
logger.log_tabular('LogBeta', average_only=True)
logger.log_tabular('Beta', average_only=True)
logger.log_tabular('PiEntropy', average_only=True)
logger.log_tabular('TestTime', average_only=True)
logger.log_tabular('EpochTime', average_only=True)
logger.log_tabular('TotalTime', time.time()-start_time)
logger.dump_tabular()