code: def make_ensemble_policy_prior(
      policy_prior_network: mbop_networks.PolicyPriorNetwork,
      spec: specs.EnvironmentSpec,
      use_round_robin: bool = True,
  ) -> PolicyPrior:
      """Creates an ensemble policy prior from its network.
  
      Args:
        policy_prior_network: The policy prior network.
        spec: Environment specification.
        use_round_robin: Whether to use round robin or mean to calculate the policy
          prior over the ensemble members.
  
      Returns:
        A policy prior.
      """
  
      def _policy_prior(
          params: networks.Params,
          key: networks.PRNGKey,
          observation_t: networks.Observation,
          action_tm1: networks.Action,
      ) -> networks.Action:
          # Regressor policies are deterministic.
          del key
          apply_fn = (
              ensemble.apply_round_robin
              if use_round_robin
              else ensemble.apply_mean
          )
          return apply_fn(
              policy_prior_network.apply,
              params,
              observation_t=observation_t,
              action_tm1=action_tm1,
          )
  
      dummy_action = utils.zeros_like(spec.actions)
      dummy_action = utils.add_batch_dim(dummy_action)
  
      return feed_forward_policy_prior_to_actor_core(
          _policy_prior, dummy_action
      )
  
file_path: acme/acme/agents/jax/mbop/models.py
imports: import functools
  from typing import Callable, Generic, Optional, Tuple
  from acme import specs
  from acme.agents.jax import actor_core
  from acme.agents.jax.mbop import ensemble
  from acme.agents.jax.mbop import networks
  from acme.jax import networks
  from acme.jax import utils
  import chex
  import jax
label_desc: Use image node_id 3 for calling a global function with example usage: make_ensemble_policy_prior(policy_prior_network, spec, use_round_robin) and returns: feed_forward_policy_prior_to_actor_core, apply_fn