class_docstr: A Conv2DTranspose decoder to generate Atari images from a latent space.
  
  Wraps an initial single linear layer with a stack of 4 Conv2DTranspose layers (with
  layer normalization) and a diag Gaussian, from which we then sample the final image.
  Sampling is done with a fixed stddev=1.0 and using the mean values coming from the
  last Conv2DTranspose layer.
docstr: Initializes a ConvTransposeAtari instance.
  
  Args:
      model_size: The "Model Size" used according to [1] Appendinx B.
          Use None for manually setting the `cnn_multiplier`.
      cnn_multiplier: Optional override for the additional factor used to multiply
          the number of filters with each CNN transpose layer. Starting with
          8 * `cnn_multiplier` filters in the first CNN transpose layer, the
          number of filters then decreases via `4*cnn_multiplier`,
          `2*cnn_multiplier`, till `1*cnn_multiplier`.
      gray_scaled: Whether the last Conv2DTranspose layer's output has only 1
          color channel (gray_scaled=True) or 3 RGB channels (gray_scaled=False).
docstr_tok: ["Initializes","a","ConvTransposeAtari","instance",".","Args",":","model_size",":","The","``","Model","Size","''","used","according","to","[","1","]","Appendinx","B",".","Use","None","for","manually","setting","the","`","cnn_multiplier","`",".","cnn_multiplier",":","Optional","override","for","the","additional","factor","used","to","multiply","the","number","of","filters","with","each","CNN","transpose","layer",".","Starting","with","8","*","`","cnn_multiplier","`","filters","in","the","first","CNN","transpose","layer",",","the","number","of","filters","then","decreases","via","`","4","*","cnn_multiplier","`",",","`","2","*","cnn_multiplier","`",",","till","`","1","*","cnn_multiplier","`",".","gray_scaled",":","Whether","the","last","Conv2DTranspose","layer","'s","output","has","only","1","color","channel","(","gray_scaled=True",")","or","3","RGB","channels","(","gray_scaled=False",")","."]
code: def __init__(
      self,
      *,
      model_size: Optional[str] = "XS",
      cnn_multiplier: Optional[int] = None,
      gray_scaled: bool,
  ):
      """Initializes a ConvTransposeAtari instance.
  
      Args:
          model_size: The "Model Size" used according to [1] Appendinx B.
              Use None for manually setting the `cnn_multiplier`.
          cnn_multiplier: Optional override for the additional factor used to multiply
              the number of filters with each CNN transpose layer. Starting with
              8 * `cnn_multiplier` filters in the first CNN transpose layer, the
              number of filters then decreases via `4*cnn_multiplier`,
              `2*cnn_multiplier`, till `1*cnn_multiplier`.
          gray_scaled: Whether the last Conv2DTranspose layer's output has only 1
              color channel (gray_scaled=True) or 3 RGB channels (gray_scaled=False).
      """
      super().__init__(name="image_decoder")
  
      self.model_size = model_size
      cnn_multiplier = get_cnn_multiplier(
          self.model_size, override=cnn_multiplier
      )
  
      # The shape going into the first Conv2DTranspose layer.
      # We start with a 4x4 channels=8 "image".
      self.input_dims = (4, 4, 8 * cnn_multiplier)
  
      self.gray_scaled = gray_scaled
  
      # See appendix B in [1]:
      # "The decoder starts with a dense layer, followed by reshaping
      # to 4 × 4 × C and then inverts the encoder architecture. ..."
      self.dense_layer = tf.keras.layers.Dense(
          units=int(np.prod(self.input_dims)),
          activation=None,
          use_bias=True,
      )
      # Inverse conv2d stack. See cnn_atari.py for corresponding Conv2D stack.
      self.conv_transpose_layers = [
          tf.keras.layers.Conv2DTranspose(
              filters=4 * cnn_multiplier,
              kernel_size=4,
              strides=(2, 2),
              padding="same",
              # No bias or activation due to layernorm.
              activation=None,
              use_bias=False,
          ),
          tf.keras.layers.Conv2DTranspose(
              filters=2 * cnn_multiplier,
              kernel_size=4,
              strides=(2, 2),
              padding="same",
              # No bias or activation due to layernorm.
              activation=None,
              use_bias=False,
          ),
          tf.keras.layers.Conv2DTranspose(
              filters=1 * cnn_multiplier,
              kernel_size=4,
              strides=(2, 2),
              padding="same",
              # No bias or activation due to layernorm.
              activation=None,
              use_bias=False,
          ),
      ]
      # Create one LayerNorm layer for each of the Conv2DTranspose layers.
      self.layer_normalizations = []
      for _ in range(len(self.conv_transpose_layers)):
          self.layer_normalizations.append(
              tf.keras.layers.LayerNormalization()
          )
  
      # Important! No activation or layer norm for last layer as the outputs of
      # this one go directly into the diag-gaussian as parameters.
      self.output_conv2d_transpose = tf.keras.layers.Conv2DTranspose(
          filters=1 if self.gray_scaled else 3,
          kernel_size=4,
          strides=(2, 2),
          padding="same",
          activation=None,
          use_bias=True,  # Last layer does use bias (b/c has no LayerNorm).
      )
      # .. until output is 64 x 64 x 3 (or 1 for self.gray_scaled=True).
  
      # 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/rllib/algorithms/dreamerv3/tf/models/components/conv_transpose_atari.py
imports: from typing import Optional
  import numpy
  from ray.rllib.algorithms.dreamerv3.utils import get_cnn_multiplier, get_gru_units, get_num_z_categoricals, get_num_z_classes
  from ray.rllib.utils.framework import try_import_tf
label_desc: Use image node_id 1 to create a new ConvTransposeAtari object from inherited base classes: tf.keras with example: obj = ConvTransposeAtari()