class_docstr: Categorical double q learning with prioritization on TD error.
class_docstr_tok: ["Categorical","double","q","learning","with","prioritization","on","TD","error","."]
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
      probs = batch.info.probability
  
      # Forward pass.
      key1, key2, key3 = jax.random.split(key, 3)
      _, logits_tm1, atoms_tm1 = network.apply(
          params, transitions.observation, is_training=True, key=key1
      )
      _, logits_t, atoms_t = network.apply(
          target_params,
          transitions.next_observation,
          is_training=True,
          key=key2,
      )
      q_t_selector, _, _ = network.apply(
          params,
          transitions.next_observation,
          is_training=True,
          key=key3,
      )
  
      # Cast and clip rewards.
      d_t = (transitions.discount * self.discount).astype(jnp.float32)
      r_t = jnp.clip(
          transitions.reward, -self.max_abs_reward, self.max_abs_reward
      ).astype(jnp.float32)
  
      # Compute categorical double Q-learning loss.
      batch_loss_fn = jax.vmap(
          rlax.categorical_double_q_learning,
          in_axes=(None, 0, 0, 0, 0, None, 0, 0),
      )
      batch_loss = batch_loss_fn(
          atoms_tm1,
          logits_tm1,
          transitions.action,
          r_t,
          d_t,
          atoms_t,
          logits_t,
          q_t_selector,
      )
  
      # Importance weighting.
      importance_weights = (1.0 / probs).astype(jnp.float32)
      importance_weights **= self.importance_sampling_exponent
      importance_weights /= jnp.max(importance_weights)
  
      # Reweight.
      loss = jnp.mean(importance_weights * batch_loss)  # []
      extra = learning_lib.LossExtra(
          metrics={},
          reverb_priorities=jnp.abs(batch_loss).astype(jnp.float64),
      )
      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 PrioritizedCategoricalDoubleQLearning obj's underlying member method code with example usage: obj.__call__(network, params, target_params, batch, key) and returns: loss,  extra