code: def batched_recurrent_to_actor_core(
      recurrent_policy: RecurrentPolicy,
      initial_core_state: RecurrentState,
  ) -> ActorCore[
      SimpleActorCoreRecurrentState[RecurrentState],
      Mapping[str, jnp.ndarray],
  ]:
      """Returns ActorCore for a recurrent policy."""
  
      def select_action(
          params: networks_lib.Params,
          observation: networks_lib.Observation,
          state: SimpleActorCoreRecurrentState[RecurrentState],
      ):
          # TODO(b/161332815): Make JAX Actor work with batched or unbatched inputs.
          rng = state.rng
          rng, policy_rng = jax.random.split(rng)
          observation = utils.add_batch_dim(observation)
          recurrent_state = utils.add_batch_dim(state.recurrent_state)
          action, new_recurrent_state = utils.squeeze_batch_dim(
              recurrent_policy(
                  params, policy_rng, observation, recurrent_state
              )
          )
          return action, SimpleActorCoreRecurrentState(
              rng, new_recurrent_state
          )
  
      initial_core_state = utils.squeeze_batch_dim(initial_core_state)
  
      def init(
          rng: PRNGKey,
      ) -> SimpleActorCoreRecurrentState[RecurrentState]:
          return SimpleActorCoreRecurrentState(rng, initial_core_state)
  
      def get_extras(
          state: SimpleActorCoreRecurrentState[RecurrentState],
      ) -> Mapping[str, jnp.ndarray]:
          return {"core_state": state.recurrent_state}
  
      return ActorCore(
          init=init, select_action=select_action, get_extras=get_extras
      )
  
file_path: acme/acme/agents/jax/actor_core.py
imports: import dataclasses
  from typing import Callable, Generic, Mapping, Tuple, TypeVar, Union
  from acme import types
  from acme.jax import networks
  from acme.jax import utils
  from acme.jax.types import PRNGKey
  import chex
  import jax
  import jax.numpy
label_desc: Use image node_id 4 for calling a global function with example usage: batched_recurrent_to_actor_core(recurrent_policy, initial_core_state) and returns: ActorCore, action,  SimpleActorCoreRecurrentState, SimpleActorCoreRecurrentState, dict