class_docstr: MPO loss for a categorical policy (Abdolmaleki et al., 2018).
  
  (Abdolmaleki et al., 2018): https://arxiv.org/pdf/1812.02256.pdf
class_docstr_tok: ["MPO","loss","for","a","categorical","policy","(","Abdolmaleki","et","al.",",","2018",")",".","(","Abdolmaleki","et","al.",",","2018",")",":","https",":","\/\/arxiv.org\/pdf\/1812.02256.pdf"]
docstr: Computes the MPO loss for a categorical policy.
  
  Args:
    params: parameters tracking the temperature and the dual variables.
    online_action_distribution: online distribution returned by the online
      policy network; expects batch_dims of [B] and event_dims of [D].
    target_action_distribution: target distribution returned by the target
      policy network; expects same shapes as online distribution.
    actions: Unused.
    q_values: Q-values associated with every action; expects shape [D, B].
  
  Returns:
    Loss, combining the policy loss, KL penalty, and dual losses required to
      adapt the dual variables.
    Stats, for diagnostics and tracking performance.
docstr_tok: ["Computes","the","MPO","loss","for","a","categorical","policy",".","Args",":","params",":","parameters","tracking","the","temperature","and","the","dual","variables",".","online_action_distribution",":","online","distribution","returned","by","the","online","policy","network",";","expects","batch_dims","of","[","B","]","and","event_dims","of","[","D","]",".","target_action_distribution",":","target","distribution","returned","by","the","target","policy","network",";","expects","same","shapes","as","online","distribution",".","actions",":","Unused",".","q_values",":","Q-values","associated","with","every","action",";","expects","shape","[","D",",","B","]",".","Returns",":","Loss",",","combining","the","policy","loss",",","KL","penalty",",","and","dual","losses","required","to","adapt","the","dual","variables",".","Stats",",","for","diagnostics","and","tracking","performance","."]
code: def __call__(
      self,
      params: CategoricalMPOParams,
      online_action_distribution: distrax.Categorical,
      target_action_distribution: distrax.Categorical,
      actions: jnp.ndarray,  # Unused.
      q_values: jnp.ndarray,  # Shape [D, B].
  ) -> Tuple[jnp.ndarray, CategoricalMPOStats]:
      """Computes the MPO loss for a categorical policy.
  
      Args:
        params: parameters tracking the temperature and the dual variables.
        online_action_distribution: online distribution returned by the online
          policy network; expects batch_dims of [B] and event_dims of [D].
        target_action_distribution: target distribution returned by the target
          policy network; expects same shapes as online distribution.
        actions: Unused.
        q_values: Q-values associated with every action; expects shape [D, B].
  
      Returns:
        Loss, combining the policy loss, KL penalty, and dual losses required to
          adapt the dual variables.
        Stats, for diagnostics and tracking performance.
      """
  
      q_values = jnp.transpose(q_values)  # [D, B] --> [B, D].
  
      # Transform dual variables from log-space.
      # Note: using softplus instead of exponential for numerical stability.
      temperature = get_temperature_from_params(params)
      alpha = jax.nn.softplus(params.log_alpha) + _MPO_FLOAT_EPSILON
  
      # Compute the E-step logits and the temperature loss, used to adapt the
      # tempering of Q-values.
      (
          logits_e_step,
          loss_temperature,
      ) = compute_weights_and_temperature_loss(  # pytype: disable=wrong-arg-types  # jax-ndarray
          q_values=q_values,
          logits=target_action_distribution.logits,
          epsilon=self._epsilon,
          temperature=temperature,
      )
      action_distribution_e_step = distrax.Categorical(
          logits=logits_e_step
      )
  
      # Only needed for diagnostics: Compute estimated actualized KL between the
      # non-parametric and current target policies.
      kl_nonparametric = action_distribution_e_step.kl_divergence(
          target_action_distribution
      )
  
      # Compute the policy loss.
      loss_policy = action_distribution_e_step.cross_entropy(
          online_action_distribution
      )
      loss_policy = jnp.mean(loss_policy)
  
      # Compute the regularization.
      kl = target_action_distribution.kl_divergence(
          online_action_distribution
      )
      mean_kl = jnp.mean(kl, axis=0)
      loss_kl = jax.lax.stop_gradient(alpha) * mean_kl
  
      # Compute the dual loss.
      loss_alpha = alpha * (
          self._epsilon_policy - jax.lax.stop_gradient(mean_kl)
      )
  
      # Combine losses.
      loss_dual = loss_alpha + loss_temperature
      loss = loss_policy + loss_kl + loss_dual
  
      # Create statistics.
      stats = CategoricalMPOStats(  # pytype: disable=wrong-arg-types  # jnp-type
          # Dual Variables.
          dual_alpha=jnp.mean(alpha),
          dual_temperature=jnp.mean(temperature),
          # Losses.
          loss_e_step=loss_policy,
          loss_m_step=loss_kl,
          loss_dual=loss_dual,
          loss_policy=jnp.mean(loss),
          loss_alpha=jnp.mean(loss_alpha),
          loss_temperature=jnp.mean(loss_temperature),
          # KL measurements.
          kl_q_rel=jnp.mean(kl_nonparametric) / self._epsilon,
          kl_mean_rel=mean_kl / self._epsilon_policy,
          # Q measurements.
          q_min=jnp.mean(jnp.min(q_values, axis=0)),
          q_max=jnp.mean(jnp.max(q_values, axis=0)),
          entropy_online=jnp.mean(online_action_distribution.entropy()),
          entropy_target=jnp.mean(target_action_distribution.entropy()),
      )
  
      return loss, stats
  
file_path: acme/acme/agents/jax/mpo/categorical_mpo.py
imports: from typing import NamedTuple, Tuple
  import distrax
  import jax
  import jax.numpy
label_desc: Use image node_id 3 for calling the CategoricalMPO obj's underlying member method code with example usage: obj.__call__(params, online_action_distribution, target_action_distribution, actions, q_values) and returns: loss,  stats