code: def mpo_compute_weights_and_temperature_loss(
      sample_q_values: Array,
      temperature_constraint: LagrangePenalty,
      projection_operator: ProjectionOperatorFn,
      sample_axis: int = 0,
  ) -> Tuple[Array, Array, Scalar]:
      """Computes the weights and temperature loss for MPO.
  
      The E-Step computes a non-parameteric sample-based approximation of the
      current policy by reweighting the state-action value function.
  
      Here, we compute this nonparametric policy and optimize the temperature
      parameter used in the reweighting.
  
      Args:
        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.
        projection_operator: Function to project temperature into the positive
          range.
        sample_axis: Axis in sample_q_values containing sampled actions.
  
      Returns:
        The temperature loss, normalized weights and number of actions samples per
        state.
      """
      chex.assert_rank(temperature_constraint.epsilon, 0)
      chex.assert_type(
          [
              sample_q_values,
              temperature_constraint.alpha,
              temperature_constraint.epsilon,
          ],
          float,
      )
  
      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}"
          )
  
      n_action_samples = sample_q_values.shape[sample_axis]
  
      # Clip the temperature value (temperature must be positive).
      temperature = projection_operator(temperature_constraint.alpha)
      epsilon = temperature_constraint.epsilon
  
      # Scale the Q-values.
      scaled_sample_q_values = sample_q_values / temperature
  
      # Temperature optimization.
      q_logsumexp = jax.scipy.special.logsumexp(
          scaled_sample_q_values, axis=sample_axis, keepdims=True
      )
  
      # The temperature loss encourages the current and previous policy to stay
      # close. This loss optimizes the convex dual of an upper bound on the average
      # KL (epsilon) between the current and previous state-action values.
      temperature_loss = temperature * epsilon + (
          temperature
          * (
              jnp.squeeze(q_logsumexp, axis=sample_axis)
              - jnp.log(n_action_samples)
          )
      )
  
      # The weights corresponds to a softmax over state-action values.
      weights = jnp.exp(scaled_sample_q_values - q_logsumexp)
  
      # Normalize the weights before the M-Step
      norm_weights = weights / jnp.sum(
          weights, axis=sample_axis, keepdims=True
      )
  
      return temperature_loss, norm_weights, n_action_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 2 for calling a global function with example usage: mpo_compute_weights_and_temperature_loss(sample_q_values, temperature_constraint, projection_operator, sample_axis) and returns: temperature_loss,  norm_weights,  n_action_samples