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: Dreams trajectories of length H from batch of h- and z-states.
  
  Note that incoming data will have the shapes (BxT, ...), where the original
  batch- and time-dimensions are already folded together. Beginning from this
  new batch dim (BxT), we will unroll `timesteps_H` timesteps in a time-major
  fashion, such that the dreamed data will have shape (H, BxT, ...).
  
  Args:
      start_states: Dict of `h` and `z` states in the shape of (B, ...) and
          (B, num_categoricals, num_classes), respectively, as
          computed by a train forward pass. From each individual h-/z-state pair
          in the given batch, we will branch off a dreamed trajectory of len
          `timesteps_H`.
      start_is_terminated: Float flags of shape (B,) indicating whether the
          first timesteps of each batch row is already a terminated timestep
          (given by the actual environment).
docstr_tok: ["Dreams","trajectories","of","length","H","from","batch","of","h-","and","z-states",".","Note","that","incoming","data","will","have","the","shapes","(","BxT",",","...",")",",","where","the","original","batch-","and","time-dimensions","are","already","folded","together",".","Beginning","from","this","new","batch","dim","(","BxT",")",",","we","will","unroll","`","timesteps_H","`","timesteps","in","a","time-major","fashion",",","such","that","the","dreamed","data","will","have","shape","(","H",",","BxT",",","...",")",".","Args",":","start_states",":","Dict","of","`","h","`","and","`","z","`","states","in","the","shape","of","(","B",",","...",")","and","(","B",",","num_categoricals",",","num_classes",")",",","respectively",",","as","computed","by","a","train","forward","pass",".","From","each","individual","h-\/z-state","pair","in","the","given","batch",",","we","will","branch","off","a","dreamed","trajectory","of","len","`","timesteps_H","`",".","start_is_terminated",":","Float","flags","of","shape","(","B",",",")","indicating","whether","the","first","timesteps","of","each","batch","row","is","already","a","terminated","timestep","(","given","by","the","actual","environment",")","."]
code: def dream_trajectory(self, start_states, start_is_terminated):
      """Dreams trajectories of length H from batch of h- and z-states.
  
      Note that incoming data will have the shapes (BxT, ...), where the original
      batch- and time-dimensions are already folded together. Beginning from this
      new batch dim (BxT), we will unroll `timesteps_H` timesteps in a time-major
      fashion, such that the dreamed data will have shape (H, BxT, ...).
  
      Args:
          start_states: Dict of `h` and `z` states in the shape of (B, ...) and
              (B, num_categoricals, num_classes), respectively, as
              computed by a train forward pass. From each individual h-/z-state pair
              in the given batch, we will branch off a dreamed trajectory of len
              `timesteps_H`.
          start_is_terminated: Float flags of shape (B,) indicating whether the
              first timesteps of each batch row is already a terminated timestep
              (given by the actual environment).
      """
      # Dreamed actions (one-hot encoded for discrete actions).
      a_dreamed_t0_to_H = []
      a_dreamed_dist_params_t0_to_H = []
  
      h = start_states["h"]
      z = start_states["z"]
  
      # GRU outputs.
      h_states_t0_to_H = [h]
      # Dynamics model outputs.
      z_states_prior_t0_to_H = [z]
  
      # Compute `a` using actor network (already the first step uses a dreamed action,
      # not a sampled one).
      a, a_dist_params = self.actor(
          # We have to stop the gradients through the states. B/c we are using a
          # differentiable Discrete action distribution (straight through gradients
          # with `a = stop_gradient(sample(probs)) + probs - stop_gradient(probs)`,
          # we otherwise would add dependencies of the `-log(pi(a|s))` REINFORCE loss
          # term on actions further back in the trajectory.
          h=tf.stop_gradient(h),
          z=tf.stop_gradient(z),
      )
      a_dreamed_t0_to_H.append(a)
      a_dreamed_dist_params_t0_to_H.append(a_dist_params)
  
      for i in range(self.horizon):
          # Move one step in the dream using the RSSM.
          h = self.world_model.sequence_model(a=a, h=h, z=z)
          h_states_t0_to_H.append(h)
  
          # Compute prior z using dynamics model.
          z, _ = self.world_model.dynamics_predictor(h=h)
          z_states_prior_t0_to_H.append(z)
  
          # Compute `a` using actor network.
          a, a_dist_params = self.actor(
              h=tf.stop_gradient(h),
              z=tf.stop_gradient(z),
          )
          a_dreamed_t0_to_H.append(a)
          a_dreamed_dist_params_t0_to_H.append(a_dist_params)
  
      h_states_H_B = tf.stack(h_states_t0_to_H, axis=0)  # (T, B, ...)
      h_states_HxB = tf.reshape(
          h_states_H_B, [-1] + h_states_H_B.shape.as_list()[2:]
      )
  
      z_states_prior_H_B = tf.stack(
          z_states_prior_t0_to_H, axis=0
      )  # (T, B, ...)
      z_states_prior_HxB = tf.reshape(
          z_states_prior_H_B,
          [-1] + z_states_prior_H_B.shape.as_list()[2:],
      )
  
      a_dreamed_H_B = tf.stack(a_dreamed_t0_to_H, axis=0)  # (T, B, ...)
      a_dreamed_dist_params_H_B = tf.stack(
          a_dreamed_dist_params_t0_to_H, axis=0
      )
  
      # Compute r using reward predictor.
      r_dreamed_HxB, _ = self.world_model.reward_predictor(
          h=h_states_HxB, z=z_states_prior_HxB
      )
      r_dreamed_H_B = tf.reshape(
          inverse_symlog(r_dreamed_HxB), shape=[self.horizon + 1, -1]
      )
  
      # Compute intrinsic rewards.
      if self.use_curiosity:
          results_HxB = self.disagree_nets.compute_intrinsic_rewards(
              h=h_states_HxB,
              z=z_states_prior_HxB,
              a=tf.reshape(
                  a_dreamed_H_B,
                  [-1] + a_dreamed_H_B.shape.as_list()[2:],
              ),
          )
          # TODO (sven): Wrong? -> Cut out last timestep as we always predict z-states
          #  for the NEXT timestep and derive ri (for the NEXT timestep) from the
          #  disagreement between our N disagreee nets.
          r_intrinsic_H_B = tf.reshape(
              results_HxB["rewards_intrinsic"],
              shape=[self.horizon + 1, -1],
          )[
              1:
          ]  # cut out first ts instead
          curiosity_forward_train_outs = results_HxB[
              "forward_train_outs"
          ]
          del results_HxB
  
      # Compute continues using continue predictor.
      c_dreamed_HxB, _ = self.world_model.continue_predictor(
          h=h_states_HxB,
          z=z_states_prior_HxB,
      )
      c_dreamed_H_B = tf.reshape(c_dreamed_HxB, [self.horizon + 1, -1])
      # Force-set first `continue` flags to False iff `start_is_terminated`.
      # Note: This will cause the loss-weights for this row in the batch to be
      # completely zero'd out. In general, we don't use dreamed data past any
      # predicted (or actual first) continue=False flags.
      c_dreamed_H_B = tf.concat(
          [
              1.0
              - tf.expand_dims(
                  tf.cast(start_is_terminated, tf.float32),
                  0,
              ),
              c_dreamed_H_B[1:],
          ],
          axis=0,
      )
  
      # Loss weights for each individual dreamed timestep. Zero-out all timesteps
      # that lie past continue=False flags. B/c our world model does NOT learn how
      # to skip terminal/reset episode boundaries, dreamed data crossing such a
      # boundary should not be used for critic/actor learning either.
      dream_loss_weights_H_B = (
          tf.math.cumprod(self.gamma * c_dreamed_H_B, axis=0)
          / self.gamma
      )
  
      # Compute the value estimates.
      v, v_symlog_dreamed_logits_HxB = self.critic(
          h=h_states_HxB,
          z=z_states_prior_HxB,
          use_ema=False,
      )
      v_dreamed_HxB = inverse_symlog(v)
      v_dreamed_H_B = tf.reshape(
          v_dreamed_HxB, shape=[self.horizon + 1, -1]
      )
  
      v_symlog_dreamed_ema_HxB, _ = self.critic(
          h=h_states_HxB,
          z=z_states_prior_HxB,
          use_ema=True,
      )
      v_symlog_dreamed_ema_H_B = tf.reshape(
          v_symlog_dreamed_ema_HxB, shape=[self.horizon + 1, -1]
      )
  
      ret = {
          "h_states_t0_to_H_BxT": h_states_H_B,
          "z_states_prior_t0_to_H_BxT": z_states_prior_H_B,
          "rewards_dreamed_t0_to_H_BxT": r_dreamed_H_B,
          "continues_dreamed_t0_to_H_BxT": c_dreamed_H_B,
          "actions_dreamed_t0_to_H_BxT": a_dreamed_H_B,
          "actions_dreamed_dist_params_t0_to_H_BxT": a_dreamed_dist_params_H_B,
          "values_dreamed_t0_to_H_BxT": v_dreamed_H_B,
          "values_symlog_dreamed_logits_t0_to_HxBxT": v_symlog_dreamed_logits_HxB,
          "v_symlog_dreamed_ema_t0_to_H_BxT": v_symlog_dreamed_ema_H_B,
          # Loss weights for critic- and actor losses.
          "dream_loss_weights_t0_to_H_BxT": dream_loss_weights_H_B,
      }
  
      if self.use_curiosity:
          ret["rewards_intrinsic_t1_to_H_B"] = r_intrinsic_H_B
          ret.update(curiosity_forward_train_outs)
  
      if isinstance(self.action_space, gym.spaces.Discrete):
          ret["actions_ints_dreamed_t0_to_H_B"] = tf.argmax(
              a_dreamed_H_B, axis=-1
          )
  
      return ret
  
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 7 for calling the DreamerModel obj's underlying member method code with example usage: obj.dream_trajectory(start_states, start_is_terminated) and returns: ret