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 the world-model's dynamics and representation losses.
  
  Args:
      config: The DreamerV3Config to use.
      fwd_out: The `forward_train` outputs of the DreamerV3RLModule.
  
  Returns:
      Tuple consisting of a) dynamics loss: Trains the prior network, predicting
      z^ prior states from h-states and b) representation loss: Trains posterior
      network, predicting z posterior states from h-states and (encoded)
      observations.
docstr_tok: ["Helper","method","computing","the","world-model","'s","dynamics","and","representation","losses",".","Args",":","config",":","The","DreamerV3Config","to","use",".","fwd_out",":","The","`","forward_train","`","outputs","of","the","DreamerV3RLModule",".","Returns",":","Tuple","consisting","of","a",")","dynamics","loss",":","Trains","the","prior","network",",","predicting","z^","prior","states","from","h-states","and","b",")","representation","loss",":","Trains","posterior","network",",","predicting","z","posterior","states","from","h-states","and","(","encoded",")","observations","."]
code: def _compute_world_model_dynamics_and_representation_loss(
      self, *, config: DreamerV3Config, fwd_out: Dict[str, Any]
  ) -> Tuple[TensorType, TensorType]:
      """Helper method computing the world-model's dynamics and representation losses.
  
      Args:
          config: The DreamerV3Config to use.
          fwd_out: The `forward_train` outputs of the DreamerV3RLModule.
  
      Returns:
          Tuple consisting of a) dynamics loss: Trains the prior network, predicting
          z^ prior states from h-states and b) representation loss: Trains posterior
          network, predicting z posterior states from h-states and (encoded)
          observations.
      """
  
      # Actual distribution over stochastic internal states (z) produced by the
      # encoder.
      z_posterior_probs_BxT = fwd_out["z_posterior_probs_BxT"]
      z_posterior_distr_BxT = tfp.distributions.Independent(
          tfp.distributions.OneHotCategorical(
              probs=z_posterior_probs_BxT
          ),
          reinterpreted_batch_ndims=1,
      )
  
      # Actual distribution over stochastic internal states (z) produced by the
      # dynamics network.
      z_prior_probs_BxT = fwd_out["z_prior_probs_BxT"]
      z_prior_distr_BxT = tfp.distributions.Independent(
          tfp.distributions.OneHotCategorical(probs=z_prior_probs_BxT),
          reinterpreted_batch_ndims=1,
      )
  
      # Stop gradient for encoder's z-outputs:
      sg_z_posterior_distr_BxT = tfp.distributions.Independent(
          tfp.distributions.OneHotCategorical(
              probs=tf.stop_gradient(z_posterior_probs_BxT)
          ),
          reinterpreted_batch_ndims=1,
      )
      # Stop gradient for dynamics model's z-outputs:
      sg_z_prior_distr_BxT = tfp.distributions.Independent(
          tfp.distributions.OneHotCategorical(
              probs=tf.stop_gradient(z_prior_probs_BxT)
          ),
          reinterpreted_batch_ndims=1,
      )
  
      # Implement free bits. According to [1]:
      # "To avoid a degenerate solution where the dynamics are trivial to predict but
      # contain not enough information about the inputs, we employ free bits by
      # clipping the dynamics and representation losses below the value of
      # 1 nat ≈ 1.44 bits. This disables them while they are already minimized well to
      # focus the world model on its prediction loss"
      L_dyn_BxT = tf.math.maximum(
          1.0,
          tfp.distributions.kl_divergence(
              sg_z_posterior_distr_BxT, z_prior_distr_BxT
          ),
      )
      # Unfold time rank back in.
      L_dyn_B_T = tf.reshape(
          L_dyn_BxT,
          (config.batch_size_B_per_learner, config.batch_length_T),
      )
  
      L_rep_BxT = tf.math.maximum(
          1.0,
          tfp.distributions.kl_divergence(
              z_posterior_distr_BxT, sg_z_prior_distr_BxT
          ),
      )
      # Unfold time rank back in.
      L_rep_B_T = tf.reshape(
          L_rep_BxT,
          (config.batch_size_B_per_learner, config.batch_length_T),
      )
  
      return L_dyn_B_T, L_rep_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 6 for calling the DreamerV3TfLearner obj's underlying member method code with example usage: obj._compute_world_model_dynamics_and_representation_loss() and returns: L_dyn_B_T,  L_rep_B_T