code: def build_slateq_losses(
      policy: Policy,
      model: ModelV2,
      _,
      train_batch: SampleBatch,
  ) -> TensorType:
      """Constructs the choice- and Q-value losses for the SlateQTorchPolicy.
  
      Args:
          policy: The Policy to calculate the loss for.
          model: The Model to calculate the loss for.
          train_batch: The training data.
  
      Returns:
          The Q-value loss tensor.
      """
  
      # B=batch size
      # S=slate size
      # C=num candidates
      # E=embedding size
      # A=number of all possible slates
  
      # Q-value computations.
      # ---------------------
      observation = train_batch[SampleBatch.OBS]
      # user.shape: [B, E]
      user_obs = observation["user"]
      batch_size = tf.shape(user_obs)[0]
      # doc.shape: [B, C, E]
      doc_obs = list(observation["doc"].values())
      # action.shape: [B, S]
      actions = train_batch[SampleBatch.ACTIONS]
  
      # click_indicator.shape: [B, S]
      click_indicator = tf.cast(
          tf.stack([k["click"] for k in observation["response"]], 1),
          tf.float32,
      )
      # item_reward.shape: [B, S]
      item_reward = tf.stack(
          [k["watch_time"] for k in observation["response"]], 1
      )
      # q_values.shape: [B, C]
      q_values = model.get_q_values(user_obs, doc_obs)
      # slate_q_values.shape: [B, S]
      slate_q_values = tf.gather(
          q_values, tf.cast(actions, dtype=tf.int32), batch_dims=-1
      )
      # Only get the Q from the clicked document.
      # replay_click_q.shape: [B]
      replay_click_q = tf.reduce_sum(
          input_tensor=slate_q_values * click_indicator,
          axis=1,
          name="replay_click_q",
      )
  
      # Target computations.
      # --------------------
      next_obs = train_batch[SampleBatch.NEXT_OBS]
  
      # user.shape: [B, E]
      user_next_obs = next_obs["user"]
      # doc.shape: [B, C, E]
      doc_next_obs = list(next_obs["doc"].values())
      # Only compute the watch time reward of the clicked item.
      reward = tf.reduce_sum(
          input_tensor=item_reward * click_indicator, axis=1
      )
  
      # TODO: Find out, whether it's correct here to use obs, not next_obs!
      # Dopamine uses obs, then next_obs only for the score.
      # next_q_values = policy.target_model.get_q_values(user_next_obs, doc_next_obs)
      next_q_values = policy.target_model.get_q_values(
          user_obs, doc_obs
      )
      scores, score_no_click = score_documents(
          user_next_obs, doc_next_obs
      )
  
      # next_q_values_slate.shape: [B, A, S]
      next_q_values_slate = tf.gather(
          next_q_values, policy.slates, axis=1
      )
      # scores_slate.shape [B, A, S]
      scores_slate = tf.gather(scores, policy.slates, axis=1)
      # score_no_click_slate.shape: [B, A]
      score_no_click_slate = tf.reshape(
          tf.tile(score_no_click, tf.shape(input=policy.slates)[:1]),
          [batch_size, -1],
      )
  
      # next_q_target_slate.shape: [B, A]
      next_q_target_slate = tf.reduce_sum(
          input_tensor=next_q_values_slate * scores_slate, axis=2
      ) / (
          tf.reduce_sum(input_tensor=scores_slate, axis=2)
          + score_no_click_slate
      )
      next_q_target_max = tf.reduce_max(
          input_tensor=next_q_target_slate, axis=1
      )
  
      target = reward + policy.config["gamma"] * next_q_target_max * (
          1.0
          - tf.cast(train_batch[SampleBatch.TERMINATEDS], tf.float32)
      )
      target = tf.stop_gradient(target)
  
      clicked = tf.reduce_sum(input_tensor=click_indicator, axis=1)
      clicked_indices = tf.squeeze(
          tf.where(tf.equal(clicked, 1)), axis=1
      )
      # Clicked_indices is a vector and tf.gather selects the batch dimension.
      q_clicked = tf.gather(replay_click_q, clicked_indices)
      target_clicked = tf.gather(target, clicked_indices)
  
      td_error = tf.where(
          tf.cast(clicked, tf.bool),
          replay_click_q - target,
          tf.zeros_like(train_batch[SampleBatch.REWARDS]),
      )
      if policy.config["use_huber"]:
          loss = huber_loss(
              td_error, delta=policy.config["huber_threshold"]
          )
      else:
          loss = tf.math.square(td_error)
      loss = tf.reduce_mean(loss)
      td_error = tf.abs(td_error)
      mean_td_error = tf.reduce_mean(td_error)
  
      policy._q_values = tf.reduce_mean(q_values)
      policy._q_clicked = tf.reduce_mean(q_clicked)
      policy._scores = tf.reduce_mean(scores)
      policy._score_no_click = tf.reduce_mean(score_no_click)
      policy._slate_q_values = tf.reduce_mean(slate_q_values)
      policy._replay_click_q = tf.reduce_mean(replay_click_q)
      policy._bellman_reward = tf.reduce_mean(reward)
      policy._next_q_values = tf.reduce_mean(next_q_values)
      policy._target = tf.reduce_mean(target)
      policy._next_q_target_slate = tf.reduce_mean(next_q_target_slate)
      policy._next_q_target_max = tf.reduce_mean(next_q_target_max)
      policy._target_clicked = tf.reduce_mean(target_clicked)
      policy._q_loss = loss
      policy._td_error = td_error
      policy._mean_td_error = mean_td_error
      policy._mean_actions = tf.reduce_mean(actions)
  
      return loss
  
imports: import functools
  import logging
  from typing import Dict
  import gymnasium
  import numpy
  from rllib_slate_q.slate_q.slateq_tf_model import SlateQTFModel
  from ray.rllib.algorithms.dqn.dqn_tf_policy import clip_gradients
  from ray.rllib.models.modelv2 import ModelV2
  from ray.rllib.models.tf.tf_action_dist import SlateMultiCategorical
  from ray.rllib.policy.policy import Policy
  from ray.rllib.policy.sample_batch import SampleBatch
  from ray.rllib.policy.tf_mixins import LearningRateSchedule, TargetNetworkMixin
  from ray.rllib.policy.tf_policy_template import build_tf_policy
  from ray.rllib.utils.framework import try_import_tf
  from ray.rllib.utils.tf_utils import huber_loss
  from ray.rllib.utils.typing import AlgorithmConfigDict, TensorType
label_desc: Use image node_id 2 for calling a global function with example usage: build_slateq_losses(policy, model, _, train_batch) and returns: loss