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: Creates a function for scheduled sampling based on given hparams.
docstr_tok: ["Creates","a","function","for","scheduled","sampling","based","on","given","hparams","."]
code: def get_scheduled_sample_func(self, batch_size):
      """Creates a function for scheduled sampling based on given hparams."""
      with tf.variable_scope(
          "scheduled_sampling_func", reuse=tf.AUTO_REUSE
      ):
          iter_num = self.get_iteration_num()
  
          # Simple function to bypass scheduled sampling in gt or pred only modes.
          def scheduled_sampling_simple(
              ground_truth_x,
              generated_x,
              batch_size,
              scheduled_sample_var,
          ):
              del batch_size
              if scheduled_sample_var:
                  return ground_truth_x
              return generated_x
  
          mode = self.hparams.scheduled_sampling_mode
          if mode == "ground_truth_only":
              scheduled_sampling_func = scheduled_sampling_simple
              scheduled_sampling_func_var = True
          elif mode == "prediction_only":
              scheduled_sampling_func = scheduled_sampling_simple
              scheduled_sampling_func_var = False
          elif mode == "prob":
              decay_steps = self.hparams.scheduled_sampling_decay_steps
              probability = tf.train.polynomial_decay(
                  1.0, iter_num, decay_steps, 0.0
              )
              scheduled_sampling_func = (
                  common_video.scheduled_sample_prob
              )
              scheduled_sampling_func_var = probability
          elif mode == "prob_inverse_exp":
              decay_steps = self.hparams.scheduled_sampling_decay_steps
              probability = common_layers.inverse_exp_decay(
                  decay_steps, step=iter_num
              )
              probability *= self.hparams.scheduled_sampling_max_prob
              probability = 1.0 - probability
              scheduled_sampling_func = (
                  common_video.scheduled_sample_prob
              )
              scheduled_sampling_func_var = probability
          elif mode == "prob_inverse_lin":
              decay_steps = self.hparams.scheduled_sampling_decay_steps
              probability = common_layers.inverse_exp_decay(
                  decay_steps // 4, step=iter_num
              )  # Very low at start.
              probability *= common_layers.inverse_lin_decay(
                  decay_steps, step=iter_num
              )
              probability *= self.hparams.scheduled_sampling_max_prob
              probability = 1.0 - probability
              scheduled_sampling_func = (
                  common_video.scheduled_sample_prob
              )
              scheduled_sampling_func_var = probability
          elif mode == "count":
              # Calculate number of ground-truth frames to pass in.
              k = self.hparams.scheduled_sampling_k
              num_ground_truth = tf.to_int32(
                  tf.round(
                      tf.to_float(batch_size)
                      * (
                          k
                          / (
                              k
                              + tf.exp(
                                  tf.to_float(iter_num) / tf.to_float(k)
                              )
                          )
                      )
                  )
              )
              scheduled_sampling_func = (
                  common_video.scheduled_sample_count
              )
              scheduled_sampling_func_var = num_ground_truth
          else:
              raise ValueError(
                  "unknown scheduled sampling method: %s" % mode
              )
  
          if isinstance(scheduled_sampling_func_var, tf.Tensor):
              tf.summary.scalar(
                  "scheduled_sampling_var", scheduled_sampling_func_var
              )
          partial_func = functools.partial(
              scheduled_sampling_func,
              batch_size=batch_size,
              scheduled_sample_var=scheduled_sampling_func_var,
          )
          return partial_func
  
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 14 for calling the NextFrameBase obj's underlying member method code with example usage: obj.get_scheduled_sample_func(batch_size) and returns: partial_func, generated_x, ground_truth_x