class_docstr: WorldModel component of [1] w/ encoder, decoder, RSSM, reward/cont. predictors.
  
  See eq. 3 of [1] for all components and their respective in- and outputs.
  Note that in the paper, the "encoder" includes both the raw encoder plus the
  "posterior net", which produces posterior z-states from observations and h-states.
  
  Note: The "internal state" of the world model 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 inserted 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.
class_docstr_tok: ["WorldModel","component","of","[","1","]","w\/","encoder",",","decoder",",","RSSM",",","reward\/cont",".","predictors",".","See","eq",".","3","of","[","1","]","for","all","components","and","their","respective","in-","and","outputs",".","Note","that","in","the","paper",",","the","``","encoder","''","includes","both","the","raw","encoder","plus","the","''","posterior","net","''",",","which","produces","posterior","z-states","from","observations","and","h-states",".","Note",":","The","``","internal","state","''","of","the","world","model","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","inserted","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","."]
docstr: Returns the (current) initial state of the world model (h- and z-states).
  
  An initial state is generated using the tanh of the (learned) h-state variable
  and the dynamics predictor (or "prior net") to compute z^0 from h0. In this last
  step, it is important that we do NOT sample the z^-state (as we would usually
  do during dreaming), but rather take the mode (argmax, then one-hot again).
docstr_tok: ["Returns","the","(","current",")","initial","state","of","the","world","model","(","h-","and","z-states",")",".","An","initial","state","is","generated","using","the","tanh","of","the","(","learned",")","h-state","variable","and","the","dynamics","predictor","(","or","``","prior","net","''",")","to","compute","z^0","from","h0",".","In","this","last","step",",","it","is","important","that","we","do","NOT","sample","the","z^-state","(","as","we","would","usually","do","during","dreaming",")",",","but","rather","take","the","mode","(","argmax",",","then","one-hot","again",")","."]
code: def get_initial_state(self):
      """Returns the (current) initial state of the world model (h- and z-states).
  
      An initial state is generated using the tanh of the (learned) h-state variable
      and the dynamics predictor (or "prior net") to compute z^0 from h0. In this last
      step, it is important that we do NOT sample the z^-state (as we would usually
      do during dreaming), but rather take the mode (argmax, then one-hot again).
      """
      h = tf.expand_dims(
          tf.math.tanh(tf.cast(self.initial_h, self._comp_dtype)), 0
      )
      # Use the mode, NOT a sample for the initial z-state.
      _, z_probs = self.dynamics_predictor(h)
      z = tf.argmax(z_probs, axis=-1)
      z = tf.one_hot(z, depth=z_probs.shape[-1], dtype=self._comp_dtype)
  
      return {"h": h, "z": z}
  
file_path: ray/python/ray/rllib/algorithms/dreamerv3/tf/models/world_model.py
imports: from typing import Optional
  import gymnasium
  import tree
  from ray.rllib.algorithms.dreamerv3.tf.models.components.continue_predictor import ContinuePredictor
  from ray.rllib.algorithms.dreamerv3.tf.models.components.dynamics_predictor import DynamicsPredictor
  from ray.rllib.algorithms.dreamerv3.tf.models.components.mlp import MLP
  from ray.rllib.algorithms.dreamerv3.tf.models.components.representation_layer import RepresentationLayer
  from ray.rllib.algorithms.dreamerv3.tf.models.components.reward_predictor import RewardPredictor
  from ray.rllib.algorithms.dreamerv3.tf.models.components.sequence_model import SequenceModel
  from ray.rllib.algorithms.dreamerv3.utils import get_gru_units
  from ray.rllib.utils.framework import try_import_tf
  from ray.rllib.utils.tf_utils import symlog
label_desc: Use image node_id 2 for calling the WorldModel obj's underlying member method code with example usage: obj.get_initial_state() and returns: dict