class_docstr: Implements DreamerV3 losses and gradient-based update logic in TensorFlow.
  
  The critic EMA-copy update step can be found in the `DreamerV3Learner` base class,
  as it is framework independent.
  
  We define 3 local TensorFlow optimizers for the sub components "world_model",
  "actor", and "critic". Each of these optimizers might use a different learning rate,
  epsilon parameter, and gradient clipping thresholds and procedures.
docstr: Helper method computing all world-model related prediction losses.
  
  Prediction losses are used to train the predictors of the world model, which
  are: Reward predictor, continue predictor, and the decoder (which predicts
  observations).
  
  Args:
      config: The DreamerV3Config to use.
      rewards_B_T: The rewards batch in the shape (B, T) and of type float32.
      continues_B_T: The continues batch in the shape (B, T) and of type float32
          (1.0 -> continue; 0.0 -> end of episode).
      fwd_out: The `forward_train` outputs of the DreamerV3RLModule.
docstr_tok: ["Helper","method","computing","all","world-model","related","prediction","losses",".","Prediction","losses","are","used","to","train","the","predictors","of","the","world","model",",","which","are",":","Reward","predictor",",","continue","predictor",",","and","the","decoder","(","which","predicts","observations",")",".","Args",":","config",":","The","DreamerV3Config","to","use",".","rewards_B_T",":","The","rewards","batch","in","the","shape","(","B",",","T",")","and","of","type","float32",".","continues_B_T",":","The","continues","batch","in","the","shape","(","B",",","T",")","and","of","type","float32","(","1.0","-",">","continue",";","0.0","-",">","end","of","episode",")",".","fwd_out",":","The","`","forward_train","`","outputs","of","the","DreamerV3RLModule","."]
code: def _compute_world_model_prediction_losses(
      self,
      *,
      config: DreamerV3Config,
      rewards_B_T: TensorType,
      continues_B_T: TensorType,
      fwd_out: Dict[str, TensorType],
  ) -> Dict[str, TensorType]:
      """Helper method computing all world-model related prediction losses.
  
      Prediction losses are used to train the predictors of the world model, which
      are: Reward predictor, continue predictor, and the decoder (which predicts
      observations).
  
      Args:
          config: The DreamerV3Config to use.
          rewards_B_T: The rewards batch in the shape (B, T) and of type float32.
          continues_B_T: The continues batch in the shape (B, T) and of type float32
              (1.0 -> continue; 0.0 -> end of episode).
          fwd_out: The `forward_train` outputs of the DreamerV3RLModule.
      """
  
      # Learn to produce symlog'd observation predictions.
      # If symlog is disabled (e.g. for uint8 image inputs), `obs_symlog_BxT` is the
      # same as `obs_BxT`.
      obs_BxT = fwd_out["sampled_obs_symlog_BxT"]
      obs_distr_means = fwd_out["obs_distribution_means_BxT"]
      # In case we wanted to construct a distribution object from the fwd out data,
      # we would have to do it like this:
      # obs_distr = tfp.distributions.MultivariateNormalDiag(
      #    loc=obs_distr_means,
      #    # Scale == 1.0.
      #    # [2]: "Distributions The image predictor outputs the mean of a diagonal
      #    # Gaussian likelihood with **unit variance** ..."
      #    scale_diag=tf.ones_like(obs_distr_means),
      # )
  
      # Leave time dim folded (BxT) and flatten all other (e.g. image) dims.
      obs_BxT = tf.reshape(
          obs_BxT, shape=[-1, tf.reduce_prod(obs_BxT.shape[1:])]
      )
  
      # Squared diff loss w/ sum(!) over all (already folded) obs dims.
      # decoder_loss_BxT = SUM[ (obs_distr.loc - observations)^2 ]
      # Note: This is described strangely in the paper (stating a neglogp loss here),
      # but the author's own implementation actually uses simple MSE with the loc
      # of the Gaussian.
      decoder_loss_BxT = tf.reduce_sum(
          tf.math.square(obs_distr_means - obs_BxT), axis=-1
      )
  
      # Unfold time rank back in.
      decoder_loss_B_T = tf.reshape(
          decoder_loss_BxT,
          (config.batch_size_B_per_learner, config.batch_length_T),
      )
      L_decoder = tf.reduce_mean(decoder_loss_B_T)
  
      # The FiniteDiscrete reward bucket distribution computed by our reward
      # predictor.
      # [B x num_buckets].
      reward_logits_BxT = fwd_out["reward_logits_BxT"]
      # Learn to produce symlog'd reward predictions.
      rewards_symlog_B_T = symlog(tf.cast(rewards_B_T, tf.float32))
      # Fold time dim.
      rewards_symlog_BxT = tf.reshape(rewards_symlog_B_T, shape=[-1])
  
      # Two-hot encode.
      two_hot_rewards_symlog_BxT = two_hot(rewards_symlog_BxT)
      # two_hot_rewards_symlog_BxT=[B*T, num_buckets]
      reward_log_pred_BxT = (
          reward_logits_BxT
          - tf.math.reduce_logsumexp(
              reward_logits_BxT, axis=-1, keepdims=True
          )
      )
      # Multiply with two-hot targets and neg.
      reward_loss_two_hot_BxT = -tf.reduce_sum(
          reward_log_pred_BxT * two_hot_rewards_symlog_BxT, axis=-1
      )
      # Unfold time rank back in.
      reward_loss_two_hot_B_T = tf.reshape(
          reward_loss_two_hot_BxT,
          (config.batch_size_B_per_learner, config.batch_length_T),
      )
      L_reward_two_hot = tf.reduce_mean(reward_loss_two_hot_B_T)
  
      # Probabilities that episode continues, computed by our continue predictor.
      # [B]
      continue_distr = fwd_out["continue_distribution_BxT"]
      # -log(p) loss
      # Fold time dim.
      continues_BxT = tf.reshape(continues_B_T, shape=[-1])
      continue_loss_BxT = -continue_distr.log_prob(continues_BxT)
      # Unfold time rank back in.
      continue_loss_B_T = tf.reshape(
          continue_loss_BxT,
          (config.batch_size_B_per_learner, config.batch_length_T),
      )
      L_continue = tf.reduce_mean(continue_loss_B_T)
  
      # Sum all losses together as the "prediction" loss.
      L_pred_B_T = (
          decoder_loss_B_T + reward_loss_two_hot_B_T + continue_loss_B_T
      )
      L_pred = tf.reduce_mean(L_pred_B_T)
  
      return {
          "L_decoder_B_T": decoder_loss_B_T,
          "L_decoder": L_decoder,
          "L_reward": L_reward_two_hot,
          "L_reward_B_T": reward_loss_two_hot_B_T,
          "L_continue": L_continue,
          "L_continue_B_T": continue_loss_B_T,
          "L_prediction": L_pred,
          "L_prediction_B_T": L_pred_B_T,
      }
  
imports: from typing import Any, Dict, Tuple
  import gymnasium
  from ray.rllib.algorithms.dreamerv3.dreamerv3 import DreamerV3Config
  from ray.rllib.algorithms.dreamerv3.dreamerv3_learner import DreamerV3Learner
  from ray.rllib.core.learner.learner import ParamDict
  from ray.rllib.core.learner.tf.tf_learner import TfLearner
  from ray.rllib.policy.sample_batch import DEFAULT_POLICY_ID, SampleBatch
  from ray.rllib.utils.annotations import override
  from ray.rllib.utils.deprecation import deprecation_warning
  from ray.rllib.utils.framework import try_import_tf, try_import_tfp
  from ray.rllib.utils.tf_utils import symlog, two_hot, clip_gradients
  from ray.rllib.utils.typing import ModuleID, TensorType
label_desc: Use image node_id 5 for calling the DreamerV3TfLearner obj's underlying member method code with example usage: obj._compute_world_model_prediction_losses() and returns: dict