code: def get_policy(
      policy_network: networks_lib.FeedForwardNetwork,
      normalization_apply_fn,
  ) -> actor_core_lib.FeedForwardPolicy:
      """Returns a function that computes actions."""
  
      def apply(
          params: networks_lib.Params,
          key: networks_lib.PRNGKey,
          obs: networks_lib.Observation,
      ) -> Tuple[networks_lib.Action, Dict[str, jnp.ndarray]]:
          del key
          params_key, policy_params, normalization_params = params
          normalized_obs = normalization_apply_fn(
              obs, normalization_params
          )
          action = policy_network.apply(policy_params, normalized_obs)
          return action, {
              "params_key": jax.tree_map(
                  lambda x: jnp.expand_dims(x, axis=0), params_key
              )
          }
  
      return apply
  
imports: from typing import Dict, Iterator, List, Optional, Tuple
  import acme
  from acme import adders
  from acme import core
  from acme import specs
  from acme.adders import reverb
  from acme.agents.jax import actor_core
  from acme.agents.jax import actors
  from acme.agents.jax import builders
  from acme.agents.jax.ars import config
  from acme.agents.jax.ars import learning
  from acme.jax import networks
  from acme.jax import running_statistics
  from acme.jax import utils
  from acme.jax import variable_utils
  from acme.utils import counting
  from acme.utils import loggers
  import jax
  import jax.numpy
  import numpy
  import reverb
label_desc: Use image node_id 1 for calling a global function with example usage: get_policy(policy_network, normalization_apply_fn) and returns: apply, action,  dict