code: def vmpo_loss(
      sample_log_probs: Array,
      advantages: Array,
      temperature_constraint: LagrangePenalty,
      kl_constraints: Sequence[Tuple[Array, LagrangePenalty]],
      projection_operator: ProjectionOperatorFn = functools.partial(
          jnp.clip, a_min=_EPSILON
      ),
      restarting_weights: Optional[Array] = None,
      importance_weights: Optional[Array] = None,
      top_k_fraction: float = 0.5,
      policy_loss_weight: float = 1.0,
      temperature_loss_weight: float = 1.0,
      kl_loss_weight: float = 1.0,
      alpha_loss_weight: float = 1.0,
      axis_name: Optional[str] = None,
      use_stop_gradient: bool = True,
  ) -> Tuple[Array, MpoOutputs]:
      """Calculates the V-MPO policy improvement loss.
  
      Note: This is a per-example loss which works on any shape inputs as long as
      they are consistent. We denote the shape of the examples E* for ease of
      reference.
  
      Args:
        sample_log_probs: Log probabilities of actions for each example. Shape E*.
        advantages: Advantages for the E-step. Shape E*.
        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* or [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.
        restarting_weights: Optional restarting weights, shape E*, 0 means that this
          step is the start of a new episode and we ignore losses at this step
          because the agent cannot influence these.
        importance_weights: Optional importance weights, shape E*.
        top_k_fraction: Fraction of samples to use in the E-step.
        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.
        axis_name: Optional axis name for `pmap`. If `None`, computations
          are performed locally on each device.
        use_stop_gradient: bool indicating whether or not to apply stop gradient.
  
      Returns:
        Per example `loss` with same shape E* as array inputs, and additional data
        including the components of this loss and the normalized weights in the
        AdditionalOutputs.
      """
      # Define default restarting weights and importance weights.
      if restarting_weights is None:
          restarting_weights = jnp.ones_like(sample_log_probs)
      if importance_weights is None:
          importance_weights = jnp.ones_like(sample_log_probs)
  
      # Check shapes.
      chex.assert_equal_shape(
          [
              advantages,
              sample_log_probs,
              restarting_weights,
              importance_weights,
          ]
      )
  
      chex.assert_rank(temperature_constraint.epsilon, 0)
      chex.assert_type(
          [
              sample_log_probs,
              advantages,
              restarting_weights,
              importance_weights,
              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, advantages.ndim + 1)
              chex.assert_equal_shape_prefix(
                  [kl, advantages], advantages.ndim
              )
          else:
              chex.assert_equal_shape([kl, advantages])
  
      # E-step: Calculate the reweighting and the temperature loss.
      (
          temperature_loss,
          norm_weights,
          num_samples,
      ) = vmpo_compute_weights_and_temperature_loss(
          advantages,
          restarting_weights,
          importance_weights,
          temperature_constraint,
          projection_operator,
          top_k_fraction,
          axis_name=axis_name,
          use_stop_gradient=use_stop_gradient,
      )
  
      # M-step: Supervised learning of reweighted trajectories using the weights
      # from the E-step, with additional KL constraints.
      # The weights are normalized so that the sum is 1. We multiply by the number
      # of examples so that we can give a policy loss per example and take the mean,
      # and we assume `restarting_weights` are already included.
      if axis_name:
          num_examples = jax.lax.all_gather(
              sample_log_probs, axis_name=axis_name
          ).size
      else:
          num_examples = sample_log_probs.size
      policy_loss = -sample_log_probs * norm_weights * num_examples
  
      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])
  
      # Calculate the total policy improvement loss.
      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 4 for calling a global function with example usage: vmpo_loss(sample_log_probs, advantages, temperature_constraint, kl_constraints, projection_operator, restarting_weights, importance_weights, top_k_fraction, policy_loss_weight, temperature_loss_weight, kl_loss_weight, alpha_loss_weight, axis_name, use_stop_gradient) and returns: loss,  MpoOutputs