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: Performs a forward step for inference (e.g. environment stepping).
  
  Works analogous to `forward_train`, except that all inputs are provided
  for a single timestep in the shape of [B, ...] (no time dimension!).
  
  Args:
      observations: The batch (B, ...) of observations to be passed through
          the encoder network to yield the inputs to the representation layer
          (which then can compute the z-states).
      previous_states: A dict with `h`, `z`, and `a` keys mapping to the
          respective previous states/actions. All of the shape (B, ...), no time
          rank.
      is_first: The batch (B) of `is_first` flags.
  
  Returns:
      The next deterministic h-state (h(t+1)) as predicted by the sequence model.
docstr_tok: ["Performs","a","forward","step","for","inference","(","e.g",".","environment","stepping",")",".","Works","analogous","to","`","forward_train","`",",","except","that","all","inputs","are","provided","for","a","single","timestep","in","the","shape","of","[","B",",","...","]","(","no","time","dimension","!",")",".","Args",":","observations",":","The","batch","(","B",",","...",")","of","observations","to","be","passed","through","the","encoder","network","to","yield","the","inputs","to","the","representation","layer","(","which","then","can","compute","the","z-states",")",".","previous_states",":","A","dict","with","`","h","`",",","`","z","`",",","and","`","a","`","keys","mapping","to","the","respective","previous","states\/actions",".","All","of","the","shape","(","B",",","...",")",",","no","time","rank",".","is_first",":","The","batch","(","B",")","of","`","is_first","`","flags",".","Returns",":","The","next","deterministic","h-state","(","h","(","t+1",")",")","as","predicted","by","the","sequence","model","."]
code: def forward_inference(
      self, observations, previous_states, is_first, training=None
  ):
      """Performs a forward step for inference (e.g. environment stepping).
  
      Works analogous to `forward_train`, except that all inputs are provided
      for a single timestep in the shape of [B, ...] (no time dimension!).
  
      Args:
          observations: The batch (B, ...) of observations to be passed through
              the encoder network to yield the inputs to the representation layer
              (which then can compute the z-states).
          previous_states: A dict with `h`, `z`, and `a` keys mapping to the
              respective previous states/actions. All of the shape (B, ...), no time
              rank.
          is_first: The batch (B) of `is_first` flags.
  
      Returns:
          The next deterministic h-state (h(t+1)) as predicted by the sequence model.
      """
      observations = tf.cast(observations, self._comp_dtype)
  
      initial_states = tree.map_structure(
          lambda s: tf.repeat(s, tf.shape(observations)[0], axis=0),
          self.get_initial_state(),
      )
  
      # If first, mask it with initial state/actions.
      previous_h = self._mask(
          previous_states["h"], 1.0 - is_first
      )  # zero out
      previous_h = previous_h + self._mask(
          initial_states["h"], is_first
      )  # add init
  
      previous_z = self._mask(
          previous_states["z"], 1.0 - is_first
      )  # zero out
      previous_z = previous_z + self._mask(
          initial_states["z"], is_first
      )  # add init
  
      # Zero out actions (no special learnt initial state).
      previous_a = self._mask(previous_states["a"], 1.0 - is_first)
  
      # Compute new states.
      h = self.sequence_model(a=previous_a, h=previous_h, z=previous_z)
      z = self.compute_posterior_z(
          observations=observations, initial_h=h
      )
  
      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 3 for calling the WorldModel obj's underlying member method code with example usage: obj.forward_inference(observations, previous_states, is_first, training) and returns: dict