class_docstr: The "sequence model" of the RSSM, computing ht+1 given (ht, zt, at).
  
  Note: The "internal state" always consists of:
  The actions `a` (initially, this is a zeroed-out action), `h`-states (deterministic,
  continuous), and `z`-states (stochastic, discrete).
  There are two versions of z-states: "posterior" for world model training and "prior"
  for creating the dream data.
  
  Initial internal state values (`a`, `h`, and `z`) are used where ever a new episode
  starts within a batch row OR at the beginning of each train batch's B rows,
  regardless of whether there was an actual episode boundary or not. Thus, internal
  states are not required to be stored in or retrieved from the replay buffer AND
  retrieved batches from the buffer must not be zero padded.
  
  Initial `a` is the zero "one hot" action, e.g. [0.0, 0.0] for Discrete(2), initial
  `h` is a separate learned variable, and initial `z` are computed by the "dynamics"
  (or "prior") net, using only the initial-h state as input.
  
  The GRU in this SequenceModel always produces the next h-state, then.
docstr: Args:
      a: The previous action (already one-hot'd if applicable). (B, ...).
      h: The previous deterministic hidden state of the sequence model.
          (B, num_gru_units)
      z: The previous stochastic discrete representations of the original
          observation input. (B, num_categoricals, num_classes_per_categorical).
docstr_tok: ["Args",":","a",":","The","previous","action","(","already","one-hot","'d","if","applicable",")",".","(","B",",","...",")",".","h",":","The","previous","deterministic","hidden","state","of","the","sequence","model",".","(","B",",","num_gru_units",")","z",":","The","previous","stochastic","discrete","representations","of","the","original","observation","input",".","(","B",",","num_categoricals",",","num_classes_per_categorical",")","."]
code: def call(self, a, h, z):
      """
  
      Args:
          a: The previous action (already one-hot'd if applicable). (B, ...).
          h: The previous deterministic hidden state of the sequence model.
              (B, num_gru_units)
          z: The previous stochastic discrete representations of the original
              observation input. (B, num_categoricals, num_classes_per_categorical).
      """
      # Flatten last two dims of z.
      z_shape = tf.shape(z)
      z = tf.reshape(z, shape=(z_shape[0], -1))
      out = tf.concat([z, a], axis=-1)
      out.set_shape(
          [
              None,
              (
                  get_num_z_categoricals(self.model_size)
                  * get_num_z_classes(self.model_size)
                  + (
                      self.action_space.n
                      if isinstance(
                          self.action_space, gym.spaces.Discrete
                      )
                      else int(np.prod(self.action_space.shape))
                  )
              ),
          ]
      )
      # Pass through pre-GRU layer.
      out = self.pre_gru_layer(out)
      # Pass through (time-major) GRU.
      h_next = self.gru_unit(
          tf.expand_dims(out, axis=0), initial_state=h
      )
      # Return the GRU's output (the next h-state).
      return h_next
  
file_path: ray/rllib/algorithms/dreamerv3/tf/models/components/sequence_model.py
imports: from typing import Optional
  import gymnasium
  import numpy
  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
label_desc: Use image node_id 2 for calling the SequenceModel obj's underlying member method code with example usage: obj.call(a, h, z) and returns: h_next