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
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.
    
tensor2tensor/models/video/savp.py