code: def select_action(
      network_def,
      params,
      state,
      rng,
      num_actions,
      eval_mode,
      epsilon_eval,
      epsilon_train,
      epsilon_decay_period,
      training_steps,
      min_replay_history,
      epsilon_fn,
      support,
  ):
      """Select an action from the set of available actions."""
      epsilon = jnp.where(
          eval_mode,
          epsilon_eval,
          epsilon_fn(
              epsilon_decay_period,
              training_steps,
              min_replay_history,
              epsilon_train,
          ),
      )
  
      rng, rng1, rng2, rng3 = jax.random.split(rng, num=4)
  
      @functools.partial(jax.vmap, in_axes=(0, 0), axis_name="batch")
      def q_function(state, key):
          q_values = network_def.apply(
              params,
              state,
              key=key,
              eval_mode=eval_mode,
              support=support,
          ).q_values
          return q_values
  
      q_values = q_function(
          state, jax.random.split(rng2, state.shape[0])
      )
  
      best_actions = jnp.argmax(q_values, axis=-1)
      random_actions = jax.random.randint(
          rng3, (state.shape[0],), 0, num_actions
      )
      p = jax.random.uniform(rng1, shape=(state.shape[0],))
      return rng, jnp.where(p <= epsilon, random_actions, best_actions)
  
imports: import copy
  import functools
  from absl import logging
  from dopamine.jax import networks
  from dopamine.jax.agents.full_rainbow import full_rainbow_agent
  import gin
  import jax
  import jax.numpy
  import numpy
  import tensorflow
label_desc: Use image node_id 1 for calling a global function with example usage: select_action(network_def, params, state, rng, num_actions, eval_mode, epsilon_eval, epsilon_train, epsilon_decay_period, training_steps, min_replay_history, epsilon_fn, support) and returns: rng,  jnp, q_values