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 training.
  
  1) Forwards all observations [B, T, ...] through the encoder network to yield
  o_processed[B, T, ...].
  2) Uses initial state (h0/z^0/a0[B, 0, ...]) and sequence model (RSSM) to
  compute the first internal state (h1 and z^1).
  3) Uses action a[B, 1, ...], z[B, 1, ...] and h[B, 1, ...] to compute the
  next h-state (h[B, 2, ...]), etc..
  4) Repeats 2) and 3) until t=T.
  5) Uses all h[B, T, ...] and z[B, T, ...] to compute predicted/reconstructed
  observations, rewards, and continue signals.
  6) Returns predictions from 5) along with all z-states z[B, T, ...] and
  the final h-state (h[B, ...] for t=T).
  
  Should we encounter is_first=True flags in the middle of a batch row (somewhere
  within an ongoing sequence of length T), we insert this world model's initial
  state again (zero-action, learned init h-state, and prior-computed z^) and
  simply continue (no zero-padding).
  
  Args:
      observations: The batch (B, T, ...) of observations to be passed through
          the encoder network to yield the inputs to the representation layer
          (which then can compute the posterior z-states).
      actions: The batch (B, T, ...) of actions to be used in combination with
          h-states and computed z-states to yield the next h-states.
      is_first: The batch (B, T) of `is_first` flags.
