class_docstr: Tests for the MPO losses.
class_docstr_tok: ["Tests","for","the","MPO","losses","."]
docstr: Test that importance weights have the correct effect.
docstr_tok: ["Test","that","importance","weights","have","the","correct","effect","."]
code: def test_importance_weights(
      self, advantages, importance_weights, expected_temperature_loss
  ):
      """Test that importance weights have the correct effect."""
      (
          temperature_loss,
          _,
          _,
      ) = mpo_ops.vmpo_compute_weights_and_temperature_loss(
          advantages,
          np.ones_like(importance_weights),
          importance_weights,
          mpo_ops.LagrangePenalty(jnp.array(1.0), _EPSILON_BOUND),
          functools.partial(np.clip, a_min=1e-8, a_max=None),
          1.0,
      )
      self.assertAlmostEqual(
          temperature_loss, expected_temperature_loss, places=4
      )
  
imports: import functools
  import math
  from typing import Tuple
  from absl.testing import absltest
  from absl.testing import parameterized
  import chex
  import haiku
  import jax
  import jax.numpy
  import numpy
  import optax
  from rlax._src import distributions
  from rlax._src import mpo_ops
label_desc: Use image node_id 9 for calling the MPOTest obj's underlying member method code with example usage: obj.test_importance_weights(advantages, importance_weights, expected_temperature_loss) without return types