class_docstr: A representation (z-state) generating layer.
  
  The value for z is the result of sampling from a categorical distribution with
  shape B x `num_classes`. So a computed z-state consists of `num_categoricals`
  one-hot vectors, each of size `num_classes_per_categorical`.
docstr: Initializes a RepresentationLayer instance.
  
  Args:
      model_size: The "Model Size" used according to [1] Appendinx B.
          Use None for manually setting the different parameters.
      num_categoricals: Overrides the number of categoricals used in the z-states.
          In [1], 32 is used for any model size.
      num_classes_per_categorical: Overrides the number of classes within each
          categorical used for the z-states. In [1], 32 is used for any model
          dimension.
docstr_tok: ["Initializes","a","RepresentationLayer","instance",".","Args",":","model_size",":","The","``","Model","Size","''","used","according","to","[","1","]","Appendinx","B",".","Use","None","for","manually","setting","the","different","parameters",".","num_categoricals",":","Overrides","the","number","of","categoricals","used","in","the","z-states",".","In","[","1","]",",","32","is","used","for","any","model","size",".","num_classes_per_categorical",":","Overrides","the","number","of","classes","within","each","categorical","used","for","the","z-states",".","In","[","1","]",",","32","is","used","for","any","model","dimension","."]
code: def __init__(
      self,
      *,
      model_size: Optional[str] = "XS",
      num_categoricals: Optional[int] = None,
      num_classes_per_categorical: Optional[int] = None,
  ):
      """Initializes a RepresentationLayer instance.
  
      Args:
          model_size: The "Model Size" used according to [1] Appendinx B.
              Use None for manually setting the different parameters.
          num_categoricals: Overrides the number of categoricals used in the z-states.
              In [1], 32 is used for any model size.
          num_classes_per_categorical: Overrides the number of classes within each
              categorical used for the z-states. In [1], 32 is used for any model
              dimension.
      """
      self.num_categoricals = get_num_z_categoricals(
          model_size, override=num_categoricals
      )
      self.num_classes_per_categorical = get_num_z_classes(
          model_size, override=num_classes_per_categorical
      )
  
      super().__init__(
          name=f"z{self.num_categoricals}x{self.num_classes_per_categorical}"
      )
  
      self.z_generating_layer = tf.keras.layers.Dense(
          self.num_categoricals * self.num_classes_per_categorical,
          activation=None,
      )
  
imports: from typing import Optional
  from ray.rllib.algorithms.dreamerv3.utils import get_num_z_categoricals, get_num_z_classes
  from ray.rllib.utils.framework import try_import_tf, try_import_tfp
label_desc: Use image node_id 1 to create a new RepresentationLayer object from inherited base classes: tf.keras.layers with example: obj = RepresentationLayer()