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: Initializes a WorldModel instance.
  
  Args:
       model_size: The "Model Size" used according to [1] Appendinx B.
          Use None for manually setting the different network sizes.
       observation_space: The observation space of the environment used.
       action_space: The action space of the environment used.
       batch_length_T: The length (T) of the sequences used for training. The
          actual shape of the input data (e.g. rewards) is then: [B, T, ...],
          where B is the "batch size", T is the "batch length" (this arg) and
          "..." is the dimension of the data (e.g. (64, 64, 3) for Atari image
          observations). Note that a single row (within a batch) may contain data
          from different episodes, but an already on-going episode is always
          finished, before a new one starts within the same row.
      encoder: The encoder Model taking observations as inputs and
          outputting a 1D latent vector that will be used as input into the
          posterior net (z-posterior state generating layer). Inputs are symlogged
          if inputs are NOT images. For images, we use normalization between -1.0
          and 1.0 (x / 128 - 1.0)
      decoder: The decoder Model taking h- and z-states as inputs and generating
          a (possibly symlogged) predicted observation. Note that for images,
          the last decoder layer produces the exact, normalized pixel values
          (not a Gaussian as described in [1]!).
      num_gru_units: The number of GRU units to use. If None, use
          `model_size` to figure out this parameter.
      symlog_obs: Whether to predict decoded observations in symlog space.
          This should be False for image based observations.
          According to the paper [1] Appendix E: "NoObsSymlog: This ablation
          removes the symlog encoding of inputs to the world model and also
          changes the symlog MSE loss in the decoder to a simple MSE loss.
          *Because symlog encoding is only used for vector observations*, this
          ablation is equivalent to DreamerV3 on purely image-based environments".
