class_docstr: The main tf-keras model containing all necessary components for DreamerV3.
  
  Includes:
  - The world model with encoder, decoder, sequence-model (RSSM), dynamics
  (generates prior z-state), and "posterior" model (generates posterior z-state).
  Predicts env dynamics and produces dreamed trajectories for actor- and critic
  learning.
  - The actor network (policy).
  - The critic network for value function prediction.
class_docstr_tok: ["The","main","tf-keras","model","containing","all","necessary","components","for","DreamerV3",".","Includes",":","-","The","world","model","with","encoder",",","decoder",",","sequence-model","(","RSSM",")",",","dynamics","(","generates","prior","z-state",")",",","and","``","posterior","''","model","(","generates","posterior","z-state",")",".","Predicts","env","dynamics","and","produces","dreamed","trajectories","for","actor-","and","critic","learning",".","-","The","actor","network","(","policy",")",".","-","The","critic","network","for","value","function","prediction","."]
docstr: Performs a training forward pass given observations and actions.
  
  Note that all input data must have a time rank (batch-major: [B, T, ...]).
  
  Args:
      observations: The environment observations with shape (B, T, ...). Thus,
          the batch has B rows of T timesteps each. Note that it's ok to have
          episode boundaries (is_first=True) within a batch row. DreamerV3 will
          simply insert an initial state before these locations and continue the
          sequence modelling (with the RSSM). Hence, there will be no zero
          padding.
      actions: The actions actually taken in the environment with shape
          (B, T, ...). See `observations` docstring for details on how B and T are
          handled.
      is_first: Batch of is_first flags. These should be True:
          - if a new episode has been started at the current timestep (meaning
          `observations` is the reset observation from the environment).
          - in each batch row at T=0 (first timestep of each of the B batch
          rows), regardless of whether the actual env had an episode boundary
          there or not.
docstr_tok: ["Performs","a","training","forward","pass","given","observations","and","actions",".","Note","that","all","input","data","must","have","a","time","rank","(","batch-major",":","[","B",",","T",",","...","]",")",".","Args",":","observations",":","The","environment","observations","with","shape","(","B",",","T",",","...",")",".","Thus",",","the","batch","has","B","rows","of","T","timesteps","each",".","Note","that","it","'s","ok","to","have","episode","boundaries","(","is_first=True",")","within","a","batch","row",".","DreamerV3","will","simply","insert","an","initial","state","before","these","locations","and","continue","the","sequence","modelling","(","with","the","RSSM",")",".","Hence",",","there","will","be","no","zero","padding",".","actions",":","The","actions","actually","taken","in","the","environment","with","shape","(","B",",","T",",","...",")",".","See","`","observations","`","docstring","for","details","on","how","B","and","T","are","handled",".","is_first",":","Batch","of","is_first","flags",".","These","should","be","True",":","-","if","a","new","episode","has","been","started","at","the","current","timestep","(","meaning","`","observations","`","is","the","reset","observation","from","the","environment",")",".","-","in","each","batch","row","at","T=0","(","first","timestep","of","each","of","the","B","batch","rows",")",",","regardless","of","whether","the","actual","env","had","an","episode","boundary","there","or","not","."]
code: def forward_train(self, observations, actions, is_first):
      """Performs a training forward pass given observations and actions.
  
      Note that all input data must have a time rank (batch-major: [B, T, ...]).
  
      Args:
          observations: The environment observations with shape (B, T, ...). Thus,
              the batch has B rows of T timesteps each. Note that it's ok to have
              episode boundaries (is_first=True) within a batch row. DreamerV3 will
              simply insert an initial state before these locations and continue the
              sequence modelling (with the RSSM). Hence, there will be no zero
              padding.
          actions: The actions actually taken in the environment with shape
              (B, T, ...). See `observations` docstring for details on how B and T are
              handled.
          is_first: Batch of is_first flags. These should be True:
              - if a new episode has been started at the current timestep (meaning
              `observations` is the reset observation from the environment).
              - in each batch row at T=0 (first timestep of each of the B batch
              rows), regardless of whether the actual env had an episode boundary
              there or not.
      """
      return self.world_model.forward_train(
          observations=observations,
          actions=actions,
          is_first=is_first,
      )
  
file_path: ray/rllib/algorithms/dreamerv3/tf/models/dreamer_model.py
imports: import re
  import gymnasium
  import numpy
  from ray.rllib.algorithms.dreamerv3.tf.models.disagree_networks import DisagreeNetworks
  from ray.rllib.algorithms.dreamerv3.tf.models.actor_network import ActorNetwork
  from ray.rllib.algorithms.dreamerv3.tf.models.critic_network import CriticNetwork
  from ray.rllib.algorithms.dreamerv3.tf.models.world_model import WorldModel
  from ray.rllib.algorithms.dreamerv3.utils import get_gru_units, get_num_z_categoricals, get_num_z_classes
  from ray.rllib.utils.framework import try_import_tf
  from ray.rllib.utils.tf_utils import inverse_symlog
label_desc: Use image node_id 5 for calling the DreamerModel obj's underlying member method code with example usage: obj.forward_train(observations, actions, is_first) and returns: self