class_docstr: Tests for the MPO losses.
class_docstr_tok: ["Tests","for","the","MPO","losses","."]
docstr: Test loss functions are equivalent regardless of axis order.
docstr_tok: ["Test","loss","functions","are","equivalent","regardless","of","axis","order","."]
code: def test_vmpo_input_axis_order_equivalence(self, per_dimension):
      """Test loss functions are equivalent regardless of axis order."""
      key = jax.random.PRNGKey(_RANDOM_SEED)
      key, new_key = jax.random.split(key)
      params = _init_params(new_key)
      out, vmpo_inputs = get_common_loss_fn_inputs(
          params, key, "advantages"
      )
      kl_constraints = get_coupled_kl_constraints(
          out, params, per_dimension=per_dimension
      )
      vmpo_inputs.update({"kl_constraints": kl_constraints})
  
      # Original loss fn inputs are [T B],
      tb_loss, tb_outputs = mpo_ops.vmpo_loss(**vmpo_inputs)
      mean_tb_loss = jnp.mean(tb_loss)
  
      # Swap axes and try [B T]
      vmpo_inputs.update(
          {
              "sample_log_probs": jnp.swapaxes(
                  vmpo_inputs["sample_log_probs"], 0, 1
              ),
              "advantages": jnp.swapaxes(
                  vmpo_inputs["advantages"], 0, 1
              ),
              "kl_constraints": [
                  (
                      jnp.swapaxes(kl, 0, 1),
                      mpo_ops.LagrangePenalty(
                          alpha=jnp.swapaxes(pen.alpha, 0, 1),
                          epsilon=pen.epsilon,
                          per_dimension=pen.per_dimension,
                      ),
                  )
                  for (kl, pen) in kl_constraints
              ],
          }
      )
      bt_loss, bt_outputs = mpo_ops.vmpo_loss(**vmpo_inputs)
      mean_bt_loss = jnp.mean(bt_loss)
  
      self.assertAlmostEqual(mean_tb_loss, mean_bt_loss, places=4)
      self.assertEqual(tb_outputs.num_samples, bt_outputs.num_samples)
  
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 11 for calling the MPOTest obj's underlying member method code with example usage: obj.test_vmpo_input_axis_order_equivalence(per_dimension) without return types