docstr_tok: ["Initializes","a","WorldModel","instance",".","Args",":","model_size",":","The","``","Model","Size","''","used","according","to","[","1","]","Appendinx","B",".","Use","None","for","manually","setting","the","different","network","sizes",".","observation_space",":","The","observation","space","of","the","environment","used",".","action_space",":","The","action","space","of","the","environment","used",".","batch_length_T",":","The","length","(","T",")","of","the","sequences","used","for","training",".","The","actual","shape","of","the","input","data","(","e.g",".","rewards",")","is","then",":","[","B",",","T",",","...","]",",","where","B","is","the","``","batch","size","''",",","T","is","the","``","batch","length","''","(","this","arg",")","and","``","...","''","is","the","dimension","of","the","data","(","e.g",".","(","64",",","64",",","3",")","for","Atari","image","observations",")",".","Note","that","a","single","row","(","within","a","batch",")","may","contain","data","from","different","episodes",",","but","an","already","on-going","episode","is","always","finished",",","before","a","new","one","starts","within","the","same","row",".","encoder",":","The","encoder","Model","taking","observations","as","inputs","and","outputting","a","1D","latent","vector","that","will","be","used","as","input","into","the","posterior","net","(","z-posterior","state","generating","layer",")",".","Inputs","are","symlogged","if","inputs","are","NOT","images",".","For","images",",","we","use","normalization","between","-1.0","and","1.0","(","x","\/","128","-","1.0",")","decoder",":","The","decoder","Model","taking","h-","and","z-states","as","inputs","and","generating","a","(","possibly","symlogged",")","predicted","observation",".","Note","that","for","images",",","the","last","decoder","layer","produces","the","exact",",","normalized","pixel","values","(","not","a","Gaussian","as","described","in","[","1","]","!",")",".","num_gru_units",":","The","number","of","GRU","units","to","use",".","If","None",",","use","`","model_size","`","to","figure","out","this","parameter",".","symlog_obs",":","Whether","to","predict","decoded","observations","in","symlog","space",".","This","should","be","False","for","image","based","observations",".","According","to","the","paper","[","1","]","Appendix","E",":","``","NoObsSymlog",":","This","ablation","removes","the","symlog","encoding","of","inputs","to","the","world","model","and","also","changes","the","symlog","MSE","loss","in","the","decoder","to","a","simple","MSE","loss",".","*","Because","symlog","encoding","is","only","used","for","vector","observations","*",",","this","ablation","is","equivalent","to","DreamerV3","on","purely","image-based","environments","''","."]
code: def __init__(
      self,
      *,
      model_size: str = "XS",
      observation_space: gym.Space,
      action_space: gym.Space,
      batch_length_T: int = 64,
      encoder: tf.keras.Model,
      decoder: tf.keras.Model,
      num_gru_units: Optional[int] = None,
      symlog_obs: bool = True,
  ):
      """Initializes a WorldModel instance.
  
      Args:
           model_size: The "Model Size" used according to [1] Appendinx B.
              Use None for manually setting the different network sizes.
           observation_space: The observation space of the environment used.
           action_space: The action space of the environment used.
           batch_length_T: The length (T) of the sequences used for training. The
              actual shape of the input data (e.g. rewards) is then: [B, T, ...],
              where B is the "batch size", T is the "batch length" (this arg) and
              "..." is the dimension of the data (e.g. (64, 64, 3) for Atari image
              observations). Note that a single row (within a batch) may contain data
              from different episodes, but an already on-going episode is always
              finished, before a new one starts within the same row.
          encoder: The encoder Model taking observations as inputs and
              outputting a 1D latent vector that will be used as input into the
              posterior net (z-posterior state generating layer). Inputs are symlogged
              if inputs are NOT images. For images, we use normalization between -1.0
              and 1.0 (x / 128 - 1.0)
          decoder: The decoder Model taking h- and z-states as inputs and generating
              a (possibly symlogged) predicted observation. Note that for images,
              the last decoder layer produces the exact, normalized pixel values
              (not a Gaussian as described in [1]!).
          num_gru_units: The number of GRU units to use. If None, use
              `model_size` to figure out this parameter.
          symlog_obs: Whether to predict decoded observations in symlog space.
              This should be False for image based observations.
              According to the paper [1] Appendix E: "NoObsSymlog: This ablation
              removes the symlog encoding of inputs to the world model and also
              changes the symlog MSE loss in the decoder to a simple MSE loss.
              *Because symlog encoding is only used for vector observations*, this
              ablation is equivalent to DreamerV3 on purely image-based environments".
      """
      super().__init__(name="world_model")
  
      self.model_size = model_size
      self.batch_length_T = batch_length_T
      self.symlog_obs = symlog_obs
      self.observation_space = observation_space
      self.action_space = action_space
      self._comp_dtype = (
          tf.keras.mixed_precision.global_policy().compute_dtype
          or tf.float32
      )
  
      # Encoder (latent 1D vector generator) (xt -> lt).
      self.encoder = encoder
  
      # Posterior predictor consisting of an MLP and a RepresentationLayer:
      # [ht, lt] -> zt.
      self.posterior_mlp = MLP(
          model_size=self.model_size,
          output_layer_size=None,
          # In Danijar's code, the posterior predictor only has a single layer,
          # no matter the model size:
          num_dense_layers=1,
          name="posterior_mlp",
      )
      # The (posterior) z-state generating layer.
      self.posterior_representation_layer = RepresentationLayer(
          model_size=self.model_size,
      )
  
      # Dynamics (prior z-state) predictor: ht -> z^t
      self.dynamics_predictor = DynamicsPredictor(
          model_size=self.model_size
      )
  
      # GRU for the RSSM: [at, ht, zt] -> ht+1
      self.num_gru_units = get_gru_units(
          model_size=self.model_size,
          override=num_gru_units,
      )
      # Initial h-state variable (learnt).
      # -> tanh(self.initial_h) -> deterministic state
      # Use our Dynamics predictor for initial stochastic state, BUT with greedy
      # (mode) instead of sampling.
      self.initial_h = tf.Variable(
          tf.zeros(shape=(self.num_gru_units,)),
          trainable=True,
          name="initial_h",
      )
      # The actual sequence model containing the GRU layer.
      self.sequence_model = SequenceModel(
          model_size=self.model_size,
          action_space=self.action_space,
          num_gru_units=self.num_gru_units,
      )
  
      # Reward Predictor: [ht, zt] -> rt.
      self.reward_predictor = RewardPredictor(
          model_size=self.model_size
      )
      # Continue Predictor: [ht, zt] -> ct.
      self.continue_predictor = ContinuePredictor(
          model_size=self.model_size
      )
  
      # Decoder: [ht, zt] -> x^t.
      self.decoder = decoder
  
      # Trace self.call.
      self.forward_train = tf.function(
          input_signature=[
              tf.TensorSpec(
                  shape=[None, None]
                  + list(self.observation_space.shape)
              ),
              tf.TensorSpec(
                  shape=[None, None]
                  + (
                      [self.action_space.n]
                      if isinstance(action_space, gym.spaces.Discrete)
                      else list(self.action_space.shape)
                  )
              ),
              tf.TensorSpec(shape=[None, None], dtype=tf.bool),
          ]
      )(self.forward_train)
  
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 1 to create a new WorldModel object from inherited base classes: tf.keras with example: obj = WorldModel()