class_docstr: Quantile Regression DQN.
  
  https://arxiv.org/abs/1710.10044
class_docstr_tok: ["Quantile","Regression","DQN",".","https",":","\/\/arxiv.org\/abs\/1710.10044"]
docstr: Calculate a loss on a single batch of data.
docstr_tok: ["Calculate","a","loss","on","a","single","batch","of","data","."]
code: def __call__(
      self,
      network: networks_lib.TypedFeedForwardNetwork,
      params: networks_lib.Params,
      target_params: networks_lib.Params,
      batch: reverb.ReplaySample,
      key: networks_lib.PRNGKey,
  ) -> Tuple[jax.Array, learning_lib.LossExtra]:
      """Calculate a loss on a single batch of data."""
      transitions: types.Transition = batch.data
      key1, key2 = jax.random.split(key)
      _, dist_q_tm1 = network.apply(
          params, transitions.observation, is_training=True, key=key1
      )
      _, dist_q_target_t = network.apply(
          target_params,
          transitions.next_observation,
          is_training=True,
          key=key2,
      )
      batch_size = len(transitions.observation)
      chex.assert_shape(
          dist_q_tm1,
          (
              batch_size,
              None,
              self.num_atoms,
          ),
          custom_message=f"Expected (batch_size, num_actions, num_atoms), got: {dist_q_tm1.shape}",
          include_default_message=True,
      )
      chex.assert_shape(
          dist_q_target_t,
          (
              batch_size,
              None,
              self.num_atoms,
          ),
          custom_message=f"Expected (batch_size, num_actions, num_atoms), got: {dist_q_target_t.shape}",
          include_default_message=True,
      )
      # Swap distribution and action dimension, since
      # rlax.quantile_q_learning expects it that way.
      dist_q_tm1 = jnp.swapaxes(dist_q_tm1, 1, 2)
      dist_q_target_t = jnp.swapaxes(dist_q_target_t, 1, 2)
      quantiles = (
          jnp.arange(self.num_atoms, dtype=jnp.float32) + 0.5
      ) / self.num_atoms
      batch_quantile_q_learning = jax.vmap(
          rlax.quantile_q_learning,
          in_axes=(0, None, 0, 0, 0, 0, 0, None),
      )
      losses = batch_quantile_q_learning(
          dist_q_tm1,
          quantiles,
          transitions.action,
          transitions.reward,
          transitions.discount,
          dist_q_target_t,  # No double Q-learning here.
          dist_q_target_t,
          self.huber_param,
      )
      loss = jnp.mean(losses)
      chex.assert_shape(losses, (batch_size,))
      extra = learning_lib.LossExtra(metrics={"mean_loss": loss})
      return loss, extra
  
imports: import dataclasses
  from typing import Tuple
  from acme import types
  from acme.agents.jax.dqn import learning_lib
  from acme.jax import networks
  import chex
  import jax
  import jax.numpy
  import reverb
  import rlax
label_desc: Use image node_id 1 for calling the QrDqn obj's underlying member method code with example usage: obj.__call__(network, params, target_params, batch, key) and returns: loss,  extra