docstr_tok: ["Performs","a","forward","step","for","training",".","1",")","Forwards","all","observations","[","B",",","T",",","...","]","through","the","encoder","network","to","yield","o_processed","[","B",",","T",",","...","]",".","2",")","Uses","initial","state","(","h0\/z^0\/a0","[","B",",","0",",","...","]",")","and","sequence","model","(","RSSM",")","to","compute","the","first","internal","state","(","h1","and","z^1",")",".","3",")","Uses","action","a","[","B",",","1",",","...","]",",","z","[","B",",","1",",","...","]","and","h","[","B",",","1",",","...","]","to","compute","the","next","h-state","(","h","[","B",",","2",",","...","]",")",",","etc","..","4",")","Repeats","2",")","and","3",")","until","t=T",".","5",")","Uses","all","h","[","B",",","T",",","...","]","and","z","[","B",",","T",",","...","]","to","compute","predicted\/reconstructed","observations",",","rewards",",","and","continue","signals",".","6",")","Returns","predictions","from","5",")","along","with","all","z-states","z","[","B",",","T",",","...","]","and","the","final","h-state","(","h","[","B",",","...","]","for","t=T",")",".","Should","we","encounter","is_first=True","flags","in","the","middle","of","a","batch","row","(","somewhere","within","an","ongoing","sequence","of","length","T",")",",","we","insert","this","world","model","'s","initial","state","again","(","zero-action",",","learned","init","h-state",",","and","prior-computed","z^",")","and","simply","continue","(","no","zero-padding",")",".","Args",":","observations",":","The","batch","(","B",",","T",",","...",")","of","observations","to","be","passed","through","the","encoder","network","to","yield","the","inputs","to","the","representation","layer","(","which","then","can","compute","the","posterior","z-states",")",".","actions",":","The","batch","(","B",",","T",",","...",")","of","actions","to","be","used","in","combination","with","h-states","and","computed","z-states","to","yield","the","next","h-states",".","is_first",":","The","batch","(","B",",","T",")","of","`","is_first","`","flags","."]
code: def forward_train(self, observations, actions, is_first):
      """Performs a forward step for training.
  
      1) Forwards all observations [B, T, ...] through the encoder network to yield
      o_processed[B, T, ...].
      2) Uses initial state (h0/z^0/a0[B, 0, ...]) and sequence model (RSSM) to
      compute the first internal state (h1 and z^1).
      3) Uses action a[B, 1, ...], z[B, 1, ...] and h[B, 1, ...] to compute the
      next h-state (h[B, 2, ...]), etc..
      4) Repeats 2) and 3) until t=T.
      5) Uses all h[B, T, ...] and z[B, T, ...] to compute predicted/reconstructed
      observations, rewards, and continue signals.
      6) Returns predictions from 5) along with all z-states z[B, T, ...] and
      the final h-state (h[B, ...] for t=T).
  
      Should we encounter is_first=True flags in the middle of a batch row (somewhere
      within an ongoing sequence of length T), we insert this world model's initial
      state again (zero-action, learned init h-state, and prior-computed z^) and
      simply continue (no zero-padding).
  
      Args:
          observations: The batch (B, T, ...) of observations to be passed through
              the encoder network to yield the inputs to the representation layer
              (which then can compute the posterior z-states).
          actions: The batch (B, T, ...) of actions to be used in combination with
              h-states and computed z-states to yield the next h-states.
          is_first: The batch (B, T) of `is_first` flags.
      """
      if self.symlog_obs:
          observations = symlog(observations)
  
      # Compute bare encoder outs (not z; this is done later with involvement of the
      # sequence model and the h-states).
      # Fold time dimension for CNN pass.
      shape = tf.shape(observations)
      B, T = shape[0], shape[1]
      observations = tf.reshape(
          observations, shape=tf.concat([[-1], shape[2:]], axis=0)
      )
  
      encoder_out = self.encoder(
          tf.cast(observations, self._comp_dtype)
      )
      # Unfold time dimension.
      encoder_out = tf.reshape(
          encoder_out,
          shape=tf.concat([[B, T], tf.shape(encoder_out)[1:]], axis=0),
      )
      # Make time major for faster upcoming loop.
      encoder_out = tf.transpose(
          encoder_out,
          perm=[1, 0]
          + list(range(2, len(encoder_out.shape.as_list()))),
      )
      # encoder_out=[T, B, ...]
  
      initial_states = tree.map_structure(
          lambda s: tf.repeat(s, B, axis=0), self.get_initial_state()
      )
  
      # Make actions and `is_first` time-major.
      actions = tf.transpose(
          tf.cast(actions, self._comp_dtype),
          perm=[1, 0]
          + list(range(2, tf.shape(actions).shape.as_list()[0])),
      )
      is_first = tf.transpose(
          tf.cast(is_first, self._comp_dtype), perm=[1, 0]
      )
  
      # Loop through the T-axis of our samples and perform one computation step at
      # a time. This is necessary because the sequence model's output (h(t+1)) depends
      # on the current z(t), but z(t) depends on the current sequence model's output
      # h(t).
      z_t0_to_T = [initial_states["z"]]
      z_posterior_probs = []
      z_prior_probs = []
      h_t0_to_T = [initial_states["h"]]
      for t in range(self.batch_length_T):
          # If first, mask it with initial state/actions.
          h_tm1 = self._mask(
              h_t0_to_T[-1], 1.0 - is_first[t]
          )  # zero out
          h_tm1 = h_tm1 + self._mask(
              initial_states["h"], is_first[t]
          )  # add init
  
          z_tm1 = self._mask(
              z_t0_to_T[-1], 1.0 - is_first[t]
          )  # zero out
          z_tm1 = z_tm1 + self._mask(
              initial_states["z"], is_first[t]
          )  # add init
  
          # Zero out actions (no special learnt initial state).
          a_tm1 = self._mask(actions[t - 1], 1.0 - is_first[t])
  
          # Perform one RSSM (sequence model) step to get the current h.
          h_t = self.sequence_model(a=a_tm1, h=h_tm1, z=z_tm1)
          h_t0_to_T.append(h_t)
  
          posterior_mlp_input = tf.concat(
              [encoder_out[t], h_t], axis=-1
          )
          repr_input = self.posterior_mlp(posterior_mlp_input)
          # Draw one z-sample (z(t)) and also get the z-distribution for dynamics and
          # representation loss computations.
          z_t, z_probs = self.posterior_representation_layer(repr_input)
          # z_t=[B, num_categoricals, num_classes]
          z_posterior_probs.append(z_probs)
          z_t0_to_T.append(z_t)
  
          # Compute the predicted z_t (z^) using the dynamics model.
          _, z_probs = self.dynamics_predictor(h_t)
          z_prior_probs.append(z_probs)
  
      # Stack at time dimension to yield: [B, T, ...].
      h_t1_to_T = tf.stack(h_t0_to_T[1:], axis=1)
      z_t1_to_T = tf.stack(z_t0_to_T[1:], axis=1)
  
      # Fold time axis to retrieve the final (loss ready) Independent distribution
      # (over `num_categoricals` Categoricals).
      z_posterior_probs = tf.stack(z_posterior_probs, axis=1)
      z_posterior_probs = tf.reshape(
          z_posterior_probs,
          shape=[-1] + z_posterior_probs.shape.as_list()[2:],
      )
      # Fold time axis to retrieve the final (loss ready) Independent distribution
      # (over `num_categoricals` Categoricals).
      z_prior_probs = tf.stack(z_prior_probs, axis=1)
      z_prior_probs = tf.reshape(
          z_prior_probs,
          shape=[-1] + z_prior_probs.shape.as_list()[2:],
      )
  
      # Fold time dimension for parallelization of all dependent predictions:
      # observations (reproduction via decoder), rewards, continues.
      h_BxT = tf.reshape(
          h_t1_to_T, shape=[-1] + h_t1_to_T.shape.as_list()[2:]
      )
      z_BxT = tf.reshape(
          z_t1_to_T, shape=[-1] + z_t1_to_T.shape.as_list()[2:]
      )
  
      obs_distribution_means = tf.cast(
          self.decoder(h=h_BxT, z=z_BxT), tf.float32
      )
  
      # Compute (predicted) reward distributions.
      rewards, reward_logits = self.reward_predictor(h=h_BxT, z=z_BxT)
  
      # Compute (predicted) continue distributions.
      continues, continue_distribution = self.continue_predictor(
          h=h_BxT, z=z_BxT
      )
  
      # Return outputs for loss computation.
      # Note that all shapes are [BxT, ...] (time axis already folded).
      return {
          # Obs.
          "sampled_obs_symlog_BxT": observations,
          "obs_distribution_means_BxT": obs_distribution_means,
          # Rewards.
          "reward_logits_BxT": reward_logits,
          "rewards_BxT": rewards,
          # Continues.
          "continue_distribution_BxT": continue_distribution,
          "continues_BxT": continues,
          # Deterministic, continuous h-states (t1 to T).
          "h_states_BxT": h_BxT,
          # Sampled, discrete posterior z-states and their probs (t1 to T).
          "z_posterior_states_BxT": z_BxT,
          "z_posterior_probs_BxT": z_posterior_probs,
          # Probs of the prior z-states (t1 to T).
          "z_prior_probs_BxT": z_prior_probs,
      }
  
file_path: 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 4 for calling the WorldModel obj's underlying member method code with example usage: obj.forward_train(observations, actions, is_first) and returns: dict