class_docstr: Stochastic next-frame model with 2 frames posterior.
class_docstr_tok: ["Stochastic","next-frame","model","with","2","frames","posterior","."]
code: def construct_model(self, images, actions, rewards):
      images = tf.unstack(images, axis=0)
      actions = tf.unstack(actions, axis=0)
      rewards = tf.unstack(rewards, axis=0)
  
      batch_size = common_layers.shape_list(images[0])[0]
      context_frames = self.hparams.video_num_input_frames
  
      # Predicted images and rewards.
      gen_rewards, gen_images, latent_means, latent_stds = (
          [],
          [],
          [],
          [],
      )
  
      # LSTM states.
      lstm_state = [None] * 7
  
      # Create scheduled sampling function
      ss_func = self.get_scheduled_sample_func(batch_size)
  
      pred_image = tf.zeros_like(images[0])
      pred_reward = tf.zeros_like(rewards[0])
      latent = None
      for timestep, image, action, reward in zip(
          range(len(images) - 1),
          images[:-1],
          actions[:-1],
          rewards[:-1],
      ):
          # Scheduled Sampling
          done_warm_start = timestep > context_frames - 1
          groundtruth_items = [image, reward]
          generated_items = [pred_image, pred_reward]
          input_image, input_reward = self.get_scheduled_sample_inputs(
              done_warm_start,
              groundtruth_items,
              generated_items,
              ss_func,
          )
  
          # Latent
          # TODO(mbz): should we use input_image iunstead of image?
          latent_images = tf.stack(
              [image, images[timestep + 1]], axis=0
          )
          latent_mean, latent_std = self.construct_latent_tower(
              latent_images, time_axis=0
          )
          latent = common_video.get_gaussian_tensor(
              latent_mean, latent_std
          )
          latent_means.append(latent_mean)
          latent_stds.append(latent_std)
  
          # Prediction
          pred_image, lstm_state, _ = self.construct_predictive_tower(
              input_image, input_reward, action, lstm_state, latent
          )
  
          if self.hparams.reward_prediction:
              pred_reward = self.reward_prediction(
                  pred_image, input_reward, action, latent
              )
              pred_reward = common_video.decode_to_shape(
                  pred_reward,
                  common_layers.shape_list(input_reward),
                  "reward_dec",
              )
          else:
              pred_reward = input_reward
  
          gen_images.append(pred_image)
          gen_rewards.append(pred_reward)
  
      gen_images = tf.stack(gen_images, axis=0)
      gen_rewards = tf.stack(gen_rewards, axis=0)
  
      return gen_images, gen_rewards, latent_means, latent_stds
  
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.layers import discretization
  from tensor2tensor.models.video import base
  from tensor2tensor.models.video import base_vae
  from tensor2tensor.utils import contrib
  from tensor2tensor.utils import registry
  import tensorflow.compat.v1
label_desc: Use image node_id 1 for calling the NextFrameSv2pTwoFrames obj's underlying member method code with example usage: obj.construct_model(images, actions, rewards) and returns: gen_images,  gen_rewards,  latent_means,  latent_stds