code: def quantile_q_learning(
      dist_q_tm1: Array,
      tau_q_tm1: Array,
      a_tm1: Numeric,
      r_t: Numeric,
      discount_t: Numeric,
      dist_q_t_selector: Array,
      dist_q_t: Array,
      huber_param: float = 0.0,
      stop_target_gradients: bool = True,
  ) -> Numeric:
      """Implements Q-learning for quantile-valued Q distributions.
  
      See "Distributional Reinforcement Learning with Quantile Regression" by
      Dabney et al. (https://arxiv.org/abs/1710.10044).
  
      Args:
        dist_q_tm1: Q distribution at time t-1.
        tau_q_tm1: Q distribution probability thresholds.
        a_tm1: action index at time t-1.
        r_t: reward at time t.
        discount_t: discount at time t.
        dist_q_t_selector: Q distribution at time t for selecting greedy action in
          target policy. This is separate from dist_q_t as in Double Q-Learning, but
          can be computed with the target network and a separate set of samples.
        dist_q_t: target Q distribution at time t.
        huber_param: Huber loss parameter, defaults to 0 (no Huber loss).
        stop_target_gradients: bool indicating whether or not to apply stop gradient
          to targets.
  
      Returns:
        Quantile regression Q learning loss.
      """
      chex.assert_rank(
          [
              dist_q_tm1,
              tau_q_tm1,
              a_tm1,
              r_t,
              discount_t,
              dist_q_t_selector,
              dist_q_t,
          ],
          [2, 1, 0, 0, 0, 2, 2],
      )
      chex.assert_type(
          [
              dist_q_tm1,
              tau_q_tm1,
              a_tm1,
              r_t,
              discount_t,
              dist_q_t_selector,
              dist_q_t,
          ],
          [float, float, int, float, float, float, float],
      )
  
      # Only update the taken actions.
      dist_qa_tm1 = dist_q_tm1[:, a_tm1]
  
      # Select target action according to greedy policy w.r.t. dist_q_t_selector.
      q_t_selector = jnp.mean(dist_q_t_selector, axis=0)
      a_t = jnp.argmax(q_t_selector)
      dist_qa_t = dist_q_t[:, a_t]
  
      # Compute target, do not backpropagate into it.
      dist_target = r_t + discount_t * dist_qa_t
      dist_target = jax.lax.select(
          stop_target_gradients,
          jax.lax.stop_gradient(dist_target),
          dist_target,
      )
  
      return quantile_regression_loss(
          dist_qa_tm1, tau_q_tm1, dist_target, huber_param
      )
  
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 19 for calling a global function with example usage: quantile_q_learning(dist_q_tm1, tau_q_tm1, a_tm1, r_t, discount_t, dist_q_t_selector, dist_q_t, huber_param, stop_target_gradients) and returns: quantile_regression_loss