# 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 typing import Tuple

from absl.testing import absltest
from absl.testing import parameterized
import chex
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

Array = chex.Array
Numeric = chex.Numeric

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=jnp.array(_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 _decoupled_multivariate_normal_kl_divergence(
    mu_0: Array, sigma_0: Numeric, mu_1: Array, sigma_1: Numeric,
    per_dimension: bool = False
) -> Tuple[Array, Array]:
  """Compute the KL between diagonal Gaussians decomposed into mean and covar-e.

  Args:
    mu_0: array like of mean values for policy 0
    sigma_0: array like of std values for policy 0
    mu_1: array like of mean values for policy 1
    sigma_1: array like of std values for policy 1
    per_dimension: Whether to return a separate kl divergence for each dimension
      on the last axis.

  Returns:
    the kl divergence between the distributions decomposed into mean and
    covariance.
  """
  # Support scalar and vector `sigma`. If vector, mu.shape==sigma.shape.
  sigma_1 = jnp.ones_like(mu_1) * sigma_1
  sigma_0 = jnp.ones_like(mu_0) * sigma_0
  v1 = jnp.clip(sigma_1**2, 1e-6, 1e6)
  v0 = jnp.clip(sigma_0**2, 1e-6, 1e6)
  mu_diff = mu_1 - mu_0
  kl_mean = 0.5 * jnp.divide(mu_diff**2, v1)
  kl_cov = 0.5 * (jnp.divide(v0, v1) - jnp.ones_like(mu_1) + jnp.log(v1) -
                  jnp.log(v0))
  if not per_dimension:
    kl_mean = jnp.sum(kl_mean, axis=-1)
    kl_cov = jnp.sum(kl_cov, axis=-1)

  return kl_mean, kl_cov


def get_decoupled_kl_constraints(out, params, per_dimension):
  """Factorises KL for Gaussian."""
  kl_mean, kl_covariance = (
      _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 = (
      _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 = jnp.array(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(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)

  @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(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)

  @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()
