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.
code: def compute_loss_for_module(
      self,
      module_id: ModuleID,
      config: DreamerV3Config,
      batch: SampleBatch,
      fwd_out: Dict[str, TensorType],
  ) -> TensorType:
      # World model losses.
      prediction_losses = self._compute_world_model_prediction_losses(
          config=config,
          rewards_B_T=batch[SampleBatch.REWARDS],
          continues_B_T=(
              1.0 - tf.cast(batch["is_terminated"], tf.float32)
          ),
          fwd_out=fwd_out,
      )
  
      (
          L_dyn_B_T,
          L_rep_B_T,
      ) = self._compute_world_model_dynamics_and_representation_loss(
          config=config, fwd_out=fwd_out
      )
      L_dyn = tf.reduce_mean(L_dyn_B_T)
      L_rep = tf.reduce_mean(L_rep_B_T)
      # Make sure values for L_rep and L_dyn are the same (they only differ in their
      # gradients).
      tf.assert_equal(L_dyn, L_rep)
  
      # Compute the actual total loss using fixed weights described in [1] eq. 4.
      L_world_model_total_B_T = (
          1.0 * prediction_losses["L_prediction_B_T"]
          + 0.5 * L_dyn_B_T
          + 0.1 * L_rep_B_T
      )
  
      # In the paper, it says to sum up timesteps, and average over
      # batch (see eq. 4 in [1]). But Danijar's implementation only does
      # averaging (over B and T), so we'll do this here as well. This is generally
      # true for all other loss terms as well (we'll always just average, no summing
      # over T axis!).
      L_world_model_total = tf.reduce_mean(L_world_model_total_B_T)
  
      # Register world model loss stats.
      self.register_metrics(
          module_id=module_id,
          metrics_dict={
              "WORLD_MODEL_learned_initial_h": (
                  self.module[module_id].world_model.initial_h
              ),
              # Prediction losses.
              # Decoder (obs) loss.
              "WORLD_MODEL_L_decoder": prediction_losses["L_decoder"],
              # Reward loss.
              "WORLD_MODEL_L_reward": prediction_losses["L_reward"],
              # Continue loss.
              "WORLD_MODEL_L_continue": prediction_losses["L_continue"],
              # Total.
              "WORLD_MODEL_L_prediction": prediction_losses[
                  "L_prediction"
              ],
              # Dynamics loss.
              "WORLD_MODEL_L_dynamics": L_dyn,
              # Representation loss.
              "WORLD_MODEL_L_representation": L_rep,
              # Total loss.
              "WORLD_MODEL_L_total": L_world_model_total,
          },
      )
      if config.report_individual_batch_item_stats:
          self.register_metrics(
              module_id=module_id,
              metrics_dict={
                  "WORLD_MODEL_L_decoder_B_T": prediction_losses[
                      "L_decoder_B_T"
                  ],
                  "WORLD_MODEL_L_reward_B_T": prediction_losses[
                      "L_reward_B_T"
                  ],
                  "WORLD_MODEL_L_continue_B_T": prediction_losses[
                      "L_continue_B_T"
                  ],
                  "WORLD_MODEL_L_prediction_B_T": (
                      prediction_losses["L_prediction_B_T"]
                  ),
                  "WORLD_MODEL_L_dynamics_B_T": L_dyn_B_T,
                  "WORLD_MODEL_L_representation_B_T": L_rep_B_T,
                  "WORLD_MODEL_L_total_B_T": L_world_model_total_B_T,
              },
          )
  
      # Dream trajectories starting in all internal states (h + z_posterior) that were
      # computed during world model training.
      # Everything goes in as BxT: We are starting a new dream trajectory at every
      # actually encountered timestep in the batch, so we are creating B*T
      # trajectories of len `horizon_H`.
      dream_data = self.module[
          module_id
      ].dreamer_model.dream_trajectory(
          start_states={
              "h": fwd_out["h_states_BxT"],
              "z": fwd_out["z_posterior_states_BxT"],
          },
          start_is_terminated=tf.reshape(
              batch["is_terminated"], [-1]
          ),  # ->BxT
      )
      if config.report_dream_data:
          # To reduce this massive mount of data a little, slice out a T=1 piece
          # from each stats that has the shape (H, BxT), meaning convert e.g.
          # `rewards_dreamed_t0_to_H_BxT` into `rewards_dreamed_t0_to_H_Bx1`.
          # This will reduce the amount of data to be transferred and reported
          # by the factor of `batch_length_T`.
          self.register_metrics(
              module_id,
              {
                  # Replace 'T' with '1'.
                  "DREAM_DATA_"
                  + key[:-1]
                  + "1": (value[:, config.batch_size_B_per_learner])
                  for key, value in dream_data.items()
                  if key.endswith("H_BxT")
              },
          )
  
      value_targets_t0_to_Hm1_BxT = self._compute_value_targets(
          config=config,
          # Learn critic in symlog'd space.
          rewards_t0_to_H_BxT=dream_data["rewards_dreamed_t0_to_H_BxT"],
          intrinsic_rewards_t1_to_H_BxT=(
              dream_data["rewards_intrinsic_t1_to_H_B"]
              if config.use_curiosity
              else None
          ),
          continues_t0_to_H_BxT=dream_data[
              "continues_dreamed_t0_to_H_BxT"
          ],
          value_predictions_t0_to_H_BxT=dream_data[
              "values_dreamed_t0_to_H_BxT"
          ],
      )
      self.register_metric(
          module_id, "VALUE_TARGETS_H_BxT", value_targets_t0_to_Hm1_BxT
      )
  
      CRITIC_L_total = self._compute_critic_loss(
          module_id=module_id,
          config=config,
          dream_data=dream_data,
          value_targets_t0_to_Hm1_BxT=value_targets_t0_to_Hm1_BxT,
      )
      if config.train_actor:
          ACTOR_L_total = self._compute_actor_loss(
              module_id=module_id,
              config=config,
              dream_data=dream_data,
              value_targets_t0_to_Hm1_BxT=value_targets_t0_to_Hm1_BxT,
          )
      else:
          ACTOR_L_total = 0.0
  
      # Return the total loss as a sum of all individual losses.
      return L_world_model_total + CRITIC_L_total + ACTOR_L_total
  
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 4 for calling the DreamerV3TfLearner obj's underlying member method code with example usage: obj.compute_loss_for_module(module_id, config, batch, fwd_out) and returns: unknown