code: def categorical_double_q_learning(
      q_atoms_tm1: Array,
      q_logits_tm1: Array,
      a_tm1: Numeric,
      r_t: Numeric,
      discount_t: Numeric,
      q_atoms_t: Array,
      q_logits_t: Array,
      q_t_selector: Array,
      stop_target_gradients: bool = True,
  ) -> Numeric:
      """Implements double Q-learning for categorical Q distributions.
  
      See "A Distributional Perspective on Reinforcement Learning", by
      Bellemere, Dabney and Munos (https://arxiv.org/pdf/1707.06887.pdf)
      and "Double Q-learning" by van Hasselt.
      (https://papers.nips.cc/paper/3964-double-q-learning.pdf).
  
      Args:
        q_atoms_tm1: atoms of Q distribution at time t-1.
        q_logits_tm1: logits of Q distribution at time t-1.
        a_tm1: action index at time t-1.
        r_t: reward at time t.
        discount_t: discount at time t.
        q_atoms_t: atoms of Q distribution at time t.
        q_logits_t: logits of Q distribution at time t.
        q_t_selector: selector Q-values at time t.
        stop_target_gradients: bool indicating whether or not to apply stop gradient
          to targets.
  
      Returns:
        Categorical double Q-learning loss (i.e. temporal difference error).
      """
      chex.assert_rank(
          [
              q_atoms_tm1,
              q_logits_tm1,
              a_tm1,
              r_t,
              discount_t,
              q_atoms_t,
              q_logits_t,
              q_t_selector,
          ],
          [1, 2, 0, 0, 0, 1, 2, 1],
      )
      chex.assert_type(
          [
              q_atoms_tm1,
              q_logits_tm1,
              a_tm1,
              r_t,
              discount_t,
              q_atoms_t,
              q_logits_t,
              q_t_selector,
          ],
          [float, float, int, float, float, float, float, float],
      )
  
      # Scale and shift time-t distribution atoms by discount and reward.
      target_z = r_t + discount_t * q_atoms_t
  
      # Select logits for greedy action in state s_t and convert to distribution.
      p_target_z = jax.nn.softmax(q_logits_t[q_t_selector.argmax()])
  
      # Project using the Cramer distance and maybe stop gradient flow to targets.
      target = categorical_l2_project(target_z, p_target_z, q_atoms_tm1)
      target = jax.lax.select(
          stop_target_gradients, jax.lax.stop_gradient(target), target
      )
  
      # Compute loss (i.e. temporal difference error).
      logit_qa_tm1 = q_logits_tm1[a_tm1]
      return distributions.categorical_cross_entropy(
          labels=target, logits=logit_qa_tm1
      )
  
imports: from typing import Union
  import chex
  import jax
  import jax.numpy
  from rlax._src import base
  from rlax._src import clipping
  from rlax._src import distributions
  from rlax._src import multistep
label_desc: Use image node_id 17 for calling a global function with example usage: categorical_double_q_learning(q_atoms_tm1, q_logits_tm1, a_tm1, r_t, discount_t, q_atoms_t, q_logits_t, q_t_selector, stop_target_gradients) and returns: distributions