code: def categorical_dist_qlearning(
      atoms_tm1,
      logits_q_tm1,
      a_tm1,
      r_t,
      pcont_t,
      atoms_t,
      logits_q_t,
      name="CategoricalDistQLearning",
  ):
      """Implements Distributional Q-learning as TensorFlow ops.
  
      The function assumes categorical value distributions parameterized by logits.
  
      See "A Distributional Perspective on Reinforcement Learning" by Bellemare,
      Dabney and Munos. (https://arxiv.org/abs/1707.06887).
  
      Args:
        atoms_tm1: 1-D tensor containing atom values for first timestep,
          shape `[num_atoms]`.
        logits_q_tm1: Tensor holding logits for first timestep in a batch of
          transitions, shape `[B, num_actions, num_atoms]`.
        a_tm1: Tensor holding action indices, shape `[B]`.
        r_t: Tensor holding rewards, shape `[B]`.
        pcont_t: Tensor holding pcontinue values, shape `[B]`.
        atoms_t: 1-D tensor containing atom values for second timestep,
          shape `[num_atoms]`.
        logits_q_t: Tensor holding logits for second timestep in a batch of
          transitions, shape `[B, num_actions, num_atoms]`.
        name: name to prefix ops created by this function.
  
      Returns:
        A namedtuple with fields:
  
        * `loss`: a tensor containing the batch of losses, shape `[B]`.
        * `extra`: a namedtuple with fields:
            * `target`: a tensor containing the values that `q_tm1` at actions
            `a_tm1` are regressed towards, shape `[B, num_atoms]`.
  
      Raises:
        ValueError: If the tensors do not have the correct rank or compatibility.
      """
      # Rank and compatibility checks.
      assertion_lists = [
          [logits_q_tm1, logits_q_t],
          [a_tm1, r_t, pcont_t],
          [atoms_tm1, atoms_t],
      ]
      base_ops.wrap_rank_shape_assert(assertion_lists, [3, 1, 1], name)
  
      # Categorical distributional Q-learning op.
      with tf.name_scope(
          name,
          values=[
              atoms_tm1,
              logits_q_tm1,
              a_tm1,
              r_t,
              pcont_t,
              atoms_t,
              logits_q_t,
          ],
      ):
          with tf.name_scope("target"):
              # Scale and shift time-t distribution atoms by discount and reward.
              target_z = (
                  r_t[:, None] + pcont_t[:, None] * atoms_t[None, :]
              )
  
              # Convert logits to distribution, then find greedy action in state s_t.
              q_t_probs = tf.nn.softmax(logits_q_t)
              q_t_mean = tf.reduce_sum(q_t_probs * atoms_t, 2)
              pi_t = tf.argmax(q_t_mean, 1, output_type=tf.int32)
  
              # Compute distribution for greedy action.
              p_target_z = _slice_with_actions(q_t_probs, pi_t)
  
              # Project using the Cramer distance
              target = tf.stop_gradient(
                  _l2_project(target_z, p_target_z, atoms_tm1)
              )
  
          logit_qa_tm1 = _slice_with_actions(logits_q_tm1, a_tm1)
  
          loss = tf.nn.softmax_cross_entropy_with_logits(
              logits=logit_qa_tm1, labels=target
          )
  
          return base_ops.LossOutput(loss, Extra(target))
  
file_path: trfl/trfl/dist_value_ops.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import collections
  import tensorflow.compat.v1
  from trfl import base_ops
  from trfl import distribution_ops
label_desc: Use image node_id 2 for calling a global function with example usage: categorical_dist_qlearning(atoms_tm1, logits_q_tm1, a_tm1, r_t, pcont_t, atoms_t, logits_q_t, name) and returns: base_ops