code: def cql_stats(
      policy: Policy, train_batch: SampleBatch
  ) -> Dict[str, TensorType]:
      # Get SAC loss stats.
      stats_dict = stats(policy, train_batch)
  
      # Add CQL loss stats to the dict.
      stats_dict["cql_loss"] = torch.mean(
          torch.stack(*policy.get_tower_stats("cql_loss"))
      )
  
      if policy.config["lagrangian"]:
          stats_dict["log_alpha_prime_value"] = torch.mean(
              torch.stack(
                  policy.get_tower_stats("log_alpha_prime_value")
              )
          )
          stats_dict["alpha_prime_value"] = torch.mean(
              torch.stack(policy.get_tower_stats("alpha_prime_value"))
          )
          stats_dict["alpha_prime_loss"] = torch.mean(
              torch.stack(policy.get_tower_stats("alpha_prime_loss"))
          )
      return stats_dict
  
imports: import numpy
  import gymnasium
  import logging
  import tree
  from typing import Dict, List, Tuple, Type, Union
  import ray
  import ray.experimental.tf_utils
  from ray.rllib.algorithms.sac.sac_tf_policy import postprocess_trajectory, validate_spaces
  from ray.rllib.algorithms.sac.sac_torch_policy import _get_dist_class, stats, build_sac_model_and_action_dist, optimizer_fn, ComputeTDErrorMixin, setup_late_mixins, action_distribution_fn
  from ray.rllib.models.torch.torch_action_dist import TorchDistributionWrapper
  from ray.rllib.models.modelv2 import ModelV2
  from ray.rllib.policy.policy_template import build_policy_class
  from ray.rllib.policy.policy import Policy
  from ray.rllib.policy.torch_mixins import TargetNetworkMixin
  from ray.rllib.policy.sample_batch import SampleBatch
  from ray.rllib.utils.framework import try_import_torch
  from ray.rllib.utils.metrics.learner_info import LEARNER_STATS_KEY
  from ray.rllib.utils.typing import LocalOptimizer, TensorType, AlgorithmConfigDict
  from ray.rllib.utils.torch_utils import apply_grad_clipping, convert_to_torch_tensor, concat_multi_gpu_td_errors
label_desc: Use image node_id 5 for calling a global function with example usage: cql_stats(policy, train_batch) and returns: stats_dict