class_docstr: Stochastic Variational Video Prediction Without Learned Prior.
class_docstr_tok: ["Stochastic","Variational","Video","Prediction","Without","Learned","Prior","."]
docstr: Builds the stochastic model.
  
  The model first encodes all the images (x_t) in the sequence
  using the encoder. Let"s call the output e_t. Then it predicts the
  latent state of the next frame using a recurrent posterior network
  z ~ q(z|e_{0:t}) = N(mu(e_{0:t}), sigma(e_{0:t})).
  Another recurrent network predicts the embedding of the next frame
  using the approximated posterior e_{t+1} = p(e_{t+1}|e_{0:t}, z)
  Finally, the decoder decodes e_{t+1} into x_{t+1}.
  Skip connections from encoder to decoder help with reconstruction.
  
  Args:
    images: tensor of ground truth image sequences
    actions: list of action tensors
    rewards: NOT used list of reward tensors
  
  Returns:
    gen_images: generated images
    fakr_rewards: input rewards as reward prediction!
    pred_mu: predited means of posterior
    pred_logvar: predicted log(var) of posterior
docstr_tok: ["Builds","the","stochastic","model",".","The","model","first","encodes","all","the","images","(","x_t",")","in","the","sequence","using","the","encoder",".","Let","''","s","call","the","output","e_t",".","Then","it","predicts","the","latent","state","of","the","next","frame","using","a","recurrent","posterior","network","z","~","q","(","z|e_","{","0",":","t","}",")","=","N","(","mu","(","e_","{","0",":","t","}",")",",","sigma","(","e_","{","0",":","t","}",")",")",".","Another","recurrent","network","predicts","the","embedding","of","the","next","frame","using","the","approximated","posterior","e_","{","t+1","}","=","p","(","e_","{","t+1","}","|e_","{","0",":","t","}",",","z",")","Finally",",","the","decoder","decodes","e_","{","t+1","}","into","x_","{","t+1","}",".","Skip","connections","from","encoder","to","decoder","help","with","reconstruction",".","Args",":","images",":","tensor","of","ground","truth","image","sequences","actions",":","list","of","action","tensors","rewards",":","NOT","used","list","of","reward","tensors","Returns",":","gen_images",":","generated","images","fakr_rewards",":","input","rewards","as","reward","prediction","!","pred_mu",":","predited","means","of","posterior","pred_logvar",":","predicted","log","(","var",")","of","posterior"]
code: def construct_model(self, images, actions, rewards):
      """Builds the stochastic model.
  
      The model first encodes all the images (x_t) in the sequence
      using the encoder. Let"s call the output e_t. Then it predicts the
      latent state of the next frame using a recurrent posterior network
      z ~ q(z|e_{0:t}) = N(mu(e_{0:t}), sigma(e_{0:t})).
      Another recurrent network predicts the embedding of the next frame
      using the approximated posterior e_{t+1} = p(e_{t+1}|e_{0:t}, z)
      Finally, the decoder decodes e_{t+1} into x_{t+1}.
      Skip connections from encoder to decoder help with reconstruction.
  
      Args:
        images: tensor of ground truth image sequences
        actions: list of action tensors
        rewards: NOT used list of reward tensors
  
      Returns:
        gen_images: generated images
        fakr_rewards: input rewards as reward prediction!
        pred_mu: predited means of posterior
        pred_logvar: predicted log(var) of posterior
      """
      # model does not support action conditioned and reward prediction
      fake_reward_prediction = rewards
      del rewards
      action_repeat = self.hparams.action_repeat
      action_type = self.hparams.action_type
  
      assert action_type in [
          "",
          "image",
          "vector",
      ], "Invalid action type."
      if not action_type:
          a_dim = 0
      elif action_type == "image":
          a_dim = self.hparams.g_dim
      else:
          assert (
              action_repeat > 0
          ), "Action repeat has to be positive integer."
          actions = tf.tile(actions, (1, 1, action_repeat))
          a_dim = actions.shape[-1]
  
      z_dim = self.hparams.z_dim
      g_dim = self.hparams.g_dim
      rnn_size = self.hparams.rnn_size
      prior_rnn_layers = self.hparams.prior_rnn_layers
      posterior_rnn_layers = self.hparams.posterior_rnn_layers
      predictor_rnn_layers = self.hparams.predictor_rnn_layers
      context_frames = self.hparams.video_num_input_frames
      has_batchnorm = self.hparams.has_batchnorm
  
      (
          seq_len,
          batch_size,
          _,
          _,
          color_channels,
      ) = common_layers.shape_list(images)
  
      # LSTM initial sizesstates.
      prior_states = [None] * prior_rnn_layers
      posterior_states = [None] * posterior_rnn_layers
      predictor_states = [None] * predictor_rnn_layers
  
      tf.logging.info(">>>> Encoding")
      # Encoding:
      enc_images, enc_skips = [], []
      enc_actions = []
      images = tf.unstack(images, axis=0)
      actions = tf.unstack(actions, axis=0)
      for i, image in enumerate(images):
          with tf.variable_scope("encoder", reuse=tf.AUTO_REUSE):
              enc, skips = self.encoder(
                  image, g_dim, has_batchnorm=has_batchnorm
              )
              enc = tfl.flatten(enc)
              enc_images.append(enc)
              enc_skips.append(skips)
              if action_type == "image":
                  enc_action, _ = self.encoder(
                      actions[i], g_dim, has_batchnorm=has_batchnorm
                  )
                  enc_action = tfl.flatten(enc_action)
                  enc_actions.append(enc_action)
  
      tf.logging.info(">>>> Prediction")
      # Prediction
      pred_mu_pos = []
      pred_logvar_pos = []
      pred_mu_prior = []
      pred_logvar_prior = []
      gen_images = []
      for i in range(1, seq_len):
          with tf.variable_scope("encoder", reuse=tf.AUTO_REUSE):
              # current encoding
              if self.is_training or len(gen_images) < context_frames:
                  h_current = enc_images[i - 1]
              else:
                  h_current, _ = self.encoder(gen_images[-1], g_dim)
                  h_current = tfl.flatten(h_current)
  
              # target encoding
              h_target = enc_images[i]
  
              if action_type == "image":
                  h_current = tf.concat(
                      [h_current, enc_actions[i - 1]], axis=1
                  )
                  h_target = tf.concat(
                      [h_target, enc_actions[i]], axis=1
                  )
              elif action_type == "vector":
                  h_current = tf.concat(
                      [h_current, actions[i - 1]], axis=1
                  )
                  h_target = tf.concat([h_target, actions[i]], axis=1)
  
          with tf.variable_scope("prediction", reuse=tf.AUTO_REUSE):
              # Prior parameters
              if self.hparams.learned_prior:
                  (
                      mu_prior,
                      logvar_prior,
                      prior_states,
                  ) = self.lstm_gaussian(
                      h_current,
                      prior_states,
                      rnn_size,
                      z_dim,
                      prior_rnn_layers,
                      "prior",
                  )
              else:
                  mu_prior = tf.zeros((batch_size, z_dim))
                  logvar_prior = tf.zeros((batch_size, z_dim))
  
              # Only use Posterior if it's training time
              if self.hparams.stochastic_model and (
                  self.is_training or len(gen_images) < context_frames
              ):
                  (
                      mu_pos,
                      logvar_pos,
                      posterior_states,
                  ) = self.lstm_gaussian(
                      h_target,
                      posterior_states,
                      rnn_size,
                      z_dim,
                      posterior_rnn_layers,
                      "posterior",
                  )
                  # Sample z from posterior distribution
                  z = common_video.get_gaussian_tensor(
                      mu_pos, logvar_pos
                  )
              else:
                  mu_pos = tf.zeros_like(mu_prior)
                  logvar_pos = tf.zeros_like(logvar_prior)
                  z = common_video.get_gaussian_tensor(
                      mu_prior, logvar_prior
                  )
  
              # Predict output encoding
              h_pred, predictor_states = self.stacked_lstm(
                  tf.concat([h_current, z], axis=1),
                  predictor_states,
                  rnn_size,
                  g_dim,
                  predictor_rnn_layers,
              )
  
              pred_mu_pos.append(tf.identity(mu_pos, "mu_pos"))
              pred_logvar_pos.append(
                  tf.identity(logvar_pos, "logvar_pos")
              )
              pred_mu_prior.append(tf.identity(mu_prior, "mu_prior"))
              pred_logvar_prior.append(
                  tf.identity(logvar_prior, "logvar_prior")
              )
  
          with tf.variable_scope("decoding", reuse=tf.AUTO_REUSE):
              skip_index = min(context_frames - 1, i - 1)
              if action_type == "vector":
                  h_pred = tf.concat([h_pred, actions[i - 1]], axis=-1)
              elif action_type == "image":
                  h_pred = tf.concat(
                      [h_pred, enc_actions[i - 1]], axis=-1
                  )
              h_pred = tf.reshape(
                  h_pred, [batch_size, 1, 1, g_dim + a_dim]
              )
              if self.hparams.has_skips:
                  x_pred = self.decoder(
                      h_pred,
                      color_channels,
                      skips=enc_skips[skip_index],
                      has_batchnorm=has_batchnorm,
                  )
              else:
                  x_pred = self.decoder(
                      h_pred,
                      color_channels,
                      has_batchnorm=has_batchnorm,
                  )
              gen_images.append(x_pred)
  
      tf.logging.info(">>>> Done")
      gen_images = tf.stack(gen_images, axis=0)
      return {
          "gen_images": gen_images,
          "fake_reward_prediction": fake_reward_prediction,
          "pred_mu_pos": pred_mu_pos,
          "pred_logvar_pos": pred_logvar_pos,
          "pred_mu_prior": pred_mu_prior,
          "pred_logvar_prior": pred_logvar_prior,
      }
  
file_path: tensor2tensor/tensor2tensor/models/video/emily.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from tensor2tensor.layers import common_layers
  from tensor2tensor.layers import common_video
  from tensor2tensor.models.video import sv2p
  from tensor2tensor.models.video import sv2p_params
  from tensor2tensor.utils import contrib
  from tensor2tensor.utils import registry
  import tensorflow.compat.v1
label_desc: Use image node_id 5 for calling the NextFrameEmily obj's underlying member method code with example usage: obj.construct_model(images, actions, rewards) and returns: dict