code: def compute_weights_and_temperature_loss(
      q_values: jnp.ndarray,
      logits: jnp.ndarray,
      epsilon: float,
      temperature: jnp.ndarray,
  ) -> Tuple[jnp.ndarray, jnp.ndarray]:
      """Computes normalized importance weights for the policy optimization.
  
      Args:
        q_values: Q-values associated with the actions sampled from the target
          policy; expected shape [B, D].
        logits: Parameters to the categorical distribution with respect to which the
          expectations are going to be computed.
        epsilon: Desired constraint on the KL between the target and non-parametric
          policies.
        temperature: Scalar used to temper the Q-values before computing normalized
          importance weights from them. This is really the Lagrange dual variable in
          the constrained optimization problem, the solution of which is the
          non-parametric policy targeted by the policy loss.
  
      Returns:
        Normalized importance weights, used for policy optimization.
        Temperature loss, used to adapt the temperature.
      """
  
      # Temper the given Q-values using the current temperature.
      tempered_q_values = jax.lax.stop_gradient(q_values) / temperature
  
      # Compute the E-step normalized logits.
      unnormalized_logits = tempered_q_values + jax.nn.log_softmax(
          logits, axis=-1
      )
      logits_e_step = jax.nn.log_softmax(unnormalized_logits, axis=-1)
  
      # Compute the temperature loss (dual of the E-step optimization problem).
      # Note that the log normalizer will be the same for all actions, so we choose
      # only the first one.
      log_normalizer = unnormalized_logits[:, 0] - logits_e_step[:, 0]
      loss_temperature = temperature * (
          epsilon + jnp.mean(log_normalizer)
      )
  
      return logits_e_step, loss_temperature
  
file_path: acme/acme/agents/jax/mpo/categorical_mpo.py
imports: from typing import NamedTuple, Tuple
  import distrax
  import jax
  import jax.numpy
label_desc: Use image node_id 1 for calling a global function with example usage: compute_weights_and_temperature_loss(q_values, logits, epsilon, temperature) and returns: logits_e_step,  loss_temperature