code: # Copyright 2020 DeepMind Technologies Limited. All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  # ==============================================================================
  """Tests for mpo_ops.py."""
  
  import functools
  import math
  
  from absl.testing import absltest
  from absl.testing import parameterized
  import haiku as hk
  import jax
  import jax.numpy as jnp
  import numpy as np
  import optax
  from rlax._src import distributions
  from rlax._src import mpo_ops
  
  NUM_SAMPLES = 10
  ACTION_DIM = 3
  TIME_DIM = 8
  BATCH_DIM = 100
  
  # NOTE: These are not typical values used for MPO. In the test case, we know the
  # Q function perfectly so we loosen the bound on the mean to zone in to the
  # optimal policy very quickly. Similarly, we maintain a high variance to sample
  # distinct actions to explore and learn from.
  _INIT_TEMPERATURE = 0.2
  _INIT_ALPHA_MEAN = 0.001
  _INIT_ALPHA_COVARIANCE = float(1e6)
  
  _EPSILON_BOUND = 0.01
  _EPSILON_MEAN_BOUND = 10.0
  _EPSILON_COVARIANCE_BOUND = 1e-12
  
  _NUM_ITERATIONS = 5000
  _TARGET_UPDATE_PERIOD = 100
  _RANDOM_SEED = 42
  
  # The offset to ensure initially the policy is not close to 0
  _MEAN_OFFSET = 2.0
  
  # The final action should optimize down to be close to 0.0
  _MAX_ACTION_ERROR = 0.2
  _MAX_KL_ERROR = 1e-6
  
  _DIAGONAL_GAUSSIAN_DIST = distributions.gaussian_diagonal()
  _PROJECTION_OPERATOR = functools.partial(jnp.clip, a_min=1e-10)
  
  
  def _hk_mock_policy_params(s_tm1):
    """Returns mock policy params."""
    # Outputs of the network are mu and sigma. Both shaped [B, ACTION_DIM].
    pi_out = hk.nets.MLP(
        output_sizes=[2 * ACTION_DIM],
        w_init=hk.initializers.VarianceScaling(1e-3),
        activation=jnp.tanh,
        activate_final=False,
        name='online_policy')(s_tm1)
    pi_mean, pi_cov = jnp.split(pi_out, 2, axis=-1)
    pi_cov = jax.nn.softplus(pi_cov)
    pi_mean = pi_mean + _MEAN_OFFSET
    return {'mean': pi_mean, 'stddev': pi_cov}
  
  
  def _init_params(key):
    init_fn, _ = hk.transform(_hk_mock_policy_params)
    key_seq = hk.PRNGSequence(key)
    s_tm1 = jax.random.normal(
        next(key_seq), (TIME_DIM, BATCH_DIM, ACTION_DIM), jnp.float32)
    online_params = init_fn(next(key_seq), s_tm1)
    return dict(
        online=online_params,
        target=online_params,
        mpo=dict(
            temperature=_INIT_TEMPERATURE,
            alpha_mean=_INIT_ALPHA_MEAN,
            alpha_covariance=_INIT_ALPHA_COVARIANCE),
    )
  
  
  def _mock_outputs(online_params, target_params, key, target_name):
    """Returns mock network outputs."""
    _, policy_params_fn = hk.transform(_hk_mock_policy_params)
    key_seq = hk.PRNGSequence(key)
  
    state_size = ACTION_DIM
  
    # Input state: [TIME_DIM, BATCH_DIM, DIM_STATE]
    s_tm1 = jax.random.normal(
        next(key_seq), (TIME_DIM, BATCH_DIM, state_size), jnp.float32)
    policy_params = policy_params_fn(online_params, None, s_tm1)
    target_policy_params = policy_params_fn(target_params, None, s_tm1)
  
    # Shape for actions: [NUM_SAMPLES, TIME_DIM, BATCH_DIM, ACTION_DIM]
    mean, stddev = target_policy_params['mean'], target_policy_params['stddev']
    mean_repeated = jnp.repeat(
        mean.reshape((1,) + mean.shape), NUM_SAMPLES, axis=0)
    stddev_repeated = jnp.repeat(
        stddev.reshape((1,) + stddev.shape), NUM_SAMPLES, axis=0)
    target_actions = _DIAGONAL_GAUSSIAN_DIST.sample(
        next(key_seq), mean_repeated, stddev_repeated)
    # If the target is advantages then num samples is 1.
    if target_name == 'advantages':
      target_actions = target_actions[0, ...]
  
    # Shape for Q: [NUM_SAMPLES, TIME_DIM, BATCH_DIM]
    # Setting Q = -a_t * tf.transpose(a_t) where a_t = s_t + a.
    # The solution to optimizing this is basically for the policy to output
    # 0 actions thereby minimizing the cost. Since this is a convex
    # optimization problem, the algorithm should get to a good solution quickly.
  
    # First compute a_t = s_t + a with shape: [NUM_SAMPLES, TIME_DIM, BATCH_DIM,
    # ACTION_DIM] since action dim is the same as shape dim here and then compute
    # the quadratic form.
    a_t = target_actions + jnp.expand_dims(s_tm1, 0)
    sample_q_values = -jnp.sum(a_t ** 2, axis=-1)
    # Set the advantage to the same as the q value.
    # Shape for advantages: [TIME_DIM, BATCH_DIM]
    advantages = sample_q_values[0, :, :]
  
    return dict(
        pi_params=policy_params,
        target_pi_params=target_policy_params,
        sample_q_values=sample_q_values,
        advantages=advantages,
        target_actions=target_actions,
    )
  
  
  def get_common_loss_fn_inputs(params, key, target_name):
    out = _mock_outputs(params['online'], params['target'], key, target_name)
    pi_sample_log_probs = _DIAGONAL_GAUSSIAN_DIST.logprob(
        out['target_actions'], out['pi_params']['mean'],
        out['pi_params']['stddev'])
  
    return out, {
        'sample_log_probs': pi_sample_log_probs,
        target_name: out[target_name],
        'temperature_constraint': mpo_ops.LagrangePenalty(
            params['mpo']['temperature'], _EPSILON_BOUND)}
  
  
  def get_decoupled_kl_constraints(out, params, per_dimension):
    """Factorises KL for Gaussian."""
    kl_mean, kl_covariance = (
        distributions.decoupled_multivariate_normal_kl_divergence(
            out['target_pi_params']['mean'], out['target_pi_params']['stddev'],
            out['pi_params']['mean'], out['pi_params']['stddev'],
            per_dimension=per_dimension))
    alpha_mean = params['mpo']['alpha_mean'] * jnp.ones_like(kl_mean)
    alpha_covariance = params['mpo']['alpha_covariance'] * jnp.ones_like(
        kl_covariance)
  
    return [
        (kl_mean, mpo_ops.LagrangePenalty(
            alpha=alpha_mean, epsilon=_EPSILON_MEAN_BOUND,
            per_dimension=per_dimension)),
        (kl_covariance, mpo_ops.LagrangePenalty(
            alpha=alpha_covariance, epsilon=_EPSILON_COVARIANCE_BOUND,
            per_dimension=per_dimension)),
    ]
  
  
  def get_coupled_kl_constraints(out, params, per_dimension):
    kl_mean, kl_covariance = (
        distributions.decoupled_multivariate_normal_kl_divergence(
            out['target_pi_params']['mean'], out['target_pi_params']['stddev'],
            out['pi_params']['mean'], out['pi_params']['stddev'],
            per_dimension=per_dimension))
    alpha_mean = params['mpo']['alpha_mean'] * jnp.ones_like(kl_mean)
    return [
        (kl_mean + kl_covariance, mpo_ops.LagrangePenalty(
            alpha=alpha_mean,
            epsilon=_EPSILON_MEAN_BOUND + _EPSILON_COVARIANCE_BOUND,
            per_dimension=per_dimension))
    ]
  
  
  def vmpo_e_step_without_restarting_or_importance_weights(advantages, **kwargs):
    restarting_weights = jnp.ones_like(advantages)
    importance_weights = jnp.ones_like(advantages)
    return mpo_ops.vmpo_compute_weights_and_temperature_loss(
        advantages=advantages, restarting_weights=restarting_weights,
        importance_weights=importance_weights, **kwargs)
  
  
  class MPOTest(parameterized.TestCase):
    """Tests for the MPO losses."""
  
    @parameterized.parameters(
        {'target_name': 'sample_q_values',
         'loss_fn': mpo_ops.mpo_loss,
         'get_kl_constraints': get_decoupled_kl_constraints,
         'per_dimension': False},
        {'target_name': 'advantages',
         'loss_fn': mpo_ops.vmpo_loss,
         'get_kl_constraints': get_decoupled_kl_constraints,
         'per_dimension': False},
        {'target_name': 'sample_q_values',
         'loss_fn': mpo_ops.mpo_loss,
         'get_kl_constraints': get_coupled_kl_constraints,
         'per_dimension': False},
        {'target_name': 'advantages',
         'loss_fn': mpo_ops.vmpo_loss,
         'get_kl_constraints': get_coupled_kl_constraints,
         'per_dimension': False},
        {'target_name': 'sample_q_values',
         'loss_fn': mpo_ops.mpo_loss,
         'get_kl_constraints': get_decoupled_kl_constraints,
         'per_dimension': True},
        {'target_name': 'advantages',
         'loss_fn': mpo_ops.vmpo_loss,
         'get_kl_constraints': get_decoupled_kl_constraints,
         'per_dimension': True},
        {'target_name': 'sample_q_values',
         'loss_fn': mpo_ops.mpo_loss,
         'get_kl_constraints': get_coupled_kl_constraints,
         'per_dimension': True},
        {'target_name': 'advantages',
         'loss_fn': mpo_ops.vmpo_loss,
         'get_kl_constraints': get_coupled_kl_constraints,
         'per_dimension': True},
    )
    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)
  
    @parameterized.parameters(
        {'e_step_fn': mpo_ops.mpo_compute_weights_and_temperature_loss,
         'additional_inputs': {},
         # dL/dq == 1 and dL/dt == epsilon (for one sample)
         'expected_deriv_of_target': [[[1]]],
         'sample_dimension': True},
        {'e_step_fn': vmpo_e_step_without_restarting_or_importance_weights,
         'additional_inputs': {'top_k_fraction': 1.0},
         'expected_deriv_of_target': [[1]],
         'sample_dimension': False},
    )
    def test_e_step_gradient_computation(
        self, e_step_fn, additional_inputs, expected_deriv_of_target,
        sample_dimension):
      """Tests the gradients from the E-step against the analytic ones."""
      # Target has shape [NUM_SAMPLES, T, B] => [1, 1, 1]
      target = jnp.array([[3]], jnp.float32)
      if sample_dimension:
        target = jnp.expand_dims(target, axis=0)
      temperature = jnp.array(0.1, jnp.float32)
      def fn(target_, temperature_):
        temperature_constraint = mpo_ops.LagrangePenalty(
            temperature_, _EPSILON_BOUND)
        temperature_loss, _, _ = e_step_fn(
            target_, temperature_constraint=temperature_constraint,
            projection_operator=_PROJECTION_OPERATOR,
            **additional_inputs)
        return jnp.mean(temperature_loss)
      grad = jax.grad(fn, argnums=(0, 1))(target, temperature)
  
      np.testing.assert_almost_equal(np.array(grad[0]), np.array(
          expected_deriv_of_target, np.float32), decimal=4)
      self.assertAlmostEqual(grad[1], _EPSILON_BOUND, places=4)
  
    @parameterized.parameters(
        {'e_step_fn': mpo_ops.mpo_compute_weights_and_temperature_loss,
         'additional_inputs': {},
         'sample_dimension': True},
        {'e_step_fn': vmpo_e_step_without_restarting_or_importance_weights,
         'additional_inputs': {'top_k_fraction': 1.0},
         'sample_dimension': False},
    )
    def test_e_step_stop_gradient(
        self, e_step_fn, additional_inputs, sample_dimension):
      """Tests no gradients flow through `weights` in the E-Step."""
      # Target has shape [NUM_SAMPLES, T, B] => [1, 1, 1]
      target = jnp.array([[3]], jnp.float32)
      if sample_dimension:
        target = jnp.expand_dims(target, axis=0)
      temperature = 0.1
      # pylint: disable=g-long-lambda
      def mean_weights_fn(target_, temperature_):
        temperature_constraint = mpo_ops.LagrangePenalty(
            temperature_, _EPSILON_BOUND)
        _, weights, _ = e_step_fn(
            target_, temperature_constraint=temperature_constraint,
            projection_operator=_PROJECTION_OPERATOR,
            **additional_inputs)
        return jnp.mean(weights)
      grad = jax.grad(mean_weights_fn, argnums=(0, 1))(target, temperature)
      np.testing.assert_almost_equal(
          np.array(grad[0]), np.zeros_like(grad[0]), decimal=4)
      self.assertAlmostEqual(grad[1], 0., places=4)
  
    def test_kl_constraint_loss_gradients(self):
      """Tests the gradients in the `_kl_constraint_loss` method."""
      kl = jnp.array(1., jnp.float32)
      alpha = jnp.array(1., 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)
  
    def test_kl_constraint_loss_stop_gradients(self):
      """Tests the stop gradients in the `kl_constraint_loss` function.
  
        The `alpha_loss` term should not affect the KL and the `kl` term should
        not affect `alpha`.
      """
      kl = jnp.array(1., jnp.float32)
      alpha = jnp.array(1., jnp.float32)
      _, _, alpha = mpo_ops.kl_constraint_loss(kl, mpo_ops.LagrangePenalty(
          alpha=alpha, epsilon=_EPSILON_MEAN_BOUND, per_dimension=False),
                                               _PROJECTION_OPERATOR)
  
      def kl_loss_fn(alpha_):
        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)(alpha)
  
      def alpha_loss_fn(kl_):
        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)(kl)
  
      # Test that there are no gradients of KL w.r.t alpha
      self.assertEqual(kl_gradients, 0.)
  
      # Test that there are no gradients of alpha w.r.t kl
      self.assertEqual(alpha_gradients, 0.)
  
    @parameterized.parameters(
        # With restarting weights of 1 (and temperature of 1) the weights should
        # be e^-1, 1, max advantage is 2 and num samples is 2 so temperature loss
        # is log(1 + e^-1) + 2 - log(2) + temperature epsilon
        {'advantages': np.array([[1.0, 2.0]]),
         'restarting_weights': np.array([[1.0, 1.0]]),
         'expected_temperature_loss': (math.log(1.0 + math.exp(-1.0)) + 2.0 -
                                       math.log(2.0) + _EPSILON_BOUND)},
        # With the second restarting weight set to 0 the weights become 1, 0
        # max advantage is 1 and num samples is 1 so temperature loss is
        # log(1) + 1 - log(1) + temperature epsilon
        {'advantages': np.array([[1.0, 2.0]]),
         'restarting_weights': np.array([[1.0, 0.0]]),
         'expected_temperature_loss': 1.0 + _EPSILON_BOUND},
    )
    def test_restarting_weights(
        self, advantages, restarting_weights, expected_temperature_loss):
      """Test that calculation is correct if restarting weight is set to 0."""
      temperature_loss, _, _ = mpo_ops.vmpo_compute_weights_and_temperature_loss(
          advantages, restarting_weights, np.ones_like(restarting_weights),
          mpo_ops.LagrangePenalty(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)
  
    @parameterized.parameters(
        # When the top k fraction is 1.0 all of the weights should be 1
        {'top_k_fraction': 1.0,
         'scaled_advantages': np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]),
         'expected_top_k_weights': np.array([[1.0, 1.0, 1.0], [1.0, 1.0, 1.0]])},
        # When the top k fraction is 0.5 it will take the bottom row as these are
        # the highest.
        {'top_k_fraction': 0.5,
         'scaled_advantages': np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]),
         'expected_top_k_weights': np.array([[0.0, 0.0, 0.0], [1.0, 1.0, 1.0]])}
    )
    def test_top_k_fraction(
        self, top_k_fraction, scaled_advantages, expected_top_k_weights):
      """Test that only the top k fraction are used."""
      top_k_weights = mpo_ops.get_top_k_weights(
          top_k_fraction, jnp.ones_like(scaled_advantages), scaled_advantages)
      np.testing.assert_allclose(top_k_weights, expected_top_k_weights)
  
    def test_top_k_fraction_too_low(self):
      """Test if the top k fraction returns 0 advantages we raise an error."""
      with self.assertRaises(ValueError):
        mpo_ops.get_top_k_weights(0.01, jnp.ones((3, 2)), jnp.ones((3, 2)))
  
    @parameterized.parameters(
        # With importance weights of 1 (and temperature of 1) the weights should
        # be e^-1, 1, max advantage is 2 and num samples is 2 so temperature loss
        # is log(1 + e^-1) + 2 - log(2) + temperature epsilon
        {'advantages': np.array([[1.0, 2.0]]),
         'importance_weights': np.array([[1.0, 1.0]]),
         'expected_temperature_loss': (math.log(1.0 + math.exp(-1.0)) + 2.0 -
                                       math.log(2.0) + _EPSILON_BOUND)},
        # If the second importance weight is 0.5 temperature loss becomes
        # log(0.5 + e^-1) + 2 - log(2) + temperature epsilon
        {'advantages': np.array([[1.0, 2.0]]),
         'importance_weights': np.array([[1.0, 0.5]]),
         'expected_temperature_loss': (math.log(0.5 + math.exp(-1.0)) + 2.0 -
                                       math.log(2.0) + _EPSILON_BOUND)},
    )
    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(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)
  
    @parameterized.parameters({'per_dimension': True}, {'per_dimension': False})
    def test_mpo_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, mpo_inputs = get_common_loss_fn_inputs(params, key, 'sample_q_values')
      kl_constraints = get_coupled_kl_constraints(out, params,
                                                  per_dimension=per_dimension)
      mpo_inputs.update({'kl_constraints': kl_constraints})
  
      # Original loss fn inputs are [S T B],
      stb_loss, stb_outputs = mpo_ops.mpo_loss(**mpo_inputs)
      mean_stb_loss = jnp.mean(stb_loss)
  
      # Swap axes and try [S B T]
      mpo_inputs.update({
          'sample_log_probs': jnp.swapaxes(mpo_inputs['sample_log_probs'], 1, 2),
          'sample_q_values': jnp.swapaxes(mpo_inputs['sample_q_values'], 1, 2),
          '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],
      })
      sbt_loss, sbt_outputs = mpo_ops.mpo_loss(**mpo_inputs)
      mean_sbt_loss = jnp.mean(sbt_loss)
  
      # Try [T B S] denoting sample_axis at 2 instead of 0.
      mpo_inputs.update({
          'sample_log_probs': jnp.swapaxes(mpo_inputs['sample_log_probs'], 0, 2),
          'sample_q_values': jnp.swapaxes(mpo_inputs['sample_q_values'], 0, 2),
          'kl_constraints': kl_constraints,  # T B
          'sample_axis': 2
      })
      tbs_loss, tbs_outputs = mpo_ops.mpo_loss(**mpo_inputs)
      mean_tbs_loss = jnp.mean(tbs_loss)
  
      self.assertAlmostEqual(mean_stb_loss, mean_sbt_loss, places=4)
      self.assertAlmostEqual(mean_tbs_loss, mean_sbt_loss, places=4)
      self.assertEqual(tbs_outputs.num_samples, sbt_outputs.num_samples)
      self.assertEqual(tbs_outputs.num_samples, stb_outputs.num_samples)
  
    @parameterized.parameters({'per_dimension': True}, {'per_dimension': False})
    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)
  
  
  if __name__ == '__main__':
    absltest.main()
  
apis: numpy.array numpy.zeros_like numpy.ones_like numpy.testing.assert_allclose