code: def mpo_loss(
      sample_log_probs: Array,
      sample_q_values: Array,
      temperature_constraint: LagrangePenalty,
      kl_constraints: Sequence[Tuple[Array, LagrangePenalty]],
      projection_operator: ProjectionOperatorFn = functools.partial(
          jnp.clip, a_min=_EPSILON
      ),
      policy_loss_weight: float = 1.0,
      temperature_loss_weight: float = 1.0,
      kl_loss_weight: float = 1.0,
      alpha_loss_weight: float = 1.0,
      sample_axis: int = 0,
      use_stop_gradient: bool = True,
  ) -> Tuple[Array, MpoOutputs]:
      """Implements the MPO loss with a KL bound.
  
      This loss implements the MPO algorithm for policies with a bound for the KL
      between the current and target policy.
  
      Note: This is a per-example loss which works on any shape inputs as long as
      they are consistent. We denote this shape E* for ease of reference. Args
      sample_log_probs and sample_q_values are shape E + an extra sample axis that
      contains the sampled actions' log probs and q values respectively. For
      example, if sample_axis = 0, the shapes expected will be [S, E*]. Or if
      E* = [T, B] and sample_axis = 1, the shapes expected will be [T, S, B].
  
      Args:
        sample_log_probs: An array of shape E* + a sample axis inserted at
          sample_axis containing the log probabilities of the sampled actions under
          the current policy.
        sample_q_values: An array of shape E* + a sample axis inserted at
          sample_axis containing the q function values evaluated on the sampled
          actions.
        temperature_constraint: Lagrange constraint for the E-step temperature
          optimization.
        kl_constraints: KL and variables for applying Lagrangian penalties to bound
          them in the M-step, KLs are  [E*, A?]. Here A is the action dimension
          in the case of per-dimension KL constraints.
        projection_operator: Function to project dual variables (temperature and kl
          constraint alphas) into the positive range.
        policy_loss_weight: Weight for the policy loss.
        temperature_loss_weight: Weight for the temperature loss.
        kl_loss_weight: Weight for the KL loss.
        alpha_loss_weight: Weight for the alpha loss.
        sample_axis: Axis in sample_log_probs and sample_q_values that contains the
          sampled actions' log probs and q values respectively. For example, if
          sample_axis = 0, the shapes expected will be [S, E*]. Or if E* = [T, B]
          and sample_axis = 1, the shapes expected will be [T, S, B].
        use_stop_gradient: bool indicating whether or not to apply stop gradient.
  
      Returns:
        Per example `loss` with shape E*, and additional data including
        the components of this loss and the normalized weights in the
        AdditionalOutputs.
      """
      chex.assert_equal_shape([sample_log_probs, sample_q_values])
      chex.assert_rank(temperature_constraint.epsilon, 0)
      chex.assert_type(
          [
              sample_log_probs,
              sample_q_values,
              temperature_constraint.alpha,
              temperature_constraint.epsilon,
          ],
          float,
      )
  
      for kl, penalty in kl_constraints:
          chex.assert_rank(penalty.epsilon, 0)
          chex.assert_type([kl, penalty.alpha, penalty.epsilon], float)
          if penalty.per_dimension:
              chex.assert_rank(kl, sample_q_values.ndim)
          else:
              chex.assert_rank(kl, sample_q_values.ndim - 1)
  
      if sample_axis < 0:
          sample_axis += sample_q_values.ndim
      if not 0 <= sample_axis < sample_q_values.ndim:
          raise ValueError(
              f"`sample_axis` {sample_axis} not in array rank {sample_q_values.ndim}"
          )
  
      # E-Step. Compute temperature loss, weights, and temperature.
      (
          temperature_loss,
          norm_weights,
          num_samples,
      ) = mpo_compute_weights_and_temperature_loss(
          sample_q_values,
          temperature_constraint,
          projection_operator,
          sample_axis=sample_axis,
      )
  
      norm_weights = jax.lax.select(
          use_stop_gradient,
          jax.lax.stop_gradient(norm_weights),
          norm_weights,
      )
  
      # M-Step. Supervised learning on reweighted probabilities using the weights
      # from the E-Step under an additional KL constraint.
      policy_loss = -jnp.sum(
          norm_weights * sample_log_probs, axis=sample_axis
      )
      kl_loss, alpha_loss = compute_parametric_kl_penalty_and_dual_loss(
          kl_constraints, projection_operator, use_stop_gradient
      )
  
      chex.assert_equal_shape([policy_loss, kl_loss, alpha_loss])
  
      # Combine all loss components. The final loss is of shape E*.
      loss = (
          policy_loss_weight * policy_loss
          + temperature_loss_weight * temperature_loss
          + kl_loss_weight * kl_loss
          + alpha_loss_weight * alpha_loss
      )
      return loss, MpoOutputs(
          temperature_loss=temperature_loss,
          policy_loss=policy_loss,
          kl_loss=kl_loss,
          alpha_loss=alpha_loss,
          normalized_weights=norm_weights,
          num_samples=num_samples,
      )
  
file_path: rlax/rlax/_src/mpo_ops.py
imports: import functools
  from typing import Callable, NamedTuple, Optional, Sequence, Tuple
  from absl import logging
  import chex
  import jax
  import jax.numpy
  from rlax._src import base
label_desc: Use image node_id 1 for calling a global function with example usage: mpo_loss(sample_log_probs, sample_q_values, temperature_constraint, kl_constraints, projection_operator, policy_loss_weight, temperature_loss_weight, kl_loss_weight, alpha_loss_weight, sample_axis, use_stop_gradient) and returns: loss,  MpoOutputs