class_docstr: Deep q learning.
  
  This matches the original DQN loss: https://arxiv.org/abs/1312.5602.
  It differs by two aspects that improve it on the optimization side
    - it uses Adam instead of RMSProp as an optimizer
    - it uses a square loss instead of the Huber one.
class_docstr_tok: ["Deep","q","learning",".","This","matches","the","original","DQN","loss",":","https",":","\/\/arxiv.org\/abs\/1312.5602",".","It","differs","by","two","aspects","that","improve","it","on","the","optimization","side","-","it","uses","Adam","instead","of","RMSProp","as","an","optimizer","-","it","uses","a","square","loss","instead","of","the","Huber","one","."]
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
  
      # Forward pass.
      key1, key2 = jax.random.split(key)
      q_tm1 = network.apply(
          params, transitions.observation, is_training=True, key=key1
      )
      q_t = network.apply(
          target_params,
          transitions.next_observation,
          is_training=True,
          key=key2,
      )
  
      # 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 Q-learning TD-error.
      batch_error = jax.vmap(rlax.q_learning)
      td_error = batch_error(q_tm1, transitions.action, r_t, d_t, q_t)
      batch_loss = jnp.square(td_error)
  
      loss = jnp.mean(batch_loss)
      extra = learning_lib.LossExtra(metrics={})
      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 QLearning obj's underlying member method code with example usage: obj.__call__(network, params, target_params, batch, key) and returns: loss,  extra