class_docstr: Multi-objective MPO loss with decoupled KL constraints.
  
  This implementation of the MO-MPO loss is based on the approach proposed in
  (Abdolmaleki, Huang et al., 2020). The following features are included as
  options:
  - Satisfying the KL-constraint on a per-dimension basis (on by default)
  
  (Abdolmaleki, Huang et al., 2020): https://arxiv.org/pdf/2005.07513.pdf
class_docstr_tok: ["Multi-objective","MPO","loss","with","decoupled","KL","constraints",".","This","implementation","of","the","MO-MPO","loss","is","based","on","the","approach","proposed","in","(","Abdolmaleki",",","Huang","et","al.",",","2020",")",".","The","following","features","are","included","as","options",":","-","Satisfying","the","KL-constraint","on","a","per-dimension","basis","(","on","by","default",")","(","Abdolmaleki",",","Huang","et","al.",",","2020",")",":","https",":","\/\/arxiv.org\/pdf\/2005.07513.pdf"]
docstr: Computes the decoupled MO-MPO loss.
  
  Args:
    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: actions sampled from the target policy; expects shape [N, B, D].
    q_values: Q-values associated with each action; expects shape [N, B, K].
  
  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","decoupled","MO-MPO","loss",".","Args",":","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",":","actions","sampled","from","the","target","policy",";","expects","shape","[","N",",","B",",","D","]",".","q_values",":","Q-values","associated","with","each","action",";","expects","shape","[","N",",","B",",","K","]",".","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,
      online_action_distribution: Union[
          tfd.MultivariateNormalDiag, tfd.Independent
      ],
      target_action_distribution: Union[
          tfd.MultivariateNormalDiag, tfd.Independent
      ],
      actions: tf.Tensor,  # Shape [N, B, D].
      q_values: tf.Tensor,  # Shape [N, B, K].
  ) -> Tuple[tf.Tensor, Dict[str, tf.Tensor]]:
      """Computes the decoupled MO-MPO loss.
  
      Args:
        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: actions sampled from the target policy; expects shape [N, B, D].
        q_values: Q-values associated with each action; expects shape [N, B, K].
  
      Returns:
        Loss, combining the policy loss, KL penalty, and dual losses required to
          adapt the dual variables.
        Stats, for diagnostics and tracking performance.
      """
  
      # Make sure the Q-values are per-objective
      q_values.get_shape().assert_has_rank(3)
      if q_values.get_shape()[-1] != self._num_objectives:
          raise ValueError(
              "Q-values do not match expected number of objectives."
          )
  
      # Cast `MultivariateNormalDiag`s to Independent Normals.
      # The latter allows us to satisfy KL constraints per-dimension.
      if isinstance(
          target_action_distribution, tfd.MultivariateNormalDiag
      ):
          target_action_distribution = tfd.Independent(
              tfd.Normal(
                  target_action_distribution.mean(),
                  target_action_distribution.stddev(),
              )
          )
          online_action_distribution = tfd.Independent(
              tfd.Normal(
                  online_action_distribution.mean(),
                  online_action_distribution.stddev(),
              )
          )
  
      # Infer the shape and dtype of dual variables.
      scalar_dtype = q_values.dtype
      if self._per_dim_constraining:
          dual_variable_shape = (
              target_action_distribution.distribution.kl_divergence(
                  online_action_distribution.distribution
              ).shape[1:]
          )  # Should be [D].
      else:
          dual_variable_shape = (
              target_action_distribution.kl_divergence(
                  online_action_distribution
              ).shape[1:]
          )  # Should be [1].
  
      # Create dual variables for the KL constraints; only happens the first call.
      self.create_dual_variables_once(dual_variable_shape, scalar_dtype)
  
      # Project dual variables to ensure they stay positive.
      min_log_temperature = tf.constant(-18.0, scalar_dtype)
      min_log_alpha = tf.constant(-18.0, scalar_dtype)
      self._log_temperature.assign(
          tf.maximum(min_log_temperature, self._log_temperature)
      )
      self._log_alpha_mean.assign(
          tf.maximum(min_log_alpha, self._log_alpha_mean)
      )
      self._log_alpha_stddev.assign(
          tf.maximum(min_log_alpha, self._log_alpha_stddev)
      )
  
      # Transform dual variables from log-space.
      # Note: using softplus instead of exponential for numerical stability.
      temperature = (
          tf.math.softplus(self._log_temperature) + _MPO_FLOAT_EPSILON
      )
      alpha_mean = (
          tf.math.softplus(self._log_alpha_mean) + _MPO_FLOAT_EPSILON
      )
      alpha_stddev = (
          tf.math.softplus(self._log_alpha_stddev) + _MPO_FLOAT_EPSILON
      )
  
      # Get online and target means and stddevs in preparation for decomposition.
      online_mean = online_action_distribution.distribution.mean()
      online_scale = online_action_distribution.distribution.stddev()
      target_mean = target_action_distribution.distribution.mean()
      target_scale = target_action_distribution.distribution.stddev()
  
      # Compute normalized importance weights, used to compute expectations with
      # respect to the non-parametric policy; and the temperature loss, used to
      # adapt the tempering of Q-values.
      (
          normalized_weights,
          loss_temperature,
      ) = compute_weights_and_temperature_loss(
          q_values, self._epsilons, temperature
      )  # Shapes [N, B, K] and [1, K].
      normalized_weights_sum = tf.reduce_sum(
          normalized_weights, axis=-1
      )
      loss_temperature_mean = tf.reduce_mean(loss_temperature)
  
      # Only needed for diagnostics: Compute estimated actualized KL between the
      # non-parametric and current target policies.
      kl_nonparametric = (
          mpo.compute_nonparametric_kl_from_normalized_weights(
              normalized_weights
          )
      )
  
      # Decompose the online policy into fixed-mean & fixed-stddev distributions.
      # This has been documented as having better performance in bandit settings,
      # see e.g. https://arxiv.org/pdf/1812.02256.pdf.
      fixed_stddev_distribution = tfd.Independent(
          tfd.Normal(loc=online_mean, scale=target_scale)
      )
      fixed_mean_distribution = tfd.Independent(
          tfd.Normal(loc=target_mean, scale=online_scale)
      )
  
      # Compute the decomposed policy losses.
      loss_policy_mean = mpo.compute_cross_entropy_loss(
          actions, normalized_weights_sum, fixed_stddev_distribution
      )
      loss_policy_stddev = mpo.compute_cross_entropy_loss(
          actions, normalized_weights_sum, fixed_mean_distribution
      )
  
      # Compute the decomposed KL between the target and online policies.
      if self._per_dim_constraining:
          kl_mean = (
              target_action_distribution.distribution.kl_divergence(
                  fixed_stddev_distribution.distribution
              )
          )  # Shape [B, D].
          kl_stddev = (
              target_action_distribution.distribution.kl_divergence(
                  fixed_mean_distribution.distribution
              )
          )  # Shape [B, D].
      else:
          kl_mean = target_action_distribution.kl_divergence(
              fixed_stddev_distribution
          )  # Shape [B].
          kl_stddev = target_action_distribution.kl_divergence(
              fixed_mean_distribution
          )  # Shape [B].
  
      # Compute the alpha-weighted KL-penalty and dual losses to adapt the alphas.
      (
          loss_kl_mean,
          loss_alpha_mean,
      ) = mpo.compute_parametric_kl_penalty_and_dual_loss(
          kl_mean, alpha_mean, self._epsilon_mean
      )
      (
          loss_kl_stddev,
          loss_alpha_stddev,
      ) = mpo.compute_parametric_kl_penalty_and_dual_loss(
          kl_stddev, alpha_stddev, self._epsilon_stddev
      )
  
      # Combine losses.
      loss_policy = loss_policy_mean + loss_policy_stddev
      loss_kl_penalty = loss_kl_mean + loss_kl_stddev
      loss_dual = (
          loss_alpha_mean + loss_alpha_stddev + loss_temperature_mean
      )
      loss = loss_policy + loss_kl_penalty + loss_dual
  
      stats = {}
      # Dual Variables.
      stats["dual_alpha_mean"] = tf.reduce_mean(alpha_mean)
      stats["dual_alpha_stddev"] = tf.reduce_mean(alpha_stddev)
      # Losses.
      stats["loss_policy"] = tf.reduce_mean(loss)
      stats["loss_alpha"] = tf.reduce_mean(
          loss_alpha_mean + loss_alpha_stddev
      )
      # KL measurements.
      stats["kl_mean_rel"] = (
          tf.reduce_mean(kl_mean, axis=0) / self._epsilon_mean
      )
      stats["kl_stddev_rel"] = (
          tf.reduce_mean(kl_stddev, axis=0) / self._epsilon_stddev
      )
      # If the policy has standard deviation, log summary stats for this as well.
      pi_stddev = online_action_distribution.distribution.stddev()
      stats["pi_stddev_min"] = tf.reduce_mean(
          tf.reduce_min(pi_stddev, axis=-1)
      )
      stats["pi_stddev_max"] = tf.reduce_mean(
          tf.reduce_max(pi_stddev, axis=-1)
      )
  
      # Condition number of the diagonal covariance (actually, stddev) matrix.
      stats["pi_stddev_cond"] = tf.reduce_mean(
          tf.reduce_max(pi_stddev, axis=-1)
          / tf.reduce_min(pi_stddev, axis=-1)
      )
  
      # Log per-objective values.
      for i, name in enumerate(self._objective_names):
          stats["{}_dual_temperature".format(name)] = temperature[i]
          stats["{}_loss_temperature".format(name)] = loss_temperature[
              i
          ]
          stats["{}_kl_q_rel".format(name)] = (
              tf.reduce_mean(kl_nonparametric[:, i]) / self._epsilons[i]
          )
  
          # Q measurements.
          stats["{}_q_min".format(name)] = tf.reduce_mean(
              tf.reduce_min(q_values, axis=0)[:, i]
          )
          stats["{}_q_mean".format(name)] = tf.reduce_mean(
              tf.reduce_mean(q_values, axis=0)[:, i]
          )
          stats["{}_q_max".format(name)] = tf.reduce_mean(
              tf.reduce_max(q_values, axis=0)[:, i]
          )
  
      return loss, stats
  
file_path: acme/acme/tf/losses/mompo.py
imports: import dataclasses
  from typing import Dict, Sequence, Tuple, Union
  from acme.tf.losses import mpo
  import sonnet
  import tensorflow
  import tensorflow_probability
label_desc: Use image node_id 4 for calling the MultiObjectiveMPO obj's underlying member method code with example usage: obj.__call__(online_action_distribution, target_action_distribution, actions, q_values) and returns: loss,  stats