class_docstr: Base class for next_frame models.
  
  This is the base class for the models that given the previous frames
  can predict the next frame. They may also support reward prediction
  and action condition prediction which enables them to be run as
  a world model in model-based RL pipeline.
  
  The API supports both recurrent and stacked frames models. Please look
  at the documents for next_frame function for the API.
  
  If you are implementing a next frame prediction model consider
  following the API presented in this class. But if the API
  is too limiting for your models, feel free to override lower level
  functions and/or inheret from T2TModel directly.
class_docstr_tok: ["Base","class","for","next_frame","models",".","This","is","the","base","class","for","the","models","that","given","the","previous","frames","can","predict","the","next","frame",".","They","may","also","support","reward","prediction","and","action","condition","prediction","which","enables","them","to","be","run","as","a","world","model","in","model-based","RL","pipeline",".","The","API","supports","both","recurrent","and","stacked","frames","models",".","Please","look","at","the","documents","for","next_frame","function","for","the","API",".","If","you","are","implementing","a","next","frame","prediction","model","consider","following","the","API","presented","in","this","class",".","But","if","the","API","is","too","limiting","for","your","models",",","feel","free","to","override","lower","level","functions","and\/or","inheret","from","T2TModel","directly","."]
docstr: Main video processing function.
docstr_tok: ["Main","video","processing","function","."]
code: def __process(
      self, all_frames, all_actions, all_rewards, all_raw_frames
  ):
      """Main video processing function."""
      hparams = self.hparams
      all_frames_copy = [tf.identity(frame) for frame in all_frames]
      orig_frame_shape = common_layers.shape_list(all_frames[0])
      batch_size = orig_frame_shape[0]
      ss_func = self.get_scheduled_sample_func(batch_size)
      target_frames = []
      extra_loss = 0.0
  
      # Any extra info required by the model goes into here.
      video_features = self.video_features(
          all_frames, all_actions, all_rewards, all_raw_frames
      )
  
      num_frames = len(all_frames)
      if self.is_recurrent_model:
          input_index_range = range(num_frames - 1)
      else:
          input_index_range = range(hparams.video_num_target_frames)
  
      # Setup the internal states as well as an auxiliary tf op
      # to enforce syncronization between prediction steps.
      if self.internal_states is None:
          internal_states = None
          sync_op = tf.no_op()
      else:
          internal_states = self.load_internal_states_ops()
          with tf.control_dependencies(flat_lists(internal_states)):
              sync_op = tf.no_op()
  
      (
          res_frames,
          sampled_frames,
          res_rewards,
          res_policies,
          res_values,
      ) = ([], [], [], [], [])
      for i in input_index_range:
          with tf.control_dependencies([sync_op]):
              (
                  frames,
                  actions,
                  rewards,
                  target_index,
              ) = self.__get_next_inputs(
                  i, all_frames, all_actions, all_rewards
              )
              target_frame = all_frames[target_index]
              target_frames.append(tf.identity(target_frame))
  
              with tf.variable_scope(
                  tf.get_variable_scope(), reuse=tf.AUTO_REUSE
              ):
                  float_frames = [
                      tf.to_float(frame) for frame in frames
                  ]
                  func_out = self.next_frame(
                      float_frames,
                      actions,
                      rewards,
                      tf.to_float(target_frame),
                      internal_states,
                      video_features,
                  )
                  (
                      res_frame,
                      res_reward,
                      res_policy,
                      res_value,
                      res_extra_loss,
                      internal_states,
                  ) = func_out
                  res_frames.append(res_frame)
                  res_rewards.append(res_reward)
                  res_policies.append(res_policy)
                  res_values.append(res_value)
                  extra_loss += res_extra_loss / float(
                      len(input_index_range)
                  )
  
                  # Syncronizing the internals states
                  # Some Tensflow Magic to make sure everything happens as it should.
                  with tf.control_dependencies([res_frame]):
                      sync_op = tf.no_op()
                      if (
                          self.is_predicting
                          and self.is_recurrent_model
                          and i == 0
                      ):
                          # The internal state save happens at the end of the 1st iteration
                          # which essentially allows recurrent models to continue
                          # running after one prediction.
                          # Necessary for planning/rl applications.
                          save_ops = self.save_internal_states_ops(
                              internal_states
                          )
                          with tf.control_dependencies(
                              flat_lists(save_ops)
                          ):
                              sync_op = tf.no_op()
  
              # Only for Softmax loss: sample frame so we can keep iterating.
              sampled_frame = self.get_sampled_frame(res_frame)
              sampled_frames.append(sampled_frame)
  
              # Check whether we are done with context frames or not
              if self.is_recurrent_model:
                  done_warm_start = (
                      i >= hparams.video_num_input_frames - 1
                  )
              else:
                  done_warm_start = (
                      True  # Always true for non-reccurent networks.
                  )
  
              if self.is_predicting and done_warm_start:
                  all_frames[target_index] = sampled_frame
  
              # Scheduled sampling during training.
              if self.is_training:
                  groundtruth_items = [tf.to_float(target_frame)]
                  generated_items = [sampled_frame]
                  (ss_frame,) = self.get_scheduled_sample_inputs(
                      done_warm_start,
                      groundtruth_items,
                      generated_items,
                      ss_func,
                  )
                  all_frames[target_index] = ss_frame
  
      video_extra_loss = self.video_extra_loss(
          sampled_frames, target_frames, internal_states, video_features
      )
      tf.summary.scalar("video_extra_loss", video_extra_loss)
      extra_loss += video_extra_loss
  
      if self.is_recurrent_model:
          has_input_predictions = hparams.video_num_input_frames > 1
          if (
              self.is_training
              and hparams.internal_loss
              and has_input_predictions
          ):
              # add the loss for input frames as well.
              extra_gts = all_frames_copy[
                  1 : hparams.video_num_input_frames
              ]
              extra_raw_gts = all_raw_frames[
                  1 : hparams.video_num_input_frames
              ]
              extra_pds = res_frames[
                  : hparams.video_num_input_frames - 1
              ]
              recon_loss = self.get_extra_internal_loss(
                  extra_raw_gts, extra_gts, extra_pds
              )
              extra_loss += recon_loss
          # Cut the predicted input frames.
          res_frames = res_frames[hparams.video_num_input_frames - 1 :]
          res_rewards = res_rewards[
              hparams.video_num_input_frames - 1 :
          ]
          res_policies = res_policies[
              hparams.video_num_input_frames - 1 :
          ]
          res_values = res_values[hparams.video_num_input_frames - 1 :]
          sampled_frames = sampled_frames[
              hparams.video_num_input_frames - 1 :
          ]
          target_frames = target_frames[
              hparams.video_num_input_frames - 1 :
          ]
  
      self.visualize_predictions(
          sampled_frames, [tf.to_float(f) for f in target_frames]
      )
  
      output_frames = tf.stack(res_frames, axis=1)
      targets = output_frames
  
      if any((self.has_rewards, self.has_policies, self.has_values)):
          targets = {"targets": output_frames}
          if self.has_rewards:
              targets["target_reward"] = tf.stack(res_rewards, axis=1)
          if self.has_policies:
              targets["target_policy"] = tf.stack(res_policies, axis=1)
          if self.has_values:
              targets["target_value"] = tf.stack(res_values, axis=1)
  
      return targets, extra_loss
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import functools
  import six
  from tensor2tensor.layers import common_hparams
  from tensor2tensor.layers import common_layers
  from tensor2tensor.layers import common_video
  from tensor2tensor.layers import discretization
  from tensor2tensor.layers import modalities
  from tensor2tensor.utils import registry
  from tensor2tensor.utils import t2t_model
  import tensorflow.compat.v1
label_desc: Use image node_id 20 for calling the NextFrameBase obj's underlying member method code with example usage: obj.__process(all_frames, all_actions, all_rewards, all_raw_frames) and returns: targets,  extra_loss