class_docstr: Tests for the MPO losses.
class_docstr_tok: ["Tests","for","the","MPO","losses","."]
docstr: Tests the gradients in the `_kl_constraint_loss` method.
docstr_tok: ["Tests","the","gradients","in","the","`","_kl_constraint_loss","`","method","."]
code: def test_kl_constraint_loss_gradients(self):
      """Tests the gradients in the `_kl_constraint_loss` method."""
      kl = jnp.array(1.0, jnp.float32)
      alpha = jnp.array(1.0, jnp.float32)
      _, _, alpha = mpo_ops.kl_constraint_loss(
          kl,
          mpo_ops.LagrangePenalty(
              alpha=alpha,
              epsilon=_EPSILON_MEAN_BOUND,
              per_dimension=False,
          ),
          _PROJECTION_OPERATOR,
      )
  
      def alpha_loss_fn(alpha_):
          penalty = mpo_ops.LagrangePenalty(
              alpha=alpha_,
              epsilon=_EPSILON_MEAN_BOUND,
              per_dimension=False,
          )
          _, alpha_loss, _ = mpo_ops.kl_constraint_loss(
              kl, penalty, _PROJECTION_OPERATOR
          )
          return alpha_loss
  
      alpha_gradients = jax.grad(alpha_loss_fn)(alpha)
      actual_alpha_gradients = _EPSILON_MEAN_BOUND - kl
  
      def kl_loss_fn(kl_):
          penalty = mpo_ops.LagrangePenalty(
              alpha=alpha,
              epsilon=_EPSILON_MEAN_BOUND,
              per_dimension=False,
          )
          kl_loss, _, _ = mpo_ops.kl_constraint_loss(
              kl_, penalty, _PROJECTION_OPERATOR
          )
          return kl_loss
  
      kl_gradients = jax.grad(kl_loss_fn)(kl)
      actual_kl_gradients = alpha
  
      self.assertAlmostEqual(kl_gradients, actual_kl_gradients)
      self.assertAlmostEqual(alpha_gradients, actual_alpha_gradients)
  
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 4 for calling the MPOTest obj's underlying member method code with example usage: obj.test_kl_constraint_loss_gradients() and returns: alpha_loss, kl_loss