class_docstr: The world-model network sub-component used to predict the `continue` flags .
  
  Predicted continue flags are used to produce "dream data" to learn the policy in.
  
  The continue flags are predicted via a linear output used to parameterize a
  Bernoulli distribution, from which simply the mode is used (no stochastic
  sampling!). In other words, if the sigmoid of the output of the linear layer is
  >0.5, we predict a continuation of the episode, otherwise we predict an episode
  terminal.
docstr: Initializes a ContinuePredictor instance.
  
  Args:
      model_size: The "Model Size" used according to [1] Appendinx B.
          Determines the exact size of the underlying MLP.
docstr_tok: ["Initializes","a","ContinuePredictor","instance",".","Args",":","model_size",":","The","``","Model","Size","''","used","according","to","[","1","]","Appendinx","B.","Determines","the","exact","size","of","the","underlying","MLP","."]
code: def __init__(self, *, model_size: Optional[str] = "XS"):
      """Initializes a ContinuePredictor instance.
  
      Args:
          model_size: The "Model Size" used according to [1] Appendinx B.
              Determines the exact size of the underlying MLP.
      """
      super().__init__(name="continue_predictor")
      self.model_size = model_size
      self.mlp = MLP(model_size=model_size, output_layer_size=1)
  
      # Trace self.call.
      dl_type = (
          tf.keras.mixed_precision.global_policy().compute_dtype
          or tf.float32
      )
      self.call = tf.function(
          input_signature=[
              tf.TensorSpec(
                  shape=[None, get_gru_units(model_size)], dtype=dl_type
              ),
              tf.TensorSpec(
                  shape=[
                      None,
                      get_num_z_categoricals(model_size),
                      get_num_z_classes(model_size),
                  ],
                  dtype=dl_type,
              ),
          ]
      )(self.call)
  
file_path: ray/python/ray/rllib/algorithms/dreamerv3/tf/models/components/continue_predictor.py
imports: from typing import Optional
  from ray.rllib.algorithms.dreamerv3.tf.models.components.mlp import MLP
  from ray.rllib.algorithms.dreamerv3.utils import get_gru_units, get_num_z_classes, get_num_z_categoricals
  from ray.rllib.utils.framework import try_import_tf, try_import_tfp
label_desc: Use image node_id 1 to create a new ContinuePredictor object from inherited base classes: tf.keras with example: obj = ContinuePredictor()