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)
self.replay_buffer = ReplayBuffer(self.buffer_size)
with tf.variable_scope("input", reuse=False):
# Create policy and target TF objects
self.policy_tf = self.policy(self.sess, self.observation_space, self.action_space,
**self.policy_kwargs)
self.target_policy = self.policy(self.sess, self.observation_space, self.action_space,
**self.policy_kwargs)
# Initialize Placeholders
self.observations_ph = self.policy_tf.obs_ph
# Normalized observation for pixels
self.processed_obs_ph = self.policy_tf.processed_obs
self.next_observations_ph = self.target_policy.obs_ph
self.processed_next_obs_ph = self.target_policy.processed_obs
self.action_target = self.target_policy.action_ph
self.terminals_ph = tf.placeholder(tf.float32, shape=(None, 1), name='terminals')
self.rewards_ph = tf.placeholder(tf.float32, shape=(None, 1), name='rewards')
self.actions_ph = tf.placeholder(tf.float32, shape=(None,) + self.action_space.shape,
name='actions')
self.learning_rate_ph = tf.placeholder(tf.float32, [], name="learning_rate_ph")
with tf.variable_scope("model", reuse=False):
# Create the policy
# first return value corresponds to deterministic actions
# policy_out corresponds to stochastic actions, used for training
# logp_pi is the log probability of actions taken by the policy
self.deterministic_action, policy_out, logp_pi = self.policy_tf.make_actor(self.processed_obs_ph)
# Monitor the entropy of the policy,
# this is not used for training
self.entropy = tf.reduce_mean(self.policy_tf.entropy)
# Use two Q-functions to improve performance by reducing overestimation bias.
qf1, qf2, value_fn = self.policy_tf.make_critics(self.processed_obs_ph, self.actions_ph,
create_qf=True, create_vf=True)
qf1_pi, qf2_pi, _ = self.policy_tf.make_critics(self.processed_obs_ph,
policy_out, create_qf=True, create_vf=False,
reuse=True)
# Target entropy is used when learning the entropy coefficient
if self.target_entropy == 'auto':
# automatically set target entropy if needed
self.target_entropy = -np.prod(self.env.action_space.shape).astype(np.float32)
else:
# Force conversion
# this will also throw an error for unexpected string
self.target_entropy = float(self.target_entropy)
# The entropy coefficient or entropy can be learned automatically
# see Automating Entropy Adjustment for Maximum Entropy RL section
# of https://arxiv.org/abs/1812.05905
if isinstance(self.ent_coef, str) and self.ent_coef.startswith('auto'):
# Default initial value of ent_coef when learned
init_value = 1.0
if '_' in self.ent_coef:
init_value = float(self.ent_coef.split('_')[1])
assert init_value > 0., "The initial value of ent_coef must be greater than 0"
self.log_ent_coef = tf.get_variable('log_ent_coef', dtype=tf.float32,
initializer=np.log(init_value).astype(np.float32))
self.ent_coef = tf.exp(self.log_ent_coef)
else:
# Force conversion to float
# this will throw an error if a malformed string (different from 'auto')
# is passed
self.ent_coef = float(self.ent_coef)
with tf.variable_scope("target", reuse=False):
# Create the value network
_, _, value_target = self.target_policy.make_critics(self.processed_next_obs_ph,
create_qf=False, create_vf=True)
self.value_target = value_target
with tf.variable_scope("loss", reuse=False):
# Take the min of the two Q-Values (Double-Q Learning)
min_qf_pi = tf.minimum(qf1_pi, qf2_pi)
# Target for Q value regression
q_backup = tf.stop_gradient(
self.rewards_ph +
(1 - self.terminals_ph) * self.gamma * self.value_target
)
# Compute Q-Function loss
# TODO: test with huber loss (it would avoid too high values)
qf1_loss = 0.5 * tf.reduce_mean((q_backup - qf1) ** 2)
qf2_loss = 0.5 * tf.reduce_mean((q_backup - qf2) ** 2)
# Compute the entropy temperature loss
# it is used when the entropy coefficient is learned
ent_coef_loss, entropy_optimizer = None, None
if not isinstance(self.ent_coef, float):
ent_coef_loss = -tf.reduce_mean(
self.log_ent_coef * tf.stop_gradient(logp_pi + self.target_entropy))
entropy_optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate_ph)
# Compute the policy loss
# Alternative: policy_kl_loss = tf.reduce_mean(logp_pi - min_qf_pi)
policy_kl_loss = tf.reduce_mean(self.ent_coef * logp_pi - qf1_pi)
# NOTE: in the original implementation, they have an additional
# regularization loss for the Gaussian parameters
# this is not used for now
# policy_loss = (policy_kl_loss + policy_regularization_loss)
policy_loss = policy_kl_loss
# Target for value fn regression
# We update the vf towards the min of two Q-functions in order to
# reduce overestimation bias from function approximation error.
v_backup = tf.stop_gradient(min_qf_pi - self.ent_coef * logp_pi)
value_loss = 0.5 * tf.reduce_mean((value_fn - v_backup) ** 2)
values_losses = qf1_loss + qf2_loss + value_loss
# Policy train op
# (has to be separate from value train op, because min_qf_pi appears in policy_loss)
policy_optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate_ph)
policy_train_op = policy_optimizer.minimize(policy_loss, var_list=tf_util.get_trainable_vars('model/pi'))
# Value train op
value_optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate_ph)
values_params = tf_util.get_trainable_vars('model/values_fn')
source_params = tf_util.get_trainable_vars("model/values_fn/vf")
target_params = tf_util.get_trainable_vars("target/values_fn/vf")
# Polyak averaging for target variables
self.target_update_op = [
tf.assign(target, (1 - self.tau) * target + self.tau * source)
for target, source in zip(target_params, source_params)
]
# Initializing target to match source variables
target_init_op = [
tf.assign(target, source)
for target, source in zip(target_params, source_params)
]
# Control flow is used because sess.run otherwise evaluates in nondeterministic order
# and we first need to compute the policy action before computing q values losses
with tf.control_dependencies([policy_train_op]):
train_values_op = value_optimizer.minimize(values_losses, var_list=values_params)
self.infos_names = ['policy_loss', 'qf1_loss', 'qf2_loss', 'value_loss', 'entropy']
# All ops to call during one training step
self.step_ops = [policy_loss, qf1_loss, qf2_loss,
value_loss, qf1, qf2, value_fn, logp_pi,
self.entropy, policy_train_op, train_values_op]
# Add entropy coefficient optimization operation if needed
if ent_coef_loss is not None:
with tf.control_dependencies([train_values_op]):
ent_coef_op = entropy_optimizer.minimize(ent_coef_loss, var_list=self.log_ent_coef)
self.infos_names += ['ent_coef_loss', 'ent_coef']
self.step_ops += [ent_coef_op, ent_coef_loss, self.ent_coef]
# Monitor losses and entropy in tensorboard
tf.summary.scalar('policy_loss', policy_loss)
tf.summary.scalar('qf1_loss', qf1_loss)
tf.summary.scalar('qf2_loss', qf2_loss)
tf.summary.scalar('value_loss', value_loss)
tf.summary.scalar('entropy', self.entropy)
if ent_coef_loss is not None:
tf.summary.scalar('ent_coef_loss', ent_coef_loss)
tf.summary.scalar('ent_coef', self.ent_coef)
tf.summary.scalar('learning_rate', tf.reduce_mean(self.learning_rate_ph))
# Retrieve parameters that must be saved
self.params = tf_util.get_trainable_vars("model")
self.target_params = tf_util.get_trainable_vars("target/values_fn/vf")
# Initialize Variables and target network
with self.sess.as_default():
self.sess.run(tf.global_variables_initializer())
self.sess.run(target_init_op)
self.summary = tf.summary.merge_all()
def setup_model(self):
with SetVerbosity(self.verbose):
assert issubclass(self.policy, ActorCriticPolicy), "Error: the input policy for the ACKTR model must be " \
"an instance of common.policies.ActorCriticPolicy."
if isinstance(self.action_space, Box):
raise NotImplementedError(
"WIP: ACKTR does not support Continuous actions yet.")
self.graph = tf.Graph()
with self.graph.as_default():
self.sess = tf_util.make_session(num_cpu=self.nprocs,
graph=self.graph)
n_batch_step = None
n_batch_train = None
if issubclass(self.policy, LstmPolicy):
n_batch_step = self.n_envs
n_batch_train = self.n_envs * self.n_steps
self.model = step_model = self.policy(self.sess,
self.observation_space,
self.action_space,
self.n_envs,
1,
n_batch_step,
reuse=False)
self.params = params = find_trainable_variables("model")
with tf.variable_scope(
"train_model",
reuse=True,
custom_getter=tf_util.outer_scope_getter(
"train_model")):
self.model2 = train_model = self.policy(
self.sess,
self.observation_space,
self.action_space,
self.n_envs,
self.n_steps,
n_batch_train,
reuse=True)
with tf.variable_scope(
"loss",
reuse=False,
custom_getter=tf_util.outer_scope_getter("loss")):
self.advs_ph = advs_ph = tf.placeholder(tf.float32, [None])
self.rewards_ph = rewards_ph = tf.placeholder(
tf.float32, [None])
self.pg_lr_ph = pg_lr_ph = tf.placeholder(tf.float32, [])
self.action_ph = action_ph = train_model.pdtype.sample_placeholder(
[None])
logpac = tf.nn.sparse_softmax_cross_entropy_with_logits(
logits=train_model.policy, labels=action_ph)
self.logits = train_model.policy
# training loss
pg_loss = tf.reduce_mean(advs_ph * logpac)
self.entropy = entropy = tf.reduce_mean(
calc_entropy(train_model.policy))
self.pg_loss = pg_loss = pg_loss - self.ent_coef * entropy
self.vf_loss = vf_loss = mse(
tf.squeeze(train_model.value_fn), rewards_ph)
train_loss = pg_loss + self.vf_coef * vf_loss
# Fisher loss construction
self.pg_fisher = pg_fisher_loss = -tf.reduce_mean(logpac)
sample_net = train_model.value_fn + tf.random_normal(
tf.shape(train_model.value_fn))
self.vf_fisher = vf_fisher_loss = -self.vf_fisher_coef * tf.reduce_mean(
tf.pow(
train_model.value_fn -
tf.stop_gradient(sample_net), 2))
self.joint_fisher = pg_fisher_loss + vf_fisher_loss
tf.summary.scalar('entropy_loss', self.entropy)
tf.summary.scalar('policy_gradient_loss', pg_loss)
tf.summary.scalar('policy_gradient_fisher_loss',
pg_fisher_loss)
tf.summary.scalar('value_function_loss', self.vf_loss)
tf.summary.scalar('value_function_fisher_loss',
vf_fisher_loss)
tf.summary.scalar('loss', train_loss)
self.grads_check = tf.gradients(train_loss, params)
with tf.variable_scope("input_info", reuse=False):
tf.summary.scalar('discounted_rewards',
tf.reduce_mean(self.rewards_ph))
tf.summary.histogram('discounted_rewards', self.rewards_ph)
tf.summary.scalar('learning_rate',
tf.reduce_mean(self.pg_lr_ph))
tf.summary.histogram('learning_rate', self.pg_lr_ph)
tf.summary.scalar('advantage',
tf.reduce_mean(self.advs_ph))
tf.summary.histogram('advantage', self.advs_ph)
if len(self.observation_space.shape) == 3:
tf.summary.image('observation', train_model.obs_ph)
else:
tf.summary.histogram('observation', train_model.obs_ph)
with tf.variable_scope(
"kfac",
reuse=False,
custom_getter=tf_util.outer_scope_getter("kfac")):
with tf.device('/gpu:0'):
self.optim = optim = kfac.KfacOptimizer(
learning_rate=pg_lr_ph,
clip_kl=self.kfac_clip,
momentum=0.9,
kfac_update=1,
epsilon=0.01,
stats_decay=0.99,
async_eigen_decomp=self.async_eigen_decomp,
cold_iter=10,
max_grad_norm=self.max_grad_norm,
verbose=self.verbose)
optim.compute_and_apply_stats(self.joint_fisher,
var_list=params)
self.train_model = train_model
self.step_model = step_model
self.step = step_model.step
self.proba_step = step_model.proba_step
self.value = step_model.value
self.initial_state = step_model.initial_state
tf.global_variables_initializer().run(session=self.sess)
self.summary = tf.summary.merge_all()
def setup_model(self):
with SetVerbosity(self.verbose):
assert issubclass(self.policy, ActorCriticPolicy), "Error: the input policy for the PPO2 model must be " \
"an instance of common.policies.ActorCriticPolicy."
self.n_batch = self.n_envs * self.n_steps
n_cpu = multiprocessing.cpu_count()
if sys.platform == 'darwin':
n_cpu //= 2
self.graph = tf.Graph()
with self.graph.as_default():
self.sess = tf_util.make_session(num_cpu=n_cpu, graph=self.graph)
n_batch_step = None
n_batch_train = None
if issubclass(self.policy, RecurrentActorCriticPolicy):
assert self.n_envs % self.nminibatches == 0, "For recurrent policies, "\
"the number of environments run in parallel should be a multiple of nminibatches."
n_batch_step = self.n_envs
n_batch_train = self.n_batch // self.nminibatches
act_model = self.policy(self.sess, self.observation_space, self.action_space, self.n_envs, 1,
n_batch_step, reuse=False, **self.policy_kwargs)
with tf.variable_scope("train_model", reuse=True,
custom_getter=tf_util.outer_scope_getter("train_model")):
train_model = self.policy(self.sess, self.observation_space, self.action_space,
self.n_envs // self.nminibatches, self.n_steps, n_batch_train,
reuse=True, **self.policy_kwargs)
with tf.variable_scope("loss", reuse=False):
self.action_ph = train_model.pdtype.sample_placeholder([None], name="action_ph")
self.advs_ph = tf.placeholder(tf.float32, [None], name="advs_ph")
self.rewards_ph = tf.placeholder(tf.float32, [None], name="rewards_ph")
self.old_neglog_pac_ph = tf.placeholder(tf.float32, [None], name="old_neglog_pac_ph")
self.old_vpred_ph = tf.placeholder(tf.float32, [None], name="old_vpred_ph")
self.learning_rate_ph = tf.placeholder(tf.float32, [], name="learning_rate_ph")
self.clip_range_ph = tf.placeholder(tf.float32, [], name="clip_range_ph")
neglogpac = train_model.proba_distribution.neglogp(self.action_ph)
self.entropy = tf.reduce_mean(train_model.proba_distribution.entropy())
vpred = train_model.value_flat
# Value function clipping: not present in the original PPO
if self.cliprange_vf is None:
# Default behavior (legacy from OpenAI baselines):
# use the same clipping as for the policy
self.clip_range_vf_ph = self.clip_range_ph
self.cliprange_vf = self.cliprange
elif isinstance(self.cliprange_vf, (float, int)) and self.cliprange_vf < 0:
# Original PPO implementation: no value function clipping
self.clip_range_vf_ph = None
else:
# Last possible behavior: clipping range
# specific to the value function
self.clip_range_vf_ph = tf.placeholder(tf.float32, [], name="clip_range_vf_ph")
if self.clip_range_vf_ph is None:
# No clipping
vpred_clipped = train_model.value_flat
else:
# Clip the different between old and new value
# NOTE: this depends on the reward scaling
vpred_clipped = self.old_vpred_ph + \
tf.clip_by_value(train_model.value_flat - self.old_vpred_ph,
- self.clip_range_vf_ph, self.clip_range_vf_ph)
vf_losses1 = tf.square(vpred - self.rewards_ph)
vf_losses2 = tf.square(vpred_clipped - self.rewards_ph)
self.vf_loss = .5 * tf.reduce_mean(tf.maximum(vf_losses1, vf_losses2))
ratio = tf.exp(self.old_neglog_pac_ph - neglogpac)
self.ratio_max = tf.reduce_max(ratio)
pg_losses = -self.advs_ph * ratio
pg_losses2 = -self.advs_ph * tf.clip_by_value(ratio, 1.0 - self.clip_range_ph, 1.0 +
self.clip_range_ph)
self.pg_loss = tf.reduce_mean(tf.maximum(pg_losses, pg_losses2))
self.approxkl = .5 * tf.reduce_mean(tf.square(neglogpac - self.old_neglog_pac_ph))
self.clipfrac = tf.reduce_mean(tf.cast(tf.greater(tf.abs(ratio - 1.0),
self.clip_range_ph), tf.float32))
loss = self.pg_loss - self.entropy * self.ent_coef + self.vf_loss * self.vf_coef
tf.summary.scalar('entropy_loss', self.entropy)
tf.summary.scalar('policy_gradient_loss', self.pg_loss)
tf.summary.scalar('value_function_loss', self.vf_loss)
tf.summary.scalar('approximate_kullback-leibler', self.approxkl)
tf.summary.scalar('clip_factor', self.clipfrac)
tf.summary.scalar('loss', loss)
with tf.variable_scope('model'):
self.params = tf.trainable_variables()
if self.full_tensorboard_log:
for var in self.params:
tf.summary.histogram(var.name, var)
grads = tf.gradients(loss, self.params)
if self.max_grad_norm is not None:
grads, _grad_norm = tf.clip_by_global_norm(grads, self.max_grad_norm)
grads = list(zip(grads, self.params))
trainer = tf.train.AdamOptimizer(learning_rate=self.learning_rate_ph, epsilon=1e-5)
self._train = trainer.apply_gradients(grads)
self.loss_names = ['policy_loss', 'value_loss', 'policy_entropy', 'approxkl', 'clipfrac', 'max_neglogp', 'mean_neglogp', 'ratio_max']
with tf.variable_scope("input_info", reuse=False):
tf.summary.scalar('discounted_rewards', tf.reduce_mean(self.rewards_ph))
tf.summary.scalar('learning_rate', tf.reduce_mean(self.learning_rate_ph))
tf.summary.scalar('advantage', tf.reduce_mean(self.advs_ph))
tf.summary.scalar('clip_range', tf.reduce_mean(self.clip_range_ph))
if self.clip_range_vf_ph is not None:
tf.summary.scalar('clip_range_vf', tf.reduce_mean(self.clip_range_vf_ph))
tf.summary.scalar('old_neglog_action_probabilty', tf.reduce_mean(self.old_neglog_pac_ph))
tf.summary.scalar('old_value_pred', tf.reduce_mean(self.old_vpred_ph))
if self.full_tensorboard_log:
tf.summary.histogram('discounted_rewards', self.rewards_ph)
tf.summary.histogram('learning_rate', self.learning_rate_ph)
tf.summary.histogram('advantage', self.advs_ph)
tf.summary.histogram('clip_range', self.clip_range_ph)
tf.summary.histogram('old_neglog_action_probabilty', self.old_neglog_pac_ph)
tf.summary.histogram('old_value_pred', self.old_vpred_ph)
if tf_util.is_image(self.observation_space):
tf.summary.image('observation', train_model.obs_ph)
else:
tf.summary.histogram('observation', train_model.obs_ph)
self.train_model = train_model
self.act_model = act_model
self.step = act_model.step
self.proba_step = act_model.proba_step
self.value = act_model.value
self.initial_state = act_model.initial_state
tf.global_variables_initializer().run(session=self.sess) # pylint: disable=E1101
self.summary = tf.summary.merge_all()
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 setup_model(self):
with SetVerbosity(self.verbose):
assert issubclass(self.policy, ActorCriticPolicy), "Error: the input policy for the PPO2 model must be " \
"an instance of common.policies.ActorCriticPolicy."
self.n_batch = self.n_envs * self.n_steps
n_cpu = multiprocessing.cpu_count()
if sys.platform == 'darwin':
n_cpu //= 2
self.graph = tf.Graph()
with self.graph.as_default():
self.sess = tf_util.make_session(num_cpu=n_cpu, graph=self.graph)
n_batch_step = None
n_batch_train = None
if issubclass(self.policy, LstmPolicy):
assert self.n_envs % self.nminibatches == 0, "For recurrent policies, "\
"the number of environments run in parallel should be a multiple of nminibatches."
n_batch_step = self.n_envs
n_batch_train = self.n_batch // self.nminibatches
act_model = self.policy(self.sess, self.observation_space, self.action_space, self.n_envs, 1,
n_batch_step, reuse=False)
with tf.variable_scope("train_model", reuse=True,
custom_getter=tf_util.outer_scope_getter("train_model")):
train_model = self.policy(self.sess, self.observation_space, self.action_space,
self.n_envs // self.nminibatches, self.n_steps, n_batch_train,
reuse=True)
with tf.variable_scope("loss", reuse=False):
self.action_ph = train_model.pdtype.sample_placeholder([None], name="action_ph")
self.advs_ph = tf.placeholder(tf.float32, [None], name="advs_ph")
self.rewards_ph = tf.placeholder(tf.float32, [None], name="rewards_ph")
self.old_neglog_pac_ph = tf.placeholder(tf.float32, [None], name="old_neglog_pac_ph")
self.old_vpred_ph = tf.placeholder(tf.float32, [None], name="old_vpred_ph")
self.learning_rate_ph = tf.placeholder(tf.float32, [], name="learning_rate_ph")
self.clip_range_ph = tf.placeholder(tf.float32, [], name="clip_range_ph")
neglogpac = train_model.proba_distribution.neglogp(self.action_ph)
self.entropy = tf.reduce_mean(train_model.proba_distribution.entropy())
vpred = train_model.value_fn
vpredclipped = self.old_vpred_ph + tf.clip_by_value(
train_model.value_fn - self.old_vpred_ph, - self.clip_range_ph, self.clip_range_ph)
vf_losses1 = tf.square(vpred - self.rewards_ph)
vf_losses2 = tf.square(vpredclipped - self.rewards_ph)
self.vf_loss = .5 * tf.reduce_mean(tf.maximum(vf_losses1, vf_losses2))
ratio = tf.exp(self.old_neglog_pac_ph - neglogpac)
pg_losses = -self.advs_ph * ratio
pg_losses2 = -self.advs_ph * tf.clip_by_value(ratio, 1.0 - self.clip_range_ph, 1.0 +
self.clip_range_ph)
self.pg_loss = tf.reduce_mean(tf.maximum(pg_losses, pg_losses2))
self.approxkl = .5 * tf.reduce_mean(tf.square(neglogpac - self.old_neglog_pac_ph))
self.clipfrac = tf.reduce_mean(tf.to_float(tf.greater(tf.abs(ratio - 1.0), self.clip_range_ph)))
loss = self.pg_loss - self.entropy * self.ent_coef + self.vf_loss * self.vf_coef
tf.summary.scalar('entropy_loss', self.entropy)
tf.summary.scalar('policy_gradient_loss', self.pg_loss)
tf.summary.scalar('value_function_loss', self.vf_loss)
tf.summary.scalar('approximate_kullback-leiber', self.approxkl)
tf.summary.scalar('clip_factor', self.clipfrac)
tf.summary.scalar('loss', loss)
with tf.variable_scope('model'):
self.params = tf.trainable_variables()
grads = tf.gradients(loss, self.params)
if self.max_grad_norm is not None:
grads, _grad_norm = tf.clip_by_global_norm(grads, self.max_grad_norm)
grads = list(zip(grads, self.params))
trainer = tf.train.AdamOptimizer(learning_rate=self.learning_rate_ph, epsilon=1e-5)
self._train = trainer.apply_gradients(grads)
self.loss_names = ['policy_loss', 'value_loss', 'policy_entropy', 'approxkl', 'clipfrac']
with tf.variable_scope("input_info", reuse=False):
tf.summary.scalar('discounted_rewards', tf.reduce_mean(self.rewards_ph))
tf.summary.histogram('discounted_rewards', self.rewards_ph)
tf.summary.scalar('learning_rate', tf.reduce_mean(self.learning_rate_ph))
tf.summary.histogram('learning_rate', self.learning_rate_ph)
tf.summary.scalar('advantage', tf.reduce_mean(self.advs_ph))
tf.summary.histogram('advantage', self.advs_ph)
tf.summary.scalar('clip_range', tf.reduce_mean(self.clip_range_ph))
tf.summary.histogram('clip_range', self.clip_range_ph)
tf.summary.scalar('old_neglog_action_probabilty', tf.reduce_mean(self.old_neglog_pac_ph))
tf.summary.histogram('old_neglog_action_probabilty', self.old_neglog_pac_ph)
tf.summary.scalar('old_value_pred', tf.reduce_mean(self.old_vpred_ph))
tf.summary.histogram('old_value_pred', self.old_vpred_ph)
if len(self.observation_space.shape) == 3:
tf.summary.image('observation', train_model.obs_ph)
else:
tf.summary.histogram('observation', train_model.obs_ph)
self.train_model = train_model
self.act_model = act_model
self.step = act_model.step
self.proba_step = act_model.proba_step
self.value = act_model.value
self.initial_state = act_model.initial_state
tf.global_variables_initializer().run(session=self.sess) # pylint: disable=E1101
self.summary = tf.summary.merge_all()
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)
self.replay_buffer = ReplayBuffer(self.buffer_size)
with tf.compat.v1.variable_scope("input", reuse=False):
# Create policy and target TF objects
self.policy_tf = self.policy(self.sess,
self.observation_space,
self.action_space,
**self.policy_kwargs)
self.target_policy_tf = self.policy(
self.sess, self.observation_space, self.action_space,
**self.policy_kwargs)
# Initialize Placeholders
self.observations_ph = self.policy_tf.obs_ph
# Normalized observation for pixels
self.processed_obs_ph = self.policy_tf.processed_obs
self.next_observations_ph = self.target_policy_tf.obs_ph
self.processed_next_obs_ph = self.target_policy_tf.processed_obs
self.action_target = self.target_policy_tf.action_ph
self.terminals_ph = tf.compat.v1.placeholder(
tf.float32, shape=(None, 1), name='terminals')
self.rewards_ph = tf.compat.v1.placeholder(tf.float32,
shape=(None, 1),
name='rewards')
self.actions_ph = tf.compat.v1.placeholder(
tf.float32,
shape=(None, ) + self.action_space.shape,
name='actions')
self.learning_rate_ph = tf.compat.v1.placeholder(
tf.float32, [], name="learning_rate_ph")
with tf.compat.v1.variable_scope("model", reuse=False):
# Create the policy
self.policy_out = policy_out = self.policy_tf.make_actor(
self.processed_obs_ph)
# Use two Q-functions to improve performance by reducing overestimation bias
qf1, qf2 = self.policy_tf.make_critics(
self.processed_obs_ph, self.actions_ph)
# Q value when following the current policy
qf1_pi, _ = self.policy_tf.make_critics(
self.processed_obs_ph, policy_out, reuse=True)
with tf.compat.v1.variable_scope("target", reuse=False):
# Create target networks
target_policy_out = self.target_policy_tf.make_actor(
self.processed_next_obs_ph)
# Target policy smoothing, by adding clipped noise to target actions
target_noise = tf.random.normal(
tf.shape(input=target_policy_out),
stddev=self.target_policy_noise)
target_noise = tf.clip_by_value(target_noise,
-self.target_noise_clip,
self.target_noise_clip)
# Clip the noisy action to remain in the bounds [-1, 1] (output of a tanh)
noisy_target_action = tf.clip_by_value(
target_policy_out + target_noise, -1, 1)
# Q values when following the target policy
qf1_target, qf2_target = self.target_policy_tf.make_critics(
self.processed_next_obs_ph, noisy_target_action)
with tf.compat.v1.variable_scope("loss", reuse=False):
# Take the min of the two target Q-Values (clipped Double-Q Learning)
min_qf_target = tf.minimum(qf1_target, qf2_target)
# Targets for Q value regression
q_backup = tf.stop_gradient(self.rewards_ph +
(1 - self.terminals_ph) *
self.gamma * min_qf_target)
# Compute Q-Function loss
qf1_loss = tf.reduce_mean(input_tensor=(q_backup - qf1)**2)
qf2_loss = tf.reduce_mean(input_tensor=(q_backup - qf2)**2)
qvalues_losses = qf1_loss + qf2_loss
# Policy loss: maximise q value
self.policy_loss = policy_loss = -tf.reduce_mean(
input_tensor=qf1_pi)
# Policy train op
# will be called only every n training steps,
# where n is the policy delay
policy_optimizer = tf.compat.v1.train.AdamOptimizer(
learning_rate=self.learning_rate_ph)
policy_train_op = policy_optimizer.minimize(
policy_loss,
var_list=tf_util.get_trainable_vars('model/pi'))
self.policy_train_op = policy_train_op
# Q Values optimizer
qvalues_optimizer = tf.compat.v1.train.AdamOptimizer(
learning_rate=self.learning_rate_ph)
qvalues_params = tf_util.get_trainable_vars(
'model/values_fn/')
# Q Values and policy target params
source_params = tf_util.get_trainable_vars("model/")
target_params = tf_util.get_trainable_vars("target/")
# Polyak averaging for target variables
self.target_ops = [
tf.compat.v1.assign(target, (1 - self.tau) * target +
self.tau * source)
for target, source in zip(target_params, source_params)
]
# Initializing target to match source variables
target_init_op = [
tf.compat.v1.assign(target, source)
for target, source in zip(target_params, source_params)
]
train_values_op = qvalues_optimizer.minimize(
qvalues_losses, var_list=qvalues_params)
self.infos_names = ['qf1_loss', 'qf2_loss']
# All ops to call during one training step
self.step_ops = [
qf1_loss, qf2_loss, qf1, qf2, train_values_op
]
# Monitor losses and entropy in tensorboard
tf.compat.v1.summary.scalar('policy_loss', policy_loss)
tf.compat.v1.summary.scalar('qf1_loss', qf1_loss)
tf.compat.v1.summary.scalar('qf2_loss', qf2_loss)
tf.compat.v1.summary.scalar(
'learning_rate',
tf.reduce_mean(input_tensor=self.learning_rate_ph))
# Retrieve parameters that must be saved
self.params = tf_util.get_trainable_vars("model")
self.target_params = tf_util.get_trainable_vars("target/")
# Initialize Variables and target network
with self.sess.as_default():
self.sess.run(tf.compat.v1.global_variables_initializer())
self.sess.run(target_init_op)
self.summary = tf.compat.v1.summary.merge_all()
def setup_model(self):
with SetVerbosity(self.verbose):
assert issubclass(self.policy, ActorCriticPolicy), "Error: the input policy for the A2C model must be an " \
"instance of common.policies.ActorCriticPolicy."
self.graph = tf.Graph()
with self.graph.as_default():
self.sess = tf_util.make_session(graph=self.graph)
self.n_batch = self.n_envs * self.n_steps
n_batch_step = None
n_batch_train = None
if issubclass(self.policy, LstmPolicy):
n_batch_step = self.n_envs
n_batch_train = self.n_envs * self.n_steps
step_model = self.policy(self.sess,
self.observation_space,
self.action_space,
self.n_envs,
1,
n_batch_step,
reuse=False,
**self.policy_kwargs)
with tf.variable_scope(
"train_model",
reuse=True,
custom_getter=tf_util.outer_scope_getter(
"train_model")):
train_model = self.policy(self.sess,
self.observation_space,
self.action_space,
self.n_envs,
self.n_steps,
n_batch_train,
reuse=True,
**self.policy_kwargs)
with tf.variable_scope("loss", reuse=False):
self.actions_ph = train_model.pdtype.sample_placeholder(
[None], name="action_ph")
self.advs_ph = tf.placeholder(tf.float32, [None],
name="advs_ph")
self.rewards_ph = tf.placeholder(tf.float32, [None],
name="rewards_ph")
self.learning_rate_ph = tf.placeholder(
tf.float32, [], name="learning_rate_ph")
neglogpac = train_model.proba_distribution.neglogp(
self.actions_ph)
self.entropy = tf.reduce_mean(
train_model.proba_distribution.entropy())
self.pg_loss = tf.reduce_mean(self.advs_ph * neglogpac)
self.vf_loss = mse(tf.squeeze(train_model.value_fn),
self.rewards_ph)
loss = self.pg_loss - self.entropy * self.ent_coef + self.vf_loss * self.vf_coef
tf.summary.scalar('entropy_loss', self.entropy)
tf.summary.scalar('policy_gradient_loss', self.pg_loss)
tf.summary.scalar('value_function_loss', self.vf_loss)
tf.summary.scalar('loss', loss)
self.params = find_trainable_variables("model")
grads = tf.gradients(loss, self.params)
if self.max_grad_norm is not None:
grads, _ = tf.clip_by_global_norm(
grads, self.max_grad_norm)
grads = list(zip(grads, self.params))
with tf.variable_scope("input_info", reuse=False):
tf.summary.scalar('discounted_rewards',
tf.reduce_mean(self.rewards_ph))
tf.summary.histogram('discounted_rewards', self.rewards_ph)
tf.summary.scalar('learning_rate',
tf.reduce_mean(self.learning_rate))
tf.summary.histogram('learning_rate', self.learning_rate)
tf.summary.scalar('advantage',
tf.reduce_mean(self.advs_ph))
tf.summary.histogram('advantage', self.advs_ph)
if len(self.observation_space.shape) == 3:
tf.summary.image('observation', train_model.obs_ph)
else:
tf.summary.histogram('observation', train_model.obs_ph)
trainer = tf.train.RMSPropOptimizer(
learning_rate=self.learning_rate_ph,
decay=self.alpha,
epsilon=self.epsilon)
self.apply_backprop = trainer.apply_gradients(grads)
self.train_model = train_model
self.step_model = step_model
self.step = step_model.step
self.proba_step = step_model.proba_step
self.value = step_model.value
self.initial_state = step_model.initial_state
tf.global_variables_initializer().run(session=self.sess)
self.summary = tf.summary.merge_all()
def setup_model(self):
# prevent import loops
from stable_baselines.trpo_mpi.utils 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.compat.v1.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.compat.v1.variable_scope("loss", reuse=False):
atarg = tf.compat.v1.placeholder(dtype=tf.float32, shape=[None]) # Target advantage function (if applicable)
ret = tf.compat.v1.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(input_tensor=kloldnew)
meanent = tf.reduce_mean(input_tensor=ent)
entbonus = self.entcoeff * meanent
explained_variance = explained_variance_tensor(self.policy_pi.value_flat, ret)
variance = variance_tensor(ret)
vferr = self.vf_coef*tf.reduce_mean(input_tensor=tf.square(self.policy_pi.value_flat - ret))\
+ self.avec_coef*(1-explained_variance)*variance
# advantage * pnew / pold
ratio = tf.exp(self.policy_pi.proba_distribution.logp(action) -
old_policy.proba_distribution.logp(action))
surrgain = tf.reduce_mean(input_tensor=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(ys=dist, xs=var_list)
flat_tangent = tf.compat.v1.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(input_tensor=grad * tangent)
for (grad, tangent) in zipsame(klgrads, tangents)]) # pylint: disable=E1111
# Fisher vector products
fvp = tf_util.flatgrad(gvp, var_list)
tf.compat.v1.summary.scalar('entropy_loss', meanent)
tf.compat.v1.summary.scalar('policy_gradient_loss', optimgain)
tf.compat.v1.summary.scalar('value_function_loss', surrgain)
tf.compat.v1.summary.scalar('approximate_kullback-leibler', meankl)
tf.compat.v1.summary.scalar('loss', optimgain + meankl + entbonus + surrgain + meanent)
self.assign_old_eq_new = \
tf_util.function([], [], updates=[tf.compat.v1.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.compat.v1.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.compat.v1.variable_scope("input_info", reuse=False):
tf.compat.v1.summary.scalar('discounted_rewards', tf.reduce_mean(input_tensor=ret))
tf.compat.v1.summary.scalar('learning_rate', tf.reduce_mean(input_tensor=self.vf_stepsize))
tf.compat.v1.summary.scalar('advantage', tf.reduce_mean(input_tensor=atarg))
tf.compat.v1.summary.scalar('kl_clip_range', tf.reduce_mean(input_tensor=self.max_kl))
if self.full_tensorboard_log:
tf.compat.v1.summary.histogram('discounted_rewards', ret)
tf.compat.v1.summary.histogram('learning_rate', self.vf_stepsize)
tf.compat.v1.summary.histogram('advantage', atarg)
tf.compat.v1.summary.histogram('kl_clip_range', self.max_kl)
if tf_util.is_image(self.observation_space):
tf.compat.v1.summary.image('observation', observation)
else:
tf.compat.v1.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.compat.v1.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):
assert issubclass(self.policy, ActorCriticPolicy), "Error: the input policy for the A2C model must be an " \
"instance of common.policies.ActorCriticPolicy."
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)
self.n_batch = self.n_envs * self.n_steps
n_batch_step = None
n_batch_train = None
if issubclass(self.policy, RecurrentActorCriticPolicy):
n_batch_step = self.n_envs
n_batch_train = self.n_envs * self.n_steps
step_model = self.policy(self.sess, self.observation_space, self.action_space, self.n_envs, 1,
n_batch_step, reuse=False, **self.policy_kwargs)
with tf.variable_scope("train_model", reuse=True,
custom_getter=tf_util.outer_scope_getter("train_model")):
train_model = self.policy(self.sess, self.observation_space, self.action_space, self.n_envs,
self.n_steps, n_batch_train, reuse=True, **self.policy_kwargs)
with tf.variable_scope("loss", reuse=False):
self.actions_ph = train_model.pdtype.sample_placeholder([None], name="action_ph")
self.advs_ph = tf.placeholder(tf.float32, [None], name="advs_ph")
self.rewards_ph = tf.placeholder(tf.float32, [None], name="rewards_ph")
self.learning_rate_ph = tf.placeholder(tf.float32, [], name="learning_rate_ph")
neglogpac = train_model.proba_distribution.neglogp(self.actions_ph)
self.entropy = tf.reduce_mean(train_model.proba_distribution.entropy())
self.pg_loss = tf.reduce_mean(self.advs_ph * neglogpac)
self.vf_loss = mse(tf.squeeze(train_model.value_flat), self.rewards_ph)
# https://arxiv.org/pdf/1708.04782.pdf#page=9, https://arxiv.org/pdf/1602.01783.pdf#page=4
# and https://github.com/dennybritz/reinforcement-learning/issues/34
# suggest to add an entropy component in order to improve exploration.
loss = self.pg_loss - self.entropy * self.ent_coef + self.vf_loss * self.vf_coef
tf.summary.scalar('entropy_loss', self.entropy)
tf.summary.scalar('policy_gradient_loss', self.pg_loss)
tf.summary.scalar('value_function_loss', self.vf_loss)
tf.summary.scalar('loss', loss)
self.params = tf_util.get_trainable_vars("model")
grads = tf.gradients(loss, self.params)
if self.max_grad_norm is not None:
grads, _ = tf.clip_by_global_norm(grads, self.max_grad_norm)
grads = list(zip(grads, self.params))
with tf.variable_scope("input_info", reuse=False):
tf.summary.scalar('discounted_rewards', tf.reduce_mean(self.rewards_ph))
tf.summary.scalar('learning_rate', tf.reduce_mean(self.learning_rate_ph))
tf.summary.scalar('advantage', tf.reduce_mean(self.advs_ph))
if self.full_tensorboard_log:
tf.summary.histogram('discounted_rewards', self.rewards_ph)
tf.summary.histogram('learning_rate', self.learning_rate_ph)
tf.summary.histogram('advantage', self.advs_ph)
if tf_util.is_image(self.observation_space):
tf.summary.image('observation', train_model.obs_ph)
else:
tf.summary.histogram('observation', train_model.obs_ph)
trainer = tf.train.RMSPropOptimizer(learning_rate=self.learning_rate_ph, decay=self.alpha,
epsilon=self.epsilon)
self.apply_backprop = trainer.apply_gradients(grads)
self.train_model = train_model
self.step_model = step_model
self.step = step_model.step
self.proba_step = step_model.proba_step
self.value = step_model.value
self.initial_state = step_model.initial_state
tf.global_variables_initializer().run(session=self.sess)
self.summary = tf.summary.merge_all()
def setup_model(self):
with SetVerbosity(self.verbose):
assert issubclass(self.policy, ActorCriticPolicy), "Error: the input policy for the ACKTR model must be " \
"an instance of common.policies.ActorCriticPolicy."
# Enable continuous actions tricks (normalized advantage)
self.continuous_actions = isinstance(self.action_space, Box)
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)
n_batch_step = None
n_batch_train = None
if issubclass(self.policy, RecurrentActorCriticPolicy):
n_batch_step = self.n_envs
n_batch_train = self.n_envs * self.n_steps
step_model = self.policy(self.sess,
self.observation_space,
self.action_space,
self.n_envs,
1,
n_batch_step,
reuse=False,
**self.policy_kwargs)
self.params = params = tf_util.get_trainable_vars("model")
with tf.variable_scope(
"train_model",
reuse=True,
custom_getter=tf_util.outer_scope_getter(
"train_model")):
train_model = self.policy(self.sess,
self.observation_space,
self.action_space,
self.n_envs,
self.n_steps,
n_batch_train,
reuse=True,
**self.policy_kwargs)
with tf.variable_scope(
"loss",
reuse=False,
custom_getter=tf_util.outer_scope_getter("loss")):
self.advs_ph = advs_ph = tf.placeholder(tf.float32, [None])
self.rewards_ph = rewards_ph = tf.placeholder(
tf.float32, [None])
self.learning_rate_ph = learning_rate_ph = tf.placeholder(
tf.float32, [])
self.actions_ph = train_model.pdtype.sample_placeholder(
[None])
neg_log_prob = train_model.proba_distribution.neglogp(
self.actions_ph)
# training loss
pg_loss = tf.reduce_mean(advs_ph * neg_log_prob)
self.entropy = entropy = tf.reduce_mean(
train_model.proba_distribution.entropy())
self.pg_loss = pg_loss = pg_loss - self.ent_coef * entropy
self.vf_loss = vf_loss = mse(
tf.squeeze(train_model.value_fn), rewards_ph)
train_loss = pg_loss + self.vf_coef * vf_loss
# Fisher loss construction
self.pg_fisher = pg_fisher_loss = -tf.reduce_mean(
neg_log_prob)
sample_net = train_model.value_fn + tf.random_normal(
tf.shape(train_model.value_fn))
self.vf_fisher = vf_fisher_loss = -self.vf_fisher_coef * tf.reduce_mean(
tf.pow(
train_model.value_fn -
tf.stop_gradient(sample_net), 2))
self.joint_fisher = pg_fisher_loss + vf_fisher_loss
tf.summary.scalar('entropy_loss', self.entropy)
tf.summary.scalar('policy_gradient_loss', pg_loss)
tf.summary.scalar('policy_gradient_fisher_loss',
pg_fisher_loss)
tf.summary.scalar('value_function_loss', self.vf_loss)
tf.summary.scalar('value_function_fisher_loss',
vf_fisher_loss)
tf.summary.scalar('loss', train_loss)
self.grads_check = tf.gradients(train_loss, params)
with tf.variable_scope("input_info", reuse=False):
tf.summary.scalar('discounted_rewards',
tf.reduce_mean(self.rewards_ph))
tf.summary.scalar('learning_rate',
tf.reduce_mean(self.learning_rate_ph))
tf.summary.scalar('advantage',
tf.reduce_mean(self.advs_ph))
if self.full_tensorboard_log:
tf.summary.histogram('discounted_rewards',
self.rewards_ph)
tf.summary.histogram('learning_rate',
self.learning_rate_ph)
tf.summary.histogram('advantage', self.advs_ph)
if tf_util.is_image(self.observation_space):
tf.summary.image('observation', train_model.obs_ph)
else:
tf.summary.histogram('observation',
train_model.obs_ph)
with tf.variable_scope(
"kfac",
reuse=False,
custom_getter=tf_util.outer_scope_getter("kfac")):
with tf.device('/gpu:0'):
self.optim = optim = kfac.KfacOptimizer(
learning_rate=learning_rate_ph,
clip_kl=self.kfac_clip,
momentum=0.9,
kfac_update=self.kfac_update,
epsilon=0.01,
stats_decay=0.99,
async_eigen_decomp=self.async_eigen_decomp,
cold_iter=10,
max_grad_norm=self.max_grad_norm,
verbose=self.verbose)
optim.compute_and_apply_stats(self.joint_fisher,
var_list=params)
self.train_model = train_model
self.step_model = step_model
self.step = step_model.step
self.proba_step = step_model.proba_step
self.value = step_model.value
self.initial_state = step_model.initial_state
tf.global_variables_initializer().run(session=self.sess)
self.summary = tf.summary.merge_all()
