code: def get_actor_core(
      networks: impala_networks.IMPALANetworks,
      environment_spec: specs.EnvironmentSpec,
      evaluation: bool = False,
  ) -> ImpalaPolicy:
      """Creates an Impala ActorCore."""
  
      dummy_logits = jnp.zeros(environment_spec.actions.num_values)
  
      def init(
          rng: jax_types.PRNGKey,
      ) -> ImpalaActorState[actor_core_lib.RecurrentState]:
          rng, init_state_rng = jax.random.split(rng)
          initial_state = networks.init_recurrent_state(
              init_state_rng, None
          )
          return ImpalaActorState(
              rng=rng,
              logits=dummy_logits,
              recurrent_state=initial_state,
              prev_recurrent_state=initial_state,
          )
  
      def select_action(
          params: networks_lib.Params,
          observation: networks_lib.Observation,
          state: ImpalaActorState[actor_core_lib.RecurrentState],
      ) -> Tuple[
          networks_lib.Action,
          ImpalaActorState[actor_core_lib.RecurrentState],
      ]:
          rng, apply_rng, policy_rng = jax.random.split(state.rng, 3)
          (logits, _), new_recurrent_state = networks.apply(
              params,
              apply_rng,
              observation,
              state.recurrent_state,
          )
  
          if evaluation:
              action = jnp.argmax(logits, axis=-1)
          else:
              action = jax.random.categorical(policy_rng, logits)
  
          return action, ImpalaActorState(
              rng=rng,
              logits=logits,
              recurrent_state=new_recurrent_state,
              prev_recurrent_state=state.recurrent_state,
          )
  
      def get_extras(
          state: ImpalaActorState[actor_core_lib.RecurrentState],
      ) -> ImpalaExtras:
          return {
              "logits": state.logits,
              "core_state": state.prev_recurrent_state,
          }
  
      return actor_core_lib.ActorCore(
          init=init, select_action=select_action, get_extras=get_extras
      )
  
file_path: acme/acme/agents/jax/impala/acting.py
imports: from typing import Generic, Mapping, Tuple
  from acme import specs
  from acme.agents.jax import actor_core
  from acme.agents.jax.impala import networks
  from acme.jax import networks
  from acme.jax import types
  import chex
  import jax
  import jax.numpy
label_desc: Use image node_id 1 for calling a global function with example usage: get_actor_core(networks, environment_spec, evaluation) and returns: actor_core_lib, ImpalaActorState, action,  ImpalaActorState, dict