class_docstr: Stochastic Adversarial Video Prediction.
docstr: Model that takes in images and returns predictions.
  
  Args:
    images: list of 4-D Tensors indexed by time.
            (batch_size, width, height, channels)
    actions: list of action tensors
             each action should be in the shape ?x1xZ
    rewards: list of reward tensors
             each reward should be in the shape ?x1xZ
  
  Returns:
    video: list of 4-D predicted frames.
    all_rewards: predicted rewards.
    latent_means: list of gaussian means conditioned on the input at
                  every frame.
    latent_stds: list of gaussian stds conditioned on the input at
                 every frame.
  
  Raises:
    ValueError: If not exactly one of self.hparams.vae or self.hparams.gan
                is set to True.
docstr_tok: ["Model","that","takes","in","images","and","returns","predictions",".","Args",":","images",":","list","of","4-D","Tensors","indexed","by","time",".","(","batch_size",",","width",",","height",",","channels",")","actions",":","list","of","action","tensors","each","action","should","be","in","the","shape","?","x1xZ","rewards",":","list","of","reward","tensors","each","reward","should","be","in","the","shape","?","x1xZ","Returns",":","video",":","list","of","4-D","predicted","frames",".","all_rewards",":","predicted","rewards",".","latent_means",":","list","of","gaussian","means","conditioned","on","the","input","at","every","frame",".","latent_stds",":","list","of","gaussian","stds","conditioned","on","the","input","at","every","frame",".","Raises",":","ValueError",":","If","not","exactly","one","of","self.hparams.vae","or","self.hparams.gan","is","set","to","True","."]
code: def construct_model(self, images, actions, rewards):
      """Model that takes in images and returns predictions.
  
      Args:
        images: list of 4-D Tensors indexed by time.
                (batch_size, width, height, channels)
        actions: list of action tensors
                 each action should be in the shape ?x1xZ
        rewards: list of reward tensors
                 each reward should be in the shape ?x1xZ
  
      Returns:
        video: list of 4-D predicted frames.
        all_rewards: predicted rewards.
        latent_means: list of gaussian means conditioned on the input at
                      every frame.
        latent_stds: list of gaussian stds conditioned on the input at
                     every frame.
  
      Raises:
        ValueError: If not exactly one of self.hparams.vae or self.hparams.gan
                    is set to True.
      """
      if not self.hparams.use_vae and not self.hparams.use_gan:
          raise ValueError(
              "Set at least one of use_vae or use_gan to be True"
          )
      if self.hparams.gan_optimization not in ["joint", "sequential"]:
          raise ValueError(
              "self.hparams.gan_optimization should be either joint "
              "or sequential got %s" % self.hparams.gan_optimization
          )
  
      images = tf.unstack(images, axis=0)
      actions = tf.unstack(actions, axis=0)
      rewards = tf.unstack(rewards, axis=0)
  
      latent_dims = self.hparams.z_dim
      context_frames = self.hparams.video_num_input_frames
      seq_len = len(images)
      input_shape = common_layers.shape_list(images[0])
      batch_size = input_shape[0]
  
      # Model does not support reward-conditioned frame generation.
      fake_rewards = rewards[:-1]
  
      # Concatenate x_{t-1} and x_{t} along depth and encode it to
      # produce the mean and standard deviation of z_{t-1}
      image_pairs = tf.concat(
          [images[: seq_len - 1], images[1:seq_len]], axis=-1
      )
  
      z_mu, z_log_sigma_sq = self.encoder(image_pairs)
      # Unstack z_mu and z_log_sigma_sq along the time dimension.
      z_mu = tf.unstack(z_mu, axis=0)
      z_log_sigma_sq = tf.unstack(z_log_sigma_sq, axis=0)
      iterable = zip(
          images[:-1], actions[:-1], fake_rewards, z_mu, z_log_sigma_sq
      )
  
      # Initialize LSTM State
      lstm_state = [None] * 7
      (
          gen_cond_video,
          gen_prior_video,
          all_rewards,
          latent_means,
          latent_stds,
      ) = ([], [], [], [], [])
      pred_image = tf.zeros_like(images[0])
      prior_latent_state, cond_latent_state = None, None
      train_mode = self.hparams.mode == tf_estimator.ModeKeys.TRAIN
  
      # Create scheduled sampling function
      ss_func = self.get_scheduled_sample_func(batch_size)
  
      with tf.variable_scope("prediction", reuse=tf.AUTO_REUSE):
          for step, (
              image,
              action,
              reward,
              mu,
              log_sigma_sq,
          ) in enumerate(
              iterable
          ):  # pylint:disable=line-too-long
              # Sample latents using a gaussian centered at conditional mu and std.
              latent = common_video.get_gaussian_tensor(
                  mu, log_sigma_sq
              )
  
              # Sample prior latents from isotropic normal distribution.
              prior_latent = tf.random_normal(
                  tf.shape(latent), dtype=tf.float32
              )
  
              # LSTM that encodes correlations between conditional latents.
              # Pg 22 in https://arxiv.org/pdf/1804.01523.pdf
              (
                  enc_cond_latent,
                  cond_latent_state,
              ) = common_video.basic_lstm(
                  latent,
                  cond_latent_state,
                  latent_dims,
                  name="cond_latent",
              )
  
              # LSTM that encodes correlations between prior latents.
              (
                  enc_prior_latent,
                  prior_latent_state,
              ) = common_video.basic_lstm(
                  prior_latent,
                  prior_latent_state,
                  latent_dims,
                  name="prior_latent",
              )
  
              # Scheduled Sampling
              done_warm_start = step > context_frames - 1
              groundtruth_items = [image]
              generated_items = [pred_image]
              (input_image,) = self.get_scheduled_sample_inputs(
                  done_warm_start,
                  groundtruth_items,
                  generated_items,
                  ss_func,
              )
  
              all_latents = tf.concat(
                  [enc_cond_latent, enc_prior_latent], axis=0
              )
              all_image = tf.concat([input_image, input_image], axis=0)
              all_action = tf.concat([action, action], axis=0)
              all_rewards = tf.concat([reward, reward], axis=0)
  
              (
                  all_pred_images,
                  lstm_state,
                  _,
              ) = self.construct_predictive_tower(
                  all_image,
                  all_rewards,
                  all_action,
                  lstm_state,
                  all_latents,
                  concat_latent=True,
              )
  
              cond_pred_images, prior_pred_images = (
                  all_pred_images[:batch_size],
                  all_pred_images[batch_size:],
              )
  
              if train_mode and self.hparams.use_vae:
                  pred_image = cond_pred_images
              else:
                  pred_image = prior_pred_images
  
              gen_cond_video.append(cond_pred_images)
              gen_prior_video.append(prior_pred_images)
              latent_means.append(mu)
              latent_stds.append(log_sigma_sq)
  
      gen_cond_video = tf.stack(gen_cond_video, axis=0)
      self.gen_prior_video = tf.stack(gen_prior_video, axis=0)
      fake_rewards = tf.stack(fake_rewards, axis=0)
  
      if train_mode and self.hparams.use_vae:
          return gen_cond_video, fake_rewards, latent_means, latent_stds
      else:
          return (
              self.gen_prior_video,
              fake_rewards,
              latent_means,
              latent_stds,
          )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import numbers
  import numpy
  from tensor2tensor.layers import common_layers
  from tensor2tensor.layers import common_video
  from tensor2tensor.models.video import savp_params
  from tensor2tensor.models.video import sv2p
  from tensor2tensor.utils import contrib
  from tensor2tensor.utils import registry
  from tensor2tensor.utils import update_ops_hook
  import tensorflow.compat.v1
  from tensorflow.compat.v1 import estimator
  import tensorflow_gan
label_desc: Use image node_id 1 for calling the NextFrameSAVP obj's underlying member method code with example usage: obj.construct_model(images, actions, rewards) and returns: gen_cond_video,  fake_rewards,  latent_means,  latent_stds, self,  fake_rewards,  latent_means,  latent_stds