def sample_posterior(self, images, actions, step_types, features=None):
sequence_length = step_types.shape[1].value - 1
actions = actions[:, :sequence_length]
if features is None:
features = self.compressor(images)
# swap batch and time axes
features = tf.transpose(features, [1, 0, 2])
actions = tf.transpose(actions, [1, 0, 2])
step_types = tf.transpose(step_types, [1, 0])
latent_dists = []
latent_samples = []
for t in range(sequence_length + 1):
if t == 0:
latent_dist = self.latent_first_posterior(features[t])
latent_sample = latent_dist.sample()
else:
reset_mask = tf.equal(step_types[t], ts.StepType.FIRST)
latent_first_dist = self.latent_first_posterior(features[t])
latent_dist = self.latent_posterior(features[t], latent_samples[t-1], actions[t-1])
latent_dist = nest_utils.map_distribution_structure(
functools.partial(tf.where, reset_mask), latent_first_dist, latent_dist)
latent_sample = latent_dist.sample()
latent_dists.append(latent_dist)
latent_samples.append(latent_sample)
latent_dists = nest_utils.map_distribution_structure(lambda *x: tf.stack(x, axis=1), *latent_dists)
latent_samples = tf.stack(latent_samples, axis=1)
return latent_samples, latent_dists
def sample_prior_or_posterior(self, actions, step_types=None, images=None):
"""Samples from the prior, except for the first time steps in which conditioning images are given."""
if step_types is None:
batch_size = tf.shape(actions)[0]
sequence_length = actions.shape[1].value # should be statically defined
step_types = tf.fill(
[batch_size, sequence_length + 1], ts.StepType.MID)
else:
sequence_length = step_types.shape[1].value - 1
actions = actions[:, :sequence_length]
if images is not None:
features = self.compressor(images)
# swap batch and time axes
actions = tf.transpose(actions, [1, 0, 2])
step_types = tf.transpose(step_types, [1, 0])
if images is not None:
features = tf.transpose(features, [1, 0, 2])
latent_dists = []
latent_samples = []
for t in range(sequence_length + 1):
is_conditional = images is not None and (t < images.shape[1].value)
if t == 0:
if is_conditional:
latent_dist = self.latent_first_posterior(features[t])
else:
latent_dist = self.latent_first_prior(step_types[t]) # step_types is only used to infer batch_size
latent_sample = latent_dist.sample()
else:
reset_mask = tf.equal(step_types[t], ts.StepType.FIRST)
if is_conditional:
latent_first_dist = self.latent_first_posterior(features[t])
latent_dist = self.latent_posterior(features[t], latent_samples[t-1], actions[t-1])
else:
latent_first_dist = self.latent_first_prior(step_types[t])
latent_dist = self.latent_prior(latent_samples[t-1], actions[t-1])
latent_dist = nest_utils.map_distribution_structure(
functools.partial(tf.where, reset_mask), latent_first_dist, latent_dist)
latent_sample = latent_dist.sample()
latent_dists.append(latent_dist)
latent_samples.append(latent_sample)
latent_dists = nest_utils.map_distribution_structure(lambda *x: tf.stack(x, axis=1), *latent_dists)
latent_samples = tf.stack(latent_samples, axis=1)
return latent_samples, latent_dists
def compute_loss(self,
images,
actions,
step_types,
rewards=None,
discounts=None,
latent_posterior_samples_and_dists=None):
sequence_length = step_types.shape[1].value - 1
if latent_posterior_samples_and_dists is None:
latent_posterior_samples_and_dists = self.sample_posterior(
images, actions, step_types)
(latent1_posterior_samples, latent2_posterior_samples), (
latent1_posterior_dists,
latent2_posterior_dists) = (latent_posterior_samples_and_dists)
(latent1_prior_samples,
latent2_prior_samples), _ = self.sample_prior_or_posterior(
actions, step_types) # for visualization
(latent1_conditional_prior_samples, latent2_conditional_prior_samples
), _ = self.sample_prior_or_posterior(
actions, step_types, images=images[:, :1]
) # for visualization. condition on first image only
def where_and_concat(reset_masks, first_prior_tensors,
after_first_prior_tensors):
after_first_prior_tensors = tf.where(reset_masks[:, 1:],
first_prior_tensors[:, 1:],
after_first_prior_tensors)
prior_tensors = tf.concat(
[first_prior_tensors[:, 0:1], after_first_prior_tensors],
axis=1)
return prior_tensors
reset_masks = tf.concat([
tf.ones_like(step_types[:, 0:1], dtype=tf.bool),
tf.equal(step_types[:, 1:], ts.StepType.FIRST)
],
axis=1)
latent1_reset_masks = tf.tile(reset_masks[:, :, None],
[1, 1, self.latent1_size])
latent1_first_prior_dists = self.latent1_first_prior(step_types)
# these distributions start at t=1 and the inputs are from t-1
latent1_after_first_prior_dists = self.latent1_prior(
latent2_posterior_samples[:, :sequence_length],
actions[:, :sequence_length])
latent1_prior_dists = nest_utils.map_distribution_structure(
functools.partial(where_and_concat, latent1_reset_masks),
latent1_first_prior_dists, latent1_after_first_prior_dists)
latent2_reset_masks = tf.tile(reset_masks[:, :, None],
[1, 1, self.latent2_size])
latent2_first_prior_dists = self.latent2_first_prior(
latent1_posterior_samples)
# these distributions start at t=1 and the last 2 inputs are from t-1
latent2_after_first_prior_dists = self.latent2_prior(
latent1_posterior_samples[:, 1:sequence_length + 1],
latent2_posterior_samples[:, :sequence_length],
actions[:, :sequence_length])
latent2_prior_dists = nest_utils.map_distribution_structure(
functools.partial(where_and_concat, latent2_reset_masks),
latent2_first_prior_dists, latent2_after_first_prior_dists)
outputs = {}
if self.kl_analytic:
latent1_kl_divergences = tfd.kl_divergence(latent1_posterior_dists,
latent1_prior_dists)
else:
latent1_kl_divergences = (
latent1_posterior_dists.log_prob(latent1_posterior_samples) -
latent1_prior_dists.log_prob(latent1_posterior_samples))
latent1_kl_divergences = tf.reduce_sum(latent1_kl_divergences, axis=1)
outputs.update({
'latent1_kl_divergence':
tf.reduce_mean(latent1_kl_divergences),
})
if self.latent2_posterior == self.latent2_prior:
latent2_kl_divergences = 0.0
else:
if self.kl_analytic:
latent2_kl_divergences = tfd.kl_divergence(
latent2_posterior_dists, latent2_prior_dists)
else:
latent2_kl_divergences = (
latent2_posterior_dists.log_prob(latent2_posterior_samples)
- latent2_prior_dists.log_prob(latent2_posterior_samples))
latent2_kl_divergences = tf.reduce_sum(latent2_kl_divergences,
axis=1)
outputs.update({
'latent2_kl_divergence':
tf.reduce_mean(latent2_kl_divergences),
})
outputs.update({
'kl_divergence':
tf.reduce_mean(latent1_kl_divergences + latent2_kl_divergences),
})
likelihood_dists = self.decoder(latent1_posterior_samples,
latent2_posterior_samples)
likelihood_log_probs = likelihood_dists.log_prob(images)
likelihood_log_probs = tf.reduce_sum(likelihood_log_probs, axis=1)
reconstruction_error = tf.reduce_sum(
tf.square(images - likelihood_dists.distribution.loc),
axis=list(range(-len(likelihood_dists.event_shape), 0)))
reconstruction_error = tf.reduce_sum(reconstruction_error, axis=1)
outputs.update({
'log_likelihood':
tf.reduce_mean(likelihood_log_probs),
'reconstruction_error':
tf.reduce_mean(reconstruction_error),
})
# summed over the time dimension
elbo = likelihood_log_probs - latent1_kl_divergences - latent2_kl_divergences
if self.model_reward:
reward_dists = self.reward_predictor(
latent1_posterior_samples[:, :sequence_length],
latent2_posterior_samples[:, :sequence_length],
actions[:, :sequence_length],
latent1_posterior_samples[:, 1:sequence_length + 1],
latent2_posterior_samples[:, 1:sequence_length + 1])
reward_valid_mask = tf.cast(
tf.not_equal(step_types[:, :sequence_length],
ts.StepType.LAST), tf.float32)
reward_log_probs = reward_dists.log_prob(
rewards[:, :sequence_length])
reward_log_probs = tf.reduce_sum(reward_log_probs *
reward_valid_mask,
axis=1)
reward_reconstruction_error = tf.square(
rewards[:, :sequence_length] - reward_dists.loc)
reward_reconstruction_error = tf.reduce_sum(
reward_reconstruction_error * reward_valid_mask, axis=1)
outputs.update({
'reward_log_likelihood':
tf.reduce_mean(reward_log_probs),
'reward_reconstruction_error':
tf.reduce_mean(reward_reconstruction_error),
})
elbo += reward_log_probs
if self.model_discount:
discount_dists = self.discount_predictor(
latent1_posterior_samples[:, 1:sequence_length + 1],
latent2_posterior_samples[:, 1:sequence_length + 1])
discount_log_probs = discount_dists.log_prob(
discounts[:, :sequence_length])
discount_log_probs = tf.reduce_sum(discount_log_probs, axis=1)
discount_accuracy = tf.cast(
tf.equal(tf.cast(discount_dists.mode(), tf.float32),
discounts[:, :sequence_length]), tf.float32)
discount_accuracy = tf.reduce_sum(discount_accuracy, axis=1)
outputs.update({
'discount_log_likelihood':
tf.reduce_mean(discount_log_probs),
'discount_accuracy':
tf.reduce_mean(discount_accuracy),
})
elbo += discount_log_probs
# average over the batch dimension
loss = -tf.reduce_mean(elbo)
posterior_images = likelihood_dists.mean()
prior_images = self.decoder(latent1_prior_samples,
latent2_prior_samples).mean()
conditional_prior_images = self.decoder(
latent1_conditional_prior_samples,
latent2_conditional_prior_samples).mean()
outputs.update({
'elbo': tf.reduce_mean(elbo),
'images': images,
'posterior_images': posterior_images,
'prior_images': prior_images,
'conditional_prior_images': conditional_prior_images,
})
return loss, outputs
def model_loss(self,
images,
actions,
step_types,
rewards,
discounts,
latent_posterior_samples_and_dists=None,
weights=None):
with tf.name_scope('model_loss'):
if self._model_batch_size is not None:
# Allow model batch size to be smaller than the batch size of the
# other losses. This is because the model loss already gets a lot of
# supervision from having a loss over all time steps.
images, actions, step_types, rewards, discounts = tf.nest.map_structure(
lambda x: x[:self._model_batch_size],
(images, actions, step_types, rewards, discounts))
if latent_posterior_samples_and_dists is not None:
latent_posterior_samples, latent_posterior_dists = latent_posterior_samples_and_dists
latent_posterior_samples = tf.nest.map_structure(
lambda x: x[:self._model_batch_size],
latent_posterior_samples)
latent_posterior_dists = slac_nest_utils.map_distribution_structure(
lambda x: x[:self._model_batch_size],
latent_posterior_dists)
latent_posterior_samples_and_dists = (
latent_posterior_samples, latent_posterior_dists)
model_loss, outputs = self._model_network.compute_loss(
images,
actions,
step_types,
rewards=rewards,
discounts=discounts,
latent_posterior_samples_and_dists=
latent_posterior_samples_and_dists)
for name, output in outputs.items():
if output.shape.ndims == 0:
tf.contrib.summary.scalar(name, output)
elif output.shape.ndims == 5:
fps = 10 if self._control_timestep is None else int(
np.round(1.0 / self._control_timestep))
if self._debug_summaries:
_gif_summary(name + '/original',
output[:self._num_images_per_summary],
fps,
step=self.train_step_counter)
_gif_summary(name,
output[:self._num_images_per_summary],
fps,
saturate=True,
step=self.train_step_counter)
else:
raise NotImplementedError
if weights is not None:
model_loss *= weights
model_loss = tf.reduce_mean(input_tensor=model_loss)
if self._debug_summaries:
common.generate_tensor_summaries('model_loss', model_loss,
self.train_step_counter)
return model_loss
