def train(self):
"""Perform training."""
assert (None not in (self._algorithm, self._driver)) and self._envs, \
"Trainer not initialized"
self._restore_checkpoint()
common.enable_summary(True)
run_under_record_context(self._train,
summary_dir=self._train_dir,
summary_interval=self._summary_interval,
flush_millis=self._summaries_flush_mills,
summary_max_queue=self._summary_max_queue)
self._save_checkpoint()
self._close_envs()
def _train(self, experience, num_updates, mini_batch_size,
mini_batch_length, update_counter_every_mini_batch,
should_summarize):
"""Train using experience."""
experience = nest_utils.params_to_distributions(
experience, self.experience_spec)
experience = self.transform_timestep(experience)
experience = self.preprocess_experience(experience)
experience = nest_utils.distributions_to_params(experience)
length = experience.step_type.shape[1]
mini_batch_length = (mini_batch_length or length)
assert length % mini_batch_length == 0, (
"length=%s not a multiple of mini_batch_length=%s" %
(length, mini_batch_length))
if len(tf.nest.flatten(
self.train_state_spec)) > 0 and not self._use_rollout_state:
if mini_batch_length == 1:
logging.fatal(
"Should use TrainerConfig.use_rollout_state=True "
"for off-policy training of RNN when minibatch_length==1.")
else:
common.warning_once(
"Consider using TrainerConfig.use_rollout_state=True "
"for off-policy training of RNN.")
experience = tf.nest.map_structure(
lambda x: tf.reshape(
x, common.concat_shape([-1, mini_batch_length],
tf.shape(x)[2:])), experience)
batch_size = tf.shape(experience.step_type)[0]
mini_batch_size = (mini_batch_size or batch_size)
def _make_time_major(nest):
"""Put the time dim to axis=0."""
return tf.nest.map_structure(lambda x: common.transpose2(x, 0, 1),
nest)
scope = get_current_scope()
for u in tf.range(num_updates):
if mini_batch_size < batch_size:
indices = tf.random.shuffle(
tf.range(tf.shape(experience.step_type)[0]))
experience = tf.nest.map_structure(
lambda x: tf.gather(x, indices), experience)
for b in tf.range(0, batch_size, mini_batch_size):
if update_counter_every_mini_batch:
common.get_global_counter().assign_add(1)
is_last_mini_batch = tf.logical_and(
tf.equal(u, num_updates - 1),
tf.greater_equal(b + mini_batch_size, batch_size))
do_summary = tf.logical_or(is_last_mini_batch,
update_counter_every_mini_batch)
common.enable_summary(do_summary)
batch = tf.nest.map_structure(
lambda x: x[b:tf.minimum(batch_size, b + mini_batch_size)],
experience)
batch = _make_time_major(batch)
# Tensorflow graph mode loses the original name scope here. We
# need to restore the original name scope
with tf.name_scope(scope):
training_info, loss_info, grads_and_vars = self._update(
batch,
weight=tf.cast(
tf.shape(batch.step_type)[1], tf.float32) /
float(mini_batch_size))
if should_summarize:
if do_summary:
# Putting `if do_summary` under the above `with` statement
# does not help. Somehow `if` statement will also lose
# the original name scope.
with tf.name_scope(scope):
self.summarize_train(training_info, loss_info,
grads_and_vars)
train_steps = batch_size * mini_batch_length * num_updates
return train_steps
def train(self,
experience: Experience,
num_updates=1,
mini_batch_size=None,
mini_batch_length=None):
"""Train using `experience`.
Args:
experience (Experience): experience from replay_buffer. It is
assumed to be batch major.
num_updates (int): number of optimization steps
mini_batch_size (int): number of sequences for each minibatch
mini_batch_length (int): the length of the sequence for each
sample in the minibatch
"""
experience = self._algorithm.preprocess_experience(experience)
length = experience.step_type.shape[1]
mini_batch_length = (mini_batch_length or length)
assert length % mini_batch_length == 0, (
"length=%s not a multiple of mini_batch_length=%s" %
(length, mini_batch_length))
experience = tf.nest.map_structure(
lambda x: tf.reshape(x, [-1, mini_batch_length] + list(x.shape[2:])
), experience)
batch_size = experience.step_type.shape[0]
mini_batch_size = (mini_batch_size or batch_size)
# The reason of this constraint is tf.reshape at L233
# TODO: remove this constraint.
assert batch_size % mini_batch_size == 0, (
"batch_size=%s not a multiple of mini_batch_size=%s" %
(batch_size, mini_batch_size))
def _make_time_major(nest):
"""Put the time dim to axis=0"""
return tf.nest.map_structure(lambda x: common.transpose2(x, 0, 1),
nest)
for u in tf.range(num_updates):
if mini_batch_size < batch_size:
indices = tf.random.shuffle(
tf.range(experience.step_type.shape[0]))
experience = tf.nest.map_structure(
lambda x: tf.gather(x, indices), experience)
for b in tf.range(0, batch_size, mini_batch_size):
batch = tf.nest.map_structure(
lambda x: x[b:tf.minimum(batch_size, b + mini_batch_size)],
experience)
# Make the shape explicit. The shapes of tensors from the
# previous line depend on tensor `b`, which is replaced with
# None by tf. This makes some operations depending on the shape
# of tensor fail. (Currently, it's alf.common.tensor_extend)
# TODO: Find a way to work around with shapes containing None
# at common.tensor_extend()
batch = tf.nest.map_structure(
lambda x: tf.reshape(x, [mini_batch_size] + list(x.shape)[
1:]), batch)
batch = _make_time_major(batch)
is_last_mini_batch = tf.logical_and(
tf.equal(u, num_updates - 1),
tf.greater_equal(b + mini_batch_size, batch_size))
common.enable_summary(is_last_mini_batch)
training_info, loss_info, grads_and_vars = self._update(
batch, weight=batch.step_type.shape[1] / mini_batch_size)
if is_last_mini_batch:
self._training_summary(training_info, loss_info,
grads_and_vars)
self._train_step_counter.assign_add(1)
def train(self,
experience: Experience,
num_updates=1,
mini_batch_size=None,
mini_batch_length=None):
"""Train using `experience`.
Args:
experience (Experience): experience from replay_buffer. It is
assumed to be batch major.
num_updates (int): number of optimization steps
mini_batch_size (int): number of sequences for each minibatch
mini_batch_length (int): the length of the sequence for each
sample in the minibatch
Returns:
train_steps (int): the actual number of time steps that have been
trained (a step might be trained multiple times)
"""
experience = self._algorithm.preprocess_experience(experience)
length = experience.step_type.shape[1]
mini_batch_length = (mini_batch_length or length)
assert length % mini_batch_length == 0, (
"length=%s not a multiple of mini_batch_length=%s" %
(length, mini_batch_length))
if len(tf.nest.flatten(self._algorithm.train_state_spec)
) > 0 and not self._use_rollout_state:
if mini_batch_length == 1:
logging.fatal(
"Should use TrainerConfig.use_rollout_state=True "
"for off-policy training of RNN when minibatch_length==1.")
else:
warning_once(
"Consider using TrainerConfig.use_rollout_state=True "
"for off-policy training of RNN.")
experience = tf.nest.map_structure(
lambda x: tf.reshape(
x, common.concat_shape([-1, mini_batch_length],
tf.shape(x)[2:])), experience)
batch_size = tf.shape(experience.step_type)[0]
mini_batch_size = (mini_batch_size or batch_size)
def _make_time_major(nest):
"""Put the time dim to axis=0."""
return tf.nest.map_structure(lambda x: common.transpose2(x, 0, 1),
nest)
for u in tf.range(num_updates):
if mini_batch_size < batch_size:
indices = tf.random.shuffle(
tf.range(tf.shape(experience.step_type)[0]))
experience = tf.nest.map_structure(
lambda x: tf.gather(x, indices), experience)
for b in tf.range(0, batch_size, mini_batch_size):
batch = tf.nest.map_structure(
lambda x: x[b:tf.minimum(batch_size, b + mini_batch_size)],
experience)
batch = _make_time_major(batch)
is_last_mini_batch = tf.logical_and(
tf.equal(u, num_updates - 1),
tf.greater_equal(b + mini_batch_size, batch_size))
common.enable_summary(is_last_mini_batch)
training_info, loss_info, grads_and_vars = self._update(
batch,
weight=tf.cast(tf.shape(batch.step_type)[1], tf.float32) /
float(mini_batch_size))
if is_last_mini_batch:
self._training_summary(training_info, loss_info,
grads_and_vars)
self._train_step_counter.assign_add(1)
train_steps = batch_size * mini_batch_length * num_updates
return train_steps