class_docstr: Tests for the MPO losses.
class_docstr_tok: ["Tests","for","the","MPO","losses","."]
docstr: Tests that the policy optimization works correctly.
docstr_tok: ["Tests","that","the","policy","optimization","works","correctly","."]
code: def test_optimization(
      self, target_name, loss_fn, get_kl_constraints, per_dimension
  ):
      """Tests that the policy optimization works correctly."""
  
      def _loss(params, key):
          out, loss_fn_inputs = get_common_loss_fn_inputs(
              params, key, target_name
          )
          kl_constraints = get_kl_constraints(
              out, params, per_dimension
          )
          loss_fn_inputs.update({"kl_constraints": kl_constraints})
          loss, mpo_stats = loss_fn(**loss_fn_inputs)
          loss = jnp.mean(loss)
          temperature_bound = jnp.mean(
              mpo_stats.normalized_weights
              * jnp.log(
                  mpo_stats.num_samples * mpo_stats.normalized_weights
                  + 1e-8
              )
          )
          return loss, {
              "outputs": out,
              "temperature_bound": temperature_bound,
          }
  
      key = jax.random.PRNGKey(_RANDOM_SEED)
      grad_fn = jax.jit(jax.grad(_loss, has_aux=True))
      optimizer = optax.adam(1e-3)
      key, new_key = jax.random.split(key)
      params = _init_params(new_key)
      opt_state = optimizer.init((params["online"], params["mpo"]))
  
      @jax.jit
      def _update(params_, opt_state_, key_):
          next_key, key_ = jax.random.split(key_)
          grad, stats = grad_fn(params_, key_)
          updates, opt_state_ = optimizer.update(
              (grad["online"], grad["mpo"]), opt_state_
          )
          online_params, mpo_params = optax.apply_updates(
              (params_["online"], params_["mpo"]), updates
          )
          params_["online"] = online_params
          params_["mpo"] = mpo_params
          return params_, opt_state_, stats, next_key
  
      for iter_idx in range(_NUM_ITERATIONS):
          params, opt_state, extra, key = _update(
              params, opt_state, key
          )
          if iter_idx % _TARGET_UPDATE_PERIOD == 0:
              params["target"] = params["online"]
  
      # Test the bounds are within tolerance.
      key, new_key = jax.random.split(key)
      _, extra = _loss(params, new_key)
      action_mean = jnp.mean(extra["outputs"]["pi_params"]["mean"])
      # Check action mean is close to 0.
      self.assertBetween(
          action_mean, -_MAX_ACTION_ERROR, _MAX_ACTION_ERROR
      )
  
      # Check the temperature are within the bounds.
      self.assertLess(extra["temperature_bound"], _EPSILON_BOUND)
  
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 1 for calling the MPOTest obj's underlying member method code with example usage: obj.test_optimization(target_name, loss_fn, get_kl_constraints, per_dimension) and returns: loss,  dict, params_,  opt_state_,  stats,  next_key