code: def main(_):
      key = jax.random.PRNGKey(FLAGS.seed)
      key_demonstrations, key_learner = jax.random.split(key, 2)
  
      # Create an environment and grab the spec.
      environment = gym_helpers.make_environment(task=FLAGS.env_name)
      environment_spec = specs.make_environment_spec(environment)
  
      # Get a demonstrations dataset with next_actions extra.
      transitions = tfds.get_tfds_dataset(
          FLAGS.dataset_name, FLAGS.num_demonstrations
      )
      double_transitions = rlds.transformations.batch(
          transitions, size=2, shift=1, drop_remainder=True
      )
      transitions = double_transitions.map(_add_next_action_extras)
      demonstrations = tfds.JaxInMemoryRandomSampleIterator(
          transitions,
          key=key_demonstrations,
          batch_size=FLAGS.batch_size,
      )
  
      # Create the networks to optimize.
      networks = crr.make_networks(environment_spec)
  
      # CRR policy loss function.
      policy_loss_coeff_fn = crr.policy_loss_coeff_advantage_exp
  
      # Create the learner.
      learner = crr.CRRLearner(
          networks=networks,
          random_key=key_learner,
          discount=FLAGS.discount,
          target_update_period=FLAGS.target_update_period,
          policy_loss_coeff_fn=policy_loss_coeff_fn,
          iterator=demonstrations,
          policy_optimizer=optax.adam(FLAGS.policy_learning_rate),
          critic_optimizer=optax.adam(FLAGS.critic_learning_rate),
          grad_updates_per_batch=FLAGS.grad_updates_per_batch,
          use_sarsa_target=FLAGS.use_sarsa_target,
      )
  
      def evaluator_network(
          params: hk.Params, key: jax.Array, observation: jax.Array
      ) -> jax.Array:
          dist_params = networks.policy_network.apply(
              params, observation
          )
          return networks.sample_eval(dist_params, key)
  
      actor_core = actor_core_lib.batched_feed_forward_to_actor_core(
          evaluator_network
      )
      variable_client = variable_utils.VariableClient(
          learner, "policy", device="cpu"
      )
      evaluator = actors.GenericActor(
          actor_core, key, variable_client, backend="cpu"
      )
  
      eval_loop = acme.EnvironmentLoop(
          environment=environment,
          actor=evaluator,
          logger=loggers.TerminalLogger("evaluation", time_delta=0.0),
      )
  
      # Run the environment loop.
      while True:
          for _ in range(FLAGS.evaluate_every):
              learner.step()
          eval_loop.run(FLAGS.evaluation_episodes)
  
file_path: acme/examples/offline/run_crr_jax.py
imports: from absl import app
  from absl import flags
  import acme
  from acme import specs
  from acme.agents.jax import actor_core
  from acme.agents.jax import actors
  from acme.agents.jax import crr
  from acme.datasets import tfds
  from acme.examples.offline import helpers
  from acme.jax import variable_utils
  from acme.types import Transition
  from acme.utils import loggers
  import haiku
  import jax
  import optax
  import rlds
label_desc: Use image node_id 2 for calling a global function with example usage: main(_) and returns: networks