def get_kl_loss(self, means, log_vars, means_p=None, log_vars_p=None):
"""Get KL loss for all the predicted Gaussians."""
kl_loss = 0.0
if means_p is None:
means_p = tf.unstack(tf.zeros_like(means))
if log_vars_p is None:
log_vars_p = tf.unstack(tf.zeros_like(log_vars))
enumerated_inputs = enumerate(zip(means, log_vars, means_p,
log_vars_p))
if self.is_training and self.hparams.stochastic_model:
for i, (mean, log_var, mean_p, log_var_p) in enumerated_inputs:
# Condition to compute kl divergence with learned prior
kl_loss += common_layers.kl_divergence(mean, log_var, mean_p,
log_var_p)
tf.summary.histogram("posterior_mean_%d" % i, mean)
tf.summary.histogram("posterior_log_var_%d" % i, log_var)
tf.summary.histogram("prior_mean_%d" % i, mean_p)
tf.summary.histogram("prior_log_var_%d" % i, log_var_p)
tf.summary.scalar("kl_raw", tf.reduce_mean(kl_loss))
beta = self.get_beta(kl_loss)
# information capacity from "Understanding disentangling in beta-VAE"
if self.hparams.information_capacity > 0.0:
kl_loss = tf.abs(kl_loss - self.hparams.information_capacity)
return beta * kl_loss
def get_extra_loss(self, mean, std):
"""Losses in addition to the default modality losses."""
beta = self.get_beta()
kl_loss = common_layers.kl_divergence(mean, std)
tf.summary.histogram("posterior_mean", mean)
tf.summary.histogram("posterior_std", std)
tf.summary.scalar("kl_raw", tf.reduce_mean(kl_loss))
return beta * kl_loss
def get_extra_loss(self, mean, std):
"""Losses in addition to the default modality losses."""
kl_loss = common_layers.kl_divergence(mean, std)
beta = self.get_beta(kl_loss)
tf.summary.histogram("posterior_mean", mean)
tf.summary.histogram("posterior_std", std)
tf.summary.scalar("kl_raw", tf.reduce_mean(kl_loss))
# information capacity from "Understanding disentangling in beta-VAE"
if self.hparams.information_capacity > 0.0:
kl_loss = tf.abs(kl_loss - self.hparams.information_capacity)
return beta * kl_loss
def get_extra_loss(self, latent_means=None, latent_stds=None,
true_frames=None, gen_frames=None, beta=1.0):
"""Losses in addition to the default modality losses."""
kl_loss = 0.0
if self.is_training:
for i, (mean, std) in enumerate(zip(latent_means, latent_stds)):
kl_loss += common_layers.kl_divergence(mean, std)
tf.summary.histogram("posterior_mean_%d" % i, mean)
tf.summary.histogram("posterior_std_%d" % i, std)
tf.summary.scalar("kl_raw", tf.reduce_mean(kl_loss))
return beta * kl_loss
def get_kl_loss(self, means, stds):
"""Get KL loss for all the predicted Gaussians."""
kl_loss = 0.0
if self.is_training and self.hparams.stochastic_model:
for i, (mean, std) in enumerate(zip(means, stds)):
kl_loss += common_layers.kl_divergence(mean, std)
tf.summary.histogram("posterior_mean_%d" % i, mean)
tf.summary.histogram("posterior_std_%d" % i, std)
tf.summary.scalar("kl_raw", tf.reduce_mean(kl_loss))
beta = self.get_beta(kl_loss)
# information capacity from "Understanding disentangling in beta-VAE"
if self.hparams.information_capacity > 0.0:
kl_loss = tf.abs(kl_loss - self.hparams.information_capacity)
return beta * kl_loss
def get_extra_loss(self,
latent_means=None, latent_stds=None,
true_frames=None, gen_frames=None):
"""Losses in addition to the default modality losses."""
del true_frames
del gen_frames
kl_loss = 0.0
if self.is_training and self.hparams.stochastic_model:
for i, (mean, std) in enumerate(zip(latent_means, latent_stds)):
kl_loss += common_layers.kl_divergence(mean, std)
tf.summary.histogram("posterior_mean_%d" % i, mean)
tf.summary.histogram("posterior_std_%d" % i, std)
tf.summary.scalar("kl_raw", tf.reduce_mean(kl_loss))
beta = self.get_beta(kl_loss)
extra_loss = beta * kl_loss
return extra_loss
def get_kl_loss(self, means, log_vars, means_p=None, log_vars_p=None):
"""Get KL loss for all the predicted Gaussians."""
kl_loss = 0.0
if means_p is None:
means_p = tf.unstack(tf.zeros_like(means))
if log_vars_p is None:
log_vars_p = tf.unstack(tf.zeros_like(log_vars))
enumerated_inputs = enumerate(zip(means, log_vars, means_p, log_vars_p))
if self.is_training and self.hparams.stochastic_model:
for i, (mean, log_var, mean_p, log_var_p) in enumerated_inputs:
kl_loss += common_layers.kl_divergence(mean, log_var, mean_p, log_var_p)
tf.summary.histogram("posterior_mean_%d" % i, mean)
tf.summary.histogram("posterior_log_var_%d" % i, log_var)
tf.summary.histogram("prior_mean_%d" % i, mean_p)
tf.summary.histogram("prior_log_var_%d" % i, log_var_p)
tf.summary.scalar("kl_raw", tf.reduce_mean(kl_loss))
beta = self.get_beta(kl_loss)
# information capacity from "Understanding disentangling in beta-VAE"
if self.hparams.information_capacity > 0.0:
kl_loss = tf.abs(kl_loss - self.hparams.information_capacity)
return beta * kl_loss
def body(self, features):
hparams = self.hparams
batch_size = common_layers.shape_list(features["inputs"])[0]
# Swap time and batch axes.
input_frames = common_video.swap_time_and_batch_axes(
features["inputs"])
target_frames = common_video.swap_time_and_batch_axes(
features["targets"])
# Get actions if exist otherwise use zeros
input_actions = self.get_input_if_exists(
features, "input_action", batch_size,
hparams.video_num_input_frames)
target_actions = self.get_input_if_exists(
features, "target_action", batch_size,
hparams.video_num_target_frames)
# Get rewards if exist otherwise use zeros
input_rewards = self.get_input_if_exists(
features, "input_reward", batch_size,
hparams.video_num_input_frames)
target_rewards = self.get_input_if_exists(
features, "target_reward", batch_size,
hparams.video_num_target_frames)
all_actions = tf.concat([input_actions, target_actions], axis=0)
all_rewards = tf.concat([input_rewards, target_rewards], axis=0)
all_frames = tf.concat([input_frames, target_frames], axis=0)
# Each image is being used twice, in latent tower and main tower.
# This is to make sure we are using the *same* image for both, ...
# ... given how TF queues work.
# NOT sure if this is required at all. Doesn"t hurt though! :)
all_frames = tf.identity(all_frames)
gen_images, gen_rewards, latent_means, latent_stds = self.construct_model(
images=all_frames,
actions=all_actions,
rewards=all_rewards,
)
beta = self.get_beta()
kl_loss = 0.0
if self.is_training:
for i, (mean, std) in enumerate(zip(latent_means, latent_stds)):
kl_loss += common_layers.kl_divergence(mean, std)
tf.summary.histogram("posterior_mean_%d" % i, mean)
tf.summary.histogram("posterior_std_%d" % i, std)
tf.summary.scalar("beta", beta)
tf.summary.scalar("kl_raw", tf.reduce_mean(kl_loss))
extra_loss = beta * kl_loss
# Ignore the predictions from the input frames.
# This is NOT the same as original paper/implementation.
predictions = gen_images[hparams.video_num_input_frames - 1:]
reward_pred = gen_rewards[hparams.video_num_input_frames - 1:]
reward_pred = tf.squeeze(reward_pred,
axis=2) # Remove undeeded dimension.
# TODO(mbz): clean this up!
def fix_video_dims_and_concat_on_x_axis(x):
x = tf.transpose(x, [1, 3, 4, 0, 2])
x = tf.reshape(x, [batch_size, 64, 3, -1])
x = tf.transpose(x, [0, 3, 1, 2])
return x
frames_gd = fix_video_dims_and_concat_on_x_axis(target_frames)
frames_pd = fix_video_dims_and_concat_on_x_axis(predictions)
side_by_side_video = tf.concat([frames_gd, frames_pd], axis=2)
tf.summary.image("full_video", side_by_side_video)
# Swap back time and batch axes.
predictions = common_video.swap_time_and_batch_axes(predictions)
reward_pred = common_video.swap_time_and_batch_axes(reward_pred)
return_targets = predictions
if "target_reward" in features:
return_targets = {
"targets": predictions,
"target_reward": reward_pred
}
return return_targets, extra_loss
def body(self, features):
hparams = self.hparams
batch_size = common_layers.shape_list(features["inputs"])[0]
# Swap time and batch axes.
input_frames = common_video.swap_time_and_batch_axes(
features["inputs"])
target_frames = common_video.swap_time_and_batch_axes(
features["targets"])
# Get actions if exist otherwise use zeros
input_actions = self.get_input_if_exists(
features, "input_action", batch_size,
hparams.video_num_input_frames)
target_actions = self.get_input_if_exists(
features, "target_action", batch_size,
hparams.video_num_target_frames)
# Get rewards if exist otherwise use zeros
input_rewards = self.get_input_if_exists(
features, "input_reward", batch_size,
hparams.video_num_input_frames)
target_rewards = self.get_input_if_exists(
features, "target_reward", batch_size,
hparams.video_num_target_frames)
all_actions = tf.concat([input_actions, target_actions], axis=0)
all_rewards = tf.concat([input_rewards, target_rewards], axis=0)
all_frames = tf.concat([input_frames, target_frames], axis=0)
# Each image is being used twice, in latent tower and main tower.
# This is to make sure we are using the *same* image for both, ...
# ... given how TF queues work.
# NOT sure if this is required at all. Doesn"t hurt though! :)
all_frames = tf.identity(all_frames)
gen_images, gen_rewards, latent_means, latent_stds = self.construct_model(
images=all_frames,
actions=all_actions,
rewards=all_rewards,
)
step_num = tf.train.get_global_step()
# TODO(mbz): what should it be if it"s undefined?
if step_num is None:
step_num = _LARGE_STEP_NUMBER
schedule = self.hparams.latent_loss_multiplier_schedule
second_stage = self.hparams.num_iterations_2nd_stage
# TODO(mechcoder): Add log_annealing schedule.
if schedule == "constant":
beta = tf.cond(tf.greater(step_num, second_stage),
lambda: self.hparams.latent_loss_multiplier,
lambda: 0.0)
elif schedule == "linear_anneal":
# Linearly anneal beta from 0.0 to self.hparams.latent_loss_multiplier.
# between self.hparams.num_iterations_2nd_stage to anneal_end.
# beta = latent_loss * (1 - (global_step - 2nd_stage) / (anneal_end - 2nd_stage)) # pylint:disable=line-too-long
anneal_end = self.hparams.anneal_end
latent_multiplier = self.hparams.latent_loss_multiplier
if anneal_end < second_stage:
raise ValueError("Expected hparams.num_iterations_2nd_stage < "
"hparams.anneal_end %d, got %d." %
(second_stage, anneal_end))
def anneal_loss(step_num):
step_num = tf.cast(step_num, dtype=tf.float32)
fraction = (float(anneal_end) - step_num) / (anneal_end -
second_stage)
return self.hparams.latent_loss_multiplier * (1 - fraction)
beta = tf.case(pred_fn_pairs={
tf.less(step_num, second_stage):
lambda: 0.0,
tf.greater(step_num, anneal_end):
lambda: latent_multiplier
},
default=lambda: anneal_loss(step_num))
kl_loss = 0.0
if self.is_training:
for i, (mean, std) in enumerate(zip(latent_means, latent_stds)):
kl_loss += common_layers.kl_divergence(mean, std)
tf.summary.histogram("posterior_mean_%d" % i, mean)
tf.summary.histogram("posterior_std_%d" % i, std)
tf.summary.scalar("beta", beta)
tf.summary.scalar("kl_raw", tf.reduce_mean(kl_loss))
extra_loss = beta * kl_loss
# Ignore the predictions from the input frames.
# This is NOT the same as original paper/implementation.
predictions = gen_images[hparams.video_num_input_frames - 1:]
reward_pred = gen_rewards[hparams.video_num_input_frames - 1:]
reward_pred = tf.squeeze(reward_pred,
axis=2) # Remove undeeded dimension.
# TODO(mbz): clean this up!
def fix_video_dims_and_concat_on_x_axis(x):
x = tf.transpose(x, [1, 3, 4, 0, 2])
x = tf.reshape(x, [batch_size, 64, 3, -1])
x = tf.transpose(x, [0, 3, 1, 2])
return x
frames_gd = fix_video_dims_and_concat_on_x_axis(target_frames)
frames_pd = fix_video_dims_and_concat_on_x_axis(predictions)
side_by_side_video = tf.concat([frames_gd, frames_pd], axis=2)
tf.summary.image("full_video", side_by_side_video)
# Swap back time and batch axes.
predictions = common_video.swap_time_and_batch_axes(predictions)
reward_pred = common_video.swap_time_and_batch_axes(reward_pred)
return_targets = predictions
if "target_reward" in features:
return_targets = {
"targets": predictions,
"target_reward": reward_pred
}
return return_targets, extra_loss
