class_docstr: Munchausen q learning.
  
  Implements M-DQN: https://arxiv.org/abs/2007.14430.
class_docstr_tok: ["Munchausen","q","learning",".","Implements","M-DQN",":","https",":","\/\/arxiv.org\/abs\/2007.14430","."]
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, key3 = jax.random.split(key, 3)
      q_online_s = network.apply(
          params, transitions.observation, is_training=True, key=key1
      )
      action_one_hot = jax.nn.one_hot(
          transitions.action, q_online_s.shape[-1]
      )
      q_online_sa = jnp.sum(action_one_hot * q_online_s, axis=-1)
      q_target_s = network.apply(
          target_params,
          transitions.observation,
          is_training=True,
          key=key2,
      )
      q_target_next = network.apply(
          target_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)
  
      # Munchausen term : tau * log_pi(a|s)
      munchausen_term = self.entropy_temperature * jax.nn.log_softmax(
          q_target_s / self.entropy_temperature, axis=-1
      )
      munchausen_term_a = jnp.sum(
          action_one_hot * munchausen_term, axis=-1
      )
      munchausen_term_a = jnp.clip(
          munchausen_term_a, a_min=self.clip_value_min, a_max=0.0
      )
  
      # Soft Bellman operator applied to q
      next_v = self.entropy_temperature * jax.nn.logsumexp(
          q_target_next / self.entropy_temperature, axis=-1
      )
      target_q = jax.lax.stop_gradient(
          r_t
          + self.munchausen_coefficient * munchausen_term_a
          + d_t * next_v
      )
  
      batch_loss = rlax.huber_loss(
          target_q - q_online_sa, self.huber_loss_parameter
      )
      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 MunchausenQLearning obj's underlying member method code with example usage: obj.__call__(network, params, target_params, batch, key) and returns: loss,  extra