code: def __init__(
      self,
      q_t_selected: TensorType,
      q_logits_t_selected: TensorType,
      q_tp1_best: TensorType,
      q_probs_tp1_best: TensorType,
      importance_weights: TensorType,
      rewards: TensorType,
      done_mask: TensorType,
      gamma=0.99,
      n_step=1,
      num_atoms=1,
      v_min=-10.0,
      v_max=10.0,
      loss_fn=huber_loss,
  ):
      if num_atoms > 1:
          # Distributional Q-learning which corresponds to an entropy loss
          z = torch.arange(0.0, num_atoms, dtype=torch.float32).to(
              rewards.device
          )
          z = v_min + z * (v_max - v_min) / float(num_atoms - 1)
  
          # (batch_size, 1) * (1, num_atoms) = (batch_size, num_atoms)
          r_tau = torch.unsqueeze(
              rewards, -1
          ) + gamma**n_step * torch.unsqueeze(
              1.0 - done_mask, -1
          ) * torch.unsqueeze(
              z, 0
          )
          r_tau = torch.clamp(r_tau, v_min, v_max)
          b = (r_tau - v_min) / ((v_max - v_min) / float(num_atoms - 1))
          lb = torch.floor(b)
          ub = torch.ceil(b)
  
          # Indispensable judgement which is missed in most implementations
          # when b happens to be an integer, lb == ub, so pr_j(s', a*) will
          # be discarded because (ub-b) == (b-lb) == 0.
          floor_equal_ceil = ((ub - lb) < 0.5).float()
  
          # (batch_size, num_atoms, num_atoms)
          l_project = F.one_hot(lb.long(), num_atoms)
          # (batch_size, num_atoms, num_atoms)
          u_project = F.one_hot(ub.long(), num_atoms)
          ml_delta = q_probs_tp1_best * (ub - b + floor_equal_ceil)
          mu_delta = q_probs_tp1_best * (b - lb)
          ml_delta = torch.sum(
              l_project * torch.unsqueeze(ml_delta, -1), dim=1
          )
          mu_delta = torch.sum(
              u_project * torch.unsqueeze(mu_delta, -1), dim=1
          )
          m = ml_delta + mu_delta
  
          # Rainbow paper claims that using this cross entropy loss for
          # priority is robust and insensitive to `prioritized_replay_alpha`
          self.td_error = softmax_cross_entropy_with_logits(
              logits=q_logits_t_selected, labels=m.detach()
          )
          self.loss = torch.mean(self.td_error * importance_weights)
          self.stats = {
              # TODO: better Q stats for dist dqn
          }
      else:
          q_tp1_best_masked = (1.0 - done_mask) * q_tp1_best
  
          # compute RHS of bellman equation
          q_t_selected_target = (
              rewards + gamma**n_step * q_tp1_best_masked
          )
  
          # compute the error (potentially clipped)
          self.td_error = q_t_selected - q_t_selected_target.detach()
          self.loss = torch.mean(
              importance_weights.float() * loss_fn(self.td_error)
          )
          self.stats = {
              "mean_q": torch.mean(q_t_selected),
              "min_q": torch.min(q_t_selected),
              "max_q": torch.max(q_t_selected),
          }
  
imports: from typing import Dict, List, Tuple
  import gymnasium
  import ray
  from ray.rllib.algorithms.dqn.dqn_tf_policy import PRIO_WEIGHTS, Q_SCOPE, Q_TARGET_SCOPE, postprocess_nstep_and_prio
  from ray.rllib.algorithms.dqn.dqn_torch_model import DQNTorchModel
  from ray.rllib.models.catalog import ModelCatalog
  from ray.rllib.models.modelv2 import ModelV2
  from ray.rllib.models.torch.torch_action_dist import get_torch_categorical_class_with_temperature, TorchDistributionWrapper
  from ray.rllib.policy.policy import Policy
  from ray.rllib.policy.policy_template import build_policy_class
  from ray.rllib.policy.sample_batch import SampleBatch
  from ray.rllib.policy.torch_mixins import LearningRateSchedule, TargetNetworkMixin
  from ray.rllib.utils.error import UnsupportedSpaceException
  from ray.rllib.utils.exploration.parameter_noise import ParameterNoise
  from ray.rllib.utils.framework import try_import_torch
  from ray.rllib.utils.torch_utils import apply_grad_clipping, concat_multi_gpu_td_errors, FLOAT_MIN, huber_loss, l2_loss, reduce_mean_ignore_inf, softmax_cross_entropy_with_logits
  from ray.rllib.utils.typing import TensorType, AlgorithmConfigDict
label_desc: Use image node_id 1 to create a new QLoss object  with example: obj = QLoss(q_t_selected, q_logits_t_selected, q_tp1_best, q_probs_tp1_best, importance_weights, rewards, done_mask, gamma, n_step, num_atoms, v_min, v_max, loss_fn)