code: def default_sample_fn(
      action_values: networks_lib.NetworkOutput,
      key: types.PRNGKey,
      epsilon: Epsilon,
  ) -> networks_lib.Action:
      return rlax.epsilon_greedy(epsilon).sample(key, action_values)
  
imports: import dataclasses
  from typing import Callable, Optional
  from acme.jax import networks
  from acme.jax import types
  import rlax
label_desc: Use image node_id 1 for calling a global function with example usage: default_sample_fn(action_values, key, epsilon) and returns: rlax