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