code: def mppi_planner(
      config: MPPIConfig,
      world_model: models.WorldModel,
      policy_prior: models.PolicyPrior,
      n_step_return: models.NStepReturn,
      world_model_params: networks.Params,
      policy_prior_params: networks.Params,
      n_step_return_params: networks.Params,
      random_key: networks.PRNGKey,
      observation: networks.Observation,
      previous_trajectory: jnp.ndarray,
  ) -> jnp.ndarray:
      """MPPI-extended trajectory optimizer.
  
      This implements the trajectory optimizer described in MBOP
      [https://arxiv.org/abs/2008.05556] (Alg. 2) which is an extended version of
      MPPI that adds support for arbitrary sampling distributions and extends the
      return horizon using an approximate model of returns.  There are a couple
      notation changes for readability:
      A -> action_trajectories
      T -> action_trajectory
  
      If the horizon is set to 0, the planner will simply call the policy_prior
      and average the action over the ensemble heads.
  
      Args:
        config: Base configuration parameters of MPPI.
        world_model: Corresponds to `f_m(s_t, a_t)_s` in Alg. 2.
        policy_prior: Corresponds to `f_b(s_t, a_tm1)` in Alg. 2.
        n_step_return: Corresponds to `f_R(s_t, a_t)` in Alg. 2.
        world_model_params: Parameters for world model.
        policy_prior_params: Parameters for policy prior.
        n_step_return_params: Parameters for n_step return.
        random_key: JAX random key seed.
        observation: Normalized current observation from the environment, `s` in
          Alg. 2.
        previous_trajectory: Normalized previous action trajectory. `T` in Alg 2.
          Shape is [n_trajectories, horizon, action_dims].
  
      Returns:
        jnp.ndarray: Average action trajectory of shape [horizon, action_dims].
      """
      action_trajectory_tm1 = previous_trajectory
      policy_prior_state = policy_prior.init(random_key)
  
      # Broadcast so that we have n_trajectories copies of each:
      observation_t = jax.tree_map(
          functools.partial(_repeat_n, config.n_trajectories),
          observation,
      )
      action_tm1 = jnp.broadcast_to(
          action_trajectory_tm1[0],
          (config.n_trajectories,) + action_trajectory_tm1[0].shape,
      )
  
      if config.previous_trajectory_clip is not None:
          action_tm1 = jnp.clip(
              action_tm1,
              a_min=-config.previous_trajectory_clip,
              a_max=config.previous_trajectory_clip,
          )
  
      # First check if planning is unnecessary:
      if config.horizon == 0:
          if hasattr(policy_prior_state, "action_tm1"):
              policy_prior_state = policy_prior_state.replace(
                  action_tm1=action_tm1
              )
          action_set, _ = policy_prior.select_action(
              policy_prior_params, observation_t, policy_prior_state
          )
          # Need to re-create an action trajectory from a single action.
          return jnp.broadcast_to(
              jnp.mean(action_set, axis=0), (1, action_set.shape[-1])
          )
  
      # Accumulators for returns and trajectories:
      cum_reward = jnp.zeros((config.n_trajectories, 1))
  
      # Generate noise once:
      random_key, noise_key = random.split(random_key)
      action_noise = config.sigma * random.normal(
          noise_key, ((config.horizon,) + action_tm1.shape)
      )
  
      # Initialize empty set of action trajectories for concatenation in loop:
      action_trajectories = jnp.zeros(
          (config.n_trajectories, 0) + action_trajectory_tm1[0].shape
      )
  
      for t in range(config.horizon):
          # Query policy prior for proposed action:
          if hasattr(policy_prior_state, "action_tm1"):
              policy_prior_state = policy_prior_state.replace(
                  action_tm1=action_tm1
              )
          action_t, policy_prior_state = policy_prior.select_action(
              policy_prior_params, observation_t, policy_prior_state
          )
          # Add action noise:
          action_t = action_t + action_noise[t]
          # Mix action with previous trajectory's corresponding action:
          action_t = (
              1 - config.beta
          ) * action_t + config.beta * action_trajectory_tm1[t]
  
          # Query world model to get next observation and reward:
          observation_tp1, reward_t = world_model(
              world_model_params, observation_t, action_t
          )
          cum_reward += reward_t
  
          # Insert actions into trajectory matrix:
          action_trajectories = jnp.concatenate(
              [action_trajectories, jnp.expand_dims(action_t, axis=1)],
              axis=1,
          )
          # Bump variable timesteps for next loop:
          observation_t = observation_tp1
          action_tm1 = action_t
  
      # De-normalize and append the final n_step return prediction:
      n_step_return_t = n_step_return(
          n_step_return_params, observation_t, action_t
      )
      cum_reward += n_step_return_t
  
      # Average the set of `n_trajectories` trajectories into a single trajectory.
      return config.action_aggregation_fn(
          action_trajectories, cum_reward
      )
  
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 5 for calling a global function with example usage: mppi_planner(config, world_model, policy_prior, n_step_return, world_model_params, policy_prior_params, n_step_return_params, random_key, observation, previous_trajectory) and returns: config, jnp