code: def quantile_expected_sarsa(
      dist_q_tm1: Array,
      tau_q_tm1: Array,
      a_tm1: Numeric,
      r_t: Numeric,
      discount_t: Numeric,
      dist_q_t: Array,
      probs_a_t: Array,
      huber_param: float = 0.0,
      stop_target_gradients: bool = True,
  ) -> Numeric:
      """Implements Expected SARSA for quantile-valued Q distributions.
  
      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: target Q distribution at time t.
        probs_a_t: action probabilities 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 Expected SARSA learning loss.
      """
      chex.assert_rank(
          [
              dist_q_tm1,
              tau_q_tm1,
              a_tm1,
              r_t,
              discount_t,
              dist_q_t,
              probs_a_t,
          ],
          [2, 1, 0, 0, 0, 2, 1],
      )
      chex.assert_type(
          [
              dist_q_tm1,
              tau_q_tm1,
              a_tm1,
              r_t,
              discount_t,
              dist_q_t,
              probs_a_t,
          ],
          [float, float, int, float, float, float, float],
      )
  
      # Only update the taken actions.
      dist_qa_tm1 = dist_q_tm1[:, a_tm1]
  
      # Compute target, do not backpropagate into it.
      dist_target = r_t + discount_t * dist_q_t
      dist_target = jax.lax.select(
          stop_target_gradients,
          jax.lax.stop_gradient(dist_target),
          dist_target,
      )
      probs_a_t = jax.lax.select(
          stop_target_gradients,
          jax.lax.stop_gradient(probs_a_t),
          probs_a_t,
      )
  
      per_action_qr = jax.vmap(
          quantile_regression_loss, in_axes=(None, None, 1, None)
      )
      per_action_loss = per_action_qr(
          dist_qa_tm1, tau_q_tm1, dist_target, huber_param
      )
  
      return jnp.dot(per_action_loss, probs_a_t)
  
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 20 for calling a global function with example usage: quantile_expected_sarsa(dist_q_tm1, tau_q_tm1, a_tm1, r_t, discount_t, dist_q_t, probs_a_t, huber_param, stop_target_gradients) and returns: jnp