code: def return_top_k_average(
      action_trajectories: jnp.ndarray,
      cum_reward: jnp.ndarray,
      k: int = 10,
  ) -> jnp.ndarray:
      r"""Calculates the top-k average over all trajectories.
  
      This will calculate the top-k average over a set of trajectories as
      defined in the POIR Paper:
  
      Note: top-k average is more numerically stable than the weighted average.
  
      Args:
        action_trajectories: (n_trajectories, horizon, action_dim) tensor of action
          trajectories.
        cum_reward: (n_trajectories) vector of corresponding cumulative rewards
          (returns) for each trajectory.
        k: the number of trajectories to average.
  
      Returns:
        Single action trajectory corresponding to the average of the k best
          trajectories.
      """
      top_k_trajectories = action_trajectories[
          jnp.argsort(jnp.squeeze(cum_reward))[-int(k) :]
      ]
      return jnp.mean(top_k_trajectories, 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 2 for calling a global function with example usage: return_top_k_average(action_trajectories, cum_reward, k) and returns: jnp