class_docstr: MPO loss with decoupled KL constraints as in (Abdolmaleki et al., 2018).
  
  This implementation of the MPO loss includes the following features, as
  options:
  - Satisfying the KL-constraint on a per-dimension basis (on by default);
  - Penalizing actions that fall outside of [-1, 1] (on by default) as a
      special case of multi-objective MPO (MO-MPO; Abdolmaleki et al., 2020).
  For best results on the control suite, keep both of these on.
  
  (Abdolmaleki et al., 2018): https://arxiv.org/pdf/1812.02256.pdf
  (Abdolmaleki et al., 2020): https://arxiv.org/pdf/2005.07513.pdf
class_docstr_tok: ["MPO","loss","with","decoupled","KL","constraints","as","in","(","Abdolmaleki","et","al.",",","2018",")",".","This","implementation","of","the","MPO","loss","includes","the","following","features",",","as","options",":","-","Satisfying","the","KL-constraint","on","a","per-dimension","basis","(","on","by","default",")",";","-","Penalizing","actions","that","fall","outside","of","[","-1",",","1","]","(","on","by","default",")","as","a","special","case","of","multi-objective","MPO","(","MO-MPO",";","Abdolmaleki","et","al.",",","2020",")",".","For","best","results","on","the","control","suite",",","keep","both","of","these","on",".","(","Abdolmaleki","et","al.",",","2018",")",":","https",":","\/\/arxiv.org\/pdf\/1812.02256.pdf","(","Abdolmaleki","et","al.",",","2020",")",":","https",":","\/\/arxiv.org\/pdf\/2005.07513.pdf"]
docstr: Computes the decoupled 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].
  
  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","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","]",".","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].
  ) -> Tuple[tf.Tensor, Dict[str, tf.Tensor]]:
      """Computes the decoupled 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].
  
      Returns:
        Loss, combining the policy loss, KL penalty, and dual losses required to
          adapt the dual variables.
        Stats, for diagnostics and tracking performance.
      """
  
      # 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._epsilon, temperature
      )
  
      # Only needed for diagnostics: Compute estimated actualized KL between the
      # non-parametric and current target policies.
      kl_nonparametric = (
          compute_nonparametric_kl_from_normalized_weights(
              normalized_weights
          )
      )
  
      if self._action_penalization:
          # Project and transform action penalization temperature.
          self._log_penalty_temperature.assign(
              tf.maximum(
                  min_log_temperature, self._log_penalty_temperature
              )
          )
          penalty_temperature = (
              tf.math.softplus(self._log_penalty_temperature)
              + _MPO_FLOAT_EPSILON
          )
  
          # Compute action penalization cost.
          # Note: the cost is zero in [-1, 1] and quadratic beyond.
          diff_out_of_bound = actions - tf.clip_by_value(
              actions, -1.0, 1.0
          )
          cost_out_of_bound = -tf.norm(diff_out_of_bound, axis=-1)
  
          (
              penalty_normalized_weights,
              loss_penalty_temperature,
          ) = compute_weights_and_temperature_loss(
              cost_out_of_bound,
              self._epsilon_penalty,
              penalty_temperature,
          )
  
          # Only needed for diagnostics: Compute estimated actualized KL between the
          # non-parametric and current target policies.
          penalty_kl_nonparametric = (
              compute_nonparametric_kl_from_normalized_weights(
                  penalty_normalized_weights
              )
          )
  
          # Combine normalized weights.
          normalized_weights += penalty_normalized_weights
          loss_temperature += loss_penalty_temperature
      # 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 = compute_cross_entropy_loss(
          actions, normalized_weights, fixed_stddev_distribution
      )
      loss_policy_stddev = compute_cross_entropy_loss(
          actions, normalized_weights, 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,
      ) = compute_parametric_kl_penalty_and_dual_loss(
          kl_mean, alpha_mean, self._epsilon_mean
      )
      (
          loss_kl_stddev,
          loss_alpha_stddev,
      ) = 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
      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)
      stats["dual_temperature"] = tf.reduce_mean(temperature)
      # Losses.
      stats["loss_policy"] = tf.reduce_mean(loss)
      stats["loss_alpha"] = tf.reduce_mean(
          loss_alpha_mean + loss_alpha_stddev
      )
      stats["loss_temperature"] = tf.reduce_mean(loss_temperature)
      # KL measurements.
      stats["kl_q_rel"] = (
          tf.reduce_mean(kl_nonparametric) / self._epsilon
      )
  
      if self._action_penalization:
          stats["penalty_kl_q_rel"] = (
              tf.reduce_mean(penalty_kl_nonparametric)
              / self._epsilon_penalty
          )
  
      stats["kl_mean_rel"] = (
          tf.reduce_mean(kl_mean) / self._epsilon_mean
      )
      stats["kl_stddev_rel"] = (
          tf.reduce_mean(kl_stddev) / self._epsilon_stddev
      )
      if self._per_dim_constraining:
          # When KL is constrained per-dimension, we also log per-dimension min and
          # max of mean/std of the realized KL costs.
          stats["kl_mean_rel_min"] = (
              tf.reduce_min(tf.reduce_mean(kl_mean, axis=0))
              / self._epsilon_mean
          )
          stats["kl_mean_rel_max"] = (
              tf.reduce_max(tf.reduce_mean(kl_mean, axis=0))
              / self._epsilon_mean
          )
          stats["kl_stddev_rel_min"] = (
              tf.reduce_min(tf.reduce_mean(kl_stddev, axis=0))
              / self._epsilon_stddev
          )
          stats["kl_stddev_rel_max"] = (
              tf.reduce_max(tf.reduce_mean(kl_stddev, axis=0))
              / self._epsilon_stddev
          )
      # Q measurements.
      stats["q_min"] = tf.reduce_mean(tf.reduce_min(q_values, axis=0))
      stats["q_max"] = tf.reduce_mean(tf.reduce_max(q_values, axis=0))
      # 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)
      )
  
      return loss, stats
  
file_path: acme/acme/tf/losses/mpo.py
imports: from typing import Dict, Tuple, Union
  import sonnet
  import tensorflow
  import tensorflow_probability
label_desc: Use image node_id 3 for calling the MPO obj's underlying member method code with example usage: obj.__call__(online_action_distribution, target_action_distribution, actions, q_values) and returns: loss,  stats