class_docstr: MPO learner (discrete or continuous, distributional or not).
class_docstr_tok: ["MPO","learner","(","discrete","or","continuous",",","distributional","or","not",")","."]
code: def _loss_fn(
      self,
      params: mpo_networks.MPONetworkParams,
      dual_params: mpo_types.DualParams,
      # TODO(bshahr): clean up types: Step is not a great type for sequences.
      sequence: adders.Step,
      target_params: mpo_networks.MPONetworkParams,
      key: jax_types.PRNGKey,
  ) -> Tuple[jnp.ndarray, mpo_types.LogDict]:
      # Compute the model predictions at the root and for the rollouts.
      predictions = self._compute_predictions(
          params=params, sequence=sequence
      )
  
      # Compute the targets to use for the losses.
      targets = self._compute_targets(
          target_params=target_params,
          dual_params=dual_params,
          sequence=sequence,
          online_policy=predictions.policy,
          key=key,
      )
  
      # TODO(abef): mask policy loss at terminal states or use uniform targets
      # is_terminal = sequence.discount == 0.
  
      # Compute MPO policy loss on each state in the sequence.
      policy_loss, policy_stats = self._policy_loss_module(
          params=dual_params,
          online_action_distribution=predictions.policy,  # [T, ...].
          target_action_distribution=targets.policy,  # [T, ...].
          actions=targets.a_improvement,  # Unused in discrete case; [N, T].
          q_values=targets.q_improvement,
      )  # [N, T]
  
      # Compute the critic loss on the states in the sequence.
      critic_loss = self._distributional_loss(
          prediction=predictions.q_value,  # [T-1, 1, ...]
          target=targets.q_value,
      )  # [N, Z, T-1]
  
      loss = (
          self._loss_scales.policy * policy_loss
          + self._loss_scales.critic * critic_loss
      )
      loss_logging_dict = {
          "loss": loss,
          "root_policy_loss": policy_loss,
          "root_critic_loss": critic_loss,
          "policy_loss": policy_loss,
          "critic_loss": critic_loss,
      }
  
      # Append MPO statistics.
      loss_logging_dict.update(
          {
              f"policy/root/{k}": v
              for k, v in policy_stats._asdict().items()
          }
      )
  
      # Compute rollout losses.
      if self._model_rollout_length > 0:
          (
              model_rollout_loss,
              rollout_logs,
          ) = self._model_rollout_loss_fn(
              params,
              dual_params,
              sequence,
              predictions.embedding,
              targets,
              key,
          )
          loss += model_rollout_loss
          loss_logging_dict.update(rollout_logs)
          loss_logging_dict.update(
              {
                  "policy_loss": policy_loss
                  + rollout_logs["rollout_policy_loss"],
                  "critic_loss": critic_loss
                  + rollout_logs["rollout_critic_loss"],
                  "loss": loss,
              }
          )
  
      return loss, loss_logging_dict
  
imports: import dataclasses
  import functools
  import time
  from typing import Any, Dict, Iterator, List, NamedTuple, Optional, Sequence, Tuple, Union
  from absl import logging
  import acme
  from acme import specs
  from acme import types
  from acme.adders import reverb
  from acme.agents.jax.mpo import categorical_mpo
  from acme.agents.jax.mpo import networks
  from acme.agents.jax.mpo import rollout_loss
  from acme.agents.jax.mpo import types
  from acme.agents.jax.mpo import utils
  from acme.jax import networks
  from acme.jax import types
  from acme.jax import utils
  import acme.jax.losses.mpo
  from acme.utils import counting
  from acme.utils import loggers
  import chex
  import jax
  import jax.numpy
  import numpy
  import optax
  import reverb
  import rlax
  import tree
label_desc: Use image node_id 5 for calling the MPOLearner obj's underlying member method code with example usage: obj._loss_fn(params, dual_params, sequence, target_params, key) and returns: loss,  loss_logging_dict