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,
momentum=self.momentum)
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()