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 critic's loss terms.
  
  Args:
      module_id: The ModuleID for which to compute the critic loss.
      config: The DreamerV3Config to use.
      dream_data: The data generated by dreaming for H steps (horizon) starting
          from any BxT state (sampled from the buffer for the train batch).
      value_targets_t0_to_Hm1_BxT: The computed value function targets of the
          shape (t0 to H-1, BxT).
  
  Returns:
      The total critic loss tensor.
docstr_tok: ["Helper","method","computing","the","critic","'s","loss","terms",".","Args",":","module_id",":","The","ModuleID","for","which","to","compute","the","critic","loss",".","config",":","The","DreamerV3Config","to","use",".","dream_data",":","The","data","generated","by","dreaming","for","H","steps","(","horizon",")","starting","from","any","BxT","state","(","sampled","from","the","buffer","for","the","train","batch",")",".","value_targets_t0_to_Hm1_BxT",":","The","computed","value","function","targets","of","the","shape","(","t0","to","H-1",",","BxT",")",".","Returns",":","The","total","critic","loss","tensor","."]
code: def _compute_critic_loss(
      self,
      *,
      module_id: ModuleID,
      config: DreamerV3Config,
      dream_data: Dict[str, TensorType],
      value_targets_t0_to_Hm1_BxT: TensorType,
  ) -> TensorType:
      """Helper method computing the critic's loss terms.
  
      Args:
          module_id: The ModuleID for which to compute the critic loss.
          config: The DreamerV3Config to use.
          dream_data: The data generated by dreaming for H steps (horizon) starting
              from any BxT state (sampled from the buffer for the train batch).
          value_targets_t0_to_Hm1_BxT: The computed value function targets of the
              shape (t0 to H-1, BxT).
  
      Returns:
          The total critic loss tensor.
      """
      # B=BxT
      H, B = dream_data["rewards_dreamed_t0_to_H_BxT"].shape[:2]
      Hm1 = H - 1
  
      # Note that value targets are NOT symlog'd and go from t0 to H-1, not H, like
      # all the other dream data.
  
      # From here on: B=BxT
      value_targets_t0_to_Hm1_B = tf.stop_gradient(
          value_targets_t0_to_Hm1_BxT
      )
      value_symlog_targets_t0_to_Hm1_B = symlog(
          value_targets_t0_to_Hm1_B
      )
      # Fold time rank (for two_hot'ing).
      value_symlog_targets_HxB = tf.reshape(
          value_symlog_targets_t0_to_Hm1_B, (-1,)
      )
      value_symlog_targets_two_hot_HxB = two_hot(
          value_symlog_targets_HxB
      )
      # Unfold time rank.
      value_symlog_targets_two_hot_t0_to_Hm1_B = tf.reshape(
          value_symlog_targets_two_hot_HxB,
          shape=[Hm1, B, value_symlog_targets_two_hot_HxB.shape[-1]],
      )
  
      # Get (B x T x probs) tensor from return distributions.
      value_symlog_logits_HxB = dream_data[
          "values_symlog_dreamed_logits_t0_to_HxBxT"
      ]
      # Unfold time rank and cut last time index to match value targets.
      value_symlog_logits_t0_to_Hm1_B = tf.reshape(
          value_symlog_logits_HxB,
          shape=[H, B, value_symlog_logits_HxB.shape[-1]],
      )[:-1]
  
      values_log_pred_Hm1_B = (
          value_symlog_logits_t0_to_Hm1_B
          - tf.math.reduce_logsumexp(
              value_symlog_logits_t0_to_Hm1_B, axis=-1, keepdims=True
          )
      )
      # Multiply with two-hot targets and neg.
      value_loss_two_hot_H_B = -tf.reduce_sum(
          values_log_pred_Hm1_B
          * value_symlog_targets_two_hot_t0_to_Hm1_B,
          axis=-1,
      )
  
      # Compute EMA regularization loss.
      # Expected values (dreamed) from the EMA (slow critic) net.
      # Note: Slow critic (EMA) outputs are already stop_gradient'd.
      value_symlog_ema_t0_to_Hm1_B = tf.stop_gradient(
          dream_data["v_symlog_dreamed_ema_t0_to_H_BxT"]
      )[:-1]
      # Fold time rank (for two_hot'ing).
      value_symlog_ema_HxB = tf.reshape(
          value_symlog_ema_t0_to_Hm1_B, (-1,)
      )
      value_symlog_ema_two_hot_HxB = two_hot(value_symlog_ema_HxB)
      # Unfold time rank.
      value_symlog_ema_two_hot_t0_to_Hm1_B = tf.reshape(
          value_symlog_ema_two_hot_HxB,
          shape=[Hm1, B, value_symlog_ema_two_hot_HxB.shape[-1]],
      )
  
      # Compute ema regularizer loss.
      # In the paper, it is not described how exactly to form this regularizer term
      # and how to weigh it.
      # So we follow Danijar's repo here:
      # `reg = -dist.log_prob(sg(self.slow(traj).mean()))`
      # with a weight of 1.0, where dist is the bucket'ized distribution output by the
      # fast critic. sg=stop gradient; mean() -> use the expected EMA values.
      # Multiply with two-hot targets and neg.
      ema_regularization_loss_H_B = -tf.reduce_sum(
          values_log_pred_Hm1_B * value_symlog_ema_two_hot_t0_to_Hm1_B,
          axis=-1,
      )
  
      L_critic_H_B = (
          value_loss_two_hot_H_B + ema_regularization_loss_H_B
      )
  
      # Mask out everything that goes beyond a predicted continue=False boundary.
      L_critic_H_B *= tf.stop_gradient(
          dream_data["dream_loss_weights_t0_to_H_BxT"]
      )[:-1]
  
      # Reduce over H- (time) axis (sum) and then B-axis (mean).
      L_critic = tf.reduce_mean(L_critic_H_B)
  
      self.register_metrics(
          module_id=module_id,
          metrics_dict={
              "CRITIC_L_total": L_critic,
              "CRITIC_L_neg_logp_of_value_targets": tf.reduce_mean(
                  value_loss_two_hot_H_B
              ),
              "CRITIC_L_slow_critic_regularization": tf.reduce_mean(
                  ema_regularization_loss_H_B
              ),
          },
      )
      if config.report_individual_batch_item_stats:
          self.register_metrics(
              module_id=module_id,
              metrics_dict={
                  # Symlog'd value targets. Critic learns to predict symlog'd values.
                  "VALUE_TARGETS_symlog_H_BxT": value_symlog_targets_t0_to_Hm1_B,
                  # Critic loss terms.
                  "CRITIC_L_total_H_BxT": L_critic_H_B,
                  "CRITIC_L_neg_logp_of_value_targets_H_BxT": value_loss_two_hot_H_B,
                  "CRITIC_L_slow_critic_regularization_H_BxT": (
                      ema_regularization_loss_H_B
                  ),
              },
          )
  
      return L_critic
  
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 8 for calling the DreamerV3TfLearner obj's underlying member method code with example usage: obj._compute_critic_loss() and returns: L_critic