docstr: Compares method with a local version of the exp-weighted averaging.
docstr_tok: ["Compares","method","with","a","local","version","of","the","exp-weighted","averaging","."]
code: def test_weighted_averages(self, cum_reward, kappa):
      """Compares method with a local version of the exp-weighted averaging."""
      action_trajectories = jnp.reshape(
          jnp.arange(3 * 10 * 4), (3, 10, 4), order="F"
      )
      averaged_trajectory = mppi.return_weighted_average(
          action_trajectories=action_trajectories,
          cum_reward=cum_reward,
          kappa=kappa,
      )
      exp_weights = jnp.exp(kappa * cum_reward)
      # Verify single-value averaging lines up with the global averaging call:
      for i in range(10):
          for j in range(4):
              np.testing.assert_allclose(
                  averaged_trajectory[i, j],
                  jnp.sum(exp_weights * action_trajectories[:, i, j])
                  / jnp.sum(exp_weights),
                  atol=1e-5,
                  rtol=1e-5,
              )
  
imports: import functools
  from typing import Any
  from acme import specs
  from acme.agents.jax.mbop import ensemble
  from acme.agents.jax.mbop import models
  from acme.agents.jax.mbop import mppi
  from acme.jax import networks
  import jax
  import jax.numpy
  import numpy
  from absl.testing import absltest
  from absl.testing import parameterized
label_desc: Use image node_id 1 for calling the WeightedAverageTests obj's underlying member method code with example usage: obj.test_weighted_averages(cum_reward, kappa) without return types