def test_mpi_adam():
"""
tests the MpiAdam object's functionality
"""
np.random.seed(0)
tf.set_random_seed(0)
a_var = tf.Variable(np.random.randn(3).astype('float32'))
b_var = tf.Variable(np.random.randn(2, 5).astype('float32'))
loss = tf.reduce_sum(tf.square(a_var)) + tf.reduce_sum(tf.sin(b_var))
learning_rate = 1e-2
update_op = tf.train.AdamOptimizer(learning_rate).minimize(loss)
do_update = tf_utils.function([], loss, updates=[update_op])
tf.get_default_session().run(tf.global_variables_initializer())
for step in range(10):
print(step, do_update())
tf.set_random_seed(0)
tf.get_default_session().run(tf.global_variables_initializer())
var_list = [a_var, b_var]
lossandgrad = tf_utils.function(
[], [loss, tf_utils.flatgrad(loss, var_list)], updates=[update_op])
adam = MpiAdam(var_list)
for step in range(10):
loss, grad = lossandgrad()
adam.update(grad, learning_rate)
print(step, loss)
def __init__(self, epsilon=1e-2, shape=()):
"""
calulates the running mean and std of a data stream
https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Parallel_algorithm
:param epsilon: (float) helps with arithmetic issues
:param shape: (tuple) the shape of the data stream's output
"""
self._sum = tf.get_variable(dtype=tf.float64,
shape=shape,
initializer=tf.constant_initializer(0.0),
name="runningsum",
trainable=False)
self._sumsq = tf.get_variable(
dtype=tf.float64,
shape=shape,
initializer=tf.constant_initializer(epsilon),
name="runningsumsq",
trainable=False)
self._count = tf.get_variable(
dtype=tf.float64,
shape=(),
initializer=tf.constant_initializer(epsilon),
name="count",
trainable=False)
self.shape = shape
self.mean = tf.cast(self._sum / self._count, tf.float32)
self.std = tf.sqrt(
tf.maximum(
tf.cast(self._sumsq / self._count, tf.float32) -
tf.square(self.mean), 1e-2))
newsum = tf.placeholder(shape=self.shape, dtype=tf.float64, name='sum')
newsumsq = tf.placeholder(shape=self.shape,
dtype=tf.float64,
name='var')
newcount = tf.placeholder(shape=[], dtype=tf.float64, name='count')
self.incfiltparams = tf_util.function(
[newsum, newsumsq, newcount], [],
updates=[
tf.assign_add(self._sum, newsum),
tf.assign_add(self._sumsq, newsumsq),
tf.assign_add(self._count, newcount)
])
def build_act(q_func, ob_space, ac_space, stochastic_ph, update_eps_ph, sess):
"""
Creates the act function:
:param q_func: (DQNPolicy) the policy
:param ob_space: (Gym Space) The observation space of the environment
:param ac_space: (Gym Space) The action space of the environment
:param stochastic_ph: (TensorFlow Tensor) the stochastic placeholder
:param update_eps_ph: (TensorFlow Tensor) the update_eps placeholder
:param sess: (TensorFlow session) The current TensorFlow session
:return: (function (TensorFlow Tensor, bool, float): TensorFlow Tensor, (TensorFlow Tensor, TensorFlow Tensor)
act function to select and action given observation (See the top of the file for details),
A tuple containing the observation placeholder and the processed observation placeholder respectively.
"""
eps = tf.get_variable("eps", (), initializer=tf.constant_initializer(0))
policy = q_func(sess, ob_space, ac_space, 1, 1, None)
obs_phs = (policy.obs_ph, policy.processed_obs)
deterministic_actions = tf.argmax(policy.q_values, axis=1)
batch_size = tf.shape(policy.obs_ph)[0]
n_actions = ac_space.nvec if isinstance(ac_space, MultiDiscrete) else ac_space.n
random_actions = tf.random_uniform(tf.stack([batch_size]), minval=0, maxval=n_actions, dtype=tf.int64)
chose_random = tf.random_uniform(tf.stack([batch_size]), minval=0, maxval=1, dtype=tf.float32) < eps
stochastic_actions = tf.where(chose_random, random_actions, deterministic_actions)
output_actions = tf.cond(stochastic_ph, lambda: stochastic_actions, lambda: deterministic_actions)
update_eps_expr = eps.assign(tf.cond(update_eps_ph >= 0, lambda: update_eps_ph, lambda: eps))
_act = tf_util.function(inputs=[policy.obs_ph, stochastic_ph, update_eps_ph],
outputs=output_actions,
givens={update_eps_ph: -1.0, stochastic_ph: True},
updates=[update_eps_expr])
def act(obs, stochastic=True, update_eps=-1):
return _act(obs, stochastic, update_eps)
return act, obs_phs
def build_train(q_func, ob_space, ac_space, optimizer, sess, grad_norm_clipping=None,
gamma=1.0, double_q=True, scope="deepq", reuse=None,
param_noise=False, param_noise_filter_func=None, full_tensorboard_log=False):
"""
Creates the train function:
:param q_func: (DQNPolicy) the policy
:param ob_space: (Gym Space) The observation space of the environment
:param ac_space: (Gym Space) The action space of the environment
:param reuse: (bool) whether or not to reuse the graph variables
:param optimizer: (tf.train.Optimizer) optimizer to use for the Q-learning objective.
:param sess: (TensorFlow session) The current TensorFlow session
:param grad_norm_clipping: (float) clip gradient norms to this value. If None no clipping is performed.
:param gamma: (float) discount rate.
:param double_q: (bool) if true will use Double Q Learning (https://arxiv.org/abs/1509.06461). In general it is a
good idea to keep it enabled.
:param scope: (str or VariableScope) optional scope for variable_scope.
:param reuse: (bool) whether or not the variables should be reused. To be able to reuse the scope must be given.
:param param_noise: (bool) whether or not to use parameter space noise (https://arxiv.org/abs/1706.01905)
:param param_noise_filter_func: (function (TensorFlow Tensor): bool) function that decides whether or not a
variable should be perturbed. Only applicable if param_noise is True. If set to None, default_param_noise_filter
is used by default.
:param full_tensorboard_log: (bool) enable additional logging when using tensorboard
WARNING: this logging can take a lot of space quickly
:return: (tuple)
act: (function (TensorFlow Tensor, bool, float): TensorFlow Tensor) function to select and action given
observation. See the top of the file for details.
train: (function (Any, numpy float, numpy float, Any, numpy bool, numpy float): numpy float)
optimize the error in Bellman's equation. See the top of the file for details.
update_target: (function) copy the parameters from optimized Q function to the target Q function.
See the top of the file for details.
step_model: (DQNPolicy) Policy for evaluation
"""
n_actions = ac_space.nvec if isinstance(ac_space, MultiDiscrete) else ac_space.n
with tf.variable_scope("input", reuse=reuse):
stochastic_ph = tf.placeholder(tf.bool, (), name="stochastic")
update_eps_ph = tf.placeholder(tf.float32, (), name="update_eps")
with tf.variable_scope(scope, reuse=reuse):
if param_noise:
act_f, obs_phs = build_act_with_param_noise(q_func, ob_space, ac_space, stochastic_ph, update_eps_ph, sess,
param_noise_filter_func=param_noise_filter_func)
else:
act_f, obs_phs = build_act(q_func, ob_space, ac_space, stochastic_ph, update_eps_ph, sess)
# q network evaluation
with tf.variable_scope("step_model", reuse=True, custom_getter=tf_util.outer_scope_getter("step_model")):
step_model = q_func(sess, ob_space, ac_space, 1, 1, None, reuse=True, obs_phs=obs_phs)
q_func_vars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope=tf.get_variable_scope().name + "/model")
# target q network evaluation
with tf.variable_scope("target_q_func", reuse=False):
target_policy = q_func(sess, ob_space, ac_space, 1, 1, None, reuse=False)
target_q_func_vars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES,
scope=tf.get_variable_scope().name + "/target_q_func")
# compute estimate of best possible value starting from state at t + 1
double_q_values = None
double_obs_ph = target_policy.obs_ph
if double_q:
with tf.variable_scope("double_q", reuse=True, custom_getter=tf_util.outer_scope_getter("double_q")):
double_policy = q_func(sess, ob_space, ac_space, 1, 1, None, reuse=True)
double_q_values = double_policy.q_values
double_obs_ph = double_policy.obs_ph
with tf.variable_scope("loss", reuse=reuse):
# set up placeholders
act_t_ph = tf.placeholder(tf.int32, [None], name="action")
rew_t_ph = tf.placeholder(tf.float32, [None], name="reward")
done_mask_ph = tf.placeholder(tf.float32, [None], name="done")
importance_weights_ph = tf.placeholder(tf.float32, [None], name="weight")
# q scores for actions which we know were selected in the given state.
q_t_selected = tf.reduce_sum(step_model.q_values * tf.one_hot(act_t_ph, n_actions), axis=1)
# compute estimate of best possible value starting from state at t + 1
if double_q:
q_tp1_best_using_online_net = tf.argmax(double_q_values, axis=1)
q_tp1_best = tf.reduce_sum(target_policy.q_values * tf.one_hot(q_tp1_best_using_online_net, n_actions), axis=1)
else:
q_tp1_best = tf.reduce_max(target_policy.q_values, axis=1)
q_tp1_best_masked = (1.0 - done_mask_ph) * q_tp1_best
# compute RHS of bellman equation
q_t_selected_target = rew_t_ph + gamma * q_tp1_best_masked
# compute the error (potentially clipped)
td_error = q_t_selected - tf.stop_gradient(q_t_selected_target)
errors = tf_util.huber_loss(td_error)
weighted_error = tf.reduce_mean(importance_weights_ph * errors)
tf.summary.scalar("td_error", tf.reduce_mean(td_error))
tf.summary.scalar("loss", weighted_error)
if full_tensorboard_log:
tf.summary.histogram("td_error", td_error)
# update_target_fn will be called periodically to copy Q network to target Q network
update_target_expr = []
for var, var_target in zip(sorted(q_func_vars, key=lambda v: v.name),
sorted(target_q_func_vars, key=lambda v: v.name)):
update_target_expr.append(var_target.assign(var))
update_target_expr = tf.group(*update_target_expr)
# compute optimization op (potentially with gradient clipping)
gradients = optimizer.compute_gradients(weighted_error, var_list=q_func_vars)
if grad_norm_clipping is not None:
for i, (grad, var) in enumerate(gradients):
if grad is not None:
gradients[i] = (tf.clip_by_norm(grad, grad_norm_clipping), var)
with tf.variable_scope("input_info", reuse=False):
tf.summary.scalar('rewards', tf.reduce_mean(rew_t_ph))
tf.summary.scalar('importance_weights', tf.reduce_mean(importance_weights_ph))
if full_tensorboard_log:
tf.summary.histogram('rewards', rew_t_ph)
tf.summary.histogram('importance_weights', importance_weights_ph)
if tf_util.is_image(obs_phs[0]):
tf.summary.image('observation', obs_phs[0])
elif len(obs_phs[0].shape) == 1:
tf.summary.histogram('observation', obs_phs[0])
optimize_expr = optimizer.apply_gradients(gradients)
summary = tf.summary.merge_all()
# Create callable functions
train = tf_util.function(
inputs=[
obs_phs[0],
act_t_ph,
rew_t_ph,
target_policy.obs_ph,
double_obs_ph,
done_mask_ph,
importance_weights_ph
],
outputs=[summary, td_error],
updates=[optimize_expr]
)
update_target = tf_util.function([], [], updates=[update_target_expr])
return act_f, train, update_target, step_model
def build_act_with_param_noise(q_func, ob_space, ac_space, stochastic_ph, update_eps_ph, sess,
param_noise_filter_func=None):
"""
Creates the act function with support for parameter space noise exploration (https://arxiv.org/abs/1706.01905):
:param q_func: (DQNPolicy) the policy
:param ob_space: (Gym Space) The observation space of the environment
:param ac_space: (Gym Space) The action space of the environment
:param stochastic_ph: (TensorFlow Tensor) the stochastic placeholder
:param update_eps_ph: (TensorFlow Tensor) the update_eps placeholder
:param sess: (TensorFlow session) The current TensorFlow session
:param param_noise_filter_func: (function (TensorFlow Tensor): bool) function that decides whether or not a
variable should be perturbed. Only applicable if param_noise is True. If set to None, default_param_noise_filter
is used by default.
:return: (function (TensorFlow Tensor, bool, float): TensorFlow Tensor, (TensorFlow Tensor, TensorFlow Tensor)
act function to select and action given observation (See the top of the file for details),
A tuple containing the observation placeholder and the processed observation placeholder respectively.
"""
if param_noise_filter_func is None:
param_noise_filter_func = default_param_noise_filter
update_param_noise_threshold_ph = tf.placeholder(tf.float32, (), name="update_param_noise_threshold")
update_param_noise_scale_ph = tf.placeholder(tf.bool, (), name="update_param_noise_scale")
reset_ph = tf.placeholder(tf.bool, (), name="reset")
eps = tf.get_variable("eps", (), initializer=tf.constant_initializer(0))
param_noise_scale = tf.get_variable("param_noise_scale", (), initializer=tf.constant_initializer(0.01),
trainable=False)
param_noise_threshold = tf.get_variable("param_noise_threshold", (), initializer=tf.constant_initializer(0.05),
trainable=False)
# Unmodified Q.
policy = q_func(sess, ob_space, ac_space, 1, 1, None)
obs_phs = (policy.obs_ph, policy.processed_obs)
# Perturbable Q used for the actual rollout.
with tf.variable_scope("perturbed_model", reuse=False):
perturbable_policy = q_func(sess, ob_space, ac_space, 1, 1, None, obs_phs=obs_phs)
def perturb_vars(original_scope, perturbed_scope):
"""
We have to wrap this code into a function due to the way tf.cond() works.
See https://stackoverflow.com/questions/37063952/confused-by-the-behavior-of-tf-cond for a more detailed
discussion.
:param original_scope: (str or VariableScope) the original scope.
:param perturbed_scope: (str or VariableScope) the perturbed scope.
:return: (TensorFlow Operation)
"""
all_vars = scope_vars(absolute_scope_name(original_scope))
all_perturbed_vars = scope_vars(absolute_scope_name(perturbed_scope))
assert len(all_vars) == len(all_perturbed_vars)
perturb_ops = []
for var, perturbed_var in zip(all_vars, all_perturbed_vars):
if param_noise_filter_func(perturbed_var):
# Perturb this variable.
operation = tf.assign(perturbed_var,
var + tf.random_normal(shape=tf.shape(var), mean=0.,
stddev=param_noise_scale))
else:
# Do not perturb, just assign.
operation = tf.assign(perturbed_var, var)
perturb_ops.append(operation)
assert len(perturb_ops) == len(all_vars)
return tf.group(*perturb_ops)
# Set up functionality to re-compute `param_noise_scale`. This perturbs yet another copy
# of the network and measures the effect of that perturbation in action space. If the perturbation
# is too big, reduce scale of perturbation, otherwise increase.
with tf.variable_scope("adaptive_model", reuse=False):
adaptive_policy = q_func(sess, ob_space, ac_space, 1, 1, None, obs_phs=obs_phs)
perturb_for_adaption = perturb_vars(original_scope="model", perturbed_scope="adaptive_model/model")
kl_loss = tf.reduce_sum(
tf.nn.softmax(policy.q_values) *
(tf.log(tf.nn.softmax(policy.q_values)) - tf.log(tf.nn.softmax(adaptive_policy.q_values))),
axis=-1)
mean_kl = tf.reduce_mean(kl_loss)
def update_scale():
"""
update the scale expression
:return: (TensorFlow Tensor) the updated scale expression
"""
with tf.control_dependencies([perturb_for_adaption]):
update_scale_expr = tf.cond(mean_kl < param_noise_threshold,
lambda: param_noise_scale.assign(param_noise_scale * 1.01),
lambda: param_noise_scale.assign(param_noise_scale / 1.01),
)
return update_scale_expr
# Functionality to update the threshold for parameter space noise.
update_param_noise_thres_expr = param_noise_threshold.assign(
tf.cond(update_param_noise_threshold_ph >= 0, lambda: update_param_noise_threshold_ph,
lambda: param_noise_threshold))
# Put everything together.
perturbed_deterministic_actions = tf.argmax(perturbable_policy.q_values, axis=1)
deterministic_actions = tf.argmax(policy.q_values, axis=1)
batch_size = tf.shape(policy.obs_ph)[0]
n_actions = ac_space.nvec if isinstance(ac_space, MultiDiscrete) else ac_space.n
random_actions = tf.random_uniform(tf.stack([batch_size]), minval=0, maxval=n_actions, dtype=tf.int64)
chose_random = tf.random_uniform(tf.stack([batch_size]), minval=0, maxval=1, dtype=tf.float32) < eps
perturbed_stochastic_actions = tf.where(chose_random, random_actions, perturbed_deterministic_actions)
stochastic_actions = tf.where(chose_random, random_actions, deterministic_actions)
perturbed_output_actions = tf.cond(stochastic_ph, lambda: perturbed_stochastic_actions,
lambda: deterministic_actions)
output_actions = tf.cond(stochastic_ph, lambda: stochastic_actions, lambda: deterministic_actions)
update_eps_expr = eps.assign(tf.cond(update_eps_ph >= 0, lambda: update_eps_ph, lambda: eps))
updates = [
update_eps_expr,
tf.cond(reset_ph, lambda: perturb_vars(original_scope="model", perturbed_scope="perturbed_model/model"),
lambda: tf.group(*[])),
tf.cond(update_param_noise_scale_ph, lambda: update_scale(), lambda: tf.Variable(0., trainable=False)),
update_param_noise_thres_expr,
]
_act = tf_util.function(inputs=[policy.obs_ph, stochastic_ph, update_eps_ph],
outputs=output_actions,
givens={update_eps_ph: -1.0, stochastic_ph: True},
updates=[update_eps_expr])
_perturbed_act = tf_util.function(
inputs=[policy.obs_ph, stochastic_ph, update_eps_ph, reset_ph, update_param_noise_threshold_ph,
update_param_noise_scale_ph],
outputs=perturbed_output_actions,
givens={update_eps_ph: -1.0, stochastic_ph: True, reset_ph: False, update_param_noise_threshold_ph: False,
update_param_noise_scale_ph: False},
updates=updates)
def act(obs, reset=None, update_param_noise_threshold=None, update_param_noise_scale=None, stochastic=True,
update_eps=-1):
"""
get the action from the current observation
:param obs: (Any) Observation that can be feed into the output of make_obs_ph
:param reset: (bool) reset the perturbed policy by sampling a new perturbation
:param update_param_noise_threshold: (float) the desired threshold for the difference between
non-perturbed and perturbed policy
:param update_param_noise_scale: (bool) whether or not to update the scale of the noise for the next time
it is re-perturbed
:param stochastic: (bool) if set to False all the actions are always deterministic (default False)
:param update_eps: (float) update epsilon a new value, if negative not update happens
(default: no update)
:return: (TensorFlow Tensor) tensor of dtype tf.int64 and shape (BATCH_SIZE,) with an action to be
performed for every element of the batch.
"""
if reset is None or update_param_noise_threshold is None or update_param_noise_scale is None:
return _act(obs, stochastic, update_eps)
else:
return _perturbed_act(obs, stochastic, update_eps, reset, update_param_noise_threshold,
update_param_noise_scale)
return act, obs_phs
def setup_model(self):
# prevent import loops
from stable_baselines_custom.gail.adversary import TransitionClassifier
with SetVerbosity(self.verbose):
assert issubclass(self.policy, ActorCriticPolicy), "Error: the input policy for the TRPO model must be " \
"an instance of common.policies.ActorCriticPolicy."
self.nworkers = MPI.COMM_WORLD.Get_size()
self.rank = MPI.COMM_WORLD.Get_rank()
np.set_printoptions(precision=3)
self.graph = tf.Graph()
with self.graph.as_default():
self.set_random_seed(self.seed)
self.sess = tf_util.make_session(num_cpu=self.n_cpu_tf_sess, graph=self.graph)
if self.using_gail:
self.reward_giver = TransitionClassifier(self.observation_space, self.action_space,
self.hidden_size_adversary,
entcoeff=self.adversary_entcoeff)
# Construct network for new policy
self.policy_pi = self.policy(self.sess, self.observation_space, self.action_space, self.n_envs, 1,
None, reuse=False, **self.policy_kwargs)
# Network for old policy
with tf.variable_scope("oldpi", reuse=False):
old_policy = self.policy(self.sess, self.observation_space, self.action_space, self.n_envs, 1,
None, reuse=False, **self.policy_kwargs)
with tf.variable_scope("loss", reuse=False):
atarg = tf.placeholder(dtype=tf.float32, shape=[None]) # Target advantage function (if applicable)
ret = tf.placeholder(dtype=tf.float32, shape=[None]) # Empirical return
observation = self.policy_pi.obs_ph
action = self.policy_pi.pdtype.sample_placeholder([None])
kloldnew = old_policy.proba_distribution.kl(self.policy_pi.proba_distribution)
ent = self.policy_pi.proba_distribution.entropy()
meankl = tf.reduce_mean(kloldnew)
meanent = tf.reduce_mean(ent)
entbonus = self.entcoeff * meanent
vferr = tf.reduce_mean(tf.square(self.policy_pi.value_flat - ret))
# advantage * pnew / pold
ratio = tf.exp(self.policy_pi.proba_distribution.logp(action) -
old_policy.proba_distribution.logp(action))
surrgain = tf.reduce_mean(ratio * atarg)
optimgain = surrgain + entbonus
losses = [optimgain, meankl, entbonus, surrgain, meanent]
self.loss_names = ["optimgain", "meankl", "entloss", "surrgain", "entropy"]
dist = meankl
all_var_list = tf_util.get_trainable_vars("model")
var_list = [v for v in all_var_list if "/vf" not in v.name and "/q/" not in v.name]
vf_var_list = [v for v in all_var_list if "/pi" not in v.name and "/logstd" not in v.name]
self.get_flat = tf_util.GetFlat(var_list, sess=self.sess)
self.set_from_flat = tf_util.SetFromFlat(var_list, sess=self.sess)
klgrads = tf.gradients(dist, var_list)
flat_tangent = tf.placeholder(dtype=tf.float32, shape=[None], name="flat_tan")
shapes = [var.get_shape().as_list() for var in var_list]
start = 0
tangents = []
for shape in shapes:
var_size = tf_util.intprod(shape)
tangents.append(tf.reshape(flat_tangent[start: start + var_size], shape))
start += var_size
gvp = tf.add_n([tf.reduce_sum(grad * tangent)
for (grad, tangent) in zipsame(klgrads, tangents)]) # pylint: disable=E1111
# Fisher vector products
fvp = tf_util.flatgrad(gvp, var_list)
tf.summary.scalar('entropy_loss', meanent)
tf.summary.scalar('policy_gradient_loss', optimgain)
tf.summary.scalar('value_function_loss', surrgain)
tf.summary.scalar('approximate_kullback-leibler', meankl)
tf.summary.scalar('loss', optimgain + meankl + entbonus + surrgain + meanent)
self.assign_old_eq_new = \
tf_util.function([], [], updates=[tf.assign(oldv, newv) for (oldv, newv) in
zipsame(tf_util.get_globals_vars("oldpi"),
tf_util.get_globals_vars("model"))])
self.compute_losses = tf_util.function([observation, old_policy.obs_ph, action, atarg], losses)
self.compute_fvp = tf_util.function([flat_tangent, observation, old_policy.obs_ph, action, atarg],
fvp)
self.compute_vflossandgrad = tf_util.function([observation, old_policy.obs_ph, ret],
tf_util.flatgrad(vferr, vf_var_list))
@contextmanager
def timed(msg):
if self.rank == 0 and self.verbose >= 1:
print(colorize(msg, color='magenta'))
start_time = time.time()
yield
print(colorize("done in {:.3f} seconds".format((time.time() - start_time)),
color='magenta'))
else:
yield
def allmean(arr):
assert isinstance(arr, np.ndarray)
out = np.empty_like(arr)
MPI.COMM_WORLD.Allreduce(arr, out, op=MPI.SUM)
out /= self.nworkers
return out
tf_util.initialize(sess=self.sess)
th_init = self.get_flat()
MPI.COMM_WORLD.Bcast(th_init, root=0)
self.set_from_flat(th_init)
with tf.variable_scope("Adam_mpi", reuse=False):
self.vfadam = MpiAdam(vf_var_list, sess=self.sess)
if self.using_gail:
self.d_adam = MpiAdam(self.reward_giver.get_trainable_variables(), sess=self.sess)
self.d_adam.sync()
self.vfadam.sync()
with tf.variable_scope("input_info", reuse=False):
tf.summary.scalar('discounted_rewards', tf.reduce_mean(ret))
tf.summary.scalar('learning_rate', tf.reduce_mean(self.vf_stepsize))
tf.summary.scalar('advantage', tf.reduce_mean(atarg))
tf.summary.scalar('kl_clip_range', tf.reduce_mean(self.max_kl))
if self.full_tensorboard_log:
tf.summary.histogram('discounted_rewards', ret)
tf.summary.histogram('learning_rate', self.vf_stepsize)
tf.summary.histogram('advantage', atarg)
tf.summary.histogram('kl_clip_range', self.max_kl)
if tf_util.is_image(self.observation_space):
tf.summary.image('observation', observation)
else:
tf.summary.histogram('observation', observation)
self.timed = timed
self.allmean = allmean
self.step = self.policy_pi.step
self.proba_step = self.policy_pi.proba_step
self.initial_state = self.policy_pi.initial_state
self.params = tf_util.get_trainable_vars("model") + tf_util.get_trainable_vars("oldpi")
if self.using_gail:
self.params.extend(self.reward_giver.get_trainable_variables())
self.summary = tf.summary.merge_all()
self.compute_lossandgrad = \
tf_util.function([observation, old_policy.obs_ph, action, atarg, ret],
[self.summary, tf_util.flatgrad(optimgain, var_list)] + losses)
def setup_model(self):
with SetVerbosity(self.verbose):
self.graph = tf.Graph()
with self.graph.as_default():
self.set_random_seed(self.seed)
self.sess = tf_util.make_session(num_cpu=self.n_cpu_tf_sess,
graph=self.graph)
# Construct network for new policy
self.policy_pi = self.policy(self.sess,
self.observation_space,
self.action_space,
self.n_envs,
1,
None,
reuse=False,
**self.policy_kwargs)
# Network for old policy
with tf.variable_scope("oldpi", reuse=False):
old_pi = self.policy(self.sess,
self.observation_space,
self.action_space,
self.n_envs,
1,
None,
reuse=False,
**self.policy_kwargs)
with tf.variable_scope("loss", reuse=False):
# Target advantage function (if applicable)
atarg = tf.placeholder(dtype=tf.float32, shape=[None])
# Empirical return
ret = tf.placeholder(dtype=tf.float32, shape=[None])
# learning rate multiplier, updated with schedule
lrmult = tf.placeholder(name='lrmult',
dtype=tf.float32,
shape=[])
# Annealed cliping parameter epislon
clip_param = self.clip_param * lrmult
obs_ph = self.policy_pi.obs_ph
action_ph = self.policy_pi.pdtype.sample_placeholder(
[None])
kloldnew = old_pi.proba_distribution.kl(
self.policy_pi.proba_distribution)
ent = self.policy_pi.proba_distribution.entropy()
meankl = tf.reduce_mean(kloldnew)
meanent = tf.reduce_mean(ent)
pol_entpen = (-self.entcoeff) * meanent
# pnew / pold
ratio = tf.exp(
self.policy_pi.proba_distribution.logp(action_ph) -
old_pi.proba_distribution.logp(action_ph))
# surrogate from conservative policy iteration
surr1 = ratio * atarg
surr2 = tf.clip_by_value(ratio, 1.0 - clip_param,
1.0 + clip_param) * atarg
# PPO's pessimistic surrogate (L^CLIP)
pol_surr = -tf.reduce_mean(tf.minimum(surr1, surr2))
vf_loss = tf.reduce_mean(
tf.square(self.policy_pi.value_flat - ret))
total_loss = pol_surr + pol_entpen + vf_loss
losses = [pol_surr, pol_entpen, vf_loss, meankl, meanent]
self.loss_names = [
"pol_surr", "pol_entpen", "vf_loss", "kl", "ent"
]
tf.summary.scalar('entropy_loss', pol_entpen)
tf.summary.scalar('policy_gradient_loss', pol_surr)
tf.summary.scalar('value_function_loss', vf_loss)
tf.summary.scalar('approximate_kullback-leibler', meankl)
tf.summary.scalar('clip_factor', clip_param)
tf.summary.scalar('loss', total_loss)
self.params = tf_util.get_trainable_vars("model")
self.assign_old_eq_new = tf_util.function(
[], [],
updates=[
tf.assign(oldv, newv) for (oldv, newv) in zipsame(
tf_util.get_globals_vars("oldpi"),
tf_util.get_globals_vars("model"))
])
with tf.variable_scope("Adam_mpi", reuse=False):
self.adam = MpiAdam(self.params,
epsilon=self.adam_epsilon,
sess=self.sess)
with tf.variable_scope("input_info", reuse=False):
tf.summary.scalar('discounted_rewards',
tf.reduce_mean(ret))
tf.summary.scalar('learning_rate',
tf.reduce_mean(self.optim_stepsize))
tf.summary.scalar('advantage', tf.reduce_mean(atarg))
tf.summary.scalar('clip_range',
tf.reduce_mean(self.clip_param))
if self.full_tensorboard_log:
tf.summary.histogram('discounted_rewards', ret)
tf.summary.histogram('learning_rate',
self.optim_stepsize)
tf.summary.histogram('advantage', atarg)
tf.summary.histogram('clip_range', self.clip_param)
if tf_util.is_image(self.observation_space):
tf.summary.image('observation', obs_ph)
else:
tf.summary.histogram('observation', obs_ph)
self.step = self.policy_pi.step
self.proba_step = self.policy_pi.proba_step
self.initial_state = self.policy_pi.initial_state
tf_util.initialize(sess=self.sess)
self.summary = tf.summary.merge_all()
self.lossandgrad = tf_util.function(
[obs_ph, old_pi.obs_ph, action_ph, atarg, ret, lrmult],
[self.summary,
tf_util.flatgrad(total_loss, self.params)] + losses)
self.compute_losses = tf_util.function(
[obs_ph, old_pi.obs_ph, action_ph, atarg, ret, lrmult],
losses)
def __init__(self,
observation_space,
action_space,
hidden_size,
entcoeff=0.001,
scope="adversary",
normalize=True):
"""
Reward regression from observations and transitions
:param observation_space: (gym.spaces)
:param action_space: (gym.spaces)
:param hidden_size: ([int]) the hidden dimension for the MLP
:param entcoeff: (float) the entropy loss weight
:param scope: (str) tensorflow variable scope
:param normalize: (bool) Whether to normalize the reward or not
"""
# TODO: support images properly (using a CNN)
self.scope = scope
self.observation_shape = observation_space.shape
self.actions_shape = action_space.shape
if isinstance(action_space, gym.spaces.Box):
# Continuous action space
self.discrete_actions = False
self.n_actions = action_space.shape[0]
elif isinstance(action_space, gym.spaces.Discrete):
self.n_actions = action_space.n
self.discrete_actions = True
else:
raise ValueError(
'Action space not supported: {}'.format(action_space))
self.hidden_size = hidden_size
self.normalize = normalize
self.obs_rms = None
# Placeholders
self.generator_obs_ph = tf.placeholder(
observation_space.dtype, (None, ) + self.observation_shape,
name="observations_ph")
self.generator_acs_ph = tf.placeholder(action_space.dtype,
(None, ) + self.actions_shape,
name="actions_ph")
self.expert_obs_ph = tf.placeholder(observation_space.dtype,
(None, ) + self.observation_shape,
name="expert_observations_ph")
self.expert_acs_ph = tf.placeholder(action_space.dtype,
(None, ) + self.actions_shape,
name="expert_actions_ph")
# Build graph
generator_logits = self.build_graph(self.generator_obs_ph,
self.generator_acs_ph,
reuse=False)
expert_logits = self.build_graph(self.expert_obs_ph,
self.expert_acs_ph,
reuse=True)
# Build accuracy
generator_acc = tf.reduce_mean(
tf.cast(tf.nn.sigmoid(generator_logits) < 0.5, tf.float32))
expert_acc = tf.reduce_mean(
tf.cast(tf.nn.sigmoid(expert_logits) > 0.5, tf.float32))
# Build regression loss
# let x = logits, z = targets.
# z * -log(sigmoid(x)) + (1 - z) * -log(1 - sigmoid(x))
generator_loss = tf.nn.sigmoid_cross_entropy_with_logits(
logits=generator_logits, labels=tf.zeros_like(generator_logits))
generator_loss = tf.reduce_mean(generator_loss)
expert_loss = tf.nn.sigmoid_cross_entropy_with_logits(
logits=expert_logits, labels=tf.ones_like(expert_logits))
expert_loss = tf.reduce_mean(expert_loss)
# Build entropy loss
logits = tf.concat([generator_logits, expert_logits], 0)
entropy = tf.reduce_mean(logit_bernoulli_entropy(logits))
entropy_loss = -entcoeff * entropy
# Loss + Accuracy terms
self.losses = [
generator_loss, expert_loss, entropy, entropy_loss, generator_acc,
expert_acc
]
self.loss_name = [
"generator_loss", "expert_loss", "entropy", "entropy_loss",
"generator_acc", "expert_acc"
]
self.total_loss = generator_loss + expert_loss + entropy_loss
# Build Reward for policy
self.reward_op = -tf.log(1 - tf.nn.sigmoid(generator_logits) + 1e-8)
var_list = self.get_trainable_variables()
self.lossandgrad = tf_util.function([
self.generator_obs_ph, self.generator_acs_ph, self.expert_obs_ph,
self.expert_acs_ph
], self.losses + [tf_util.flatgrad(self.total_loss, var_list)])