code: def return_weighted_average(
      action_trajectories: jnp.ndarray,
      cum_reward: jnp.ndarray,
      kappa: float,
  ) -> jnp.ndarray:
      r"""Calculates return-weighted average over all trajectories.
  
      This will calculate the return-weighted average over a set of trajectories as
      defined on l.17 of Alg. 2 in the MBOP paper:
      [https://arxiv.org/abs/2008.05556].
  
      Note: Clipping will be performed for `cum_reward` values > 80 to avoid NaNs.
  
      Args:
        action_trajectories: (n_trajectories, horizon, action_dim) tensor of action
          trajectories, corresponds to `A` in Alg. 2.
        cum_reward: (n_trajectories) vector of corresponding cumulative rewards
          (returns) for each trajectory. Corresponds to `\mathcal{R}` in Alg. 2.
        kappa: `\kappa` constant, changes the 'peakiness' of the exponential
          averaging.
  
      Returns:
        Single action trajectory corresponding to the return-weighted average of the
          trajectories.
      """
      # Substract maximum reward to avoid NaNs:
      cum_reward = cum_reward - cum_reward.max()
      # Remove the batch dimension of cum_reward allows for an implicit broadcast in
      # jnp.average:
      exp_cum_reward = jnp.exp(kappa * jnp.squeeze(cum_reward))
      return jnp.average(
          action_trajectories, weights=exp_cum_reward, axis=0
      )
  
file_path: acme/acme/agents/jax/mbop/mppi.py
imports: import dataclasses
  import functools
  from typing import Callable, Optional
  from acme import specs
  from acme.agents.jax.mbop import models
  from acme.jax import networks
  import jax
  from jax import random
  import jax.numpy
label_desc: Use image node_id 1 for calling a global function with example usage: return_weighted_average(action_trajectories, cum_reward, kappa) and returns: jnp