code: def vmpo_compute_weights_and_temperature_loss(
      advantages: Array,
      restarting_weights: Array,
      importance_weights: Array,
      temperature_constraint: LagrangePenalty,
      projection_operator: ProjectionOperatorFn,
      top_k_fraction: float,
      axis_name: Optional[str] = None,
      use_stop_gradient: bool = True,
  ) -> Tuple[Array, Array, Array]:
      """Computes the weights and temperature loss for V-MPO.
  
      Args:
        advantages: Advantages for the E-step. Shape E*.
        restarting_weights: Restarting weights, 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. Shape E*.
        importance_weights: Optional importance weights. Shape E*
        temperature_constraint: Lagrange constraint for the E-step temperature
          optimization.
        projection_operator: Function to project dual variables (temperature and kl
          constraint alphas) into the positive range.
        top_k_fraction: Fraction of samples to use in the E-step.
        axis_name: Optional axis name for `pmap` or 'vmap'. If `None`, computations
          are performed locally on each device.
        use_stop_gradient: bool indicating whether or not to apply stop gradient.
  
      Returns:
        The temperature loss, normalized weights and number of samples used.
      """
      chex.assert_equal_shape(
          [advantages, restarting_weights, importance_weights]
      )
      chex.assert_rank(temperature_constraint.epsilon, 0)
      chex.assert_type(
          [
              advantages,
              restarting_weights,
              importance_weights,
              temperature_constraint.alpha,
              temperature_constraint.epsilon,
          ],
          float,
      )
  
      importance_weights = jax.lax.select(
          use_stop_gradient,
          jax.lax.stop_gradient(importance_weights),
          importance_weights,
      )
  
      # Lagrange constraint.
      temperature = projection_operator(temperature_constraint.alpha)
      epsilon_temperature = temperature_constraint.epsilon
  
      # Scale the advantages.
      scaled_advantages = restarting_weights * advantages / temperature
      max_scaled_advantage = jnp.max(scaled_advantages)
      # If the axis_name is not None find the maximum across all devices.
      if axis_name:
          assert use_stop_gradient  # Cant differentiate through pmax.
          max_scaled_advantage = jax.lax.stop_gradient(
              max_scaled_advantage
          )
          max_scaled_advantage = jax.lax.pmax(
              max_scaled_advantage, axis_name=axis_name
          )
      else:
          max_scaled_advantage = jax.lax.select(
              use_stop_gradient,
              jax.lax.stop_gradient(max_scaled_advantage),
              max_scaled_advantage,
          )
      # Maybe don't use all of the advantages.
      top_k_restarting_weights = get_top_k_weights(
          top_k_fraction,
          restarting_weights,
          scaled_advantages,
          axis_name,
          use_stop_gradient,
      )
  
      all_sum = base.AllSum(axis_name)
  
      # Reweight the old trajectories.
      unnormalized_weights = (
          top_k_restarting_weights
          * importance_weights
          * jnp.exp(scaled_advantages - max_scaled_advantage)
      )
      # If the axis_name is not None these sums will be taken across all devices.
      sum_weights = all_sum(unnormalized_weights) + _EPSILON
      num_samples = all_sum(top_k_restarting_weights) + _EPSILON
  
      normalized_weights = unnormalized_weights / sum_weights
      normalized_weights = jax.lax.select(
          use_stop_gradient,
          jax.lax.stop_gradient(normalized_weights),
          normalized_weights,
      )
  
      # Calculate the temperature loss.
      log_mean_weights = (
          jnp.log(sum_weights)
          + max_scaled_advantage
          - jnp.log(num_samples)
      )
      temperature_loss = temperature * (
          epsilon_temperature + log_mean_weights
      )
  
      return temperature_loss, normalized_weights, 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 6 for calling a global function with example usage: vmpo_compute_weights_and_temperature_loss(advantages, restarting_weights, importance_weights, temperature_constraint, projection_operator, top_k_fraction, axis_name, use_stop_gradient) and returns: temperature_loss,  normalized_weights,  num_samples